多向量检索省下四分之三向量：Token Pooling 训练法

多向量检索（ColBERT 这一类）把文档拆成一组 token 向量，召回效果好，但索引大得吓人。Hugging Face 博客在 2026 年 11 月 11 日发了一篇 Token Pooling-Aware Training，讲的不只是「压缩」，而是「训练模型去适应压缩」：同样的向量预算，保留率从 94.6% 提到 97.8%。

这个思路和 Matryoshka 嵌入很像：不是等模型训完再截断，而是让训练目标提前知道「将来会被截断」，于是模型学会把信息放在截断后仍能用的位置上。

为什么多向量索引这么贵

单向量模型把整篇文档压成一个向量；多向量模型保留每个 token 的上下文向量，用 MaxSim 让查询的不同部分匹配文档的不同部分——问「续航长的笔记本」时，前半句可以命中一处，后半句命中另一处。

代价是向量数量。文章算了一笔账：按每个残差坐标 4 bit 量化，每篇文档 200 个 128 维向量要 12.8 KB，一百万篇就是 12.8 GB；池化到每篇 50 个向量，同样的维度和精度下变成 3.2 GB。

关键观察很简单：很多 token 向量本来就相似。同一段话里描述相近概念的几个 token，没有必要各存一份。

池化：把相似 token 合成一个

• 用 Ward 层次聚类：每个 token 一个组，每轮合并「合并后组内方差增加最小」的两个组，直到达到预算。
• 每个组用均值代表，再重新归一化，之后照常用 MaxSim 打分。
• 默认保护第一个 token，并且只对文档向量生效，查询向量完全不动——所以查询侧的逻辑不用改。

预算有三种控制方式：按比例（pool_factor=4 大致保留四分之一向量）、按上限（max_tokens=16 表示每篇最多 16 个，含受保护的 token）、按共享阈值（distance_threshold，内容越杂保留越多）。比例适合文档长度接近的语料，阈值适合长度和内容差异都大的语料。

训练：让模型学会在压缩后检索

池化本身对普通模型也能用，问题出在「合并会抹掉对检索有用的区别」，而原训练目标从没要求模型在压缩表示下工作。文章的解法是双损失：

• 每个 batch 算两个检索损失：原始文档向量一份，压缩到约四分之一后一份，各占一半。
• Ward 聚类是离散决策，训练时用 straight-through estimator 把梯度从池化向量传回成员 token，不去微分聚类归属。
• 实现上是 MultiVectorTokenPoolingLoss 包住现成的对比损失；池化变体复用同一次前向的 token 向量，不用多跑一遍编码器。
• 实验设置：两个模型都从 ModernBERT-base 出发，各学一套 128 维 token 投影，跑同样的 120,000 条 MS MARCO 三元组一个 epoch，seed 42，只有损失不同，没做超参搜索。
• 训练脚本用 mini-batch 64、训练 batch 256 的梯度缓存；token 上限 300，短查询补到至少 32 个位置。

结果

NanoBEIR 上的平均 nDCG@10（推理时按不同预算池化）：

几个读法：

• 在 pool_factor=4 上，nDCG@10 从 49.03 到 51.23，同时少存 74.8% 的文档向量；保留率从 94.6% 升到 97.8%。
• 压缩越狠，训练带来的增益越大：上限压到每篇 4 个向量时，普通模型掉到 37.22，池化感知模型还有 43.00。
• 自适应池化在平均 29.58 个向量时保留 97.3%，少存 84.7%；普通模型同一预算只保留 93.3%。
• 训练只用了「完整表示 + 比例池化」两种视图，但收益能迁移到固定上限和自适应阈值上，说明学到的是「怎么安排 token 向量」，不是背下了某个预算。

更早的论文给出了另一组参照：同样的聚类池化，索引体积减半几乎不掉点；缩到原来的四分之一到三分之一时，多数数据集退化仍低于 5%，而且不需要改架构、不需要查询时处理。这篇文章的价值是给这套「事后压缩」补上了「训练时就考虑压缩」这一环。

怎么自己测

• 装 sentence-transformers[train,kernels]，用 uv 跑仓库里的示例脚本，先跑 --baseline 再跑池化感知版本；两个模型除损失外必须同数据、同种子。
• 换成你自己的语料，但三元组要同一批，否则比的不是损失，是数据。
• 评估至少覆盖 pool_factor 的 2、4、6 三档，画「每篇平均向量数 vs nDCG」曲线，而不是只看一个点。
• 上生产前再叠加量化（例如 4 bit 残差），因为池化和量化可以叠加，收益不是简单相加。
• 想要更快，可以开 kernel：文章说 Ward 聚类的 CUDA 实现比 SciPy 快 30 到 150 倍，具体取决于设置。

边界

• 示例是 ModernBERT-base、128 维投影、120,000 条三元组、单 epoch、没有调参搜索；结论不能直接外推到所有编码器。
• NanoBEIR 是小基准，分数绝对值不代表你的领域；增益幅度也会随数据变化。
• 自适应阈值要在你的语料上校准，换语料后平均向量数会漂移。
• 训练目标只优化检索，不保证其他下游任务；用它做分类或聚类前要重新评估。

本文基于以下公开报道整理。

参考来源

• Token Pooling-Aware Training for Multi-Vector Embedding Models — Hugging Face Blog（2026-11-11）
• Reducing the Footprint of Multi-Vector Retrieval with Minimal Performance Impact via Token Pooling — arXiv（2024-09-23）

本文由 AI 辅助整理，方法与数字来自 Hugging Face 博客的正文与结果表、以及 arXiv 摘要；我们未独立复现训练，表中结果均为原作者实验。