CAUSALMIX: Data Mixture as Causal Inference for Language Model Training

Summary (Overview)

• Core Contribution: CAUSALMIX reformulates LLM data mixture optimization as a causal inference problem, estimating state-conditioned marginal returns of domain proportions rather than fitting static global mappings from mixtures to performance.
• Key Innovation: The framework uses Double Machine Learning (DML) and causal forests to orthogonalize treatment effects from confounding data-state covariates, enabling extrapolation to unseen data pools and larger models without proxy retraining.
• Empirical Results: CAUSALMIX consistently outperforms baselines (RegMix, DoReMi, ODM, DMO, Grid) across data scales (100K–800K) and model sizes (0.5B–7B), achieving the highest Avg_Dev scores in all settings.
• Transferability: The framework successfully generalizes to entirely unseen LongCoT data (AM-Thinking-v1-Distilled-math&code) with a different model family (Qwen3-4B), demonstrating robust cross-dataset and cross-architecture transfer.
• Interpretability: CATE Interpreter analysis reveals "skill conflicts" between factual knowledge and logical reasoning, as well as how data quality thresholds modulate the effectiveness of math/coding data.

Introduction and Theoretical Foundation

Background and Motivation

Data mixing—the relative proportion of different domains (instruction following, math reasoning, coding, etc.) in SFT training data—substantially impacts downstream LLM performance. However, determining the optimal mixture is challenging because:

1. Cost: Training LLMs is expensive, making exhaustive grid search over the continuous simplex of mixture weights intractable.
2. Static assumptions: Existing methods like RegMix fit global mappings from mixture weights to loss, assuming static data distributions. These fail when data pools shift and require costly retraining.
3. State-dependence: A single static optimal mixture does not exist—the optimal mixture depends on the inherent complexity, quality, and difficulty of the specific data pool.

Theoretical Foundation

The paper builds on several key theoretical frameworks:

• Potential Outcomes Framework (Rubin, 2005; Imbens & Rubin, 2015): Each training run with a prescribed mixture is viewed as a treatment, and downstream performance as the outcome.
• Double Machine Learning (DML) (Chernozhukov et al., 2018): Used to orthogonalize treatment and outcome variables with respect to data-state covariates.
• Causal Forests (Wager & Athey, 2018; Oprescu et al., 2019): Non-parametric estimation of heterogeneous treatment effects.
• Robinson's Transformation (Robinson, 1988): Partially linear approximation of the response surface.

The key insight is to ask a localized causal question: How does a relative change in domain proportions causally affect downstream performance under the current data state? This contrasts with the traditional approach of seeking a universal mapping from mixture proportions to absolute performance.

Methodology

Problem Formulation

Given $K$ data domains and a fixed training budget, a mixture is represented as:

$$T = (T_1, \dots, T_K), \qquad T_k \geq 0, \qquad \sum_{k=1}^{K} T_k = 1.$$

The continuous treatment uses a log-mixture representation:

$$Z = \log(T + \varepsilon),$$

where the logarithm is applied element-wise and $\varepsilon > 0$ is a small smoothing constant.

For each historical proxy run $i$, the observed triplet is $(X_i, T_i, Y_i)$ where:

• $X_i$: covariates (data state available before training)
• $T_i$: treatment (mixture fixed before training)
• $Y_i$: outcome (downstream performance after training)

Identification

The conditional response function is defined as:

$$\mu(x, z) = \mathbb{E}[Y(t) \mid X = x].$$

Using a partially linear approximation:

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

where $g(x)$ captures state-dependent baseline performance and $\theta_0(x) \in \mathbb{R}^K$ is the state-conditioned marginal data return (generalized CATE for multidimensional continuous treatments).

Identification assumptions:

• Consistency: $Y_i = Y_i(T_i)$
• Ignorability: $Y(t) \perp T \mid X$
• Local overlap and smoothness within historical treatment support

Under these assumptions:

$$\mathbb{E}[Y \mid X = x, Z = z] = \mathbb{E}[Y(z) \mid X = x] = \mu(x, z).$$

Orthogonal Estimation via DML

Nuisance functions are defined as:

$$m_0(X) = \mathbb{E}[Y \mid X], \qquad e_0(X) = \mathbb{E}[Z \mid X],$$

with residuals:

$$\widetilde{Y} = Y - m_0(X), \qquad \widetilde{Z} = Z - e_0(X).$$

The marginal return is estimated from:

$$\widetilde{Y} \approx \theta_0(X)^{\top} \widetilde{Z}.$$

Cross-fitting is used to avoid overfitting, and the heterogeneous effect model is learned by minimizing the orthogonal (R-loss) objective:

$$\hat{\theta} = \arg \min_{\theta} \sum_{i} \left(\widetilde{Y}_i - \theta(X_i)^{\top} \widetilde{Z}_i\right)^2.$$

Mixture Policy Extraction

