sentence-transformers 多向量编码器 Late-Interaction 评分指南：ColBERT MaxSim、MeanMaxSim 与 XTR 全局 Top-k

sentence_transformers.multi_vector_encoder.scoring 是 sentence-transformers 多向量编码器（MultiVectorEncoder，即 ColBERT 风格模型）的评分模块，为训练损失提供 late-interaction 相似度计算：ColBERT 风格的 MaxSim、长度归一化的 MeanMaxSim，以及 XTR 风格的全局 top-k 评分。通过把其中一个函数（或配置好的可调用对象）作为损失函数的 similarity_fct 参数传入，可以在不改动损失逻辑的前提下切换评分策略。读完本文，你将掌握该模块全部 11 个公开 API 的形状约定、参数语义与底层原理，并能正确地在 MultiVectorMultipleNegativesRankingLoss、MultiVectorDistillKLDivLoss、MultiVectorMarginMSELoss 三种损失中使用它们。

模块定位：损失与评分解耦的 similarity_fct 接口

多向量编码器的损失（如 in-batch negatives 对比损失、蒸馏损失）只关心"拿到一批查询与文档的 token 级嵌入，返回一个分数张量"，而不关心分数具体怎么算。scoring 模块正是这个解耦点：所有评分函数共享同一套关键字约定（queries_embeddings、documents_embeddings、queries_mask、documents_mask），损失只负责调用并消费返回值。

模块的公开接口在 sentence_transformers/multi_vector_encoder/scoring/init.py 中统一导出，共 11 个符号，按功能分为四组：

分组	函数 / 类	输出形状	典型用途
ColBERT scoring	colbert_scores、colbert_scores_pairwise、colbert_kd_scores	(Q, Q*N) / (Q,) / (Q, N)	MaxSim 对比 / 成对 / 蒸馏
MeanMaxSim scoring	mean_colbert_scores、mean_colbert_scores_pairwise、mean_colbert_kd_scores	同上	长度归一化版本，默认 length_normalize=True
XTRScores	XTRScores、xtr_scores、xtr_scores_pairwise、xtr_kd_scores	(Q, Q*N) / (Q,) / (Q, N)	XTR 全局 top-k 评分
XTRKDScores	XTRKDScores	(Q, N)	可复用配置化的 XTR 蒸馏评分器

所有输入输出统一为 float32：即便传入 bf16/fp16 或整数量化嵌入，中间计算也会先提升为 float32（见下文 XTR 部分对整数嵌入的处理），保证分数精度与数值稳定性。

ColBERT scoring：MaxSim 晚期交互

ColBERT 风格的 MaxSim 评分是 ColBERTv2 检索模型的经典 late-interaction 打分：对查询中每个 token，取其与文档所有 token 相似度的最大值，再按查询 token 求和。三个函数分别面向三种损失接口，实现均位于 sentence_transformers/multi_vector_encoder/scoring/colbert.py。

colbert_scores：in-batch 负样本对比评分

对比评分是 in-batch negatives 训练的默认评分器（MultiVectorMultipleNegativesRankingLoss 的默认 similarity_fct）。其输入输出约定是理解整个模块的关键：

• 查询嵌入：(Q_query, q_tokens, dim)；
• 文档嵌入：(Q_doc, N, d_tokens, dim)，即按查询分组的堆叠文档——每个查询对应 N 个候选文档（一个正样本加若干负样本）；
• 返回：完整的 (Q_query, Q_doc * N) 分数矩阵，query-major 排序：scores[i, j*N + n] 表示查询 i 与第 j 个查询组中第 n 个文档的 MaxSim。当 Q_query == Q_doc 时，查询 i 的正样本恰好位于第 i*N 列。

实现上，文档轴按组逐个迭代（for j in range(N)），每次只让一个 (Q_query, Q_doc, q_tokens, d_tokens) 四维中间张量存活，从而把峰值内存限制在单组规模（见 colbert.py）。底层委托给 sentence_transformers/util/similarity.py 的 maxsim，其评分公式为 sum_i max_j (a_i . b_j)。

colbert_scores_pairwise：成对评分

面向匹配好的 (query_i, document_i) 对，接收 (batch_size, q_tokens, dim) 查询与 (batch_size, d_tokens, dim) 文档，返回长度为 batch_size 的一维 float32 向量。它是 colbert_scores 的成对变体，直接委托给 maxsim_pairwise。该函数是 MultiVectorMarginMSELoss 的默认评分器，可与 xtr_scores_pairwise 互换——两者都遵守统一的 similarity_fct 关键字约定。

colbert_kd_scores：知识蒸馏评分

