MAXIM：用于图像处理的多轴 MLP

摘要

Transformer 和多层感知器 (MLP) 模型的最新进展为计算机视觉任务提供了新的网络架构设计。 尽管这些模型被证明在图像识别等许多视觉任务中是有效的，但将它们适应低级视觉仍然存在挑战。 支持高分辨率图像的不灵活性和局部注意力的限制可能是主要瓶颈。 在这项工作中，我们提出了一种名为 MAXIM 的基于多轴 MLP 的架构，它可以作为图像处理任务的高效灵活的通用视觉主干。 MAXIM 使用 UNet 形状的层次结构，并支持由空间门控 MLP 实现的远程交互。 具体来说，MAXIM 包含两个基于 MLP 的构建模块：一个多轴门控 MLP，允许局部和全局视觉线索进行高效且可扩展的空间混合；以及一个交叉门控模块，它是交叉注意力的替代方案，它考虑了交叉注意力-功能调节。 这两个模块都完全基于 MLP，但也受益于全局和“全卷积”，这是图像处理所需的两个属性。 我们广泛的实验结果表明，所提出的 MAXIM 模型在一系列图像处理任务（包括去噪、去模糊、去雨、去雾和增强）的十多个基准测试中实现了最先进的性能，同时需要更少或相当数量的参数和 FLOPs 优于竞争模型。 源代码和经过训练的模型可在 https://github.com/google-research/maxim 上获取。

1简介

图像处理任务，例如恢复和增强，是重要的计算机视觉问题，其目的是从降级的输入中产生所需的输出。 各种类型的退化可能需要不同的图像增强处理，例如去噪、去模糊、超分辨率、去雾、低光增强等。 鉴于精选的大规模训练数据集的可用性不断增加，最近的高性能方法[110, 111, 125, 50, 20, 15, 60, 61, 17, 52, 22]基于高度设计的卷积神经网络（CNN）在许多任务上展示了最先进的（SOTA）性能。

图1： 我们提出的 MAXIM 模型在五项图像处理任务的 PSNR 方面大大提高了最先进的性能：1）去噪（$+$0.24 dB，SIDD [2]）；2）去模糊（$+$0.15 dB on GoPro [62]) 3) 去毛刺（$+$0.86 dB on Rain100L [105]), 4) 去光斑（$+$0.94 dB，RESIDE [46]），以及 5) 润饰（增强）（$+$1.15 dB，FiveK [8]）。

改进底层模型的架构设计是提高大多数计算机视觉任务（包括图像恢复）性能的关键之一。 许多研究人员发明或借用了单独的模块或构建块，并将其实施到低级视觉任务中，包括残差学习[120, 43, 95]、密集连接[121, 95]，层次结构[42,41,37]，多阶段框架[111,16,34,113]，以及注意力机制[110 、111、66、91]。

最近对 Vision Transformers (ViT) [24,11,57] 的研究探索证明了它们作为首选 CNN 模型替代品的巨大潜力。 ViT [24] 的优雅也激发了使用更简单的全局运算符的类似模型设计，例如 MLP-Mixer [87]、gMLP [54]、GFNet [76] 和 FNet [44] 等。 尽管成功应用于许多高级任务[24,57,89,4,104,102,85]，这些全局模型在低级增强方面的功效且恢复问题尚未得到广泛研究。 低级视觉 Transformers 的开创性工作 [15, 10] 直接应用了完全自注意力，仅接受固定大小的相对较小的 patch（例如 48$\times$48 ）。 当使用裁剪[15]应用于较大图像时，这种策略将不可避免地导致补丁边界伪影。 基于局部注意力的 Transformer [52, 97] 改善了这个问题，但它们也受到感受野大小的限制，或者失去非局部性[93, 24]，这是 Transformer 和 MLP 模型相对于分层 CNN 的一个引人注目的特性。

为了克服这些问题，我们提出了一种通用图像处理网络，称为MAXIM，用于低级视觉任务。 MAXIM 的一个关键设计元素是使用多轴方法(秒 3.2)局部和全局相互作用并行。 通过在每个分支的单轴上混合信息，这种基于 MLP 的算子变成“全卷积”并根据图像大小线性缩放，这显着提高了其灵活性用于密集的图像处理任务。 我们还定义并构建了一个纯基于 MLP 的交叉选通模块，该模块使用相同的多轴方法自适应选通 MAXIM 颈部的跳跃连接，并进一步提高性能。 受最近恢复模型的启发，我们开发了一种简单但有效的多阶段、多尺度架构，由一堆 MAXIM 主干网络组成。 MAXIM 在一系列图像处理任务上实现了强大的性能，同时需要很少的参数和 FLOP。 我们的贡献是：

• 一种新颖且通用的图像处理架构，称为 MAXIM，使用一堆编码器-解码器主干，由多尺度、多阶段损失进行监督。

• 专为低级视觉任务量身定制的多轴门控 MLP 模块，始终享有全局感受野，并且具有相对于图像大小的线性复杂度。

• 交叉门控块，对两个独立的特征进行交叉条件处理，也是全局且全卷积的。

• 大量实验表明，MAXIM 在 10 多个数据集上取得了 SOTA 结果，包括去噪、去模糊、去雨、去雾和增强。

2相关工作

修复模型。 在最近建立视觉基准的巨大努力的推动下，基于学习的模型，特别是 CNN 模型，已经开发出来，在各种图像增强任务上实现了最先进的性能[111,16,37, 50、81、17、15、52]。 这些性能提升主要归功于新颖的架构设计和/或特定于任务的模块和单元。 例如，UNet [80] 已经孵化了许多成功的编码器-解码器设计[111, 37, 20]，用于图像恢复，这些设计改进了早期的单尺度特征处理模型[120, 45]。 为高级视觉任务开发的先进组件也已被引入低级视觉任务中。 残差和密集连接[120,43,95,121,95]，多尺度特征学习[41,97,20]，注意力机制[ 110, 111, 66, 91, 121] 和非本地网络[93, 53, 121] 就是很好的例子。 最近，多阶段网络[111,16,34,113]相对于上述单阶段模型在具有挑战性的去模糊和去雨任务[23,34,111]。 这些多阶段框架通常受到其在姿态估计 [18, 48]、动作分割 [25, 47] 和图像等更高级别问题上的成功的启发。一代[116, 117]。

低级视觉变形金刚。 Transformer 最初是针对 NLP 任务[90]提出的，其中多头自注意力和前馈 MLP 层堆叠在一起以捕获单词之间的非局部交互。 Dosovitskiy 等人.创造了术语Vision Transformer (ViT)[24]，并展示了第一个用于图像识别的纯Transformer模型。 最近的几项研究探讨了 Transformer 解决低级视觉问题的问题，例如。，开创性的预训练图像处理 Transformer (IPT)[15] 。 与 ViT 类似，IPT 直接将普通 Transformer 应用于图像补丁。 [10] 的作者提出了一种时空卷积自注意力网络，该网络利用局部信息实现视频超分辨率。 最近，Swin-IR [52] 和 UFormer [97] 在一系列图像恢复任务中应用了高效的基于窗口的局部注意力模型。

MLP 视觉模型。 最近，一些作者认为，当在 ViT 中使用基于补丁的架构时，复杂的自注意力机制的必要性变得值得怀疑。 例如，MLP-Mixer [87] 采用简单的 Token 混合 MLP 来代替 ViT 中的自注意力，从而形成全 MLP 架构。 [54]的作者提出了gMLP，它将空间门控单元应用于视觉标记。 ResMLP [88] 采用仿射变换代替层标准化来实现加速。 最近的技术，例如 FNet [44] 和 GFNet [76] 证明简单的傅立叶变换可以用作自注意力或 MLP 的竞争替代方案。

图2： 马克西姆架构。 我们采用 (a) 一个编码器-解码器主干，每个 (b) 编码器、解码器和瓶颈包含一个多轴门控 MLP 块(图 3 )以及剩余通道注意力块。 (c) 交叉门控模块进一步增强了该模型，该交叉门控模块允许全局上下文特征对跳跃连接进行门控。 更详细的描述可以在Sec.A.2中找到。

