技術要約：Diffusion Transformerのための品質認識型モジュレーション

問題提起

現代のテキスト・トゥ・イメージ（text-to-image）拡散モデル、特にStable Diffusion 3.5 (SD3.5) のようなDiffusion Transformer (DiT) は、Adaptive Layer Normalization (AdaLN) を介してデノイジング過程を制御するために、条件付け信号（具体的にはタイムステップ埋め込みとプールド・プロンプト埋め込み）に依存しています。これらの信号は、現在のノイズレベルや意味的な意図を効果的に伝達しますが、**品質認識情報（quality-aware information）**を欠いています。その結果、モジュレーション機構は、潜在画像が進化するにつれて発生するエラーの蓄積、プロンプトの誤った解釈、あるいは幻覚（hallucination）による特徴といった、サンプル固有の失敗モードに対して直接反応することができません。ベースラインの条件付けは現在の潜在状態に対して独立しているため（$\partial M_{base} / \partial x_t = 0$）、モデルは中間段階の画像の実際の品質に基づいて忠実度の問題を修正することができません。

手法

著者らは、学習済みのバックボーンを修正することなく、拡散プロセスに品質認識信号を注入するために設計された、軽量なトランスフォーマーベースのアダプターである**品質表現モジュール（Quality Representation Module: QRM）**を提案しています。

アーキテクチャ

QRMは、ベースラインのAdaLNパラメータ（$M_{base}$）に対する加法的なモジュレーション・オフセット（$M_{qrm}$）を予測する、トランスフォーマー・デコーダーとして実装されています。

• 入力: モジュールは、現在の潜在画像（$x_t$）、プールド・プロンプト埋め込み（$p$）、タイムステップ埋め込み（$t$）、およびベースラインのモジュレーション・パラメータ（$M_{base}$）を受け取ります。
• 処理:
  1. 潜在画像 $x_t$ は、学習可能な $1\times1$ 畳み込みと適応型プーリングを通じて、潜在トークンへと投影されます。
  2. ベースラインのモジュレーション・パラメータは正規化され、クエリ・トークン（$Z_0$）を形成するために投影されます。
  3. 潜在トークン、プロンプト、およびタイムステップ埋め込みから、キー・バリュー（KV）メモリが構築されます。
  4. クロスアテンションおよびセルフアテンション層を通じて、クエリ・トークンはKVメモリにアテンションを向け、洗練された表現を生成します。
  5. 出力は、ネイティブのAdaLN幅へと再投影され、オフセット・ベクトル $M_{qrm}$ を生成します。
• 適用: トランスフォーマー・ブロックに適用される最終的なモジュレーションは $M_t = M_{base,t} + M_{qrm,t}$ となります。

学習戦略 (ReFL)

QRMは、**報酬フィードバック学習（Reward Feedback Learning: ReFL）**戦略を用いて学習されます。

• 凍結されたコンポーネント: 拡散バックボーン（デノイザー、テキストエンコーダー、VAE）、サンプラー、および報酬モデルはすべて凍結されたままです。QRMのパラメータのみが更新されます。
• 目的関数: このモジュールは、中間デノイジング・ステップからデコードされた画像に適用される、凍結された報酬モデル（例：HPSv2.1またはCLIPScore）から導出されるスカラー報酬スコアを最大化するように学習します。
• 損失関数: 本論文では、報酬最大化損失、またはベースラインを特定のマージンで上回るように促すマージン追求型損失のいずれかを使用しています。
• タイムステップの選択: 効率性と構造形成への集中を高めるため、QRMは全拡散スケジュールの間ではなく、定義された特定の高ノイズ・タイムステップのサブセット（例：$t \in \{50, \dots, 26\}$）においてのみ活性化されます。

推論

