CausalMix: Data Mixture as Causal Inference for Language Model Training

Zinan Tang; Yukun Zhang; Shaomian Zheng; Zhuoshi Pan; Qizhi Pei; Dingnan Jin; Jun Zhou; Yujun Wang; Biqing Huang

Source-linked AI summary

CausalMix: Data Mixture as Causal Inference for Language Model Training

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

arXiv:2607.01104v1Jul 1, 2026cs.LGcs.AIcs.CL

TL;DR

Static mixture-optimization methods do not account for changing data-pool states, making optimal LLM training mixtures difficult to determine at scale. CausalMix estimates state-conditioned causal returns to infer mixtures that consistently outperform strong baselines across model scales and data budgets, including unseen LongCoT pools.

Problem
Existing data-mixing methods struggle to identify optimal LLM training mixtures because exhaustive search is expensive and static mappings ignore data-pool state.
Method
CausalMix treats data-pool statistics as covariates and domain mixtures as treatments to estimate state-conditioned causal marginal returns for mixture optimization.
Results
CausalMix consistently outperforms strong baselines across model scales and data budgets and transfers to unseen LongCoT data pools.
Takeaways & Limitations
CausalMix provides a transferable and interpretable framework for optimizing LLM data mixtures according to the state of the data pool.
Takeaways & Limitations
Using too many covariates degrades causal-estimation performance, partly because the 512-run historical meta-dataset is vulnerable to the curse of dimensionality.

Abstract

from arXiv · show

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.

1 Introduction

CausalMix reframes data-mixture optimization as causal marginal-return estimation conditioned on data-pool state, rather than black-box hyperparameter search. It uses causal modeling to isolate confounding effects, update mixture weights conservatively, and analyze state-dependent treatment effects.

Motivation: Training-data composition substantially affects LLM downstream performance, including during supervised fine-tuning across instruction-following, mathematical-reasoning, and coding domains.
Limitations of prior methods: Existing methods such as RegMix fit global mappings from mixture weights to validation loss, but loss-centric approaches often falter during supervised fine-tuning.RegMix treats historical training runs as independent samples and optimizes weights using a reference dataset.
CausalMix framework: CausalMix conditions mixture treatments on data-state covariates, including normalized loss, entropy, and writing style, to estimate localized marginal effects.The framework replaces a universal mapping from mixture proportions to absolute performance with a state-conditioned causal question.
Causal estimation: Double Machine Learning and causal forests orthogonalize treatments and outcomes against data state, isolating marginal returns from confounding caused by inherent data-pool quality.
Optimization: A trust-region-constrained policy update adjusts mixture weights conservatively after identifying the causal direction.
Interpretability and transferability: CATE analysis exposes skill conflicts between factual knowledge and complex logical reasoning and shows that data-quality thresholds affect the effectiveness of math and coding data.

2 Related works

Prior data-mixture optimization methods mainly target LLM pre-training using proxy models or loss modeling, while SFT mixture optimization remains relatively underexplored. Causal inference offers a framework for mitigating spurious correlations and distribution shifts through potential outcomes and causal graphical models.

Data mixture optimization: Existing offline data-mixture methods primarily optimize pre-training domain weights with proxy models or training-loss models, whereas SFT optimization remains underexplored.The cited methods include Xie et al. (2023b), Albalak et al. (2023), Liu et al. (2025), Fan et al. (2024), Ye et al. (2025), and Chen et al. (2025).
Causal inference in machine learning: Causal inference integrated with machine learning helps address spurious correlations and distribution shifts using potential outcomes and causal graphical models.This research line is grounded in the potential outcomes framework and causal graphical models.

3 Methodology

CausalMix formulates data-mixture optimization as state-conditioned causal marginal-return estimation, using historical proxy runs to relate data-state covariates, mixture treatments, and downstream outcomes. It estimates heterogeneous local mixture effects with orthogonal causal learning and converts target-state returns into feasible simplex mixtures.