3我们的方法：MAXIM

据我们所知，我们提出了第一个用于低级视觉的有效通用 MLP 架构，我们称之为 Multi-AXIs M 用于图像处理的LP (MAXIM)。 与之前的低级 Transformer [15, 10, 52, 97] 不同，MAXIM 具有几个所需的属性，使其对于图像处理任务很有吸引力。 首先，MAXIM以线性复杂度表达任意大图像上的全局感受野；其次，它直接支持任意输入分辨率，即.，是全卷积的；最后，它提供了局部(Conv)和全局(MLP)块的平衡设计，优于 SOTA 方法，无需大规模预训练[ 15]。

3.1 主干网

MAXIM 主干网(图 2a）遵循源自 UNet [80] 的编码器-解码器设计原则>。 我们观察到，诸如 Conv3x3 等占用空间较小的算子对于类 UNet 网络的性能至关重要。 因此，我们依赖于每个块的混合模型设计(图 2b）-本地的Conv ，以及用于远程交互的MLP——充分利用它们。

为了允许不同尺度的远程空间混合，我们将多轴门控 MLP 块 (MAB) 插入每个编码器、解码器和瓶颈中(图 2b)，带有残余通道注意块 (RCAB) [100, 111] (LayerNorm-Conv-LeakyReLU-Conv-SE [31])随后堆叠。 受跳跃连接[67, 71]的门控过滤的启发，我们扩展门控MLP（gMLP）来构建交叉门控块（CGB，图) 2c)，这是交叉注意力（三阶相关性）的有效二阶替代方案，用于交互或调节两个不同的特征。 我们利用瓶颈(图2a）的全局特征来门控跳跃连接，同时传播将全局功能向上细化到下一个 CGB。 多尺度特征融合[84,110,20](红色和蓝色线）用于聚合编码器中的多级信息$\to$CGB 和 CGB$\to$解码器数据流。

3.2多轴门控 MLP

图3： 多轴门控 MLP 模块 （彩色效果最佳）。 输入首先被投影到 $[6,4,C]$ 特征，然后分成两个头。 在本地分支中，半头被阻挡成$3\times 2$不重叠的$[2,2,C/2]$补丁，而我们网格 另一半在全局分支中使用$2\times 2$网格。 我们仅在每个分支的单个轴上应用gMLP块[54]（如右侧gMLP块所示） - 局部分支的第二轴和全局分支的第一轴，同时沿其他空间维度共享。 并行运行的 gMLP 运算符对应于局部和全局（扩张）参与区域，如不同颜色所示(即。，相同颜色使用空间混合gMLP 算子）。 我们提出的块表达了任意输入分辨率的全局和局部感受野。

我们的工作受到[123]中提出的多轴阻塞自注意力的启发，它在多个单轴上进行注意力。 在块图像上的两个轴上执行的注意力对应于两种形式的稀疏自注意力，即区域注意力和扩张注意力。 尽管并行捕获局部和全局信息，但该模块无法适应测试图像通常具有任意大小的图像恢复或增强任务。

我们改进了图像处理任务中的"多轴"概念，构建了一个（分头）多轴门控 MLP 块（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$。 为了使其完全卷积，我们仅在每个分支的单个轴上应用门控MLP（gMLP）块[54] - $2$nd axis 用于本地分支，$1$st axis 用于全局分支 - 同时在其他空间上共享参数轴。 直观地说，并行应用多轴 gMLP 分别对应于空间信息的局部和全局（扩张）混合。 最后，将处理后的头连接并投影以减少通道数量，并使用输入的长跳跃连接进一步组合通道数量。 值得注意的是，与处理固定大小图像补丁 [15] 的方法相比，这种方法通过避免补丁边界伪影为我们的模型提供了优势。

复杂性分析。 我们提出的多轴 gMLP 模块 (MAB) 的计算复杂度为：

$$\mathrm{\Omega }\left(\mathrm{MAB}\right)=\underset{\mathrm{Global}\mathrm{gMLP}}{\underbrace{d^{2}HWC}}+\underset{\mathrm{Local}\mathrm{gMLP}}{\underbrace{b^{2}HWC}}+\underset{\mathrm{Dense}\mathrm{layers}}{\underbrace{10HW{C}^2}},$$ (1)

相对于图像大小 $HW$ 是线性，而 ViT、Mixer 和 gMLP 等其他全局模型是二次。

多轴方法的通用性。 我们提出的并行多轴模块(图 3)提出了一种以可扩展的方式在 2D 图像上应用 1D 算子的原理方法。 它还具有显着的灵活性和通用性。 例如，用空间 MLP [87]、自注意力 [24] 甚至傅里叶变换 [76, 44] 直接替换 gMLP 导致一系列 MAXIM 变体（参见Sec. 4.3D），所有变体都共享全局性和全卷积性。 它还可以轻松扩展到任何可能在例如.语言模型上定义的未来一维运算符。

3.3 交叉门控MLP模块

相对于 UNet 的一个常见改进是利用上下文特征来选择性地gate跳过连接[67, 71]中的特征传播，这通常是通过使用交叉注意力[90, 13]。 在这里，我们构建了一个有效的替代方案，即交叉门控块（CGB，图 2c），作为 MAB 的扩展(Sec. 3.2)，只能处理单个特征。 CGB 可以被视为一个更通用的调节层，与多个特征[90,70,13]交互。 我们遵循与 MAB 中使用的类似设计模式。

更具体地说，令$𝐗,𝐘$为两个输入特征，${𝐗}_1,{𝐘}_1\in {ℝ}^{H\times W\times C}$为图中第一个Dense层之后投影的特征. 2c．然后应用输入投影：

$${𝐗}_2=\sigma \left({𝐖}_1\mathrm{𝖫𝖭}\left({𝐗}_1\right)\right),{𝐘}_2=\sigma \left({𝐖}_2\mathrm{𝖫𝖭}\left({𝐘}_1\right)\right)$$ (2)

其中 $\sigma$ 是 $\mathrm{𝖦𝖤𝖫𝖴}$ 激活 [30],$\mathrm{𝖫𝖭}$ 是层归一化 [5]， ${𝐖}_1,{𝐖}_2$ 是 MLP 投影矩阵。 多轴分块门控权重分别根据 ${𝐗}_2,{𝐘}_2$ 计算，但相反应用：

$$\widehat{𝐗}={𝐗}_2\odot \mathrm{G}\left({𝐘}_2\right),\widehat{𝐘}={𝐘}_2\odot \mathrm{G}\left({𝐗}_2\right)$$ (3)

其中 $\odot$ 表示逐元素乘法，函数 $\mathrm{G}(\cdot )$ 使用我们提出的多轴方法从输入中提取多轴交叉选通权重(秒）。 3.2)：

$$\mathrm{G}\left(𝐱\right)={𝐖}_5\left([{𝐖}_3{\mathrm{𝖡𝗅𝗈𝖼𝗄}}_b\left({𝐳}_{\mathrm{𝟏}}\right),{𝐖}_4{\mathrm{𝖦𝗋𝗂𝖽}}_d\left({𝐳}_{\mathrm{𝟐}}\right)]\right)$$ (4)

其中 $[\cdot ,\cdot ]$ 表示串联。 这里 $({𝐳}_{\mathrm{𝟏}},{𝐳}_{\mathrm{𝟐}})$ 是沿着通道维度从 $𝐳$ 分裂出来的两个独立头，其中 $𝐳$ 表示激活后的投影特征 $𝐱$：

$$[{𝐳}_{\mathrm{𝟏}},{𝐳}_{\mathrm{𝟐}}]=𝐳=\sigma \left({𝐖}_6\mathrm{𝖫𝖭}\left(𝐱\right)\right),$$ (5)

和 ${𝐖}_3,{𝐖}_4$ 是应用于具有固定窗口大小 $b\times b$ 的块/网格特征的 2nd 和 1st 轴的空间投影矩阵(${\mathrm{𝖡𝗅𝗈𝖼𝗄}}_b$) 和固定网格大小 $d\times d$ (${\mathrm{𝖦𝗋𝗂𝖽}}_d$)。 最后，我们采用来自输入的残差连接，遵循输出通道投影，保持与输入相同的通道尺寸 (${𝐗}_1,{𝐘}_1$)，使用投影矩阵 ${𝐖}_7$、${𝐖}_8$，表示为

$${𝐗}_3={𝐗}_1+{𝐖}_7\widehat{𝐗},{𝐘}_3={𝐘}_1+{𝐖}_8\widehat{𝐘}.$$ (6)

CGB 的复杂性也受到 Eq. 1 的严格限制。

3.4多阶段多尺度框架

我们进一步采用多阶段框架，因为我们发现它比扩大模型宽度或高度更有效（参见 ablation Sec. 4.3A）。 我们认为全分辨率处理[69,77,16]是比多面片层次结构[83,111,113]更好的方法，因为后者可能会导致边界跨补丁的影响。 为了实施更强的监督，我们在每个阶段应用多尺度方法[48,18,20]来帮助网络学习。 我们利用监督注意力模块[111]沿着阶段逐步传播注意力特征。 我们利用跨门块(Sec. 3.3)进行跨阶段特征融合。 详细信息请读者参考图9。

	SIDD [2]		DND [72]		Average
Method	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$
DnCNN [120]	23.66	0.583	32.43	0.790	28.04	0.686
MLP [7]	24.71	0.641	34.23	0.833	29.47	0.737
BM3D [21]	35.65	0.685	34.51	0.851	35.08	0.768
CBDNet* [29]	30.78	0.801	38.06	0.942	34.42	0.872
RIDNet* [3]	38.71	0.951	39.26	0.953	38.99	0.952
AINDNet* [38]	38.95	0.952	39.37	0.951	39.16	0.952
VDN [107]	39.28	0.956	39.38	0.952	39.33	0.954
SADNet* [12]	39.46	0.957	39.59	0.952	39.53	0.955
CycleISP* [109]	39.52	0.957	39.56	0.956	39.54	0.957
MIRNet [110]	39.72	0.959	39.88	0.956	39.80	0.958
MPRNet [111]	39.71	0.958	39.80	0.954	39.76	0.956
MAXIM-3S	39.96	0.960	39.84	0.954	39.90	0.957

表1： 去噪结果。 我们的模型仅在 SIDD [2] 上进行训练，并在 SIDD [2] 和 DND [72] 上进行评估，其中 * 表示使用附加训练数据的方法。

形式上，给定输入图像$𝐈\in {ℝ}^{H\times W\times 3}$，我们首先通过缩小尺寸来提取其多尺度变体：${𝐈}_n,n=1,\mathrm{\dots },N$。 MAXIM 预测 $S$ 阶段中每个阶段 $s$ 的多尺度恢复输出，总共产生 $S\times N$ 输出：${𝐑}_{s,n}$。 尽管是多阶段的，MAXIM 仍进行端到端训练，损失在各个阶段和规模上累积：

$$ℒ=\sum _{s=1}^S\sum _{n=1}^N\left[{ℒ}_{char}({𝐑}_{s,n},{𝐓}_n)+\lambda {ℒ}_{freq}({𝐑}_{s,n},{𝐓}_n)\right],$$ (7)

其中 ${𝐓}_n$ 表示（双线性重新缩放）多尺度目标图像，${ℒ}_{char}$ 是 Charbonnier 损失 [111]：

$${ℒ}_{char}(𝐑,𝐓)=\sqrt{{\Vert 𝐑-𝐓\Vert }^2+{ϵ}^2},$$ (8)

我们设置$ϵ={10}^{-3}$的地方。 ${ℒ}_{freq}$ 是加强高频细节的频率重建损失[35, 20]：

$${ℒ}_{freq}(𝐑,𝐓)={\Vert ℱ\left(𝐑\right)-ℱ\left(𝐓\right)\Vert }_1$$ (9)

其中 $ℱ(\cdot )$ 表示 2D 快速傅立叶变换。 我们在所有实验中使用$\lambda =0.1$作为权重因子。

4实验

我们的目标是为广泛的图像处理任务构建通用主干。 因此，我们在 17 不同数据集上评估了五个不同任务上的 MAXIM：(1) 去噪、(2) 去模糊、(3) 去雨、(4) 去雾和 (5) 增强（修饰）（总结于 选项卡。 8. 更全面的结果和可视化可以在Sec.A.6中找到。

4.1 实验设置

数据集和指标。 我们测量了真实图像和预测图像之间的 PSNR 和 SSIM [96] 指标，以进行定量比较。 我们使用SIDD [2]和DND [72]进行去噪，GoPro [62]，HIDE [81] 和 RealBlur [79] 用于去模糊，组合数据集 Rain13k 用于 [111] 中用于去雨。 RESIDE [46] 用于去雾，而 Five-K[8] 和 LOL [98] 用于评估增强。

Input	Target	VDN [107]	DANet [108]	MIRNet [110]	CycleISP [109]	MPRNet [111]	MAXIM-3S

图4： 去噪比较。 SIDD [2] 中的示例表明我们的方法可以产生更清晰的去噪结果。

	GoPro [62]		HIDE [81]		Average
Method	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$
DeblurGAN [40]	28.70	0.858	24.51	0.871	26.61	0.865
Nah et al. [62]	29.08	0.914	25.73	0.874	27.41	0.894
Zhang et al. [118]	29.19	0.931	-	-	-	-
DeblurGAN-v2 [41]	29.55	0.934	26.61	0.875	28.08	0.905
SRN [86]	30.26	0.934	28.36	0.915	29.31	0.925
Shen et al. [81]	-	-	28.89	0.930	-	-
Gao et al. [28]	30.90	0.935	29.11	0.913	30.01	0.924
DBGAN [119]	31.10	0.942	28.94	0.915	30.02	0.929
MT-RNN [69]	31.15	0.945	29.15	0.918	30.15	0.932
DMPHN [113]	31.20	0.940	29.09	0.924	30.15	0.932
Suin et al. [83]	31.85	0.948	29.98	0.930	30.92	0.939
MPRNet [111]	32.66	0.959	30.96	0.939	31.81	0.949
Pretrained-IPT [15]	32.58	-	-	-	-	-
MIMO-UNet+ [20]	32.45	0.957	29.99	0.930	31.22	0.944
HINet [16]	32.71	0.959	30.32	0.932	31.52	0.946
MAXIM-3S	32.86	0.961	32.83	0.956	32.85	0.959

表2： 去模糊结果。 我们的模型在 GoPro [62] 上进行训练，并在 GoPro 和 HIDE 数据集 [81] 上进行评估。

	RealBlur-R [79]		RealBlur-J [79]		Average
Method	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$	PSNR$\uparrow$	SSIM$\uparrow$
Hu et al. [33]	33.67	0.916	26.41	0.803	30.04	0.860
Nah et al. [62]	32.51	0.841	27.87	0.827	30.19	0.834
DeblurGAN [40]	33.79	0.903	27.97	0.834	30.88	0.869
Pan et al. [68]	34.01	0.916	27.22	0.790	30.62	0.853
Xu et al. [103]	34.46	0.937	27.14	0.830	30.8	0.884
DeblurGAN-v2 [41]	35.26	0.944	28.70	0.866	31.98	0.905
Zhang et al. [118]	35.48	0.947	27.80	0.847	31.64	0.897
SRN [86]	35.66	0.947	28.56	0.867	32.11	0.907
DMPHN [113]	35.70	0.948	28.42	0.860	32.06	0.904
MPRNet [111]	35.99	0.952	28.70	0.873	32.35	0.913
MAXIM-3S	35.78	0.947	28.83	0.875	32.31	0.911
†DeblurGAN-v2	36.44	0.935	29.69	0.870	33.07	0.903
†SRN [86]	38.65	0.965	31.38	0.909	35.02	0.937
†MPRNet [111]	39.31	0.972	31.76	0.922	35.54	0.947
†MIMO-UNet+ [20]	-	-	32.05	0.921	-	-
†MAXIM-3S	39.45	0.962	32.84	0.935	36.15	0.949

表3： RealBlur [79] 上的去模糊结果。 ^† 表示在 RealBlur 上训练的方法，而没有 ^† 的方法表示仅在 GoPro 上训练的方法。

Input	Target	DMPHN [113]	Suin et al. [83]	MPRNet [111]	HINet [16]	MIMO-UNet [20]	MAXIM-3S

图5： 去模糊比较。 第一行显示来自 GoPro [62] 的示例，而第二行显示来自 HIDE [81] 的示例。

训练细节。 我们提出的 MAXIM 模型是端到端可训练的，既不需要大规模预训练，也不需要渐进训练。 网络在 $256\times 256$ 随机裁剪的补丁上进行训练。 我们为每个任务训练不同的迭代。 我们使用随机水平和垂直翻转、${90}^{\circ }$ 旋转以及概率 $0.5$ 的 MixUp [112] 来进行数据增强。 我们使用 Adam 优化器 [39] ，初始学习率为 $2\times {10}^{-4}$，随着余弦退火衰减 [ 稳步下降到 ${10}^{-7}$ 59]。 测试时，我们使用两侧对称填充将输入图像填充为 $64\times 64$ 的倍数。 经过推理，我们将填充后的图像裁剪回原始大小。 有关每项任务的更多训练详细信息，请参阅Sec.A.1。

建筑配置。 我们设计了两种 MAXIM 变体：两级模型 MAXIM-2S 和三级模型 MAXIM-3S，用于不同的任务。 我们从 $32$ 初始通道开始进行特征提取，具有 3 个下采样层，其中特征从 ${256}^2\times 32$、${128}^2\times 64$、${64}^2\times 128$ 收缩，到${32}^2\times 256$由两个瓶颈处理(图2a)，然后对称扩展回全分辨率。 当应用于$256\times 256$图像时，MAXIM-2S和MAXIM-3S的参数数量和所需的FLOPs显示在最后两行 选项卡。 7A。

4.2 主要结果

去噪。 我们报道 选项卡。 1 SIDD [2] 和 DND [72] 数据集上的数值比较。 可以看出，我们的方法比之前的 SOTA 技术(例如.、MIRNet [110] 性能优于 0.24 dB SIDD 上的 PSNR，同时在 DND 上获得有竞争力的 PSNR (39.84 dB)。 图 4显示了SIDD上的视觉结果。 我们的方法清楚地消除了真实的噪音，同时保留了精细的细节，与其他方法相比，产生了视觉上令人愉悦的结果。

去模糊。 选项卡。 2 显示了 MAXIM-3S 与 SOTA 去模糊方法在两个合成模糊数据集上的定量比较：GoPro [62] 和 HIDE [81]。 与之前的最佳模型 HINet [16] 相比，我们的方法在 PSNR 方面实现了 0.15 dB 增益。 值得注意的是，GoPro 训练的 MAXIM-3S 模型在 HIDE 数据集上泛化得非常好，设置了新的 SOTA PSNR 值：32.83 dB。 我们还在两种设置下对来自 RealBlur [79] 的真实世界模糊图像进行了评估：(1) 直接在 RealBlur 上应用 GoPro 训练的模型，以及 (2) 在 RealBlur 上微调模型。 在设置(1)下，MAXIM-3S在RealBlur-J子集上排名第一，同时在RealBlur-R上获得前两名的性能。 图 5显示了 GoPro [62]、HIDE [81 上评估模型的视觉比较] 和 RealBlur [79] 分别。 可以观察到，我们的模型恢复文本的效果非常好，这可能归因于在每个块中使用多轴 MLP 模块，该模块全局聚合了不同尺度的重复模式。

Input	Target	RESCAN [49]	PreNet [77]	MSPFN [34]	MPRNet [111]	HINet [16]	MAXIM-2S

图6： 消除比较。 顶行和底行分别展示了来自 Rain100L [105] 和 Test100 [115] 的示例，展示了 MAXIM 消除雨纹同时恢复更多细节的能力，从而产生视觉上更令人愉悦的结果。

Input	Target	GCANet [14]	GridDehaze [55]	DuRN [56]	MSBDN [23]	FFA-Net [74]	MAXIM-2S

图7： 去雾比较。 顶行和底行示例了 SOTS 室内和室外场景 [46] 的视觉结果。

Rain100L [105] Rain100H [105] Test100 [115] Test1200 [114] Test2800 [27] Average Method PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ DerainNet [26] 27.03 0.884 14.92 0.592 22.77 0.810 23.38 0.835 24.31 0.861 22.48 0.796 SEMI [99] 25.03 0.842 16.56 0.486 22.35 0.788 26.05 0.822 24.43 0.782 22.88 0.744 DIDMDN [114] 25.23 0.741 17.35 0.524 22.56 0.818 29.65 0.901 28.13 0.867 24.58 0.770 UMRL [106] 29.18 0.923 26.01 0.832 24.41 0.829 30.55 0.910 29.97 0.905 28.02 0.880 RESCAN [49] 29.80 0.881 26.36 0.786 25.00 0.835 30.51 0.882 31.29 0.904 28.59 0.857 PreNet [77] 32.44 0.950 26.77 0.858 24.81 0.851 31.36 0.911 31.75 0.916 29.42 0.897 MSPFN [34] 32.40 0.933 28.66 0.860 27.50 0.876 32.39 0.916 32.82 0.930 30.75 0.903 MPRNet [111] 36.40 0.965 30.41 0.890 30.27 0.897 32.91 0.916 33.64 0.938 32.73 0.921 HINet [16] 37.20 0.969 30.63 0.893 30.26 0.905 33.01 0.918 33.87 0.940 33.00 0.925 MAXIM-2S 38.06 0.977 30.81 0.903 31.17 0.922 32.37 0.922 33.80 0.943 33.24 0.933 Table 4: Deraining comparisons. Our method consistently yields better quality metrics with respect to both PSNR or SSIM on all the tested datasets: Rain100L [105], Rain100H [105], Test100 [115], Test1200 [114], Test2800 [27]

SOTS-Indoor SOTS-Outdoor Method PSNR$\uparrow$ SSIM$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$ DehazeNet [9] 21.14 0.847 22.46 0.851 GFN [78] 22.30 0 880 21.55 0.844 GCANet [14] 30.23 0.959 19.98 0.704 GridDehaze [55] 32.14 0.983 30.86 0.981 GMAN [58] 27.93 0.896 28.47 0.944 MSBDN [23] 33.79 0.984 23.36 0.875 DuRN [56] 32.12 0.980 24.47 0.839 FFA-Net [74] 36.39 0.989 33.57 0.984 AECR-Net [101] 37.17 0.990 - - MAXIM-2S 38.11 0.991 34.19 0.985 Table 5: Dehazing comparisons. Our model achieved the best results on both indoor and outdoor scenes.

Input	Target	CycleGAN [124]	Exposure [32]	DPE [19]	EnlightenGAN	UEGAN [65]	MAXIM-2S
Input	Target	Retinex [98]	GLAD [92]	KinD [122]	EnlightenGAN	MIRNet [110]	MAXIM-2S

图8： 修饰和低光增强比较。 顶行显示了来自 MIT-Adobe FiveK 数据集 [8] 的示例，而底行则举例说明了来自 LOL [98] 的比较。 我们的模型生成了杂色且更自然的颜色（顶部）进行修饰，同时在底部示例中实现了更清晰、更明亮的视觉增强。

除雨。 继之前的工作[34, 111]之后，我们使用 Y 通道（在 YCbCr 颜色空间中）计算性能指标。 选项卡。 5 显示与以前的方法的定量比较。 可以看出，我们的模型在所有数据集上的 SOTA 性能都有所提高。 我们的模型相对于之前最佳模型 HINet [16] 的平均 PSNR 增益为 0.24 dB。 我们在图6中展示了一些具有挑战性的示例，这表明我们的方法始终如一地提供了忠实恢复的图像，而没有引入任何明显的视觉伪影。

去雾。 我们报告了与 SOTA 模型的比较 选项卡。 5. 我们的模型在 SOTS [46] 室内和室外设备上的 PSNR 超过了之前的最佳模型 0.94 dB 和 0.62 dB。 图 7显示我们的模型在平坦区域和纹理上恢复了更好质量的图像，同时实现了和谐的全局色调。

增强/修饰。 作为 选项卡。 6 如图所示，我们的模型分别在 FiveK [8] 和 LOL [98] 上实现了最佳 PSNR 和 SSIM 值。 正如图 8顶行所示，与其他技术相比，MAXIM 恢复了多种自然色彩。 关于底部的示例，虽然 MIRNet [110] 获得了更高的 PSNR，但我们始终观察到我们的模型获得了更好的视觉质量，细节更清晰，噪声更少。 此外，感知相关性更高的 SSIM 指数表明 MAXIM-2S 相对于 MIRNet 具有显着优势。

其他基准。 由于篇幅限制，我们在秒中详细介绍了 REDS 去模糊[63]和雨滴去除任务[73]的实验结果。 A.5。

4.3消融

我们进行了广泛的消融研究，以验证所提出的多轴门控 MLP 模块、交叉门控模块和多级多尺度架构。 评估是在 GoPro 数据集 [62] 上进行的，该数据集是在大小为 $256\times 256$ 的图像块上进行 ${10}^6$ 次迭代训练的。 我们使用 MAXIM-2S 模型作为 Ablation-A 和 -B 的测试平台。

A。 单独的组件。 我们通过逐步添加 (1) 级间交叉门控模块 (CGB_IS)、(2) 监督注意力模块 (SAM)、(3) 跨级交叉门控模块来进行消融（CGB_CS，以及（4）多尺度监督（MS-Sp）。 选项卡。 7A 表示每个相应组件的 PSNR 增益为 0.25、0.63、0.36、0.26 dB。

	FiveK [8]			LOL [98]
Method	PSNR$\uparrow$	SSIM$\uparrow$	Method	PSNR$\uparrow$	SSIM$\uparrow$
CycleGAN [124]	18.23	0.835	Retinex [98]	16.77	0.559
Exposure [32]	22.35	0.861	GLAD [92]	19.71	0.703
EnlightenGAN	17.74	0.828	EnlightenGAN	17.48	0.657
DPE [19]	24.08	0.922	KinD [122]	20.37	0.804
UEGAN [65]	25.00	0.929	MIRNet [110]	24.14	0.830
MAXIM-2S	26.15	0.945	MAXIM-2S	23.43	0.863

表6： FiveK [8] 和 LOL [98] 的增强结果。

B. 多轴方法的效果。 我们进一步研究了我们提出的多轴方法的必要性，如图所示 选项卡。 7B. 我们在（1）基线 UNet、（2）通过添加 MAB 的本地分支（MAB_{$\mathrm{\ell }$})、（3）通过添加 MAB 的全局分支（MAB_$g$ ），（4）添加CGB本地分支（CGB_{$\mathrm{\ell }$})，（5）添加CGB全局分支（CGB_$g$)。 请注意，通过添加 MAB_{$\mathrm{\ell }$}，PSNR 的巨大跳跃（+1.04 dB）很大程度上归因于输入和输出通道投影层的添加，因为我们还观察到 31.42< 的高性能/t2> dB PSNR（如果仅添加 MAB_$g$)。 总体而言，我们观察到添加 MAB 时取得了重大改进，而添加 CGB 时则获得了相对较小的改进。

