用因果推断优化大模型训练数据混合,动态适配数据分布变化。
CausalMix: Data Mixture as Causal Inference for Language Model Training

- 将数据混合问题转为因果推断,以统计特征为协变量、领域混合为处理因素。
- 在512次实验基础上外推800K数据池最优混合方案,训练7B模型性能更优。
- 可解释性强,适用于动态数据分布和长链推理任务,适合大规模训练场景。
在大语言模型训练中,数据混合对模型性能至关重要。现有方法通过代理模型优化混合权重,但依赖静态数据分布假设,当数据池变化时需从头重新训练,难以扩展至更大规模数据和模型。本文提出CausalMix,将数据混合优化视为因果推断问题:将数据池的统计特征作为协变量,领域混合作为处理因素。在512次Qwen2.5-0.5B运行中拟合因果模型,估计条件平均处理效应(CATE),并外推至800K数据池,应用于7B模型训练。同时,该框架成功推广至Qwen3-4B-Base的长链思维数据。通过因果建模剥离混杂偏差,CausalMix能动态推断状态相关的最优数据混合。大量实验表明,基于CausalMix的混合策略在多个下游任务中持续提升性能,优于RegMix等基线。此外,使用CATE解释器提供混合策略的可视化分析。总体而言,CausalMix为优化大模型数据混合提供了因果且可解释的框架。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。