CausalMix: Data Mixture as Causal Inference for Language Model Training

Abstract: In 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.

• The paper introduces a causal inference framework, CausalMix, to optimize data mixtures for language model training using state-conditioned effect estimation.
• It employs double machine learning to orthogonalize treatment and outcome, isolating true marginal returns across high-dimensional covariates.
• Empirical results show CausalMix outperforms established baselines across domains, emphasizing its robustness and transferability for adaptive mixture policies.

CausalMix: A Causal Inference Framework for Data Mixture Optimization in LLM Training

Motivation and Problem Statement

Optimal domain mixture selection is critical for maximizing LLM performance during supervised fine-tuning (SFT). Traditional strategies rely on proxy model-based loss minimization or regression over mixtures, which generally assume a static data distribution and ignore state-dependent data effects. These approaches are suboptimal when datasets or data complexity shift, and they lack interpretability and transferability to new domains or model scales. The paper "CausalMix: Data Mixture as Causal Inference for LLM Training" (2607.01104) formalizes data mixture optimization as a state-conditioned causal inference problem, introducing a pipeline that isolates the true marginal effect of data mixture alterations across complex, high-dimensional covariate spaces.

CausalMix Pipeline: Causal Modeling of Data Mixtures

CausalMix reframes mixture selection as estimating the Conditional Average Treatment Effect (CATE) of domain mixture changes on downstream metrics, conditional on quantified data characteristics (covariates). The framework leverages historical proxy-model training runs, where each run is characterized by a triplet: pre-training data-state covariates $X_i$ (capturing quality, difficulty, entropy, etc.), the assigned domain mixture $T_i$, and post-training outcome $Y_i$ (typically aggregate downstream evaluation). The core task is to estimate, for any target covariate profile, the state-conditioned marginal return $\theta_0(x)$ for an infinitesimal log-change in domain mixture proportions.

Figure 1: Overview of the CausalMix pipeline. Historical proxy runs provide data-state covariates, mixture assignments, and downstream outcomes, which are used to estimate state-conditioned marginal data returns through orthogonal causal learning.

Identification relies on the assumption of ignorability by design, since all mixtures are set a priori and do not depend on downstream observations. Mixture proportions are log-transformed for compositional stability, and the estimand is operationalized via a partially linear approximation.

CausalMix employs the Double Machine Learning (DML) framework for orthogonalization: both treatments (mixtures) and outcomes are residualized with respect to the covariates via flexible regressors (LightGBM), and the state-conditioned marginal effect is estimated by regressing residual outcomes on residualized treatments. This orthogonalization is critical to separate the effect of mixture changes from inherent data pool quality, as confirmed by ablation.

Mixture Policy Extraction

After causal effect estimation, mixture selection proceeds via either (a) an analytic mapping that assigns positive marginal return domains mixture weights in proportion to their estimated effect (and zeros out non-beneficial domains), or (b) a search-based strategy that averages over top-performing candidate mixtures, smoothing inference noise.

The analytic policy is precisely characterized by the solution of a KKT-regularized optimization over the probability simplex, and maps to a rectified, normalized assignment of domain mixture weights (see supplemental proof).

Empirical Results and Robustness

CausalMix is evaluated on multi-domain SFT using standardized datasets (e.g., Tulu-3-SFT-Mixture, AM-Thinking-v1-Distilled), proxy-trained on Qwen2.5-0.5B, and then extrapolated to larger models (Qwen2.5-7B, Qwen3-4B). Covariates are drawn from the OpenDataArena-scored-data, explicitly including metrics for normalized cross-entropy loss, top-token entropy (HES), and stylistic quality. Evaluation spans knowledge recall, reasoning, math, coding, instruction following, and safety.

Key findings include:

• CausalMix outperforms RegMix, DMO, DoReMi, and other SFT-oriented baselines under all data budgets (100K–800K) and model scales in domain-aggregated downstream metrics.
• The search-based variant exhibits stronger robustness on out-of-domain, unseen sets due to variance smoothing.
• CausalMix demonstrates strong cross-dataset and cross-model transfer: when fit on proxy runs from one mixture pool and extrapolated to a distinct pool and model family, the state-aware mixture yields the highest averaged downstream scores.

Notably, ablation studies confirm two indispensable components: - Removing covariates (i.e., reverting to context-agnostic mixture regression) substantially degrades generalization; the mixture-over-state mapping is non-stationary. - Omitting DML orthogonalization (i.e., vanilla regression) reduces effect estimation fidelity, decreasing absolute performance.

Covariate and Model Selection

CausalMix's performance is contingent on both the expressiveness and parsimony of covariate selection. Empirical analysis via Spearman rank correlation reveals an optimal regime with three covariates (normalized loss, HES, writing style). Larger covariate sets decrease rank correlation, reflecting sample complexity and bias-variance tradeoff, and highlighting the curse of dimensionality in causal effect estimation.

Figure 2: Spearman rank correlation under different covariate combinations; performance deteriorates when including more than three covariates, underscoring the importance of judicious covariate selection.

Among causal estimators supporting continuous, multidimensional treatments, tree-based CausalForestDML with LightGBM first-stage models achieved superior R-loss metrics, outperforming linear and polynomial competitors.

Interpretability: CATE Model Analysis

CausalMix enables fine-grained interpretability of data mixture dynamics. Tree interpreter analysis of the learned causal model reveals:

• Instruction following data (IF) consistently yields strong, positive marginal returns across most covariate subspaces, underlining its centrality to downstream alignment.
• Knowledge data produces negative marginal effects in high-complexity (high normalized loss and HES) target domains, indicating adverse "skill conflict" with logical reasoning—a claim corroborated by independent analyses [wu2025knowledge, balappanawar2025if].
• Math, coding, and safety data are only beneficial when stylistic quality and entropy are moderate; otherwise, they induce noise in low-quality pools.

  Figure 3: Simplified visualization of the CATE model tree interpreter. The sign and magnitude of marginal returns depend strongly on data quality and complexity covariates.

Implications and Future Directions

Practically, CausalMix unlocks scalable, state-aware data mixture policies that (1) remain robust to large-scale data pool changes without requiring exhaustive retraining, (2) generalize across model and dataset shifts, and (3) provide interpretable, causal attribution for mixture policy choices. Theoretically, this causal formalism enables an actionable bridge from mixture effect identification to bounded policy optimization, grounded in statistical treatment effect literature.

Broader implications center on the emergence of second-order mixture effects—e.g., "skill conflict" and "synergy" regimes—which can inform future work on dynamic curriculum learning, skill interpolation, and the explicit management of trade-offs in LLM behavior. The framework also prompts future research on high-dimensional CATE estimation under sample constraints and causal discovery for dynamic or temporally evolving mixture policies.

Conclusion

CausalMix demonstrates that recasting LLM data mixture selection as a causal inference problem, with explicit state conditioning and orthogonalized effect estimation, yields robust and transferable mixture policies with significant gains over state-of-the-art baselines. The interpretability of the causal marginal returns exposes domain interactions not accessible via traditional regression or loss minimization. As LLM training moves toward more dynamic and adaptive paradigms, causal approaches such as CausalMix will be increasingly central to maximizing model capabilities under realistic computational, distributional, and generalization constraints.