{"slug": "token-pooling-aware-training-for-multi-vector-embedding-models", "title": "Token Pooling-Aware Training for Multi-Vector Embedding Models", "summary": "Token pooling-aware training raises a multi-vector embedding model's retention of unpooled NanoBEIR nDCG@10 to 97.8% at pool_factor=4, up from 94.6% for a normally trained model, according to a Hugging Face blog post by tomaarsen. The technique trains the encoder to preserve information when similar token vectors are merged via hierarchical Ward clustering, cutting document vectors 4x; the post cites Stefan Josef's Learn to Pool work and LightOn's pooling regularization study as prior research. The post notes that 200 vectors of 128 dimensions at 4 bits per residual coordinate require 12.8 KB per document, or 12.8 GB across a million documents, versus 3.2 GB when pooled to 50 vectors.", "body_md": "[Feature Extraction •  0.1B • Updated   •  88](https://huggingface.co/tomaarsen/mve-ModernBERT-base-token-pooling)    \n\n# \n\t\tToken Pooling-Aware Training for Multi-Vector Embedding Models\n\t\n\n [Update on GitHub](https://github.com/huggingface/blog/blob/main/token-pooling.md)\n\nMulti-vector embedding models give each query token its own opportunity to find a match in a document. This makes them wonderfully expressive for search, but those document vectors add up quickly. A corpus with millions of documents can contain hundreds of millions of vectors!\n\nToken pooling makes these representations smaller by merging similar token vectors. The more aggressively I compress them, however, the easier it becomes to lose something useful. Pooling-aware training teaches the encoder to preserve the information that matters when its vectors are merged, helping you retrieve better from the same storage budget.\n\nIn this post, I'll explain how pooling works, train a model for compressed retrieval with Sentence Transformers, and explore the resulting quality and storage tradeoff. With `pool_factor=4`, my normally trained model retains 94.6% of its unpooled NanoBEIR nDCG@10, with a **4x** reduction in document vectors. Token pooling-aware training raises that retention to 97.8% at the same storage budget! Let's see how.\n\nThis post builds on pooling-aware training studied in Stefan Josef's [Learn to Pool](https://stefan-jo.github.io/learn-to-pool/downloads/paper.pdf) and [LightOn's pooling regularization study](https://huggingface.co/blog/lightonai/lateon-hpool-regularization), which also explores targeted training and adaptive budgets. Here, I'll focus on understanding the method and putting it to work with a short, reproducible Sentence Transformers example.\n\n## \n\t\tTable of Contents\n\t\n\n## \n\t\tWhy Multi-Vector Retrieval Needs So Many Vectors\n\t\n\nA single-vector embedding model compresses an entire document into one vector. [Multi-vector models](https://huggingface.co/blog/multi-vector-encoder) such as ColBERT instead keep a collection of contextual token vectors, allowing different parts of a query to match different parts of a document.\n\nFor example, a query about a laptop with a long battery life can find evidence about the laptop in one part of a document and battery endurance in another. With MaxSim scoring, each query vector selects its best matching document vector, and those scores are added together:\n\nEven with the 2-bit or 4-bit residual quantization supported by [Fast-PLAID](https://github.com/lightonai/fast-plaid), storing hundreds of vectors per document adds up. At 4 bits per residual coordinate, 200 vectors of 128 dimensions require 12.8 KB of quantized residuals per document, or 12.8 GB across a million documents. Pooling each document to 50 vectors reduces that to 3.2 GB, at the same dimensionality and precision.\n\nThe useful observation is that many document vectors can be similar. Several tokens might describe closely related concepts, creating an opportunity to represent them together.\n\n## \n\t\tToken Pooling\n\t\n\n[HierarchicalTokenPooling](https://sbert.net/docs/package_reference/multi_vector_encoder/modules.html#hierarchicaltokenpooling) groups similar token embeddings using [Ward hierarchical clustering](https://scikit-learn.org/stable/modules/clustering.html#hierarchical-clustering). Starting with one group per token vector, it repeatedly merges the two groups that cause the smallest increase in how spread out the vectors are within their groups. This spread is measured by the sum of squared distances from each vector to its group's mean, so Ward's method favors compact groups of similar vectors.\n\nThese merges build a hierarchy of progressively larger groups, letting us choose a level of compression that fits our budget. We then represent each group by its mean and normalize the resulting vectors again, ready for MaxSim scoring.\n\nYou can make compression proportional to document length:\n\n``` python\nfrom sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling\n\npooling = HierarchicalTokenPooling(pool_factor=4)\n```\n\nWith `pool_factor=4`, documents retain approximately one quarter of their original vectors, so longer documents can keep proportionally more capacity. By default, the first token is protected from merging and pooling applies only to document embeddings, leaving query embeddings unchanged.\n\n### \n\t\tChoosing a Pooling Budget\n\t\n\n`HierarchicalTokenPooling` offers three alternative controls:\n\n| Setting | What it expresses | \n|---|---|\n| `max_tokens=16` | A maximum of 16 vectors per document, including protected tokens. | \n| `pool_factor=4` | Approximately one quarter of each document's original vectors. | \n| `distance_threshold=...` | A shared Ward merge-distance threshold, allowing vector counts to vary with embedding geometry. | \n\nThis distinction matters when document lengths vary. With `pool_factor=4`, a longer document can retain approximately one quarter of its vectors, allowing its representation to grow with its content. With `distance_threshold`, the vector count can grow when more distinct embeddings need to be preserved.\n\n### \n\t\tCalibrating Adaptive Pooling\n\t\n\nAdaptive pooling uses a shared distance threshold to merge similar embeddings more aggressively while preserving more vectors for documents with distinct content. It can be used with any multi-vector model, *including one trained without a pooling-aware objective*.\n\nRather than choosing a threshold manually, you can calibrate it for a target average vector count. Given a representative sample of your corpus in `calibration_documents`:\n\n``` python\nfrom sentence_transformers import MultiVectorEncoder\nfrom sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling\n\nmodel = MultiVectorEncoder(\"lightonai/LateOn\")\ncalibration_embeddings = model.encode_document(calibration_documents)\nthreshold = HierarchicalTokenPooling.calibrate_distance_threshold(calibration_embeddings, target_tokens=16)\npooling = HierarchicalTokenPooling(distance_threshold=threshold)\n\n# Apply pooling, targeting an average of 16 vectors per document\ndocument_embeddings = model.encode_document(documents_to_index, token_pooling=pooling)\n```\n\nCalibration only needs document embeddings and returns a threshold you can reuse when indexing. The target is an average, so individual documents can receive more or fewer than 16 vectors. The achieved average can also change on another corpus.\n\n## \n\t\tToken Pooling-Aware Training\n\t\n\nPooling can already be useful on an ordinary multi-vector model. The difficulty appears when a document contains a distinction that matters for retrieval, but merging its vectors makes that distinction harder to recover. The original training objective never required the model to work with these compressed representations. But, you can teach it to do exactly that.\n\nIf you've used [Matryoshka embeddings](https://huggingface.co/blog/matryoshka), the training principle will feel familiar. Matryoshka applies a retrieval objective at several embedding widths, teaching smaller prefixes to remain useful. Pooling-aware training applies the objective to representations with fewer token vectors.\n\nFor my example, each training batch contributes two losses:\n\n1. Retrieval using the original document embeddings.\n2. Retrieval after compressing documents to approximately one quarter of their original vectors.\n\nI average them:\n\nThe first term teaches the model to retrieve with its full representations. The second teaches it to arrange its token embeddings so that useful matches survive pooling. Both objectives update the same encoder, giving it feedback about retrieval before and after compression.\n\nWard clustering makes discrete decisions about which tokens belong together. During training, `HierarchicalTokenPooling` uses a straight-through estimator to pass gradients from pooled vectors back to their member tokens, without differentiating the cluster assignments. The encoder can then learn from retrieval errors made with its pooled representations.\n\n### \n\t\tA Loss Wrapper in Sentence Transformers\n\t\n\n[`MultiVectorTokenPoolingLoss`](https://sbert.net/docs/package_reference/multi_vector_encoder/losses.html#multivectortokenpoolingloss) wraps an existing multi-vector retrieval loss and applies it using your chosen poolers:\n\n``` python\nfrom sentence_transformers import MultiVectorEncoder\nfrom sentence_transformers.multi_vector_encoder.losses import (\n    CachedMultiVectorMultipleNegativesRankingLoss,\n    MultiVectorTokenPoolingLoss,\n)\nfrom sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling\n\nmodel = MultiVectorEncoder(\"answerdotai/ModernBERT-base\", similarity_fn_name=\"maxsim\")\npooling = HierarchicalTokenPooling(pool_factor=4)\n\nbase_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64)\nloss = MultiVectorTokenPoolingLoss(\n    model,\n    base_loss,\n    token_poolings=[None, pooling],\n    weights=[0.5, 0.5],\n)\n```\n\n`CachedMultiVectorMultipleNegativesRankingLoss` supplies the retrieval objective. It teaches each query to score its relevant document above the negative documents in the batch. By default, it uses `model.similarity`, using the model's similarity setting (MaxSim in this example) and kernel settings. A model configured with `similarity_fn_name=\"meanmaxsim\"` would also use MeanMaxSim during training.\n\nGradient caching allows the encoder to process small mini-batches while retaining a larger set of in-batch negatives. My training script uses mini-batches of 64 with a training batch size of 256.\n\n`MultiVectorTokenPoolingLoss` adds the compression objective. Here, `None` means that one variant uses the original embeddings, while `pooling` supplies the factor-4 variant. The two weights average their losses. Pooling variants reuse the computed token embeddings, so each additional pooler adds clustering and scoring work without requiring a separate encoder pass for that variant.\n\nThe poolers can also be custom `BaseTokenPooling` implementations that preserve gradients through their input embeddings. You can adapt the training objective to the representation you want to deploy.\n\nReplace the factor-4 pooler with your chosen deployment configuration, or include several poolers and choose weights for their losses.\n\nThe complete [training script](https://github.com/huggingface/sentence-transformers/blob/main/examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py) includes the dataset, trainer, validation loss, and small retrieval evaluations during training. From a checkout of the repository, run both variants locally with [uv](https://docs.astral.sh/uv/getting-started/installation/):\n\n```\nuv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py --baseline\nuv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py\n```\n\nThe script declares its training and kernel dependencies, so uv installs them for you. Compatible kernels are enabled automatically.\n\nYou can also train on a cloud GPU with [Hugging Face Jobs](https://huggingface.co/docs/hub/jobs-quickstart). Install the CLI and log in with a token that can run Jobs and upload models:\n\n```\npip install -U huggingface_hub\nhf auth login\n\nhf jobs uv run \\\n    --flavor a10g-small \\\n    --timeout 4h \\\n    --secrets HF_TOKEN \\\n    -- https://raw.githubusercontent.com/huggingface/sentence-transformers/main/examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py\n```\n\nThis runs the pooling-aware variant on an A10G GPU with a four-hour time limit. Passing `HF_TOKEN` as a secret lets the script upload the trained model to your Hub account when training finishes. To run the ordinary-loss baseline as well, repeat the command with `--baseline` after the script URL.\n\n### \n\t\tOptional Kernel Acceleration\n\t\n\nYou can accelerate MaxSim scoring, Ward clustering, and supported backbone operations with [kernels](https://sbert.net/docs/multi_vector_encoder/usage/kernels.html). Install the training and kernel dependencies:\n\n```\npip install -U \"sentence-transformers[train,kernels]\"\n```\n\nIn the training example above, replace the model and pooler initialization before constructing the loss:\n\n```\nmodel = MultiVectorEncoder(\n    \"answerdotai/ModernBERT-base\",\n    device=\"cuda\",\n    similarity_fn_name=\"maxsim\",\n    use_kernels=True,\n)\npooling = HierarchicalTokenPooling(pool_factor=4)\n```\n\nHierarchical poolers inherit the model's kernel settings by default, so you only need to enable kernels on the model.\n\nSentence Transformers selects compatible kernels for your setup and uses the existing implementation wherever a replacement isn't available. For example, the [default MaxSim kernel](https://huggingface.co/kernels/sentence-transformers/late-interaction-kernels) scores queries against documents without building the full token similarity matrix in memory, but it does require CUDA and Triton.\n\nThe default [Ward clustering kernel](https://huggingface.co/kernels/sentence-transformers/ward-clustering) considerably speeds up hierarchical clustering by running on CUDA instead of the CPU. We're talking about 30-150x faster than SciPy depending on the setting.\n\n## \n\t\tUsing the Trained Model\n\t\n\nLet's load the pooling-aware model and retrieve with compressed document embeddings. Kernel settings aren't saved with the model, so add `use_kernels=True` when loading if you'd like to use them for inference too:\n\n``` python\nfrom sentence_transformers import MultiVectorEncoder\nfrom sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling\n\nmodel = MultiVectorEncoder(\"tomaarsen/mve-ModernBERT-base-token-pooling\")\npooling = HierarchicalTokenPooling(pool_factor=4)\n\nqueries = [\"Which planet is known as the Red Planet?\"]\ndocuments = [\n    \"Venus is often called Earth's twin because of its similar size and proximity.\",\n    \"Mars, known for its reddish appearance, is often referred to as the Red Planet.\",\n    \"Jupiter, the largest planet in our solar system, has a prominent red spot.\",\n]\n\nquery_embeddings = model.encode_query(queries)\ndocument_embeddings = model.encode_document(documents, token_pooling=pooling)\n\nprint([len(embedding) for embedding in document_embeddings])\n# [5, 6, 5]\nprint(model.similarity(query_embeddings, document_embeddings))\n# tensor([[16.0037, 25.9879, 22.8159]], device='cuda:0')\n\n# For comparison, the same model without token pooling:\n# [18, 21, 19]\n# tensor([[16.6115, 27.6060, 24.1766]], device='cuda:0')\n```\n\nThe pooler compresses the document embeddings after the encoder's forward pipeline, reducing the total from 58 vectors to 16 while leaving query embeddings unchanged. Mars still receives the highest similarity score, followed by Jupiter and Venus, so the ranking is preserved despite storing far fewer document vectors.\n\nYou can change `pool_factor` at inference to explore a different storage budget with the same trained model. For repeated use, you can also pass `token_pooling=pooling` when initializing `MultiVectorEncoder` to make it the runtime default. The pooler's `tasks` still control which inputs are compressed (e.g. only documents).\n\n## \n\t\tEvaluation Results\n\t\n\nTo investigate, I trained two models from `answerdotai/ModernBERT-base`, each with a fresh 128-dimensional token projection:\n\n- [tomaarsen/mve-ModernBERT-base-token-pooling-baseline](https://huggingface.co/tomaarsen/mve-ModernBERT-base-token-pooling-baseline) , using only the cached contrastive loss.\n- [tomaarsen/mve-ModernBERT-base-token-pooling](https://huggingface.co/tomaarsen/mve-ModernBERT-base-token-pooling) , using the cached contrastive loss wrapped with the unpooled and factor-4 objectives above.\n\nBoth runs use one epoch over the same 120,000 MS MARCO triplets, seed 42, and matching initialization and optimization settings. The training token cap is 300, and short queries are expanded to at least 32 token positions to provide additional contextual query vectors. Only the loss differs between the two experiments, and I didn't run a hyperparameter search.\n\nI evaluated both models on NanoBEIR@13, reporting mean nDCG@10 (%). The figure compares ordinary and pooling-aware training with pool factors of 2, 3, 4, and 6, alongside unpooled retrieval. The table also includes fixed caps and adaptive thresholds, showing how the training benefits transfer to other pooling settings.\n\n| Pooling at inference | Mean vectors | Ordinary model | Pooling-aware model | Gain (points) | \n|---|---|---|---|---|\n| Unpooled | 193.22 | 51.84 | 52.39 | +0.55 | \n| `pool_factor=2` | 96.86 | 50.10 | 51.74 | +1.64 | \n| `pool_factor=3` | 64.74 | 49.70 | 51.45 | +1.76 | \n| `pool_factor=4` | 48.68 | 49.03 | 51.23 | +2.20 | \n| `pool_factor=5` | 39.04 | 48.13 | 50.28 | +2.15 | \n| `pool_factor=6` | 32.62 | 47.07 | 49.85 | +2.78 | \n| `max_tokens=64` | 54.76 | 49.27 | 50.89 | +1.63 | \n| `max_tokens=32` | 29.58 | 48.25 | 50.12 | +1.87 | \n| `max_tokens=16` | 15.60 | 46.61 | 49.26 | +2.65 | \n| `max_tokens=8` | 7.99 | 43.12 | 46.61 | +3.48 | \n| `max_tokens=4` | 4.00 | 37.22 | 43.00 | +5.78 | \n| `distance_threshold` (calibrated) | 54.76 | 50.10 | 51.24 | +1.14 | \n| `distance_threshold` (calibrated) | 29.58 | 48.38 | 50.96 | +2.58 | \n| `distance_threshold` (calibrated) | 15.60 | 47.23 | 49.63 | +2.41 | \n\nFor adaptive pooling, I calibrated a separate threshold for each encoder and NanoBEIR corpus, using its unpooled document embeddings to match the average vector counts under caps of 64, 32, and 16. Short documents bring those averages below their respective caps.\n\nAt the trained `pool_factor=4`, nDCG@10 improves from 49.03 to 51.23, while storing 74.8% fewer document vectors than unpooled retrieval. This increases score retention from 94.6% to 97.8%, where retention is the pooled mean nDCG@10 divided by each model's own unpooled mean. Unpooled retrieval also improves, from 51.84 to 52.39.\n\nThe gains extend to fixed caps and adaptive pooling, even though training used only the full representation and proportional pooling. At an average of 29.58 vectors per document, adaptive pooling retains 97.3% of the pooling-aware model's unpooled score while storing 84.7% fewer document vectors. The normally trained model retains 93.3% at the same budget.\n\n## \n\t\tCombining Compression Techniques\n\t\n\nIf index size is your main concern, I would reach for these options in this order:\n\n1. Embedding quantization to 4, 2, or 1 bit(s), for example via [Fast-PLAID](https://github.com/lightonai/fast-plaid) , to reduce storage per coordinate.\n2. Token pooling to reduce the number of document vectors, optionally with pooling-aware fine-tuning as shown in this blog.\n3. [Matryoshka token embeddings](https://sbert.net/examples/multi_vector_encoder/training/matryoshka/README.html) to reduce the dimensionality of each token vector, optionally with fine-tuning to preserve retrieval quality at smaller widths.\n\nIf the model's default dimensionality is larger than 128, I would consider Matryoshka truncation and fine-tuning earlier as well. These approaches can be combined, so evaluate the resulting index size and retrieval quality with your search engine and corpus.\n\n`MultiVectorMatryoshkaLoss` trains token embeddings to remain useful after truncation. You can wrap it with `MultiVectorTokenPoolingLoss` to train for both smaller dimensions and fewer document vectors, reusing the encoder's embeddings across the compression settings.\n\nTry the [complete training example](https://sbert.net/examples/multi_vector_encoder/training/token_pooling/README.html), choose a budget that matters for your application, and compare it with a matched ordinary-loss run. The accompanying [evaluation script](https://github.com/huggingface/sentence-transformers/blob/main/examples/multi_vector_encoder/training/token_pooling/evaluate_token_pooling.py) measures unpooled retrieval, caps of 64, 32, 16, 8, and 4 vectors, and pool factors of 2 through 6, so you can plot the quality and storage tradeoff for your own model.\n\nWith a few extra lines in the training objective, a smaller document representation becomes something the model actively learns to use. That opens up an exciting amount of room to make multi-vector retrieval more practical!\n\n## \n\t\tReferences\n\t\n\n- Stefan Josef. [Learn to Pool: Lightweight Fine-Tuning for Flexible Multi-Vector Compression](https://stefan-jo.github.io/learn-to-pool) .\n- Benjamin Clavié, Antoine Chaffin, and Griffin Adams. [Reducing the Footprint of Multi-Vector Retrieval with Minimal Performance Impact via Token Pooling](https://huggingface.co/papers/2409.14683) .\n- Antoine Chaffin. [After the party comes the free lunch: regularizing ColBERT models to enhance pooling capabilities and reduce index footprint](https://huggingface.co/blog/lightonai/lateon-hpool-regularization) .\n- Luyu Gao, Yunyi Zhang, Jiawei Han, and Jamie Callan. [Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup](https://huggingface.co/papers/2101.06983) .\n- [Introduction to Matryoshka Embedding Models](https://huggingface.co/blog/matryoshka) .", "url": "https://wpnews.pro/news/token-pooling-aware-training-for-multi-vector-embedding-models", "canonical_source": "https://huggingface.co/blog/token-pooling", "published_at": "2026-11-11 00:00:00+00:00", "updated_at": "2026-11-11 14:54:35.412552+00:00", "lang": "en", "topics": ["large-language-models", "ai-research", "machine-learning", "artificial-intelligence"], "entities": ["Hugging Face", "tomaarsen", "Sentence Transformers", "NanoBEIR", "Stefan Josef", "LightOn", "HierarchicalTokenPooling", "Fast-PLAID"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/token-pooling-aware-training-for-multi-vector-embedding-models", "markdown": "https://wpnews.pro/news/token-pooling-aware-training-for-multi-vector-embedding-models.md", "text": "https://wpnews.pro/news/token-pooling-aware-training-for-multi-vector-embedding-models.txt", "jsonld": "https://wpnews.pro/news/token-pooling-aware-training-for-multi-vector-embedding-models.jsonld"}}