TiKMiX: Incorporando a Influência de Dados em Misturas Dinâmicas para Pré-treinamento de Modelos de Linguagem

TiKMiX: Take Data Influence into Dynamic Mixture for Language Model Pre-training

August 25, 2025

Autores: Yifan Wang, Binbin Liu, Fengze Liu, Yuanfan Guo, Jiyao Deng, Xuecheng Wu, Weidong Zhou, Xiaohuan Zhou, Taifeng Wang

cs.AI

Resumo

A mistura de dados utilizada no pré-treinamento de um modelo de linguagem é um pilar fundamental para seu desempenho final. No entanto, uma estratégia de mistura estática é subótima, uma vez que as preferências de aprendizado do modelo para diversos domínios de dados mudam dinamicamente ao longo do treinamento. Crucialmente, observar essas preferências em evolução de forma computacionalmente eficiente permanece um desafio significativo. Para abordar isso, propomos o TiKMiX, um método que ajusta dinamicamente a mistura de dados de acordo com as preferências em evolução do modelo. O TiKMiX introduz o conceito de Influência de Grupo, uma métrica eficiente para avaliar o impacto dos domínios de dados no modelo. Essa métrica permite formular o problema de mistura de dados como uma busca por uma distribuição ótima que maximize a influência. Resolvemos isso por meio de duas abordagens: o TiKMiX-D para otimização direta e o TiKMiX-M, que utiliza um modelo de regressão para prever uma mistura superior. Treinamos modelos com diferentes números de parâmetros, em até 1 trilhão de tokens. O TiKMiX-D supera o desempenho de métodos state-of-the-art como o REGMIX, utilizando apenas 20% dos recursos computacionais. O TiKMiX-M resulta em um ganho médio de desempenho de 2% em 9 benchmarks de tarefas downstream. Nossos experimentos revelam que as preferências de dados de um modelo evoluem com o progresso do treinamento e a escala, e demonstramos que ajustar dinamicamente a mistura de dados com base na Influência de Grupo, uma medida direta dessas preferências, melhora significativamente o desempenho ao mitigar a subdigestão de dados observada com proporções estáticas.

English

The data mixture used in the pre-training of a language model is a cornerstone of its final performance. However, a static mixing strategy is suboptimal, as the model's learning preferences for various data domains shift dynamically throughout training. Crucially, observing these evolving preferences in a computationally efficient manner remains a significant challenge. To address this, we propose TiKMiX, a method that dynamically adjusts the data mixture according to the model's evolving preferences. TiKMiX introduces Group Influence, an efficient metric for evaluating the impact of data domains on the model. This metric enables the formulation of the data mixing problem as a search for an optimal, influence-maximizing distribution. We solve this via two approaches: TiKMiX-D for direct optimization, and TiKMiX-M, which uses a regression model to predict a superior mixture. We trained models with different numbers of parameters, on up to 1 trillion tokens. TiKMiX-D exceeds the performance of state-of-the-art methods like REGMIX while using just 20% of the computational resources. TiKMiX-M leads to an average performance gain of 2% across 9 downstream benchmarks. Our experiments reveal that a model's data preferences evolve with training progress and scale, and we demonstrate that dynamically adjusting the data mixture based on Group Influence, a direct measure of these preferences, significantly improves performance by mitigating the underdigestion of data seen with static ratios.