CausalMix : Mélange de données comme inférence causale pour l'entraînement de modèles de langage

Résumé

Dans l'entraînement des grands modèles de langage (LLM), le mélange de données joue un rôle crucial dans la détermination des performances du modèle. Les méthodes récentes optimisent les poids de mélange via des modèles proxy, mais elles reposent sur l'hypothèse de distributions de données statiques. Par conséquent, lorsque le pool de données sous-jacent évolue, ces méthodes nécessitent un réentraînement coûteux depuis zéro. Cette limitation restreint leur capacité à passer à l'échelle de manière fluide, depuis de petites configurations jusqu'à des pools de données et des tailles de modèle plus importants. Dans cet article, nous proposons CausalMix pour surmonter cette limitation en formulant l'optimisation du mélange de données comme un problème d'inférence causale. Nous définissons les caractéristiques statistiques du pool de données comme des covariables et le mélange de domaines comme le traitement. Après avoir ajusté un modèle causal sur 512 exécutions de Qwen2.5-0.5B pour estimer l'effet conditionnel moyen du traitement (CATE), nous extrapolons le mélange optimal pour un pool de 800 000 données et l'appliquons à l'entraînement d'un modèle de 7B. De plus, nous généralisons avec succès le cadre à des données de longues chaînes de pensée sur Qwen3-4B-Base. En exploitant la modélisation causale pour isoler les biais de confusion, CausalMix déduit dynamiquement des mélanges de données optimaux dépendant de l'état. Des expériences approfondies montrent que le mélange guidé par CausalMix améliore systématiquement les performances sur plusieurs tâches en aval, surpassant RegMix et d'autres références. En outre, nous utilisons l'interpréteur CATE pour fournir une analyse visuelle de la stratégie de mélange apprise. En résumé, CausalMix offre un cadre causal et interprétable pour l'optimisation des mélanges de données pour les 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.