AIモデルとプラットフォーム
大規模言語モデルをマルチトークン予測で高性能化する
Unite.AI を Google の優先ソースに追加
大規模言語モデル (LLM) である GPT、LLaMA などのモデルは、人間のようなテキストを理解し、生成する驚くべき能力で世界を驚かせました。ただし、これらのモデルの標準的なトレーニング方法である「次のトークン予測」には、いくつかの内在的な制限があります。
次のトークン予測では、モデルは、前の単語の後に続く単語を予測するようにトレーニングされます。 このアプローチは成功を収めてきましたが、長距離の依存関係や複雑な推論タスクに苦労するモデルにつながる可能性があります。さらに、教師強制トレーニング方式と推論時の自己回帰生成プロセスの不一致は、サブオプティマルなパフォーマンスにつながる可能性があります。
Meta AI の Gloeckle ら (2024) による最近の研究論文では、次のトークン予測の限界を解決しようとする新しいトレーニングパラダイムである「マルチトークン予測」が紹介されています。このブログ投稿では、この革新的な研究の核心概念、技術的詳細、および潜在的な影響について深く掘り下げます。
シングルトークン予測：従来のアプローチ
マルチトークン予測の詳細に踏み込む前に、大規模言語モデルのトレーニングにおいて数年間「主力」として使用されてきた従来のアプローチ、つまりシングルトークン予測（次のトークン予測とも呼ばれる）を理解することが重要です。
次のトークン予測パラダイム
次のトークン予測パラダイムでは、言語モデルは、前のコンテキストに基づいて、シーケンス内の次の単語を予測するようにトレーニングされます。より正式に、モデルは、前のトークン x1、x2、…、xt が与えられた場合に、次のトークン xt+1 の確率を最大化するように設計されます。これは、次のクロスエントロピー損失を最小化することで実現されます:
L = -Σt log P(xt+1 | x1, x2, …, xt)
このシンプルながら強力なトレーニングオブジェクトは、GPT (Radford et al., 2018)、BERT (Devlin et al., 2019) などの多くの成功した大規模言語モデルの基礎となりました。
教師強制と自己回帰生成
次のトークン予測は、教師強制というトレーニング技術に依存しています。ここで、モデルは、トレーニング中の各将来のトークンに対して、正解ラベルを提供されます。これにより、モデルは正しいコンテキストとターゲットシーケンスから学び、より安定したトレーニングが可能になります。
しかし、推論または生成中に、モデルは自己回帰的に動作し、前の生成されたトークンに基づいて、1 つずつトークンを予測します。トレーニング方式 (教師強制) と推論方式 (自己回帰生成) の不一致は、特に長いシーケンスや複雑な推論タスクの場合、潜在的な不一致やサブオプティマルなパフォーマンスにつながる可能性があります。
次のトークン予測の限界
次のトークン予測は驚くべき成功を収めてきましたが、いくつかの内在的な限界があります:
- 短距離依存: 次のトークンだけを予測することで、モデルは長距離の依存関係やテキストの全体的な構造と一貫性を捉えるのに苦労する可能性があります。これにより、不一致または無意味な生成につながる可能性があります。
- ローカルパターン依存: 次のトークン予測モデルは、トレーニングデータのローカルパターンに依存する可能性があり、分布外のシナリオや抽象的な推論を必要とするタスクに一般化するのが難しい場合があります。
- 推論能力: 複数ステップの推論、アルゴリズム的思考、または複雑な論理演算を必要とするタスクの場合、次のトークン予測では、十分な帰納的バイアスや表現を提供できない可能性があります。
- サンプルの非効率性: 次のトークン予測のローカル性のため、モデルは必要な知識や推論スキルを習得するために、より大きなトレーニングデータセットを必要とする可能性があり、潜在的なサンプルの非効率性につながる可能性があります。
これらの限界は、研究者が代替のトレーニングパラダイム、たとえばマルチトークン予測を探求する動機となりました。これは、これらの限界の一部を解決し、大規模言語モデルの新しい能力を解放することを目指しています。
従来の次のトークン予測アプローチと新しいマルチトークン予測テクニックを対比することで、読者は後者の動機と潜在的な利点をよりよく理解できます。これにより、この革新的な研究のより深い探求の準備が整います。
マルチトークン予測とは
マルチトークン予測の背後にある基本的な考え方は、言語モデルを単に次のトークンを予測するのではなく、同時に複数の将来のトークンを予測するようにトレーニングすることです。具体的には、トレーニング中、モデルは、トレーニングコーパス内の各位置で、n 個の独立した出力ヘッドを使用して、次の n 個のトークンを予測するように設計されます。
たとえば、4 トークン予測のセットアップでは、モデルは、前のコンテキストに基づいて、1 回に 4 つの次のトークンを予測するようにトレーニングされます。このアプローチにより、モデルは長距離の依存関係を捉え、テキストの全体的な構造と一貫性をよりよく理解するよう促します。
簡単な例
マルチトークン予測の概念をよりよく理解するために、簡単な例を考えてみましょう。
「The quick brown fox jumps over the lazy dog.」
従来の次のトークン予測アプローチでは、モデルは前のコンテキストに基づいて、次の単語を予測するようにトレーニングされます。たとえば、コンテキスト「The quick brown fox jumps over the」が与えられた場合、モデルは次の単語「lazy」を予測するように設計されます。
しかし、マルチトークン予測では、モデルは同時に複数の将来のトークンを予測するようにトレーニングされます。たとえば、n=4 の場合、モデルは、同じコンテキスト「The quick brown fox jumps over the」に基づいて、1 回に 4 つの次のトークンを予測するように設計されます。結果として、モデルは「lazy dog .」(ドットは文の終わりを示す) のシーケンスを予測するように設計されます。
モデルを複数の将来のトークンを同時に予測するようにトレーニングすることで、長距離の依存関係を捉え、テキストの全体的な構造と一貫性をよりよく理解するよう促します。
技術的詳細
著者は、マルチトークン予測を実装するためのシンプルながら有効なアーキテクチャを提案しています。モデルは、共有トランスフォーマートランクと、次のトークンを予測するための n 個の独立したトランスフォーマーレイヤー (出力ヘッド) で構成されます。
トレーニング中、フォワードパスとバックワードパスは、GPU メモリフットプリントを最小限に抑えるように慎重に調整されます。共有トランクは潜在的な表現を計算し、各出力ヘッドは順番にフォワードパスとバックワードパスを実行し、トランクレベルで勾配を蓄積します。このアプローチにより、すべてのログリットベクトルとその勾配を同時にマテリアライズする必要がなくなり、ピークの GPU メモリ使用量が O(nV + d) から O(V + d) に削減されます。ここで、V は語彙サイズ、d は潜在的な表現の次元です。
メモリ効率の高い実装
マルチトークン予測モデルのトレーニングにおける課題の 1 つは、GPU メモリ使用量を削減することです。語彙サイズ (V) が潜在的な表現の次元 (d) よりも大きいため、ログリットベクトルが GPU メモリ使用量のボトルネックになります。
この課題に対処するために、著者は、フォワードパスとバックワードパスのシーケンスを慎重に調整することで、メモリ効率の高い実装を提案しています。すべてのログリットとその勾配を同時にマテリアライズするのではなく、各独立した出力ヘッドのフォワードパスとバックワードパスを順番に計算し、トランクレベルで勾配を蓄積します。
このアプローチにより、すべてのログリットベクトルとその勾配を同時にメモリに保持する必要がなくなり、ピークの GPU メモリ使用量が O(nV + d) から O(V + d) に削減されます。ここで、n は予測される将来のトークンの数です。
マルチトークン予測の利点
研究論文では、マルチトークン予測を使用することによるいくつかの利点が提示されています:
- サンプルの効率性の向上: モデルを複数の将来のトークンを同時に予測するようにトレーニングすることで、マルチトークン予測はサンプルの効率性を向上させます。著者は、コード理解と生成タスクで、13B パラメータを持つモデルが平均で約 15% 多くの問題を解決することを実証しています。
- 推論の高速化: マルチトークン予測でトレーニングされた追加の出力ヘッドは、自己推測デコーディングのバリアントである並列トークン予測に利用できます。これにより、幅広いバッチサイズに対して最大 3 倍の高速な推論時間が実現します。
- 長距離依存の促進: マルチトークン予測は、モデルが長距離の依存関係やパターンを捉えることを促します。これは、より大きなコンテキストで理解と推論を必要とするタスクに特に有益です。
- アルゴリズム的推論: 著者は、合成タスクの実験を提示し、マルチトークン予測モデルが、小さいモデルサイズでも、帰納的頭部とアルゴリズム的推論能力の開発において優れていることを示しています。
- 一貫性と整合性: モデルを複数の将来のトークンを同時に予測するようにトレーニングすることで、マルチトークン予測は、一貫性と整合性のある表現の開発を促します。これは、長いテキストの生成、物語、創造的な文章、または説明的なマニュアルの生成などのタスクに特に有益です。
- 一般化の向上: 著者らの実験は、マルチトークン予測モデルが、特に分布外の設定で、より優れた一般化能力を示すことを示唆しています。これは、モデルの長距離のパターンと依存関係を捉える能力により、未知のシナリオに効果的に外挿できるためであると考えられます。
例と直感
マルチトークン予測がなぜうまく機能するのかについて、より多くの直感を提供するために、いくつかの例を考えてみましょう:
- コード生成: コード生成のコンテキストでは、複数のトークンを同時に予測することで、モデルはより複雑なコード構造を理解して生成できます。たとえば、関数定義を生成する場合、単に次のトークンを予測するだけで、関数シグネチャを正確に生成するのに十分なコンテキストが得られない可能性があります。ただし、複数のトークンを同時に予測することで、モデルは関数名、パラメータ、戻り値の型などの依存関係をよりよく捉えることができ、より正確で一貫性のあるコード生成につながります。
- 自然言語推論: 自然言語推論タスクでは、モデルは複数のステップや情報のピースを推論する必要があります。複数のトークンを同時に予測することで、モデルは推論プロセスのさまざまなコンポーネント間の依存関係をよりよく捉えることができ、より一貫性のある応答につながります。
- 長文生成: 長文の生成、たとえば物語、記事、またはレポートの場合、次のトークン予測でトレーニングされたモデルは、長い期間にわたって一貫性と整合性を維持するのに苦労する可能性があります。マルチトークン予測は、モデルがテキストの全体的な構造と流れを捉えることを促します。これは、より一貫性のある長文の生成につながります。
限界と将来の方向
論文で提示された結果は印象的ですが、いくつかの限界と将来の研究のための開かれた質問があります:
- 最適なトークンの数: 論文では、さまざまな n (予測される将来のトークンの数) の値を探索し、n=4 が多くのタスクでうまく機能することを発見しました。ただし、最適な n の値は、特定のタスク、データセット、モデルサイズに依存する可能性があります。最適な n を決定するための原則的な方法の開発は、さらにパフォーマンスの向上につながる可能性があります。
- 語彙サイズとトークン化: 著者は、マルチトークン予測モデルに対する最適な語彙サイズとトークン化戦略が、次のトークン予測モデルと異なる可能性があることを指摘しています。これを探求することで、シーケンスの長さと計算効率性の間のより良いトレードオフにつながる可能性があります。
- 補助的な予測損失: 著者は、彼らの研究が、大規模言語モデルに対する新しい補助的な予測損失の開発に興味をもたらす可能性があることを示唆しています。代替的な補助的な損失とマルチトークン予測との組み合わせを探求することは、興味深い研究方向です。
- 理論的な理解: 論文では、マルチトークン予測がなぜうまく機能するのかについて、いくつかの直感と実証的な証拠が提示されています。ただし、このアプローチがなぜうまく機能するのかについて、より深い理論的な理解は貴重です。
結論
Gloeckle ら (2024) の研究論文「マルチトークン予測を使用した大規模言語モデルの高速化と高性能化」は、新しいトレーニングパラダイムを紹介し、大規模言語モデルのパフォーマンスと能力を大幅に改善する可能性があります。モデルを複数の将来のトークンを同時に予測するようにトレーニングすることで、マルチトークン予測は、長距離の依存関係、アルゴリズム的推論能力、およびサンプルの効率性の向上を促します。
著者によって提案された技術的実装は、シンプルながら計算効率性が高く、大規模言語モデルのトレーニングに適用可能です。さらに、自己推測デコーディングのためにマルチトークン予測でトレーニングされた出力ヘッドを利用できることは、実用的には大きな利点です。
まだ開かれた質問や探求すべき領域があるものの、この研究は、大規模言語モデルの分野における重要な進歩を表しています。大規模言語モデルの需要が増え続ける中、マルチトークン予測は、これらの強力な AI システムの次の世代の重要なコンポーネントとなる可能性があります。
