PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

次にくる言葉を予測しようとしている場面を想像してみてください。長い間、最も賢いコンピュータ（大規模言語モデルと呼ばれます）は、本を音読するように、一単語ごとに読み進め、次の単語を考えるために一回ずつ立ち止まってきました。これは正確ですが、時間がかかります。最近、科学者たちは「拡散（ディフュージョン）」モデルを用いた、物語の新しい書き方を発見しました。これは、一歩ずつ細かく削っていくのではなく、大理石の塊から一度に彫刻を彫り上げる彫刻家のようなものです。これにより、コンピュータは多くの単語を同時に推測することができ、驚異的に高速になります。

しかし、落とし穴があります。一度に多くの単語を推測すると、いくつか間違えてしまうことがあります。これを修正するために、「投機的デコーディング（speculative decoding）」と呼ばれる巧妙なトリックがあります。これは、素早いジュニア・アシスタントが次の数単語を推測し、その後、非常に賢い上司がその推測が正しいかどうかをチェックするようなものです。もし上司が同意すれば、その単語の束全体を瞬時に受け入れ、大幅な時間を節約できます。問題は、ジュニア・アシスタント（拡散モデル）は個々の単語を推測することには長けていますが、それらの単語が特定の順序でどのように組み合わさるかを必ずしも理解していないことです。それは、アシスタントがケーキの材料を個別に選ぶのは得意ですが、どの組み合わせが実際に美味しい味になるのかまでは分からない、という状況に似ています。

ここで、新しい論文が登場します。Zheng Wang氏率いる研究者たちは、現在の「速いアシスタント」の使い方では、多くのスピードを無駄にしていることに気づきました。彼らは、拡散モデルが膨大な種類の単語の組み合わせを生成できる一方で、現在の手法ではたった一つの経路しかチェックしていないことを見出しました。それは、まるで一本の廊下を歩き、その突き当たりにドアが開いていることを祈るようなものです。著者らは、PRESTO（Prefix-Aligned Tree Drafting）と呼ばれる新システムを提案しています。一本の廊下を進む代わりに、PRESTOは可能性の「木（ツリー）」を構築し、一度に多くの経路を探索します。しかし、ここが魔法のポイントです。このシステムは、アシスタントの自信の測定方法における根本的な欠陥を修正します。アシスタントの元の自信は「プレフィックス・ブラインド（接頭辞に盲目）」であり、前にある単語を気にしません。PRESTOは「プレフィックス・アラインド（接頭辞に整合した）」スコアを追加し、それがコンパスのように機能することで、探索される経路が、上司によって受け入れられる可能性が最も高いものであることを保証します。

その結果、システムは大幅に高速化されました。テストにおいて、PRESTOはコンピュータが推測の各ラウンドでより多くの単語を受け入れるのを助けました。既存の優れたセットアップと比較して、プロセス全体を1.5倍高速化しました。他のセットアップでは、1.12倍のスピードアップを実現しました。この論文は、推測のプロセスを一本の直線ではなく、木が枝分かれする冒険として扱うことで、拡散のスピードと慎重なチェックの正確さの両立が可能になると示唆しています。

問題点： 「一本の経路」の罠

なぜPRESTOが必要なのかを理解するために、あなたが友人と「マッドリップス（穴埋めゲーム）」をしている場面を想像してください。友人は、欠けている単語を推測しようとしています。この友人は拡散モデルです。彼らは空白を見て、「ここに入る単語は『猫』かな！」「あるいは『犬』かな？」「それとも『ロケット』かな？」と、これらすべての選択肢を同時に叫ぶことができます。

しかし、この友人の使い方は非常に硬直的です。現在の方法は、友人のトップの推測を取り、それを書き込み、次に「上司（ターゲットモデル）」にその推測が正しいかどうかを尋ねます。もし上司が「ノー」と言えば、すべてが破棄され、最初からやり直しになります。もし上司が「イエス」と言えば、次の単語に進み、これを繰り返します。これは**線形ドラフティング（linear drafting）**と呼ばれます。それは、森の中を歩きながら、目の前にある道だけを見ているようなものです。

著者らは、このアプローチが非効率的であると指摘しました。拡散モデルは同時に多くの選択肢を生成するため、そこには膨大な「組合せ空間」が存在します。それは、地図上に千もの異なる小道があるのに、一つの道しか歩くことが許されないようなものです。論文によれば、単一の経路に固執することで、多くの有効なルートを見逃していることが示されています。実際、数学の問題（GSM8K）において、現在の手法では平均して約6.5単語しか受け入れられませんでしたが、研究者たちは、もしすべての最適な経路をチェックできれば、10単語近くを受け入れられたはずだと計算しました。これは大きなギャップです！

不一致： 「盲目な」コンパス

研究者たちはさらに深く掘り下げ、単に（木構造で）より多くの経路をチェックするだけでは、なぜ完璧に機能しないのかという具体的な理由を見つけました。彼らは「根本的な不一致」を特定しました。

標準的なAI（自己回帰モデル）の世界では、単語の信頼度スコアは、その前にある単語に大きく依存します。もし文章が「猫がマットの上に座って……」であれば、モデルは「マット」が非常にありそうな次の単語であることを知っていますが、「ピザ」はそうではありません。これは**プレフィックス・アラインド（接頭辞に整合している）**な状態です。

しかし、拡散モデルは異なります。彼らは各ポジションに対して独立した「周辺確率」を生成します。これは、モデルが「ポジション5において、『猫』である確率は80%だ」と言うようなもので、ポジション4が「The」であっても「The quick brown」であっても気にしません。これは**プレフィックス・ブラインド（接頭辞に盲目）**な状態です。

