💬 NLP

K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling

K-Forcingは、教師モデルを複数の未来のトークンを単一のフォワードパスで共同デコード可能な条件付きマッピングへと蒸留することで、自己回帰的なテキスト生成を加速させる新しいプッシュフォワード言語モデリングパラダイムであり、わずかな品質低下のみで、高負荷なバッチサービング下における大幅な推論速度向上を実現する。

原著者： Zhiwei Tang, Yuanyu He, Yizheng Han, Wangbo Zhao, Jiasheng Tang, Fan Wang, Bohan Zhuang

公開日 2026-06-10

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Zhiwei Tang, Yuanyu He, Yizheng Han, Wangbo Zhao, Jiasheng Tang, Fan Wang, Bohan Zhuang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは物語を書こうとしていると想像してください。しかし、あなたの横には非常に厳格で、かつ天才的な編集者がいます。その編集者は、あなたが一度に一つの単語しか書くことを許しません。

単語を一つ書くたびに、あなたは手を止めて、紙を編集者に渡し、編集者が全体を読み終えるのを待ち、次の単語を考え、そしてまた紙を返してもらう必要があります。それから次の単語を書き、手を止め、また繰り返すのです。

これが現在のAI（自己回帰型、または「AR」モデルと呼ばれます）の仕組みです。これらは非常に賢いのですが、「ストップ・アンド・ゴー」のリズムに縛られているため、非常に低速です。長い物語を書こうとすれば、編集者のデスクへ何千回も往復しなければなりません。

問題点：「一度に一単語ずつ」というボトルネック

この論文は、この「一度に一単語ずつ」というアプローチが、まるでティースプーンでスイミングプールを満たそうとしているようなものだと主張しています。たとえスプーンが速かったとしても、プロセスは「何回の往復ができるか」によって制限されてしまいます。コンピュータの用語で言えば、AIは次の計算（スプーン）を行う前に、メモリ（プール）の準備ができるのを待っている状態なのです。これは、多くの人が同時にテキストを生成しようとしているとき、非常に非効率的になります。

旧来の解決策：なぜそれらは完全には機能しなかったのか

科学者たちは、二つの主要なアイデアでこれを修正しようとしましたが、どちらにも欠陥がありました。

「下書きとチェック」方式（投機的デコーディング / Speculative Decoding）： 学生が次の数単語を推測しようとし、その後、先生がそれらをチェックする場面を想像してください。もし先生が同意すれば成功です。しかし、もし同意できなければ、学生は最初からやり直しになります。
- 欠陥： 学生が5単語を推測することもあり、時には1単語しか推測できないこともあります。これがスケジュールを乱します。大勢の人がこれを行っている場合、全員のタイミングがズレてしまい、システムは再び低速化します。
「拡散（ディフュージョン）」方式（Diffusion Method）： 白紙のキャンバスから始めて、部分的にパーツを少しずつ明らかにしていく絵画を描く場面を想像してください。ただし、一度に複数のパーツを推測しようとします。
- 欠陥： この論文は、複数のパーツを独立して推測すること（例えば、空と草を別々に推測すること）は、空と草が一致しない、バラバラな絵を作ってしまう可能性があると主張しています。完璧な絵を作るためには、結局のところ、一度に非常に小さなパーツを一つずつ明らかにする必要があり、スピードアップという目的が台無しになってしまいます。

新しい解決策：K-Forcing（「魔法の設計図」）

著者らはK-Forcingを導入しました。編集者に一単語を求める代わりに、AIに魔法の設計図を見せ、一度の往復で複数の単語を一度に（例えば4単語）書かせる方法です。

その仕組みを、簡単な比喩を使って説明します。

1. 「前向きのマップ」（設計図）
ランダムな数字（サイコロの目のようなもの）を即座に特定の文章へと変換するマシンがあると想像してください。

従来の方法： サイコロを振り、「3」が出ると、マシンは「The」と言います。次にまた振って「5」が出ると、「cat」と言います。
K-Forcingの方法： 一度に4つのサイコロを振ります。マシンは設計図を見て、「よし、これら4つの数字は『The cat sat down』というフレーズに対応している」と判断します。そして、これら4つの単語を瞬時に出力します。

2. どうやって設計図を手に入れるのか？（Progressive Self-Forcing）
設計図をただ推測するだけではいけません。それは完璧である必要があります。そこで、彼らは「教師と生徒」のゲームを利用します。

ステップ1： 低速だが完璧な「教師」AIを用意します。教師にランダムな数字を与え、一単語を書かせます。そして、「ランダムな数字 0.45」＝「The」というペアを記録します。
ステップ2： 「生徒」AIがこの繋がりを学習するように訓練します。
ステップ3（魔法のトリック）： 生徒が1単語を書くことに習熟したら、その生徒を使って、2単語を書く方法を教えさせます。次に、それを使って4単語を書く方法を教えさせます。
なぜこれが重要か： 複雑なルール全体を一度に学ぼうとするのではなく、補助輪付きの自転車から練習するように、段階的に積み上げていくのです。

