[論文レビュー] Learning A Coarse-to-Fine Diffusion Transformer for Image Restoration

Liyan Wang, Qinyu Yang+4 more|arXiv (Cornell University)|Aug 17, 2023

MRI in cancer diagnosis被引用数 5

ひとこと要約

本稿では、画像復元のための粗大から細かい拡散Transformer（C2F-DFT）を提案し、ノイズの制約に加えて、サンプリングされた出力を真値と照合することで復元品質を向上させる、新たなトレーニング方式を導入している。本手法は、時間ステップに条件付けられた拡散自己注意機構とフィードフォワードネットワークを活用することで、雨消し、ぼかし除去、ノイズ除去のタスクにおいて最先端の性能を達成している。

ABSTRACT

Recent years have witnessed the remarkable performance of diffusion models in various vision tasks. However, for image restoration that aims to recover clear images with sharper details from given degraded observations, diffusion-based methods may fail to recover promising results due to inaccurate noise estimation. Moreover, simple constraining noises cannot effectively learn complex degradation information, which subsequently hinders the model capacity. To solve the above problems, we propose a coarse-to-fine diffusion Transformer (C2F-DFT) for image restoration. Specifically, our C2F-DFT contains diffusion self-attention (DFSA) and diffusion feed-forward network (DFN) within a new coarse-to-fine training scheme. The DFSA and DFN respectively capture the long-range diffusion dependencies and learn hierarchy diffusion representation to facilitate better restoration. In the coarse training stage, our C2F-DFT estimates noises and then generates the final clean image by a sampling algorithm. To further improve the restoration quality, we propose a simple yet effective fine training scheme. It first exploits the coarse-trained diffusion model with fixed steps to generate restoration results, which then would be constrained with corresponding ground-truth ones to optimize the models to remedy the unsatisfactory results affected by inaccurate noise estimation. Extensive experiments show that C2F-DFT significantly outperforms diffusion-based restoration method IR-SDE and achieves competitive performance compared with Transformer-based state-of-the-art methods on $3$ tasks, including image deraining, image deblurring, and real image denoising. Code is available at https://github.com/wlydlut/C2F-DFT.

研究の動機と目的

• 拡散ベースの画像復元モデルが、サンプリング中に不正確なノイズ推定を経験することで復元品質が低下するという限界を是正すること。
• 単なるノイズ制約を越えて、複雑な劣化パターンを学習する二段階トレーニング方式によりモデル容量を向上させること。
• 拡散埋め込み型Transformerアーキテクチャを用いて、画像復元における長距離依存関係のモデリングと階層的特徴学習を強化すること。
• ノイズではなく、サンプリングされた復元結果を用いてモデルを最適化する微調整戦略を開発し、不正確なノイズ推定に起因する誤差を是正すること。
• 雨消し、ぼかし除去、実世界のノイズ除去を含む、複数の画像復元ベンチマークで競争力のある性能を達成すること。

提案手法

• モデルは、時間ステップ埋め込みに条件付けられた拡散自己注意（DFSA）と拡散フィードフォワードネットワーク（DFN）を備えた拡散Transformerアーキテクチャを採用しており、拡散プロセスにおける時間的ダイナミクスをモデル化する。
• 粗大から細かいトレーニング方式を導入し、(a) ノイズ予測を制約する粗大トレーニング段階と、(b) サンプリングプロセスからのモデル出力を真値画像と照合する細かいトレーニング段階から構成される。
• 細かいトレーニング段階では、ノイズではなく、粗大トレーニング済みモデルから得られたサンプリングされた復元結果を用いてモデルを最適化することで、不正確なノイズ推定に起因する誤差を是正する。
• 一貫性と効率性を確保するため、微調整プロセスでは固定された4ステップのサンプリング手順を採用し、知覚的品質と計算コストのバランスを取る。
• 粗大トレーニング段階では、類似したトレーニング時間で固定パッチベースラインを上回る汎化性能と効率性を実現するため、パッチサイクルトレーニング戦略を採用する。
• 時間ステップ条件付けをDFSAおよびDFNに埋め込むことで、モデルが異なる拡散ステップに応じて動作を適応可能とし、復元性能を顕著に向上させる。

実験結果

リサーチクエスチョン

• RQ1 ノイズではなく、サンプリングされた出力を精練することで、拡散モデルにおける画像復元品質が向上するか？
• RQ2 拡散Transformerにおける時間ステップ埋め込みは、モデルの長距離依存関係の捉え方と復元忠実度にどのように影響するか？
• RQ3 微調整段階で生成出力に基づいてモデルを制約することで、標準的なノイズ制約トレーニングに比べて性能が向上するか？
• RQ4 微調整パイプラインにおいて、復元品質と推論速度のバランスを取る最適なサンプリングステップ数は何か？
• RQ5 パッチサイクルトレーニング戦略は、固定パッチトレーニングと比較して、モデルの汎化性能と効率性の向上にどの程度効果的か？

主な発見

• 提案されたC2F-DFTは、Rain100H雨消しベンチマークで31.38 dBのPSNRと0.921のSSIMを達成し、ベースラインの粗大トレーニング（30.93 dB PSNR）およびIR-SDEを上回った。
• Test2800ぼかし除去セットでは、34.03 dBのPSNRと0.944のSSIMを達成し、粗大トレーニングのみの結果より0.29 dBの向上を示した。
• 微調整段階により視覚的品質が顕著に向上し、粗大トレーニングと比較してアーティファクトとノイズが低減された。これは図9および図10で視覚的に示されている。
• 時間ステップ埋め込みによりPSNRが4.08 dB向上（26.85 dBから30.93 dB）した。これはその重要性を裏付ける。
• 微調整で4ステップのサンプリングを採用することで、知覚的品質（LPIPS）と推論時間の間で良好なトレードオフが達成され、少ないステップや多いステップよりも優れた性能を示した。
• パッチサイクルトレーニング戦略は、類似したトレーニング時間で固定パッチトレーニングを上回る性能を達成し、その有効性を確認した。