TOC-SR: Task-Optimal Compact diffusion for Image Super Resolution

TOC-SR の技術的概要：画像超解像のためのタスク最適化コンパクト拡散モデル

問題定義
拡散モデルは、自然画像の統計をモデル化し強力な生成事前分布を提供することで、単一画像超解像（SR）のような画像復元タスクにおいて卓越した性能を示していますが、実用的な展開には大きな障壁が存在します。標準的な拡散パイプラインは、大規模なバックボーンネットワークと反復的なサンプリング手順（多くの場合数十回のノイズ除去ステップ）を必要とするため、高い遅延、メモリ使用量、計算コストが生じます。コンパクトな一歩 SR モデルを作成しようとする既存の試みは、モデルが限られた容量内で生成事前分布と復元マッピングを同時に学習しなければならないため、不安定な最適化と品質の低下に悩まされることが多いです。

手法
著者らは、生成忠実度を維持するコンパクトな拡散バックボーンをまず発見することにより、効率的な一歩 SR モデルを構築するためのフレームワーク TOC-SR を提案します。この手法は、以下の 3 つの明確な段階で進行します。

1. 潜在空間容量の拡張（4 チャンルから 16 チャンルへ）：
   Stable Diffusion の標準的な 4 チャンル潜在表現が、SR に必要な高周波詳細情報のモデル化において情報ボトルネックを生み出していることを認識し、著者らは空間圧縮係数を維持したまま潜在次元を 16 チャンルに拡張します。標準的な 16 チャンル Stable Diffusion モデルが存在しないため、高容量 VAE（FLUX モデルファミリーから）の表現能力を修正された Stable Diffusion VAE に蒸留することで、新しいバックボーンを構築します。これには、16 チャンル潜在分布に一致しつつ再構成忠実度を維持するように学生エンコーダーを訓練することが含まれます。その後、拡散 U-Net を 16 チャンルテンソルを受け入れ生成するように修正し、拡散事前分布をこの拡張された空間に適応させるために共同微調整を行います。

2. 代理モデルに基づく $\epsilon$ 制約ベイズ最適化によるコンパクトバックボーンの発見：
   拡張された 16 チャンルモデルの計算コストを削減するため、著者らはコンパクトモデルの直接訓練ではなく、探索戦略を採用します。

   • 代理モデルの構築： 教師 U-Net を逐次的な段階に分解します。各段階について、構造化プルーニングとモジュール削減（例えば、アテンションモジュールの除去や残差ブロックの簡素化）を通じて、入力 - 出力テンソルの次元性を厳密に維持しつつ、パラメータ効率の良い代理ブロックのライブラリを構築します。
   • 特徴量別生成蒸留（FwKD）： 各候補代理ブロックを、校正サブセットにおける教師の特徴量に対する特徴量別蒸留損失を用いて独立して較正します。これにより、探索段階での完全な拡散訓練の必要性を回避します。
   • $\epsilon$ 制約ベイズ最適化（BO）： 構造探索を制約付き最適化問題として定式化します。目的は、生成忠実度（$f_{acc} \leq \epsilon$）の制約条件下でパラメータ数（$f_{param}$）を最小化することであり、ここで $f_{acc}$ は候補と教師間のノイズ予測の不一致を測定します。ガウス過程代理モデルが BO プロセスを導き、$\epsilon$ 値のグリッド全体で構造空間を効率的に探索し、圧縮と忠実度の間の最適なトレードオフを提供する構造を特定します。

3. タスク特化と一歩蒸留：
   最適なコンパクトバックボーンが特定されたら、条件付き拡散訓練を通じて SR タスクに適応させます。ここでは、モデルが低解像度（LR）入力に条件付けられた潜在表現のノイズ除去を学習します。最後に、この反復的拡散 SR モデルを一歩生成器に蒸留します。学生ネットワークは、ノイズの付いた潜在表現と LR 条件から最終的なクリーンな高解像度潜在を直接予測するように訓練され、反復的なノイズ除去の必要性を排除します。

主要な貢献

• SR のための潜在空間拡張： 本論文は、潜在空間を 4 チャンルから 16 チャンルに拡張することが、高周波詳細のモデル化におけるボトルネックを緩和し、標準的なアーキテクチャよりも圧縮のためのより強力な出発点を提供することを示しています。
• 代理モデルに基づく構造探索： 著者らは、特徴量別生成蒸留と $\epsilon$ 制約ベイズ最適化を組み合わせた新規フレームワークを導入し、コンパクトな拡散バックボーンを発見します。このアプローチは、探索中にタスク固有の訓練を必要とすることなく、候補アーキテクチャが教師の生成ダイナミクスに忠実であることを保証しつつ、パラメータを体系的に削減します。
• 効率的な一歩 SR： 結果として得られるフレームワーク TOC-SR は、コンパクトな拡散事前分布を一歩推論メカニズムと成功裡に統合し、反復サンプリングなしで高品質な復元を達成します。

実験結果
実験は、最先端の拡散ベース SR 手法（StableSR、DiffBIR、SeeSR など）および既存のコンパクト/一歩モデル（SinSR、OSEDiff、PocketSR など）と TOC-SR を比較する標準ベンチマーク（DIV2K、RealSR、DRealSR）上で実施されました。

• 効率性： 拡張された 16 チャンル拡散モデルと比較して、発見された TOC-SR バックボーンは、パラメータを 6.6 倍、GMACs を 2.8 倍 削減しています。
• 性能： DIV2K 検証セットにおいて、TOC-SR は 24.26 dB の PSNR を達成し、リストされた手法の中で最良の LPIPS スコア（0.242）を記録しました。これらはすべて単一ステップで動作しています。
• 実世界への汎化性： 実世界データセット（RealSR および DRealSR）において、TOC-SR は知覚指標（MUSIQ）において他のコンパクトなベースラインを上回り、競争力のある PSNR を維持することで、実世界の劣化に対する頑健性を示しています。
• 定性的分析： 視覚的比較により、TOC-SR は他のコンパクトモデルよりも微細なテクスチャや細い構造をよりよく保持し、過度な平滑化や不自然な微細テクスチャといった一般的なアーティファクトを回避することが示されています。

意義と主張
本論文は、TOC-SR が高品質かつ効率的な復元システムを構築するための安定した有効な経路を提供すると主張しています。コンパクトな生成事前分布の発見をタスク固有の適応から分離することで、このフレームワークは、コンパクトな SR モデルを直接訓練する際にしばしば見られる最適化の不安定性を回避します。著者らは、このアプローチが拡散モデルの生成力と実用的展開に必要な効率性のバランスに成功し、計算コスト（数十ステップから 1 ステップへ）を大幅に削減しながら強力な復元品質を維持していると述べています。この研究は、制約付き最適化を通じてコンパクトな拡散事前分布を発見することが、エッジ対応画像復元のための viable な戦略であることを示唆しています。