3. 結果：「バッチ処理」という超能力
K-Forcingは、実行されるたびに常に決まった数の単語（例：常にちょうど4単語）を生成するため、コンピュータは混乱しません。100人のリクエストがあっても、全員が全く同じタイミングで自分の4単語を受け取ることができます。

スピード： この論文は、多くのリクエストを処理する際、これが2.4倍から3.5倍高速になることを示しています。
品質： テキストの品質は、低速な「教師」よりもわずかに劣ります（おそらく5%程度の低下）が、それでも非常に高品質であり、スピードの向上は絶大です。

まとめ

K-Forcingを、一度に一つの荷物を届ける配達員から、一度の停車でパレット全体の荷物をドロップオフする配送トラックへのアップグレードだと考えてください。これは「何を」届けるか（言葉）を変えるのではなく、「どのように」届けるかを変えるものであり、忙しい時間帯においてシステム全体をより高速かつ効率的にします。

この論文は、AIが未来について「考える」方法（単語のブロックを予測するようにする）を変えることで、新しいハードウェアを必要とせずに、劇的なスピードアップを実現できることを証明しています。

技術要約: K-Forcing

問題提起

自己回帰（AR）言語モデルはテキスト生成の支配的なパラダイムであるが、逐次的なトークンごとのデコーディングという根本的な効率性のボトルネックを抱えている。このプロセスはメモリ帯域に制約されており、 $L$ 個のトークン列を生成するために $L$ 回のフォワードパスを必要とするため、ハードウェアアクセラレータの利用率が低下し、特に高負荷のバッチサービングシナリオにおいて高いレイテンシを引き起こす。

既存の加速手法は、産業規模のデプロイメントにおいてこのボトルネックを効果的に解決することに苦慮している：

投機的デコーディング（ドラフト・アンド・ベリファイ）: 単一リクエストのレイテンシ削減には効果的であるが、これらの手法は可変長の受理済み出力を生成する。これが規則的なバッチ処理を妨げ、同期の問題（例：不揃いなテンソル/ragged tensors）や、高負荷時にスループットの向上を相殺してしまうオーバーヘッドを生じさせる。
拡散言語モデル（DLM）: 複数のトークンを反復的なデノイジングによってデコードする。しかし、標準的な定式化（例：Masked Diffusion Language Models）は、位置ごとの周辺予測器を学習する。理論的分析によれば、結合条件付き分布ではなく、これらの独立した周辺分布からサンプリングすることは、分布的な誤差を導入する。品質を維持するためには、多くの場合、多くの精緻化ステップを必要とし、ARと比較してフォワード評価回数（NFE）を大幅に削減することに失敗する。

手法: K-Forcing

著者らは、同時次- $k$ -トークン・デコーディングのために設計されたプッシュフォワード言語モデリング・パラダイムである K-Forcing を提案する。核心となる目的は、固定長の出力を保持し、規則的なバッチングを維持しながら、独立した一様ノイズ変数を単一のフォワードパスで $k$ 個の未来トークンの結合サンプルへと変換する条件付き写像を学習することである。

1. プッシュフォワード定式化

結合分布 $p(x_{t+1}, \dots, x_{t+k} | x_{\le t})$ を明示的にモデル化する代わりに（これは大きな語彙サイズに対して計算量的に困難である）、K-Forcing は決定論的な写像 $G_\theta: \mathcal{V}^t \times [0, 1]^k \to \mathcal{V}^k$ を学習する。

入力: コンテキスト $x_{\le t}$ および $k$ 個の独立同一分布（i.i.d.）ノイズ変数 $z \sim \text{Uniform}(0, 1)^k$ 。
出力: $k$ 個の未来トークンの結合サンプル $(\hat{x}_{t+1}, \dots, \hat{x}_{t+k})$ 。
理論的根拠: 本論文では、ARオーラクルによる逆累積分布関数（CDF）をアンロールすることで、閉形式のプッシュフォワード写像が存在することを証明している。この写像は、結合確率テーブルを明示的にパラメータ化することなく、理論的にARティーチャーの正確な結合条件付き分布を再現できる。

2. プログレッシブ・セルフフォーシング蒸留（Progressive Self-Forcing Distillation）

この写像を事前学習済みARティーチャーから学習するために、著者らは Progressive Self-Forcing Distillation を導入し、従来のノイズ反転ベースラインの2つの主要な制限（学習と推論のミスマッチ、および数値的な脆弱性）に対処する。

