パリは[M] [M]都市だ。[M]はまだ決めきれずに伏せておいた空欄だ。この2マスをどう埋めるかによって、文は「有名な観光都市」にも「美しい歴史都市」にもなる。
拡散型言語モデルを速く動かす一つの方法は、完成した文どうしを比べることではない。空欄をもう一マスずつ埋めた中間候補を複数つくり、それをまとめて検査することだ。パリは有名な[M]都市だのように、あと一マスだけ開いた候補を並べて「モデルは本当にこう埋めるのか」を一緒に問う。
こうしてまとめて検査すれば、検査の「回数」は減る。代わりに候補が増えるほど、検査1回が重くなる。デコーディングのコストはおおまかに「検査回数 × 検査1回あたりのコスト」に分かれるが、回数を減らした分だけ1回あたりのコストが膨らむトレードオフというわけだ。
SpecFoldが手を入れるのは回数ではなく、検査1回の重みだ。出発点は単純な事実にある。一緒に検査される中間候補どうしは一マスしか違わないので、内部計算の大半が重なる。SpecFoldはその重なる計算を折りたたんで(fold)再利用し、実際に変わったマスだけをあらためて計算する。言うのは簡単だが、どこまでが「変わったマス」なのかを見分けることがこの手法のほぼすべてだ。そこから解きほぐしていこう。
まずは用語を三つだけ手元に置こう。トークンはモデルが文章を処理する小さな単位、ブロックはそのトークンを複数まとめた区間、隠れ状態は各マスの意味を担うモデル内部の表現だ。
最近よくある言語モデルは、左から右へトークンを一つずつ吐き出す。拡散型言語モデルは違う。文章をブロックに分け、一つのブロックをまるごと伏せた([M])状態で始める。複数のマスを同時に予測し、そのうち確信の高いいくつかだけを公開する([M]を外す)。残った伏せマスをまた埋める作業を、ブロックが完成するまで繰り返す。
一度に複数のマスを埋めるこの並列性が、速度の余地を開く。問題は品質だ。使える文を得るには同じブロックに対してモデルを何度も回さなければならず、この反復が推論コストの大半を食う。だから「反復をどう減らすか」が長年の宿題だった。
投機的デコーディング(speculative decoding)は、将来出てきそうな候補を前もってつくり、まとめて検証しておく手法だ。検証を通った候補が後で実際に使われれば、前もってやっておいた計算を再利用して、これからやるはずの検査を飛ばす。
もう一度パリは[M] [M]都市だに戻ろう。これを親候補と呼ぼう。一マスを追加で公開したパリは有名な[M]都市だは子候補だ。二つは新しく開いた「有名な」の一マスだけが違う。拡散型モデルのアテンションは前後をすべて参照するので(双方向)、候補一つはトークン一つではなく、ブロック全体の状態をなす。
だとすれば、子が親と重なる部分は親の計算を受け継げばよさそうに思える。まさにここが落とし穴だ。その話は少し後にして、まずはなぜ1回あたりのコストが問題なのかから見よう。
複数の候補ブランチを一度に検証するやり方(マルチブランチ)は、検査回数をうまく減らす。ところが一度検査するたびに、現在のブランチとすべての候補ブランチを一つも飛ばさずに全部計算する。ブランチが増えるほど1回あたりのコストがそのまま膨らむ。論文はこのやり方の代表格であるSpiffy(補正を使う投機的デコーディングのベースライン)を自前で実装して測り直したが、ドラフト予算が大きくなるほど検査1回あたりの平均遅延が上がり続けた(Fig. 3、Nemotron-Labs-Diffusion-3B、H100)。SpecFoldが狙うのはまさにこの1回あたりのコストだ。
観察から出発する。子は親とほぼ同じで、新しく公開した少数のマスでだけ違う。著者らがFast-dLLM-v2-7BでGSM8Kのサンプルを解きながら、候補・層・位置・ステップごとに親と子の隠れ状態の差(同じ層の入力で測った子と親の相対差 ‖H子−H親‖/‖H親‖)を観察したところ(Fig. 2)、親と子の隠れ状態は大半が似ており、大きな差は新しく公開された位置の近くにだけ集まっていた。
では重なるものをそのまま受け継げば終わりだろうか。いや、違う。一マスの変化はそのマスだけにとどまらない。アテンションが前後を見るので、新しく開いた一マスの変化がアテンションを伝って隣のマスへ広がる。論文はこの広がりを「残差の波紋(residual ripple)」と呼ぶ(4.4節)。だから「親の計算を再利用してよいか」はブランチごと、マスごと、層ごとに別々に判断しなければならない。従来のやり方はこの判断をまったくせず、すべてを計算し直していた。SpecFoldはこの判断を自分の仕事にする。
段落の校正にたとえてみよう。一段落だけ直したからといって本全体を組み直すのは無駄だ。変わった段落だけ組み直し、残りはそのままにする。ただし直した段落が長くなると後ろの配置がずれるように、修正の影響は周囲へ広がる。どこまでずれたかを見きわめ、ちょうどその分だけ直す――それがSpecFoldだ。
具体的には、一つのマスは差の測定 → 親の結果を再利用 → 自分の差を保存 → 必要なら再計算という道をたどって動く。
‖H子−H親‖/‖H親‖)だけを測る。この測定は安い。重い層の計算を回して子の出力を求めてから比べるのではなく、層の入力で二つのブランチの隠れ状態の差を見るだけだ。比較に使う値と省略する計算がここで分かれる。パリは有名な[M]都市だの「有名な」が生んだ小さな差が、そうして保存される。折りたたみで節約するには、選ぶコスト自体が安くなければならない。そこで、τを超えて計算し直すマスをすべてのブランチからかき集めて一つの束に圧縮したうえで、重い層をその束に対して一度だけ回す。ばらけたまま何度も呼ばない。また、ブランチどうしが共有する前半の計算キャッシュは、ブランチ数だけコピーせず、一つを参照だけ切り替えて使う。コピーそのものが高価なデータ移動だからだ。省いて節約した時間が、こうして選び移すコストを上回ったときに初めて実際に速くなる。
一つははっきりさせておこう。差の保存と自動解除は、候補の多様性と出力品質を守ろうとする仕組みだ。仕組みがあることと、最終的な品質が実際に保たれたことは別の話だ。実際の精度は後述の数値で別に見る。
実装メモ(実装者向け) 各層はGPUカーネルを直接書く言語Tritonで実装した。ゲートカーネルがブランチ間の隠れ状態の差を測り、折りたたんだマスの出所を記録すると、すべてのブランチの再計算位置を一つの束(packed batch)に圧縮する。出所解釈カーネルは折りたたみチェーンをたどり、各折りたたみマスを実際の親状態につなぐ。アテンションはFlashAttention式のストリーミングカーネルで処理し、折りたたんだマスのアテンション出力はそのブランチ自身のキー・バリューで合わせ直す(補正)。共有キャッシュはコピーせず参照だけ切り替える方式(zero-copy strided view)で展開し、キー・バリューを複数のクエリヘッドが共有するアテンション変種(GQA)もサポートする。
どんな再利用なのか トランスフォーマー推論で冗長性を減らす道はたいてい二つだった。一つは時間軸の再利用――続くステップが前のステップの計算を受け継ぐ(KVキャッシュが代表)。もう一つは一度の実行の中で不要なトークン・層を飛ばすことだ。SpecFoldが狙うのは毛色が違う。同じ検証ステップで並んで評価される候補ブランチ「どうし」の、横方向の冗長性だ。対象が違うので、ドラフトのつくり方には手をつけず、Spiffyのような手法の上にそのまま乗る。競合ではなく補完だ。
表を読む前に二つだけ。速度はTPS(毎秒生成トークン数、高いほど速い)で測る。比較相手は二つだ。投機的ドラフトをまったく使わない基本のブロックデコーディングVanilla、そして先に挙げたSpiffy。タスクは性格別にコード生成(HumanEval、MBPP)、数学推論(GSM8K、MATH-500)、指示追従(IFEval)に分かれる。
実験設計の核心は一つだ。SpiffyとSpecFoldに同じドラフトグラフを食わせる。ただし同じドラフトグラフを使っても、実行経路や検査回数(NFE)が完全に同じになるわけではない。実際、表1のFast-dLLM-v2-7B・HumanEvalでNFEはSpiffy 117、SpecFold 108と違い、精度も分かれる。それでも二つが飛ばす検査回数は近い水準にそろうので、速度差のかなりの部分を1回あたりのコスト削減で説明できる。1回あたりのコストを下げることがSpecFoldの核心的な貢献だが、速度差の全体をそれだけで完全に切り分けられたと断定することはできない。
まずは最大の事例。もっとも小さいモデルでもっとも大きい利得が出た。
| モデル・タスク | Vanilla | Spiffy | SpecFold |
|---|---|---|---|
| Fast-dLLM-v2-1.5B · MBPP | 121 TPS | 148 TPS | 242 TPS |
SpecFoldはVanillaの1.99倍、Spiffyの1.64倍だ。ヘッドラインの数字はここから出る。ただしこれは全体の最大値であって、期待値ではない。
全体の傾向は別に読む必要がある。同じ1.5Bモデルの五つのベンチマークでは、Vanilla比+54~99%と幅が広かった。モデルが大きくなると利得は急速に縮み、NLD(Nemotron-Labs-Diffusion系列)-14Bの二つのタスクではVanilla比+10~24%にとどまり、Spiffyとは98対102、110対115 TPSと事実上互角だった。モデルが小さいほど、そしてより遅いVanillaと比べるほど倍率が大きくなるという意味だ。
精度は性格が違うので別に見る。各行がどのベースラインと比較した値なのかも併記した。
| モデル・タスク | ベースライン → SpecFold | 方向 |
|---|---|---|
| NLD-14B · HumanEval | 71.34 → 74.39 | 上昇 |
| NLD-14B · MATH | 84.00 → 85.60 | 上昇 |
| Fast-dLLM-v2-7B · HumanEval | Vanilla 65.85 · Spiffy 64.63 → 61.59 | 低下 |
| NLD-3B · HumanEval | Spiffy 75.61 → 70.12 | 低下 |
平均で見ればベースラインと近い。だがタスクごとに両方向へ揺れる。NLD-14Bのようにむしろ上がるタスクがある一方、コード生成(HumanEval)では7B・3Bともに目立って下がった。平均ではなく、自分のタスクがどちらに当たるかが重要だ。
評価条件。 二系列・五モデル(Fast-dLLM-v2 1.5B・7B、NLD 3B・8B・14B)。すべてゼロショット、特記なければH100 80GB 1枚基準。コードタスクはEvalPlus、残りはLM-Evalで採点。Spiffyの補正サンプルはMBPP・MATHの訓練分割から取り、評価データと分離した。
結論から。省いて節約した時間が、選び(選別)・合わせ(補正)・移す(データ移動)コストを上回ったときにだけ実際に速くなる。
観測された事実は二つだ。一つ、モデルが大きくなるほど(NLD-14B)向上幅が数パーセントに収束した。二つ、コード生成の一部タスクで精度が下がった。
理由は仮説にとどめる。重要なのは、どちらも測定で見分けられるという点だ。
実行中の再計算割合と選別・補正時間を別々に測れば、遅い理由が「節約する側が小さいから」なのか「コスト側が大きいから」なのかが、その測定で分かれる。
導入を検討するなら条件がつく。利得がモデル規模に敏感で、コード生成では実際に精度が下がったタスクがあるので、タスクごとのベンチマークは選択肢ではない。再計算のしきい値τが速度と精度のトレードオフを握っている点も覚えておくべきだ。τを低く取れば折りたたみが少なく安全だが遅く、高くすればその逆になる。運用タスクごとにこの曲線を描いて安全地点を探す必要がある。最後に、この高速化はカスタムカーネルに結びついているので、評価された二系列(Fast-dLLM-v2、NLD)の外のモデルや別のアテンション変種なら、まずカーネルの互換性から確認しなければならない。
ONTLEのスタックは拡散型モデルではなく、SpecFoldのカーネルを持ってきて使うこともできない。この論文は適用するコードではなく、投げかける問い一つとして読むのが正しい。私たちが候補を複数つくるとき、その候補どうしのあいだで実際に重なる作業はどれくらいあるのか。
一つの場面で見てみよう(いま実際にこう動いているという意味ではなく、構造を示すための仮の例だ)。ONTLEのPMエージェントが週次計画を立てる。先週の進捗・滞っているタスク・チームリソースという共通の文脈を敷き、その上で攻撃的な案・安全な案・リソース最小の案の三つを出して比べ、一つを選ぶ。三案は同じ文脈を同じように参照する。しかし素朴に組むと、案ごとにその文脈を最初からモデルに通し直す。参照は共有するのに計算は共有できないのだ。候補が増えるほど、トークンと遅延がその分膨らむ。
SpecFoldがくれるヒントは順序だ。共有する計算を一度だけ済ませ、残りの案がその結果を受け継ぐようにする。そして案ごとに変わる小さな部分が結果にちゃんと反映されたかを確認する。これを抜かすと三案が互いに似てきて、そもそも複数の案をつくった理由が消える。残差の波紋をブランチごとに持ち運んでいた、あの悩みとまさに同じだ。
効果があったかどうかは、同じ条件で四つを測れば分かれる。案一つあたりの遅延、案一つあたりのコスト、候補の多様性、そして最終的に選んだ案の品質。共有計算を折りたたんで前の二つが減りつつ後の二つが保たれるなら、そのファンアウト(一つの入力から複数の枝へ伸びる構造)は折りたたむ価値がある。とくに四つ目、最終品質はこの論文でもタスクによって揺れた指標だ。速度だけ見て品質の測定を飛ばしてはいけない、という教訓がここでもそのまま効く。
著者らは直列実行(serialized-execution)で点検し、実装そのものは正しいと述べている(付録C)。揺れの一部は計算順序から生じる微細な数値差、すなわち浮動小数点ノイズで説明する。ただしあくまで「一部」なので、コード生成タスクの低下を全部ノイズに帰すことはできない。タスクごとの実測が必要な理由だ。
論文はサンプル一つを解きながらモデルを通過させる回数(NFE)も併せて測る。同じドラフトグラフを食わせてもNFEが完全に同じにはならないが(表1の7B・HumanEvalはSpiffy 117対SpecFold 108)近い水準にそろっており、それでもTPSが開くなら、その差のかなりの部分は検査1回あたりのコストから来たと読める。1回あたりのコストを切り出して見ようとした設計とかみ合う。
ドラフト予算とドラフティング戦略を変えながら、再計算割合とスループットがどう動くかは付録Dでさらに追える。折りたたむ余地が減れば再計算が増えて利得が削られるという本文の仮説を見きわめる手がかりになる。