研究者らは、既存の方法における静的分布仮定の限界に対処するため、大規模言語モデルのデータ混合最適化を因果推論問題として捉える手法であるCausalMixを提案する。このアプローチは統計的特徴を共変量、ドメイン混合を治療として定式化し、Qwen2.5-0.5Bの512回のランから条件付き平均治療効果(CATE)を推定して、より大規模なモデルに対する最適な混合を外挿する。

  • CausalMixは因果モデリングを活用して交絡バイアスを分離し、状態依存の最適なデータ混合を動的に推論する。
  • このフレームワークはQwen3-4B-Baseにおける長いchain-of-thoughtデータに対して正常に汎化できる。
  • 実験により、CausalMixによって誘導された混合は複数の下流タスク全体で一貫して性能を向上させ、RegMixや他のベースラインを上回ることを示した。
  • 学習された混合戦略の視覚的分析のためにCATE Interpreterが提供される。

CausalMixは、高コストな再トレーニングを行うことなく、小規模な設定からより大規模なデータプールへのシームレスなスケーリングを可能にする、因果的かつ解釈可能なLLMデータ混合最適化のためのフレームワークを提供する。