Causal formulation: CausalMix models each prescribed mixture as a continuous treatment and uses a partially linear response to estimate local marginal returns within historical treatment support.The log-mixture representation accommodates compositional data on a probability simplex and diminishing marginal returns.
Causal formulation: Historical proxy runs provide data-state covariates, mixture assignments, and downstream outcomes for estimating state-conditioned marginal data returns.The data state is measured before training and evaluation, while the mixture is fixed before training and the outcome is downstream performance.
Causal formulation: The state-conditioned marginal return θ0(x) generalizes CATE: positive θ0,k(x) favors increasing domain k, whereas negative values indicate possible negative transfer.The covariate state can encode data-pool quality, difficulty, complexity, or stylistic statistics, but excludes post-training model information and downstream evaluation results.
Orthogonal causal learning: Double Machine Learning residualizes outcomes and treatments against covariates, while cross-fitting and an orthogonal loss estimate heterogeneous causal responses rather than absolute-score prediction.This isolates mixture effects from state-dependent baseline performance and uses residual treatment variation to explain residual outcome variation.
Mixture extraction: At the target state, CausalMix maps estimated log-mixture returns to a feasible simplex mixture either analytically or by averaging top-scoring raw candidate mixtures.Search-based extraction transforms each candidate into log-treatment space before causal scoring; averaging high-scoring candidates reduces inference noise and smooths the policy.

4 Experiments

Experiments evaluate CAUSALMIX across data scales, model sizes, and LongCoT data, using controlled comparisons with established mixture baselines. Results indicate strong downstream performance, transferability across datasets and models, and the importance of covariates and DML orthogonalization.

Experimental setup: The study samples 512 100K-instance sub-datasets from five Tulu 3 domains and treats their domain proportions as the treatment T.The domains are Coding, Instruction Following, Math Reasoning, Knowledge Recall, and Safety & NonCompliance.
Experimental setup: Qwen2.5-0.5B serves as the proxy model, with downstream evaluation grouped into Knowledge, Reasoning, Math, Coding, IF, and Safety capabilities.Training uses LlamaFactory, while evaluation follows the Tulu 3 protocol with OpenCompass.
Experimental setup: CAUSALMIX fits a causal model on 512 historical (X, T, Y) triplets using LightGBM and CausalForestDML within EconML.The experiments compare against Grid, RegMix, DoReMi, ODM, and DMO, re-implementing automated baselines for fair comparison.
Main results: CAUSALMIX achieves strong AvgDev performance and generalizes strongly on SUns, while CAUSALMIX-S outperforms CAUSALMIX-A on AvgUns.The paper attributes the latter improvement potentially to averaging the top-100 candidate mixtures, which may smooth idiosyncratic variance.
Generalization: The transferability experiment applies a causal predictor trained on Tulu 3 historical data to the entirely unseen AM-Thinking-v1-Distilled dataset and different model architectures.The outcome is defined as average downstream performance in coding and math, supporting extrapolation without costly proxy-model retraining and extending evaluation to LongCoT data.
Ablation studies: Ablations show performance degradation when DML orthogonalization is removed or covariates are discarded, highlighting both components’ necessity.The w/o X variant uses LightGBM to learn a global treatment-to-outcome mapping, ˆY = g(T).

5 Analysis

The analysis selects causal and first-stage models, identifies a three-covariate feature profile, and interprets how data domains affect downstream alignment across feature subspaces.

