arXiv:2606.14971cs.LGcs.AI2026-06

用梯度下降自动优化数据混合比例,提升训练效率

FastMix: Fast Data Mixture Optimization via Gradient Descent

论文配图:FastMix: Fast Data Mixture Optimization via Gradient Descent
图 1 · 摘自论文原文
  • 将数据混合优化转化为可微分的双层优化问题
  • 仅训练一个代理模型,搜索成本大幅降低
  • 适合需要高效数据组合的预训练与后训练场景

尽管大规模多样化数据集推动了大模型的进展,但确定预训练和后训练的最佳数据混合仍是一个开放难题。本文提出FASTMIX框架,通过仅训练一个代理模型,自动化发现最优数据混合。不同于依赖预设启发式或资源密集型模拟的方法,FASTMIX联合优化混合系数与模型参数,显著提升效率与可扩展性。其核心是将混合选择重构为双层优化问题,证明优化混合比例等价于在均匀采样下对各数据源分配损失权重。这一设计将混合系数嵌入可微迭代优化目标,实现混合与模型的梯度优化。FASTMIX采用近似迭代优化过程,交替执行:(i) 根据当前混合比例采样数据更新模型参数(内循环);(ii) 基于验证反馈更新混合比例(外循环)。在预训练与后训练任务中,FASTMIX均优于基线方法,同时大幅降低搜索成本。代码已开源。

原文摘要 · Abstract (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)

数据混合优化算法模型训练

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