跳到主要内容

数据截至 (上游 commit 5cc889d47547)

05 · 稀疏与多向量变体(SPLADE / ColBERT)

这一章讲什么: 本库除稠密双塔外的两个特化家族。它们共享第 1 章的模块链与第 2 章的训练骨架,差别集中在「输出的向量是什么形态」与「相似度怎么算」。读完你会知道什么时候该放弃「一句一向量」。


1. 它要解决的小问题

稠密双塔的两个已知短板:

  • 词项不可解释、精确匹配弱。 一个 768 维稠密向量看不出「含哪些词」;对型号、编号、罕见专有名词这类词面必须命中的场景,稠密向量常常不如倒排索引。
  • 一句一向量是有损压缩。 长文档的所有信息压进一个向量,细粒度匹配(query 的某个词命中 doc 的某个词)被平均掉了。

两个变体分别回应:稀疏向量把「维度」变回「词项」,多向量干脆不压缩


2. SparseEncoder:SPLADE 稀疏向量

2.1 思路:维度即词项

SPLADE(SParse Lexical AnD Expansion)的洞察:MLM 模型(如 BERT)的输出头本来就为每个 token 位置输出词表大小的 logits——「这个位置最可能是哪个词」。把序列上所有位置的这个词表分布聚合(取 max)起来,就得到一个词表大小的向量:第 i 维的值 ≈ 「词 i 对这段文本的相关度」。

这个向量是稀疏的(绝大多数维是 0),而且天然带词扩展能力:文本里没出现的词也可以有非零权重(MLM 学到的联想)。于是它直接兼容倒排索引——检索时只需遍历 query 非零维对应的倒排表。

2.2 模块链与池化

默认链(sparse_encoder/model.py:1062_load_default_modules):骨干是 ForMaskedLM 架构时拼 Transformer(transformer_task="fill-mask") + SpladePooling(pooling_strategy="max")(:1108-1123)。

SpladePooling.forward(sparse_encoder/modules/splade_pooling.py:66 起)的公式很短:

token_embeddings = MLM logits,形状 (batch, seq_len, vocab)
→ 乘 attention_mask
→ relu → log1p # log(1 + relu(logits)),压掉长尾
→ 沿 seq_len 取 max(或 sum)
→ sentence_embedding,形状 (batch, vocab)

真实代码在 splade_pooling.py:104-140(分块沿序列维处理以控内存,OOM 时会提示你调小 chunk_size,:129-136)。激活函数可选 log1p_relu(再套一层 log1p,OpenSearch 系模型的做法,docstring :33-39)。FA2 展平输入走 _forward_flattened(:143),逐段变换——注释说明原因:对整个扁平 logits 先变换会把峰值显存翻倍,而那正是 SPLADE 训练的瓶颈(:146-148)。

输出可以直接 .to_sparse() 成 CSR 稀疏张量(sparse_encoder/model.py:768-769),默认相似度是点积(sparse_encoder/model.py:818-820)——因为权重本身已带尺度,归一化反而丢信息。

2.3 可解释性:decode 回词

稀疏向量最大的福利是可解码回人类可读的词。SparseEncoder.decode(sparse_encoder/model.py:1444)把向量的非零维翻成 (token, weight) 列表并按权重降序——调试「模型到底学到了什么关联」全靠它。另有 SparseEncoder.sparsity(sparse_encoder/model.py:1229)统计一批嵌入的稀疏度。

2.4 训练:效果损失 + 稀疏度正则

稀疏模型训出来「太密」就失去意义——倒排索引的开销正比于非零维数。所以 SpladeLoss(sparse_encoder/losses/splade.py:16)是两项之和:

作用实现
主损失(如 SparseMultipleNegativesRankingLoss)检索效果splade.py:143-153,内部损失任意换
文档正则(默认 FlopsLoss)× λ_d把文档嵌入压稀疏splade.py:154-159
查询正则(可选)× λ_q把查询嵌入压稀疏splade.py:161-163

FlopsLoss(sparse_encoder/losses/flops.py:11)的公式是「batch 均值向量的 L2 平方」:sum(mean(embeddings, dim=0) ** 2)(flops.py:50-57)——均值平方惩罚的是「很多样本在同一维上都有小值」,逼模型把质量集中到少数维。threshold 参数(l0 masking)跳过已经足够稀疏的样本。注意它不能单独当损失用:forward 直接 raise NotImplementedError(flops.py:45-48),只通过 compute_loss_from_embeddings 被 SpladeLoss 调用。

