稀疏检索中词汇表选择的关键作用与优化策略

1. 词汇表在稀疏检索中的核心作用

在构建稀疏检索系统（特别是SPLADE类模型）时，词汇表远不止是简单的单词列表。它实际上定义了模型的交互空间——词汇表的大小直接决定了稀疏表示的维度。这个发现源于我在训练韩语SPLADE模型时遇到的一系列问题，这些问题最终都追溯到词汇表的选择上。

通过与sentence-transformers维护者和OpenSearch社区的讨论（参见GitHub issue #3431），我们确认了一个关键结论：在基于学习的稀疏检索中，词汇表的选择对系统性能具有决定性影响。这不同于传统的信息检索系统，在那些系统中，词汇表主要影响索引构建效率，而在现代稀疏嵌入模型中，词汇表质量直接关系到表示能力。

关键发现：当使用jhu-clsp/mmBERT-base这类多语言模型时，如果目标语言（如韩语）在词汇表中覆盖不足，会导致表示空间完全崩溃——所有激活值归零，使得检索系统彻底失效。

2. 实验设计与模型选型

2.1 骨干模型对比

我选择了四种具有不同词汇表特性的预训练模型进行对比实验，所有模型都在约90万条本地数据集上训练：

选择这组模型的目的是验证一个假设：词汇表与目标语言的匹配度比模型规模更重要。前两个模型是韩语优化的，后两个则是通用的多语言模型，其中mmBERT-base覆盖语言最多。

2.2 训练配置细节

所有模型都采用对比学习框架训练，这是当前训练稀疏检索器最高效的方法之一。具体配置如下：

• 数据集结构：使用三元组格式<查询, 正例, 难负例>，通过BM25采样构建
• 批大小：8（受限GPU显存）
• 最大长度：512 tokens
• 正则化权重：
  • 查询侧：5e-5
  • 文档侧：3e-5
• 精度：bfloat16混合精度训练
• 硬件：klue/roberta-base使用4块GPU，其他模型使用8块GPU

这里特别说明正则化权重的选择：文档侧权重略低于查询侧，这是因为文档通常比查询长，需要更温和的稀疏约束以避免过度惩罚。

3. 实验结果与分析

3.1 训练动态观察

所有模型的训练损失都呈现良好收敛趋势，但出现了几个关键现象：

1. gte-multilingual-base的初始震荡：

   • 这个模型在训练初期出现明显的正则化损失峰值
   • 原因分析：其25万的大词汇量导致初始激活过于密集，模型需要时间学习如何抑制不相关语言的token

2. mmBERT-base的灾难性崩溃：

   • 验证集Recall@10先上升后骤降
   • 最终查询和文档的激活维度都归零
   • 对应GitHub issue中描述的"表示崩溃"现象

实战技巧：当发现验证指标突然下跌时，应立即检查激活维度统计。如果同时出现查询和文档激活归零，很可能是词汇表不匹配导致的崩溃前兆。

3.2 最终评估表现

在MTEB-ko-retrieval基准上的测试结果如下：

几个关键发现：

1. 韩语优化模型的优势：

   • A.X-Encoder-base和roberta-base表现最佳
   • 它们的激活维度适中（几百到几十万），说明学习到了有效的稀疏模式

2. 多语言模型的挑战：

   • gte-multilingual-base虽然能训练，但激活维度过高（数百万）
   • 表明模型难以确定哪些token对韩语检索真正重要

3. 完全失败的案例：

   • mmBERT-base所有指标接近零
   • 激活维度归零证实了表示空间的完全崩溃

3.3 mmBERT崩溃的深度分析

通过检查训练日志，我们发现mmBERT-base的崩溃过程分为两个阶段：

阶段一（前2844步）：

• 查询激活：~200维度
• 文档激活：~5000维度
• 验证指标正常上升

阶段二（4266步后）：

• 所有激活突然归零
• 验证指标同步暴跌

根本原因有三方面：

1. 分词器-语言不匹配：

   • gemma2分词器将韩语过度分割为罕见子词
   • 导致文本无法有意义地映射到词汇空间

2. 稀疏正则化的副作用：

   • 当表示本身已经很弱时，L1/FLOPS惩罚会加速归零
   • "全零"成为损失函数的最优点

3. 归一化问题：

   • SPLADE使用log(1 + ReLU(Wx))激活
   • 不合适的LayerNorm导致大多数值被ReLU置零

4. 实践建议与经验总结

4.1 词汇表选择指南

基于这些实验，我总结出以下词汇表选择原则：

1. 语言覆盖优先于模型规模：

   • 首先确保词汇表对目标语言有良好覆盖
   • 不要盲目选择参数量更大的模型

2. 警惕过度多语言模型：

   • 覆盖100+语言的模型可能不适合特定语种任务
   • 检查分词器在目标语言上的表现

3. 评估词汇表质量的方法：

   • 统计目标语料中的OOV(未登录词)率
   • 检查常见词是否被合理分割

4.2 训练过程中的监控

为了避免出现不可逆的表示崩溃，建议：

1. 实时跟踪这些指标：

   • 查询/文档的平均激活维度
   • 非零激活的比例变化
   • 正则化损失的相对大小

2. 设置早期停止条件：

   • 当激活维度连续多个checkpoint下降时暂停训练
   • 检查分词质量并考虑调整正则化强度

4.3 模型适配技巧

对于必须使用多语言模型的情况：

1. 词汇表过滤：

   • 根据目标语言筛选出相关token子集
   • 可以显著降低噪声和计算量

2. 渐进式正则化：

   • 初始阶段使用较弱正则化
   • 随着训练逐步增强稀疏约束

3. 温度调节：

   • 在softmax激活前加入温度参数
   • 帮助模型更平缓地学习稀疏模式

5. 理论启示与未来方向

这些实验揭示了一个更深层的理论观点：在稀疏检索中，词汇表实际上定义了查询和文档的交互空间。与传统稠密检索不同，这里的表示空间是离散且由词汇表显式定义的。

这一发现对模型设计有几个重要启示：

1. 词汇表即归纳偏置：

   • 它决定了模型可以学习哪些交互模式
   • 应该被视为重要的超参数而非实现细节

2. 跨语言迁移的挑战：

   • 直接迁移多语言模型可能适得其反
   • 需要考虑词汇表在不同语言间的平衡

3. 动态词汇的可能性：

   • 未来可能探索可学习的词汇表机制
   • 例如在训练过程中逐步调整token重要性

在实际项目中，我现在会首先花时间评估候选模型的词汇表质量，而不是比较它们的架构或参数量。一个简单但有效的做法是：用目标语料中的典型查询和文档样例，检查模型的分词结果是否保持语义完整性。如果发现大量不合理的子词分割，这个模型很可能不适合作为稀疏检索的骨干网络。