CausalMix: Data Mixture as Causal Inference for Language Model Training

• URL: http://arxiv.org/abs/2607.01104v1
• Date: Wed, 01 Jul 2026 15:56:11 GMT
• Title: CausalMix: Data Mixture as Causal Inference for Language Model Training
• Abstract summary: In Large Language Model (LLM) training, data mixing plays a pivotal role in determining model performance.<n>Recent methods optimize mixture weights via proxy models, but they rely on the assumption of static data distributions.<n>We propose CausalMix to address this limitation by casting data mixture optimization as a causal inference problem.
• Score: 25.091228095854415
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Data Mixing as Mixture Experiment: Response Surface Methodology and Optimal Design for Large Language Model Pretraining [1.0312968200748118]
  We show the structure of a classical mixture experiment using proxy-based methods.<n>Data domains are mixture components, token shares are component proportions, proxy-training runs are experimental design points, and validation loss defines a response surface over the probability simplex.<n>Using RegMix as an empirical case study, we demonstrate how the framework can both interpret observed mixture responses and design more efficient proxy experiments.
  arXiv  Detail & Related papers  (2026-08-24T23:56:38Z)
• FastMix: Fast Data Mixture Optimization via Gradient Descent [55.109758242985556]
  Fastmix is a novel framework that automates data mixture discovery while training only a single proxy model.<n>We show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling.<n>FASTmix outperforms baselines while drastically reducing search cost across pre-intensive and post-training.
  arXiv  Detail & Related papers  (2026-06-12T21:46:57Z)
• Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization [20.220685778194156]
  We introduce a compute-efficient pipeline for data mixture scaling.<n>First, we propose CAMEL, a capacity-aware mixture law that models validation loss.<n>We also introduce a loss-to-benchmark prediction law that estimates benchmark accuracy from validation loss.
  arXiv  Detail & Related papers  (2026-03-09T06:58:00Z)
• MergeMix: Optimizing Mid-Training Data Mixtures via Learnable Model Merging [72.00014675808228]
  textbfMix determines optimal data mixing ratios by repurposing model merging weights as a high-fidelity, low-cost performance proxy.<n>Experiments on models with 8B and 16B parameters validate that MergeMix achieves performance comparable to or surpassing exhaustive manual tuning.
  arXiv  Detail & Related papers  (2026-01-25T14:31:57Z)
• CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training [63.07024608399447]
  We propose an automated framework that discovers, evaluates, and refines data mixtures in a pre-training setting.<n>We introduce ClimbLab, a filtered 1.2-trillion-token corpus with 20 clusters as a research playground, and ClimbMix, a compact yet powerful 400-billion-token dataset.
  arXiv  Detail & Related papers  (2025-04-17T17:58:13Z)
• Optimizing Pre-Training Data Mixtures with Mixtures of Data Expert Models [24.396525123797073]
  We propose a method to optimize language model pre-training data mixtures through efficient approximation of the cross-entropy loss corresponding to each candidate mixture.<n>We use this approximation as a source of additional features in a regression model, trained from observations of model loss for a small number of mixtures.
  arXiv  Detail & Related papers  (2025-02-21T21:27:48Z)
• RegMix: Data Mixture as Regression for Language Model Pre-training [40.45464495981735]
  We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task.<n>RegMix trains many small models on diverse data mixtures, uses regression to predict performance of unseen mixtures, and applies the best predicted mixture to train a large-scale model with orders of magnitude more compute.
  arXiv  Detail & Related papers  (2024-07-01T17:31:03Z)
• BiMix: A Bivariate Data Mixing Law for Language Model Pretraining [47.77701041534746]
  The impact of pretraining data composition on model performance remains poorly understood.<n>$textbfBiMix$ provides a systematic framework for understanding and optimizing data mixtures.<n>Our work contributes both theoretical insights into data mixing dynamics and practical tools for enhancing LLM training efficiency.
  arXiv  Detail & Related papers  (2024-05-23T09:44:02Z)
• Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance [55.872926690722714]
  We study the predictability of model performance regarding the mixture proportions in function forms.<n>We propose nested use of the scaling laws of training steps, model sizes, and our data mixing law.<n>Our method effectively optimize the training mixture of a 1B model trained for 100B tokens in RedPajama.
  arXiv  Detail & Related papers  (2024-03-25T17:14:00Z)
• Efficient Online Data Mixing For Language Model Pre-Training [101.45242332613944]
  Existing data selection methods suffer from slow and computationally expensive processes. Data mixing, on the other hand, reduces the complexity of data selection by grouping data points together. We develop an efficient algorithm for Online Data Mixing (ODM) that combines elements from both data selection and data mixing.
  arXiv  Detail & Related papers  (2023-12-05T00:42:35Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.