SpladeLoss.forward 返回 dict 而非标量(splade.py:138-166)——各分量由 trainer 的 track_loss_components 分别记录日志(第 2 章 §2.1),你能看到「效果项」与「稀疏项」各自的曲线。

2.5 另一条路:CSR / SparseAutoEncoder

骨干不是 MLM 时,默认链退化为 Transformer(feature-extraction) + Pooling(mean) + SparseAutoEncoder(sparse_encoder/model.py:1125-1131 起)——先用普通池化出稠密向量,再用一个自编码器把它扩成稀疏表示。这是与 SPLADE 不同的稀疏化路线,本章不展开。


3. MultiVectorEncoder:ColBERT 迟交互

3.1 思路:不池化,留着 token 级向量

ColBERT 在双塔与 CrossEncoder 之间取了一个中间点:

  • 编码仍是双塔式:query 和 document 各自独立过模型——所以文档侧仍可离线预计算。
  • 但不池化:每个 token 保留一个向量(通常投影到 128 维)。
  • 打分用 MaxSim(迟交互):对每个 query token,取它与文档所有 token 相似度的最大值,再对 query token 求和。

maxsim 的实现与形状约定在 sentence_transformers/util/similarity.py:376(docstring 给出公式 sum_i max_j (a_i · b_j));默认相似度名就是 "maxsim",另有 "meanmaxsim"(除以 query 真实 token 数,消除长度依赖,multi_vector_encoder/model.py:105-106)。

3.2 直觉对照

双塔ColBERTCrossEncoder
每条文本的表示1 个向量T 个向量(每 token 一个)无(逐对现算)
交互时机编码后(向量点积)编码后(token 级 max-sum)编码中(逐层交叉注意力)
文档可预计算是(但存 T×128)
精度

3.3 模块链:投影 + 归一化 + 掩码

裸骨干的默认链(multi_vector_encoder/model.py:1195_load_default_modules,三条路径在 :1204-1215 的 docstring):普通 transformer 拼 Transformer + Dense(投影到 128) + Normalize(逐 token) + MultiVectorMask;Stanford-NLP 原版 ColBERT 存档(HF_ColBERT)会读其 metadata 复刻特殊 token 与长度配置,并从仓库根部的 linear.weight 加载投影(:1207-1210);transformers 原生 *ForRetrieval(ColPali/ColQwen2 等)则投影归一化都在模型内部,只补一个掩码(:1249-1268)。

链尾两个模块的职责:

  • Normalize(module_input_name="token_embeddings"):逐 token L2 归一(normalize.py:14-21 的 docstring 明说这就是给 MaxSim 准备的——归一化后内积即余弦,分数落在 [−1,1])。
  • MultiVectorMask(multi_vector_encoder/modules/multi_vector_mask.py:18):决定哪些 token 参与打分。三类规则(forward,:124-152):
    1. 沿用 attention_mask(去掉 padding);
    2. skiplist_words(标点等)对应的 token id 从文档侧剔除——词在构造时给、模型加载后由 on_model_ready 钩子解析成 id(:89-122,解析失败的词被丢弃并警告,避免误删真 [UNK] token);
    3. query 侧把 query_expansion_positions(ColBERT 的 [MASK] 扩充位)强制算入(:136-140)。

3.4 训练与打分族

多向量的损失是稠密损失的同构改写(multi_vector_encoder/losses/):MultipleNegativesRankingLossCachedMultipleNegativesRankingLossMarginMSELossDistillKLDivLoss 各有一份 MaxSim 版,共享打分函数族(multi_vector_encoder/scoring/colbert.py):

函数场景形状
colbert_scores(:97)对比训练(in-batch 负例)(Q_query, Q_doc×N) 全矩阵
colbert_scores_pairwise(:65)逐对打分(MarginMSE 用)(batch,)
colbert_kd_scores(:10)listwise 蒸馏教师分数(batch, n_ways)

