[论文解读] 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

一句话总结

该论文提出了一种无需卷积或多分辨率结构的图像到图像MLP-mixer架构，在去噪和压缩感知等图像重建任务中表现优于U-Net和视觉Transformer（ViT）。通过保留空间块的位置信息，该模型在参数量更少且参数量与图像分辨率呈线性增长的情况下，实现了最先进性能，即使在中等规模数据集上训练也能取得优异结果。

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、视觉Transformer及经典方法（如BM3D）在性能和参数效率方面的表现。

提出的方法

• 该模型将输入图像划分为非重叠的P×P×3块，并将每个块线性投影到维度为C的可学习嵌入空间中，同时保留其相对空间位置。
• 采用改进的MLP-mixer层，通过共享多层感知机（使用GeLU激活函数）分别对高度、宽度和通道维度进行特征混合。
• 通过保持块的空间顺序，引入位置归纳偏置，从而在性能上优于原始的、不考虑位置信息的MLP-mixer。
• 网络采用标准的MLP-mixer块结构，但通过保留空间结构并跨块共享权重，将其适配为图像到图像翻译任务。
• 模型通过标准优化方法端到端训练，用于图像重建任务，包括高斯噪声和真实世界噪声去噪，以及压缩感知。
• 与原始MLP-mixer不同，该架构的参数量随图像分辨率呈线性增长，而非二次增长。

实验结果

研究问题

• RQ1 无卷积或多分辨率设计的MLP-mixer架构是否能在图像重建任务中实现最先进性能？
• RQ2 保留图像块之间的相对空间位置信息，对基于MLP的图像重建模型性能有何影响？
• RQ3 在中等规模数据集上训练时，所提出的图像到图像MLP-mixer是否优于U-Net和视觉Transformer？
• RQ4 与U-Net及其他基线方法相比，该架构的参数效率和推理吞吐量如何？
• RQ5 非卷积模型是否能在不依赖大规模训练数据集的情况下，与CNN在图像重建任务中竞争？

主要发现

• 在合成和真实世界去噪任务中，图像到图像MLP-mixer在性能上略优于U-Net，在SIDD数据集上达到33.66 dB，而U-Net为33.13 dB。
• 在高斯噪声去噪任务中，该模型达到29.87 dB，优于U-Net（29.87 dB）、ViT（29.14 dB）和原始MLP-mixer（27.36 dB）。
• 在压缩感知MRI重建任务中，所有模型表现相近，表明在该特定任务中无显著架构优势。
• 为达到与U-Net相当的去噪性能，该图像到图像MLP-mixer所需参数更少，且参数量与图像分辨率呈线性增长。
• 吞吐量测量显示，该模型在批量大小为1时比DnCNN更快，与ViT相当，但慢于U-Net和原始MLP-mixer。
• 随着训练数据集和模型规模增大，图像到图像MLP-mixer与U-Net之间的性能差距进一步扩大，表明其在大规模场景下具有更强的数据效率优势。