Sentence Transformers 详解多向量嵌入模型的 Token 池化感知训练

多向量嵌入模型让查询中的每个 token 都能在文档中寻找匹配项。这大幅增强了搜索的表达能力，但也让文档向量的数量急剧膨胀。一个包含数百万篇文档的语料库，往往会产生数以亿计的向量！

Token 池化（Token pooling）通过合并相似的 token 向量来缩减表示体积。不过，压缩得越狠，丢失有用信息的风险就越高。感知池化训练（Pooling-aware training）能让编码器学会在向量合并时保留关键信息，从而在同样的存储预算下获得更好的检索表现。

本文中，我将介绍池化的工作原理，使用 Sentence Transformers 训练一个用于压缩检索的模型，并探讨检索质量与存储成本之间的权衡。在文档向量减少 4 倍的前提下，借助 pool_factor=4，常规训练的模型能保留未池化时 NanoBEIR nDCG@10 分数的 94.6%。而在同样的存储预算下，经过 Token 感知池化训练的模型把这一指标提升到了 97.8%！接下来看看具体是怎么实现的。

本文基于 Stefan Josef 的 Learn to Pool 以及 LightOn 的池化正则化研究中探讨的感知池化训练，后者还探索了定向训练与自适应预算。本文将重点梳理该方法的核心逻辑，并结合一个简短、可复现的 Sentence Transformers 示例带大家上手实践。

在相同的向量预算下，感知池化训练能显著提升检索质量。详情见下文。

目录

• 为什么多向量检索需要这么多向量
• Token Pooling
  • 选择池化预算
  • 校准自适应池化
• Token Pooling-Aware Training
  • Sentence Transformers 中的损失函数包装器
  • 可选的算子加速
• 使用训练好的模型
• 评估结果
• 结合多种压缩技术
• 参考文献
• 代码与模型

为什么多向量检索需要这么多向量

单向量嵌入模型会将整篇文档压缩成一个向量。而像 ColBERT 这样的多向量模型则会保留一组上下文 token 向量，使查询的不同部分能够分别匹配文档的不同部分。

例如，在查询“长续航笔记本电脑”时，查询可以在文档的某一处匹配到有关笔记本的信息，并在另一处匹配到电池续航的内容。采用 MaxSim 打分时，每个查询向量都会找出与之最匹配的文档向量，然后将这些匹配得分相加：

MaxSim⁡(Q,D)=∑q∈Qmax⁡d∈Dq⊤d \operatorname{MaxSim}(Q, D) = \sum_{q \in Q} \max_{d \in D} q^\top d

即便使用 Fast-PLAID 支持的 2 比特或 4 比特残差量化，单篇文档存储数百个向量的开销依旧相当可观。在每个残差坐标占 4 比特的情况下，每篇文档包含 200 个 128 维向量需要 12.8 KB 的量化残差空间，换算到 100 万篇文档就是 12.8 GB。如果在相同维度和精度下将每篇文档池化至 50 个向量，存储占用就能降至 3.2 GB。

一个有价值的发现是：很多文档向量彼此相似。多个 token 可能在描述密切相关的概念，这就为将它们合并表示提供了可能。

Token 池化

HierarchicalTokenPooling 利用 Ward 层次聚类将相似的 token 嵌入分组。它最初将每个 token 向量划为一个独立的组，随后不断合并两组向量——每次都挑选合并后组内离散程度增幅最小的那两组。离散程度是用各向量到组均值的距离平方和来衡量的，因此 Ward 聚类法倾向于生成由相似向量构成的紧凑分组。

这种合并操作会构建出一个由小及大的层次化分组结构，便于我们按计算预算自由选择压缩程度。接着，我们用各组的均值向量代表该组，并重新进行归一化，以便直接用于 MaxSim 评分。

你也可以让压缩比例与文档长度成正比：

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor=4)

在设置 pool_factor=4 时，文档会保留大约四分之一的原始向量，因此长文档可以按比例保留更多容量。默认情况下，首个 token 会受到保护而不参与合并，且池化仅作用于文档嵌入，查询（query）嵌入则保持不变。

选择池化预算

HierarchicalTokenPooling 提供了三种控制方式：

