Token Pooling-Conscious Coaching for Multi-Vector Embedding Fashions

Multi-vector embedding fashions give every question token its personal alternative to discover a match in a doc. This makes them splendidly expressive for search, however these doc vectors add up shortly. A corpus with thousands and thousands of paperwork can include lots of of thousands and thousands of vectors!

Token pooling makes these representations smaller by merging related token vectors. The extra aggressively I compress them, nevertheless, the better it turns into to lose one thing helpful. Pooling-aware coaching teaches the encoder to protect the data that issues when its vectors are merged, serving to you retrieve higher from the identical storage price range.

On this publish, I am going to clarify how pooling works, prepare a mannequin for compressed retrieval with Sentence Transformers, and discover the ensuing high quality and storage tradeoff. With pool_factor=4, my usually educated mannequin retains 94.6% of its unpooled NanoBEIR nDCG@10, with a 4x discount in doc vectors. Token pooling-aware coaching raises that retention to 97.8% on the similar storage price range! Let’s have a look at how.

This publish builds on pooling-aware coaching studied in Stefan Josef’s Be taught to Pool and LightOn’s pooling regularization examine, which additionally explores focused coaching and adaptive budgets. Right here, I am going to give attention to understanding the strategy and placing it to work with a brief, reproducible Sentence Transformers instance.

Pooling-aware coaching improves retrieval high quality on the similar vector price range. Particulars under.

Desk of Contents

Why Multi-Vector Retrieval Wants So Many Vectors

A single-vector embedding mannequin compresses a whole doc into one vector. Multi-vector fashions equivalent to ColBERT as an alternative preserve a set of contextual token vectors, permitting totally different components of a question to match totally different components of a doc.

For instance, a question a couple of laptop computer with an extended battery life can discover proof concerning the laptop computer in a single a part of a doc and battery endurance in one other. With MaxSim scoring, every question vector selects its finest matching doc vector, and people scores are added collectively:

$$\mathrm{MaxSim}(Q,D)=\sum _{q\in Q}\underset{d\in D}{\max }{q}^{\mathrm{\top }}doperatorname\{MaxSim\}(Q,\; D)\; =\; sum\_\{q\; in\; Q\}\; max\_\{d\; in\; D\}\; q^prime\; d$$

Even with the 2-bit or 4-bit residual quantization supported by Quick-PLAID, storing lots of of vectors per doc provides up. At 4 bits per residual coordinate, 200 vectors of 128 dimensions require 12.8 KB of quantized residuals per doc, or 12.8 GB throughout 1,000,000 paperwork. Pooling every doc to 50 vectors reduces that to three.2 GB, on the similar dimensionality and precision.

The helpful statement is that many doc vectors might be related. A number of tokens may describe intently associated ideas, creating a possibility to signify them collectively.

Token Pooling

HierarchicalTokenPooling teams related token embeddings utilizing Ward hierarchical clustering. Beginning with one group per token vector, it repeatedly merges the 2 teams that trigger the smallest improve in how unfold out the vectors are inside their teams. This unfold is measured by the sum of squared distances from every vector to its group’s imply, so Ward’s technique favors compact teams of comparable vectors.

These merges construct a hierarchy of progressively bigger teams, letting us select a stage of compression that matches our price range. We then signify every group by its imply and normalize the ensuing vectors once more, prepared for MaxSim scoring.

You can also make compression proportional to doc size:

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor=4)

With pool_factor=4, paperwork retain roughly one quarter of their authentic vectors, so longer paperwork can preserve proportionally extra capability. By default, the primary token is protected against merging and pooling applies solely to doc embeddings, leaving question embeddings unchanged.

Selecting a Pooling Funds

HierarchicalTokenPooling affords three various controls:

This distinction issues when doc lengths range. With pool_factor=4, an extended doc can retain roughly one quarter of its vectors, permitting its illustration to develop with its content material. With distance_threshold, the vector rely can develop when extra distinct embeddings must be preserved.