推論時、学習済みのQRMモジュールはプラグイン・アダプターとして機能します。それは現在の潜在状態に基づいてモジュレーション・オフセットを予測し、それをベースラインのパラメータに加算します。これにより、バックボーンを再学習したりサンプリング・スケジュールを変更したりすることなく、モデルはデノイジングの軌道をより高品質な状態へと導くことが可能になります。

主な貢献

1. QRMアーキテクチャ: 既存のモデル入力（潜在、プロンプト、タイムステップ）から品質認識表現を学習し、モジュレーション・パラメータを洗練させる軽量なトランスフォーマー・モジュールの導入。
2. 非侵襲的な統合: 基礎となる拡散バックボーン、VAE、またはサンプリング・スケジュールを変更することなく、潜在画像の品質信号を用いて拡散トランスフォーマーのモジュレーションを拡張するメカニズム。
3. 報酬誘導型モジュレーション: 予測されたモジュレーション更新（$M_{qrm}$）が、各タイムステップにおいてデノイジング過程を効果的に洗練させ、プロンプトへの忠実度と画像の品質を向上させられることを実証。
4. 包括的な評価: QRMのアーキテクチャサイズ、活性化タイムステップ範囲、および学習報酬メトリクスの影響を分析する広範なアブレーション研究。

実験結果

著者らは、Parti-Promptsベンチマークにおける**Stable Diffusion 3.5 (SD3.5)**を用いてQRMを評価しました。

• 定量的パフォーマンス:
  • SD3.5のベースラインと比較して、報酬信号としてHPSv2.1を用いて学習されたQRMは、評価されたすべての手法（プロンプト洗練、ファインチューニング、ReNegアプローチを含む）の中で最高のスコアを達成しました。
  • 改善度: Aestheticスコアで+0.20、HPSv2.1スコアで+0.82の向上。
  • CLIPScoreで学習した場合、QRMは知覚的な品質を維持しつつ、人間による好みの整合性（HPSv2.1）を向上させました。
• 定性的分析:
  • デノイジングの軌道の可視化により、初期の高ノイズ・タイムステップにおいてQRMを適用することで、ベースラインよりも大幅に早く、認識可能な意味的構造が出現することが示されました。
  • 中期または後期の段階でのみQRMを適用した場合、改善は限定的であり、グローバルな構造は拡散プロセスにおける早い段階で確立されることが確認されました。
• アブレーションによる知見:
  • 容量（Capacity）: より大きなQRMモデルは、特定の学習報酬（例：HPSv2.1）に対してより強く最適化されますが、他のメトリクス（例：CLIPScore）との間でわずかなトレードオフが見られました。
  • タイムステップ範囲: 初期のタイムステップ（$50 \to 26$）でのモジュレーションは、後期の段階でのモジュレーションよりも一貫して整合性の面で効果的でした。
  • 報酬メトリクス: 報酬モデル（CLIPScore vs HPSv2.1）の選択は、QRMが学習する知覚的なバイアスを直接決定しました。

重要性と主張

本論文は、モジュレーション・レベルの介入が、大規模な拡散モデルを人間の好みの信号に整合させるための、スケーラブルかつ効果的な経路を提供すると主張しています。バックボーン全体の高価なファインチューニング（LoRAやフル・ファインチューニングなど）を必要とする手法とは異なり、QRMは進化する潜在表現に条件付けられた加法的なモジュレーション更新を通じて完全に動作することで、一貫した品質向上を実現します。

著者らは、QRMを、異なる報酬信号によって制御可能であり、様々な拡散トランスフォーマー・アーキテクチャに適用可能な、柔軟で品質認識型のアダプターとして位置付けています。QRMは学習に使用した報酬モデルのバイアスを継承することを認めており、今後の研究では、堅牢性と汎用性をさらに高めるために、マルチ報酬学習や適応的なタイムステップ選択を探索できる可能性があるとしています。本研究は、最先端モデルの学習済み重みを変更することなく、画像忠実度とプロンプト整合性の大幅な向上を実現できることを示しています。