当文档长短不一时，这种差异就很关键。在 pool_factor=4 模式下，较长的文档可以保留大约四分之一的向量，其表示能力能够随内容增加而扩展；而在 distance_threshold 模式下，只有当需要保留更多差异较大的嵌入时，向量数量才会增加。

校准自适应池化

自适应池化采用通用的距离阈值，在更积极地合并相似嵌入的同时，为内容差异较大的文档保留更多向量。该方法适用于任何多向量模型，即使模型在训练时并未加入池化感知目标也可以使用。

除了手动选择阈值，你还可以根据目标平均向量数量来进行校准。只要在 calibration_documents 中提供语料库的代表性样本：

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("lightonai/LateOn")
calibration_embeddings = model.encode_document(calibration_documents)
threshold = HierarchicalTokenPooling.calibrate_distance_threshold(calibration_embeddings, target_tokens=16)
pooling = HierarchicalTokenPooling(distance_threshold=threshold)

# Apply pooling, targeting an average of 16 vectors per document
document_embeddings = model.encode_document(documents_to_index, token_pooling=pooling)

校准过程只需用到文档嵌入，并会返回一个可在建立索引时复用的阈值。由于设定的目标只是平均值，因此具体到单篇文档，其向量数量可能多于或少于 16 个；若换到其他语料库上，实际达到的平均值也可能有所变化。

Token 池化感知训练

普通的多向量模型直接使用池化就已经能见效了。真正的难点在于：当文档中包含对检索至关重要的细微差异时，合并向量往往会让这些差异难以还原。模型最初的训练目标从未要求过它配合这种压缩表示一起工作。不过，你可以通过训练让它专门具备这种能力。

如果你了解过 俄罗斯套娃嵌入（Matryoshka embeddings），就会对这里的训练原理感到熟悉。套娃嵌入在多个不同的嵌入维度上施加检索目标，引导模型让较短的前缀维度依然保持有效；而池化感知训练则是将检索目标直接施加在 token 向量数量更少的压缩表示上。

在我的示例中，每个训练批次都会产生两项损失：

1. 使用原始文档嵌入进行检索的损失。
2. 将文档压缩至原始向量数量约四分之一后再进行检索的损失。

我对两者取平均值：

L=12Lunpooled+12Lpooled \mathcal{L} = \tfrac{1}{2}\mathcal{L}_{\mathrm{unpooled}} + \tfrac{1}{2}\mathcal{L}_{\mathrm{pooled}}

第一项损失引导模型学会使用完整表征进行检索，第二项损失则引导模型优化 token 嵌入的排布，确保池化后依然保留关键的匹配信息。这两个目标更新的是同一个编码器，让它能同时获取压缩前后的检索效果反馈。

Ward 聚类在决定哪些 token 归为一类时属于离散决策。在训练期间，HierarchicalTokenPooling 利用直通估计器（straight-through estimator）将梯度从池化向量直接传回其包含的 token，无需对聚类分配过程求导。这样一来，编码器就能根据池化表征产生的检索误差持续优化。

Sentence Transformers 中的损失包装器

MultiVectorTokenPoolingLoss 封装了现有的多向量检索损失，并通过你选定的池化器来应用它：

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import (
    CachedMultiVectorMultipleNegativesRankingLoss,
    MultiVectorTokenPoolingLoss,
)
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("answerdotai/ModernBERT-base", similarity_fn_name="maxsim")
pooling = HierarchicalTokenPooling(pool_factor=4)

base_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64)
loss = MultiVectorTokenPoolingLoss(
    model,
    base_loss,
    token_poolings=[None, pooling],
    weights=[0.5, 0.5],
)

CachedMultiVectorMultipleNegativesRankingLoss 提供了检索目标，促使每个查询对相关文档的打分高于批次内的负例。它在默认情况下使用 model.similarity，并沿用模型的相似度设置（本例中为 MaxSim）及内核配置。如果模型配置了 similarity_fn_name="meanmaxsim"，训练时也会采用 MeanMaxSim。

借助梯度缓存（Gradient caching），编码器可以在处理较小 mini-batch 的同时保留更多批次内负例。我的训练脚本采用了大小为 64 的 mini-batch，实际训练批次大小为 256。

