技術要約：固定点マスク生成モデリング (CoFRe)

問題提起

MDLM（テキスト用）やMaskGIT（画像用）のようなマスク生成モデル（MGMs）は、並列デコーディングと多様なモダリティにわたる強力な生成品質を可能にします。しかし、これらは重大な効率性のボトルネックに直面しています。標準的なMGMは、各リファインメントステップにおいてフルシーケンスの双方向トランスフォーマー・パスを必要とするため、高い学習コスト（VRAMおよび時間）と、低いサンプリング予算（限定された総フォワードパス数）における生成品質の低下を招きます。

既存の効率化手法には以下の3つのカテゴリがありますが、いずれもこの問題を完全には解決できていません：

1. サンプラーの改善： トークン選択戦略を変更するが、ステップごとのデノイザーのコストは固定されたままとなる。
2. 効率的な固定深度アーキテクチャ： 構造的な無駄（例：マスクされたトークンのスキップ）を削減するが、依然として固定容量のデノイザーに依存している。
3. 適応型/ループ型モデル： 制御可能な深度を提供するが、離散的なマスク生成への効果的な転移に失敗することが多い。これは、連続的な拡散の仮定を乱す、離散的な状態変化（顕在化したトークン vs マスクされたトークン）によるものです。

核心となる課題は、特に標準的なMGMが苦戦する低予算のレジームにおいて、生成品質を犠牲にすることなくデノイザーの計算コストを削減し、適応的な深度を可能にすることです。

手法：CoFRe フレームワーク

著者らは、固定点マスク生成モデル (FP-MGMs) と、CoFRe (Consistency, Fixed-Point, and Reuse) と呼ばれる完全な学習・推論フレームワークを導入しています。このアプローチは、標準的なMGMにおける個別のトランスフォーマー層のミドルスタックを、重みを共有する固定点ブロックに置き換えるものです。

1. 固定点デノイジング・アーキテクチャ

デノイザーは、深い一連のユニークな層の代わりに、以下のように分解されます：

• 前処理スタック ($P_{\theta_P}$): 初期隠れ状態生成のための明示的な層。
• 固定点ブロック ($F_{\theta_F}$): 平衡状態 $h^*_t = F_{\theta_F}(h^*_t; \tilde{h}_t, t)$ を見つけるために反復的に適用される共有ブロック。実効的な深度は、パラメータ数ではなく、ソルバーの反復回数 ($N$) によって制御されます。
• 後処理スタック ($H_{\theta_H}$): 平衡状態をロジットへと写像する明示的な層。

この設計は、重みを共有することでパラメータ数を削減し、推論中に実効的な深度を動的に調整することを可能にします。学習には、全ソルバートラジェトリの活性化を保存することを避けるため、Stochastic Jacobian-Free Backpropagation (SJFB) を使用します。これは、勾配追跡の有無でランダムな回数の反復をアンロールすることで勾配を近似します。

2. 三状態再利用 (Three-State Reuse: 3SR)

標準的なウォームスタート（前のステップの解を次のステップの初期状態として使用すること）は、MGMにおいては非効率的です。なぜなら、トークンの状態は一様に変化しないからです（一部のトークンは可視のまま、一部はマスクされたまま、そして一部は新たに顕在化します）。

• メカニズム: ソルバーの初期化 $h^0_t$ は、前の固定点解 ($h^*_{t+1}$) と現在の前層出力 ($h_{pre,t}$) の加重結合となります。
• トークンごとの変調:
  • 変化のない可視トークン: 前の解を完全に再利用する ($\gamma = 1$)。
  • 依然としてマスクされているトークン: 解を部分的に再利用し、デノイジングが進むにつれて再利用係数を増加させる。
  • 新たに顕在化したトークン: 前の解への依存を最小限にし ($\gamma = 0.2$)、現在の入力注入への依存度を高める。
    この戦略は、新たに顕在化したトークンが最大の表現シフトを経験する一方で、安定したトークンは最小限の再計算で済むという事実を考慮に入れています。

3. ステップ間一貫性正則化 ($L_{CONS}$)

固定点アーキテクチャ単体では、特に低予算の場合、ステップ間での表現のドリフトが生じる可能性があります。

