CausalMix: 인과 추론을 활용한 언어 모델 학습을 위한 데이터 혼합 기법

CausalMix: Data Mixture as Causal Inference for Language Model Training

대규모 언어 모델(LLM) 학습에서 데이터 혼합은 모델 성능에 중요한 역할을 합니다. 최근 연구에서는 프록시 모델을 사용하여 혼합 가중치를 최적화하지만, 이는 데이터 분포가 고정되어 있다는 가정에 의존합니다. 결과적으로, 기본 데이터 풀이 변경되면 이러한 방법들은 상당한 비용이 드는 재학습이 필요합니다. 이로 인해 작은 규모의 설정에서 대규모 데이터 풀 및 모델 크기로 원활하게 확장하는 데 제약이 있습니다. 본 논문에서는 데이터 혼합 최적화를 인과 추론 문제로 간주하여 이러한 제한 사항을 해결하기 위한 CausalMix를 제안합니다. 우리는 데이터 풀의 통계적 특징을 공변량으로, 도메인 혼합을 처치 변수로 정의했습니다. Qwen2.5-0.5B 모델 512번 실행에 대한 인과 모델을 학습하여 조건부 평균 처치 효과(CATE)를 추정하고, 이를 통해 800K 데이터 풀에 대한 최적의 혼합 비율을 예측하고 7B 모델 학습에 적용했습니다. 또한, 이 프레임워크를 Qwen3-4B-Base 모델의 긴 연쇄 사고(chain-of-thought) 데이터로 성공적으로 확장했습니다. CausalMix는 인과 모델링을 활용하여 교란 변수를 분리함으로써 상태에 따라 달라지는 최적의 데이터 혼합을 동적으로 추론합니다. 광범위한 실험 결과, CausalMix가 안내하는 혼합 방식은 RegMix 및 기타 기준 모델보다 여러 하위 작업에서 일관되게 성능을 향상시키는 것으로 나타났습니다. 또한, CATE 해석기를 사용하여 학습된 혼합 전략에 대한 시각적 분석을 제공합니다. 전반적으로, CausalMix는 LLM 데이터 혼합을 최적화하기 위한 인과적이고 해석 가능한 프레임워크를 제공합니다.

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.