Block Diffusion: AR と DLLM を結ぶ連続体

LLaDA (Nie ほか 2025年) の semi-autoregressive サンプリングは「全並列に生成するか、それともブロック単位で AR 的に進めるか」というブロック幅のパラメータを推論時に導入した。しかしそこではブロック構造は 訓練時には意識されておらず、推論時の便宜として導入されたサンプリング戦略にすぎなかった。BD3-LMsBlock Discrete Denoising Diffusion Language Models)(Arriola ほか 2025年) はこのブロック構造を 訓練段階で第一級の市民に格上げし、ブロック幅 \(K\) という単一のハイパーパラメータで「純 AR (\(K=1\))」から「フル DLLM (\(K=L\))」までを連続的に内挿する。

本章ではこの block diffusion の定式化を整理し、なぜ訓練時にブロック構造を入れることが効くのかを論じる。あわせて、別軸での内挿として近年提案された GIDD(Generalized Interpolating Discrete Diffusion)(Rütte ほか 2025年) にも触れ、DLLM 周辺で進む「内挿型の研究」の地形を地図化する。

なぜこの章を読むべきか

LLaDA を読んだ読者には、semi-AR サンプリングが「KV-cache を活かしつつ大域的整合性を保つ」工夫であったことが既に見えているはずである。しかしそこには次の不整合が残る。

  • 訓練と推論の乖離: LLaDA は系列全体に対する masked diffusion で訓練されている。推論時にブロック単位で逐次化するのは、訓練分布とは異なる条件で生成していることになる
  • ブロック幅の選択の根拠が薄い: 推論時のみブロック幅を決めるため、「どのブロック幅が最適か」は経験則に頼らざるを得ない

BD3-LMs は、これらに対して「訓練時にブロック構造を埋め込む」という素直な答えを与える。さらに、ブロック幅をハイパーパラメータとして明示することで、AR と DLLM の間の連続体が一本のスペクトラムとして見えるようになる。本章の目標は、このスペクトラムを設計空間として把握できるようにすることである。

ノートBD3-LMs の位置付け

BD3-LMs は MDLM (Sahoo ほか 2024年) の within-block formulation と AR の across-block factorization を組み合わせた ハイブリッドな尤度モデルである。新規の損失設計(分散低減)と学習可能なノイズスケジュールにより、当時の DLLM の中で最良の尤度を達成した。LLaDA の semi-AR を「サンプリング工夫」から「訓練レシピ」へ昇格させた、と理解すると見通しが良い。

Block diffusion の定式化

系列の分割

長さ \(L\) の系列 \(x_0 = (x_0^1, \dots, x_0^L)\) を、幅 \(K\) のブロック \(M = L/K\) 個に分割する。

\[ x_0 = (B_1, B_2, \dots, B_M), \qquad B_m = (x_0^{(m-1)K+1}, \dots, x_0^{mK}) \]

ブロック幅 \(K\) は連続パラメータと見ることができる。

  • \(K = 1\) → 各ブロックが 1 トークン。across-block の AR factorization が次トークン予測と同型になり、純粋な AR LLM に退化する
  • \(K = L\) → 単一の巨大ブロック。within-block の masked diffusion が系列全体に適用され、フル DLLM(MDLM / LLaDA 相当) に退化する
  • \(1 < K < L\) → AR と DLLM のハイブリッド

Joint likelihood

BD3-LMs は系列尤度を「ブロック間 AR」×「ブロック内 diffusion」として因数分解する。

\[ \log p_\theta(x_0) = \sum_{m=1}^{M} \log p_\theta(B_m \mid B_{<m}) \]

\(\log p_\theta(B_m \mid B_{<m})\) には、MDLM 的な diffusion ELBO による変分下界が当てられる。

\[ \log p_\theta(B_m \mid B_{<m}) \geq \mathcal{L}_{\text{diff}}(B_m \mid B_{<m}; \theta) \]

ここで \(\mathcal{L}_{\text{diff}}\) は MDLM と同じ「重み付き masked cross-entropy」の形を取る。総合損失は

\[ \mathcal{L}_{\text{BD3}}(\theta) = - \sum_{m=1}^{M} \mathcal{L}_{\text{diff}}(B_m \mid B_{<m}; \theta) \]

となる。across-block の確率は AR 的に厳密に書き下されるため、最終的な尤度は 真の対数尤度の上界(つまり下界の負)として扱える。\(K=1\) では \(\mathcal{L}_{\text{diff}}\) が単一トークンの予測に縮退し、これが通常の next-token CE と一致する。\(K=L\) では単一ブロックの diffusion ELBO となり、MDLM の損失そのものになる。

Attention マスク

