Rho-1: すべてのトークンが必要なわけではない

Rho-1: Not All Tokens Are What You Need

April 11, 2024

著者: Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, Weizhu Chen

cs.AI

要旨

従来の言語モデルの事前学習手法では、すべての訓練トークンに対して一様に次のトークン予測損失を適用してきました。この常識に挑戦し、我々は「コーパス内のすべてのトークンが言語モデルの訓練において同等に重要ではない」と主張します。我々の初期分析では、言語モデルのトークンレベルの訓練ダイナミクスを探り、異なるトークンに対する損失パターンの違いを明らかにしました。これらの知見を活用して、我々は新しい言語モデルRho-1を提案します。従来の言語モデルがコーパス内のすべての次のトークンを予測するのとは異なり、Rho-1はSelective Language Modeling（SLM）を採用し、目的の分布に沿った有用なトークンを選択的に訓練します。このアプローチでは、参照モデルを使用して事前学習トークンをスコアリングし、その後、過剰損失が高いトークンに焦点を当てた損失で言語モデルを訓練します。15BのOpenWebMathコーパスで継続的に事前学習を行うと、Rho-1は9つの数学タスクにおいて、最大30%の絶対的なfew-shot精度の向上をもたらしました。ファインチューニング後、Rho-1-1Bと7Bは、それぞれMATHデータセットで40.6%と51.8%の最先端の結果を達成し、事前学習トークンのわずか3%でDeepSeekMathに匹敵する性能を示しました。さらに、80Bの一般トークンで事前学習を行うと、Rho-1は15の多様なタスクにわたって平均6.8%の向上を達成し、言語モデルの事前学習の効率と性能の両方を向上させました。

English

Previous language model pre-training methods have uniformly applied a next-token prediction loss to all training tokens. Challenging this norm, we posit that "Not all tokens in a corpus are equally important for language model training". Our initial analysis delves into token-level training dynamics of language model, revealing distinct loss patterns for different tokens. Leveraging these insights, we introduce a new language model called Rho-1. Unlike traditional LMs that learn to predict every next token in a corpus, Rho-1 employs Selective Language Modeling (SLM), which selectively trains on useful tokens that aligned with the desired distribution. This approach involves scoring pretraining tokens using a reference model, and then training the language model with a focused loss on tokens with higher excess loss. When continual pretraining on 15B OpenWebMath corpus, Rho-1 yields an absolute improvement in few-shot accuracy of up to 30% in 9 math tasks. After fine-tuning, Rho-1-1B and 7B achieved state-of-the-art results of 40.6% and 51.8% on MATH dataset, respectively - matching DeepSeekMath with only 3% of the pretraining tokens. Furthermore, when pretraining on 80B general tokens, Rho-1 achieves 6.8% average enhancement across 15 diverse tasks, increasing both efficiency and performance of the language model pre-training.