[논문 리뷰] Image-to-Image MLP-mixer for Image Reconstruction

Youssef Mansour, Kang Lin외 1명|arXiv (Cornell University)|2022. 02. 04.

Image and Signal Denoising Methods인용 수 8

한 줄 요약

이 논문은 컨볼루션 또는 다중 해상도 구조를 사용하지 않고도 U-Net과 비전 트랜스포머보다 이미지 복원 작업(예: 노이즈 제거 및 압축 감지)에서 최신 기술 수준을 초월하는 이미지 간 MLP-믹서 아키텍처를 제안한다. 공간 패치 위치를 유지함으로써, 이미지 해상도에 비례하여 선형적으로 증가하는 파라미터 수를 가진 상태에서, 중간 크기의 데이터셋으로 훈련된 경우에도 최고의 성능을 달성한다.

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-믹서의 성능 및 파라미터 효율성을 U-Net, 비전 트랜스포머 및 BM3D와 같은 전통적 방법과 비교하는 것.

제안 방법

• 모델은 입력 이미지를 겹치지 않는 P×P×3 패치로 나누고, 각 패치를 차원 C의 학습 가능한 임bed딩 공간으로 선형 투영하며, 상대적인 공간 위치를 유지한다.
• GeLU 활성화 함수를 사용하는 공유 다층퍼셉트론을 활용해 높이, 너비, 채널 차원을 별도로 혼합하는 수정된 MLP-믹서 레이어를 적용한다.
• 패치의 공간 순서를 유지함으로써 위치 유도적 편향을 도입함으로써, 원래의 위치 무관 MLP-믹서보다 성능이 향상된다.
• 표준 MLP-믹서 블록 구조를 사용하지만, 공간 구조를 유지하고 패치 간에 공유 가중치를 사용함으로써 이미지 간 번역에 적합하게 변형한다.
• 이미지 복원 작업(예: 가우시안 노이즈 및 실제 세계 노이즈 제거, 압축 감지 포함)을 위한 표준 최적화 방법을 사용해 엔드 투 엔드로 훈련한다.
• 원래의 MLP-믹서와 달리, 파라미터 수가 이미지 해상도에 비례하여 선형적으로 증가한다.

실험 결과

연구 질문

• RQ1 컨볼루션 또는 다중 해상도 설계 없이도 MLP-믹서 아키텍처가 이미지 복원 작업에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ2 이미지 패치의 상대적인 공간 위치를 유지하는 것이 MLP 기반 이미지 복원 모델의 성능에 어떤 영향을 미치는가?
• RQ3 중간 크기의 데이터셋으로 훈련된 경우 제안된 이미지 간 MLP-믹서가 U-Net과 비전 트랜스포머를 능가하는가?
• RQ4 제안된 아키텍처의 파라미터 효율성과 추론 처리량은 U-Net 및 기타 기준 모델과 비교해 어떻게 되는가?
• RQ5 비컨볼루션 모델이 거대한 훈련 데이터셋이 필요 없이도 CNN과 경쟁할 수 있는가?

주요 결과

• 합성 및 실제 세계 노이즈 제거 작업에서 이미지 간 MLP-믹서는 U-Net을 약간 앞서며, SIDD 데이터셋에서 33.66 dB의 성능을 기록한 반면 U-Net은 33.13 dB을 기록했다.
• 가우시안 노이즈 제거에서 29.87 dB의 성능을 기록하여 U-Net(29.87 dB), ViT(29.14 dB), 원래의 MLP-믹서(27.36 dB)를 모두 능가했다.
• 압축 감지 MRI 복원 작업에서는 모든 모델이 유사한 성능을 보이며, 특정 아키텍처에 유의미한 이점이 없다는 것을 시사했다.
• 동일한 노이즈 제거 성능을 달성하기 위해 U-Net보다 더 적은 파라미터를 사용했으며, 이미지 해상도에 비례하여 파라미터 수가 선형적으로 증가했다.
• 처리량 측정 결과, 배치 크기가 1일 때 DnCNN보다 빠르고 ViT와 유사한 성능를 보였지만, U-Net과 원래의 믹서보다는 느렸다.
• 더 큰 훈련 데이터셋과 모델 크기에서 이미지 간 MLP-믹서와 U-Net 간의 성능 격차가 커지며, 스케일이 증가할수록 더 강한 데이터 효율성 향상 효과를 보였다.