訓練時の attention マスクは、ブロック構造を反映してハイブリッドになる。

  • Within-block: ブロック内では双方向 attention。MDLM 同様、[MASK] 位置が他の位置(マスク・非マスクを問わず)を参照できる
  • Across-block: ブロック間は因果(causal)マスク。\(B_m\)\(B_{<m}\) のみを参照でき、\(B_{>m}\) は見えない

この attention 構造は、ブロックを 1 つの「拡張トークン」と見れば AR のそれと同型になる。違いは、各「拡張トークン」が次の 1 トークンではなく \(K\) トークンの集合であり、その集合内部を masked diffusion で生成する、という点にある。

flowchart LR
    subgraph B1["B_1 (双方向)"]
        t1[t_1] <--> t2[t_2]
        t2 <--> t3[t_3]
        t1 <--> t3
    end
    subgraph B2["B_2 (双方向)"]
        t4[t_4] <--> t5[t_5]
        t5 <--> t6[t_6]
        t4 <--> t6
    end
    subgraph B3["B_3 (双方向)"]
        t7[t_7] <--> t8[t_8]
        t8 <--> t9[t_9]
        t7 <--> t9
    end
    B1 -.causal.-> B2
    B2 -.causal.-> B3
    B1 -.causal.-> B3
図 1: BD3-LMs の attention 構造。ブロック内は双方向、ブロック間は因果。図はブロック数 M=3、ブロック幅 K=3 の例

勾配分散と学習可能スケジュール

Diffusion ELBO の高分散問題

BD3-LMs の主要な技術的貢献は、目的関数自体の設計ではなく、勾配分散の制御である。MDLM 的な ELBO は、時刻 \(t \in [0,1]\) のサンプリングと、その時刻でのマスクパターンのサンプリングという二重のランダム性を持つ。AR の next-token CE と比べて、\(\mathcal{L}_{\text{diff}}\) の確率勾配は分散が高くなりやすい。

ブロック構造を入れると、各ブロックが短くなるため within-block の分布が薄くなる。具体的には、

  • 1 ブロックの中で「全マスク」(\(t=1\))になる確率が無視できなくなる
  • 全マスクのブロックは training signal をほとんど提供しないが、ELBO の重み \(1/t\) で増幅される

この組み合わせが勾配分散を膨らませ、フル DLLM では問題にならなかった水準の不安定さが顕在化する。

分散低減推定量

BD3-LMs は ELBO の Monte Carlo 推定量を再設計し、上記の分散を抑える。具体的には、

  • マスク数の clipping: ブロック内の全マスクをアンサンプル領域から外し、少なくとも 1 つの観測トークンが残る条件付き分布を扱う
  • 時刻分布のリパラメトライズ: \(t\) の分布をブロック幅に応じて調整し、有意義な training signal が得られる範囲に集中させる

これは数学的には ELBO の同じ族の中の異なる推定量を取ることに相当し、訓練の収束先は変わらないが分散が下がる。

データ駆動なノイズスケジュール

さらに BD3-LMs は、勾配分散を最小化するようにノイズスケジュールそのものを学習する。直観的には、

  • ブロック内のどの \(t\) で training signal が最も有用かはデータに依存する
  • それを事前に決め打ちするより、データから推定する方が分散を下げられる

論文ではこの「データ駆動スケジュール」の学習が、固定スケジュール(線形・cosine)に対して尤度面で明確に勝つことを示している。

ヒント「分散低減が効く」ことの意味

BD3-LMs の貢献は「新しいモデル族の提案」というより、既存モデル族を実用ラインまで持ち上げる工学である。block diffusion 自体は LLaDA の semi-AR から自然に想像できる構造だが、それを訓練して効かせるためには勾配分散の制御が必要だった、というのが論文の核心である。離散拡散の他のレシピでも、分散低減は共通する課題として現れる。

推論時の利点

KV-cache の自然な活用

完成済みブロック \(B_{<m}\) は固定であり、その K/V は再計算する必要がない。AR LLM の KV-cache と同じ仕組みで、

  • \(B_m\) を生成する間、\(B_{<m}\) の K/V はキャッシュから読む
  • \(B_m\) が完成したら、その K/V をキャッシュに追加
  • \(B_{m+1}\) の生成に進む

このパターンは、フル DLLM では難しい(毎ステップ全位置の forward が必要なため)が、block diffusion では across-block の causal 構造のおかげで自然に成立する。KV-cache のヒット率は \(K\) が大きいほど低くなる(フル DLLM に近づくため)が、\(K \ll L\) の領域では AR LLM に近い恩恵が得られる。

可変長生成

