RegMix solved pretraining mixture optimization with 512 tiny proxy models plus LightGBM regression, but its assumption is static: fit once, get one globally optimal mixture. When the underlying data pool shifts (domains added or removed, distribution drift), the entire regression needs to be rerun. At the SFT stage this problem is worse, because SFT data pools and domain definitions change frequently.

CausalMix ( Data Mixture as Causal Inference for Language Model Training , 2026/07) recasts mixture optimization as a causal inference problem. The difference: RegMix learns a global mapping $T \to Y$ (mixture to performance), while CausalMix learns a conditional causal effect $\theta_0(X)$ (the marginal return of each domain under the current data state $X$). By orthogonalizing confounding effects of the data state via Double Machine Learning, the causal model does not need new proxy experiments when the data pool changes.

After fitting the causal model on 512 Qwen2.5-0.5B runs over tulu-3-sft-mixture, extrapolating to 800K samples for a 7B model, CausalMix-S reaches AvgDev 62.28 (RegMix 60.14, DMO 60.35) and CausalMix-A reaches AvgUns 49.09 (RegMix 48.12, DMO 48.98). On transfer to a completely different Qwen3-4B + AM-Thinking long-CoT dataset, CausalMix averages 66.66, versus RegMix 61.40 and DMO 63.47.

Background and motivation

RegMix’s regression target is validation loss on a reference dataset: given a mixture $T$, predict Pile-CC validation loss. This works for pretraining but hits two difficulties at the SFT stage:

1. SFT evaluates on downstream task accuracy rather than loss; the loss-to-downstream correspondence is less stable than in pretraining
2. SFT data pools change frequently (swap an instruction set, add a batch of code data), meaning RegMix-style regression needs a full re-run of 512 proxies for each new pool

CausalMix’s motivation: if we can learn “the causal marginal return of the mixture under the current data state,” then even when the pool changes (reflected as a change in data state $X$), the causal model only needs the new $X$ plugged in to produce a new optimal mixture, without rerunning proxy experiments.

Causal formulation

The paper models the problem as: given proxy experiments $(X_i, T_i, Y_i)$, where $X$ is the data state (specifically three metrics: HES for reasoning complexity, Normalized_Loss for data difficulty, Writing_Style for text quality), $T$ is the mixture vector, and $Y$ is the average downstream score on the Development set. The goal is to estimate conditional marginal returns:

$$\mu(x, Z) \approx g(x) + \theta_0(x)^\top Z$$

where $Z = \log(T + \epsilon)$ is the log-mixture representation, $g(x)$ is baseline performance, and $\theta_0(x) \in \mathbb{R}^K$ is the CATE (Conditional Average Treatment Effect): how performance changes when log-mixture $Z$ is perturbed under data state $x$.

If $\theta_{0,k}(x) > 0$, increasing domain $k$’s proportion improves performance under the current data state; if $< 0$, there is negative transfer.

DML orthogonalization

Directly regressing $Y = f(X, T)$ conflates the confounding effects of data state with the causal effect of the mixture. The paper uses Double Machine Learning (DML) for orthogonalization: first estimate $\hat{m}(X) = E[Y|X]$ and $\hat{e}(X) = E[Z|X]$ with two LightGBMs, then estimate the causal effect from residuals:

$$\tilde{Y} = Y - \hat{m}(X), \quad \tilde{Z} = Z - \hat{e}(X)$$

$$\hat{\theta} = \arg\min_\theta \sum_i |\tilde{Y}_i - \theta(X_i)^\top \tilde{Z}_i|^2$$

This R-loss objective does not predict absolute scores; it estimates “how residual variation in the mixture explains residual variation in performance.” The final CATE learner is CausalForestDML, a causal random forest.

Ablations confirm the necessity of both orthogonalization and data state $X$: removing orthogonalization (w/o Orth.) makes the model degenerate into naive supervised regression over concatenated $(X, T)$, dropping AvgDev to 59.65, below RegMix’s 60.14; removing $X$ (w/o X) gives AvgDev 61.30, also below full CausalMix-A’s 61.84.

From CATE to final mixture

Given $\hat{\theta}(X_{\text{tar}})$, the paper provides two methods to convert marginal returns into a mixture:

Analytical (CausalMix-A): directly normalize positive marginal returns, $T_k \propto [\hat{\theta}_k(X_{\text{tar}})]_+$

Search-based (CausalMix-S): sample 100,000 candidate mixtures from a Dirichlet distribution, score each with the causal model, average the top-100. This matches RegMix’s search strategy; the difference is that the scoring function switches from a regression model to a causal model.

In experiments CausalMix-S slightly outperforms CausalMix-A on the unseen set; the paper attributes this to the top-100 averaging smoothing single-point noise.

Experimental results

Data is tulu-3-sft-mixture, split into five domains: Coding, Instruction Following, Math Reasoning, Knowledge Recall, Safety. 512 subsets of 100K samples each, proxy model is Qwen2.5-0.5B. Downstream evaluation splits into Development set and Unseen set.

Main comparisons run at three data scales (100K/400K/800K) and two model scales (0.5B/7B). At 800K + 7B, CausalMix-S AvgDev is 62.28 (best), CausalMix-A AvgUns is 49.09 (best); RegMix is 60.14 / 48.12, DMO is 60.35 / 48.98. The two variants trade wins: S is better on the Development set, A is better on the Unseen set.

The transfer experiment is more telling: the causal model trained on tulu-3-sft-mixture is directly applied to AM-Thinking-v1 long-CoT data (a completely different data pool), training and evaluating on Qwen3-4B-Base. CausalMix averages 66.66 on Math+Code, versus RegMix 61.40, DoReMi 62.00, DMO 63.47. The improvement in the transfer setting (about 3 to 5 points) is larger than in the main experiment, confirming that the causal model is more robust to data pool changes than global regression.

CATE interpretability

The paper visualizes the CausalForestDML tree structure. Notable findings:

1. Instruction Following data has positive marginal returns across nearly all data states: a safe bet
2. Knowledge data has negative marginal returns under high-difficulty, high-complexity states, reflecting “skill conflicts” between factual knowledge and logical reasoning
3. Math, Coding, and Safety data have negative marginal returns on low-quality data (introducing distributional noise), but synergistic gains on moderate-quality data

These insights are invisible from RegMix’s LightGBM feature importances, because RegMix does not condition on data state.

Limitations

Proxy scale is still limited to 512 runs of 0.5B models. Three covariates work best; adding more degrades performance, which the paper attributes to the meta-dataset being too small (only 512 rows) for the causal estimator to handle higher dimensions without excessive variance.

The causal identification assumption requires that the mixture is determined before training and evaluation, ruling out online data mixing approaches.

The definition of $X$ currently relies on pre-computed scores from OpenDataArena; applying CausalMix to a dataset without these scores requires running the scoring pipeline first.

Takeaway

CausalMix upgrades RegMix from “global regression” to “conditional causal effect estimation”: adding data state as a covariate, orthogonalizing confounding via DML, and estimating CATE with a causal forest. The payoff is twofold: slightly better performance on the same data pool (1 to 2 points), and direct transfer to different data pools without rerunning proxy experiments (about 3 to 5 points in the transfer experiment). The latter matters more for practical engineering.

For readers already familiar with RegMix: RegMix fits a global regression line from $T$ to $Y$; CausalMix adds “what the data looks like” as a conditioning variable on top of that line, using causal orthogonalization to ensure adding $X$ does not introduce new bias.

Extended reading: RegMix paper summary