蒸馏场景中，每个查询已由教师模型检索并打分了 N 个候选文档（通常是一个正样本加若干负样本）。函数接收 (batch_size, q_tokens, dim) 查询与 (batch_size, n_ways, d_tokens, dim) 堆叠文档，返回 (batch_size, n_ways) 的 MaxSim 分数，可直接与教师分数做 KL 蒸馏。实现是对每个 n_ways 列调用一次 maxsim_pairwise 再堆叠（见 colbert.py），因此它是 MultiVectorDistillKLDivLoss 的默认 similarity_fct。

公共参数

三个 ColBERT 函数共享同一套参数，语义与 maxsim / maxsim_pairwise 完全一致：

参数	默认值	说明
queries_embeddings / documents_embeddings	必填	list、np.ndarray 或 torch.Tensor
queries_mask	None	可选 (batch, q_tokens) 掩码，None 时由全零行自动推断为 padding
documents_mask	None	可选文档掩码，同上
chunk_elements	None	元素预算，控制分块打分的内存峰值；None 采用 maxsim 的 1 亿元素预算（fp32 下约 400 MB，bf16/fp16 减半）
length_normalize	False	是否除以查询真实 token 数（True 即 MeanMaxSim）

关于 chunk_elements 有一个容易踩坑的细节：预算覆盖的是填充后的嵌入加上评分中间张量。当查询侧很小时，padding 才是占内存的大头，只按中间张量预算会严重低估实际占用（源码注释明确提示 "Counting only the intermediate under-provisions"）。因此降低 chunk_elements 是裁剪训练显存的直接手段。

MeanMaxSim scoring：长度归一化的 MaxSim

MaxSim 是对每个查询 token 累加一个相似度，分数尺度随查询长度增长，不同查询之间的分数不可比。MeanMaxSim 改为除以每个查询的真实 token 数（掩码之和），分数落在单 token 相似度量级（归一化嵌入下约 [-1, 1]），跨查询长度可比。由于除数是逐行常数，同一查询内的排序完全不变。

mean_colbert_scores、mean_colbert_scores_pairwise、mean_colbert_kd_scores 分别是三个 ColBERT 函数的归一化封装，唯一区别是 length_normalize 默认值为 True；显式传 False 即恢复为普通 ColBERT 函数，因此整个 ColBERT 家族共享相同关键字（见 colbert.py）。

训练与评估的一致性：model.similarity_fn_name

文档特别强调：用 MeanMaxSim 训练，就必须让评估也用 MeanMaxSim 打分。做法是在模型上设置：

model.similarity_fn_name = "meanmaxsim"

MultiVectorEncoder.similarity_fn_name 属性（见 model.py）决定 similarity() 与 similarity_pairwise() 使用 maxsim 还是 mean_maxsim（首次访问且未显式设置时默认 "maxsim"）。评估器（如 information retrieval 评估器）在调用时通过 SimilarityFunction.to_similarity_fn(model.similarity_fn_name) 解析评分函数（见 evaluation/information_retrieval.py），模型卡模板也会把 similarity_fn_name 写入卡片（见 model_card_template.md），从而保证训练、评估、模型卡三方口径一致。

注意：similarity_fn_name 只支持 "maxsim" 与 "meanmaxsim"。若试图设置为 "xtr"，模型会直接抛出 ValueError，并解释原因：XTR 是训练期评分，其全局 top-k 依赖 batch 组成，逐对分数没有定义良好的语义（见 model.py）。

XTRScores：XTR 风格全局 top-k 评分

XTR（eXtendable Token Retrieval，Lee et al. 2023）与 ColBERT 的关键区别在于 token 检索的视野：MaxSim 是"查询 token 对单个文档内所有文档 token 取最大"；XTR 则是"查询 token 对 batch 内所有文档的所有 token 全局取 top-k"，模拟从索引中检索的行为。实现位于 sentence_transformers/multi_vector_encoder/scoring/xtr.py。

xtr_scores：全局 top-k 对比评分

签名与 colbert_scores 相同（(Q, q_tokens, dim) 查询、(Q, N, d_tokens, dim) 堆叠文档、返回 query-major 的 (Q, Q*N)），额外增加 top_k 参数：

• top_k：每个查询 token 在全部 Q*N 个文档 token 中保留的匹配数，默认 256，必须是正整数（传入 0、负数或布尔值会抛出 ValueError，有对应测试用例 test_xtr_scores_rejects_non_positive_topk 覆盖，见 tests/multi_vector_encoder/test_model.py）。

内部流程（xtr.py）：