C. 为什么要多级？ 为了理解这一点，我们在宽度（通道）、深度（缩小步骤）和级数方面扩大了 MAXIM。 选项卡。 7C 表明，与使其更宽或更深相比，将主干打包成多级可产生最佳性能vs。复杂性权衡（32.44 dB、22.2 M、339.2 G） 。

D. 超越 gMLP：MAXIM 系列。 如Sec. 3.2中所述，我们提出的多轴方法(图 3)提供了一种在（高分辨率）图像上应用任何一维运算符的可扩展方法，具有相对于图像大小的线性复杂性，同时保持完全卷积。 我们在 SIDD [2] 上使用 MAXIM-1S 和 -2S 进行了试点研究，以探索 MAXIM 系列：MAXIM-FFT、-MLP ，-gMLP（在本文中建模），-SA，我们使用傅里叶变换滤波器 [76, 44]，空间 MLP [87]，gMLP [54]，以及使用相同多轴方法在空间轴上的自注意力[24](图 3)。 作为 选项卡。 7D 显示，gMLP 和自注意力变体实现了最佳性能，而 FFT 和 MLP 系列的计算效率更高。 我们将更深入的探索留给未来的工作。

CGBIS SAM CGBCS MS-Sp PSNR 30.73 ✓ 30.98 ✓ ✓ 31.61 ✓ ✓ ✓ 31.97 ✓ ✓ ✓ ✓ 32.23 A. Individual components. MAB$\mathrm{\ell }$ MAB$g$ CGB$\mathrm{\ell }$ CGB$g$ PSNR 30.48 ✓ 31.52 ✓ ✓ 31.68 ✓ ✓ ✓ 31.84 ✓ ✓ ✓ ✓ 31.91 B. Effects of multi-axis approach.

