CausalMix: Mistura de Dados como Inferência Causal para o Treinamento de Modelos de Linguagem

Resumo

No treinamento de Modelos de Linguagem de Grande Escala (LLMs), a mistura de dados desempenha um papel fundamental na determinação do desempenho do modelo. Métodos recentes otimizam pesos de mistura por meio de modelos proxy, mas dependem da suposição de distribuições de dados estáticas. Como resultado, quando o conjunto subjacente de dados muda, esses métodos exigem um retreinamento custoso do zero. Essa limitação restringe sua capacidade de escalar de forma contínua de configurações pequenas para conjuntos de dados e tamanhos de modelo maiores. Neste artigo, propomos o CausalMix para abordar essa limitação, enquadrando a otimização da mistura de dados como um problema de inferência causal. Formulamos as características estatísticas do conjunto de dados como covariáveis e a mistura de domínios como o tratamento. Após ajustar um modelo causal em 512 execuções do Qwen2.5-0.5B para estimar o Efeito Médio do Tratamento Condicional (CATE), extrapolamos a mistura ideal para um conjunto de dados de 800K e a aplicamos para treinar um modelo de 7B. Além disso, generalizamos com sucesso a estrutura para dados de cadeia de pensamento longa no Qwen3-4B-Base. Ao utilizar a modelagem causal para isolar vieses de confusão, o CausalMix infere dinamicamente misturas de dados ideais dependentes do estado. Experimentos extensivos mostram que a mistura guiada pelo CausalMix melhora consistentemente o desempenho em várias tarefas downstream, superando o RegMix e outras baselines. Além disso, usamos o Interpretador CATE para fornecer uma análise visual da estratégia de mistura aprendida. Em suma, o CausalMix oferece uma estrutura causal e interpretável para otimizar misturas de dados em LLMs.

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.