CausalMix: Datenmischung als kausale Inferenz für das Training von Sprachmodellen

Zusammenfassung

Im Training großer Sprachmodelle (LLMs) spielt die Datenmischung eine zentrale Rolle für die Modellleistung. Aktuelle Methoden optimieren Mischungsgewichte über Proxymodelle, basieren jedoch auf der Annahme statischer Datenverteilungen. Wenn sich der zugrunde liegende Datenpool verschiebt, erfordern diese Methoden daher kostspieliges erneutes Training von Grund auf. Diese Einschränkung verhindert eine nahtlose Skalierung von kleinen Umgebungen hin zu größeren Datenpools und Modellgrößen. In dieser Arbeit schlagen wir CausalMix vor, um diese Einschränkung zu beheben, indem wir die Optimierung von Datenmischungen als Problem der kausalen Inferenz betrachten. Wir formulieren die statistischen Merkmale des Datenpools als Kovariaten und die Domänenmischung als Behandlung. Nachdem wir ein kausales Modell auf 512 Läufen von Qwen2.5-0.5B zur Schätzung des bedingten durchschnittlichen Behandlungseffekts (CATE) angepasst haben, extrapolieren wir die optimale Mischung für einen 800K-Datenpool und wenden sie auf das Training eines 7B-Modells an. Darüber hinaus verallgemeinern wir das Framework erfolgreich auf lange Gedankenkettendaten (Chain-of-Thought) auf Qwen3-4B-Base. Durch die Nutzung kausaler Modellierung zur Isolierung von Störfaktoren (Confounding) leitet CausalMix dynamisch zustandsabhängige optimale Datenmischungen ab. Umfangreiche Experimente zeigen, dass die durch CausalMix geleitete Mischung die Leistung über mehrere nachgelagerte Aufgaben hinweg konsistent verbessert und RegMix sowie andere Basislinien übertrifft. Zusätzlich verwenden wir den CATE-Interpreter zur visuellen Analyse der erlernten Mischstrategie. Insgesamt bietet CausalMix ein kausales und interpretierbares Framework zur Optimierung von LLM-Datenmischungen.

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.