arXiv:2411.05735cs.LGcs.AI2024-11ICLR被引 36

提出新方法Aioli,动态优化多类型数据混合比例,显著提升语言模型性能。

Aioli: A Unified Optimization Framework for Language Model Data Mixing

  • 统一现有方法为同一优化框架,揭示其依赖隐含混合规律
  • 新方法在线估计混合规律参数,6个数据集上平均降低0.27困惑度
  • 适用于计算资源受限场景,可大幅改进短周期训练后的长期性能

语言模型性能依赖于最优数据组混合比例(如法律、代码、数学)。此前方法通过回归建模或动态调整比例来学习混合策略,但令人惊讶的是,无一方法在平均测试困惑度上持续优于简单的分层采样基线。我们发现,现有方法本质均等价于求解一个共同优化问题:在特定混合规律约束下最小化平均损失。该框架表明,评估方法的混合规律拟合精度可预测其表现。实证显示,现有方法对混合规律参数设定不准,导致性能不稳定。基于此,我们提出新方法Aioli,通过训练中直接估计混合规律参数并动态调整比例。Aioli在6个数据集上全部超越分层采样,平均降低0.27困惑度;而其他方法无法一致超越,最差时高出6.9困惑度。在计算受限的实际场景中,即使仅用短周期训练学习初始比例,Aioli仍能全程动态调整,性能提升最高达12.012困惑度点。

原文摘要 · Abstract (English)

Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over training runs to dynamically updating proportions throughout training. Surprisingly, we find that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity. To understand this inconsistency, we unify existing methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law -- an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of a method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we observe. Using this insight, we derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.27 test perplexity points, whereas existing methods fail to consistently beat stratified sampling, doing up to 6.9 points worse. Moreover, in a practical setting where proportions are learned on shorter runs due to computational constraints, Aioli can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.012 test perplexity points.

语言模型数据混合优化框架动态调整

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。