Aioli: A Unified Optimization Framework for Language Model Data Mixing

要約

言語モデルのパフォーマンスは、トレーニングに使用するデータ グループ (法律、コード、数学など) の最適な組み合わせを特定することに依存します。
これまでの研究では、トレーニング実行全体にわたる回帰モデルのフィッティングからトレーニング全体での比率の動的更新まで、混合比率を効率的に学習するためのさまざまな方法が提案されてきました。
驚くべきことに、グループごとの平均テストの複雑さの点で、単純な層化サンプリング ベースラインを一貫して上回る既存の方法はないことがわかりました。
このペーパーでは、既存の手法を標準の最適化フレームワークに統合することによって、この不一致の原因を研究します。
我々は、すべての手法が手法固有の混合法則、つまり損失が混合比の関数であるという仮定に従って、総損失を最小限に抑えるように割合を設定していることを示します。
混合則の既存のパラメータ化は、経験的に真の損失比率関係を表現できるが、手法自体は混合則パラメータを不正確に設定することが多く、その結果、パフォーマンスが低下し、一貫性がなくなることがわかりました。
最後に、フレームワークからの洞察を活用して、Aioli という名前の新しいオンライン手法を導き出します。これは、トレーニング全体を通じて混合則パラメータを直接推定し、それらを使用して比率を動的に調整します。
経験的に、Aioli は 6 つのデータセットのうち 6 つで層化サンプリングよりもテストの複雑さのポイントが平均 0.28 ポイント上回っていますが、既存の方法では常に層化サンプリングを上回ることができず、最大 6.9 ポイント劣っています。
さらに、計算上の制約により比率が短い実行で学習される実際の設定では、Aioli はトレーニングの実行全体にわたってこれらの比率を動的に調整することができ、既存の手法と比較してパフォーマンスを一貫して最大 12.01 テストパープレキシティ ポイント向上させることができます。

要約(オリジナル)

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 per group. In this paper, we study the cause of this inconsistency by unifying existing methods into a standard optimization framework. We show that all methods set proportions to minimize total loss, subject to a method-specific mixing law — an assumption on how loss is a function of mixture proportions. We find that existing parameterizations of mixing laws can express the true loss-proportion relationship empirically, but the methods themselves often set the mixing law parameters inaccurately, resulting in poor and inconsistent performance. Finally, we leverage the insights from our framework to derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Empirically, Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.28 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.01 test perplexity points.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google