[論文レビュー] Image-to-Image MLP-mixer for Image Reconstruction

Youssef Mansour, Kang Lin+1 more|arXiv (Cornell University)|Feb 4, 2022

Image and Signal Denoising Methods被引用数 8

ひとこと要約

この論文では、畳み込みやマルチスケール構造を用いない画像再構成タスク（ノイズ除去や圧縮センシングなど）において、U-Net やビジョントランスフォーマーよりも優れた性能を発揮する画像対画像 MLP-mixer アーキテクチャを提案する。空間パッチの位置を保持することで、パラメータ数が少なく、画像解像度に比例して線形に増加するパrameterスケーリングを実現し、中程度のデータセットで学習された場合でも最先端の性能を達成している。

ABSTRACT

Neural networks are highly effective tools for image reconstruction problems such as denoising and compressive sensing. To date, neural networks for image reconstruction are almost exclusively convolutional. The most popular architecture is the U-Net, a convolutional network with a multi-resolution architecture. In this work, we show that a simple network based on the multi-layer perceptron (MLP)-mixer enables state-of-the art image reconstruction performance without convolutions and without a multi-resolution architecture, provided that the training set and the size of the network are moderately large. Similar to the original MLP-mixer, the image-to-image MLP-mixer is based exclusively on MLPs operating on linearly-transformed image patches. Contrary to the original MLP-mixer, we incorporate structure by retaining the relative positions of the image patches. This imposes an inductive bias towards natural images which enables the image-to-image MLP-mixer to learn to denoise images based on fewer examples than the original MLP-mixer. Moreover, the image-to-image MLP-mixer requires fewer parameters to achieve the same denoising performance than the U-Net and its parameters scale linearly in the image resolution instead of quadratically as for the original MLP-mixer. If trained on a moderate amount of examples for denoising, the image-to-image MLP-mixer outperforms the U-Net by a slight margin. It also outperforms the vision transformer tailored for image reconstruction and classical un-trained methods such as BM3D, making it a very effective tool for image reconstruction problems.

研究の動機と目的

• 畳み込みやマルチスケール構造を用いない、画像再構成に適した非畳み込み的・非マルチスケールのニューラルネットワークを構築し、最先端の畳み込みモデルと同等の性能を達成すること。
• 中規模のデータセットで学習された場合に、MLPベースのアーキテクチャが画像再構成で優れた性能を発揮できるかどうかを調査すること。
• パッチベースのMLPアーキテクチャにおいて、空間的位置情報の保持が画像再構成性能に与える影響を評価すること。
• 提案された画像対画像 MLP-mixer の性能とパラメータ効率を、U-Net やビジョントランスフォーマーや、古典的手法（BM3D など）と比較すること。

提案手法

• モデルは入力画像を重複のない P×P×3 のパッチに分割し、各パッチを次元 C の学習可能な埋め込み空間に線形変換することで、相対的な空間的位置を保持する。
• GeLU活性化関数を用いた共有された多層パーセプトロンを用いて、高さ、幅、チャネルの各次元で特徴を別々に混合する、変更を加えたMLP-mixer層を適用する。
• パッチの空間的順序を維持することで位置的インダクティブバイアスを組み込み、元の位置に依存しないMLP-mixerよりも性能が向上する。
• 標準的なMLP-mixerブロック構造を採用しているが、空間的構造を保持し、パッチ間で重みを共有することで、画像対画像変換に適応させた。
• Gaussianノイズや実世界のノイズ除去、圧縮センシングなど、画像再構成タスクのための標準的な最適化手法を用いて、エンドツーエンドで学習する。
• 元のMLP-mixerとは異なり、パramータ数が画像解像度に比例して線形に増加する。

実験結果

リサーチクエスチョン

• RQ1 畳み込みやマルチスケール設計を一切用いないMLP-mixerアーキテクチャが、画像再構成タスクで最先端の性能を達成できるか？
• RQ2 画像パッチの相対的な空間的位置を保持することで、MLPベースの画像再構成モデルの性能にどのような影響を与えるか？
• RQ3 中規模のデータセットで学習された場合に、提案された画像対画像 MLP-mixer が U-Net やビジョントランスフォーマーよりも優れた性能を発揮するか？
• RQ4 提案アーキテクチャのパラメータ効率と推論スループットは、U-Net や他のベースラインと比べてどうか？
• RQ5 畳み込みを用いないモデルが、膨大なトレーニングデータを必要とせずに、CNNと同等の性能を発揮できるか？

主な発見

• 合成ノイズと実世界ノイズの両方の除去タスクにおいて、画像対画像 MLP-mixer は U-Net よりわずかに優れており、SIDD データセットでは 33.66 dB を達成したのに対し、U-Net は 33.13 dB であった。
• Gaussianノイズ除去では 29.87 dB を達成し、U-Net（29.87 dB）、ViT（29.14 dB）、および元のMLP-mixer（27.36 dB）を上回った。
• 圧縮センシング MRI 再構成においては、すべてのモデルが同程度の性能を示し、特定のアーキテクチャに顕著な利点は認められなかった。
• U-Net と同等のノイズ除去性能を達成するためのパラメータ数が少なく、画像解像度に比例して線形にパラメータ数が増加する。
• スループット測定では、バッチサイズ1において DnCNN よりも高速で、ViT と同等の性能を示したが、U-Net や元のミキサーに比べて遅かった。
• 訓練データセットやモデルサイズが大きくなるにつれて、画像対画像 MLP-mixer と U-Net の性能差が拡大し、スケールが大きくなるほどデータ効率の向上が顕著に現れた。