FastMix: 勾配降下法による高速データ混合最適化

要旨

大規模で多様なデータセットが大規模モデルの最近の進展を牽引してきた一方で、事前学習および事後学習における最適なデータ混合比率を特定することは、依然として重要な未解決課題です。本研究では、この課題に対処するため、単一のプロキシモデルを学習するだけでデータ混合の発見を自動化する新たなフレームワークFASTMIXを提案します。FASTMIXは、事前に定義されたヒューリスティックスやリソース負荷の高いシミュレーションに依存する代わりに、混合係数とモデルパラメータを同時に最適化することで、従来手法と比較して効率性とスケーラビリティを大幅に向上させます。FASTMIXの中核は、混合選択を二段階最適化問題として再定式化することです。この再定式化の下で、混合比率の最適化は、一様なソースサンプリングにおけるソース別損失重みの割り当てと数学的に等価であることを示します。これにより、混合係数を微分可能な反復最適化目的関数に直接埋め込み、混合とモデルの両方に対して効率的な勾配ベースの最適化を可能にします。この最適化問題を解くために、FASTMIXは近似的な反復最適化手順を実装し、(i)現在の混合比率に従ってサンプリングされたデータを用いたモデルパラメータの更新（内部ループ）と、(ii)検証フィードバックに基づく混合比率の更新（外部ループ）を交互に行います。事前学習と事後学習の両方において、FASTMIXはベースラインを上回る性能を示すと同時に、探索コストを大幅に削減します。コードはhttps://github.com/hrtan/fastmixで公開されています。

English

While large and diverse datasets have driven recent advances in large models, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. We address this challenge with FASTMIX, a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource-intensive simulations, FASTMIX jointly optimizes mixture coefficients and model parameters, substantially improving efficiency and scalability over prior approaches. At the core of FASTMIX is a reformulation of mixture selection as a bilevel optimization problem. Under this reformulation, we show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling. This embeds the mixture coefficients directly into the differentiable iterative optimization objective, enabling efficient, gradient-based optimization of both mixture and model. To solve the optimization problem, FASTMIX implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop). Across pre- and post-training, FASTMIX outperforms baselines while drastically reducing search cost. Code (https://github.com/hrtan/fastmix)