多向量嵌入模型聚合意识训练提升检索质量

多向量嵌入模型的标记聚合意识训练：在压缩存储与检索质量间寻求平衡

多向量嵌入模型（如 ColBERT）为每个查询令牌提供了在文档中匹配特定部分的机会，但这种细粒度表示带来了巨大的存储开销。一个包含数百万文档的数据集可能产生数以亿计的向量。虽然通过合并相似符号向量可以减小表示体积，但过度压缩往往导致关键信息丢失。

“聚合意识训练”（Pooling-aware training）旨在解决这一矛盾。该方法教导编码器在向量被合并时保留重要信息，从而在相同的存储预算下提升检索性能。基于 Stefan Josef 的“学习共享”和 LightOn 的相关研究，Tom Aarsen 在 Hugging Face 上展示了如何利用 Sentence Transformers 实现这一技术。

实验数据显示，在 pool_factor=4（即文档向量数量减少至原来的四分之一）的设置下，经过普通训练的模型保留了未合并模型 94.6% 的 NanoBEIR nDCG@10 分数；而采用聚合意识训练后，相同存储预算下的检索质量提升至 97.8%。

为什么多向量检索需要如此多的向量

单向量嵌入模型将整个文档压缩为一个向量，而多向量模型则保留一组上下文符号向量，允许查询的不同部分分别匹配文档的不同片段。例如，针对“长电池寿命笔记本电脑”的查询，可以在文档某处找到关于“笔记本电脑”的证据，在另一处找到关于“电池耐用性”的证据。MaxSim 评分机制让每个查询向量选择最匹配的文档向量并累加得分。

即使使用 Fast-PLAID 支持的 2 位或 4 位残差量化，存储成本依然高昂。以 128 维、每文档 200 个向量为例，若每个残差坐标为 4 位，一百万份文件需占用约 12.8GB 空间。若将每个文档的向量数整合至 50 个，存储空间可降至 3.2GB。鉴于许多文档向量描述的概念高度相关，将它们合并代表是一个可行的优化方向。

代币聚合机制

HierarchicalTokenPooling 模块利用 Ward 层次聚类算法将相似的代币嵌入组合在一起。该算法反复合并两个组，使得组内向量的扩散最小化——扩散程度由每个向量到其组平均值的平方距离总和衡量。Ward 方法倾向于形成紧凑且内部相似度高的组。

选定压缩级别后，每组用平均值表示，并对结果向量进行归一化处理以备 MaxSim 评分。压缩比例可与文档长度成正比：

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor=4)

当 pool_factor=4 时，文档保留大约四分之一的原始向量，较长文档因此能保持更大的容量。默认情况下，第一个代币免受合并，且聚合仅应用于文档嵌入，不改变查询嵌入。

选择共同的预算控制方式

HierarchicalTokenPooling 提供三种控制模式：

当文档长度差异较大时，这些区别至关重要。pool_factor 使表示规模随内容增长，而 distance_threshold 则在需要保留更明显嵌入特征时增加向量数量。

校准适应性聚合

适应性聚合使用共享距离值，对相似嵌入进行更积极的合并，同时为内容独特的文档保留更多向量。此方法适用于任何多向量模型，包括未经过聚合意识训练的模型。

无需手动选择阈值，可通过代表性样本校准目标平均向量计数：

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)

document_embeddings = model.encode_document(documents_to_index, token_pooling=pooling)

校准过程仅需文档嵌入，返回的值可在索引阶段重用。目标是维持平均值，单个文档接收的向量数量可能多于或少于 16 个。

标记聚合意识培训原理

在未专门训练的多向量模型上，聚合虽有用，但当文档包含重要细微差别时，合并向量可能导致这些细节难以恢复。原始训练目标并未要求模型适应压缩后的表示。然而，可以通过训练让模型学会这一点。

这与 Matryoshka Embeddings 的原理类似：后者在不同嵌入宽度上应用检索目标，教导较小的前缀保持有用性；聚合意识训练则将目标应用于具有较少符号向量的表示。

在每个训练批次中，损失函数计算两部分：

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

这两个术语的平均值构成总损失。第一项教导模型利用完整表示进行检索，第二项则教导模型安排符号嵌入，使有用的匹配在压缩后仍能存活。两个目标更新相同的编码器，使其了解压缩前后的检索反馈。

在训练过程中，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 提供基础检索目标，教导模型给相关文档打分高于批内负面文档。梯度缓存允许编码器处理小批量数据，同时保留较大的批内负样本集合。示例脚本使用 64 的小批量大小和 256 的训练批次大小。

MultiVectorTokenPoolingLoss 增加了压缩目标：None 表示使用原始嵌入，聚合器提供因子-4 变体。两者权重平均。聚合变体重用已计算的代币嵌入，因此额外的聚合器主要增加集群和评分工作，无需单独的编码器前向传播。

完整的训练脚本支持本地运行或云端 GPU 训练：

# 本地运行
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

# 云端 Hugging Face 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 作为密钥允许脚本在训练结束后上传模型。若要运行基线对比，需在脚本 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 上运行，速度比 SciPy 快 30-150 倍。

使用训练好的模型

加载聚合意识模型并使用压缩后的文档嵌入进行检索。注意内核设置不随模型保存，推理时需重新指定：

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])
print(model.similarity(query_embeddings, document_embeddings))

池化器在编码器前向管道之后压缩文档嵌入，将总向量数从 58 减至 16，查询嵌入保持不变。尽管存储的文档向量大幅减少，Mars 仍获得最高相似度得分，排序得以保留。

用户可在推理时更改 pool_factor 以探索不同的存储预算。对于重复使用，初始化 MultiVectorEncoder 时传入 token_pooling=pooling 可将其设为运行时默认值。

评估结果

研究人员基于 answerdotai/ModernBERT-base 训练了两个模型，均带有新的 128 维令牌投影：

• tomaarsen/mve-ModernBERT-base-token-pooling-baseline：仅使用缓存对比损失。
• tomaarsen/mve-ModernBERT-base-token-pooling：使用包裹了未池化和因子-4 目标的缓存对比损失。

两次运行均在相同的 120,000 MS MARCO 三元组上进行一个 epoch 训练，种子为 42，初始化和优化设置一致。训练令牌上限为 300，短查询扩展至至少 32 个令牌位置。唯一区别在于损失函数，未进行超参数搜索。

在 NanoBEIR@13 数据集上的评估显示，聚合意识训练在不同池化因子（2, 3, 4, 6）及固定上限和自适应阈值下，均显著提升了检索质量。