Rescaling MLM-Head for Neural Sparse Retrieval

技术摘要：针对神经稀疏检索的 MLM-Head 重缩放技术

问题陈述
学习型稀疏检索（Learned Sparse Retrieval, LSR）模型，特别是 SPLADE，依赖于预训练编码器的掩码语言模型（MLM）头来生成稀疏词汇表示。尽管直觉上认为，将标准的 BERT 主干替换为更强大的现代编码器（如 ModernBERT、Ettin、RoBERTa）应该会提升检索性能，但作者观察到一个反直乎常理的失败现象：在标准的 SPLADE 训练方案下，具有较大 MLM-head L2 范数的模型往往会出现严重的性能退化或完全的训练崩溃。

论文指出其根本原因是规模失配（scale mismatch）。在 SPLADE 中，稀疏向量直接由 MLM-head 的输出（$H W_{head}^T$）构建，并通过未归一化的点积进行比较。当 MLM-head 矩阵 $W_{head}$ 的范数较大时，它会放大稀疏激活值和检索分数的幅度。这种放大扭曲了对比排序损失（contrastive ranking loss）的优化动力学，并与稀疏性正则化（FLOPS）产生冲突，从而导致训练不稳定且检索效果不佳。

方法论
为了在不改变模型架构或训练目标的情况下解决这种规模失配问题，作者提出了一种简单的、零成本的初始化时修正方法：MLM-Head 重缩放（MLM-Head Rescaling）。

在训练开始前，将 MLM-head 投影矩阵按常数因子 $k$ 进行重缩放：
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

该方法的主要特征包括：

• 零成本调整： 它不引入额外的参数，不需要改变架构，且在训练或推理过程中不会增加计算开销。
• 保留信息： 该操作保留了预训练词表投影的方向和词汇信息。
• 适用性： 对于具有权重共享（tied）MLM heads 的模型，该操作应用于共享的权重矩阵以维持权重绑定（weight tying）。

实验设置
作者使用固定的 SPLADE 方案，在 MS MARCO 三元组数据集上评估了七种编码器主干（BERT、DistilBERT、ALBERT、GTE-MLM、RoBERTa、ModernBERT 和 Ettin）。他们测试了各种重缩放因子（$k$），并从以下维度评估性能：

• 领域内（In-domain）： MS MARCO Dev 和 TREC 2019。
• 跨领域（Out-of-domain）： BEIR-13（平均 nDCG@10）。
• 稀疏性-效率权衡（Sparsity-Efficiency Trade-off）： 通过相对于激活 FLOPs 的 NanoBEIR 性能进行分析。

核心结果

1. 恢复大范数主干的性能： 标准初始化（$k=1$）导致具有大范数的现代编码器性能显著下降。例如，ModernBERT 的 BEIR-13 nDCG@10 降至 0.127（相比之下 BERT 为 0.370）。应用重缩放（$k=16$）后，ModernBERT 的性能恢复至 0.405，超越了 BERT 基线。同样，Ettin 从 0.221 提升至 0.391。
2. 非单调敏感性： 最优重缩放因子取决于主干模型的初始范数。ModernBERT 和 Ettin 需要较强的重缩放（$k=16$），而 RoBERTa 在 $k=2$ 时达到峰值。过度重缩放（例如对 RoBERta 使用 $k=16$）或对原本范数较小的模型（如 ALBERT）进行重缩放，都会导致性能崩溃。
3. 训练稳定性： 重缩放显著稳定了训练动力学。未经重缩放的 ModernBERT 和 Ettin 模型初始损失值极高，并收敛到了较差的局部最优解。重缩放后的变体则表现出更平滑的损失曲线和更低的最终损失。
4. 改进的权衡： 重缩放后的模型实现了更好的有效性-稀疏性权衡，运行在通常由成功的 BERT-SPLADE 模型占据的高有效性、低 FLOPs 区域。

意义与主张
本文认为，将预训练编码器适配到学习型稀疏检索中的瓶颈不仅在于编码器容量或预训练质量，更在于 MLM-head 的规模校准。

作者主张：

• MLM-head 的规模是 LSR 中一个关键且此前被忽视的因素。
• 更强大的编码器并不自动成为 SPLADE 中 BERT 的“即插即用”替代品；它们的投影规模必须针对稀疏检索的具体动力学进行校准。
• 简单的重缩放修正可以将不稳定训练的现代编码器转化为具有竞争力的稀疏检索器，通常能匹配甚至超越经典的 BERT-SPLADE 基线。

这项工作表明，未来将预训练编码器适配到 LSR 时，必须考虑词表投影层的几何结构和规模，而不应仅关注编码器的表示能力。