CausalMix: 언어 모델 학습을 위한 인과 추론으로서의 데이터 혼합 (Data Mixture)
요약
CausalMix는 LLM 학습 시 데이터 혼합 최적화를 인과 추론 문제로 접근하여 해결하는 새로운 프레임워크입니다. 기존의 정적인 방식과 달리 데이터 풀의 변화에 유연하게 대응하며, CATE를 통해 최적의 혼합 비율을 동적으로 추론합니다.
핵심 포인트
- 데이터 혼합 최적화를 인과 추론(Causal Inference) 문제로 공식화
- 데이터 풀 변화에 따른 재학습 비용을 줄이는 외삽(Extrapolation) 기술 적용
- CATE를 활용하여 상태 의존적인 최적의 데이터 혼합 비율 도출
- Qwen 모델 시리즈를 통해 성능 향상 및 일반화 가능성 입증
- CATE Interpreter를 통한 학습된 혼합 전략의 시각적 분석 제공
대규모 언어 모델 (LLM) 학습에서 데이터 혼합 (data mixing)은 모델 성능을 결정하는 데 중추적인 역할을 합니다. 최근의 방법들은 프록시 모델 (proxy models)을 통해 혼합 가중치를 최적화하지만, 이는 정적인 데이터 분포 (static data distributions) 가정을 기반으로 합니다. 결과적으로, 근본적인 데이터 풀 (data pool)이 변화할 때 이러한 방법들은 처음부터 다시 수행해야 하는 비용이 많이 드는 재학습을 요구합니다. 이러한 한계는 소규모 설정에서 더 큰 데이터 풀과 모델 크기로 원활하게 확장하는 능력을 제한합니다. 본 논문에서는 데이터 혼합 최적화를 인과 추론 (causal inference) 문제로 설정함으로써 이러한 한계를 해결하기 위한 CausalMix를 제안합니다. 우리는 데이터 풀의 통계적 특징을 공변량 (covariates)으로, 도메인 혼합 (domain mixture)을 처치 (treatment)로 공식화합니다. 조건부 평균 처치 효과 (Conditional Average Treatment Effect, CATE)를 추정하기 위해 Qwen2.5-0.5B의 512회 실행 결과에 인과 모델을 적합시킨 후, 800K 데이터 풀에 대한 최적의 혼합을 외삽 (extrapolate)하여 7B 모델을 학습하는 데 적용합니다. 나아가, 우리는 이 프레임워크를 Qwen3-4B-Base의 긴 사고 사슬 (long chain-of-thought) 데이터로 성공적으로 일반화했습니다. 혼란 편향 (confounding biases)을 격리하기 위해 인과 모델링을 활용함으로써, CausalMix는 상태 의존적인 최적의 데이터 혼합을 동적으로 추론합니다. 광범위한 실험을 통해 CausalMix에 의해 가이드된 혼합이 여러 다운스트림 태스크 (downstream tasks)에서 일관되게 성능을 향상시키며, RegMix 및 기타 베이스라인 (baselines)보다 뛰어난 성능을 보임을 입증했습니다. 또한, 우리는 CATE Interpreter를 사용하여 학습된 혼합 전략에 대한 시각적 분석을 제공합니다. 종합적으로, CausalMix는 LLM 데이터 혼합을 최적화하기 위한 인과적이고 해석 가능한 프레임워크를 제공합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기