colbert_scores 有个内存设计:文档轴按组循环,任何时刻只有一组的 (Q_query, Q_doc, q_tokens, d_tokens) 四维中间张量是活的(docstring :113-122);chunk_elements 预算(默认 1 亿元素,util/similarity.py:233)控制这个峰值。另有 XTR 变体(scoring/xtr.py)只让 query 的 [CLS] 与文档 token 交互,本章不展开。

3.5 成本账

ColBERT 的代价集中在存储:每文档 T 个 128 维向量,比单向量大两个数量级。缓解手段库内自带:token 池化模块(multi_vector_encoder/modules/token_pooling.pyBaseTokenPooling/HierarchicalTokenPooling,:69/:223)把相似 token 向量聚类合并,以及 query 侧用固定 query_length + query expansion 把 query 长度钉死(base/modules/transformer.py:729-741 的参数文档)。


4. 图示:三种向量形态一张图

输入: "柏林人口多少" (4 个 token)

【稠密双塔】 【稀疏 SPLADE】 【多向量 ColBERT】
Transformer MLM Transformer Transformer
│ │ │
Pooling 逐位置词表 logits Dense 投影到 128
(mean) log1p(relu)·max Normalize(逐 token)
│ │ │
1×768 向量 1×|V| 稀疏向量 4×128 向量组
│ │ │
余弦相似 点积(可倒排) MaxSim(逐 token max 再求和)

怎么读这张图: 三列共享同一个「Transformer 起头」;分叉只在「token 级输出怎么聚合成可检索的表示」。这正对应模块链哲学——换的是链尾,不是骨架。


5. 关键细节 / 坑

  • SPLADE 路线锁定 MLM 骨干。architectures 是否以 ForMaskedLM 结尾探测(sparse_encoder/model.py:1108);给非 MLM 模型会静默走 CSR 路线——两条路线的模型行为完全不同,加载时留意日志(:1121 vs :1126)。
  • 稀疏向量默认用点积,别随手换余弦。 权重幅度本身携带信息(sparse_encoder/model.py:818-820 的默认值),归一化会改变检索语义。
  • SpladePooling 是显存大户。 中间张量是 (batch, seq, vocab)——vocab 动辄 3 万;OOM 时按警告提示调小 chunk_size(splade_pooling.py:127-136)。
  • ColBERT 的 query/document 不对称是配置出来的。 query_length/document_length/query_expansion 都住在 Transformer 模块的配置里并随存档保存(base/modules/transformer.py:751-761config_keys)——训练与推理必须用同一份,换模型时不要把别的模型的这些值带过来。
  • MultiVectorMask 的 skiplist 在模型就绪时才解析成 token id(multi_vector_mask.py:89-90on_model_ready);手动 model.append(MultiVectorMask(...)) 不会触发这个钩子,要自己调 resolve_with_tokenizer(注释,module.py:389-399)。
  • 多向量训练显存随四维打分张量爆炸:(Q, D, q_tokens, d_tokens)。先调小 chunk_elements,再考虑减小 batch(colbert_scores docstring,colbert.py:113-122)。

6. 代码地图(本章)

主题文件路径符号名
稀疏模型主体sentence_transformers/sparse_encoder/model.pySparseEncoderencodedecodesparsity
SPLADE 池化sentence_transformers/sparse_encoder/modules/splade_pooling.pySpladePooling.forward_forward_flattened
SPLADE 组合损失sentence_transformers/sparse_encoder/losses/splade.pySpladeLoss.forward
稀疏度正则sentence_transformers/sparse_encoder/losses/flops.pyFlopsLoss.compute_loss_from_embeddings
稀疏版 MNRLsentence_transformers/sparse_encoder/losses/sparse_multiple_negatives_ranking.pySparseMultipleNegativesRankingLoss
多向量模型主体sentence_transformers/multi_vector_encoder/model.pyMultiVectorEncoder_load_default_modules
MaxSim 打分sentence_transformers/util/similarity.pymaxsimmaxsim_pairwise
ColBERT 打分族sentence_transformers/multi_vector_encoder/scoring/colbert.pycolbert_scorescolbert_scores_pairwisecolbert_kd_scores
token 掩码sentence_transformers/multi_vector_encoder/modules/multi_vector_mask.pyMultiVectorMask.forwardresolve_with_tokenizer
token 压缩sentence_transformers/multi_vector_encoder/modules/token_pooling.pyBaseTokenPoolingHierarchicalTokenPooling