MultiVectorTokenPoolingLoss 引入了压缩目标。这里的 None 表示保留原始嵌入的变体，pooling 则对应 4 倍压缩变体。两个损失项会按权重取平均。不同池化变体会复用算好的 token 嵌入，因此每增加一个池化器，只会额外增加聚类与打分计算，无需为该变体重新运行一遍编码器。

池化器也可以是自定义的 BaseTokenPooling 实现，只要能在输入嵌入中保留梯度即可。你可以根据最终部署所需的表征形式灵活调整训练目标。

你可以将 4 倍压缩池化器替换为实际部署时的配置，也可以同时引入多个池化器并分别为其损失设置权重。

完整的训练脚本包含了数据集、训练器、验证损失以及训练过程中的简易检索评估。克隆代码仓库后，可以直接使用 uv 在本地运行这两种变体：

uv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py --baseline
uv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py

该脚本声明了训练和内核所需的依赖，uv 会自动完成安装，兼容的内核也会自动启用。

你也可以借助 Hugging Face Jobs 在云端 GPU 上运行训练。安装 CLI 并使用具备 Jobs 运行权限与模型上传权限的 Token 登录：

pip install -U huggingface_hub
hf auth login

hf jobs uv run \
    --flavor a10g-small \
    --timeout 4h \
    --secrets HF_TOKEN \
    -- https://raw.githubusercontent.com/huggingface/sentence-transformers/main/examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py

这会在 A10G GPU 上运行感知池化的变体，运行限时 4 小时。将 HF_TOKEN 作为密钥传入后，脚本会在训练结束时自动把模型上传到你的 Hub 账户。如果还想运行采用常规损失的基准版本，只需在脚本 URL 后追加 --baseline 再次执行该命令即可。

可选的内核加速

你可以通过加速内核来提升 MaxSim 打分、Ward 聚类以及支持的骨干网络运算速度。安装训练及内核依赖项：

pip install -U "sentence-transformers[train,kernels]"

在前面的训练示例中，构建损失函数前需替换模型和池化器的初始化代码：

model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    device="cuda",
    similarity_fn_name="maxsim",
    use_kernels=True,
)
pooling = HierarchicalTokenPooling(pool_factor=4)

分层池化器默认会继承模型的内核设置，因此只需在模型上启用内核即可。

Sentence Transformers 会根据你的运行环境选择兼容的内核；若无可替代内核，则回退到现有实现。例如，默认的 MaxSim 内核在对查询和文档打分时，无需在内存中构建完整的 token 相似度矩阵，不过它需要 CUDA 和 Triton 支持。

默认的 Ward 聚类算子改在 CUDA 而非 CPU 上运行，大幅提升了层次聚类的速度。视具体配置不同，速度可达 SciPy 的 30 到 150 倍。

使用训练好的模型

接下来加载支持池化的模型，用压缩后的文档嵌入来进行检索。模型本身不会保存算子配置，如果推理时也想用，加载时加上 use_kernels=True 即可：

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling")
pooling = HierarchicalTokenPooling(pool_factor=4)

queries = ["Which planet is known as the Red Planet?"]
documents = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
]

query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents, token_pooling=pooling)

print([len(embedding) for embedding in document_embeddings])
# [5, 6, 5]
print(model.similarity(query_embeddings, document_embeddings))
# tensor([[16.0037, 25.9879, 22.8159]], device='cuda:0')

# For comparison, the same model without token pooling:
# [18, 21, 19]
# tensor([[16.6115, 27.6060, 24.1766]], device='cuda:0')

池化器在编码器的前向计算后压缩文档嵌入，向量总数从 58 个缩减到 16 个，同时保持查询嵌入不变。火星的相似度得分依然最高，木星和金星紧随其后——尽管存储的文档向量大幅减少，排序却丝毫不受影响。

你可以在推理时修改 pool_factor，用同一个训练好的模型尝试不同的存储预算。若需多次使用，也可以在初始化 MultiVectorEncoder 时传入 token_pooling=pooling，将其设为运行时默认值。池化器的 tasks 依然负责控制压缩哪些输入（例如仅压缩文档）。

评测结果

为了深入探究，我基于 answerdotai/ModernBERT-base 训练了两个模型，各自带有一个全新的 128 维 token 投影层：

• tomaarsen/mve-ModernBERT-base-token-pooling-baseline，仅使用带缓存的对比损失。
• tomaarsen/mve-ModernBERT-base-token-pooling，在带缓存的对比损失外，包装了上文提到的未池化目标以及 4 倍池化目标。

