💬 NLP

Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models

本論文は、非隣接なグループにシーケンス位置を分割して並列的なアンマスキングを行い、結合エントロピーの増加を最小化することで、品質を損なうことなく、また基礎となるデノイザーを変更することなく、マスクド拡散言語モデルにおけるテキスト生成を最大5.8倍高速化する推論のみの手法である Dilated Unmasking Scheduler (DUS) を紹介する。

原著者： Omer Luxembourg, Haim Permuter, Eliya Nachmani

公開日 2026-05-14

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

CC BY 4.0

原著者： Omer Luxembourg, Haim Permuter, Eliya Nachmani

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

巨大なジグソーパズルを解こうとしていると想像してください。ただし、完成図が見えるのではなく、すべてのピースが黒いシールで覆われた箱から始めるのです。あなたの目標は、シールを剥がして完成図を明らかにすることです。

AI テキスト生成の世界において、これが**マスクド拡散言語モデル（MDLMs）**の仕組みです。これらは、すべての単語が隠蔽（マスク）された文から始め、答えを再構築するためにそれらを一つずつ「マスク解除」しようとします。

問題点：「自信に満ちた」アプローチは遅い

従来、これらの AI モデルは非常に慎重で自信に満ちた人物のように振る舞います。彼らはパズルを見て、最も可能性が高いピースを推測し、その場所のシールだけを剥がします。その後、再び見て、次に最も自信のある場所を推測し、別のシールを剥がします。

これは正確ですが、信じられないほど遅いです。100 語を明らかにする必要がある場合、あなたは「推測マシン」（AI モデル）に 100 回別々にアクセスする必要があります。まるで、小さな筆を塗料缶に浸して 1 平方インチだけ塗り、筆を再び浸し、これを繰り返して壁を塗ろうとしているようなものです。

一部の研究者は、「では、最も自信のある場所をトップ 10 選んで、それらを一度にすべてマスク解除しよう！」と言ってみて速度向上を図りました。しかし、これはしばしば逆効果になります。なぜなら、それらの 10 の場所は自信に基づいて選ばれたため、通常は互いに隣り合っているからです。隣接する場所を同時にマスク解除することは、それらがどのように接続しているかを知ることなく、10 個の隣接する正方形を塗ろうとするようなものです。AI は混乱し、間違いを犯し、後戻りするか、無意味なものを生成することになります。

解決策：「拡散」戦略

この論文の著者らは、**拡散マスク解除スケジューラ（DUS）**と呼ばれる新しいゲームの遊び方を提案しています。

DUS を、最良の場所を推測する人物ではなく、固定された計画を持つ賢い建設現場の監督だと考えてください。AI に「どの単語が正しいと思いますか？」と尋ねる代わりに、監督は「1 番、5 番、9 番、13 番の位置の単語をマスク解除します」と言います。

ここでの比喩は以下の通りです：
長い暗い廊下に電球を埋め込むと想像してください。

古い方法（トークンごと）： 1 つの電球を点灯させ、部屋が調整されるのを待ち、次の電球を点灯させ、待ち、というのを繰り返します。これには永遠にかかります。
「自信に満ちた」並列方法： 廊下を見て、最初の 5 つの電球は推測しやすいと判断し、それらを一度にすべて点灯させます。しかし、それらがすべて固まっているため、光は均一ではなく、さらに奥の暗い場所を見逃してしまいます。
DUS 方法： 拡散スケジュールを使用します。
1. ラウンド 1： 1 番、5 番、9 番、13 番、17 番の位置の電球を点灯させます（大きな隙間を残して）。
2. ラウンド 2： それらの間の隙間を埋めます：3 番、7 番、11 番、15 番...
3. ラウンド 3： 残った小さな隙間を埋めます。

なぜこれが機能するのか

DUS の魔法は間隔にあります。初期のラウンドで AI に互いに遠く離れた単語を明らかにさせることで、「固まる」問題を回避します。

独立性： 互いに遠く離れた単語は、互いにあまり依存しません。文の 5 番目と 6 番目の単語を一緒に推測するよりも、文の最初の単語と最後の単語を独立して推測する方が簡単です。
文脈の構築： 広く間隔を空けた単語が明らかになると、それらはアンカーとして機能します。AI が 2 番目のラウンドで隙間を埋めに戻るとき、それは文のより豊かな「地図」を持って作業することになり、推測がはるかに正確になります。

結果：高速かつ高精度

この論文は、数学の問題解決（GSM8K）、コード作成（HumanEval）、一般教養クイズへの回答などの困難なタスクでこの方法をテストしました。

