技術要約：マスク型拡散モデルにおける再帰的スケーリング

問題提起

マスク型拡散モデル（MDM）は、自己回帰（AR）モデルの逐次的な性質とは対照的に、完全な並列生成能力を提供する、離散シーケンス生成のための有望なパラダイムとして台頭している。しかし、高品質な生成を実現するためにMDMをスケールアップさせる際、従来は2つの軸（パラメータ数（幅/深さ）の増加、またはデノイジングステップ数 $T$ の増加）に依存してきた。このアプローチは、訓練（より大きなモデル）と推論（より多くのフォワードパス）の両面において、多大な計算コストを伴う。再帰的計算（同一の変換を複数回適用すること）は、ARモデルや教師あり学習においてパラメータ効率を向上させる手法として有望視されているが、MDMの文脈においてはほとんど探索されていない。本研究が取り組む中心的な課題は、再帰的な「深さ」が、パラメータ数を増やすことなく生成品質を向上させる第3の直交するスケーリング軸として機能するかどうか、そして、追加のデノイジングステップを置き換えて推論時の計算量を削減できるかどうかである。

手法

著者らは、標準的なMDMの訓練パラダイムに再帰的なループ機構を組み込んだ**再帰的マスク型拡散モデル（R-MDM）**を提案する。

アーキテクチャ

R-MDMは、$KL$層のトランスフォーマーを一度だけ適用する代わりに、共有された$K$層のトランスフォーマーブロック$f_\theta$を、単一のデノイジングステップ内で$L$回連続して適用する。

• パラメータ効率: モデルは、実効的な深さが $K \times L$ であるにもかかわらず、単一パスの$K$層モデルと同じパラメータ数を持つ。
• フォワードパス: ノイズを含む入力 $x_t$ が与えられると、隠れ状態は反復的に更新される：$h^{(\ell)} = f_\theta(h^{(\ell-1)}, \ell, L)$。ループ $\ell$ の出力は、ループ $\ell+1$ の入力として機能する。
• ステップ埋め込み: モデルが再帰の異なる段階を区別できるように、ステップ埋め込み $v^{(\ell)}$ が隠れ状態に加算される。この埋め込みは、正規化されたループ進行度 $s_\ell = (\ell-1)/(L-1)$ を通じて小さなMLPから導出され、モデルが再帰の深さに応じて計算を特化させることを可能にする。

訓練目的関数

本論文では、再帰ループに対する4つの教師あり戦略を評価している：

1. 全ステップ（ALL）: 全$L$個のループにわたってクロスエントロピー損失を平均化する。これは、密な勾配信号を提供する（深く教師あり学習に着想を得ている）。
2. 最終ステップ（FINAL）: 最終ループのロジットのみを監督し、標準的なMDMの目的関数を拡張する。
3. 重み付き（Weighted）: 異なるループに対して非一様な重みを割り当てる。
4. 切り捨て（Truncated）: 最終の$k$個のループのみを監督する。

著者らは、All-steps損失が、特に低計算量レジームにおいて最も堅牢な性能を示すことを見出した。

推論とスケジューリング

推論時、モデルは訓練時の再帰深度 $L_t$ とは異なるサンプリング再帰深度 $L_s$ で評価できる。著者らは、固定深度、確率的サンプリング、およびカリキュラムベースのスケジューリングを含む、様々な訓練スケジュールを検討している。逆カリキュラム（高い再帰深度から始めて1へと減少させる）が優れていることが判明しており、これにより、モデルは訓練の初期段階で高容量の反復的な精緻化戦略を学習するように強制される。

主な貢献

1. 再帰的深さをスケーリング軸として導入: パラメータ数を増やすことなく実効的なモデルの深さを増すメカニズムとして、デノイジングステップ内での再帰的ループを提案した。
2. パラメータ効率の証明: R-MDMは、構造化された推論タスクにおいて、大幅に大きな非再帰的ベースライン（最大5倍のパラメータ）と同等またはそれを上回る性能を示すことが示された。
3. 推論時の効率性: 再帰的な精緻化がデノイジングステップを部分的に代替できることを示し、総フォワードパス数（$T \times L$）を減らして高品質な生成を実現できることを実証した。
4. タスク依存性の分析: 構造化タスク（数独、Countdown）および非構造化タスク（Text8）にわたり、再帰、パラメータ数、およびデノイジングステップのトレードオフを体系的に評価した。

実験結果

実験は、Sudoku（9x9および25x25）、Countdown（3、4、5個のオペランドを用いた算術計画）、およびText8（文字レベルの言語モデリング）の3つのデータセットを用いて行われた。

構造化推論タスク（Sudoku & Countdown）

• パラメータの置換（RQ1）: 再帰モデル $(K \otimes L)$ は、一貫してアイソFLOP（同等計算量）の単一パス・ベースライン $(KL \otimes 1)$ と同等またはそれを上回る。
  • 例: Sudoku 9x9において、$(6 \otimes 5)$ モデル（10.6Mパラメータ）は、$T=5$ において93.6%の有効パズル率（VPR）を達成し、同じステップ数で77.8%しか達成できなかった30層の単一パスモデル（53.1Mパラメータ）の性能に匹敵する。
  • 例: Countdown-4において、$(3 \otimes 5)$ モデル（5.5Mパラメータ）は、$T=20$ で70.2%の目標到達率（RTR）に達し、15層のベースライン（26.7Mパラメータ）が達成した69.4%を上回った。
• ステップの削減（RQ2）: 再帰モデルは、目標品質の閾値に達するために、より少ないデノイジングステップを必要とする。
  • ループモデル（$L_s=3, T=5$）は、合計15回のフォワードパスで95%の有効パズル率を達成するが、同等の単一パスモデルは40ステップを必要とする（計算量の2.7倍の削減）。
  • 80%のマスキングを伴うSudoku 25x25において、再帰モデル（$L=5$）は $T=1$ で91.8%のVPRを達成するが、ベースラインはこのステップ数ではほぼ完全に失敗する。
• 柔軟性: 逆カリキュラム（$L_t$ が減少する）で訓練されたモデルは、異なるサンプリング深度（$L_s$）で効果的に評価でき、推論時における柔軟な計算資源の割り当てを可能にする。

非構造化言語モデリング（Text8）

• Text8データセットにおいて、再帰モデルは尤度ベースの指標（NLL、Perplexity）において単一パス・ベースラインを下回った。
• しかし、定性的分析によれば、再帰モデルは尤度スコアが悪化しているにもかかわらず、より一貫性のあるテキストを生成していることが示唆されており、これは標準的な指標がこのレジームにおいて誤解を招く可能性があることを示している。著者らは、再帰の恩恵は明示的な反復構造を持つタスクにおいて最も顕著であることを指摘している。

意義と主張

本論文は、再帰的深さが、特に構造化生成タスクにおいて、MDMのための実用的なスケーリングメカニズムであることを主張している。その意義は主に以下の2点にある：

1. パラメータ効率: R-MDMは、わずかな割合のパラメータでより大きなモデルの性能を達成できるため、リソース制約のある環境での高品質な生成を可能にする。
2. 計算効率: 再帰的な精緻化をデノイジングステップの代わりに用いることで、推論時に必要な総フォワードパス数を減らし、生成品質とレイテンシのトレードオフを改善できる。

著者らは、再帰は構造化ドメインにおける反復的な制約伝播のための強力なメカニズムとして機能する一方で、非構造化言語モデリングにおけるその有用性は、さらなる調査を要する未解決の課題であると結論付けている。本研究は、従来のパラメータ数中心の視点を超えた、生成モデルのスケーリングに関する理解を深めるものである。