フル DLLM は固定長 \(L\) を前提とする実装が多く、可変長生成はやや扱いにくい。一方 block diffusion では、

  • 生成は \(B_1, B_2, \dots\) と進むため、任意のブロック境界で停止できる
  • EOS トークンが生成された時点でブロック単位で打ち切れる
  • 残りの位置を [MASK] のまま放置する必要がない

これは AR LLM の自然な動作と同型であり、長さの事前知識が不要になる。

推論コストのつまみ

ブロック幅 \(K\) は推論コストを連続的に動かすつまみとして機能する。

表 1: ブロック幅 \(K\) と推論時トレードオフ
\(K\) の選び方 並列性 KV-cache 効率 生成品質の典型
\(K = 1\) なし(AR と同じ) 最大 AR LLM 並み
\(K\) 小(数〜数十) バランス型
\(K\) 中(数十〜数百) DLLM 的並列性が活きる
\(K = L\) 最大(フル並列) 最小 MDLM / LLaDA 並み

表 1 は「正解の \(K\)」を示すものではなく、\(K\) を選択肢として持てること自体が block diffusion の設計上の利点であることを示している。

LLaDA の semi-AR との対比

LLaDA の semi-autoregressive サンプリングと BD3-LMs はどちらもブロック構造を扱うが、立ち位置が異なる。

表 2: LLaDA semi-AR と BD3-LMs の対比
観点 LLaDA semi-AR BD3-LMs
ブロック構造の存在 推論時のみ 訓練・推論の両方
訓練分布 系列全体の masked diffusion ブロック内 diffusion + ブロック間 AR
ブロック幅の根拠 経験則 訓練時に決めて整合させる
KV-cache の活用 部分的(推論時の便宜) 構造的(causal mask が保証)
可変長生成 やや扱いにくい 自然
ブロック幅 1 の極限 機能はするが訓練と乖離 純 AR と等価

要するに、LLaDA semi-AR は block diffusion の「訓練しないバージョン」とみなせる。BD3-LMs はそこに訓練時の整合を入れることで、ブロック幅というハイパーパラメータが「推論側だけの工夫」から「設計空間の正規の軸」へと昇格している。

逆向きに見ると、LLaDA で semi-AR サンプリングが動いてしまったこと自体が、masked diffusion 訓練が部分的にブロック構造を許容することの傍証でもある。BD3-LMs はその許容範囲を最大限に引き出すレシピを提供している。

GIDD: 別軸の内挿

block diffusion がブロック幅という軸での内挿を提供するのに対し、GIDD(Generalized Interpolating Discrete Diffusion)(Rütte ほか 2025年)ノイズ過程の軸で内挿する。

内挿される 2 つの極

GIDD は forward 過程として「マスクのみ」と「一様分布のみ」の両極を混合する。

  • Mask-only(MDLM / LLaDA): トークンは確率的に [MASK] に置換される。一度 [MASK] になったら戻らない
  • Uniform-only(D3PM の uniform transition、第 4 章 参照): トークンは確率的に語彙上の他のトークンに置換される

GIDD はこの 2 種類の遷移を混合パラメータで補間する。極端を取れば mask-only に、別の極端を取れば uniform-only に退化し、中間では両方の遷移が同時に起きる。

自己訂正能力

mask-only 訓練の弱点として、モデルが「誤ったトークン」を見たことがないという問題がある。訓練時に観測されるのは「正しいトークン」と「[MASK]」のみで、推論時に低品質のトークンを後段で訂正する経験が訓練に含まれない。

uniform 成分を混ぜると、モデルは「ノイズで置換された誤ったトークン」を入力として受け取り、それを訂正する経験を訓練時に積める。GIDD はこの self-correction 能力が、mask-only モデルにはない強みとして現れることを示している。

具体的には、GIDD で訓練したモデルが推論時にトークンの再評価・置換を行う能力を獲得することが報告されている。

Block diffusion との関係

GIDD と BD3-LMs は 直交する 2 軸の内挿である。

  • BD3-LMs: 構造の軸(ブロック幅)。AR factorization と全並列 diffusion の連続体
  • GIDD: ノイズ過程の軸(マスク vs 一様)。absorbing と uniform の連続体

両者は原理的には組み合わせ可能であり、ブロック内のノイズ過程を GIDD で、ブロック間を AR で、というハイブリッドも考えうる。ただし論文段階では両者は独立に提案されており、実装上の組み合わせはまだ open である。

AR-DLLM の連続体を描き直す

LLaDA の章では「AR と DLLM が異なる factorization である」と整理した。BD3-LMs はそれを 連続スペクトラムとして描き直すことを可能にする。