两次训练均在相同的 12 万条 MS MARCO 三元组上训练 1 个 epoch，随机种子为 42，初始化和优化参数完全一致。训练时最大 token 数设为 300，较短的查询会扩充到至少 32 个 token 位置，以提供更多的上下文查询向量。两组实验唯一的区别就是损失函数，我也没有进行超参数搜索。

我在 NanoBEIR@13 上评测了这两个模型，指标采用平均 nDCG@10（%）。下图对比了普通训练与支持池化的训练在池化系数为 2、3、4、6 以及未池化检索下的表现。表格中还加入了固定上限与自适应阈值的表现，展示训练带来的收益如何迁移到其他池化配置上。

自适应池化方面，我针对每个编码器和 NanoBEIR 语料库分别校准了阈值，利用其未池化的文档嵌入，使其平均向量数对齐 64、32 和 16 的上限。短文档会让实际平均值低于各自的上限。

在训练所用的 pool_factor=4 下，存储的文档向量比未池化检索减少了 74.8%，nDCG@10 却从 49.03 提升到了 51.23。这使分数保留率从 94.6% 跃升至 97.8%（保留率指池化后的平均 nDCG@10 除以该模型自身的未池化均值）。未池化检索的表现也有所提升，从 51.84 提高到 52.39。

尽管训练时只使用了完整表示和等比池化，但这种收益同样延伸到了固定上限和自适应池化中。在平均每篇文档 29.58 个向量的情况下，自适应池化保留了池化感知模型 97.3% 的未池化得分，同时文档向量存储量减少了 84.7%。在相同预算下，普通训练的模型仅能保留 93.3%。

结合多种压缩技术

如果你最关心的是索引体积，建议按以下顺序依次考虑：

1. 将嵌入量化为 4、2 或 1 位（例如借助 Fast-PLAID），以压缩每个坐标维度的存储开销。
2. 通过 Token 池化减少文档向量的数量；如本文所述，还可以搭配感知池化的微调方法。
3. 采用俄罗斯套娃 Token 嵌入降低每个 Token 向量的维度；同样可以配合微调，确保低维度下依然保持良好的检索质量。

如果模型的默认维度大于 128，我建议及早考虑套娃截断和微调。这些方案完全可以叠加使用，建议结合你具体的搜索引擎和语料库，实测最终的索引体积和检索表现。

MultiVectorMatryoshkaLoss 能够训练 Token 嵌入，使其截断后依然可用。你还可以给它套上一层 MultiVectorTokenPoolingLoss，同时针对更低的维度和更少的文档向量进行训练，让编码器的嵌入能在各种压缩配置下复用。

不妨试试这个完整训练示例，按你的实际业务设定资源预算，并与采用常规损失函数的对照组对比。配套的评测脚本支持未池化检索、限制最多 64、32、16、8、4 个向量，以及 2 到 6 倍池化因子的测试，方便你绘制出属于自己模型的检索质量与存储开销权衡曲线。

只需在训练目标里加上几行代码，模型就能主动学会利用更紧凑的文档表示。这为多向量检索的落地应用打开了令人兴奋的广阔空间！

参考资料

• Stefan Josef. Learn to Pool: Lightweight Fine-Tuning for Flexible Multi-Vector Compression.
• Benjamin Clavié, Antoine Chaffin, and Griffin Adams. Reducing the Footprint of Multi-Vector Retrieval with Minimal Performance Impact via Token Pooling.
• Antoine Chaffin. After the party comes the free lunch: regularizing ColBERT models to enhance pooling capabilities and reduce index footprint.
• Luyu Gao, Yunyi Zhang, Jiawei Han, and Jamie Callan. Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup.
• 俄罗斯套娃嵌入模型入门指南。

代码与模型

• 感知 Token 池化的训练示例与说明文档
• 完整训练脚本
• NanoBEIR 评测脚本
• Example checkpoints:
  • tomaarsen/mve-ModernBERT-base-token-pooling
  • tomaarsen/mve-ModernBERT-base-token-pooling-baseline
• 兼容 Sentence Transformers 的预训练多向量模型
• 多向量模型的感知套娃训练
• Fast-PLAID