S W D PSNR Params FLOPs Base 1 32 3 31.08 6.1M 93.6G Wider 1 64 3 32.09 19.4M 309.9G 1 96 3 32.31 41.7M 648.9G Deeper 1 32 4 31.17 19.8M 121.6G 1 32 5 31.43 75.0M 153.4G More 2 32 3 31.82 14.1M 216.4G stages 3 32 3 32.44 22.2M 339.2G C. Why multi-stage? Variant PSNR Params FLOPs M1-FFT 39.67 4.1M 71G M1-MLP 39.75 5.4M 83G M1-gMLP 39.80 6.1M 93G M1-SA 39.79 5.3M 111G M2-FFT 39.74 10.1M 172G M2-MLP 39.70 12.7M 195G M2-gMLP 39.83 14.1M 216G M2-SA 39.85 12.5M 250G D. Beyond gMLP.

表7： 消融研究。 子表A和B中的组件在Sec.4.3中定义。 S、W和D分别表示级数、宽度和深度。 子表D中的M1和M2分别表示MAXIM-1S和MAXIM-2S型号。

5结论

我们提出了一个用于恢复或增强任务的通用网络，称为 MAXIM，受到最近流行的基于 MLP 的全局模型的启发。 我们的工作提出了一种有效且高效的方法，将 gMLP 应用于低级视觉任务，以获得全局关注，这是基本 CNN 所缺失的属性。 我们对 MAXIM 系列的 gMLP 初始化显着提高了多项具有中等复杂度的图像增强和恢复任务的最先进水平。 我们演示了一些应用程序，但还有许多超出本工作范围的可能性，使用 MAXIM 可以显着受益。 我们未来的工作包括探索用于极高分辨率图像处理的更有效模型，以及可以适应多种任务的大型模型。

