面向多向量嵌入模型的词元池化感知训练

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

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

本文将解释池化如何工作，用 Sentence Transformers 训练一个用于压缩检索的模型，并探讨由此产生的质量与存储权衡。在 pool_factor=4 时，我正常训练的模型保留了未池化 NanoBEIR nDCG@10 的 94.6%，同时文档向量减少到 1/4。词元池化感知训练在相同存储预算下将保留率提高到 97.8%！让我们看看如何做到。

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

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

目录#

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

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

单向量嵌入模型将整个文档压缩为一个向量。ColBERT 等多向量模型则保留一组上下文词元向量，使查询的不同部分能够匹配文档的不同部分。

例如，一个关于「电池续航长的笔记本电脑」的查询，可以在文档的一部分找到关于笔记本电脑的证据，在另一部分找到关于电池续航的证据。在 MaxSim 评分中，每个查询向量选择其最佳匹配的文档向量，然后将这些得分相加：

MaxSim(Q, D) = Σ_{q∈Q} max_{d∈D} qᵀd

即使使用 Fast-PLAID 支持的 2 位或 4 位残差量化，每文档存储数百个向量也会迅速累积。在每残差坐标 4 位的情况下，200 个 128 维向量需要每文档 12.8 KB 的量化残差，即一百万文档需要 12.8 GB。将每个文档池化到 50 个向量，在相同维度和精度下可减少到 3.2 GB。

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

词元池化#

HierarchicalTokenPooling 使用 Ward 层次聚类（Ward hierarchical clustering）对相似的词元嵌入进行分组。它从每个词元向量一组开始，反复合并那些使组内向量离散程度增加最小的两组。这种离散程度用每个向量到其组均值的平方距离之和来衡量，因此 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)

# 应用池化，目标为每文档平均 16 个向量
document_embeddings = model.encode_document(documents_to_index, token_pooling=pooling)

校准只需要文档嵌入，并返回一个可在索引时复用的阈值。目标是平均值，因此单个文档可能得到多于或少于 16 个向量。在另一语料上，实际达到的平均值也可能变化。

词元池化感知训练#

池化在普通多向量模型上已经有用。困难出现在文档包含对检索重要的区分、但合并其向量使该区分更难恢复时。原始训练目标从未要求模型使用这些压缩表示。但是，你可以教会它正是这样做。

如果你用过套娃嵌入（Matryoshka embeddings），训练原理会感到熟悉。Matryoshka 在多个嵌入宽度上应用检索目标，教会较小的前缀仍然有用。池化感知训练则对具有更少词元向量的表示应用目标。

在我的示例中，每个训练批次贡献两个损失：

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

我对它们取平均：

L = 1/2 · L_unpooled + 1/2 · L_pooled

第一项教会模型用完整表示进行检索。第二项教会它安排词元嵌入，使有用的匹配在池化后仍然保留。两个目标更新同一个编码器，为它提供压缩前后检索的反馈。

Ward 聚类会就哪些词元属于一组做出离散决策。训练期间，HierarchicalTokenPooling 使用直通估计器（straight-through estimator）将梯度从池化向量传回其成员词元，而不对聚类分配做微分。编码器随后可以从其池化表示所犯的检索错误中学习。

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。

梯度缓存允许编码器处理小批量（mini-batch），同时保留更大的批内负样本（in-batch negatives）集合。我的训练脚本使用 64 的 mini-batch 和 256 的训练批次大小。

MultiVectorTokenPoolingLoss 增加压缩目标。这里的 None 表示一个变体使用原始嵌入，而 pooling 提供 factor-4 变体。两个权重对它们的损失取平均。池化变体复用已计算的词元嵌入，因此每增加一个池化器，都会增加聚类和评分工作，但不需要为该变体单独进行编码器前向传递。

池化器也可以是自定义的 BaseTokenPooling 实现，只要它们能通过输入嵌入保留梯度。你可以让训练目标适应你想要部署的表示。

将 factor-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 和上传模型的词元登录：

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 作为 secret 传入，可让脚本在训练完成后将训练好的模型上传到你的 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')

# 作为对比，同一模型不使用词元池化时：
# [18, 21, 19]
# tensor([[16.6115, 27.6060, 24.1766]], device='cuda:0')

池化器在编码器前向流水线之后压缩文档嵌入，将总数从 58 个向量减少到 16 个，同时保持查询嵌入不变。火星仍然获得最高相似度得分，其次是木星和金星，因此尽管存储的文档向量少得多，排名仍得以保留。

你可以在推理时更改 pool_factor，用同一个训练好的模型探索不同的存储预算。对于重复使用，你还可以在初始化 MultiVectorEncoder 时传入 token_pooling=pooling，使其成为运行时默认值。池化器的任务仍然控制哪些输入被压缩（例如仅文档）。

评估结果#

为进行研究，我从 answerdotai/ModernBERT-base 训练了两个模型，每个都有全新的 128 维词元投影：

• tomaarsen/mve-ModernBERT-base-token-pooling-baseline，仅使用缓存对比损失。
• tomaarsen/mve-ModernBERT-base-token-pooling，使用缓存对比损失，并用上述未池化和 factor-4 目标包装。

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

我在 NanoBEIR@13 上评估两个模型，报告平均 nDCG@10（%）。图中比较了普通训练和池化感知训练在 pool factor 为 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. 词元池化以减少文档向量数量，可选地使用本文展示的池化感知微调。
3. 套娃词元嵌入以降低每个词元向量的维度，可选地通过微调在较小宽度下保持检索质量。

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

MultiVectorMatryoshkaLoss 训练词元嵌入在截断后仍然有用。你可以用 MultiVectorTokenPoolingLoss 包装它，同时训练更小的维度和更少的文档向量，并在不同压缩设置之间复用编码器的嵌入。

尝试完整的训练示例，选择对你的应用重要的预算，并将其与匹配的普通损失运行进行比较。附带的评估脚本测量未池化检索、64/32/16/8/4 向量的上限，以及 2 到 6 的 pool factor，因此你可以为自己的模型绘制质量与存储的权衡曲线。

只需在训练目标中增加几行代码，更小的文档表示就成为模型主动学习使用的东西。这为让多向量检索更实用开辟了令人兴奋的空间！

参考文献#

• 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.
• Introduction to Matryoshka Embedding Models.

代码与模型#

• 词元池化感知训练示例和 README
• 完整训练脚本
• NanoBEIR 评估脚本
• 示例检查点：
• tomaarsen/mve-ModernBERT-base-token-pooling
• tomaarsen/mve-ModernBERT-base-token-pooling-baseline
• 与 Sentence Transformers 兼容的预训练多向量模型
• 多向量模型的 Matryoshka 感知训练
• Fast-PLAID