CausalMix: Datamenging als causale inferentie voor het trainen van taalmodellen

Samenvatting

Bij de training van Large Language Models (LLM's) speelt datamenging een cruciale rol in het bepalen van de modelprestaties. Recente methoden optimaliseren menggewichten via proxymodellen, maar ze zijn gebaseerd op de aanname van statische dataverdelingen. Wanneer de onderliggende datapool verschuift, vereisen deze methoden dientengevolge kostbare hertraining vanaf nul. Deze beperking belemmert hun vermogen om naadloos op te schalen van kleine settings naar grotere datapools en modelgroottes. In dit artikel stellen we CausalMix voor om deze beperking aan te pakken door optimalisatie van datamengsels te modelleren als een causaal inferentieprobleem. We formuleren de statistische kenmerken van de datapool als covariaten en de domeinmenging als de behandeling. Na het fitten van een causaal model op 512 runs van Qwen2.5-0.5B om het Conditionele Gemiddelde Behandelingseffect (CATE) te schatten, extrapoleren we het optimale mengsel voor een datapool van 800K en passen dit toe om een 7B-model te trainen. Bovendien generaliseren we het raamwerk succesvol naar lange chain-of-thought-data op Qwen3-4B-Base. Door causale modellering te gebruiken om vertekening door confounders te isoleren, leidt CausalMix dynamisch toestandsafhankelijke optimale datamengsels af. Uitgebreide experimenten tonen aan dat het door CausalMix geleide mengsel consistent de prestaties verbetert op meerdere downstreamtaken, en beter presteert dan RegMix en andere basislijnen. Daarnaast gebruiken we de CATE Interpreter voor visuele analyse van de geleerde mengstrategie. Samenvattend biedt CausalMix een causaal en interpreteerbaar raamwerk voor het optimaliseren van LLM-datamengsels.

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.