更广泛的影响。 所提出的模型可以用作增强和修饰日常照片的有效工具。 然而，出于隐私考虑，去噪和去模糊等增强技术很容易被恶意使用。 根据特定数据训练的模型可能会出现偏差。 这些问题应该由研究人员负责任地解决。

6致谢

我们感谢 Junjie Ke、Mauricio Delbracio、Sungjoon Choi、Irene Zhu、Innfarn Yoo、Huiwen Chang 和 Ce Liu 的宝贵讨论和反馈。

附录A附录

A.1 数据集和训练详细信息

论文中使用的所有数据集总结在 选项卡。 8. 我们在下面描述每个数据集的训练细节。 请注意，我们将 ${\mathrm{\ell }}_2$ 损失用于去雾任务，同时将主论文中定义的损失用于所有其他任务。

图像去噪。 我们在 SIDD [2] 中提供的 $320$ 高分辨率图像上训练我们的模型，并在 1,280 ($256\times 256$) 和 1,000 ($512\times 512$) 上进行评估t3>)图像分别由 SIDD [2] 和 DND [72] 作者提供。 DND结果通过在线服务器[1]获得。 我们将训练图像裁剪成步幅为 256 的 $512\times 512$ 补丁来准备训练补丁。 我们对 MAXIM-3S 模型进行了 600k 步训练，批量大小为 256。