flowchart LR
    AR["K=1<br/><b>純 AR</b><br/>GPT / Llama"]
    Small["K 小 (数〜数十)<br/>BD3-LMs (小 K)<br/>AR 寄り"]
    Mid["K 中 (数十〜数百)<br/>BD3-LMs (中 K)<br/>バランス"]
    Large["K 大 (数百〜L)<br/>BD3-LMs (大 K)<br/>DLLM 寄り"]
    DLM["K=L<br/><b>フル DLLM</b><br/>MDLM / LLaDA"]

    AR --> Small
    Small --> Mid
    Mid --> Large
    Large --> DLM
図 2: ブロック幅 K を軸とした AR-DLLM の連続体。両端点と中間領域それぞれに代表的なモデルが対応する

このスペクトラムを認めると、次のような問いが自然に立つ。

  • どの \(K\) が最も良い尤度・生成品質を与えるか。タスク依存か、スケール依存か
  • 同じ計算予算(FLOPs、メモリ)の下では、\(K\) のどの領域が Pareto 前線になるか
  • 推論時に動的に \(K\) を変える戦略は意味があるか(短い応答は小 \(K\)、長文は大 \(K\) など)

BD3-LMs の論文では、中間域に「sweet spot」が存在し、\(K=1\)(純 AR)よりも尤度が良くなる場合があることが報告されている。これは「AR が常に最強ではない」という事実の最も具体的な実証である。

中間域で尤度が改善する現象は、次のように解釈できる。AR は条件付き分布の連鎖を「左の文脈のみ」に縛るが、ブロック内の双方向 attention は ブロック内の同時依存 を捉えられる。一方でブロックが大きすぎる(\(K=L\))と、訓練信号が薄くなる / 推論ステップが増える / KV-cache が効かないなどの欠点が表面化する。両者のバランスが中間 \(K\) で取れる、という解釈である。

ただし「中間が最強」が普遍的なのか、データセット・スケール・タスクに依存するのかは未だ open であり、ここは次の節で扱う未解決問題のひとつである。

未解決の問い

block diffusion は新しい設計空間を開いたが、その全貌は明らかではない。本書の Open Problems の章 と重なる部分も多いが、block diffusion に固有の問いとして次が挙げられる。

  • 最適なブロック幅: タスク・スケール・データに対して、\(K\) の最適値はどう決まるか。scaling law 的な経験則は構築可能か
  • 動的ブロック幅: 推論時にブロック境界を動的に決める戦略はあり得るか。文の区切りや意味的単位に合わせて適応的に区切ることは有効か
  • 誘導・編集との組み合わせ: classifier-free guidance、infilling、in-place editing 等の推論時介入は、block diffusion とどう整合するか。across-block の AR 構造は guidance を素直に乗せやすそうだが、within-block での介入は別途設計が要る
  • 理論的理解: AR と DLLM の連続体において、表現力・サンプル複雑度・推論コストはどう関数化されるか。\(K\) をパラメータとしたときの汎化バウンドや収束性は未整理である
  • GIDD との組み合わせ: ノイズ過程の軸(GIDD)とブロック幅の軸(BD3-LMs)の同時最適化は、片方ずつより良い設計を生むか

これらは AR LLM 側にすでに存在する道具立てを、連続体としての DLLM に翻案する作業として整理できる。

関連章へのリンク

  • semi-AR サンプリングの源流: LLaDA
  • ブロック内 diffusion の土台: MDLM
  • GIDD の uniform transition の元: D3PM and SEDD
  • 分野全体の未解決問題: Open Problems

block diffusion は「AR vs DLLM」の二項対立を一本の連続体に組み直し、その上で訓練・推論の整合を取るための具体的なレシピを提供する。LLaDA を読んだ後にこの章を読むと、semi-AR を「便利な工夫」から「設計空間の正規の点」へと位置づけ直せるはずである。

参考文献

Arriola, Marianne, Aaron Gokaslan, Justin T. Chiu, ほか. 2025年. 「Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models」. International Conference on Learning Representations. https://arxiv.org/abs/2503.09573.
Nie, Shen, Fengqi Zhu, Zebin You, ほか. 2025年. 「Large Language Diffusion Models」. arXiv preprint arXiv:2502.09992. https://arxiv.org/abs/2502.09992.
Rütte, Dimitri von, Janis Fluri, Yuhui Ding, Antonio Orvieto, Bernhard Schölkopf, と Thomas Hofmann. 2025年. 「Generalized Interpolating Discrete Diffusion」. arXiv preprint arXiv:2503.04482. https://arxiv.org/abs/2503.04482.
Sahoo, Subham Sekhar, Marianne Arriola, Yair Schiff, ほか. 2024年. 「Simple and Effective Masked Diffusion Language Models」. Advances in Neural Information Processing Systems. https://openreview.net/forum?id=L4uaAR4ArM.