旷视low-level系列（三）：NAFNet -- Simple Baselines for Image Restoration

标题：简单基线方法用于图像修复
所属机构：旷视科技
收录于国际计算机视觉大会 ECCV2022
论文链接：https://arxiv.org/abs/2204.04676
代码库链接：https://github.com/megvii-research/NAFNet

文章目录

• 1. Motivation
• 2. Contributions
• 3. Methods
• 4. Experiments
• 5. Comments

1. Motivation

在图像恢复领域（Field of image restoration），基于端到端深度学习的传统方法表现出色（outstanding performance）。然而，在实际应用中（in practice），由于计算资源限制（resource constraints），其在实时性要求较高的场景下（high real-time requirement scenarios）难以落地实施（difficult to deploy）。针对此类问题（issues），研究者往往倾向于选择那些能够在保证较低计算复杂度的前提下（while maintaining low computational complexity）实现与现有最优算法相当甚至超越的方法（perform equivalently or better）。在本研究中（In this study），旷视科技的研究人员提出了一种新型基准模型（novel baseline model）用于图像修复任务（image restoration task），该模型不仅超越了现有最优算法（outperforms state-of-the-art methods），而且具有更高的计算效率（higher computational efficiency）。

2. Contributions

• 通过分析和提取SOTA方法的核心组件构建了一个计算开销较小的基础模型，在性能上超越了现有SOTA方案的同时其计算开销显著降低；
• 通过剔除或替代非线性激活函数简化了现有基础模型并提出了一种完全不依赖于非线性激活函数的新网络架构NAFNet在性能上与现有方案相当甚至超越之首次表明完全不含非线性激活函数的网络架构同样能达到顶尖模型水平。

3. Methods

为便于讨论起见，作者将模型系统的总复杂度划分为block间complexity（inter-block complexity）与block内部complexity（intra-block complexity）。其中(block间complexity)指的是不同block之间所存在的各种连接关系所带来的整体难度。如图所示,(a)展示了多阶段架构模式,(b)则体现了多尺度融合架构特征。在(block内部complexity)方面，则侧重于减少每个block内部组件之间的相互作用程度。
通过减少(block间complexity)，作者采用了单阶段UNet结构，并特别关注于减小每个block内部组件之间的相互作用程度。

Fig.1

为了规划block内部结构以降低其复杂度至最低水平，请问作者的具体思路是什么？具体来说，在此问题下，请问作者的大致思路如下：首先构建一个基础的plain block，并借鉴SOTA模型中已有的成熟模块，在 plain block 的基础上逐步添加或替换相关组件以形成 baseline version。随后通过消融实验验证其性能是否有所提升。最后，在确保性能不下降的情况下尽可能简化架构。

Fig.2

plain block 基于这一分析，在构建 plain block 时选择了一个更为简洁的架构方案。具体而言，在传统 SOTA 方法中引入复杂的 self-attention 结构不仅会增加模型的计算开销（运算量），而且与我们提出的 simple baseline 的目标理念存在一定的违背关系。基于此原则，在 plain block 的内部结构设计中完全摒弃了自注意力机制（self-attention）等复杂组件，在模型构建阶段仅采用卷积操作、激活函数以及 shortcut 模块进行组合，并参考 Restormer 块的设计思路确定各模块之间的连接顺序

baseline block 基于plain block，在normalization, activation和attention三个维度上进行优化

归一化操作在低级视觉任务中不适用于神经网络设计应被视为一种共识意见。在网络架构中引入归一化操作可能会对性能水平产生影响。

BN导致性能下降的原因可能包括：
1）小的batch-size带来了统计结果不够稳定；
2）与高层任务不同的是它们往往倾向于寻找一致性的表示，
而底层视觉的任务恰恰相反——它们主要通过学习图片特有的属性来增强细节恢复的效果（例如此前研究者曾尝试通过捕捉图像分布中的sigma值来提升边缘区域的表现），然而... batchNorm因其在每个batch内部进行注意力机制的应用，
容易引入其他批次的信息干扰，
这种现象使得它在处理特定图像细节时出现了不足，
因此会导致性能下降。

鉴于LN在transformer架构及多种SOTA图像恢复方法中具有广泛应用，并基于此观察可知，在plain block中添加Layer Normalization（LN）被认为是一个不可或缺的关键组件用于实现SOTA性能。进一步分析表明，在保持较高训练效率的前提下（即使学习率增加10倍），通过这一发现，在实验结果中我们观察到当学习率提升至最大值时，在所有测试集上均实现了显著的性能提升：具体而言，在SIDD数据集上提升了0.44dB，在GoPro数据集上提升了3.39dB。

• activation
  ReLU被广泛应用于计算机视觉任务中，但是目前的SOTA方法中逐渐倾向于使用GELU（高斯误差线性单元激活函数）代替ReLU，例如transformer中主要使用GELU激活函数，因此作者也将plain block中的ReLU替换为GELU。替换结果：SIDD-0.02dB，GoPro+0.21dB，去噪性能相当，但是去模糊有提升，因此选择使用GELU作为激活函数。