图像去模糊。 我们使用来自 GoPro [62] 的 2,103 个图像对训练我们的模型。 为了证明泛化能力，我们在 1,111 对 GoPro 评估集、HIDE 数据集 [81] 中的 2,025 个图像以及 RealBlur 数据集 [79]< 上评估了我们的 GoPro 训练模型。 /t1>，其中分别包含 980 个相机 JPEG 输出和 RAW 图像的配对图像。 我们将 GoPro 中的训练图像裁剪为步长为 128 的 $512\times 512$ 补丁，以生成训练补丁。 我们对 MAXIM-3S 模型进行了超过 60 万步的训练，批量大小为 256。 为了评估 RealBlur 设置 (2)（请参阅主要论文），我们加载了 GoPro 预训练检查点，并分别在 RealBlur-J 和 RealBlur-R 上微调了 70k 和 15k 迭代。 此外，我们还使用来自 NTIRE 2021 图像去模糊挑战赛道 2 JPEG 工件的 REDS 数据集的 24,000 张图像来训练我们的模型[63]。 为了进行评估，我们遵循 NTIRE 2021 图像去模糊挑战赛中的设置[64],即。，我们在验证中使用了 300 张图像一组红色。 我们在 REDS [63] 上从头开始训练 10k epoch。

Task	Dataset	#Train	#Test	Test Dubname
Denoising	SIDD [2]	320	40	SIDD
	DND [72]	0	50	DND
Deblurring	GoPro [62]	2103	1111	GoPro
	HIDE [81]	0	2025	HIDE
	RealBlur-J [79]	3758	980	RealBlur-J
	RealBlur-R [79]	3758	980	RealBlur-R
	REDS [63]	24000	300	REDS
Deraining	Rain14000 [27]	11200	2800	Test2800
	Rain1800 [105]	1800	0	-
	Rain800 [115]	700	98	Test100
	Rain100H [105]	0	100	Rain100H
	Rain100L [105]	0	100	Rain100L
	Rain1200 [114]	0	1200	Test1200
	Rain12 [51]	12	0	-
	Raindrop [73]	861	58	Raindrop-A
	Raindrop [73]	0	239	Raindrop-B
Dehazing	RESIDE-ITS [46]	13990	500	SOTS-Indoor
	RESIDE-OTS [46]	313950	500	SOTS-Outdoor
Enhancement	MIT-Adobe FiveK [8]	4500	500	FiveK
(Retouching)	LOL [98]	485	15	LOL

表8： 五种图像处理任务的数据集摘要。

图9： 我们采用通用的多阶段框架来提高 MAXIM 的性能，以应对具有挑战性的恢复任务。 受[111, 16]的启发，我们采用监督注意力模块（SAM）和跨阶段特征融合来帮助后期阶段学习。 与以前的方法不同，由于提出的多轴 MLP 方法，我们的 MAXIM 主干网络在每个阶段的每一层都获得了全局感知，使其在学习低级和高级特征的全局交互方面更加强大。

图像去雨。 在[111, 34]之后，我们使用了一个复合训练集，其中包含从多个数据集收集的13,712个干净的雨图像对[27, 105, 115, 105, 114, 51]。 在五个测试集上进行评估：Rain100H [105]、Rain100L [105]、Test100 [115]、Test1200 [114 ] 和 Test2800 [27]。 我们对 MAXIM-2S 模型进行了超过 50 万步的训练，批量大小为 512。 对于雨滴去除任务，我们在 Raindrop 数据集 [73] 中的 861 对训练图像上训练 MAXIM-2S，执行 80k 步骤，批量大小为 512，并在测试集 A（58 张图像）和分别是测试集 B（239 张图像）。

图像去雾。 RESIDE数据集[46]包含两个子集：室内训练集（ITS）包含由1399张干净图像生成的13,990张模糊图像，以及户外训练集（OTS）由8,970张合成的313,950张模糊图像组成无雾霾的室外场景。 我们在综合目标测试集 (SOTS) [46] 上评估了我们的模型：分别为 ITS 训练的模型提供 500 张室内图像，为 OTS 训练的模型提供 500 张室外图像。 我们使用 ${\mathrm{\ell }}_2$ 损失对 RESIDE-ITS 和 RESIDE-OTS 进行了 10k 和 500 epoch 的训练。

图像增强。 我们使用[65]提供的MIT-Adobe FiveK [8]数据集进行修饰评估：前4,500张图像用于训练，其余500张用于测试。 我们将训练图像裁剪为步幅为 256 的 $512\times 512$ 补丁。 我们还使用了 LOL 数据集 [98]，其中包含 500 对图像用于低光增强。 我们在 485 个训练图像上训练模型，并在 15 个测试图像上进行评估。 我们分别在 FiveK 和 LOL 上训练了 14k 和 180k 步。

A.2架构细节

Fig. 9中展示了我们提出的多阶段和多尺度通用框架，其中每个阶段都使用单阶段 MAXIM 主干网，这在主论文中有所说明。 我们利用多尺度输入输出方法[20]来深度监督每个阶段。 具体来说，给定输入图像$𝐈\in {ℝ}^{H\times W\times 3}$，我们使用最近邻降尺度方法[20]来生成多尺度输入变体：${𝐈}_n,n=1,2,3$，同时我们采用用于产生地面实况变体的双线性缩减器：${𝐓}_n,n=1,2,3$。 对于每个阶段，我们使用 Conv3x3 从每个尺度的输入中提取浅层特征。 除了第一阶段之外，我们使用交叉门控块(CGB)将浅层特征与来自先前监督注意模块（SAM）[111]的注意特征融合在一起。 我们还采用跨阶段特征融合[111, 16]来帮助后面的阶段，其中来自前一阶段的中间Encoder和Decoder特征与使用CGB在当前阶段编码的特征融合(图9中的蓝线 )。

A.2.1 配置

单级 MAXIM 的 Encoder 部分的详细规格如下所示 选项卡。 9. 我们还提供每个块和层的输入和输出形状。 这里 Conv3x3_s1_w32 表示具有 3x3 内核、步幅 1 和 32 个通道的 Conv 层。 MAB和RCAB是Encoder / Decoder / Bottleneck中的两个主要组件。 请注意，在 Bottleneck 块中，我们使用 (Conv1x1) 层来替换 RCAB 中的 Conv3x3。

Depth	Input shape	Output Shape	Layers
1	${256}^2\times 3$	${256}^2\times 32$	Conv3x3_s1_w32
1	${256}^2\times 32$	${256}^2\times 32$	CGB* ($b=d=16$)
1	${256}^2\times 32$	${256}^2\times 32$	Conv1x1_s1_w32
1	${256}^2\times 32$	${256}^2\times 32$	{ $\begin{array}{c}\mathrm{MAB}\left(b=d=16\right)\\ \mathrm{RCAB}\left(3\times 3,r=4\right)\end{array}\}\times 2$
1	${256}^2\times 32$	${128}^2\times 32$	Conv3x3_s2_w32
2	${128}^2\times 32$	${128}^2\times 64$	Conv3x3_s1_w64
2	${128}^2\times 64$	${128}^2\times 64$	CGB* ($b=d=16$)
2	${128}^2\times 64$	${128}^2\times 64$	Conv1x1_s1_w64
2	${128}^2\times 64$	${128}^2\times 64$	{ $\begin{array}{c}\mathrm{MAB}\left(b=d=16\right)\\ \mathrm{RCAB}\left(3\times 3,r=4\right)\end{array}\}\times 2$
2	${128}^2\times 64$	${64}^2\times 64$	Conv3x3_s2_w64
3	${64}^2\times 64$	${64}^2\times 128$	Conv3x3_s1_w128
3	${64}^2\times 128$	${64}^2\times 128$	CGB* ($b=d=8$)
3	${64}^2\times 128$	${64}^2\times 128$	Conv1x1_s1_w128
3	${64}^2\times 128$	${64}^2\times 128$	{ $\begin{array}{c}\mathrm{MAB}\left(b=d=8\right)\\ \mathrm{RCAB}\left(3\times 3,r=4\right)\end{array}\}\times 2$
3	${64}^2\times 128$	${32}^2\times 128$	Conv3x3_s2_w128
4	${32}^2\times 128$	${32}^2\times 256$	Conv1x1_s1_w256
4	${32}^2\times 256$	${32}^2\times 256$	{ $\begin{array}{c}\mathrm{MAB}\left(b=d=8\right)\\ \mathrm{RCAB}\left(1\times 1,r=4\right)\end{array}\}\times 2$
4	${32}^2\times 256$	${32}^2\times 256$	Conv1x1_s1_w256
4	${32}^2\times 256$	${32}^2\times 256$	{ $\begin{array}{c}\mathrm{MAB}\left(b=d=16\right)\\ \mathrm{RCAB}\left(1\times 1,r=4\right)\end{array}\}\times 2$

