Hugging Face 详解 Token Pooling-Aware Training：多向量嵌入模型压缩后仍保 97.8% 检索质量

多向量嵌入模型让查询中的每个词元都有机会在文档中找到匹配。这使它们在搜索中表现出色，但这些文档向量会迅速累积。一个包含数百万文档的语料库可能包含数亿个向量！

Token pooling（词元池化）通过合并相似的词元向量来缩小这些表示。然而，压缩得越激进，就越容易丢失有用的信息。池化感知训练（pooling-aware training）教会编码器在向量被合并时保留重要的信息，帮助你在相同的存储预算下获得更好的检索效果。

在这篇文章中，我将解释池化的工作原理，使用 Sentence Transformers 训练一个用于压缩检索的模型，并探讨由此带来的质量与存储的权衡。通过 pool_factor=4，我正常训练的模型保留了未池化时 NanoBEIR nDCG@10 的 94.6%，同时文档向量减少了 4 倍。Token pooling 感知训练在相同存储预算下将保留率提升到了 97.8%！让我们看看如何做到。

本文建立在 Stefan Josef 的 Learn to Pool 以及 LightOn 的池化正则化研究所研究的池化感知训练之上，后者还探讨了针对性训练和自适应预算。在这里，我将着重于理解该方法，并通过一个简短、可复现的 Sentence Transformers 示例将其付诸实践。

目录

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

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

单向量嵌入模型将整个文档压缩为一个向量。而 多向量模型（如 ColBERT）则保留一组上下文相关的词元向量，允许查询的不同部分与文档的不同部分相匹配。

例如，一个关于续航持久的笔记本电脑的查询，可以在文档的一个部分找到关于笔记本电脑的证据，在另一个部分找到电池续航的证据。使用 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 的量化残差，一百万个文档则需要 12.8 GB。将每个文档池化为 50 个向量，在相同维度和精度下可将其降至 3.2 GB。

一个有用的观察是，许多文档向量可能是相似的。多个词元可能描述着密切相关的概念，这为将它们一起表示提供了机会。

Token Pooling

HierarchicalTokenPooling 使用 Ward 层次聚类 将相似的词元嵌入分组。从每个词元向量各自为一组开始，它反复合并那两组使组内向量离散程度增幅最小的组。这种离散程度用每个向量到其组均值的距离平方之和来度量，因此 Ward 方法倾向于形成由相似向量构成的紧凑组。

这些合并构建出一个规模逐渐增大的分组层次结构，使我们能够选择符合预算的压缩级别。随后我们用每组的均值来表示该组，并对得到的向量再次归一化，为 MaxSim 评分做好准备。

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

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor=4)

使用 pool_factor=4 时，文档大约保留其原始向量的四分之一，因此较长的文档可以按比例保留更多容量。默认情况下，第一个词元受到保护不被合并，且池化仅应用于文档嵌入，查询嵌入保持不变。

选择池化预算

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 个向量。在另一个语料库上，实际达到的平均值也可能变化。

词元池化感知训练

池化在普通的多向量模型上就已经能发挥作用了。困难出现在文档包含对检索很重要的区分信息，而合并其向量会使这种区分更难恢复的时候。原始的训练目标从未要求模型处理这些压缩表示。不过，你可以教会模型做到这一点。

如果你用过 Matryoshka 嵌套嵌入，那么这种训练原则会很熟悉。Matryoshka 在多个嵌入宽度上施加检索目标，教会较小的前缀保持可用。池化感知训练则将该目标应用于词元向量更少的表示。

以我的示例来说，每个训练批次贡献两个损失：

1. 使用原始文档嵌入进行检索。
2. 将文档压缩到大约原始向量四分之一之后进行检索。

我对它们取平均：

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

第一个损失项教模型使用其完整表征进行检索。第二项教它排列其词元嵌入，使有用的匹配能在池化后保留下来。两个目标更新同一个编码器，让它在压缩之前和之后都获得关于检索的反馈。

Ward 聚类会就哪些词元归为一组做出离散决策。训练期间，HierarchicalTokenPooling使用直通估计器将梯度从池化后的向量传回其成员词元，而不对聚类分配求导。这样编码器就能从使用池化表征时产生的检索错误中学习。

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。

梯度缓存让编码器能够处理较小的迷你批次，同时保留更大的一组批内负例。我的训练脚本使用 64 的迷你批次，训练批次大小为 256。

MultiVectorTokenPoolingLoss添加压缩目标。这里，None表示一个变体使用原始嵌入，而pooling提供 4 倍因子变体。两个权重对其损失取平均。池化变体复用已计算的词元嵌入，因此每增加一个池化器只会增加聚类和打分的工作，而无需为该变体单独运行一次编码器。

这些池化器也可以是自定义的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 上运行池化感知变体，时间限制为四小时。将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 核可以在不于内存中构建完整词元相似度矩阵的情况下对查询与文档进行打分，但它需要 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 个,同时保持查询向量不变。Mars 仍然获得最高的相似度分数,其次是 Jupiter 和 Venus,因此尽管存储的文档向量少得多,排序仍然得以保留。

你可以在推理时更改 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 倍因子目标。

两次运行都在相同的 120,000 条 MS MARCO 三元组上进行一个 epoch 训练,种子为 42,初始化和优化设置也相同。训练 token 上限为 300,短查询会扩展到至少 32 个 token 位置,以提供额外的上下文查询向量。两个实验之间只有损失不同,我没有进行超参数搜索。

我在 NanoBEIR@13 上评估了两个模型,报告平均 nDCG@10(%)。该图比较了普通训练与池化感知训练在池化因子为 2、3、4、6 时的表现,并附带未池化检索。表中还包括固定上限和自适应阈值,展示了训练收益如何迁移到其他池化设置。

对于自适应池化,我为每个编码器和 NanoBEIR 语料库分别校准了一个阈值,使用其未池化的文档向量来匹配上限为 64、32 和 16 时的平均向量数。短文档会使这些平均值低于相应的上限。

在训练所用的 pool_factor=4 下,nDCG@10 从 49.03 提升到 51.23,同时比未池化检索少存储 74.8% 的文档向量。这将分数保留率从 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 pooling（令牌池化）可以减少文档向量的数量，还可以像本博客中展示的那样进行池化感知的微调。
3. Matryoshka 令牌嵌入可以降低每个令牌向量的维度，还可以通过微调在更小的宽度下保持检索质量。

如果模型的默认维度大于 128，我也会更早地考虑 Matryoshka 截断和微调。这些方法可以组合使用，因此请结合你的搜索引擎和语料库评估最终的索引大小和检索质量。

MultiVectorMatryoshkaLoss 训练令牌嵌入，使其在截断后仍然有效。你可以将它与 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 和 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 和 Jamie Callan。Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup。
• Matryoshka 嵌入模型入门。

代码与模型

• 令牌池化感知训练示例与 README
• 完整训练脚本
• NanoBEIR 评估脚本
• Example checkpoints:
  • tomaarsen/mve-ModernBERT-base-token-pooling
  • tomaarsen/mve-ModernBERT-base-token-pooling-baseline
• 与 Sentence Transformers 兼容的预训练多向量模型
• 面向多向量模型的 Matryoshka 感知训练
• Fast-PLAID