[Feature Request] Support for multi-stage retriever and re-ranker in OpenSearch to use late interaction embedding models like ColBert, ColPali etc.

[prasadnu]

Is your feature request related to a problem? Please describe

ColBert, Colpali like late interaction models provide multiple token level vectors for a document/query and multiple patch level vectors for an image. These patch/token level vectors when subjected to similarity search should be scored using MaxSim operation based on the late interaction mechanism. An excellent blog from Weviate on late interaction models, MaxSim etc.

These models when used for search provides better relevance due to the interaction of every query and document token which is not possible with ordinary vector search based on single vectors for the document and query, where similarity is calculated as Cosine(1 doc vector, 1 query vector). Also, late interaction models allows interpretability of vector search as in which token (for a document) or patch (for an image) is causing the document to appear in the result and you can highlight those tokens/patches for your results (semantic highlighting). Examples,

Describe the solution you'd like

The workflow can be depicted in 2 ways,

Type 1. Pooling of query token vectors:

• During ingestion, the multiple vectors per document are stored using OpenSearch nested vector fields
• Phase 1 (Retrieval): On query time, the multiple token vectors of the query will be converted into a single query vector (Max/Min Pooling) and subjected to ANN search against the document nested vector field.
• The unique documents fetched (de-duplication already handled by OpenSearch nested vector field) should enter the ranking phase.
• Phase 2 (Ranking):In the ranking phase, all the token level vectors of both the documents and query are considered to calculate the score for the whole document-query using MaxSim formula.

Type 2. No pooling done on query token vectors:

• During ingestion, the multiple vectors per document are stored using OpenSearch nested vector fields
• Phase 1 (Retrieval): On query time, each token vector of the query will be subjected to ANN search against the document nested vector field.
• The unique documents fetched (de-duplication already handled by OpenSearch nested vector field) enter the ranking phase. These documents are the union list of all candidates that resulted from ANN search of every token query vector.
• Phase 2 (Ranking): In the ranking phase, all the token level vectors of both the documents and query are considered to calculate the score for the whole document-query using MaxSim formula.

Just for similar feature comparison, the type 1 multi-phase retrieval is supported in a ElasticSearch using a new rescorer retriever query type and maxsim ranking is supported by a script_score function, maxSimDotProduct (MaxSim over dot product of full precision vectors) and maxSimInvHamming (MaxSim over Hamming distance of bit vector dimensions) over multiple vectors per document stored in a field type called rank_vectors

Related component

Search

Describe alternatives you've considered

ClientSide MaxSim scoring
OpenSearch Hybrid query to support upto ANN of 5 query token vectors against the document multi vectors in nested vector field

Additional context

No response

[rohangpatil]

Looking forward to this!

[msfroh]

@ylwu-amzn Is this the kind of thing your team would work on?

Would a search pipeline response processor be part of a valid solution?

[heemin32]

For Ingestion and Phase1(Retrieval), we might need to use average value of token embedding to avoid scalability issue. We could implement ingest processor to support late interaction embedding models in neural plugin to make life easier.

For Phase2, we might need new query type or script score to calculate MaxSim between two list of vectors: query token vectors, document token vectors.

[bouhajer]

This would be super to have in OpenSearch! Seeing many requests on the same topic, so looking forward to it <3

[ylwu-amzn]

@ylwu-amzn Is this the kind of thing your team would work on?

Would a search pipeline response processor be part of a valid solution?

Thanks @msfroh for involving me. I think one option is to extend ML inference processor to support this.

@prasadnu comes up two ways

For Retrieval phase:

• First way "Pooling of query token vectors", we can build a ML inference request processor to translate query into token vectors, then do pooling.
• Second way, no need to do query token vector pooling, just one ML inference request processor is enough.

For ranking phase: agree with @heemin32 , we may need new query type, or build new search response processor to calculate MaxSim.

[q-andy]

Great uses of ml inference search/ingest processors pointed out by @ylwu-amzn. For ingest side, this may be possible using ML inference processor too, but processing the vectors from the model output is tricky. If we index into vector store using the average value of token embeddings suggested by @heemin32, we first have to calculate the average. For initial POC it may be doable with a combination of script/foreach processors, but performance would not be good, a specialized processor would be preferable.

For retrieval side, in addition to reranking results, we'd also want a mechanism to take advantage of the token-level relevance. One of the advantages of these models is we can get token/patch level scores within a document, which means we'd also want a way to label and present that to users.

For labelling text we have the highlighter feature, but that doesn't fit this use case perfectly, since highlighting is a binary label. With these models we could potentially score tokens individually with continuous values.

@prasadnu @rohangpatil @bouhajer Thanks for expressing interest in this! Out of curiosity, could explain your potential use cases and what aspects of this strategy you're most excited to see?

[heemin32]

I was trying to identify the minimal gaps that need to be addressed for this to work. If we assume customers provide token vectors themselves—similar to how they currently use the knn field without an embedding processor—we would need to:

• Store the average of the token vectors in the knn field.
• Store the full token vectors in some format, list of vectors. (this part is currently missing).

Once vectors are ingested, we can perform a knn search using the averaged token vector, and then rescore the retrieved documents using the max similarity between the query's token vectors and the document's stored token vectors (this is also missing today).

There may be other features to consider, but I believe these two are the primary missing components needed to support this functionality.