速度： DUS を使用すると、AI は単語ごとに 1 ラウンドではなく、わずか数ラウンド（対数的時間）でテキストのブロック全体をマスク解除できます。これにより、従来の遅い方法と比較して最大5.8 倍の高速化が実現しました。
品質： 驚くべきことに、一度にマスク解除する単語の数を減らして間隔を空けることで、AI は「自信に満ちた」並列方法よりも少ない間違いしか犯しませんでした。速くなっただけでなく、同時に賢くなったのです。
再トレーニング不要： これは「プラグアンドプレイ」のアップグレードです。AI モデルを再トレーニングしたり、その脳を変更したりする必要はありません。ゲーム中に単語を明らかにするルールブックを変更するだけです。

まとめ

この論文は、単純なスケジューリングのトリックを紹介しています：最も簡単な単語を最初に推測するのではなく、最も広く分散した単語を最初に推測してください。

壁を埋める前に遠く離れた柱を据える建設プロジェクトのようにテキスト生成を扱うことで、AI は推論能力、数学問題の解決能力、コード作成能力を失うことなく、はるかに高速にテキストを生成できます。これは、新しいハードウェアやモデルのトレーニングを必要とせずに、遅いステップバイステップのプロセスを高速な並列プロセスに変換します。

技術的サマリー：速度の計画：マスクド拡散言語モデルのための拡張スケジューリング

1. 問題定義

マスクド拡散言語モデル（MDLM）は、並列かつ非自己回帰的なテキスト生成を可能にすることで、従来の自己回帰（AR）モデルに対する有望な代替手段を提供します。しかし、推論効率において重大なボトルネックが存在します。MDLM は理論的には任意の順序での生成をサポートするものの、既存のサンプリング戦略はしばしばこの潜在能力を実現できていません。

現在の手法は、モデルの信頼度やエントロピーに基づいてアンマスクするトークンを選択する自己プランナーに依存するのが一般的です。本論文は、これらの信頼度ベースのプランナーが、無意識のうちに自己回帰的な振る舞いを模倣していると主張します。最も信頼度の高いトークン（しばしば空間的に隣接している）を貪欲に選択することで、並列アンマスクステップ間の相互作用を無視してしまいます。その結果、「近視眼的」なプロセスが生じ、モデルは実質的に遅いトークンごとの生成に戻り、長さ $G$ の系列に対して $O(G)$ のノイズ除去器呼び出しを必要とします。

半自己回帰的なブロック拡散のような代替戦略は、系列を連続したスパンに分割しますが、サイズ $B$ のブロックあたり $O(B)$ のノイズ除去器呼び出しを依然として課し、並列性を十分に活用できていません。さらに、最近のアンマスクを推論時の計画問題として捉える研究は、外部のガイドモデルや複雑な学習目的を必要とすることが多く、実用的なドロップインソリューションとしての限界があります。

2. 手法：拡張アンマスクスケジューラー（DUS）

著者らは、基礎となるノイズ除去器を変更したり再学習を必要としたりすることなく、各ノイズ除去ステップにおける結合エントロピーの増加を最小化するように設計された、純粋に推論時のモデル非依存プランナーである**拡張アンマスクスケジューラー（DUS）**を提案します。

中核概念

DUS は、サイズ $B$ の系列ブロックを、非隣接の拡張グループに分割し、それらを並列にアンマスクします。スケジューラーは、扱いやすい周辺和の目的関数と、扱いにくい結合エントロピーの目的関数の間のエントロピーギャップ（ $\Delta$ ）を最小化することが、高品質な並列サンプリングの鍵であるという原理に基づいて動作します。

$\Delta(I_t; S_t) = \sum_{i \in I_t} H(X_i | S_t) - H(X_{I_t} | S_t)$

ここで、 $I_t$ はステップ $t$ でアンマスクされるインデックスの集合、 $S_t$ は現在の状態です。選択されたトークンが条件付き独立である場合、このギャップは最小化されます。

アルゴリズム設計

DUS は、決定論的かつ粗い粒度から細かい粒度へのスケジューリングを採用します：

拡張パターン: モデルのスコアに基づいてトークンを選択するのではなく、DUS は固定された拡張パターンを使用してアンマスクするインデックスを事前に定義します。
反復的洗練: サイズ $B$ $B$ のブロックについて、スケジューラーは $R = \lceil \log_a B \rceil$ $R = ⌈ lo g_{a} B ⌉$ 回の反復（ここで $a > 1$ $a > 1$ は基底、通常は 2）でブロックを完了します。
- 初期反復: 広く間隔を空けたトークンの疎な集合をアンマスクします（大きなストライド $s_t$ ）。これは、高速混合する可変長マルコフ連鎖（VLMC）の仮定の下で、同時にサンプリングされるトークン間の条件付き依存性を低減します。
- 後続の反復: 文脈 $S_t$ が露出した隣接トークンによって豊かになるにつれて、より細かい解像度で残りのギャップを埋めます。
計算量: このアプローチは、ブロックあたりのノイズ除去器呼び出し数を $O(B)$ から $O(\log B)$ に削減し、全体の計算量を $O(\frac{G}{B} \log B)$ に達成します。