Calibrating Adaptive Pooling

Adaptive pooling makes use of a shared distance threshold to merge related embeddings extra aggressively whereas preserving extra vectors for paperwork with distinct content material. It may be used with any multi-vector mannequin, together with one educated and not using a pooling-aware goal.

Quite than selecting a threshold manually, you may calibrate it for a goal common vector rely. Given a consultant pattern of your corpus in calibration_documents:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

mannequin = MultiVectorEncoder("lightonai/LateOn")
calibration_embeddings = mannequin.encode_document(calibration_documents)
threshold = HierarchicalTokenPooling.calibrate_distance_threshold(calibration_embeddings, target_tokens=16)
pooling = HierarchicalTokenPooling(distance_threshold=threshold)


document_embeddings = mannequin.encode_document(documents_to_index, token_pooling=pooling)

Calibration solely wants doc embeddings and returns a threshold you may reuse when indexing. The goal is a mean, so particular person paperwork can obtain extra or fewer than 16 vectors. The achieved common also can change on one other corpus.

Token Pooling-Conscious Coaching

Pooling can already be helpful on an abnormal multi-vector mannequin. The issue seems when a doc comprises a distinction that issues for retrieval, however merging its vectors makes that distinction more durable to get better. The unique coaching goal by no means required the mannequin to work with these compressed representations. However, you may train it to do precisely that.

When you’ve used Matryoshka embeddings, the coaching precept will really feel acquainted. Matryoshka applies a retrieval goal at a number of embedding widths, instructing smaller prefixes to stay helpful. Pooling-aware coaching applies the target to representations with fewer token vectors.

For my instance, every coaching batch contributes two losses:

1. Retrieval utilizing the unique doc embeddings.
2. Retrieval after compressing paperwork to roughly one quarter of their authentic vectors.

I common them:

$$\mathcal{L}=\frac{1}{2}{\mathcal{L}}_{\mathrm{u}\mathrm{n}\mathrm{p}\mathrm{o}\mathrm{o}\mathrm{l}\mathrm{e}\mathrm{d}}+\frac{1}{2}{\mathcal{L}}_{\mathrm{p}\mathrm{o}\mathrm{o}\mathrm{l}\mathrm{e}\mathrm{d}}mathcal\{L\}\; =\; tfrac\{1\}\{2\}mathcal\{L\}\_\{mathrm\{unpooled\}\}\; +\; tfrac\{1\}\{2\}mathcal\{L\}\_\{mathrm\{pooled\}\}$$

The primary time period teaches the mannequin to retrieve with its full representations. The second teaches it to rearrange its token embeddings in order that helpful matches survive pooling. Each goals replace the identical encoder, giving it suggestions about retrieval earlier than and after compression.

Ward clustering makes discrete choices about which tokens belong collectively. Throughout coaching, HierarchicalTokenPooling makes use of a straight-through estimator to cross gradients from pooled vectors again to their member tokens, with out differentiating the cluster assignments. The encoder can then be taught from retrieval errors made with its pooled representations.

A Loss Wrapper in Sentence Transformers

MultiVectorTokenPoolingLoss wraps an present multi-vector retrieval loss and applies it utilizing your chosen poolers:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import (
    CachedMultiVectorMultipleNegativesRankingLoss,
    MultiVectorTokenPoolingLoss,
)
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

mannequin = MultiVectorEncoder("answerdotai/ModernBERT-base", similarity_fn_name="maxsim")
pooling = HierarchicalTokenPooling(pool_factor=4)

base_loss = CachedMultiVectorMultipleNegativesRankingLoss(mannequin, mini_batch_size=64)
loss = MultiVectorTokenPoolingLoss(
    mannequin,
    base_loss,
    token_poolings=[None, pooling],
    weights=[0.5, 0.5],
)

