论文阅读：MAXIM Multi-Axis MLP for Image Processing

这是 2022 CVPR 上的一篇文章，介绍了用 MLP 做 low-level 图像处理的工作

Abstract

近年来，Transformer 和多层感知机（MLP）模型的发展为计算机视觉任务提供了新的网络架构设计。尽管这些模型在图像识别等许多视觉任务中已被证明是有效的，但在将它们应用于底层视觉任务时仍存在挑战。对高分辨率图像支持的灵活性不足以及局部注意力机制的局限性，可能是主要的瓶颈所在。在这项工作中，我们提出了一种基于多轴 MLP 的架构，称为 MAXIM，它可以作为一种高效且灵活的通用视觉主干架构，用于图像处理任务。MAXIM 采用了 U 型网络的分层结构，并通过空间门控的 MLP 实现了长距离的交互。具体来说，MAXIM 包含两个基于 MLP 的构建模块：一个多轴门控 MLP，它能够对局部和全局视觉线索进行高效且可扩展的空间混合；以及一个交叉门控模块，作为交叉注意力机制的替代方案，用于处理跨特征条件。这两个模块都完全基于 MLP，但同时也受益于其全局特性和 “全卷积” 特性，这两个特性在图像处理中是非常理想的。我们广泛的实验结果表明，所提出的 MAXIM 模型在十多个基准测试中，涵盖一系列图像处理任务，包括去噪、去模糊、去雨、去雾和增强等，均取得了当前最优的性能表现。而且，与其他具有竞争力的模型相比，它所需的参数数量和浮点运算次数（FLOPs）更少或相当。

Introduction

图像处理任务，如图像复原和增强，是重要的计算机视觉问题，其目标是从退化的输入图像中生成理想的输出图像。不同类型的图像退化可能需要不同的图像增强处理方法，例如去噪、去模糊、超分辨率、去雾、低光照增强等等。

鉴于经过精心整理的大规模训练数据集日益增多，近期基于精心设计的卷积神经网络（CNN）的高性能方法 [15, 17, 20, 22, 50, 52, 60, 61, 110, 111, 125] 在许多任务上展现出了当前最优（SOTA）的性能。

改进基础模型的架构设计是提升包括图像复原在内的大多数计算机视觉任务性能的关键因素之一。众多研究人员发明或借鉴了各种独立的模块或构建单元，并将它们应用于底层视觉任务中，这些应用涉及到残差学习 [43, 95, 120]、密集连接 [95, 121]、分层结构 [37, 41, 42]、多阶段框架 [16, 34, 111, 113] 以及注意力机制 [66, 91, 110, 111]。

最近对视觉 Transformer（ViT）[11, 24, 57] 的研究探索已经证明了它们作为常用卷积神经网络（CNN）模型替代方案的巨大潜力。ViT [24] 的出色表现也激发了类似的模型设计，这些设计采用了更简单的全局算子，比如 MLP-Mixer [87]、gMLP [54]、GFNet [76] 和 FNet [44] 等等。尽管这些模型在许多高级任务中取得了成功应用 [4, 24, 57, 85, 89, 102, 104]，但它们在底层图像增强和复原问题上的有效性尚未得到广泛研究。在将 Transformer 应用于底层视觉领域的开创性工作 [10, 15] 中，直接采用了全自注意力机制，而这种机制只接受相对较小且固定尺寸的图像块（例如 48×48）。当使用裁剪的方法将其应用于较大尺寸的图像时 [15]，这种策略不可避免地会导致图像块边界伪影的出现。基于局部注意力的 Transformer [52, 97] 改善了这一问题，但它们也受到限制，其感受野大小有限，或者会失去非局部性 [24, 93]，而非局部性是 Transformer 和 MLP 模型相对于分层 CNN 而言极具吸引力的一个特性。

为了克服这些问题，我们提出了一种通用的图像处理网络，名为 MAXIM，用于底层视觉任务。MAXIM 的一个关键设计要素是采用了多轴方法（第 3.2 节），该方法可以并行捕捉局部和全局的交互信息。通过在每个分支的单轴上混合信息，这种基于多层感知机（MLP）的算子变成了 “全卷积” 形式，并且在图像尺寸方面呈线性扩展，这显著提高了它在密集型图像处理任务中的灵活性。

