CausalMix: Data Mixture as Causal Inference for Language Model Training

2026-07-01

AI Summary

This work addresses the limitations of conventional data mixing methods, which rely on static distributional assumptions, struggle to adapt to dynamically evolving data pools, and incur prohibitive retraining costs that hinder scalability in large models. The paper introduces a novel approach by formulating data mixing optimization as a causal inference problem, treating statistical data characteristics as covariates and mixing proportions as intervention variables. By estimating the Conditional Average Treatment Effect (CATE), the method dynamically infers optimal mixing strategies that generalize across model scales and data distributions, offering both interpretability and strong empirical performance. Evaluated on a 800K-scale data pool, the framework successfully extrapolates optimal mixing ratios. Experiments with the Qwen model family demonstrate that a 7B model trained using this approach outperforms baselines such as RegMix across multiple downstream tasks, and its efficacy is further validated on long-chain-of-thought data from Qwen3-4B-Base.

Abstract

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.

Problem

• data mixture
• causal inference
• large language models
• distribution shift
• scalability

Innovation

• Causal Inference
• Data Mixture Optimization
• Conditional Average Treatment Effect
• Large Language Models
• Interpretable AI