今回の論文
今回取り上げるのは、Fabian Gloeckle、Badr Youbi Idrissi、Baptiste Roziere、David Lopez-Paz、Gabriel Synnaeve による論文「Better & Faster Large Language Models via Multi-token Prediction」です。2024年4月30日に arXiv で公開されました。公開元は arXiv、研究分野は大規模言語モデル、学習目標設計、推論高速化です。URL は https://arxiv.org/abs/2404.19737 です。
この論文を選んだ理由は、モデル本体を大きく作り替えずに、学習時の目的関数を少し拡張するだけで性能向上と推論高速化の両方に効くからです。基盤モデル開発だけでなく、コード生成、社内向け LLM、軽量推論基盤の設計にも発想を持ち込みやすく、実装のヒントが多い論文です。
どんな技術か
Multi-Token Prediction は、LLM に「次の1トークン」だけではなく「次の複数トークン」を同時に予測させる学習手法です。
通常の自己回帰型 LLM は、ある位置までの文脈を見て、次に来る1トークンだけを当てるように学習します。これはシンプルで強力ですが、学習信号がかなり局所的です。たとえばコード生成や推論では、先の展開を見据えた表現を内部で作れたほうが有利なのに、訓練では直後の1文字や1単語だけを当てる圧力が強くなります。
そこでこの論文では、共有された Transformer 本体の上に複数の予測ヘッドを載せ、2個先、4個先、場合によっては8個先までのトークンを並列に予測させます。これにより、モデルは局所的な次トークンだけでなく、少し先の流れも意識した表現を学びやすくなります。
さらに面白いのは、この追加ヘッドが学習時だけの飾りではなく、推論時には自己推測デコーディングのドラフト役としても使えることです。つまり、学習目標の工夫がそのまま推論最適化にもつながっています。
課題
この技術が解決しようとしているのは、LLM の標準的な next-token prediction だけでは、学習効率と生成性能の両方に限界があるという課題です。
何が難しいのかというと、次トークン予測は非常にローカルな教師信号だからです。現在位置から見て直後の1トークンだけを当てる訓練では、少し先まで見通した構造やアルゴリズム的なパターンを強く学ぶとは限りません。特にコード生成や手順的なテキスト生成では、先の展開を内部表現として持てるかどうかが品質に効きます。
既存の方法ではどこに限界があるのかというと、モデル性能を上げるための王道は、より大きなモデル、より多いデータ、より長い学習です。しかしこれは計算資源の要求が大きく、改善のためのコストが重いです。また、推論を速くする方法として speculative decoding のような手法はありますが、通常は別のドラフトモデルや複雑な実行パイプラインが必要になります。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは、学習コストと推論コストの両方が製品価値に直結するからです。コード補完、エージェントのツール呼び出し、社内文書の生成支援、チャット応答などでは、精度だけ高くても遅いと使いにくく、速いだけで品質が落ちても採用されません。
実際の AI システムでは、推論時の待ち時間を下げたいが、別モデル追加で運用は複雑にしたくない、という場面がよくあります。また、限られたデータや学習予算で少しでも強いモデルを作りたいという要件もあります。Multi-Token Prediction は、その両方に対して比較的筋の良い打ち手を出している点が重要です。
用語解説
- Next-Token Prediction
- 自己回帰型 LLM の基本的な学習目標です。直前までの文脈から次の1トークンだけを予測します。この記事では、この標準目標を少し先まで拡張するのが Multi-Token Prediction だと理解すると読みやすくなります。
- Output Head
- 共有されたモデル本体の上で、最終的な予測を出す層です。この論文では、1個先、2個先、3個先といった未来位置ごとに独立した head を持たせるのが核心です。
- Teacher Forcing
- 学習時に、モデルが自分で出したトークンではなく正解トークン列を入力として使う訓練方法です。通常の next-token 学習はこの枠組みで進むため、局所的な予測に寄りやすいという背景があります。
- Self-Speculative Decoding
- 同じモデル自身の補助ヘッドを使って先のトークン候補をまとめて提案し、メインヘッドで検証しながら高速生成する方法です。別ドラフトモデルが不要なので、運用を複雑化しにくい点が実務上重要です。
- pass@k
- コード生成でよく使われる指標で、k 回生成した中に正解が含まれる確率を表します。この論文は HumanEval や MBPP で pass@1 だけでなく pass@10、pass@100 も比較しており、単発精度だけでなく多様な解答能力も見ています。
技術の仕組み
Multi-Token Prediction の仕組みは、発想自体はかなりシンプルです。ただし、実際に大規模学習へ載せるにはメモリ使用量や推論とのつなぎ方に工夫があります。
基本アイデア
通常の next-token prediction では、時刻 t までの入力を見て、t+1 の1トークンだけを予測します。これに対して Multi-Token Prediction では、同じ文脈から t+1、t+2、…、t+n までを同時に予測します。
重要なのは、未来のトークン列を1本の連鎖として逐次生成するのではなく、共有された内部表現から複数の独立 head で並列予測する点です。つまり、モデル本体は共通の文脈表現を作り、その上で「1個先担当」「2個先担当」「3個先担当」という枝を生やすイメージです。
モデル構造
論文の構成は、共有 Transformer trunk と、未来位置ごとに用意された独立 head からなります。各 head は shared trunk の出力表現を受け取り、それぞれ別の未来トークンを予測します。unembedding 行列は共有されるため、全体としては大改造ではありません。
この構造のポイントは、モデル全体のパラメータを不公平に増やしていないことです。著者らは、追加した head 分だけ trunk 側の層数を減らし、比較対象の next-token モデルと総パラメータ数をそろえています。つまり、単純に大きなモデルにしたから勝った、という解釈を避けています。
学習方法
学習目標は、1ステップ先の損失だけではなく、n ステップ先までの損失を足し合わせたものです。数式で書くと未来トークン列全体の対数尤度を最大化する形ですが、実装上は「各未来位置の cross-entropy を並列に最適化する」と捉えると十分です。
この設計によって、共有 trunk の中間表現には「今この1トークンを当てる」だけでなく、「この先の流れまである程度読める」ことが求められます。論文はこれを、よりグローバルなパターン学習やアルゴリズム的な構造理解に有利だと位置づけています。
メモリ効率の工夫
一見すると、head を増やせばその分だけ語彙次元の logits も増えるため、GPU メモリを大きく食いそうです。実際、素直に全部の head の logits を同時に保持すると、ピークメモリはかなり膨らみます。
この論文では、shared trunk の forward を1回流したあと、各 head の forward と backward を順番に処理し、勾配だけ trunk 側へ積み上げる実装を使っています。これにより、巨大な語彙 logits を head ごとに解放しながら進められるので、ピークメモリを naive 実装の O(nV + d) から O(V + d) に落とせます。ここで V は語彙サイズ、d は隠れ表現の次元です。
つまり、「複数先を当てると学習が重すぎるのでは」という自然な不安に対して、著者らは実装レベルでもきちんと答えを用意しています。
推論方法
最も単純な使い方では、推論時には通常どおり next-token head だけを使います。この場合、Multi-Token Prediction は学習を強くするための補助目標として働きます。
一方で、追加された future head はそのまま自己推測デコーディングのドラフト候補としても使えます。たとえば4-token prediction で学習したモデルなら、1回の文脈処理から複数個先の候補を同時に出せます。それを main head が確認しながらまとめて受理できれば、逐次1トークンずつ生成するより速くなります。
ここで重要なのは、外部の小型ドラフトモデルを別に用意しなくてよいことです。学習時に獲得した future head をそのまま使えるため、サービング構成を複雑にしにくいのが実務上の利点です。
なぜ効くのか
論文の主張を実務的に言い換えると、Multi-Token Prediction はモデルに対して「局所的な次単語当て」より少し高い視点の内部表現を要求します。その結果、コードのように手順や構造が大事なタスクで強くなりやすく、ついでに future head を使った高速生成まで可能になります。
特にコード生成では、次の1文字よりも、次の数トークンが作る構文や処理ブロックのまとまりが重要です。この論文の改善幅がコード系ベンチマークで大きいのは、その性質と整合的です。
実験と結果
論文では、コードモデル、自然言語モデル、アルゴリズム的タスク、推論速度の4方向で検証しています。単に精度が上がるかだけでなく、どの規模で効くか、推論に転用できるかまで見ているのが特徴です。
何を検証したのか
主な検証項目は、Multi-Token Prediction が next-token prediction より下流性能を改善するか、最適な予測本数 n はどのあたりか、学習を複数エポック回しても効果が続くか、自然言語でも有効か、そして自己推測デコーディングでどれだけ速くなるかです。
どんなデータセットや評価指標を使ったのか
コード生成では MBPP、HumanEval、APPS、CodeContests を使い、指標は主に pass@k です。自然言語では summarization などの生成系評価に加え、多肢選択や negative log-likelihood も見ています。推論速度は、1秒あたりの生成トークン数で比較しています。
なお、コード生成ベンチマークの一部では、各モデルと各指標に対して最適温度を選ぶ、いわゆる oracle temperature で比較しています。したがって、結果は純粋な相対比較として読むのが適切です。
大きいモデルほど効果が出やすい
著者らは、300M から 13B までのコードモデルを学習し、モデルサイズごとの効果を比較しています。その結果、小さいモデルでは必ずしも有利ではありませんが、規模が大きくなるほど Multi-Token Prediction の利得がはっきり出ました。
論文の要約でも強調されている通り、13B モデルでは比較対象の next-token モデルより HumanEval で 12% 多く、MBPP で 17% 多くの問題を解けました。これは、単に loss が少し下がったという話ではなく、実際のコード生成能力に目に見える差が出ているということです。
4-token prediction がかなり強い
7B クラスのコードモデルを 200B tokens、32k 語彙で学習した比較では、n=4 がかなり良いバランスでした。MBPP pass@1 は baseline の 30.0 から 33.8 に上がり、HumanEval pass@1 も 22.8 から 24.0 に改善しています。HumanEval pass@100 では 62.0 から 66.1 まで伸びています。
逆に n を大きくしすぎると常に良いわけではありません。たとえば同条件で n=8 は MBPP pass@1 が 30.7、HumanEval pass@1 が 20.0 で、n=4 より弱くなっています。つまり、未来を見せれば見せるほど良いのではなく、予測距離には最適帯があるということです。
byte-level 学習では長いパターン学習に特に効く
論文は極端な設定として byte-level tokenization でも検証しています。ここでは 7B モデルに 8-byte prediction を使うと、MBPP pass@1 が 19.3 から 32.3 に大きく伸び、HumanEval pass@1 も 18.1 から 21.8 に上がりました。
著者らはこの結果を、Multi-Token Prediction がより長いパターンやアルゴリズム的構造を学びやすくしている証拠として解釈しています。byte-level のように局所予測が特に難しい条件で伸びるのは、この主張と相性が良いです。
学習後のファインチューニングでも利点が残る
CodeContests での比較では、4-token prediction で事前学習したモデルをファインチューニングすると、通常の next-token 事前学習モデルより pass@k 全体で良い結果が出ました。さらに、事前学習時は multi-token、ファインチューニング時は通常の next-token に戻す設定も強く、補助目標として学んだ表現が下流タスクで活きていることが示されています。
これは実務上かなり重要です。つまり、製品用の最終学習や蒸留の段階では標準的な next-token パイプラインを使い続けても、その前段の pretraining で Multi-Token Prediction の恩恵を受けられる可能性があります。
推論速度も改善した
4-token prediction で学習したモデルは、自己推測デコーディングを使うことで text 生成で最大 2.7 倍、論文全体の要約では最大 3 倍の高速化が報告されています。さらに 8-byte prediction の設定では 6.4 倍の速度向上も示されています。
ここで効いているのは、future head の予測精度が単なる後付けのドラフトより高いことです。事前学習の段階から「先のトークンを当てる」こと自体を学んでいるため、自己推測デコーディングの成功率が上がりやすいというわけです。
自然言語でも大きくは崩れない
論文では、Multi-Token Prediction の利点がコードだけに閉じないことも確認しています。生成系では summarization などに改善があり、多肢選択や NLL 系の標準評価では大きな劣化を起こしていません。
つまり、この手法はコード専用の特殊トリックというより、自己回帰学習の目的関数そのものを少し良くする方向の提案だと見られます。ただし改善幅はコードで特に大きいため、現時点では構造的な生成タスクでより価値が出やすいと考えるのが自然です。
何に使える？
この技術は、基盤モデルの研究だけでなく、実際の AI アプリ開発や推論基盤設計にもいくつかの使い道があります。
コード生成モデルの強化
最も直接的なのは、コード補完やコード生成モデルの事前学習です。関数やブロックの先を見据えた内部表現を学びやすいため、補完の一貫性や問題解決性能の改善が期待できます。社内向けコーディング支援やテスト生成でも活かしやすいです。
追加ドラフトモデルなしの高速推論
推論インフラの観点では、別のドラフトモデルを追加せずに speculative decoding 系の高速化を狙えるのが大きいです。モデル1本で学習と高速化をつなげられるので、サービング構成、監視、モデル更新の複雑さを抑えやすくなります。
小型モデルの品質底上げ
論文では小さいモデルで常に勝つわけではありませんが、一定規模以上のモデルでは明確な改善が出ています。中規模モデルを社内専用で育てたいケースでは、モデルサイズを大きくしすぎずに品質を伸ばす補助目標として検討価値があります。
エージェントの計画やツール呼び出し支援
エージェントは、1手先ではなく数手先の行動列が整っているかが重要です。Multi-Token Prediction をそのままエージェントへ適用するというより、「少し先まで整合する表現を学ばせる」発想が、アクション列生成や structured output の品質改善に効く可能性があります。この点は論文の直接実験ではないため推測を含みますが、技術的な方向性としては自然です。
低遅延が重要な SaaS や社内ツール
チャット、要約、作業支援 UI では、応答速度の体感が継続利用に大きく効きます。もし基盤モデルを自前で学習または継続事前学習できるなら、Multi-Token Prediction は品質と速度を同時に詰める選択肢になります。
開発や事業へのヒント
この論文から得られるヒントは、モデル改善を「アーキテクチャを大きく変える」か「GPU を増やす」かの二択で考えないことです。学習目標の設計だけでも、かなり実用的な差が出る可能性があります。
まず見るべきは loss の定義
自分で AI アプリやモデルを作るなら、モデル本体より先に学習目標を疑う価値があります。次トークンだけを当てる目標が本当にタスクに合っているかを考えるだけで、改善の余地が見つかることがあります。特にコード、計画生成、長い structured text では、その傾向が強そうです。
学習と推論を別物として切り離しすぎない
この論文の良さは、学習時の追加 head が推論時の高速化部品にもなることです。実務でも、訓練専用の工夫がそのまま運用上の価値につながる設計は強いです。たとえば補助タスクや蒸留の副産物を、検証器やドラフト器として再利用できないかを見る発想は応用しやすいです。
プロダクト差別化は基盤モデルの運用単純化にもある
精度が少し良いだけではなく、モデル数を増やさずに速くできること自体が事業上の価値になります。運用モデル数が減れば、監視、AB テスト、バージョン管理、障害切り分けも楽になります。Multi-Token Prediction は、その方向の設計思想として参考になります。
小規模チームでも試せる
超大規模事前学習を再現するのは難しくても、既存モデルへの継続事前学習やドメイン特化学習で「未来位置ごとの補助 head」を足してみる実験は可能です。すぐ商用投入できるとは限りませんが、コード生成や定型文生成のように先読み構造がある領域では、試す価値があります。
今後注目すべき方向性
今後注目したいのは、Multi-Token Prediction そのものだけでなく、「次トークン以外の学習信号」をどう設計するかです。補助目標、自己検証、将来状態予測、階層的な token planning などは、LLM の学習効率と推論効率の両面を変える可能性があります。
限界
まず、この手法はシンプルに見えても、既存学習コードへそのまま差し込めば終わりではありません。複数 head の管理、メモリ効率化、デコーディング実装まで含めると、実装の難しさはそれなりにあります。
次に、効果はモデル規模やタスクによって差があります。論文でも小さいモデルでは一貫して優位ではなく、n を増やしすぎると性能が落ちる条件がありました。したがって、常に未来トークン数を増やせば得という単純な話ではありません。
また、推論高速化は自己推測デコーディングの成功率に依存します。future head の予測が外れやすいタスクでは、理論上の高速化がそのまま出ない可能性があります。モデル品質、温度設定、生成長、バッチ条件によって実効速度は変わります。
さらに、論文は主にコードと自然言語生成を対象にしており、RAG、ツール使用、マルチモーダル生成、複雑な対話管理などで同じ程度の利得が出るかは未検証です。応用先を広げる際には追加実験が必要です。
最後に、比較結果の一部は oracle temperature に基づいています。したがって、本番運用で固定温度や安全制約を入れた場合に同じ差がそのまま出るとは限りません。実務では、評価パイプラインを自分の生成条件に寄せて再検証する必要があります。
よくある質問
Q. Multi-Token Prediction は speculative decoding と同じ技術ですか？
A. 同じではありません。Multi-Token Prediction は学習手法で、複数の未来トークンを同時に予測するようモデルを訓練します。speculative decoding は推論高速化手法です。ただしこの論文では、前者で学習した future head を後者に活用できるため、両者がきれいにつながっています。
Q. どのくらい先まで予測させるのがよいですか？
A. 論文では 32k 語彙のコード学習で n=4 がかなり良いバランスでした。一方で n=8 が常に勝つわけではありません。最適値はタスク、語彙、モデルサイズで変わるので、少数の候補を比較するのが現実的です。
Q. 既存の事前学習済みモデルにも後付けできますか？
A. 可能性はありますが、論文の主効果は事前学習段階で future head を持たせる点にあります。既存モデルへ後付けで head を載せるだけでは、同じ利得は出ないかもしれません。継続事前学習や蒸留と組み合わせて検証するのが現実的です。
Q. まず実務で試すならどの領域が向いていますか？
A. コード生成、構造化テキスト生成、定型フォーマット出力のように、数トークン先の整合性が重要な領域です。逆に、非常に短い応答や分類中心のタスクでは利点が目立ちにくい可能性があります。
Q. 小規模チームがこの論文からすぐ得られる学びは何ですか？
A. モデル改善はパラメータ数の拡大だけではない、という点です。補助 head や将来予測のような学習信号を足すことで、品質改善と推論設計を同時に見直せる可能性があります。
今日の学び
この論文は、LLM の標準である next-token prediction だけでは、先の展開を見据えた表現学習や推論高速化に限界があるという課題を扱いました。
それに対して、共有 trunk の上に複数の future head を置き、次の複数トークンを同時に予測する Multi-Token Prediction で解こうとしました。これにより、コード生成などの構造的タスクで性能が上がり、さらに自己推測デコーディングによる高速化にもつながりました。
そこから得られるヒントは、LLM の改善はモデル本体の巨大化だけではなく、何を予測させるかという学習目標の設計にも大きな余地があるということです。特に、少し先の整合性が重要なアプリでは、この発想を応用できそうです。