CachedMultiVectorMultipleNegativesRankingLoss provides the retrieval goal. It teaches every question to attain its related doc above the destructive paperwork within the batch. By default, it makes use of mannequin.similarity, utilizing the mannequin’s similarity setting (MaxSim on this instance) and kernel settings. A mannequin configured with similarity_fn_name="meanmaxsim" would additionally use MeanMaxSim throughout coaching.

Gradient caching permits the encoder to course of small mini-batches whereas retaining a bigger set of in-batch negatives. My coaching script makes use of mini-batches of 64 with a coaching batch dimension of 256.

MultiVectorTokenPoolingLoss provides the compression goal. Right here, None signifies that one variant makes use of the unique embeddings, whereas pooling provides the factor-4 variant. The 2 weights common their losses. Pooling variants reuse the computed token embeddings, so every extra pooler provides clustering and scoring work with out requiring a separate encoder cross for that variant.

The poolers will also be customized BaseTokenPooling implementations that protect gradients via their enter embeddings. You’ll be able to adapt the coaching goal to the illustration you wish to deploy.

Substitute the factor-4 pooler along with your chosen deployment configuration, or embrace a number of poolers and select weights for his or her losses.

The whole coaching script consists of the dataset, coach, validation loss, and small retrieval evaluations throughout coaching. From a checkout of the repository, run each variants domestically with uv:

uv run examples/multi_vector_encoder/coaching/token_pooling/training_cached_contrastive.py --baseline
uv run examples/multi_vector_encoder/coaching/token_pooling/training_cached_contrastive.py

The script declares its coaching and kernel dependencies, so uv installs them for you. Appropriate kernels are enabled routinely.

You too can prepare on a cloud GPU with Hugging Face Jobs. Set up the CLI and log in with a token that may run Jobs and add fashions:

pip set up -U huggingface_hub
hf auth login

hf jobs uv run 
    --flavor a10g-small 
    --timeout 4h 
    --secrets HF_TOKEN 
    -- https://uncooked.githubusercontent.com/huggingface/sentence-transformers/predominant/examples/multi_vector_encoder/coaching/token_pooling/training_cached_contrastive.py

This runs the pooling-aware variant on an A10G GPU with a four-hour time restrict. Passing HF_TOKEN as a secret lets the script add the educated mannequin to your Hub account when coaching finishes. To run the ordinary-loss baseline as properly, repeat the command with --baseline after the script URL.

Non-compulsory Kernel Acceleration

You’ll be able to speed up MaxSim scoring, Ward clustering, and supported spine operations with kernels. Set up the coaching and kernel dependencies:

pip set up -U "sentence-transformers[train,kernels]"

Within the coaching instance above, exchange the mannequin and pooler initialization earlier than developing the loss:

mannequin = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    gadget="cuda",
    similarity_fn_name="maxsim",
    use_kernels=True,
)
pooling = HierarchicalTokenPooling(pool_factor=4)

Hierarchical poolers inherit the mannequin’s kernel settings by default, so that you solely must allow kernels on the mannequin.

Sentence Transformers selects suitable kernels on your setup and makes use of the present implementation wherever a substitute is not obtainable. For instance, the default MaxSim kernel scores queries in opposition to paperwork with out constructing the total token similarity matrix in reminiscence, however it does require CUDA and Triton.

The default Ward clustering kernel significantly hastens hierarchical clustering by working on CUDA as an alternative of the CPU. We’re speaking about 30-150x quicker than SciPy relying on the setting.

Utilizing the Skilled Mannequin

Let’s load the pooling-aware mannequin and retrieve with compressed doc embeddings. Kernel settings aren’t saved with the mannequin, so add use_kernels=True when loading if you would like to make use of them for inference too:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

mannequin = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling")
pooling = HierarchicalTokenPooling(pool_factor=4)

queries = ["Which planet is known as the Red Planet?"]
paperwork = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
]

query_embeddings = mannequin.encode_query(queries)
document_embeddings = mannequin.encode_document(paperwork, token_pooling=pooling)

print([len(embedding) for embedding in document_embeddings])

