通过优化数据组织改进大语言模型预训练
训练大语言模型（LLM）时，通常会将文档拼接成“超级文档”，再分割为符合模型上下文长度的序列。这种做法虽然提高了训练效率，但常导致不必要的截断——单个文档被拆分到连续序列中。
在即将于国际机器学习会议（ICML 2024）发表的论文《减少截断提升语言建模效果》中，我们深入研究了这种常见的拼接-分块文档处理方法。研究发现，该方法严重损害模型理解上下文连贯性和事实一致性的能力，不仅影响下游任务性能，还增加了产生幻觉的风险。
为解决此问题，我们提出了最佳适配打包（best-fit packing）——一种创新的文档处理策略，通过优化文档组合来消除不必要的文本截断。在实验中，我们对比了使用最佳适配打包训练和常规方法训练的模型在六大下游任务上的表现：阅读理解、自然语言推理、上下文遵循、摘要、常识与闭卷问答以及程序合成。结果显示，最佳适配打包在22项子任务上均实现性能提升，最高达15%（程序合成）至17%（上下文遵循）。更重要的是，该方法还将闭域幻觉现象有效降低达58.3%。
截断的后果
研究分析了文档截断导致的三大问题：
- 未定义名称：在Python等编程语言中，截断可能分离变量定义与调用，引入语法错误并导致某些变量未定义，使模型学习到误导性模式。
- 无依据内容：截断损害数据完整性。例如，某个引用（“周一早上的地震”）与其前述内容分离，导致生成内容不忠实。
- 缺失知识：截断阻碍知识获取。例如，模型无法获知ICML会议的地点，因为会议名称和地点出现在不同训练序列中。
最佳适配打包方案
该方法通过优化文档到训练序列的分配来消除不必要截断，同时尽可能减少相对于拼接增加的序列数量。这实质上是经典的装箱问题（NP难问题）的变体，我们采用了一种在实践中表现良好的启发式算法——最佳适配递减（BFD）算法。
针对LLM预训练通常涉及数百万文档的特点，我们优化了BFD算法的时间复杂度，使其随数据规模线性扩展，确保适用于大规模预训练数据集。实际应用中，最佳适配打包生成的训练序列数量与传统方法相近，但显著减少了截导致的数据损失。
实现机制
最佳适配打包将每个训练序列视为容量等于LLM上下文长度的“箱子”。目标是将完整文档组合分配到各箱子中以最小化空间浪费。
具体步骤包括：
- 将大于上下文长度的文档分割成上下文长度的块加余数
- 按从大到小排序文档（和文档片段）
- 遍历排序列表，将每个文档分配到可用空间最接近文档大小的箱子中
为最大化效率，我们使用三种数据结构管理文档到箱子的分配：二叉树和两个表。这种设计基于两个观察：(1) 最大箱子大小为模型上下文长度，树不会太深；(2) 不需要区分具有相同剩余容量的箱子。
实验结果
我们在文本和代码上预训练了70亿和130亿参数的模型，上下文长度分别为2,000和8,000，比较了最佳适配打包与拼接方法的效果。在六大下游任务测试中，最佳适配打包在所有任务上平均表现更优，其中阅读理解（+4.7%）、自然语言推理（+9.3%）、上下文遵循（+16.8%）和程序合成（+15.0%）提升最为显著。
此外，最佳适配打包还有效防止了闭域幻觉，特别是在程序合成任务中，“未定义名称”错误减少达58.3%，表明模型对程序结构和逻辑的理解更加完整。
该方法还提升了模型遵循指令（如长度约束）的能力，并帮助模型获取因训练数据稀缺而对截断敏感的“长尾知识”。这一结果也为解释LLM难以学习长尾知识提供了可能的原因。
虽然实验主要关注LLM预训练，但最佳适配打包同样适用于微调阶段，其在微调中的益处将是未来研究的有趣课题。