Aioli: A Unified Optimization Framework for Language Model Data Mixing

语言模型的性能取决于识别出用于训练的最佳数据组混合比例（例如法律、代码、数学）。先前的研究提出了一系列高效学习混合比例的方法，从在训练过程中拟合回归模型到在整个训练过程中动态更新比例。令人惊讶的是，我们发现没有任何现有方法在每个组的平均测试困惑度方面始终优于简单的分层抽样基线。在这篇论文中，我们通过将现有方法统一到一个标准优化框架中来研究这种不一致的原因。我们表明，所有方法都设置比例以最小化总损失，但受制于特定方法的混合定律——即损失与混合比例之间的函数关系的假设。我们发现，现有的混合定律参数化可以在经验上表达真实的损失-比例关系，但这些方法本身往往无法准确设置混合定律参数，导致性能不佳且不一致。最后，我们利用我们框架中的见解，推导出一种新的在线方法Aioli，该方法在整个训练过程中直接估计混合定律参数，并使用这些参数动态调整比例。实证结果表明，Aioli在6个数据集中有6个优于分层抽样，平均提高了0.28个测试困惑度点，而现有方法则无法始终胜过分层抽样，最差时相差6.9个点。此外，在实际情况下，由于计算限制，比例是在较短的运行中学习的，Aioli可以在此类全程训练中动态调整这些比例，相比现有方法最多可提高12.01个测试困惑度点。

该论文探讨了如何通过优化不同数据组（如法律、代码、数学等）的混合比例来提高语言模型的性能。尽管已有多种方法试图高效地学习这些比例，但没有一种方法能一致地优于简单的分层抽样基线。

论文的关键思路是将现有的各种方法统一到一个标准的优化框架中，分析每种方法如何设置混合比例以最小化总损失，并基于特定的混合定律假设。研究发现，现有方法虽然可以表达真实的损失-比例关系，但在设置混合定律参数时不够准确，导致性能不佳和不一致。基于这一洞察，论文提出了一种新的在线方法——Aioli，该方法在训练过程中直接估计混合定律参数并动态调整比例。

论文通过实验证明，Aioli在6个数据集上均优于分层抽样，平均提高了0.28个测试困惑度点。此外，在实际应用中，当由于计算限制而需要在较短的运行中学习比例时，Aioli能够动态调整这些比例，进一步提升性能，最高可提高12.01个测试困惑度点。论文还提供了详细的实验设计和数据集信息，但未提及是否有开源代码。

近期相关研究包括：1) 通过回归模型拟合训练运行来学习混合比例的方法；2) 在训练过程中动态更新比例的方法。例如，《Optimizing Mixture Proportions for Language Models via Regression》和《Dynamic Mixture Proportions for Efficient Training of Language Models》等论文提出了类似的方法，但都没有达到一致的性能提升。