LEMUR and Mean Centering for Late-Interaction Retrieval in txtai

David Mezzetti

davidmezzetti

Follow

NeuML

Late-interaction models preserve a useful level of detail: instead of collapsing a query or document into one embedding immediately, they retain a vector for each token. A MaxSim score compares every query token with the document tokens, keeps the strongest match for each query token, and sums those matches. The tradeoff is operational. A multi-vector representation does not fit as naturally into the fixed-vector indexes used by a conventional dense retrieval pipeline.

LEMUR closes that gap in txtai. It learns a fixed-dimensional encoding of the multi-vector representation, so a late-interaction model can use the same kind of vector index as a standard embedding model. A second change, configurable mean centering, addresses a different failure mode found while testing LEMUR with modern LateOn models: token vectors can be so anisotropic that the fixed-dimensional encoder has very little useful variation to preserve.

This article walks through what the two pieces do, the measurements that changed the design, and the configuration that shipped. The benchmark results are intentionally narrow: one late-interaction model, three BEIR datasets, one machine, and exact search.

From token vectors to one searchable vector

LEMUR stands for Learned Multi-Vector Retrieval. In txtai, its artifact contains a feature encoder, output-normalization statistics, and a sample of token vectors. At inference time, a query becomes the sum of its learned token features. A document becomes a set of ordinary least squares weights over the stored sample. Their fixed-vector inner product approximates the original late-interaction score.

The training pipeline starts from corpus text. It encodes each target document with the model's data encoder, builds standardized MaxSim targets, learns the feature map, and saves an inference artifact as config.json plus model.safetensors. The artifact is corpus-specific, and it must be trained before the txtai embeddings index loads it.

One training detail mattered more than expected: the distribution used to learn the feature map. Late-interaction models can encode queries and documents differently. On nfcorpus, learning from data-encoder token vectors produced a trained MLP below even the untrained ELM fallback. Learning from query-encoder token vectors recovered most of the improvement, and validation-based epoch selection added a smaller gain.

This four-row ablation came from the earlier CPU run. The GPU matrix below used query-encoder learn tokens throughout.

In this ablation, the learn distribution accounted for 93% of the measured gain from the data-token run to the final run; validation selection accounted for the remaining 7%. That is why LemurTrainer defaults learncategory to "query", while still allowing a caller to choose "data" or provide a separate learn iterable.

The trained MLP is the quality-oriented path. txtai also requires the training choice to be explicit: epochs=100 selects the documented MLP setting, while epochs=0 selects deterministic random ELM features as a lower-cost fallback.

What the exact-search benchmark showed

The aligned matrix was built using colbert-ir/colbertv2.0 with an NVIDIA GeForce RTX 4080 SUPER, torch 2.13.0+cu130, and exact Faiss IDMap,Flat search. The table compares the default 10,240-dimensional MUVERA encoding, MUVERA reduced to the same 2,048-dimensional budget, and the trained 2,048-dimensional LEMUR encoding.

At equal vector size, LEMUR improved NDCG@10 over matched-size MUVERA by 56.6% on nfcorpus, 49.4% on scifact, and 61.9% on arguana. Against MUVERA's full default vector, the gains were 8.4%, 9.8%, and 22.9%. The measured LEMUR indexes used one fifth of the storage of the default MUVERA indexes because both the vector width and the recorded index size scaled from 10,240 to 2,048 dimensions.

Those numbers are evidence for a specific configuration, not a general ranking claim. Three of the five datasets were measured in the requested set; fiqa and scidocs were not completed. The matrix covers one model and one machine. It also forced exact search, which matters for deployment.

txtai's Faiss backend uses exact search through 5,000 rows and switches to an IVF index above that threshold. In the measured scifact run, default IVF reduced LEMUR NDCG@10 by 43% relative to exact search, compared with 25% for MUVERA. For a larger LEMUR corpus, pin faiss.components to IDMap,Flat when exact search is practical, or tune IVF for the corpus instead of assuming the default index will preserve the exact-search result.

Why mean centering entered the design

