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

2025-08-25

AI Summary

In language model pretraining, static data mixing strategies fail to adapt to the dynamic evolution of model preferences across domains and lack efficient mechanisms for quantifying domain-specific influence. This paper proposes TiKMiX, the first framework to introduce the Group Influence metric—a differentiable, computationally efficient measure of each domain’s training impact—formulating data mixing as an influence-maximization optimization problem. TiKMiX supports two adaptive strategies: gradient-driven direct optimization and regression-based prediction, integrated with dynamic resource allocation and large-scale distributed training. Evaluated on a trillion-token corpus, TiKMiX-D achieves superior performance to REGMIX at 20% lower computational cost; TiKMiX-M attains an average 2.0% improvement across nine downstream tasks, significantly mitigating the “indigestion” problem—ineffective learning from poorly mixed data—commonly observed with static mixing.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Language Grounding & Multi-modal NLP
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Dynamic data mixture adjustment for language model pre-training
• Efficiently measuring data domain influence on model preferences
• Optimizing data distribution to maximize training performance gains

Innovation

• Dynamic data mixture adjustment using Group Influence
• Two approaches: TiKMiX-D optimization and TiKMiX-M prediction
• Efficient metric evaluating data domain impact on model