CausalMix:将数据混合优化重构为因果推断问题
阅读原文· arxiv.orgCausalMix将大语言模型训练中的数据混合优化重构为因果推断问题,将数据池统计特征作为协变量、领域混合作为处理变量,在512次Qwen2.5-0.5B运行上拟合因果模型估计条件平均处理效应(CATE),外推出800K数据池的最优混合比例并用于训练7B模型。该框架还能泛化至Qwen3-4B-Base的长链式推理数据。通过因果建模隔离混杂偏差,CausalMix动态推断状态依赖的最优数据混合,在多个下游任务上优于RegMix等基线,并借助CATE解释器提供可视化分析。
In Large Language Model (LLM) training, data mixing plays a pivotal role in determining model performance. Recent methods optimize mixture weights via proxy models, but they rely on the assumption of static data distributions. As a result, when the underlying data pool shifts, these methods require costly retraining from scratch. This limitation restricts their ability to scale seamlessly from small settings to larger data pools and model sizes. In this paper, we propose CausalMix to address this limitation by casting data mixture optimization as a causal inference problem. We formulate the statistical features of the data pool as covariates and the domain mixture as the treatment. After fitting a causal model on 512 runs of Qwen2.5-0.5B to estimate the Conditional Average Treatment Effect (CATE), we extrapolate the optimal mixture for an 800K data pool and apply it to train a 7B model. Furthermore, we successfully generalize the framework to long chain-of-thought data on Qwen3-4B-Base. By leveraging causal modeling to isolate confounding biases, CausalMix dynamically infers state-dependent optimal data mixtures. Extensive experiments show that the mixture guided by CausalMix consistently improves performance across multiple downstream tasks, outperforming RegMix and other baselines. In addition, we use the CATE Interpreter to provide visual analysis of the learned mixing strategy. Overall, CausalMix offers a causal and interpretable framework for optimizing LLM data mixtures.