[論文レビュー] MAXIM: Multi-Axis MLP for Image Processing

Zhengzhong Tu, Hossein Talebi+5 more|arXiv (Cornell University)|Jan 9, 2022

Visual Attention and Saliency Detection被引用数 7

ひとこと要約

MAXIMは、ノイズ除去、ぼかし回復、雨除去、ホコリ除去、強調処理など多様な低レベルビジョンタスクにおいて最先端の性能を達成する、画像処理のための新規なマルチアックスMLPアーキテクチャを提案する。マルチアックスゲーテッドMLPとクロスゲーティングブロックを用いた完全畳み込み型で、グローバルな受容 field を持つ設計により、従来のモデルよりも少ないパラメータ数とFLOPsで優れた結果を達成している。

ABSTRACT

Recent progress on Transformers and multi-layer perceptron (MLP) models provide new network architectural designs for computer vision tasks. Although these models proved to be effective in many vision tasks such as image recognition, there remain challenges in adapting them for low-level vision. The inflexibility to support high-resolution images and limitations of local attention are perhaps the main bottlenecks. In this work, we present a multi-axis MLP based architecture called MAXIM, that can serve as an efficient and flexible general-purpose vision backbone for image processing tasks. MAXIM uses a UNet-shaped hierarchical structure and supports long-range interactions enabled by spatially-gated MLPs. Specifically, MAXIM contains two MLP-based building blocks: a multi-axis gated MLP that allows for efficient and scalable spatial mixing of local and global visual cues, and a cross-gating block, an alternative to cross-attention, which accounts for cross-feature conditioning. Both these modules are exclusively based on MLPs, but also benefit from being both global and `fully-convolutional', two properties that are desirable for image processing. Our extensive experimental results show that the proposed MAXIM model achieves state-of-the-art performance on more than ten benchmarks across a range of image processing tasks, including denoising, deblurring, deraining, dehazing, and enhancement while requiring fewer or comparable numbers of parameters and FLOPs than competitive models. The source code and trained models will be available at \url{https://github.com/google-research/maxim}.

研究の動機と目的

• 変換器と畳み込みニューラルネットワーク（CNN）が低レベルビジョンタスクにおいて抱える制限、特に高解像度画像に対する柔軟性の欠如と受容 field の制限を解決すること。
• 長距離相互作用と密度予測タスクをサポートする、柔軟で効率的かつ汎用的なビジョンバックボーンを構築すること。
• 自己注意機構を純粋なMLPベースのモジュールに置き換え、画像サイズに線形にスケーリングするグローバルな文脈を維持すること。
• マルチステージ・マルチスケールアーキテクチャを用いて、ノイズ除去、ぼかし回復、雨除去、ホコリ除去、強調処理などの画像修復タスクの性能を向上させること。

提案手法

• マルチスケール・マルチステージの監視を可能にするために、MAXIMバックボーンをスタックしたUNet型のエンコーダ・デコーダアーキテクチャを採用する。
• 空間グリッドと局所ブロックを別々に処理するマルチアックスゲーテッドMLPブロックを導入し、線形の計算複雑度でグローバルおよびローカル特徴の混合を実現する。
• 空間ゲーティングユニットは、グリッドとブロックの両方の軸に沿ってゲーテッドMLPを適用し、長距離依存性を保持するとともに、完全畳み込み型の動作を可能にする。
• クロスアテンションを置き換えるために、同じマルチアックスアプローチを用いてスカイプ接続を適応的に条件づけるクロスゲーティングブロックを導入し、グローバルでパラメータ効率の良い特徴調製を実現する。
• 異なる解像度での特徴学習を強化するために、訓練中にマルチスケール損失を用いる。
• 自己注意や畳み込みのインダクティブバイアスを一切排除した完全MLPベースのアーキテクチャであり、密度画像処理に非常に柔軟である。

実験結果

リサーチクエスチョン

• RQ1 純粋なMLPベースのアーキテクチャは、低レベル画像修復タスクにおいて、自己注意ベースのモデルを上回ることができるか？
• RQ2 マルチアックスMLP設計は、画像サイズに線形にスケーリングされるグローバルな文脈モデリングを達成できるか？
• RQ3 完全畳み込み型でMLPベースのバックボーンは、パッチ化やクロッピングのアーティファクトなしに高解像度画像処理を効果的に行えるか？
• RQ4 画像修復における特徴相互作用モデリングにおいて、クロスゲーティング機構はクロスアテンションと比較してどのように異なるか？
• RQ5 統一的で軽量なアーキテクチャは、多様な画像処理ベンチマークで最先端の性能を達成できるか？

主な発見

• MAXIMは、ノイズ除去、ぼかし回復、雨除去、ホコリ除去、強調処理の各タスクにおいて、10以上のベンチマークで最先端の性能を達成した。
• SIDDデータセットにおける画像ノイズ除去では、先行SOTAよりPSNRが+0.24 dB向上した。
• GoProデータセットにおける画像ぼかし回復では、前回最良モデルよりPSNRが+0.15 dB向上した。
• Rain100Lデータセットにおける雨除去では、PSNRが+0.86 dB向上し、先行手法を著しく上回った。
• RESIDEデータセットにおける画像ホコリ除去では、PSNRが+0.94 dB向上し、複雑な劣化に対して強い性能を示した。
• FiveKデータセットにおける画像リタッチでは、PSNRが+1.15 dB向上し、優れた一般化性能と詳細回復能力を示した。