ステージ1（フォワード蒸留）: $k=1$ のPFLMをブートストラップする。ARティーチャーは、ノイズ $z$ をサンプリングし、その逆CDFサンプラーを適用することでターゲットを生成する。生徒モデルは、 $(x_{\le t}, z)$ からトークンへの写像を学習する。これにより、ノイズ分布が推論時（一様分布）と一致することを保証し、狭いCDFビンを反転させることを回避する。
ステージ2（セルフフォーシング蒸留）: より大きな $k$ （例： $2 \to 4$ ）へスケールさせるために、学習済みのPFLM( $k$ )がPFLM( $2k$ )のティーチャーとして機能する。ティーチャーは $z_{1:k}$ を用いて最初の $k$ 個のトークンを生成し、次に生成されたトークンに条件付けて、次の $k$ 個のトークンを $z_{k+1:2k}$ を用いて生成する。生徒モデルは、全ノイズベクトルが与えられた状態で、フル $2k$ シーケンスを一度のパスで予測することを学習する。このプロセスにより、予測ウィンドウを再帰的に倍増させる（ $1 \to 2 \to 4$ ）。

3. アーキテクチャ設計

著者らは、ARバックボーンを再利用する 完全因果的アーキテクチャ（Fully Causal Architecture） を提案している：

各ノイズ変数は、シーケンス内の個別のトークンとして扱われる。
共有の予測ヘッドが未来のトークンをデコードする。
因果的アテンションは、 $j$ 番目の未来トークンがコンテキストと接頭辞ノイズ変数 $(z_1, \dots, z_j)$ のみに依存することを保証し、これは理論的な逆CDF構成を反映している。
この設計は、すべての確率的決定を単一の共有潜在トークンにルーティングするという情報ボトルネック（標準的なマルチトークン予測アーキテクチャに見られるもの）を回避する。

主な貢献

既存の限界の分析: 本論文は、ドラフト・アンド・ベリファイ手法が可変長の出力のためにバッチサービングにおいて失敗すること、および拡散モデルが結合分布ではなく周辺分布からサンプリングするためにNFEを効果的に削減できないことを特定している。
プッシュフォワード定式化: K-Forcingは、一様ノイズを結合トークンブロックへと写す暗黙的な生成モデルとして定式化されており、固定長出力と結合サンプリングの要件を満たしている。
学習戦略: ノイズ反転に伴う学習と推論のミスマッチおよび数値的な脆弱性を排除する、プログレッシブ・セルフフォーシング蒸留の導入。
実証的検証: 統計的なモデリングの変更のみによって、ARデコーディングに対するバッチサービングのスループットの大幅な向上を実現できることを示す。

実験結果

標準的な因果的Transformerバックボーンを用い、LM1B および OpenWebText で評価を行った：

スループット: $k=4$ トークンを1回のフォワードパスで生成するように構成した場合、K-Forcingは異なるバッチサイズ（4, 16, 128）において、ARと比較して 2.4〜3.5倍の高速化 を達成した。
品質: ARティーチャーと比較して、生成品質にはわずかな低下が見られる。
- LM1Bでは、 $k=4$ の場合、LLM-as-a-Judgeによるペアワイズ評価においてARに対し約42.9%の勝率となった。
- OpenWebTextでは、 $k=4$ の場合、約39.4%の勝率となった。
- より低い $k$ 値（例： $k=2$ ）は、よりスムーズなトレードオフを提供する。 $k=2$ はARとほぼ区別がつかず（勝率約50%）、かつ約1.6倍の高速化を提供する。
SOTAとの比較: K-Forcingは、MDLM（複数のトークンをアンマスクする際に品質が急激に低下する）および投機的デコーディングのベースライン（Medusa, PTP）と比較して、品質とNFEのトレードオフにおいて優れている。K-Forcingは、同等のNFEにおいてMDLMよりも低いGen-PPLと高い勝率を達成し、同等の品質を得るために投機的手法よりも少ないフォワードパスを使用する。

意義と主張

本論文は、現実世界の高負荷デプロイメントにおけるAR生成を加速するための有望なルートを提示していると主張している。

既存の手法が不規則なバッチングのために失敗しやすい、高負荷バッチサービングという特定のボトルネックに対処している。
プッシュフォワード写像による結合マルチトークンサンプリングが、効率的なKVキャッシュ管理とコンティニュアスバッチングに必要な固定長構造を保持しつつ、逐次的デコーディングに代わる実行可能な選択肢であることを示している。
著者らは、本研究をプッシュフォワード言語モデルというより広い研究方向への第一歩として位置づけている。彼らは、特に大きな $k$ において、ARティーチャーとの非自明な品質差が依然として存在することを認め、トレーニングコストを削減するためのカスタムブロック疎（block-sparse）アテンションカーネルの開発や、事前学習済みARティーチャーに依存しない学習パラダイムの探索といった将来の研究方向を特定している。