CausalMix: 数据混合作为语言模型训练的因果推断

摘要

在大语言模型（LLM）训练中，数据混合策略对模型性能起着决定性作用。现有方法通过代理模型优化混合权重，但这类方法假设数据分布保持静态。当底层数据池发生变动时，这些方法需要从头开始重新训练，成本高昂。这一局限性阻碍了其从小规模设置向更大规模数据池和模型尺寸的无缝扩展。本文提出CausalMix方法，通过将数据混合优化转化为因果推断问题来突破该限制。我们将数据池的统计特征建模为协变量，将领域混合视为处理变量。在基于512次Qwen2.5-0.5B模型运行拟合因果模型以估计条件平均处理效应（CATE）后，我们推断出适用于80万数据池的最优混合方案，并将其应用于7B模型的训练。此外，我们成功将该框架推广至Qwen3-4B-Base模型的长思维链数据。通过利用因果建模隔离混杂偏差，CausalMix能够动态推断依赖于状态的最优数据混合策略。大量实验表明，CausalMix指导的混合方案在多项下游任务中持续提升模型性能，优于RegMix及其他基线方法。同时，我们借助CATE解释器对所学混合策略进行可视化分析。总体而言，CausalMix为优化LLM数据混合提供了具有因果可解释性的框架。

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.