[논문 리뷰] MAXIM: Multi-Axis MLP for Image Processing

Zhengzhong Tu, Hossein Talebi외 5명|arXiv (Cornell University)|2022. 01. 09.

Visual Attention and Saliency Detection인용 수 7

한 줄 요약

MAXIM는 이미지 처리를 위한 새로운 다축 MLP 아키텍처를 제안하며, 노이즈 제거, 왜곡 제거, 비 온 후 제거, 안개 제거 및 향상과 같은 다양한 저수준 시각 작업에서 최신 기술 성능을 달성한다. 다축 게이팅 MLP와 크로스 게이팅 블록을 통해 전체 컨볼루션형, 전역 수용장치 설계를 구현함으로써 기존 모델보다 더 적은 파라미터 수와 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}.

연구 동기 및 목표

• Transformer와 CNN의 저수준 시각 작업에서의 한계, 특히 고해상도 이미지 처리에 대한 유연성 부족과 제한된 수용장치를 해결하기 위해.
• 장거리 상호작용과 조밀한 예측 작업을 지원하는 유연하고 효율적이며 일반적인 목적의 비전 백본을 개발하기 위해.
• 자기주의 주의 메커니즘을 순수 MLP 기반 모듈로 대체하여 이미지 크기와 선형적으로 확장되면서도 전역적 맥락을 유지하기 위해.
• 다단계, 다스케일 아키텍처를 통해 노이즈 제거, 왜곡 제거, 비 온 후 제거, 안개 제거 및 이미지 향상과 같은 이미지 복원 작업의 성능을 향상시키기 위해.

제안 방법

• MAXIM는 다중 해상도, 다중 단계 감독을 위한 MAXIM 백본 스택을 포함한 UNet 구조의 인코더-디코더 아키텍처를 사용한다.
• 공간 격자와 국소 블록을 별도로 처리함으로써 전역 및 국소 특징 혼합을 가능하게 하는 다축 게이팅 MLP 블록을 도입한다. 이는 선형 복잡도를 유지한다.
• 공간 게이팅 유닛은 격자와 블록의 두 축을 따라 게이팅 MLP를 적용하여 장거리 의존성을 유지하고 전체 컨볼루션 연산을 지원한다.
• 크로스 게이팅 블록은 동일한 다축 접근 방식을 사용하여 스킵 커넥션을 적응적으로 조건화함으로써 크로스 어텐션을 대체하며, 전역적이고 파rameter 효율적인 특징 조절을 가능하게 한다.
• 다양한 해상도에서의 특징 학습을 향상시키기 위해 훈련 중에 다스케일 손실을 사용한다.
• 아키텍처는 자기주의 주의와 컨볼루션 인덕티브 바이어스를 모두 피하기 때문에, 조밀한 이미지 처리에 매우 뛰어난 유연성을 제공한다.

실험 결과

연구 질문

• RQ1 순수 MLP 기반 아키텍처가 저수준 이미지 복원 작업에서 어텐션 기반 모델을 능가할 수 있는가?
• RQ2 다축 MLP 설계가 이미지 크기와 선형 복잡도로 전역 맥락 모델링을 달성할 수 있는가?
• RQ3 완전히 컨볼루션형, MLP 기반 백본이 패치 또는 자르기 아티팩트 없이 고해상도 이미지 처리를 효과적으로 처리할 수 있는가?
• RQ4 이미지 복원에서 특징 상호작용을 모델링할 때 크로스 게이팅 메커니즘은 크로스 어텐션보다 어떻게 비교되는가?
• RQ5 통합적이고 경량화된 아키텍처가 다양한 이미지 처리 벤치마크에서 최신 기술 성능을 달성할 수 있는가?

주요 결과

• MAXIM은 노이즈 제거, 왜곡 제거, 비 온 후 제거, 안개 제거 및 향상 작업을 포함한 10개 이상의 벤치마크에서 최신 기술 성능을 달성한다.
• 이미지 노이즈 제거를 위한 SIDD 데이터셋에서, MAXIM은 이전 최고 성능 모델 대비 PSNR를 +0.24 dB 향상시켰다.
• GoPro에서의 이미지 왜곡 제거 작업에서, MAXIM은 이전 최고 성능 모델 대비 +0.15 dB의 PSNR 향상을 달성했다.
• Rain100L에서의 비 온 후 제거 작업에서, MAXIM은 +0.86 dB의 PSNR 향상을 기록했으며, 이는 이전 방법들을 크게 앞서는 성능이다.
• RESIDE에서의 이미지 안개 제거 작업에서, MAXIM은 +0.94 dB의 PSNR 향상을 기록했으며, 이는 복잡한 오염에 대한 강력한 성능을 보여준다.
• FiveK에서의 이미지 리터치 작업에서, MAXIM은 +1.15 dB의 PSNR 향상을 달성했으며, 이는 강력한 일반화 능력과 세부 사항 복원 능력을 시사한다.