見出し画像

【論文瞬読】多トークン予測: 大規模言語モデルの新しい学習パラダイム

AI Nest

こんにちは！株式会社AI Nestです。今日は、大規模言語モデル (Large Language Models; LLMs) の訓練に関する興味深い研究について紹介したいと思います。最近、Gloeckleらによる「Better & Faster Large Language Models via Multi-token Prediction」という論文を読む機会がありました。この論文では、LLMsの訓練において、次のトークン予測の代わりに多トークン予測を用いることで、サンプル効率と推論速度を向上させる手法が提案されています。

タイトル：Better & Faster Large Language Models via Multi-token Prediction
URL：https://arxiv.org/abs/2404.19737
著者：Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

従来のLLM訓練の課題

LLMsは、大規模なテキストコーパスを用いて、次のトークン予測というタスクで訓練されるのが一般的です。つまり、モデルは与えられた文脈に基づいて、次に来るトークンを予測するように学習します。しかし、この手法には課題があります。

教師あり学習 (teacher forcing) と呼ばれる手法で訓練されたLLMsは、自己回帰的な生成の際に、訓練時とは異なる分布のデータを入力として受け取ることになります。これにより、モデルは局所的なパターンに過剰に適合し、長期的な依存関係を適切に捉えられなくなる可能性があります。

多トークン予測とは？

多トークン予測は、この課題に対処するために提案された新しい学習パラダイムです。この手法では、LLMsの訓練時に、次のトークンだけでなく、複数の将来のトークンを一度に予測します。具体的には、共有のモデルトランクと複数の出力ヘッドを用いて、各位置で次のn個のトークンを独立に予測します。

多トークン予測の概要

この方法により、モデルは長期的な依存関係を考慮しながら、効率的に言語を学習できるようになります。また、教師あり学習と自己回帰的生成の間の分布の不一致を軽減することができます。

順伝播と逆伝播の順序

実験結果が示す多トークン予測の有効性

論文では、さまざまなタスクにおいて多トークン予測の有効性が実験的に示されています。以下は、主な実験結果の概要です。

  1. モデルサイズと多トークン予測の効果の関係

    • モデルサイズを大きくするほど、多トークン予測の効果が顕著になる

    • 130億パラメータのモデルでは、コード生成タスクにおいて最大17％の性能向上

  2. 推論速度の向上

    • 推論時に追加の予測ヘッドを利用することで、最大3倍の高速化を達成

  3. バイトレベルのトークン化における効果

    • 多トークン予測がより長期的なパターンの学習を促進

    • 8バイト予測モデルが、次のバイト予測モデルに比べ、MBPPで67％、HumanEvalで20％の性能向上

  4. 多エポック学習における優位性

    • 多エポック学習でも、多トークン予測の優位性が維持される

  5. ファインチューニングタスクにおける性能

    • 多トークン予測の事前学習モデルが、CodeContestsなどのファインチューニングタスクでも優れた性能を示す

モデルサイズ別の多トークン予測モデルのMBPP結果
多トークン予測によるパフォーマンスの向上とバイトレベルの効率的な訓練の可能性
CodeContestsでのファインチューニング性能の比較

これらの結果は、多トークン予測がLLMsの性能向上に寄与する有望なアプローチであることを示しています。

多トークン予測の有効性の理論的考察

論文では、多トークン予測の有効性を説明するための理論的考察も提示されています。

  1. 教師あり学習と自己回帰的生成の分布の不一致の軽減

    • 多トークン予測は、訓練時と推論時の分布の不一致を軽減することで、モデルの汎化性能を向上させる

  2. 情報理論的な観点からの損失関数の分解

    • 多トークン予測の損失関数は、局所的なクロスエントロピー項、相互情報項、およびシフトされた次のトークンのクロスエントロピー項に分解できる

    • これにより、モデルは将来のトークンの予測に有用な特徴を事前に計算するように動機付けられる

ただし、これらの議論はまだ直感的なレベルにとどまっており、より厳密な理論的裏付けが求められます。

今後の展望と課題

多トークン予測は、LLMsの性能向上に向けた新しいアプローチとして期待できます。今後は、以下のような方向性で研究が進められることが望まれます。

  1. より大規模なモデルやより多様なタスクへの適用

    • 多トークン予測の真の潜在力を明らかにするために、さらに大規模なモデルや多様なタスクへの適用が求められる

  2. 自然言語処理タスクにおける効果の検証

    • 要約タスクでは一定の改善が見られたが、質問応答などのほかのタスクへの適用可能性についても調べる必要がある

  3. 理論的裏付けの強化

    • 多トークン予測の有効性を説明するための理論的枠組みの構築が求められる

  4. 最適な予測トークン数の探索

    • タスクやデータセットに応じた最適な予測トークン数の選択方法の確立が望まれる

7Bモデルを用いた複数トークン訓練では選択タスクのパフォーマンスが向上しない

これらの課題に取り組むことで、多トークン予測はLLMsの発展に大きく貢献することが期待されます。

まとめ

Gloeckleらの論文は、LLMsの訓練に新しい視点をもたらしました。多トークン予測というシンプルながら有望なアイデアを提示し、その有効性を実験的に示した点で高く評価できます。また、理論的考察を通じて、多トークン予測の働きを説明しようとする試みも興味深いものがありました。

この研究は、LLMsのさらなる発展に向けた重要な一歩を踏み出したと言えるでしょう。今後、より大規模なモデルやより多様なタスクへの適用、理論的裏付けの強化などを通じて、多トークン予測の真の可能性が明らかになることを期待したいと思います！