CausalMix: Mezcla de datos como inferencia causal para el entrenamiento de modelos de lenguaje

Resumen

En el entrenamiento de Modelos de Lenguaje de Gran Escala (LLM), la mezcla de datos desempeña un papel crucial en la determinación del rendimiento del modelo. Métodos recientes optimizan los pesos de la mezcla mediante modelos proxy, pero se basan en el supuesto de distribuciones de datos estáticas. Como resultado, cuando el conjunto de datos subyacente cambia, estos métodos requieren un reentrenamiento costoso desde cero. Esta limitación restringe su capacidad para escalar sin problemas desde entornos pequeños hasta conjuntos de datos y tamaños de modelo más grandes. En este artículo, proponemos CausalMix para abordar esta limitación, planteando la optimización de la mezcla de datos como un problema de inferencia causal. Formulamos las características estadísticas del conjunto de datos como covariables y la mezcla de dominios como el tratamiento. Tras ajustar un modelo causal en 512 ejecuciones de Qwen2.5-0.5B para estimar el Efecto Promedio del Tratamiento Condicional (CATE), extrapolamos la mezcla óptima para un conjunto de datos de 800K y la aplicamos para entrenar un modelo de 7B. Además, generalizamos con éxito el marco a datos de cadena de pensamiento larga en Qwen3-4B-Base. Al aprovechar el modelado causal para aislar los sesgos de confusión, CausalMix infiere dinámicamente mezclas de datos óptimas dependientes del estado. Experimentos exhaustivos muestran que la mezcla guiada por CausalMix mejora consistentemente el rendimiento en múltiples tareas posteriores, superando a RegMix y otras líneas base. Adicionalmente, utilizamos el Intérprete CATE para proporcionar un análisis visual de la estrategia de mezcla aprendida. En conjunto, CausalMix ofrece un marco causal e interpretable para optimizar las mezclas de datos en 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.