CausalMix：將數據混合視為語言模型訓練中的因果推斷

摘要

在大语言模型（LLM）的训练中，数据混合对模型性能起着关键作用。近期的方法通过代理模型优化混合权重，但这类方法依赖于静态数据分布的假设。因此，当底层数据池发生变化时，这些方法需要从零开始进行昂贵的重新训练。这一局限性限制了它们从小规模设置向更大数据池和更大模型规模无缝扩展的能力。在本文中，我们提出CausalMix，通过将数据混合优化视为因果推断问题来解决这一局限性。我们将数据池的统计特征作为协变量，将领域混合作为处理变量。在基于512次Qwen2.5-0.5B模型的运行拟合因果模型以估计条件平均处理效应（CATE）后，我们外推得到一个800K数据池的最优混合方案，并将其应用于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.