Token Pooling-Aware Training for Multi-Vector Embedding Models

This post explains token pooling for multi-vector embedding models, which merges similar document token vectors via Ward hierarchical clustering to cut storage, and introduces pooling-aware training in Sentence Transformers that applies the retrieval loss to both original and pooled representations. At pool_factor=4, the pooling-aware model retains 97.8% of its unpooled NanoBEIR nDCG@10 versus 94.6% for a normally trained model, with a 4x reduction in document vectors. It includes a reproducible training example, optional kernel acceleration, and evaluation results across pool factors, fixed caps, and calibrated adaptive thresholds.