TiKMiX: 言語モデル事前学習における動的混合へのデータ影響の取り込み

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

August 25, 2025

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

cs.AI

要旨

言語モデルの事前学習に使用されるデータの混合比率は、その最終的な性能の基盤となる。しかし、静的な混合戦略は最適ではなく、モデルの様々なデータドメインに対する学習選好はトレーニングを通じて動的に変化する。特に、これらの変化する選好を計算効率的に観察することは依然として大きな課題である。この問題に対処するため、我々はTiKMiXを提案する。これは、モデルの変化する選好に応じてデータの混合比率を動的に調整する手法である。TiKMiXは、データドメインがモデルに与える影響を評価するための効率的な指標であるGroup Influenceを導入する。この指標により、データ混合問題を、影響を最大化する最適な分布の探索として定式化することが可能となる。我々はこれを2つのアプローチで解決する：直接最適化を行うTiKMiX-Dと、より優れた混合比率を予測する回帰モデルを使用するTiKMiX-Mである。最大1兆トークンまでのデータを用いて、異なるパラメータ数のモデルをトレーニングした。TiKMiX-Dは、REGMIXのような最先端の手法の性能を上回りながら、計算リソースのわずか20%しか使用しない。TiKMiX-Mは、9つの下流ベンチマークで平均2%の性能向上をもたらす。我々の実験は、モデルのデータ選好がトレーニングの進捗とスケールに応じて進化することを明らかにし、これらの選好を直接測定するGroup Influenceに基づいてデータの混合比率を動的に調整することで、静的な比率で見られるデータの消化不足を緩和し、性能を大幅に向上させることを実証している。

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.