1. 整数（量化）嵌入先提升为 float32——否则整数分数网格无法用 torch.finfo 的 dtype 最小值做掩码填充；
2. matmul 得到 (Q, q_tokens, Q*N, d_tokens) 分数网格，用 masked_fill 把 padding token 置为 dtype 最小值；
3. 对每个查询 token，跨所有文档 token 做 topk 取全局 top-k；
4. 每个分数 = 查询检索到的各 token 最大值之和 ÷ Z，其中 Z = 至少检索到该文档一个真实 token 的查询 token 数（Lee et al. 2023 论文 eq. 5）。源码注释明确指出这与 PyLate / PrimeQA 的实现不同——后者除以的是"正 per-token 最大值计数"，会放大全负行 1000 倍以上，并在 top_k 覆盖整个 token 池时产生分歧；
5. 未检索到任何 token 的文档按 0 补全，完全被掩码的文档则写入哨兵值（_fill_empty_document_scores，约 -1e9），确保排在所有真实文档之后。

由于全局 top-k 在 batch 的 pooled token 上运行，当 top_k 小于 token 池大小时，同一对的分数会随 batch 组成变化——这正是 XTR 与 MaxSim 的本质差异，也是它不能作为模型级 similarity 的原因。

xtr_scores_pairwise：成对 XTR 评分

返回长度为 batch_size 的一维 float32 分数。内部把 3D 文档嵌入扩展为 (Q, N=1, d_tokens, dim) 后调用 xtr_kd_scores 再取对角（见 xtr.py）。作为 MultiVectorMarginMSELoss 的 similarity_fct 传入即可把该损失的成对蒸馏从 MaxSim 切换为 XTR（有测试 xtr_scores_pairwise 作为 drop-in 评分器的用例，见 tests/multi_vector_encoder/losses/test_misc.py）。

xtr_kd_scores：XTR 蒸馏评分

与 xtr_scores 相同的全局 top-k 算法，但只返回每个查询自己的 N 路文档分数 (Q, N)，匹配 MultiVectorDistillKLDivLoss 期望的接口。实现是调用 xtr_scores 后按 query-major 布局 gather 出对角组（见 xtr.py）。

配置类 XTRScores

裸函数每次调用都要重传 top_k / chunk_elements，XTRScores 类把它们存起来，作为可复用的配置化 callable：

from sentence_transformers.multi_vector_encoder.scoring import XTRScores

scorer = XTRScores(top_k=128, chunk_elements=50_000_000)
# 之后可以直接作为 similarity_fct 使用，无需每次重传参数

类实现（xtr.py）额外提供 get_config_dict()，返回 {"top_k": ..., "chunk_elements": ...}，使超参数可序列化进模型卡与训练日志。测试验证了 XTRScores(top_k=16) 与 XTRScores(top_k=256) 序列化后可以区分（见 tests/multi_vector_encoder/losses/test_misc.py）。

此外，xtr_scores 的热路径可以直接用 torch.compile 编译加速：文档与源码均给出 similarity_fct=torch.compile(xtr_scores) 的用法（见 xtr.py）。

XTRKDScores：可复用的 XTR 蒸馏评分器

XTRKDScores 继承自 XTRScores，唯一区别是 __call__ 调用的是 xtr_kd_scores 而非 xtr_scores，因此输出为 (Q, N) 而非 (Q, Q*N)（见 xtr.py）。它面向蒸馏场景的典型用法：

from sentence_transformers.multi_vector_encoder.scoring import XTRKDScores

loss = MultiVectorDistillKLDivLoss(model, similarity_fct=XTRKDScores(top_k=64))

序列化后显示为 XTRKDScores(top_k=64, chunk_elements=None)，同样可写进模型卡与日志（对应测试见 tests/multi_vector_encoder/losses/test_misc.py）。

在训练损失中接入评分函数

scoring 模块的价值在于"评分可插拔"：三种多向量损失各自持有默认评分器，传入其他函数即可切换策略而不改损失代码。各损失默认值与推荐的替换见下表（默认值定义分别见 multiple_negatives_ranking.py、distill_kl_div.py、margin_mse.py）：

损失	默认 similarity_fct	输出形状	推荐的替代
MultiVectorMultipleNegativesRankingLoss	colbert_scores	(Q, Q*N)	XTRScores、mean_colbert_scores
MultiVectorDistillKLDivLoss	colbert_kd_scores	(Q, N)	XTRKDScores、mean_colbert_kd_scores
MultiVectorMarginMSELoss	colbert_scores_pairwise	(Q,)	xtr_scores_pairwise、mean_colbert_scores_pairwise