我们还定义并构建了一个纯粹基于 MLP 的交叉门控模块，该模块使用相同的多轴方法，对 MAXIM 中间部分的跳跃连接进行自适应门控，从而进一步提升了性能。受近期图像复原模型的启发，我们开发了一种简单而有效的多阶段、多尺度架构，该架构由一系列 MAXIM 主干网络堆叠而成。

MAXIM 在一系列图像处理任务中都取得了出色的性能，同时所需的参数数量和浮点运算次数（FLOPs）极少。我们的贡献在于：

• 一种新颖且通用的图像处理架构，名为 MAXIM，采用了一系列编码器 - 解码器主干网络，并由多尺度、多阶段损失函数进行监督。
• 专为底层视觉任务量身定制的多轴门控 MLP 模块，该模块始终具备全局感受野，且相对于图像大小具有线性复杂度。
• 一个交叉门控模块，用于对两个独立的特征进行交叉条件处理，同样具备全局特性且为全卷积结构。
• 大量实验表明，MAXIM 在包括去噪、去模糊、去雨、去雾和图像增强等在内的 10 多个数据集上取得了当前最优的结果。

MAXIM

我们提出了首个用于底层视觉任务的有效通用多层感知机（MLP）架构，我们将其称为用于图像处理的多轴多层感知机（MAXIM）。与之前的底层视觉 Transformer 模型 [10,15,52,97] 不同，MAXIM 具有几个理想的特性，使其在图像处理任务中颇具吸引力。首先，MAXIM 能够以线性复杂度在任意大尺寸的图像上实现全局感受野；其次，它可以直接支持任意的输入分辨率，也就是说，它是全卷积的；最后，它提供了一种局部（卷积层）和全局（多层感知机）模块的平衡设计，在无需大规模预训练的情况下 [15]，性能超越了当前最优（SOTA）方法。

• 图 2 MAXIM 架构。我们采用（a）一个编解码器主干网络，其中每个（b）编码器、解码器以及瓶颈层都包含一个多轴门控多层感知器（MLP）模块（图 3）以及一个残差通道注意力模块。该模型通过（c）一个交叉门控模块得到进一步提升，这个交叉门控模块使得全局上下文特征能够对跳跃连接进行门控。更详细的描述见附录 A.2。

Main Backbone

MAXIM 主干网络（图 2a）遵循源自 U-Net [80] 的编码器 - 解码器设计原则。我们发现，像 3×3 卷积这样作用范围小的算子对类 U-Net 网络的性能至关重要。因此，我们对每个模块采用混合模型设计（图 2b）—— 使用卷积层处理局部信息，使用多层感知机（MLP）实现长距离交互 —— 以便充分发挥它们各自的优势。

为了在不同尺度上实现长距离的空间信息混合，我们将多轴门控多层感知机模块（MAB）插入到每个编码器、解码器以及瓶颈层中（图 2b），随后堆叠一个残差通道注意力模块（RCAB）[100, 111]（LayerNormConv-LeakyReLU-Conv-SE [31]）。
受跳跃连接的门控滤波机制 [67, 71] 的启发，我们对门控多层感知机（gMLP）进行扩展，构建了一个交叉门控模块（CGB，图 2c）。它是一种高效的二阶模块，可作为交叉注意力机制（三阶相关性）的替代方案，用于实现两种不同特征之间的交互或条件处理。我们利用来自瓶颈层（图 2a）的全局特征对跳跃连接进行门控，同时将经过优化的全局特征向上传播到下一个交叉门控模块。多尺度特征融合 [20, 84, 110]（红色和蓝色线条）被用于在编码器到交叉门控模块（Encoder $\to$ CGB）以及交叉门控模块到解码器（CGB $\to$ Decoder）的数据流中聚合多层次的信息。

Multi-Axis Gated MLP

我们的工作受到了文献 [123] 中提出的多轴分块自注意力机制的启发，该机制可在不止一个轴上执行注意力操作。在分块图像的两个轴上执行的注意力操作对应于两种形式的稀疏自注意力机制，即区域注意力和扩张注意力。尽管该模块能够并行捕捉局部和全局信息，但它无法适应图像复原或增强任务，因为在这些任务中，测试图像的尺寸往往是任意的。

