[論文レビュー] U2-Former: A Nested U-shaped Transformer for Image Restoration

Haobo Ji, Xin Feng+3 more|arXiv (Cornell University)|Dec 4, 2021

Image Enhancement Techniques被引用数 14

ひとこと要約

U2-Former は、マルチスケール特徴相互作用と特徴フィルタリング機構を活用することで、計算効率を向上させつつ、画像復元のための深くネストされたU字型のトランスフォーマー・アーキテクチャを提案する。マルチビュー対照学習を統合することで、ノイズと綺麗な画像を分離し、リフレクション除去、雨除去、ホイミングの分野で最先端の性能を達成する。

ABSTRACT

While Transformer has achieved remarkable performance in various high-level vision tasks, it is still challenging to exploit the full potential of Transformer in image restoration. The crux lies in the limited depth of applying Transformer in the typical encoder-decoder framework for image restoration, resulting from heavy self-attention computation load and inefficient communications across different depth (scales) of layers. In this paper, we present a deep and effective Transformer-based network for image restoration, termed as U2-Former, which is able to employ Transformer as the core operation to perform image restoration in a deep encoding and decoding space. Specifically, it leverages the nested U-shaped structure to facilitate the interactions across different layers with different scales of feature maps. Furthermore, we optimize the computational efficiency for the basic Transformer block by introducing a feature-filtering mechanism to compress the token representation. Apart from the typical supervision ways for image restoration, our U2-Former also performs contrastive learning in multiple aspects to further decouple the noise component from the background image. Extensive experiments on various image restoration tasks, including reflection removal, rain streak removal and dehazing respectively, demonstrate the effectiveness of the proposed U2-Former.

研究の動機と目的

• 高コストな計算とスケール間通信の悪さのため、画像復元におけるトランスフォーマーの深さが制限されている問題に対処する。
• 標準的なエンコーダ・デコーダ型トランスフォーマーの非効率性を克服し、より深い特徴学習と改善されたスケール間相互作用を可能にする。
• パッチレベル、コンテンツレベル、ノイズレベルの表現を介したマルチビュー対照学習を導入することで、ノイズと背景の分離を向上させる。
• 低品質なトークンを圧縮する特徴フィルタリング機構により、自己注意ブロックの計算効率を最適化する。
• 畳み込みニューラルネットワーク（CNN）ベースの事前知識に依存せずに、リフレクション除去、雨除去、ホイミングを含む多様な画像復元タスクで優れた性能を達成する。

提案手法

• 外側にエンコーダ・デコーダフレームワーク、内側にU字型トランスフォーマーブロックを備えたネストされたU字型アーキテクチャを採用し、深くマルチスケールの特徴相互作用を実現する。
• 自己注意計算の前に低品質なトークンを pruning する特徴フィルタリング機構を導入することで、計算負荷を低減しつつも、重要な特徴を保持する。
• 3つのポジティブペア構成を用いたマルチビュー対照学習を実装する：(1) 同一の復元画像からのパッチ、(2) 復元画像と正解の綺麗な画像、(3) 復元画像とランダムな正解画像。これによりノイズ不変の特徴を学習する。
• 対照学習に加えて正解画像を用いた標準的な監視を用い、復元と特徴の分離を同時に最適化する。
• エンコーディングおよびデコーディングパスに深くトランスフォーマーブロックをスタックさせ、スケール間でグローバルな文脈モデリングを可能にする。
• すべてのレイヤーで自己注意機構を活用し、長距離依存性を捉えることで、復元画像の空間的整合性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 ネストされたU字型の深さを持つトランスフォーマー・アーキテクチャは、画像復元におけるスケール間の特徴通信を改善できるか？
• RQ2 トランスフォーマーブロックに特徴フィルタリング機構を導入することで、計算コストを著しく低減しつつ性能を維持できるか？
• RQ3 マルチビュー対照学習は、画像復元においてノイズ成分を綺麗な画像コンテンツから効果的に分離できるか？
• RQ4 U2-Formerは、多様な画像復元タスクにおいて、PSNR および SSIM の観点で最先端の手法と比較してどのように差をつけるか？
• RQ5 ネストされたU字型設計は、標準的なU-Net や Uformer スタイルのアーキテクチャと比較して、どれほど性能を向上させるか？

主な発見

• 実世界のリフレクション除去ベンチマークでは、U2-Former は Real20 データセットで 23.67 dB の PSNR と 0.835 の SSIM を達成し、Uformer（22.02 dB PSNR）や他の最先端手法を上回った。
• Rain100L データセットにおける雨除去では、U2-Former は 39.31 dB の PSNR を達成し、Uformer（37.68 dB）と比較して 1.63 dB の向上を示した。視覚的評価でも詳細回復が優れていた。
• Rain100H データセットでは、U2-Former が 30.87 dB の PSNR を達成し、Uformer よりも 1.50 dB の向上を示し、重度の雨条件下でも優れた性能を発揮した。
• 画像ホイミングでは、屋内セットで 36.42 dB の PSNR と 0.988 の SSIM、屋外セットで 31.10 dB の PSNR と 0.976 の SSIM を達成し、Uformer よりもそれぞれ 4.51 dB と 4.58 dB の向上を示した。
• 視覚的結果から、U2-Former は色歪みや過剰な平滑化を軽減し、DCP や DA などのベースライン手法と比較して、より細かいディテールを良好に保持していた。
• アブレーションスタディの結果、ネストされたU字型構造とマルチビュー対照学習の両方が性能向上に顕著な貢献を示した。また、特徴フィルタリング機構により推論効率が向上したが、性能の低下は見られなかった。