Sentence Transformers维护者Tom Aarsen:池化感知训练让多向量嵌入4倍压缩下保留97.8%质量
多向量嵌入模型(如ColBERT)为每个查询token提供独立的文档匹配机会,检索表达力强,但代价是每篇文档产生数百个向量。一个百万级语料库可能包含数亿个向量,即使使用Fast-PLAID的4位残差量化,200个128维向量仍需每文档12.8 KB,百万文档合计12.8 GB。
Token pooling通过Ward层次聚类合并相似的token向量来压缩表示。以pool_factor=4为例,文档仅保留约四分之一的原始向量。压缩越激进,越容易丢失对检索至关重要的区分信息——而普通训练目标从未要求编码器适应这种压缩。
Sentence Transformers维护者Tom Aarsen在2026年11月11日发布的博客中介绍并实现了池化感知训练方法。其核心思路与Matryoshka嵌入类似:在多个表示粒度上同时施加检索目标。具体而言,每个训练批次贡献两个损失——原始文档向量的检索损失和池化后文档向量的检索损失,各占权重0.5取平均。
实现上,MultiVectorTokenPoolingLoss包装已有的多向量检索损失(如CachedMultiVectorMultipleNegativesRankingLoss),接受一组token_poolings参数。None代表原始向量,HierarchicalTokenPooling(pool_factor=4)代表压缩后的向量。Ward聚类在训练时通过直通估计器(straight-through estimator)将池化向量的梯度传回成员token,不对聚类分配本身求导,使编码器能从压缩后的检索错误中学习。
Aarsen在answerdotai/ModernBERT-base上训练了两个模型:一个仅使用普通对比损失,一个使用池化感知损失。两者均在一个epoch的120,000条MS MARCO三元组上训练,seed 42,初始化与优化设置完全匹配,仅损失函数不同。在NanoBEIR@13评测中,普通模型在pool_factor=4时保留94.6%的未池化nDCG@10(49.03对51.84),池化感知模型保留97.8%(51.23对52.39)。
值得注意的是,池化感知训练的收益不仅限于训练时使用的pool_factor=4。在固定上限(max_tokens=8时增益达+3.48个百分点,max_tokens=4时达+5.78个百分点)和自适应距离阈值下,增益同样存在,尽管训练仅使用了全表示和比例池化两种条件。
该方法的适用边界需要留意:所有数字来自作者单次自测,未池化基线本身较低(51.84/52.39),训练规模有限(一个epoch、120K三元组),评测集为NanoBEIR@13而非完整BEIR。Aarsen本人也说明未进行超参数搜索。方法已集成到Sentence Transformers库中,代码和模型权重(tomaarsen/mve-ModernBERT-base-token-pooling)可供复现。
信息源:huggingface.co