我们对用于图像处理任务的 “多轴” 概念进行了改进，构建了一个（分头式）多轴门控多层感知机模块（MAB），如图 3 所示。与文献 [123] 中在单层应用多轴注意力机制不同，我们首先将头部分为两半，每一半独立进行划分。在局部分支中，尺寸为 $(H,W,C/2)$ 的特征的半个头被分块成形状为 $(\frac{H}{b}\times \frac{W}{b},b\times b,C/2)$ 的张量，这表示将其划分为大小为 $(b\times b)$ 的非重叠窗口；在全局分支中，另一半头部使用固定的 $(d\times d)$ 网格划分为形状为 $(d\times d,\frac{H}{d}\times \frac{W}{d},C/2)$ 的形式，每个窗口大小为 $(\frac{H}{d}\times \frac{W}{d})$。为便于可视化，在图 3 中我们设 $b=2$，$d=2$。为使其成为全卷积结构，我们仅在每个分支的单个轴上应用门控多层感知机（gMLP）模块 [54] —— 局部分支在第二个轴上应用，全局分支在第一个轴上应用 —— 同时在其他空间轴上共享参数。直观来讲，并行应用多轴 gMLP 分别对应着空间信息的局部和全局（扩张式）混合。最后，处理后的头部被拼接起来并进行投影以减少通道数，再通过来自输入的长跳跃连接进一步合并。值得注意的是，与处理固定大小图像块的方法 [15] 相比，这种方法为我们的模型带来了优势，避免了图像块边界伪影的问题。

复杂度分析。我们所提出的多轴门控多层感知机模块（MAB）的计算复杂度为。

$$\Omega (\text{MAB})=\underset{\text{Global gMLP}}{\underbrace{d^{2}HWC}}+\underset{\text{Local gMLP}}{\underbrace{b^{2}HWC}}+\underset{\text{Dense layers}}{\underbrace{10HW{C}^2}}\text{\hspace{1em}(1)}$$

这相对于图像尺寸 HW 呈线性关系，而像视觉 Transformer（ViT）、Mixer 和门控多层感知机（gMLP）等其他全局模型则呈二次方关系。

多轴方法的通用性。我们提出的并行多轴模块（图 3）提供了一种以可扩展方式在二维图像上应用一维算子的原则性方法。它还具有显著的灵活性和通用性。例如，直接将门控多层感知机（gMLP）替换为空间多层感知机 [87]、自注意力机制 [24]，甚至傅里叶变换 [44, 76]，就能得到一系列 MAXIM 变体，这些变体都具备全局性和全卷积性。它也很容易扩展到未来可能在语言模型等中定义的任何一维算子。

• 图 3 多轴门控多层感知器（MLP）模块。输入首先被投影为一个 [6; 4; C] 的特征，然后拆分为两个头。
  在局部分支中，半个头被划分成 3×2 个不重叠的 [2; 2; C/2] 的补丁块，而在全局分支中，我们使用 2×2 的网格划分另一半。我们仅在每个分支的单一轴上应用门控多层感知器（gMLP）模块 [54]（如右侧的 gMLP 模块所示）—— 局部分支应用于第二轴，全局分支应用于第一轴，同时在其他空间维度上共享。并行运行的 gMLP 操作符对应于局部和全局（扩张的）关注区域，如不同颜色所示（即，使用 gMLP 操作符对相同颜色的区域进行空间混合）。我们提出的模块在任意输入分辨率下都能表达全局和局部感受野。

Cross Gating MLP Block

对 U-Net 的一种常见改进方法是利用上下文特征，在跳跃连接中选择性地控制特征传播 [67,71]，这通常通过交叉注意力机制来实现 [13, 90]。在这里，我们构建了一种有效的替代方案，即交叉门控模块（CGB，图 2c），它是多轴门控多层感知机模块（MAB，3.2 节）的扩展，因为 MAB 只能处理单个特征。CGB 可以看作是一种更通用的条件层，能够与多种特征进行交互 [13, 70, 90]。我们采用了与 MAB 类似的设计模式。

