CausalMix: Смешивание данных как причинно-следственный вывод для обучения языковых моделей

Аннотация

В обучении больших языковых моделей (БЯМ) смешивание данных играет ключевую роль в определении производительности модели. Современные методы оптимизируют веса смеси с помощью прокси-моделей, однако они опираются на предположение о статичности распределений данных. В результате, при изменении базового пула данных эти методы требуют дорогостоящего переобучения с нуля. Данное ограничение мешает им масштабироваться от небольших конфигураций к более крупным пулам данных и размерам моделей. В этой статье мы предлагаем метод CausalMix для преодоления этого ограничения, формулируя оптимизацию смеси данных как задачу причинно-следственного вывода. Мы представляем статистические характеристики пула данных в качестве ковариат, а доменную смесь — в качестве воздействия. После подгонки причинно-следственной модели на 512 запусках Qwen2.5-0.5B для оценки усредненного условного эффекта воздействия (CATE) мы экстраполируем оптимальную смесь для пула данных объёмом 800 тыс. токенов и применяем её для обучения модели на 7B параметров. Кроме того, мы успешно обобщаем этот подход на данные с длинными цепочками рассуждений на основе Qwen3-4B-Base. Используя причинно-следственное моделирование для устранения смещающих факторов, CausalMix динамически определяет оптимальные смеси данных, зависящие от состояния. Обширные эксперименты показывают, что смесь, полученная с помощью CausalMix, последовательно улучшает производительность на множестве целевых задач, превосходя RegMix и другие базовые методы. Помимо этого, мы используем интерпретатор CATE для визуального анализа изученной стратегии смешивания. В целом, CausalMix предлагает причинно-следственную и интерпретируемую основу для оптимизации смесей данных в БЯМ.

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.