• メカニズム: 一貫性損失は、「よりノイジーな」学生の状態 ($z_{t_s}$) と、同じ基礎シーケンスから導出された「よりクリーンな」教師の状態 ($z_{t_c}$) （入れ子状のマスクを使用）の隠れ表現を一致させます。
• 目的: $L = L_{MGM} + \lambda \|h_s - \text{sg}(h_c)\|^2_2$。ここで、$h_s$ と $h_c$ はロジット直前の最終的な隠れ状態です。
• 効果: これはクロスタイム自己蒸留として機能し、予測をシャープにし、固定点ソルバーを安定させます。これは低予算の生成において極めて重要です。

4. 学習済みモデルの変換

著者らは、フルリトレーニングを行うことなく、学習済みのMDLMチェックポイントをFP-MDLMに変換できることを示しています。CKA分析によって深い層における冗長性が示されていることをガイドとして、特定の層を前処理/後処理ブロックおよび固定点ブロックにマッピングし、KLダイバージェンス損失を用いた短い教師・学生適応ステージを適用することで、最小限の追加学習で強力な性能を実現します。

主な貢献

1. FP-MGMs: 個別のトランスフォーマー層を共有固定点ブロックに置き換え、パラメータ数を削減し、適応的な深度を可能にする新しいアーキテクチャクラス。
2. CoFRe フレームワーク: 離散的なマスク生成特有の課題に対処するために、固定点デノイジングに Three-State Reuse (3SR) と Cross-Step Consistency ($L_{CONS}$) を組み合わせた完全なレシピ。
3. 効率的な適応: 短い蒸留を通じて既存の学習済みMGMを効率的なFP-MGMに変換する手法を提供し、ゼロからの大規模な事前学習コストを回避する。
4. 実証的検証: これらのコンポーネントが、テキストおよび画像の両方のモダリティにおいて、品質とコストのトレードオフを共同で改善することを実証。

実験結果

フレームワークは、言語については OpenWebText (OWT)、画像については ImageNette で評価されました。

言語 (OpenWebText)

• 効率性: MDLMのベースラインと比較して、CoFReはパラメータを 38.8% (170M $\to$ 104M)、学習時間を 11.5%、VRAM使用量を 16.9% 削減しました。
• 品質: 96個のトランスフォーマー・ブロック・フォワードパスという低いサンプリング予算において、CoFReは生成パープレキシティ (Gen. PPL) を MDLM の 830.8 から 101.8 へと改善しました。SDTT (Self-Distillation Through Time) を備えたMDLMと比較しても、同じ予算においてCoFReは有意に低いパープレキシティ (101.8 vs 193.1) を達成しています。
• アブレーション: $L_{CONS}$ が低予算での品質向上の主要な要因であり、3SR はソルバーの初期化を最適化することでさらなる改善を提供することが特定されました。

画像 (ImageNette)

• 効率性: CoFReは、MaskGIT-Largeと比較して、学習時間を 48.6%、VRAMを 50.7% 削減しました。
• 品質: CoFReは、テストされたすべての予算においてFID（Fréchet Inception Distance）を改善しました。例えば、予算48において、FIDはMaskGIT-Largeの 174.1 からCoFReの 96.7 へと向上しました。

意義と主張

本論文は、CoFReが、より安価な学習と、限定されたサンプリング予算下でのより強力な生成を実現するための実用的なフレームワークを提供すると主張しています。

• 品質とコストのトレードオフ: 主な貢献は、マスク生成モデルのパレート境界をシフトさせ、より少ないパラメータと計算リソースで高品質な生成を可能にしたことです。
• モダリティ非依存: アプローチは離散的なテキストトークンと量子化された画像潜在変数（latents）の両方で検証されており、汎用性を示唆しています。
• 実用性: 学習済みモデルを短期間の適応によって変換できる能力により、大規模な事前学習リソースを必要とせずにこの手法を利用できます。
• 限界: 著者らは、現在の結果がOWTスケールおよびImageNetteに限定されていること、また、固定点ソルバーが制御フローのオーバーヘッドを追加するため、ハードウェアの実装によっては必ずしもウォールクロック時間での高速化に直結しない可能性があることを謙虚に述べています。また、本手法は単調なマスクデコーディング向けに設計されており、トークンを再マスクするサンプラーには適応が必要である可能性についても言及しています。

要約すると、本論文は、デノイザーを固定点問題として扱い、状態の再利用とステップ間の一貫性を注意深く管理することで、生成品質を損なうことなく、むしろ向上させながら、マスク生成モデルを大幅に効率化できると論じています。