[論文レビュー] Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Ziwei Luo, Fredrik Gustafsson+3 more|arXiv (Cornell University)|Apr 17, 2023

Advanced Image Processing Techniques被引用数 6

ひとこと要約

Refusionは、高分解能の詳細をデコーダーに接続された潜在変数を介して保持しながら、低次元の低解像度潜在空間でノイズ除去を行うU-Netベースの潜在拡散モデルを提案する。この手法により、敵対的学習やデータセットの選別を必要とせず、6000×4000の画像を処理可能であり、実世界の画像修復タスクにおいて最先端の知覚的性能を達成している。NTIRE 2023 シャドウ除去チャレンジでは2位を獲得した。

ABSTRACT

This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also propose a U-Net based latent diffusion model which performs diffusion in a low-resolution latent space while preserving high-resolution information from the original input for the decoding process. Compared to the previous latent-diffusion model which trains a VAE-GAN to compress the image, our proposed U-Net compression strategy is significantly more stable and can recover highly accurate images without relying on adversarial optimization. Importantly, these modifications allow us to apply diffusion models to various image restoration tasks, including real-world shadow removal, HR non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation. By simply replacing the datasets and slightly changing the noise network, our model, named Refusion, is able to deal with large-size images (e.g., 6000 x 4000 x 3 in HR dehazing) and produces good results on all the above restoration problems. Our Refusion achieves the best perceptual performance in the NTIRE 2023 Image Shadow Removal Challenge and wins 2nd place overall.

研究の動機と目的

• 実世界の複雑な劣化を伴う大規模な画像修復タスクにおける拡散モデルの適用性を向上させること。
• 従来の拡散モデルの限界、例えば高い計算コスト、既知の劣化パラメータへの依存、潜在空間圧縮における不安定性を克服すること。
• 敵対的最適化を必要とせず、多様な修復タスクに一般化可能な安定的で効率的かつ汎用性の高いフレームワークを開発すること。
• 低次元の潜在空間で動作することで、6000×4000などの高解像度画像修復を可能にするとともに、微細な詳細を保持すること。

提案手法

• 逆方向のノイズ除去プロセスとして、IR-SDEにインspiredされた平均回帰型確率微分方程式（SDE）を用い、劣化を暗黙的にモデル化することでタスクの一般化を可能にする。
• 従来の潜在拡散モデルで使われるVAE-GANに代わり、U-Netベースの圧縮ネットワークを採用することで、敵対的学習を必要とせず、より安定的かつ正確な潜在表現学習を実現する。
• ノイズ予測ヘッドとして、非線形活性化関数を含まないネットワーク（NAFNet）を採用し、標準的なU-Netアーキテクチャに比べて効率性と性能を向上させる。
• 潜在空間をノイズ除去に使用する一方で、デコーダーへのスキップ接続を通じて高解像度特徴を保持することで、微細な詳細の正確な再構成を可能にする。
• データセットとノイズネットワークを単に変更するだけで、さまざまな画像修復タスクに対応可能であり、柔軟性と一般化性能を確保する。
• 潜在Refusionバージョンは、潜在空間で拡散を実行することで、FLOPsを約15倍削減し、6000×4000のような大規模な画像の推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 敵対的学習やVAE-GAN圧縮に依存せずに、U-Netベースの潜在拡散モデルが安定的かつ高精細な画像修復を達成できるか？
• RQ2 U-NetノイズネットワークをNAFNetに置き換えると、画像修復タスクにおける性能と効率性にどのような影響を与えるか？
• RQ3 圧縮された潜在空間における潜在拡散が、非常に大きな画像（例：6000×4000）の効果的な修復を可能にするとともに、高解像度の詳細を保持できるか？
• RQ4 提案手法が、シャドウ除去、ヒーディング除去、ボケ効果変換などの多様な実世界の画像修復タスクにどの程度一般化できるか？
• RQ5 大規模なデータセット上で、既存の拡散ベース手法と比較して、知覚的品質と推論速度の面で、本手法はどの程度優れているか？

主な発見

• RefusionはNTIRE 2023 イメージシャドウ除去チャレンジで2位を獲得し、LPIPS（0.149）とMOS（8.94）の最高の知覚的スコアを記録し、他のすべてのチームを上回った。
• 実世界のシャドウ除去タスクにおいて、LPIPSが0.149、MOSが8.94という優れた知覚的性能を達成し、DHANなどのベースラインモデルを大きく上回った。
• 潜在Refusionバージョンは、フルRefusionモデルに比べてFLOPsを約15倍削減し、フルRefusionやIR-SDEが失敗する6000×4000の画像の推論を可能にした。
• ボケ効果変換タスクでは、潜在RefusionはRestormerよりも良いLPIPSスコアを達成し、ほぼ同等の速度で動作し、よりシャープな結果を生成した。
• U-Netベースの潜在圧縮戦略は、VAE-GANベースの手法よりも著しく安定しており、敵対的最適化を必要とせず、高精度な画像回復を実現した。
• 本手法は、メモリと計算リソースの制約により、フル解像度の拡散モデル（IR-SDE や Refusion）が非現実的となる非一様なヒーディング除去を6000×4000の画像で成功裏に実行した。