CausalMix: 언어 모델 훈련을 위한 인과 추론으로서의 데이터 혼합

초록

대규모 언어 모델(LLM) 훈련에서 데이터 혼합은 모델 성능을 결정하는 데 핵심적인 역할을 한다. 최근 방법들은 프록시 모델을 통해 혼합 가중치를 최적화하지만, 이는 데이터 분포가 정적이라는 가정에 의존한다. 결과적으로, 기저 데이터 풀이 변화할 때 이러한 방법들은 처음부터 재훈련하는 데 많은 비용이 소요된다. 이러한 한계는 소규모 환경에서 더 큰 데이터 풀과 모델 크기로 원활하게 확장하는 능력을 제약한다. 본 논문에서는 데이터 혼합 최적화를 인과 추론 문제로 전환하여 이러한 한계를 해결하는 CausalMix를 제안한다. 데이터 풀의 통계적 특징을 공변량으로, 도메인 혼합을 처리로 정식화한다. Qwen2.5-0.5B의 512회 실행 데이터에 인과 모델을 적합하여 조건부 평균 처리 효과(CATE)를 추정한 후, 800K 데이터 풀에 대한 최적 혼합을 외삽하여 7B 모델 훈련에 적용한다. 나아가, 이 프레임워크를 Qwen3-4B-Base의 긴 사고 사슬 데이터로 성공적으로 일반화한다. 인과 모델링을 활용하여 교란 편향을 분리함으로써, CausalMix는 상태에 의존적인 최적 데이터 혼합을 동적으로 추론한다. 광범위한 실험을 통해 CausalMix가 안내하는 혼합이 여러 하위 작업에서 일관되게 성능을 향상시키며, RegMix 및 다른 기준선들을 능가함을 보인다. 또한, CATE 인터프리터를 사용하여 학습된 혼합 전략에 대한 시각적 분석을 제공한다. 종합적으로, 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.