Conscious Token Pooling in Sentence Transformers

A new training device combines variants with and without token grouping to teach multi-vector models to withstand compression. Coupled with a calibrable hierarchical pooling, it aims for an average of vectors per document while maintaining the stated dimension and accuracy. Memory gains are significant at 4 bits when hundreds of vectors per document need to be indexed.

Lighter Indexing Without Repassing Through the Encoder

Pooling variants reuse already computed token embeddings; adding an additional pooler does not require a new pass through the encoder and only adds clustering and scoring work. A training script can operate with mini-batches of 64 and a batch size of 256, and the base loss is instantiated with a mini_batch_size of 64. On the indexing side, reducing each document to 50 vectors brings the footprint down to 3.2 GB for one million documents, compared to 12.8 GB when 200 vectors of 128 dimensions are stored at 4 bits per coordinate. The compression objective adds to the retrieval objective, and the chosen weights average their losses.

MultiVectorTokenPoolingLoss Combines Grouping and Learning

MultiVectorTokenPoolingLoss combines a multi-vector retrieval loss function and uses it with the specified poolers. The retrieval loss function learns to assign a higher score to its relevant document for each query compared to the negatives in the batch, using the similarity measure defined in the model by default. A model with the parameter set to meanmaxsim would also use MeanMaxSim during training. An example configuration uses token_poolings=[None, pooling] and weights [0.5, 0.5] with an encoder such as "answerdotai/ModernBERT-base" and a HierarchicalTokenPooling pooler with pool_factor=4. Gradient caching allows for processing small mini-batches while retaining more effective negatives.

Training Designed for Grouped Representations

Each batch can contribute to two weighted losses of 0.5: one without pooling and the other with pooling. The first term learns retrieval on complete representations, while the second pushes the model to structure its token embeddings so that useful matches survive merging. Both objectives update the same encoder, with Ward clustering making discrete decisions. During training, the gradients of the grouped vectors are passed directly to the member tokens without differentiating the assignments, allowing for learning from observed errors on the compressed versions. This approach transposes the idea of Matryoshka embeddings to token reduction, applying the objective to less detailed representations, in order to leverage the already useful grouping on classic multi-vector models while limiting its undesirable effects.

Hierarchical Pooling by Ward in Practice

HierarchicalTokenPooling relies on Ward's hierarchical clustering to group similar token embeddings. It starts with one cluster per token, iteratively merging the two groups whose union increases the internal spread the least, measured by the sum of squared distances to the mean. The merges build a hierarchy that allows for different compression budgets. Each cluster is represented by its mean, and then the vectors are renormalized and can be directly used with MaxSim. This mechanism fits within the multi-vector paradigm where different parts of a query can correspond to distinct segments of a document, as opposed to a single embedding for the entire text.

Setting a Budget: Ceiling, Fraction, or Distance Threshold

Three controls are proposed to manage the number of vectors per document: a maximum fixed at 16, retaining about a quarter of the initial vectors, or a shared distance threshold for Ward merges. This parameter plays a key role when document sizes fluctuate: with a pool_factor=4, a longer document retains nearly a quarter of its vectors, while using a distance threshold, the number of vectors can grow if certain distinct embeddings need to be maintained. Adaptive grouping based on a common threshold merges similarities more thoroughly and retains more vectors in areas where the content shows greater diversity. This approach works with any multi-vector model, including those not specifically trained for pooling.

Calibrating the Threshold for a Target Average of Vectors

Instead of manually setting a threshold, it can be calibrated to achieve a target average of vectors. From a representative sample, documents are encoded, a threshold is calculated with HierarchicalTokenPooling.calibrate_distance_threshold (target_tokens=16), and then this threshold is reused during indexing. Calibration requires only document embeddings and aims for an average: each text can receive more or fewer than 16 vectors, and the average obtained can vary across another corpus. This process combines with compressed indexes via residual quantization of 2 or 4 bits, a context in which the memory cost of hundreds of vectors per document accumulates quickly, as illustrated by the orders of magnitude at 4 bits.