• attention机制
  低分辨率通常难以满足高清晰度的需求。例如，在实际应用中广泛使用的1080p分辨率图像，在自注意力机制的空间复杂度与分辨率呈二次函数关系这一特性下（即随着空间分辨率增加而计算复杂度呈二次增长），其实用性受到限制。Uformer 通过限制自注意力作用范围来减少计算开销；然而这种设计虽然降低了计算负担但牺牲了全局信息捕捉能力。Restormer 提出了一种基于通道域注意力的设计方法——Multi-Dconv Head Transposed Attention模块；这种方案将空间域注意力转换为通道域注意力；通过这种方式可以在不显著增加计算开销的前提下实现对全局上下文信息的有效捕捉。受此启发作者认为：在现有的模型架构中仅需满足两个关键条件即可作为通道注意力的基础方案：1）高效的计算效率；2）具备全局信息捕获能力。而SENet 中提出的通道注意力模块则完美地满足了上述两项要求：一方面该方案能够高效地进行特征提取；另一方面它能够有效建模并捕获复杂的长距离依赖关系。因此作者决定将该模块整合到基础网络架构中以进一步提升模型性能。

通过该研究团队在该研究领域内的深入探索与系统优化，在该研究领域内的相关指标均得到了显著提升。

NAFNet模块化设计非常简洁，在此基础上提出了一个值得探讨的问题：能否在保证性能不受影响的前提下简化系统？研究者尝试从现有最优算法的一致性出发寻求解决方案，并发现大多数SOTA架构均采用了门控线性单元（GLU），这表明该组件很可能具有重要的价值。

GLU的形式如下：

其中X表示特征图，f和g为线性变换器，\sigma为非线性激活函数，如Sigmoid，⊙表示元素级乘法。添加GLU到baseline中可能会提高性能，但块内的复杂性也在增加，与“更简单”的目标相悖。为了解决该问题，作者重新审视了baseline中的GELU，发现GLU与GELU在形式上存在可转换性。GELU的形式如下：

可以看出，在某些特定条件下（例如当f,g=\text{identity}且$\sigma=\Phi时），GLU与GELU的表现是一致的。此外，在去除了σ的影响下（即即使去掉σ），通过将特征图沿通道维度分割并计算其乘积的形式仍能引入非线性特性。综上所述，在改进模型性能方面作者提出了一个创新性的变体方案

Fig.3

通过在baseline block中采用Simple Gate替代GELU后，在图像恢复任务中实施该策略可获得显著的性能提升。具体而言，在SIDD基准测试中实现了约0.08dB的增益，在GoPro测试中则提升了约0.41dB。实验结果表明该方法能够有效地替代传统GELU激活函数。

至此，在通道注意力机制中仅存少数几种非线性激活函数：sigmoid和relu函数。图穷匕见呀！作者意图是彻底消除网络中的所有非线性激活。提出首个完全摒弃非线性激活函数的图像恢复模型！

SENet中提出的channel attention的形式如下：

如果将attention的生成过程用函数表示，CA(x)可以简化为：

突然意识到其形式与GLU之间存在显著相似度。研究者指出，在这种特殊情况下，可以通过保留上述两个核心功能即可实现对原始GLU行为的近似模拟。为了实现这一目标，在模型架构中引入了Simplified Channel Attention（SCA），该机制旨在保留通道注意力模块中对全局信息的聚合能力和跨通道交互的核心作用。

虽然SCA相比CA在形式上更为简单，但是性能没有损失，SIDD+0.03dB，GoPro+0.09dB。

因该网络省略了非线性激活函数而被命名为：Nonlinear Activation Free Network（NAFNet）。

4. Experiments

在从PlainNet到Baseline展开的消融研究中，通过验证Layer Normalization（LN），ReLU网络被替换成GELU激活函数，并引入了Channel Attention机制来考察其作用的有效性。

通过从Baseline到NAFNet阶段的消融实验,验证将GELU替换成SimpleGate以及将CA替换成SCA的效果

5. Comments

• 看这篇论文的时候发现作者设计网络的思路与ConvNeXt一文很相似，首先选择一个简单的网络结构，然后借鉴SOTA方法中的一些有效组件和tricks，通过实验确定其有效性，若有效则保留，无效则弃之，靠着“边角料”的组合实现SOTA性能（真·科学炼丹？不是）。有一说一，逐步控制变量的实验思路值得学习。
• 另外一点，看到最近图像恢复领域的SOTA逐渐有抛弃transformer的趋势（可能是因为计算量太大只能抛弃，毕竟性能卷完就该卷计算量了），而且在性能上也能超越Restormer这样的transformer-based网络架构。但是值得注意的是，CNN-based SOTA在宏观和细节设计上大多都借鉴了transformer的思想，例如本文提出的NAFNet，以及23年的KBNet，24年的CGNet，都能看到transformer的影子。虽说self-attention是transformer的核心，但是组件的选择、结构上的编排以及大大小小的tricks也是不可缺失的部分。