Analytical extraction (CAUSALMIX-A):

$$T_k^{\mathrm{A}} = \frac{[\hat{\theta}_k(X_{\mathrm{tar}})]_+}{\sum_{j=1}^{K} [\hat{\theta}_j(X_{\mathrm{tar}})]_+}, \qquad [a]_+ = \max(a, 0).$$

Search-based extraction (CAUSALMIX-S): Draws 100,000 candidate mixtures from a Dirichlet distribution, evaluates them via the fitted causal model, and averages the top-100 performers:

$$T^{\mathrm{S}} = \frac{1}{K_{\mathrm{top}}} \sum_{m \in \mathrm{Top}} T^{(m)}.$$

This serves as local bagging to reduce inference noise and enhance generalization.

Experimental Setup

• Data: tulu-3-sft-mixture with 5 domains (Coding, IF, Math Reasoning, Knowledge Recall, Safety)
• Proxy model: Qwen2.5-0.5B, 512 sub-datasets of 100K instances each
• Covariates: Normalized_Loss, Writing_Style, HES (from OpenDataArena scores)
• Causal estimator: CausalForestDML with LightGBM first-stage predictors
• Evaluation: OpenCompass across 6 capabilities (Knowledge, Reasoning, Math, Coding, IF, Safety)

Empirical Validation / Results

Main Results (Table 1)

Key findings across all settings:

CAUSALMIX consistently outperforms all baselines on Avg_Dev. Notably, at 7B scale, CAUSALMIX-S achieves the best Avg_Dev (62.28), surpassing DMO (60.35) and RegMix (60.14).

LongCoT Transfer (Table 2)

On Qwen3-4B with AM-Thinking-v1-Distilled-Code&Math (20K data):

CAUSALMIX achieves the best overall performance (66.66), demonstrating robust transfer to unseen data pools and model architectures.

Ablation Study (Table 3)

• w/o X (removing covariates): Performance drops to 33.29 (0.5B) and 61.30 (7B), confirming the importance of state conditioning
• w/o Orth. (removing DML orthogonalization): Performance drops to 32.66 (0.5B) and 59.65 (7B), validating the necessity of orthogonalization
• Full CAUSALMIX-A achieves 33.94 (0.5B) and CAUSALMIX-S achieves 62.28 (7B)

Model Selection

Causal estimator (Table 4a): CausalForestDML achieves the best RScore (+0.1683), outperforming LinearDML (+0.1445) and other alternatives. The non-parametric tree-based architecture is better suited for capturing complex interactions in data mixing dynamics.

First-stage predictors (Table 4b): LightGBM for both outcome and treatment models achieves the highest RScore (0.1683), substantially outperforming RandomForest, GradientBoosting, and linear models.

Covariate selection (Figure 2): The optimal combination is three covariates: HES (complexity), Normalized_Loss (difficulty), and Writing_Style (quality). Too few covariates fail to control confounding; too many suffer from the curse of dimensionality given the limited meta-dataset size.

Theoretical and Practical Implications

Interpretable Insights (CATE Interpreter)

The tree-based analysis reveals several important findings:

1. Skill conflicts: Knowledge data has negative effects on difficult target data (high Normalized_Loss and high HES), corroborating the existence of "skill conflicts" between logical reasoning and factual knowledge injection.

2. Data quality thresholds: In low-quality regions (low Writing_Style, low HES), complex domains (Math, Coding, Safety) introduce distributional noise and degrade performance. However, at moderate Writing_Style and HES levels, these domains produce strong synergistic gains.

3. IF data as primary driver: Instruction Following yields stable positive returns across feature subspaces, serving as the backbone of downstream alignment.

Practical Implications

• Scalability: CAUSALMIX learns underlying causal dynamics rather than memorizing specific datasets, enabling extrapolation to larger data pools (800K) and model sizes (7B) without new proxy experiments.
• Transferability: Successfully generalizes from Qwen2.5-0.5B proxy to Qwen3-4B on entirely different data (LongCoT), avoiding costly proxy-model retraining.
• Interpretability: The causal framework provides principled explanations for why certain mixtures work, unlike black-box optimization methods.

Conclusion

CAUSALMIX shifts SFT data mixture optimization from static validation-loss minimization to state-conditioned causal marginal return estimation. By treating historical proxy runs as causal treatments and combining orthogonalized estimation (DML) with a conservative trust-region policy, the framework isolates the marginal utility of domain proportions from confounding data-state effects.

Key takeaways:

1. Causal inference provides a principled and interpretable framework for data mixture optimization
2. State-conditioned estimation enables transferability across data pools and model scales
3. The framework reveals quantifiable skill conflicts and quality thresholds in data mixing dynamics

Future directions:

• Scaling up the number of proxy models to support more covariates
• Extending the framework to pre-training data mixture optimization
• Further exploration of the causal mechanisms underlying data mixing dynamics in LLM training