この「盲目な」スコアを使って推測の木を作ろうとすると、ランキングの問題が発生します。最初の単語にとっては素晴らしく見える経路を選んでも、モデルが文脈の変化を認識していないため、二番目の単語にとってはひどいものになってしまう可能性があります。それは、曲がった後の道が一方通行であることを無視して、現在の通りに基づいた指示だけを出すGPSのようなものです。論文は、これらの盲目的なスコアを使用して木を構築することは、「信頼できない経路ランキング」を招き、システムが間違った枝を探索して時間を浪費することになると主張しています。

解決策： PRESTO

PRESTO（Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding）は、拡散モデルのスコアに「補正」を加えることでこれを解決します。

1. プレフィックス・アラインド・スコアリング： 著者らは、拡散モデルの強力な「周辺的（marginal）」な信号（単独での単語の尤もらしさ）と、「プレフィックス条件付き（prefix-conditioned）」の信号（前の単語に基づいた尤もらしさ）を組み合わせる必要があると気づきました。彼らは、拡散確率にシンプルなn-gramモデル（単語の組み合わせを見る軽量なツール）から導かれた補正係数を乗算する、新しいスコアリング式を作成しました。これにより、物語の流れを尊重するスコアが生まれます。
2. 優先度ベースの木探索： 単にトップの経路を選ぶのではなく、PRESTOは「木」を構築します。新しい補正済みスコアを使用して、どの枝を成長させるかを決定します。上司に受け入れられる可能性が最も高い経路を優先します。それは、単に真っ直ぐ歩くのではなく、たとえ最初のステップでは目立たなくても、頂上に到達する可能性が最も高いルートを地図を見て選ぶハイカーのようなものです。

論文では、木を成長させる2つの方法をテストしました：ビームサーチ（Beam Search）（各ステップで固定数のトップパスを保持する）と、ベストファーストサーチ（Best-First Search）（これまでに発見された単一の最良パスを常に拡張する）です。彼らは、自分たちの設定においてはビームサーチがベストファーストサーチと同等に機能することを発見し、よりシンプルで効率的なオプションを採用しました。

結果： より速く、よりスマートに

著者らは、数学の問題（GSM8K, Math500）、コーディングの課題（HumanEval, LiveCodeBench）、およびチャット会話を含む様々なタスクでPRESTOをテストしました。彼らは2種類のシステムを使用しました：

• 専用の拡散ドラフター（Dedicated Diffusion Drafters）： 大きな自己回帰モデル（dFlashなど）のために、単語を推測する小さな高速な拡散モデル。
• 自己投機的拡散LLM（Self-Speculative Diffusion LLMs）： 推測と検証を自分自身で行う単一の拡散モデル（Nemotron-Labs-Diffusionなど）。

結果は一貫していました。PRESTOは一貫して平均受け入れ長（Average Acceptance Length）、つまり一度に上司が受け入れる単語の数を増加させました。

• dFlash システム（Qwen3-8Bを使用）では、平均受け入れ長が約6.6単語から9.6単語に跳ね上がりました。これは、エンドツーエンドの速度が1.5倍になったことを意味します。
• Nemotron-Labs-Diffusion システムでは、受け入れ長が8.8から9.9単語に増加し、1.12倍のスピードアップを実現しました。

さらに印象的なことに、論文はPRESTOが「確率的（stochastic）」、つまりランダム性が高い状況でも機能することを示しました。このような状況は予測がより困難ですが、速度向上はより顕著になり、一部のベンチマークではスループットがほぼ倍増しました。

著者らはまた、この新しい手法の「コスト」についても確認しました。木を構築し、新しいスコアを計算するために必要な追加作業は極めて小さく、総時間の4%未満であることがわかりました。時間の大部分（90%以上）は、依然として上司による実際の検証に費やされています。これは、PRESTOがオーバーヘッドによってシステムを遅延させることなく、非常に効率的なアップグレードであることを意味します。

PRESTOが「行わない」こと

この論文が主張していないことも明記しておく必要があります。著者らは、単に（プレフィックス・アラインド・スコアリングなしの）素朴な木構造を適用するだけでは、最適ではないと明言しています。拡散モデルの生のスコアを使用して木を作ったとしても、スコアの「盲目的な」性質により、完全な恩恵を受けることはできません。PRESTOは、まさにそのスコアリングの不一致を修正するためのものです。

さらに、この論文は拡散モデルの問題を完全に解決したと主張しているわけではありません。彼らは、自らの手法が、補正を行うための「扱いやすいプレフィックス・アラインド信号」（使用したn-gramモデルのようなもの）に依存していることを認めています。彼らは、将来の研究ではさらに豊かな信号を探索できる可能性があると示唆していますが、現時点では、このシンプルな補正だけで大きな成果を得るには十分であるとしています。

なぜこれが重要なのか

AIをより速く、より効率的にするという競争において、わずかなスピードの差も重要です。投機的デコーディングは、最も正確なモデルを使用しながら、その全時間を支払うペナルティを回避できるため、ホットなトピックとなっています。しかし、現在の手法は、拡散モデルを線形な機械として扱っており、その「一度に多くの選択肢を生成できる」という独自の能力を無視していたために、限界がありました。

PRESTOは、拡散モデルを「マルチパス・エクスプローラー（多経路の探索者）」として扱うことで、ゲームのルールを変えます。スコアリングを上司の検証方法に適合させることで、拡散モデルの並列生成のポテンシャルを最大限に引き出します。その結果、単に少し速くなるだけでなく、大幅に効率的なシステムとなり、テキスト生成、数学問題の解決、コード作成を、これまでのタイプのモデルでは不可能と思われていたスピードで実現します。著者らが述べているように、彼らは「一本の経路」を辿る歩みを「木に基づいた遠征」へと変え、あらゆる一歩が正しい答えへの一歩となるようにしたのです。