【DL輪読会】 Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

DEEP LEARNING JP [DL Papers] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models Hiroto Osaka, Matsuo Iwasawa Lab, M1 http://deeplearning.jp/

• http://deeplearning.jp/

Paper Information Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models ▍ Conference：ICLR 2025 Oral ▍ Authors Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov ▍ Links o Project Page o https://m-arriola.com/bd3lms/ o Paper o https://openreview.net/forum?id=tyEyYT267x o Code o https://github.com/kuleshov-group/bd3lms [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 3

• https://m-arriola.com/bd3lms/
• https://openreview.net/forum?id=tyEyYT267x
• https://unidisc-diffusion.github.io/

Background（Language Models） ▍ Autoregressive Models [2] ▍ Diffusion Models [3] q 高い生成品質 q 並列生成が可能 q 任意長の系列を扱える q 系列長が固定 q 逐次的生成で速度が遅い q KV-cache が使えない [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 4

Proposed Method ▍ BD3-LM q ブロック内での離散拡散モデル q ⾃⼰回帰型（Transformer）と拡散型 （Diffusion）の両⽅の利点を合わせる q 並列化 q ブロックごとにトークンをまとめて生成 q KV-Cache q 過去ブロックの情報をキャッシュして再 利用 [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 5

Mathematics of Block Diffusion ▍ NELBO（Negative ELBO）[4] q 再構成損失 q 低ノイズ状態からの復元 q 拡散損失 q 段階的復元の KL（復元の正確さ） q 事前損失 q 最終ノイズとモデルの事前分布の整合性 6

Masked Diffusion Models ▍ マスク付き拡散モデル q 拡散行列 𝑸𝒕 q 各トークンが保持される確率 𝜶𝒕 q 𝟏 − 𝜶𝒕 の確率で [mask] に変換 q マスクトークン q モデルにノイズからの復元を学習させる q マスクされた単語を予測できるように訓練 q マスク導入により、明確な欠損を表現できるため モデルの復元学習が安定 7

Training Algorithms ▍ 学習アルゴリズム q 通常は各ノイズレベルでモデルを複数回実⾏ q 提案⼿法では2回の Forward Pass のみ q クリーン系列でKVキャッシュを計算・保存 q キャッシュを利⽤しノイズ系列を復元 q 利点：計算量が劇的に削減、訓練が効率化 8

Specialized Attention Mask ▍ Attention Mask q MBD（Block Diagonal） q 同じノイズブロック内でのみ q MOBC（Offset Block Causal） q 全てのクリーン系列と過去のノイズトーク ンを参照 q MBC（Block Causal） q クリーン系列が過去のクリーントークンの みを見る（通常のアテンションマスク） q クリーン系列とノイズ系列の一括処理 [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 9

Sampling Algorithms ▍ サンプリングアルゴリズム q ブロック単位の逐次⽣成 q 各ブロック内は拡散モデルで並列⽣成可能 q ブロック間はAR的に⽣成されるため、任意⻑ ⽣成可能 q KVキャッシュ再利⽤ q ⽣成中も過去ブロックのKVキャッシュを維持 q 利点：並列性とARモデルの品質を両⽴ 10

Experimental Settings ▍ Dataset q LM1B（One Billion Word Benchmark） q 約10億語の⼤規模テキストコーパス q ニュース記事など、幅広いジャンルを網羅した英語テキスト q OpenWebText (OWT) q Web上の英語テキスト（約8百万ドキュメント、合計40GB以上のサイズ） ▍ 評価指標 q Perplexity q 尤度の評価 q PPLが低いほど、モデルの予測精度が⾼い 11

Result1. Perplexity ▍ LM1B（One Billion Word Benchmark） ▍ OWT（OpenWebText） q 全ての離散拡散系モデルを上回る尤度 q 拡散モデルの中で最良の PPL を達成 q L’ = 4 の小さいブロックサイズが高性能 [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 12

Result2. Variable-Length Sequence Generation ▍ 可変⻑⽣成 ▍ サンプル品質 q OWT で学習したモデルで 500 ⽂書をサンプ q 最⼩の拡散ステップ数で最⾼の⽣成品質を達成 q SSD-LM より桁違いに⾼速 リング q ⽣成を [EOS] か低エントロピー区間が出現す るまで継続 [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 13

Ablation Study ▍ Noise Schedules ▍ Efficiency q 線形 / 対数方よりも Clipped が高性能 q 通常は2回のパスが必要だが、1回のパスで OK q FlexAttention [8] を使用 q 学習速度が 20%〜25% 向上 [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 14

Summary ▍ Conclusion q 既存の離散拡散モデルが抱える主要課題に対処するアプローチを提案 q 任意長のシーケンス生成が困難 q Autoregressive Models と比較して品質が劣る q Block Discrete Denoising Diffusion Language Models（BD3-LMs） q ブロック単位での Autoregressive な生成とブロック内での拡散生成の組み合わせ q 分散勾配を低減するデータ駆動型ノイズスケジュールの設計 q 高い尤度性能と柔軟な長さのシーケンス生成能力の両立 15

Limitation & Discussion ▍ Limitation q BD3-LM は通常の拡散モデルよりも訓練コストが高い q ブロック単位で逐次生成するため、ブロックの大きさによっては Autoregressive 同様遅くなる q 大きなブロックでは制御性が下がる q ブロックサイズ選定はタスクに依存。適応的に調整する必要性 ▍ Discussion q 画像などのマルチモーダルな生成への応用可能性はあるか q 強化学習などと組み合わせた制御生成は可能か 16

References 1. 2. 3. 4. 5. 6. 7. 8. Swerdlow, Alexander, et al. "Unified Multimodal Discrete Diffusion." arXiv preprint arXiv:2503.20853 (2025). Radford, Alec, et al. "Improving language understanding by generative pre-training." (2018). Li, Xiang, et al. "Diffusion-lm improves controllable text generation." Advances in neural information processing systems 35 (2022): 4328-4343. Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. pmlr, 2015. Chelba, Ciprian, et al. "One billion word benchmark for measuring progress in statistical language modeling." arXiv preprint arXiv:1312.3005 (2013). Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http: //Skylion007.github.io/OpenWebTextCorpus, 2019. Jelinek, Fred, et al. "Perplexity—a measure of the difficulty of speech recognition tasks." The Journal of the Acoustical Society of America 62.S1 (1977): S63-S63. Dong, Juechu, et al. "Flex Attention: A Programming Model for Generating Optimized Attention Kernels." arXiv preprint arXiv:2412.05496 (2024). 17

Appendix 18

Appendix 1. Likelihood Gaps q 同じ目的関数（L’=1）で学習しても、Diffusion の方が PPL が高い現象 q 拡散モデルはマスクされたトークンに「だけ」損失が かかるため、1バッチ内の学習信号は AR モデルより 少ない q 拡散モデルの性能劣化は学習効率の問題だった ▍ Batch NELBO [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 19

Appendix 2. Low-Variance Noise Scheduling q 極端なマスク率は勾配の分散を増加させる q マスク率が低すぎる → 予測が簡単 q マスク率が高すぎる → 周辺分布を学ぶだけで効果が低い ▍ Clipped Noise Schedules [1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models 20