表9： 单级 MAXIM 主干编码器部分的详细架构规范。 深度 1-3 表示编码器块，而深度 4 对应于骨干块。 请注意，在瓶颈中，我们在 RCAB 中使用 Conv1x1。 *表示第一阶段未使用的层。

MAXIM 的 Decoder 部分关于以下对称： 选项卡。 9，并且具有相同的配置。 对于CGB琴颈，我们对深度1和2使用$b=d=16$，而对深度3采用$b=d=8$。 基本上，我们将高分辨率阶段的块和网格大小设置为 $16$ (即。特征大小$\ge 128$) $8$ 用于低分辨率阶段(即。特征尺寸)。 因此，输入图像的两个维度都需要能被 64 整除，因此需要在推理过程中用 64 的乘数来填充图像。

A.2.2与其他 MLP 的比较

在图10中，我们展示了最近MLP模型之间近似有效感受野的视觉比较：MLP-Mixer[87 ]、gMLP [54]、Swin-Mixer [57] 以及我们提出的 MAXIM。 我们的方法实现了稀疏交互，以获得局部(图10c中的red)和全局扩张(绿色)空间通信。 此外，如图所示 选项卡。 10与之前的 MLP 模型不同，MAXIM 获得了全局和全卷积特性，并且复杂度与像素数 $N$ 呈线性关系。

A.3JAX 实现

这里我们提供了 MAXIM 关键组件的 JAX [6] 实现，即 算法 中的多轴门控 MLP 模块 (MAB) 1。

A.4 性能与。 复杂

我们展示了性能与。复杂性权衡 选项卡。 11 与所有任务的其他竞争方法相比。 可以看出，我们的模型以非常适中的复杂度获得了最先进的性能。 例如，在去噪方面，MAXIM-3S 只有 $21\%$ FLOPs 和 MIRNet [110] 的 $70\%$ 参数；在去模糊方面，我们的 MAXIM-3S 模型仅需要先前最佳模型 HINet [16] 参数数量的 $25\%$ 参数数量，并且仅需要 $19\%$ 参数数量Transformer 模型 IPT [15] 的参数数量。 还值得注意的是，与 IPT 不同，我们的模型不需要大规模预训练即可获得领先的性能，这使其对于数据集通常规模有限的低级任务具有吸引力。

Model	Complexity	Fully-conv	Global
MLP-Mixer [87]	$𝒪\left(N^2\right)$	✗	✓
gMLP [54]	$𝒪\left(N^2\right)$	✗	✓
Swin-Mixer [57]	$𝒪\left(N\right)$	✓	✗
MAXIM (ours)	$𝒪\left(N\right)$	✓	✓

表10： MAXIM 与其他 MLP 模型的比较。 我们的模型是全卷积的和全局的，具有相对于像素数 $N$ 的线性复杂度。

图10： (a) Mixer/gMLP、(b) Swin-Mixer 和 (c) MAXIM 的蓝色像素的有效感受野（阴影区域）的可视化。 MAXIM 获得局部(红色)和（扩张）全局(绿色)感知。 黄色像素可以通过局部和全局分支来实现。

A.5额外实验

由于主论文篇幅有限，我们还展示了去模糊和雨滴去除的实验结果。

REDS 去模糊 [63]。 选项卡。 12 显示了 MAXIM-3S 与获胜解决方案 HINet [16] 以及领先模型 MPRNet [111] 在 NTIRE 2021 图像去模糊挑战赛 REDS 数据集上的定量比较轨道 2 JPEG 伪影[63]。 这些指标是根据 300 个验证图像计算并取平均值的。 我们的 MAXIM-3S 模型的 PSNR 超过 HINet 0.1 dB。

Task	Dataset	Model	PSNR	Params	FLOPs
Denoise	SIDD [2]	MPRNet [111]	39.71	15.7M	1176G
		MIRNet [110]	39.72	31.7M	1572G
		MAXIM-3S	39.96	22.2M	339G
Deblur	GoPro [62]	MPRNet [111]	32.66	20.1M	1554G
		HINet [16]	32.71	88.7M	341G
		IPT [15]	32.58	114M	1188G
		MAXIM-3S	32.86	22.2M	339G
Derain	Rain13k (Average)	MSPFN [34]	30.75	21.7M	-
		MPRNet [111]	32.73	3.64M	297G
		MAXIM-2S	33.24	14.1M	216G
Dehaze	Indoor [46]	MSBDN [23]	33.79	31.3M	83G
		FFA-Net [74]	36.36	4.5M	576G
		MAXIM-2S	39.72	14.1M	216G
Enhance	LOL [98]	MIRNet [110]	24.14	31.7M	1572G
		MAXIM-2S	23.43	14.1M	216G

表11： 模型性能与。我们的模型与所有任务的其他竞争方法的复杂性比较。 FLOP 是根据大小为 $256\times 256$ 的输入图像计算的。

	REDS [63]
Method	PSNR	SSIM
MPRNet [111]	28.79	0.911
HINet [16]	28.83	0.862
MAXIM-3S	28.93	0.865

表12： REDS 上的去模糊比较。 我们的方法在 NTIRE 2021 图像去模糊挑战 Track 2 JPEG 伪像的 REDS 数据集上优于之前的获胜解决方案 (HINet)。 分数是根据验证集中的 300 张图像进行评估的。 结果来自[16]的作者。

	Raindrop-A [73]		Raindrop-B [73]
Method	PSNR	SSIM	PSNR	SSIM
AGAN [73]	31.62	0.921	25.05	0.811
DuRN [56]	31.24	0.926	25.32	0.817
Quan [75]	31.36	0.928	-	-
MAXIM-2S	31.87	0.935	25.74	0.827

表13： 雨滴去除数据集 [73] 的除雨比较。 我们的 MAXIM-2S 模型在 Raindrop 测试集 A 和 B 上均达到了最先进的性能。

去除雨滴 [73]。 除了主论文中报告的雨条纹去除任务之外，我们还评估了雨滴去除任务上的 MAXIM 模型。 可以看出 选项卡。 13，我们的模型在 Raindrop 测试集 A 和 B 上实现了最佳性能：31.87 dB 和 25.74 dB PSNR。

A.6 更多视觉比较

去噪。 图 12显示了我们的模型与SIDD上的SOTA模型的去噪结果[2]。 我们的模型恢复了更多细节，产生视觉上令人愉悦的输出。

去模糊。 GoPro [62]、HIDE [81]、RealBlur-J [79] 和 REDS [63] 上的视觉结果如图图13、图14、图1>152>0>、图4>16<分别为 /t15>3>。 我们的模型在合成和现实世界的去模糊基准上都优于其他竞争方法。

除雨。 我们的模型与 SOTA 方法在除雨方面的定性比较如图 图 17、图所示。 18、图、19、图0> 201>。

雨滴去除。 我们在图21中提供了雨滴测试集A和B[73]上雨滴去除任务的视觉比较> 和图22。

去雾。 我们在图23和中提供了SOTS [46]室内和室外设备的去雾比较。 图24。

修饰。 图 25 显示了我们的模型与提供的 5-K 数据集 [8] 上的竞争方法的其他比较通过 [65] 修饰结果。

低光增强。 图 26展示了LOL [98]测试集上低光增强的评估。

A.7 权重可视化