理論的保証

高速混合する VLMC の仮定の下で、本論文はインデックスを距離 $d$ だけ間隔を空けることが、それらの間の相互情報を幾何学的に減衰させることを証明しています。その結果、広く間隔を空けたトークンの結合条件付きエントロピーは、個々の条件付きエントロピーの和に近づき、並列サンプリングのための周辺エントロピーの代理変数の使用が正当化されます。

3. 主要な貢献

推論のみ、モデル非依存のデコーディング: DUS は、モデルのアーキテクチャ、重み、学習手順を一切変更せずに適用可能な「ドロップイン」戦略です。
対数関数的アンマスクスケジューリング: 局所的文脈を尊重し、 $O(\log B)$ 回の反復で結合エントロピーを最小化する、決定論的かつ粗い粒度から細かい粒度への拡張戦略を導入します。
理論的分析: 著者らは、高速混合の仮定の下で各反復において DUS が結合エントロピーの境界に近づくことを示す理論的 bound を提供しており、トークンのクラスタ化された選択により大きなエントロピーギャップを招く可能性のあるベースラインと比較しています。
相補的なハイブリッド拡張: 本論文は、拡張間隔の原理を適応的サンプリャ（エントロピー境界付きおよび信頼度境界付き）へのポストフィルタとして適用できることを示しており、計算コストを同等またはわずかに増加させる範囲で精度を向上させています。

4. 実験結果

著者らは、5 つの MDLM 変種（LLaDA-8B、Dream-7B、DiffuCoder-7B）を 4 つのベンチマークカテゴリで評価しました：

数学: GSM8K, MATH500
コード: HumanEval, MBPP
一般知識: BBH, MMLU-Pro
指示追従: IFEval

主要な発見

速度と精度のトレードオフ: DUS は、従来の信頼度ベースのプランナーを一貫して上回ります。同等の速度向上において、DUS は数学ベンチマークで最大27%、コードベンチマークで**20%**のダウンストリームタスク精度の向上をもたらします。
ウォールクロック速度向上: DUS は、ノイズ除去器を変更することなく、トークンごとの MDLM デコーディングに対して最大5.8 倍のウォールクロック速度向上を達成します。例えば、GSM8K における LLaDA-8B では、DUS（ブロックサイズ $B=32$ ）は 5.8 倍の速度向上で 49.36% の精度を達成しますが、トークンごとのベースラインは 72.63% の精度を達成するものの、著しく遅いです。
AR との比較: 特定の構成（例えば、GSM8K における DUS を使用した LLaDA-8B）では、この手法は競争力のある精度を維持しながら、自己回帰的な Llama-3-8B ベースラインよりも5.6 倍高速です。
ロバスト性: この手法は、コード生成や一般知識推論など、多様なドメインで良好に機能し、厳密な高速混合の仮定に違反する可能性のあるタスクにおいても拡張スケジューリングが有効であることを実証しています。
ハイブリッドの改善: 適応的サンプリャ（EB/CB）へのポストフィルタとして適用された場合、拡張間隔は密にクラスタリングされたトークンの選択を防ぐことで精度を向上させます（例：EB における HumanEval で +13.4%）。

5. 意義と主張

本論文は、DUS がマスクド拡散モデルに内在する「品質と速度のトレードオフ」を成功裡に解決すると主張しています。性能を低下させる一般的な加速手法とは異なり、DUS は既存の並列サンプリング戦略と比較して、推論を加速すると同時に出力品質を向上させます。

著者らは、DUS が以下を必要とすることなく、拡散デコーディングの内在的な並列性を解き放つことを強調しています：

モデルの再学習。
補助的なプランナーモデル（BERT 誘導プランナーなど）の追加。
ノイズ除去器の学習目的の変更。

品質と速度のトレードオフをブロックサイズ $B$ の決定論的かつ予測可能な関数へと変換することで、DUS は効率的かつ高品質なマスクド拡散言語モデルを実用化する現実的な道筋を提供します。この手法は、ステップごとのコストを削減することで直交的に動作する KV キャッシングなどの他の効率化技術と相補的な技術として提示されています。