---
title: "CausalMix: Data Mixture as Causal Inference for Language Model Training"
source: "arXiv AI/CL/LG（论文）"
source_tier: "t2"
url: "https://arxiv.org/abs/2607.01104v1"
permalink: "https://feed.astralor.com/items/itm_8b8206eadc49c0ff53479e22"
published: "2026-07-01T15:56:11.000Z"
score: 77
channel: "ai"
tags: ["CausalMix","Qwen2.5","Causal Inference","Large Language Models","Data Mixing Optimization","CATE"]
---

# CausalMix: Data Mixture as Causal Inference for Language Model Training

## AI commentary

> What excites me most about CausalMix is how it breaks the 'static assumption curse' in data mixture using causal inference. For teams working on large-scale pre-training or continual training, this paper directly addresses the pain point of having to re-optimize ratios whenever the data changes. While the validation on long chain-of-thought data is still shallow, the core idea is clean and impactful — a must-read for anyone interested in automatic data ratio optimization.

## At a glance

CausalMix proposes to formulate data mixture optimization for large language model training as a causal inference problem. By treating data pool statistics as covariates and domain mixture as treatment, it fits a causal model on 512 runs of Qwen2.5-0.5B to estimate Conditional Average Treatment Effect (CATE), then extrapolates the optimal mixture for an 800K data pool and successfully applies it to train a 7B model. This approach does not rely on static data distribution assumptions and dynamically infers state-dependent optimal mixtures, avoiding costly retraining from scratch. Extensive experiments show consistent improvements over RegMix and other baselines across multiple downstream tasks, along with a CATE-based visual analysis framework.

## Key points

- Formulates data mixture optimization as a causal inference problem using CATE estimation, breaking the static distribution assumption
- Fits a causal model on 512 small model runs and successfully extrapolates to larger data pools and 7B model training
- Dynamically infers state-dependent optimal mixtures, avoiding costly retraining and supporting extension from short text to long chain-of-thought data
- Provides a CATE interpreter for visual analysis, enhancing interpretability of the mixing strategy

## Full text

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.

---

Verify against the original before citing; curated by Astral Feed: https://feed.astralor.com/items/itm_8b8206eadc49c0ff53479e22