更具体地说，设 $\mathbf{X},\mathbf{Y}$ 为两个输入特征，且 ${\mathbf{X}}_1,{\mathbf{Y}}_1\in {\mathbb{R}}^{H\times W\times C}$是图 2c 中经过第一个全连接层（Dense 层）投影后的特征。然后进行输入投影操作:

$${\mathbf{X}}_2=\sigma ({\mathbf{W}}_1\text{LN}({\mathbf{X}}_1)),{\mathbf{Y}}_2=\sigma ({\mathbf{W}}_2\text{LN}({\mathbf{Y}}_1))\text{\hspace{1em}(2)}$$

其中 $\sigma$ 是高斯误差线性单元（GELU）激活函数 [30] ，LN 是层归一化 [5] ，${\mathbf{W}}_1$、${\mathbf{W}}_2$是多层感知机（MLP）投影矩阵。多轴分块门控权重分别从 ${\mathbf{X}}_2$, ${\mathbf{Y}}_2$ 计算得出，但应用方式相反：

$$\hat{\mathbf{X}}={\mathbf{X}}_2\odot G({\mathbf{Y}}_2),\hat{\mathbf{Y}}={\mathbf{Y}}_2\odot G({\mathbf{X}}_2)\text{\hspace{1em}(3)}$$

其中 $\odot$ 表示逐元素相乘，函数 $G(\cdot )$ 使用我们提出的多轴方法从输入中提取多轴交叉门控权重。

$$G(\mathbf{x})={\mathbf{W}}_5([{\mathbf{W}}_3{\text{Block}}_b({\mathbf{z}}_1),{\mathbf{W}}_4{\text{Grid}}_d({\mathbf{z}}_2)])\text{\hspace{1em}(4)}$$

其中$[\cdot ,\cdot ]$表示拼接操作。这里$({\mathbf{z}}_1,{\mathbf{z}}_2)$是沿着通道维度从$\mathbf{z}$拆分出的两个独立头部，其中$\mathbf{z}$表示经过激活后的投影特征$\mathbf{x}$。

$$[{\mathbf{z}}_1,{\mathbf{z}}_2]=\mathbf{z}=\sigma ({\mathbf{W}}_6\text{LN}(\mathbf{x}))\text{\hspace{1em}(5)}$$

${\mathbf{W}}_3$、${\mathbf{W}}_4$是空间投影矩阵，分别应用于具有固定窗口大小 $b\times b$（${\text{Block}}_b$）的分块特征的第二个轴，以及具有固定网格大小$d\times d$（${\text{Grid}}_d$）的网格化特征的第一个轴。最后，我们采用输入的残差连接，在输出通道投影之后，使用投影矩阵${\mathbf{W}}_7$、${\mathbf{W}}_8$保持与输入（${\mathbf{X}}_1$、${\mathbf{Y}}_1$）相同的通道维度，记为:

$${\mathbf{X}}_3={\mathbf{X}}_1+{\mathbf{W}}_7\hat{\mathbf{X}},{\mathbf{Y}}_3={\mathbf{Y}}_1+{\mathbf{W}}_8\hat{\mathbf{Y}}\text{\hspace{1em}(6)}$$

Multi-Stage Multi-Scale Framework

我们进一步采用了多阶段框架，因为与增加模型宽度或高度相比（见 4.3A 节消融实验），我们发现这种方法更有效。我们认为全分辨率处理 [16, 69, 77] 比多图像块分层结构 [83, 111, 113] 更好，因为后者可能会在图像块之间产生边界效应。

为了实施更强的监督，我们在每个阶段都采用多尺度方法 [18, 20, 48] 来辅助网络学习。我们利用监督注意力模块 [111] 沿各个阶段逐步传播注意力特征。我们使用交叉门控模块进行跨阶段特征融合。具体细节请读者参考图 9。

• 图 9 我们采用了一种通用的多阶段框架，以提升 MAXIM 在具有挑战性的图像修复任务中的性能。受文献 [16, 111] 的启发，我们使用了监督注意力模块（SAM）以及跨阶段特征融合，来辅助后续阶段的学习。与先前的方法不同，由于所提出的多轴多层感知器（MLP）方法，我们的 MAXIM 主干网络在每个阶段的每一层都能实现全局感知，这使得它在学习低级和高级特征中的全局交互方面更加强大。