Causal estimator selection: CausalForestDML achieves the best performance among evaluated multidimensional continuous-treatment estimators, attributed to its non-parametric tree-based recursive partitioning.The architecture is described as better suited than linear causal models to complex interactions between multidimensional covariates and treatments.
First-stage predictor selection: Using LightGBM for both outcome and treatment first-stage models achieves the highest RScore and substantially outperforms standalone regressors and linear models.The first-stage predictors estimate conditional expectations for outcomes and treatments.
Covariate selection: The best covariate combination contains HES, Normalized_Loss, and Writing_Style, representing data Complexity, Difficulty, and Quality.HES captures reasoning complexity, Normalized_Loss reflects predictability and training utility, and Writing_Style evaluates clarity and coherence.
Covariate selection: Effective causal modeling requires diverse feature profiles, whereas adding too many covariates degrades performance because the limited historical meta-dataset increases estimator vulnerability.Covariate combinations are evaluated using Spearman rank correlation against ground-truth scores.
CATE interpretation: IF data yields stable positive downstream-alignment returns across feature subspaces, while Knowledge data has negative effects on difficult targets with high Normalized_Loss and high HES.The latter pattern corroborates skill conflicts between logical reasoning and factual knowledge injection.

6 Conclusion

CAUSALMIX reframes SFT data-mixture optimization as state-conditioned causal marginal-return estimation rather than static validation-loss minimization. It outperforms strong baselines across model scales and data budgets, transfers to unseen LongCoT pools, and yields interpretable insights into skill conflicts in LLM training.

Conclusion: CAUSALMIX shifts SFT data-mixture optimization from static validation-loss minimization to state-conditioned causal marginal-return estimation.The framework treats historical proxy training runs as causal treatments and combines orthogonalized estimation with a conservative trust…
Conclusion: CAUSALMIX consistently outperforms strong baselines across different model scales and data budgets while exhibiting strong transferability to unseen LongCoT data pools.
Conclusion: Its causal framework provides interpretable insights, including quantified skill conflicts between factual knowledge injection and complex logical reasoning.These insights provide a principled foundation for future research on understanding and optimizing LLM training dynamics.

A Experimental details · A.1 Datasets · A.2 Models

The experimental setup evaluates CausalMix on two supervised fine-tuning datasets, using OpenDataArena-derived covariates and Qwen models spanning proxy, scaled, and extension experiments. The datasets provide broad domain coverage, verified reasoning data, and fine-grained sample-level scores.

A Experimental details: CausalMix is evaluated on two supervised fine-tuning datasets.
A.1 Datasets: tulu-3-sft-mixture contains 939,344 samples spanning seven domains for training the Tulu 3 model series.
A.1 Datasets: AM-Thinking-v1-Distilled contains automatically verified responses from 1.89 million queries across diverse reasoning domains.
A.1 Datasets: OpenDataArena-scored-data-2603 supplies the data-state covariates used by CausalMix.
A.1 Datasets: OpenDataArena-scored-data-2603 comprises 63 instruction-following datasets with nearly 25 million samples and 30-dimensional scores.
A.2 Models: Qwen2.5-0.5B serves as the proxy model, while the learned mixture strategy is scaled to Qwen2.5-7B and extended to Qwen3-4B-Base.
A.2 Models: Qwen2.5 provides base and instruction-tuned models ranging from 0.5B to 72B parameters, with improvements in knowledge, coding, mathematical reasoning, instruction following, and long-form generation.
A.2 Models: Qwen3 is a newer Qwen generation offering dense and mixture-of-experts models with advances in reasoning, instruction following, agent capabilities, and multilingual support.

A.3 Benchmarks · A.4 Baselines · A.5 Computing costs

The evaluation spans knowledge, reasoning, math, code, instruction-following, toxicity, and truthfulness benchmarks, partitioned into development and unseen sets. CausalMix is compared with offline mixture-optimization baselines under matched proxy-training configurations and a quantified compute budget.

