CausalMix: Data Mixture as Causal Inference for Language Model Training

CausalMix frames language-model data-mixture optimization as a causal-inference problem. Statistical characteristics of the data pool are treated as covariates, while domain mixture is the treatment. The authors fit a causal model using 512 training runs of Qwen2.5-0.5B, estimate conditional average treatment effects, and extrapolate a mixture to an 800K-data pool and 7B-model training. They also extend the framework to long chain-of-thought data with Qwen3-4B-Base. According to the paper abstract, CausalMix consistently outperforms RegMix and other baselines across downstream tasks and provides visual interpretation through a CATE Interpreter.