技術要約: TDiR (Transformer-based Diffusion for Image Restoration)

問題提起

画像復元はコンピュータビジョンにおける重要な課題であり、ノイズ、ぼけ、色被り、光の散乱などの影響を受けた劣化画像から、クリーンでオリジナルの画像を復元することを目指しています。これらの劣化は、物体検出、マッピング、分類といったダウンストリームタスクにおける画像の有用性を著しく損なわせます。非生成的な手法はハイレベルなビジョンタスクにおいて成功を収めていますが、水中再構成、デノイジング、除雨（deraining）のようなローレベルのビジョンタスクに見られる複雑でマルチモーダルな分布には、しばしば対応できません。特に水中画像は、環境条件によって変化する光の散乱、吸収、および色の変化という特有の課題を抱えており、単一の劣化した入力から単一のクリーンな出力へとマッピングするのではなく、複数の潜在的な劣化状態をナビゲートできるモデルを必要とする場合があります。

手法

本論文では、複数の画像復元タスクを統一されたフレームワーク内で扱うために設計された、トランスフォーマーベースの条件付き拡散モデルであるTDiRを提案しています。

• アーキテクチャ: TDiRの中核は、PromptIRのバックボーンを適応させたTransformerベースのU-Netアーキテクチャです。PromptIRは、個別のタスクごとに異なるモデルを必要とせずに、様々なレベルの劣化に適応するためのプロンプトモジュールを利用しています。
• 拡散プロセス: モデルは、前方拡散プロセス（ガウスノイズを段階的に追加する工程）と後方デノイジング段階からなる拡散フレームワーク内で動作します。
  • 前方プロセス: 入力画像 $y_0$ に対して、$T$ 回のイテレーションにわたってガウスノイズを加え、純粋なノイズの状態にします。
  • 後方プロセス: デノイジングネットワークは、このプロセスを逆転させることを学習し、ガウス分布からデータセットの実分布へと移動します。
• 主要なアーキテクチャの変更: 著者らは、専用のデコーダーパスウェイが、劣化した入力と拡散タイムステップの埋め込みを結合するという軽量な修正を導入しています。これにより、ネットワークは前方パス中に導入された特定のノイズレベルに基づいて、その復元プロセスを条件付けることが可能になります。
• 学習戦略: トランスフォーマーのデータ飢餓的な性質に対処するため、著者らは事前学習済みのエンコーダ（凍結）をバックボーンとして利用しています。学習中はデコーダーのパラメータのみが最適化されます。モデルは、水中強調、デノイジング、除雨の3つのタスクをカバーする結合データセットを用いて学習されます。目的関数は、予測されたノイズと画像に加えられた実際のノイズとの間の $L_1$ ロスを最小化します。
• 推論: 復元プロセスは反復的なデノイジング手順を伴い、画像ごとに複数のフォワードパスを必要とします。

主な貢献

1. 統一された条件付き拡散モデル: 事前学習済みのPromptIRバックボーンをノイズ条件付きデノイザーへと変換し、タスク固有のモデルの切り替えなしに、多様な復元タスク（水中、デノイジング、除雨）への適用を可能にするTDiRの提案。
2. 軽量なタイムステップ・コンディショニング: 劣化した入力と拡散タイムステップの埋め込みを結合するデコーダーパスウェイの導入。これにより、モデルはノイズレベルに基づいて復元戦略を動的に適応させることができます。
3. 包括的な評価: 5つのベンチマークデータセット（UIEB, Test-60, BSD68, Rain100Lなど）と4つの評価指標（PSNR, SSIM, UCIQE, UIQM）を用い、TDiRを18個の最先端技術と比較した厳格な評価。
4. 定性的分析: 純粋な識別ネットワークと比較した拡散ベースの復元プロセスの検討。これにより、残留パッチベースのアーティファクトという制約を認めつつも、マルチモーダルな劣化を扱う上での利点を明らかにしました。

実験結果

論文では、TDiRが評価されたタスク全体において現在の最先端手法を凌駕していることが報告されています。

• 水中画像強調: UIEBデータセットにおいて、TDiRは22.90 dBのPSNRと0.8724のSSIMという最高のフルリファレンス・スコアを達成し、最も近い競合モデルであるMetaUEをそれぞれ0.89 dBおよび0.011 SSIM上回りました。ノーリファレンスのTest-60サブセットでは、UIQMスコア2.73を達成し、他の手法を大幅に上回りました。定性的結果は、PromptIRや他のベースラインと比較して、優れた色の忠実度と緑色の色被りの抑制を示しています。
• 画像デノイジング: BSD68データセットにおいて、TDiRはノイズレベル $\sigma=15, 25, 50$ のすべてにおいて、すべてのベンチマーク手法（専門化されたシングルタスクモデルを含む）を上回りました。特に $\sigma=50$ において、TDiRは専門モデルであるPromptIRを上回りました。高ノイズレベルではパッチによるアーティファクトのためにSSIMがわずかに低下しましたが、モデルは視覚的にクリーンで、均一な領域における粒状のアーティファクトが少ない結果を生み出しました。
• 画像除雨: Rain100Lデータセットにおいて、TDiRは37.43のPSNRを達成し、再学習されたPromptIR*ベースラインを0.40 dB上回り、AirNetやMPRNetといった他の手法を大幅に上回りました。定性的分析により、均一な領域（空、砂など）における雨筋の完全な除去と、競合手法と比較した微細なテクスチャのより良い保持が確認されました。

意義と主張

本論文は、拡散モデルとトランスフォーマーの組み合わせが、特に複雑でマルチモーダルな劣化を伴うシナリオにおいて、画像復元のための堅牢なソリューションを提供すると主張しています。著者らは、拡散モデルが劣化した画像の多様な分布をナビゲートすることに長けており、現実世界の変動性を反映したより微細な復元を可能にすることを強調しています。これは、照明や透明度が大きく変化する水中イメージングにおいて特に有益です。

しかし、論文は自らの限界についても控えめなトーンを維持しています。以下の点を認めています：

• 計算コスト: 拡散の反復的な性質により、TDiRはシングルパスの識別モデルよりも推論コストが大幅に高く、リアルタイムアプリケーションへの適用を制限する可能性があります。
• アーティファクト: トランスフォーマーベースのデノイザーによるパッチ単位の処理は、ステッチ境界でのブロッキングアーティファクトを引き起こす可能性があり、知覚的な品質に対してPSNR/SSIMスコアをわずかに低下させることがあります。
• 汎用性のトレードオフ: 統一モデルは多用途ですが、特定のシナリオ（例：平坦な背景を持つ低深度の水中シーン）では、専門化されたシングルタスクモデルに一貫して劣っており、汎用性とピーク精度の間のトレードオフを示唆しています。
• 深度依存性: モデルは効果的な水中復元のために深度に関連する手がかりに依存しており、これは非常に困難なシーンでは欠如している場合があります。

著者らは、TDiRが複数のタスクにわたって画像品質の著しい向上を示した一方で、オーバーラップパッチ推論、より少ない拡散ステップへの蒸留、またはタスク条件付きファインチューニングなどの手法を通じて、これらの限界に対処できる可能性があると結論付けています。