图 11可视化在 GoPro 上训练的 MAXIM-3S 每个阶段的块 gMLP 和网格 gMLP 层的空间投影矩阵 [62]。 与[54]类似，我们也观察到学习后的权重表现出局部性和空间不变性。 令人惊讶的是，全局网格 gMLP 层还学习执行“局部”操作（但在统一扩张网格上）。 同一层中的块gMLP和网格gMLP的空间权重通常表现出相似或耦合的形状，这可能归因于多轴gMLP块中的平行分支设计。 然而，我们还没有观察到这些过滤器在不同阶段如何变化的明显趋势。

A.8限制和讨论

我们的模型与现有 SOTA 共有的一个潜在限制是对现实世界示例的泛化相对不充分。 这也许可以归因于现有合成图像恢复基准提供的训练示例。 通过数据生成方案创建更真实的大规模数据集[82, 94]可以改善这一缺点。 此外，我们观察到我们的模型倾向于稍微过度拟合某些基准，因为我们在训练期间没有应用强正则化(例如。，dropout）。 尽管我们发现正则化可能会导致我们的模型在我们评估的这些基准上的性能略有下降，但未来值得探索以有效提高我们的恢复模型的泛化能力。

值得一提的是，我们的模型能够生成高质量的清晰图像，在视觉上可与最先进的生成模型[123, 36]相媲美。 值得注意的是，我们的模型产生了更保守的结果，而不会产生许多不存在的细节的幻觉，从而提供比生成模型更可靠的结果。

图11： 在 GoPro [62] 上训练的 MAXIM-3S 模型的块 gMLP 和网格 gMLP 层中的空间投影权重。 每行显示一组减少的连续通道的滤波器（重新整形为 2D）。 编码器深度 1 和 2 的滤波器大小为 $16\times 16$，而编码器深度 3 和 Bottleneck1 的滤波器大小为 $8\times 8$（调整大小为相同形状以获得更好的可视化效果）。 值得注意的是，块 gMLP 层（左）的权重直接应用于局部窗口内的像素，并在特征图的每个非重叠窗口共享（类似于跨步卷积），而网格 gMLP 层（右）的权重）对应于覆盖在整个图像上的全局扩张聚合。

图12： VDN [107]、DANet [108]、MIRNet [110]< 中 SIDD [2] 图像去噪的可视化示例/t3>、CycleISP [109]、MPRNet [111] 以及提出的 MAXIM-3S。 我们的模型清楚地消除了真实噪声，同时恢复了更多细节。

图13： DMPHN [113]、Suin 等人中的 GoPro [62] 图像去模糊视觉示例。 [83]、MPRNet [111]、HINet [16]、MIMO-UNet [20] 以及我们的 MAXIM -3S。

图14： DMPHN [113]、Suin 等人在 HIDE [81] 上图像去模糊的视觉比较。 [83]、MPRNet [111]、HINet [16]、MIMO-UNet [20] 以及我们的 MAXIM -3S。

图15： 先前最佳模型 MPRNet [111] 和 MAXIM-3S 在 RealBlur-J [79] 上的图像去模糊视觉比较。

图16： 我们的模型与获胜解决方案 HINet [16] 在 REDS [63] 上的图像去模糊视觉比较，适用于 NTIRE 2021 图像去模糊挑战赛 Track 2 JPEG 的 REDS 数据集工件[63]。

图17： RESCAN [49]、PreNet [77]、MSPFN [34]< 中 Rain100L [105] 图像除雨的可视化示例/t3>、MPRNet [111]、HINet [16] 和我们的 MAXIM-2S 模型。

图18： Rain100H [105] 上图像除雨的视觉示例。 在降雨量极高的情况下，与之前的竞争方法相比，我们的模型可以恢复更多的细节和纹理。

图19： Test100 [115] 上图像除雨的视觉示例。 我们的模型消除了雨状条纹和可见的 JPEG 伪影。

图20： Test1200 [114] 上图像除雨的视觉示例。

图21： AGAN [73]、DuRN [56]、Quan [75 之间 Raindrop-A [73] 上雨滴去除的视觉比较] 和 MAXIM-2S。

图22： Raindrop 测试集 B [73] 上雨滴去除的视觉比较。

图23： GCANet [14]、GridDehaze [55]、DuRN [56 之间 SOTS 室内测试集 [46] 上图像去雾的视觉比较]、MSBDN [23]、FFA-Net [74] 和我们的 MAXIM-2S。

图24： MAXIM-2S 的 SOTS 户外测试集 [46] 上的图像去雾与其他方法的视觉比较。

图25： [65] 的作者提供的 MIT-Adobe FiveK [8] 图像修饰与 CycleGAN [124] 之间的视觉比较，Exposure [32]、DPE [19]、EnlightenGAN [37]、UEGAN [65] 和 MAXIM-2S。

图26： Retinex [98] 、GLAD [92] 、KinD [122] 、EnlightenGAN [37] 、MIRNet [110] 和 MAXIM-2S 之间在 LOL 数据集 [98] 上进行图像弱光增强的可视化示例。 与其他方法相比，我们的模型有效增强了照明，同时大大降低了噪声，从而产生了更高质量的图像。

算法1 实现多轴门控 MLP 块 (MAB) 的 JAX 代码。

⬇

from typing import Sequence

import einops

import flax.linen as nn

import jax.numpy as jnp

def block_images(x, patch_size):

n, h, w, channels = x.shape

grid_height, grid_width = h // patch_size[0], w // patch_size[1]

x = einops.rearrange(x, "n␣(gh␣fh)␣(gw␣fw)␣c␣->␣n␣(gh␣gw)␣(fh␣fw)␣c",

gh=grid_height, gw=grid_width, fh=patch_size[0], fw=patch_size[1])

return x

def unblock_images(x, grid_size, patch_size):

x = einops.rearrange(x, "n␣(gh␣gw)␣(fh␣fw)␣c␣->␣n␣(gh␣fh)␣(gw␣fw)␣c",

gh=grid_size[0], gw=grid_size[1], fh=patch_size[0], fw=patch_size[1])

return x

class SpatialGatingUnit(nn.Module):

"""Gated␣MLP␣applied␣on␣a␣specified␣axis:␣-3␣for␣grid␣and␣-2␣for␣block."""

@nn.compact

def __call__(self, x, axis=-3):

u, v = jnp.split(x, 2, axis=-1)

v = nn.LayerNorm()(v)

n = x.shape[axis] # get spatial dim at the ’grid’ or ’block’ axis

v = jnp.swapaxes(v, -1, axis)

v = nn.Dense(n)(v)

v = jnp.swapaxes(v, -1, axis)

return u * (v + 1.)

class SpatialGmlpLayer(nn.Module):

"""Gated␣MLP␣applied␣on␣a␣specified␣axis:␣-3␣for␣grid␣and␣-2␣for␣block."""

grid_size: Sequence[int]

block_size: Sequence[int]

@nn.compact

def __call__(self, x, axis=-3):

n, h, w, num_channels = x.shape

if axis=-3: # for grid gMLP layer

gh, gw = self.grid_size

fh, fw = h // gh, w // gw

elif axis=-2: # for block gMLP layer

fh, fw = self.block_size

gh, gw = h // fh, w // fw

x = block_images(x, patch_size=(fh, fw))

y = nn.LayerNorm()(x)

y = nn.Dense(num_channels * 2)(y)

y = nn.gelu(y)

y = SpatialGatingUnit()(y, axis=axis)

y = nn.Dense(num_channels)(y)

x = x + y

x = unblock_images(x, grid_size=(gh, gw), patch_size=(fh, fw))

return x

class MultiAxisGmlpBlock(nn.Module):

block_size: Sequence[int]

grid_size: Sequence[int]

@nn.compact

def __call__(self, x):

shortcut = x

n, h, w, num_channels = x.shape

x = nn.LayerNorm()(x)

x = nn.Dense(num_channels * 2)(x)

x = nn.gelu(x)

# split two heads, then applied grid gMLP and block gMLP respectively.

u, v = jnp.split(x, 2, axis=-1)

u = SpatialGmlpLayer(grid_size=self.grid_size)(u, axis=-3)

v = SpatialGmlpLayer(block_size=self.block_size)(v, axis=-2)

# Concat and output projection

x = jnp.concatenate([u, v], axis=-1)

x = nn.Dense(num_channels)(x)

x = x + shortcut

return x