Beyond Data Schedule: Fast LLM Pretraining on Data Mixtures by Loss-Reweighting
作者: Jiazheng Li, Huaqing Zhang, Jing Xu, Tianle Cai, Jingzhao Zhang
分类: cs.LG, cs.AI
发布日期: 2026-10-23
💡 一句话要点
提出LARE以解决大语言模型预训练中的知识遗忘问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 预训练 损失重加权 知识遗忘 领域适应 模型优化 机器学习 人工智能
📋 核心要点
- 现有的大语言模型预训练方法在提升能力的同时，容易导致早期知识的遗忘，影响模型的整体性能。
- 本文提出LARE方法，通过损失重加权自动调整各领域的权重，以优化模型在不同领域的学习效果。
- 实验结果表明，LARE在多个基准测试中显著提升了模型性能，尤其在GSM8K和MATH任务上分别提高了17.06和9.00个百分点。
📝 摘要（中文）
大型语言模型通常通过连续的预训练阶段获得专业能力，但这些提升可能会导致早期知识的遗忘。本文研究了通过调整各领域对学习的贡献强度来改善这种权衡的方法，同时保持数据暴露不变。在控制的知识-数学实验中，损失重加权产生的单一检查点超越了沿任一阶段轨迹获得的最佳知识和推理分数。这一发现揭示了数据混合比例无法捕获的独立控制。基于此，本文提出了LARE（固定暴露下的损失自适应重加权），该方法根据运行损失统计自动调整领域权重。通过对梯度噪声的理论分析，其归一化的逆平方根规则在不改变采样混合或额外运行的情况下平衡相对损失规模。该方法在知识、数学和代码等领域的五个评估基准上均有所提升，显示出调整领域损失规模可以在不增加领域暴露的情况下改善联合能力。
🔬 方法详解
问题定义：本文旨在解决大型语言模型在预训练过程中因领域权重不当导致的知识遗忘问题。现有方法往往无法有效平衡不同领域的学习贡献，导致模型在某些领域的性能下降。
核心思路：论文提出的LARE方法通过损失重加权，自动调整各领域的权重，旨在优化模型在多个领域的学习能力，而不增加数据暴露。
技术框架：LARE的整体架构包括数据混合的固定暴露、损失统计的实时监控和领域权重的动态调整。该方法通过分析梯度噪声，设计出一种归一化的逆平方根规则来平衡损失规模。
关键创新：LARE的核心创新在于其损失自适应重加权机制，能够在不改变数据混合的情况下，优化领域权重的分配。这一方法与传统的静态权重分配方法本质上不同，能够更灵活地应对不同领域的学习需求。
关键设计：LARE采用了实时损失统计来动态调整领域权重，使用归一化的逆平方根规则来平衡不同领域的损失规模，确保模型在各个领域的学习效果最大化。
📊 实验亮点
实验结果显示，LARE在1.5B模型上训练240B个token时，平均提升了7.84个百分点，特别是在GSM8K和MATH任务上分别提高了17.06和9.00个百分点，显著优于传统的直接混合方法。
🎯 应用场景
该研究的潜在应用场景包括教育、金融、医疗等领域，能够帮助开发更为智能的语言理解系统，提升模型在特定领域的表现。通过优化预训练过程，未来的模型可以更好地保留和利用多领域知识，增强其通用性和适应性。
📄 摘要（原文）
Large language models often acquire specialized capabilities through successive pretraining stages, but these gains can come at the cost of forgetting earlier knowledge. In this work, we investigate whether such trade-offs can be improved by changing how strongly each domain contributes to learning while keeping its data exposure fixed. We find that, in controlled knowledge--mathematics experiments, loss reweighting produces a single checkpoint that surpasses the best knowledge and reasoning scores attained along either stage trajectory. This reveals an independent control that data-mixture ratios alone do not capture. Building on this finding, we propose LARE (Loss-Adaptive Reweighting under fixed Exposure), which automatically adjusts domain weights from running loss statistics. Motivated by a theoretical analysis of gradient noise, its normalized inverse-square-root rule balances relative loss scales without changing the sampled mixture or requiring additional runs to search for weights. Across knowledge, mathematics, and code, \Method{} improves all five evaluated benchmarks for a 1.5B model trained on 240B tokens, averaging +7.84 pp over a matched direct mixture, including gains of +17.06 pp on GSM8K and +9.00 pp on MATH. These findings demonstrate that adjusting domain loss scales can improve joint capabilities without additional domain exposure.