一个完整的 XTR 对比训练示例（以 MS MARCO 对比训练为蓝本，参考 examples/multi_vector_encoder/training/msmarco/training_contrastive.py 的数据组织方式）：

from datasets import Dataset
from sentence_transformers import MultiVectorEncoder, MultiVectorEncoderTrainer
from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss
from sentence_transformers.multi_vector_encoder.scoring import XTRScores

model = MultiVectorEncoder("answerdotai/ModernBERT-base")

train_dataset = Dataset.from_dict(
    {
        "query": ["What is the capital of France?", "Who painted the Mona Lisa?"],
        "positive": ["Paris is the capital of France.", "Leonardo da Vinci painted the Mona Lisa."],
        "negative": ["Berlin is the capital of Germany.", "Van Gogh painted The Starry Night."],
    }
)

# 切换到 XTR 全局 top-k 评分
loss = MultiVectorMultipleNegativesRankingLoss(model, similarity_fct=XTRScores(top_k=128))

trainer = MultiVectorEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)
trainer.train()

若用 MeanMaxSim 训练，还需同步设置 model.similarity_fn_name = "meanmaxsim"，保证评估与模型卡与训练打分一致：

from sentence_transformers.multi_vector_encoder.scoring import mean_colbert_scores

model.similarity_fn_name = "meanmaxsim"
loss = MultiVectorMultipleNegativesRankingLoss(model, similarity_fct=mean_colbert_scores)

与 loss 内其他参数的联动

• scale（温度）：MultiVectorMultipleNegativesRankingLoss 的 scale 是 1 / temperature。MaxSim 是无界和（量级约 [0, num_query_tokens]），所以 scale=1.0（默认，与 PyLate 一致）是合理起点；改用 MeanMaxSim 后每个分数被查询 token 数归一，可以从约等于平均查询长度的 scale 起步。
• score_mini_batch_size：对查询轴分块打分，控制打分阶段瞬时显存；若要同时切分文档轴，可以在评分器里绑定预算，如 similarity_fct=partial(colbert_scores, chunk_elements=...)（见 multiple_negatives_ranking.py）。
• mini_batch_size：嵌入阶段按行分块合并文档列，单条超长文档只加宽自己的块，但激活会保留到反向传播，因此它限制的是 padding 浪费而非总内存。

底层原理与测试佐证

评分函数的最底层依赖 sentence_transformers/util/similarity.py 的 maxsim / maxsim_pairwise 与 mean_maxsim 系列，它们实现了 chunk 分块打分：文档按元素预算贪心打包（_chunk_ranges），每个 chunk 独立 padding，长尾 outlier 只撑大自己的 chunk；查询 token 归约为 sum（掩码置 0），文档 token 归约为 max（掩码置 dtype 最小值，防止 padding 在负相似度上胜出）；空文档写入 -1e9 哨兵分。SimilarityFunction 枚举（"maxsim"、"meanmaxsim"）则负责把模型卡中的名称解析为对应函数。

仓库测试对评分语义有系统性验证，可作为行为契约参考：

• 对角线等价：colbert_scores 的对角等于 colbert_scores_pairwise，XTR 版本同理（在 top_k 覆盖全 token 池时成立），见 tests/multi_vector_encoder/losses/test_misc.py；
• 归一化一致性：mean_colbert_* 是 colbert_* 的长度归一化形式，length_normalize=False 恢复原函数，见 tests/multi_vector_encoder/losses/test_misc.py；
• 健壮性边界：top_k 非正数报错、整数嵌入自动提升 float32、查询 padding 掩码自动推导，见 tests/multi_vector_encoder/test_model.py；
• 配置序列化：XTRScores(top_k=16) 与 XTRScores(top_k=256) 在模型卡与日志中可区分，见 tests/multi_vector_encoder/losses/test_misc.py。

小结

sentence_transformers.multi_vector_encoder.scoring 用一组小而完整的 API 覆盖了多向量编码器的全部主流评分范式：ColBERT MaxSim（对比 / 成对 / 蒸馏三形态）、长度归一化的 MeanMaxSim（配合 similarity_fn_name = "meanmaxsim" 保持训练评估一致）、XTR 全局 top-k（裸函数与可配置类双形态）。它们通过统一的 similarity_fct 接口与三种多向量损失解耦，并共享 chunk_elements 内存预算机制与 float32 数值稳定保证。实际使用时的要点可归纳为四句话：对比训练用 colbert_scores / XTRScores，成对蒸馏用 colbert_scores_pairwise / xtr_scores_pairwise，列表蒸馏用 colbert_kd_scores / XTRKDScores，长度归一化时给模型设置 similarity_fn_name = "meanmaxsim"。