print(mannequin.similarity(query_embeddings, document_embeddings))

The pooler compresses the doc embeddings after the encoder’s ahead pipeline, lowering the whole from 58 vectors to 16 whereas leaving question embeddings unchanged. Mars nonetheless receives the very best similarity rating, adopted by Jupiter and Venus, so the rating is preserved regardless of storing far fewer doc vectors.

You’ll be able to change pool_factor at inference to discover a special storage price range with the identical educated mannequin. For repeated use, you too can cross token_pooling=pooling when initializing MultiVectorEncoder to make it the runtime default. The pooler’s duties nonetheless management which inputs are compressed (e.g. solely paperwork).

Analysis Outcomes

To research, I educated two fashions from answerdotai/ModernBERT-base, every with a contemporary 128-dimensional token projection:

Each runs use one epoch over the identical 120,000 MS MARCO triplets, seed 42, and matching initialization and optimization settings. The coaching token cap is 300, and quick queries are expanded to at the very least 32 token positions to offer extra contextual question vectors. Solely the loss differs between the 2 experiments, and I did not run a hyperparameter search.

I evaluated each fashions on NanoBEIR@13, reporting imply nDCG@10 (%). The determine compares abnormal and pooling-aware coaching with pool elements of two, 3, 4, and 6, alongside unpooled retrieval. The desk additionally consists of fastened caps and adaptive thresholds, exhibiting how the coaching advantages switch to different pooling settings.

For adaptive pooling, I calibrated a separate threshold for every encoder and NanoBEIR corpus, utilizing its unpooled doc embeddings to match the typical vector counts below caps of 64, 32, and 16. Quick paperwork carry these averages under their respective caps.

On the educated pool_factor=4, nDCG@10 improves from 49.03 to 51.23, whereas storing 74.8% fewer doc vectors than unpooled retrieval. This will increase rating retention from 94.6% to 97.8%, the place retention is the pooled imply nDCG@10 divided by every mannequin’s personal unpooled imply. Unpooled retrieval additionally improves, from 51.84 to 52.39.

The positive aspects lengthen to fastened caps and adaptive pooling, regardless that coaching used solely the total illustration and proportional pooling. At a mean of 29.58 vectors per doc, adaptive pooling retains 97.3% of the pooling-aware mannequin’s unpooled rating whereas storing 84.7% fewer doc vectors. The usually educated mannequin retains 93.3% on the similar price range.

Combining Compression Methods

If index dimension is your predominant concern, I’d attain for these choices on this order:

1. Embedding quantization to 4, 2, or 1 bit(s), for instance through Quick-PLAID, to cut back storage per coordinate.
2. Token pooling to cut back the variety of doc vectors, optionally with pooling-aware fine-tuning as proven on this weblog.
3. Matryoshka token embeddings to cut back the dimensionality of every token vector, optionally with fine-tuning to protect retrieval high quality at smaller widths.

If the mannequin’s default dimensionality is bigger than 128, I’d contemplate Matryoshka truncation and fine-tuning earlier as properly. These approaches might be mixed, so consider the ensuing index dimension and retrieval high quality along with your search engine and corpus.

MultiVectorMatryoshkaLoss trains token embeddings to stay helpful after truncation. You’ll be able to wrap it with MultiVectorTokenPoolingLoss to coach for each smaller dimensions and fewer doc vectors, reusing the encoder’s embeddings throughout the compression settings.

Attempt the whole coaching instance, select a price range that issues on your utility, and evaluate it with a matched ordinary-loss run. The accompanying analysis script measures unpooled retrieval, caps of 64, 32, 16, 8, and 4 vectors, and pool elements of two via 6, so you may plot the standard and storage tradeoff on your personal mannequin.

With a number of further strains within the coaching goal, a smaller doc illustration turns into one thing the mannequin actively learns to make use of. That opens up an thrilling quantity of room to make multi-vector retrieval extra sensible!

References

Code and Fashions