CausalMix: Data Mixture as Causal Inference for Language Model Training

📄 arXiv: 2607.01104v1 📥 PDF

作者: Zinan Tang, Yukun Zhang, Shaomian Zheng, Zhuoshi Pan, Qizhi Pei, Dingnan Jin, Jun Zhou, Yujun Wang, Biqing Huang

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-07-01

备注: 22 pages, 3 figures

💡 一句话要点

提出CausalMix以解决LLM训练中的数据混合优化问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 因果推断 数据混合 大型语言模型 条件平均处理效应 模型训练 性能优化 可解释性

📋 核心要点

  1. 现有数据混合优化方法依赖静态数据分布，无法适应数据池的变化，导致需要重新训练。
  2. CausalMix将数据混合优化视为因果推断问题，通过因果模型动态推断最佳数据混合。
  3. 实验结果显示，CausalMix在多个下游任务中表现优异，显著超越了传统方法的性能。

📝 摘要（中文）

在大型语言模型（LLM）训练中，数据混合对模型性能至关重要。现有方法通过代理模型优化混合权重，但依赖于静态数据分布的假设，导致在数据池变化时需要昂贵的重新训练。为了解决这一限制，本文提出CausalMix，将数据混合优化视为因果推断问题，通过对512次Qwen2.5-0.5B的运行进行因果模型拟合，估计条件平均处理效应（CATE），并为800K数据池推断最佳混合。实验表明，CausalMix在多个下游任务中持续提升性能，超越了RegMix等基线方法。

🔬 方法详解

问题定义：本文旨在解决大型语言模型训练中数据混合优化的不足，现有方法在数据分布变化时需要重新训练，效率低下。

核心思路：CausalMix通过将数据混合优化视为因果推断问题，利用因果模型动态推断最佳数据混合，从而克服静态假设的限制。

技术框架：整体流程包括数据池特征的统计建模、因果模型的拟合以及最佳混合的推断。具体步骤为：首先对数据池进行特征提取，然后构建因果模型，最后根据模型输出推断最佳数据混合。

关键创新：CausalMix的核心创新在于将数据混合优化转化为因果推断问题，能够动态适应数据池的变化，与传统方法相比具有更高的灵活性和效率。

关键设计：在模型设计中，使用条件平均处理效应（CATE）来评估不同数据混合的效果，并通过CATE解释器提供可视化分析，帮助理解混合策略的学习过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CausalMix在多个下游任务中均显著提升了模型性能，相较于RegMix等基线方法，性能提升幅度达到X%（具体数据未知），展示了其在数据混合优化中的有效性和优势。

🎯 应用场景

CausalMix的研究成果在大型语言模型的训练中具有广泛的应用潜力，能够有效提升模型在多种下游任务中的表现。其因果推断框架为数据混合优化提供了新的思路，未来可扩展至其他机器学习领域，推动模型训练的效率和效果提升。

📄 摘要（原文）

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.