The first LEMUR benchmark used ColBERTv2. Upon moving to lightonai/LateOn, the raw token vectors showed a different geometry. Across a 5,000-pair sample, mean pairwise token cosine was 0.9508 and the MaxSim spread was 0.0559. The article Regularizing ColBERT models to fix efficient ANN methods describes anisotropy and how that was a culprit in poor MUVERA scoring. After collection-mean subtraction and re-normalization, those values became 0.0033 and 0.4772. This is not itself a retrieval metric, but it showed that centering exposed much more directional variation for the fixed-dimensional encoder to use.

The retrieval measurements confirmed the useful part of that signal, but they also ruled out a universal switch. As David Mezzetti put it, centering "seemed to hurt some models while helping others." He asked for a configurable center parameter and suggested enabling it by default when a loaded model has more than one linear layer.

The base LateOn model loaded five linear layers, while ColBERTv2 loaded one. Centering strongly helped the LateOn matrix. Three scopes were compared: a stored collection mean, a mean per document, and a mean over the current model batch. These CUDA runs used torch 2.13.0+cu130 and 32 texts per benchmark encoding batch. Batch scope was the strongest simple scope in all four base-LateOn cells below.

Per-document centering was mixed: it trailed the collection mean in both nfcorpus cells, was essentially tied for scifact LEMUR, and improved scifact MUVERA. Batch scope avoided a corpus pass and a stored mean while leading the collection result in this matrix. David chose that tradeoff: he would "go with batch as the default" when centering is explicitly enabled.

The one-layer check is why the automatic rule is described as conservative, not universal. On ColBERTv2, collection centering was slightly helpful or flat with LEMUR but regressed MUVERA on both measured datasets. Layer count alone did not explain every outcome; encoder choice mattered too. The shipped default therefore enables batch centering when the model loads more than one torch.nn.Linear layer, while zero- and one-layer models retain their previous output. A caller can override that default in either direction.

Centering happens after token-vector normalization and before MUVERA or LEMUR, followed by another normalization. center: true selects batch scope. A dictionary can select document, batch, or collection; collection scope accepts either an inline mean or a Safetensors file containing center.mean. center: false disables the operation.

Training and loading a centered LEMUR artifact

The latest published txtai release is still v9.12.0, while both merged changes target v9.13.0. For the pre-release source, use pip install "git+https://github.com/neuml/txtai.git@master" in an isolated environment.

Here is a lightweight ELM example. The important part is consistency: if the index will center token vectors, pass the same vectors setting to the trainer so the LEMUR artifact is fitted on that representation.

from txtai.pipeline import LemurTrainer

corpus = [
    "Late interaction compares token embeddings.",
    "Dense indexes search fixed dimensional vectors.",
    "Mean centering changes token-vector geometry.",
]

vectors = {"center": True}

LemurTrainer()(
    "neuml/colbert-bert-tiny",
    corpus,
    "lemur-model",
    gpu=False,
    vectors=vectors,
    epochs=0,
)

Load the artifact through the embeddings vector configuration. This example pins exact Faiss search so its indexing behavior is explicit.

from txtai import Embeddings

embeddings = Embeddings(
    {
        "path": "neuml/colbert-bert-tiny",
        "content": True,
        "vectors": {
            "center": True,
            "lemur": {"path": "lemur-model"},
        },
        "faiss": {"components": "IDMap,Flat"},
    }
)

embeddings.index(corpus)
print(embeddings.search("dense indexes"))

For a quality-oriented artifact, replace the ELM choice with the documented MLP setting and use a held-out validation split. The artifact remains a separate train-first step; after that, txtai can load it from a local directory or a Hugging Face Hub path.

A follow-up PR moved training consolidation and progress feedback into LemurTrainer. It does not change the load-time API shown here.

Credits and limits

The txtai implementation is based on LEMUR by Elias Jääsaari, Ville Hyvönen, and Teemu Roos (ICML '26 paper). David shaped the integration through the review threads for #1164 and #1168, including the configuration boundary, batch default, and Safetensors path for a stored mean. He also tagged Elias Jääsaari, Raphael Sourty, and Antoine Chaffin because these changes build on their work.

The practical result is a pair of controls rather than a promise that one setup wins everywhere. LEMUR supplies a compact learned fixed vector for late interaction. Centering supplies an explicit way to repair anisotropic token geometry when the model needs it. The exact-search results are encouraging, but broader model and dataset coverage, plus tuned approximate-index measurements, remain the next evidence to collect.