具体来说，给定一个输入图像 $I\in {\mathbb{R}}^{H\times W\times 3}$，我们首先通过下采样提取其多尺度变体：$I_n$，$n=1,\cdots ,N$ 。MAXIM 在S个阶段中的每个阶段s预测多尺度的复原输出，总共产生$S\times N$个输出: $R_{s,n}$ 。尽管 MAXIM 是多阶段的，但它是端到端训练的，损失会在各个阶段和尺度上累积:

$$\mathcal{L}=\sum _{s=1}^S\sum _{n=1}^N[{\mathcal{L}}_{char}({\mathbf{R}}_{s,n},{\mathbf{T}}_n)+\lambda {\mathcal{L}}_{freq}({\mathbf{R}}_{s,n},{\mathbf{T}}_n)]\text{\hspace{1em}(7)}$$

其中${\mathbf{T}}_n$表示（经过双线性缩放的）多尺度目标图像，${\mathcal{L}}_{char}$是 Charbonnier 损失函数 [111]：

$${\mathcal{L}}_{char}(\mathbf{R},\mathbf{T})=\sqrt{\Vert \mathbf{R}-\mathbf{T}{\Vert }^2+{ϵ}^2}\text{\hspace{1em}(8)}$$

这里我们将$ϵ$设为$10^{-3}$ 。${\mathcal{L}}_{freq}$是频率重建损失函数，用于强化高频细节 [20, 35]：
$${\mathcal{L}}_{freq}(\mathbf{R},\mathbf{T})=\Vert \mathcal{F}(\mathbf{R})-\mathcal{F}(\mathbf{T}){\Vert }_1\text{\hspace{1em}(9)}$$

其中$\mathcal{F}(\cdot )$表示二维快速傅里叶变换。在所有实验中，我们使用$\lambda =0.1$作为加权因子。

Experiments

我们旨在为广泛的图像处理任务构建一个通用的骨干网络。因此，我们在 17 个不同的数据集上针对五种不同的任务对 MAXIM 进行了评估：（1）去噪、（2）去模糊、（3）去雨、（4）去雾，以及（5）增强（润饰）（汇总于表 8 ）。

• 表 8

Experimental Setup

• 数据集与评估指标
  我们通过测量真实图像与预测图像之间的峰值信噪比（PSNR）和结构相似性指数（SSIM）[96] 来进行定量比较。去噪任务使用 SIDD [2] 和 DND [72] 数据集；去模糊任务使用 GoPro [62]、HIDE [81] 和 RealBlur [79] 数据集；去雨任务使用 [111] 中用到的 Rain13k 组合数据集；去雾任务使用 RESIDE [46] 数据集；图像增强（润饰）任务则在 Five - K [8] 和 LOL [98] 数据集上进行评估。
• 训练细节
  我们提出的 MAXIM 模型可进行端到端训练，既无需大规模预训练，也无需渐进式训练。网络在 256×256 的随机裁剪图像块上进行训练。针对不同任务，训练的迭代次数有所不同。数据增强方面，我们采用随机水平和垂直翻转、90° 旋转，以及概率为 0.5 的 MixUp [112] 方法。优化器使用 Adam [39]，初始学习率设为 2×10⁻⁴ ，并通过余弦退火衰减 [59] 逐步降至 10⁻⁷ 。测试时，我们使用对称填充将输入图像填充为 64×64 的倍数，推理完成后，再将填充后的图像裁剪回原始尺寸。关于各任务的更多训练细节见附录 A.1。
• 架构配置
  我们设计了两种 MAXIM 变体：用于不同任务的两阶段模型 MAXIM - 2S 和三阶段模型 MAXIM - 3S 。特征提取初始通道数设为 32，包含 3 个下采样层，特征从 256²×32 经 128²×64 、64²×128 ，收缩至 32²×256 ，并由两个瓶颈层处理（图 2a），然后对称扩展回全分辨率。当应用于 256×256 图像时，MAXIM - 2S 和 MAXIM - 3S 的参数量和所需浮点运算次数（FLOPs）列于表 7A 的最后两行。