[论文解读] 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 和交叉门控模块，MAXIM 在参数量和 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 采用类似 UNet 的编码器-解码器架构，并通过堆叠 MAXIM 主干网络实现多尺度、多阶段的监督。
• 引入一种多轴门控 MLP 模块，分别处理空间网格和局部块，实现全局与局部特征混合，且复杂度为线性。
• 空间门控单元沿两个轴（网格和块）应用门控 MLP，以保留长距离依赖关系并支持全卷积操作。
• 交叉门控模块通过相同的多轴方法自适应地调节跳跃连接，替代交叉注意力机制，实现全局、参数高效的特征调制。
• 训练过程中使用多尺度损失，以增强不同分辨率下的特征学习能力。
• 整个架构完全基于 MLP 构建，避免自注意力和卷积归纳偏置，使其在密集图像处理任务中具有高度灵活性。

实验结果

研究问题

• RQ1 纯 MLP 架构能否在低层次图像恢复任务中超越基于注意力机制的模型？
• RQ2 多轴 MLP 设计能否以与图像尺寸线性相关的复杂度实现全局上下文建模？
• RQ3 全卷积的、基于 MLP 的主干网络能否在不引入分块或裁剪伪影的情况下有效处理高分辨率图像？
• RQ4 在图像恢复中，交叉门控机制与交叉注意力机制在建模特征交互方面有何对比？
• RQ5 统一的轻量化架构能否在多种图像处理基准上实现最先进性能？

主要发现

• MAXIM 在超过十个去噪、去模糊、去雨、去雾和增强任务的基准测试中均达到最先进性能。
• 在 SIDD 图像去噪数据集上，MAXIM 相较于之前的最先进方法，PSNR 提升了 +0.24 dB。
• 在 GoPro 图像去模糊任务中，MAXIM 相较于先前最佳模型，PSNR 提升了 +0.15 dB。
• 在 Rain100L 去雨数据集上，MAXIM 实现了 +0.86 dB 的 PSNR 提升，显著优于先前方法。
• 在 RESIDE 图像去雾数据集上，MAXIM 实现了 +0.94 dB 的 PSNR 提升，展现出在复杂退化情况下的强大性能。
• 在 FiveK 图像润饰任务中，MAXIM 实现了 +1.15 dB 的 PSNR 提升，表明其具备出色的泛化能力和细节恢复能力。