技術要約: MOSAIK

問題提起

ピクセル空間拡散モデル（PDM）は、変分オートエンコーダ（VAE）による再構成の天井を回避することで、微細な高周波空間詳細を保持できるという、潜在拡散モデル（LDM）に対する大きな利点を提供します。しかし、ピクセル空間で直接生成することは、トークン数の大幅な増加を招きます。トランスフォーマーの自己注意機構（self-attention mechanism）はトークン数に対して二次的な計算量を持つため、PDMは法外な計算コストを課されることになります。

既存の効率化手法は、デノイジングのタイムステップやネットワークの深さに応じてパッチサイズ（単一のトークンが表す空間領域）を変化させることで、これを軽減しようと試みています。これらのアプローチは総トークン数を削減しますが、任意のステップにおいて画像全体に空間的に一様なパッチレイアウトを適用します。この一様性は、根本的な非効率性を生み出します。つまり、単純で情報の少ない領域に対して計算を使いすぎるか、逆に、微細な詳細が不可欠な敏感な領域において表現能力を制限してしまうかのいずれかを選択することになります。本論文は、視覚情報の密度は本質的に不均一であり、現在の手法には欠けている動的な領域ごとの割り当て戦略が必要であると主張しています。

手法: MOSAIK

MOSAIK（Multi-patch cOntent-aware Spatial Allocation of Image toKens）は、粗粒化（coarsening）によって予測される「ダメージ」（忠実度の損失）に基づいて、異なる画像領域およびデノイジングステップ間でパッチサイズを変化させるように設計されたフレームワークです。このアプローチはPixelDiTバックボーンに基づいて構築されており、3段階のトレーニングパイプラインを利用しています。

ステージ1: マルチパッチ流暢性（Multi-Patch Fluency）

目的は、元の生成的な事前分布（prior）に対する破滅的忘却を起こすことなく、複数の一様なパッチサイズ（$P \in \{16, 32, 64\}$）で動作できるようにモデルを装備することです。

• アーキテクチャ: バックボーンの重みは凍結されます。低ランク適応（LoRA）モジュールがDiTブロックに注入されます。$p32$および$p64$用の個別のトークナイザが導入され、これらは事前学習済みの$p16$投影をバイリニアダウンサンプリングで合成することによって初期化されます。
• トレーニング: モデルは自己蒸留（self-distillation）を行い、すべての単一パッチサイズにわたって、凍結された$p16$ティーチャーモデルの速度予測に一致するように最適化されます。

ステージ2: ミックスパッチ適応（Mixed-Patch Adaptation）

このステージでは、モデルが単一のフォワードパス内で、任意かつ不均一なパッチレイアウト（異なる領域が異なるパッチサイズを使用する構成）を処理することを可能にします。

• 位置エンコーディング: 標準的な回転位置エンコーディング（RoPE）は一様なグリッドを想定しています。MOSAIKは、正規化された$p16$座標フレーム内における各パッチの幾何学的中心に各トークンを固定することで、ミックスパッチによる空間的な不一致に対処します。
• トレーニング: ジェネレータは、ランダムにサンプリングされた不均一なレイアウトでトレーニングされ、パッチサイズの空間配置を画像の内容から意図的に切り離します。これにより、モデルは、配置が視覚的特徴とどのように整列しているかにかかわらず、あらゆるレイアウトを実行することを学習します。

ステージ3: ダメージ誘導型レイアウト予測器（Damage-Guided Layout Predictor）

このステージでは、特定の領域を粗粒化することによって引き起こされる劣化（degradation）を推定し、トークンを動的に割り当てる軽量な予測器を導入します。

• ダメージの定義: 「ダメージ」は、一様な$p64$レイアウトの予測と、高忠実度な$p16$ティーチャーとの間の二乗誤差として定義され、重なり合わない$64 \times 64$の領域にわたって平均化されます。
• 予測器のトレーニング: FiLMコンディショニングを備えた軽量な畳み込みネットワークが、中間的なバックボーン特徴量（具体的にはDiTブロック6の活性化値）とテキスト埋め込みから、このダメージマップを直接回帰するようにトレーニングされます。これにより、推論時の二次的なフォワードパスを回避できます。
• 推論戦略: 各デノイジングステップにおいて、予測器はダメージマップを推定します。予測ダメージが高い領域には細かいパッチ（$p16$）が割り当てられ、ダメージが低い領域には粗いパッチ（$p32$または$p64$）が割り当てられます。割り当てはユーザー指定のトークン予算によって制約され、追加されたトークンあたりのダメージ減少が最も大きいリファインメントを優先します。

