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)