CausalMix: 言語モデル訓練のための因果推論としてのデータ混合

要旨

大規模言語モデル（LLM）の訓練において、データ混合はモデルの性能を決定づける重要な役割を果たす。近年の手法では、プロキシモデルを用いて混合比率を最適化するが、それらはデータ分布が静的であるという仮定に依存している。そのため、基盤となるデータプールが変化すると、これらの手法ではコストのかかる再訓練を一から行う必要が生じる。この制約により、小規模な設定からより大規模なデータプールやモデルサイズへのシームレスなスケーリングが困難になる。本稿では、データ混合最適化を因果推論問題として捉えることでこの制約に対処するCausalMixを提案する。データプールの統計的特徴を共変量、ドメイン混合を処置として定式化する。Qwen2.5-0.5Bを用いた512回の実行で因果モデルを適合させ、条件付き平均処置効果（CATE）を推定した後、80万データプールに対する最適混合比率を外挿し、それを7Bモデルの訓練に適用する。さらに、このフレームワークをQwen3-4B-Base上の長い思考連鎖データにまで一般化することに成功した。因果モデリングを活用して交絡バイアスを分離することにより、CausalMixは状態に依存した最適なデータ混合比率を動的に推定する。広範な実験により、CausalMixによって導かれた混合比率が複数の下流タスクで一貫して性能向上をもたらし、RegMixなどのベースラインを凌駕することが示された。さらに、CATE Interpreterを用いて学習された混合戦略の可視化分析を提供する。総じて、CausalMixはLLMのデータ混合を最適化するための因果的で解釈可能なフレームワークを提供する。

English

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.