A.3 Benchmarks: The benchmark suite covers knowledge, reasoning, math, code, instruction-following, toxicity, and truthfulness, with development and unseen partitions.SDev includes MMLU, MMLU-Pro, BBH, DROP, GSM8K, MATH, HumanEval, MBPP, IFEval, and TruthfulQA; SUns includes GPQA, AGIEval, OlympiadBench, HumanEval+, IFBench, and ToxiGen.
A.3 Benchmarks: MMLU uses zero-shot chain-of-thought prompting and reports the macro average across all subjects as its final metric.The benchmark includes questions requiring factual recall, logical reasoning, and problem-solving.
A.3 Benchmarks: ToxiGen evaluates zero-shot unnormalized accuracy on 274k toxic and benign statements about 13 minority groups, while TruthfulQA uses zero-shot mc1 on 817 questions across 38 categories.TruthfulQA covers categories including health, law, finance, and politics.
A.4 Baselines: CausalMix is compared with RegMix, DoReMi, ODM, and DMO, which respectively use regression, Group DRO, multi-armed bandits, and validation-loss optimization for data mixing.RegMix trains many small models to predict unseen-mixture performance; DoReMi derives domain weights without downstream-task access; ODM optimizes proportions during training; DMO models effective data transfer using scaling laws.
A.5 Computing costs: 512 proxy models with 0.5B parameters are trained on 100K SFT examples each, averaging approximately 4096 tokens and totaling 5.53 × 10^20 estimated FLOPs.Because CausalMix is state-aware, it maintains strong OOD generalization without retraining the proxy models.
A.5 Computing costs: Baseline proxy training is matched to the same configuration for RegMix with 512 models, while DoReMi and ODM each use one model.DMO uses mixture proportions reported in the original paper because it is trained on the same data.

A.6 Hyperparameters

The experiments use fixed seeds and NVIDIA H800 GPUs, with training configurations following prior work and model-specific hyperparameters. Evaluation uses different maximum token limits depending on whether training includes LongCoT-style reasoning.

All experiments use random seed 42 and NVIDIA H800 GPUs.
Training follows DMO for Qwen2.5-0.5B and Qwen2.5-7B, and OpenDataArena for Qwen3-4B-Base.All training hyperparameters are listed in Table 5.
Evaluation sets max_tokens to 4096 for Qwen2.5-0.5B and Qwen2.5-7B, versus 32,768 for Qwen3-4B-Base.The difference is determined by whether the training data includes LongCoT-style reasoning; evaluation hyperparameters are listed in Table 6.
Qwen2.5-7B training uses learning_rate 5.0e-6, 3 epochs, 8 GPUs, and per-device batch size 16.It also uses gradient_accumulation_steps 2, cosine scheduling, warmup_ratio 0.1, cutoff_len 4096, DeepSpeed z2, FlashAttention fa2, Liger kernel, and bf16.
Qwen3-4B training uses learning_rate 5.0e-5, 3 epochs, 8 GPUs, and per-device batch size 2.It also uses gradient_accumulation_steps 2, cosine scheduling, warmup_ratio 0.1, cutoff_len 32768, DeepSpeed z2, FlashAttention fa2, Liger kernel, and bf16.

B Proof of the analytical mixture policy

The analytical CAUSALMIX-A policy solves a simplex-constrained optimization using KKT conditions. Its exact closed form assigns zero weight to non-positive causal returns and proportionally normalized weight to positive returns.

Optimization formulation: The policy finds the raw mixture T* that maximizes expected causal performance gain under the Level-Log formulation over the probability simplex.The constraints require Σ_k T_k = 1 and T_k ≥ 0 for every domain k.
KKT derivation: KKT conditions handle the non-negativity constraints by introducing a Lagrange multiplier for the equality constraint and nonnegative multipliers for domain inequalities.The derivation minimizes −J(T) and imposes stationarity, feasibility, and complementary slackness.
Sign-based allocation: Domains with non-positive estimated marginal returns receive zero optimal allocation.For negative returns, the logarithmic objective diverges positively as allocation approaches zero, forcing the optimum to the boundary; the result also includes zero returns.
Sign-based allocation: Domains with positive estimated marginal returns form the active set and receive allocations proportional to their estimated marginal returns.Positive-return domains must have positive allocation, and feasibility determines the normalization factor.
Closed-form policy: The global optimum is analytically equivalent to applying ReLU to causal marginal returns followed by L1 normalization.This closed-form policy is mathematically exact under the simplex constraint.