Skip to content
AI Primer
release

Sentence Transformers v6.3 adds pooling-aware training for retrieval with fewer vectors

Sentence Transformers v6.3 adds pooling-aware training for multi-vector retrieval. Maintainer tests retained 97.3% of unpooled retrieval quality with 84.7% fewer document vectors; code and checkpoints are available.

5 min read
Sentence Transformers v6.3 adds pooling-aware training for retrieval with fewer vectors
Sentence Transformers v6.3 adds pooling-aware training for retrieval with fewer vectors

TL;DR

  • Sentence Transformers v6.3 trains multi-vector retrievers to work with fewer document vectors, as announced in the release thread.
  • Adaptive pooling retained 97.3% of unpooled nDCG@10 with 84.7% fewer document vectors in the maintainer's experiment.
  • CUDA pooling took 1.02 seconds versus 7.00 seconds with fastcluster on 5,000 SciFact documents in the kernel benchmark.

The published checkpoint accepts different pooling budgets at inference, while adaptive calibration needs only document embeddings. The release notes also change MeanMaxSim loss scoring and the treatment of missed positives in evaluation.

Pooling-aware training

HierarchicalTokenPooling uses Ward clustering to merge similar document-token vectors into normalized cluster means. By default, queries remain unpooled and the first token is protected from merging.

The example wraps a cached contrastive retrieval loss with MultiVectorTokenPoolingLoss, averaging two objectives:

  1. Retrieval with the original document embeddings, weighted 0.5.
  2. Retrieval after factor-4 document pooling, weighted 0.5.

Both objectives reuse the same encoder embeddings, a pleasantly small change to the training loop. Additional pooling settings add clustering and scoring work without another encoder pass.

A straight-through estimator passes gradients from pooled vectors back to their member tokens without differentiating the discrete cluster assignments.

NanoBEIR results

Sentence Transformers maintainer Tom Aarsen trained two ModernBERT-base models with fresh 128-dimensional token projections on the same 120,000 MS MARCO triplets. Only the loss changed; the published evaluation used one seed, no hyperparameter search, and mean nDCG@10 across 13 NanoBEIR datasets.

| Inference representation | Mean document vectors | Ordinary training | Pooling-aware training |
| --- | ---: | ---: | ---: |
| Unpooled | 193.22 | 51.84 | 52.39 |
| Factor-4 pooling | 48.68 | 49.03 | 51.23 |
| Adaptive pooling | 29.58 | 48.38 | 50.96 |
| Maximum 16 vectors | 15.60 | 46.61 | 49.26 |

At factor 4, pooling-aware training gains 2.20 nDCG points over ordinary training. Retention divides the pooled score by each model's own unpooled score: 97.8% for pooling-aware training versus 94.6% for ordinary training.

The reported savings count document vectors, rather than measured total index bytes or query latency.

Document vector budgets

The checkpoint supports three alternative controls at encoding time:

  • pool_factor=4: retain approximately one quarter of each document's vectors.
  • max_tokens=16: retain at most 16 vectors, including protected tokens.
  • distance_threshold=...: use a shared Ward merge-distance threshold, allowing documents with more distinct embeddings to retain more vectors.

calibrate_distance_threshold targets an average vector count on sample document embeddings, rather than a per-document cap. Another corpus can produce a different average.

For the benchmark, Aarsen calibrated a separate threshold for each encoder and NanoBEIR corpus to match average counts obtained under fixed caps. Adaptive budgets worked despite training using only unpooled and factor-4 representations.

Hierarchical versus adjacent pooling

Aarsen also compared hierarchical pooling with averaging adjacent tokens on LateOn, a separate model from the ModernBERT checkpoints above.

Matryoshka compression

v6.3 adds MultiVectorMatryoshkaLoss for training shorter token-vector prefixes to preserve retrieval quality.

  • Nesting the Matryoshka loss inside MultiVectorTokenPoolingLoss trains for fewer vectors and fewer dimensions together, with pooling applied before truncation.
  • Quantization can stack with both, reducing bits per coordinate.

In a separate 30,000-triplet experiment, 64-dimensional retrieval scored 47.11 with Matryoshka training versus 45.46 with ordinary training, a 1.65-point gain. Relative to each model's 128-dimensional score, retention rose from 93.8% to 96.2%.

CUDA kernels

The optional [kernels] dependency and use_kernels=True enable compatible Hub implementations, according to the kernel announcement, for:

  • MaxSim and MeanMaxSim scoring.
  • Ward clustering for token pooling.
  • Supported transformer-backbone operations.

Including encoding, the release benchmark improved from 23.14 to 18.91 seconds, an 18% reduction. That run used mLateOn, BF16, batch size 32, factor-2 pooling, and an RTX 3090.

Default MaxSim acceleration requires CUDA and Triton. Operations without compatible kernels fall back to reference implementations.

Runtime and scoring changes

Outside compression, v6.3 adds native AWS Neuron inference through device="neuron", faster Hugging Face Dataset inputs, and more efficient community detection.

The compatibility notes spell out four operational changes:

  • Saved settings: constructor truncate_dim is saved, while token_pooling and kernel settings remain runtime options and are not restored with the model.
  • Loss scoring: built-in multi-vector contrastive and distillation losses follow the model's configured similarity even with kernels disabled. MeanMaxSim models now use means where these losses previously used MaxSim sums.
  • Evaluation: CrossEncoder and RRF MAP/nDCG normalize against all known positives. Unretrieved positives receive no credit, changing historical scores.
  • Normalization: newly constructed hierarchical poolers normalize outputs by default. Older saved configurations without normalize_embeddings retain their previous unnormalized behavior.

Further reading

Discussion across the web

Where this story is being discussed, in original context.

On X· 5 threads
TL;DR2 posts
NanoBEIR results1 post
Document vector budgets1 post
Hierarchical versus adjacent pooling1 post
CUDA kernels1 post
Share on X