主な貢献

1. 空間適応型パッチ化: MOSAIKは、ピクセル空間拡散モデルにおいて、従来の効率化手法の空間的一様性の制約を超え、不均一なパッチレイアウトを適用する最初のフレームワークです。
2. ダメージ誘導型割り当て: 著者らは、パッチ割り当てをダメージ予測問題として定式化し、軽量な予測器を使用して、粗粒化が最も大きな忠実度損失を引き起こす領域を特定します。
3. 3段階のトレーニングパイプライン: ルーターを再学習することなくジェネレータをガイドするために、まずマルチパッチ流暢性を確立し、次にミックスレイアウトに適応させ、最後に個別の予測器をトレーニングするという新しいトレーニング戦略を採用しています。
4. 統合可能性: 本フレームワークは、特徴量キャッシュのような時間的効率化手法と直交するように設計されており、結合された最適化が可能です。

結果

1024x1024の画像生成に関する実験がGenEvalおよびDPG-Benchベンチマークで行われ、一様パッチ削減、時間的パッチスケジューリング（DDiT）、および特徴量キャッシュ（TeaCache, TaylorSeer）を含むベースラインと比較されました。

• 効率 vs 品質: MOSAIKは、GenEvalにおいてフル計算量のPixelDiTと同等のスコア（0.74）を維持しながら、70%のFLOPs削減と83%のトークン数削減を達成し、DPG-Benchではわずか1.0ポイントの低下（82.5 vs 83.5）にとどまりました。
• 時間的スケジューリングとの比較: 積極的な計算体制（例：60-70%のFLOPs削減）において、MOSAIKは一貫してDDiTを上回ります。DDiTは単一のパッチサイズをグローバルに適用しますが、MOSAIKは計算リソースを敏感な領域（例：被写体の境界）に動的に再配分することで、構造的な完全性とテクスチャやエッジなどの微細な詳細を保持します。
• キャッシュとの比較: キャッシュベースのベースライン（TeaCache, TaylorSeeler）と比較して、MOSAIKは、特に高度に制約された条件下において、同等の予算で優れた性能を発揮します。
• 相補性: キャッシング手法と組み合わせた場合、MOSAIKは、キャッシュ単体では大幅な品質崩壊が起こるような高い圧縮レベル（最大80%のFLOPs削減）においても、高いパフォーマンスを維持します。
• アブレーション: ダメージ誘導型予測器をランダムまたはコンテンツ盲目的な割り当て器に置き換えると、大幅な性能低下が見られました。これは、予測されたダメージのランキングが、空間的な感度に関する真の情報を保持していることを検証しています。

意義と主張

本論文は、MOSAIKが、時間的または深さ方向の適応に焦点を当てた手法で見落とされてきた、空間的な不均一性が効率的な生成における重要な要因であることを示していると主張しています。

• 最適なリソース配分: 著者らは、一様な細かいパッチ化は単純な領域に対して計算を浪費し、一様な粗いパッチ化は複雑な領域の細部を破壊すると主張しています。MOSAIKは、最も必要とされる場所にリソースを割り当てることで、このギャップを埋めます。
• ピクセル空間の利点: 本フレームワークは、ピクセル空間モデルが持つ固有の空間的冗長性（潜在モデルではVAEによる圧縮がすでに冗長性を消費しているのと対照的）を活用し、ルーターとジェネレータの端から端までの共同最適化を必要とせずに、動的な粗粒化を利用します。
• スケーラビリティ: このアプローチは、中程度の削減から積極的な制約に至るまで、幅広い計算予算において安定したパフォーマンスを可能にし、異なる予算レベルのための再学習を必要としません。

結論として、MOSAIKはピクセル空間拡散における効率性の新しい標準を確立し、時間的スケジューリングやキャッシュに対して非常に競争力のある選択肢を提供しており、将来の研究として、空間と時間の両方における冗長性を活用したビデオ拡散への拡張を示唆しています。