NAFNet:Simple Baselines for Image Restoration

Abstract

近年来，图像复原技术取得了长足的进步，但现有的图像复原方法（SOTA）系统复杂度也在不断增加，不利于对各种方法的分析和比较。在本文中，我们提出了一种简单的基线，它超越了SOTA方法，并且计算效率很高。为了进一步简化基线，我们揭示了非线性激活函数，例如Sigmoid、ReLU、GELU、Softmax等不是必需的：它们可以被乘法替换或去除。因此，我们从基线推导出一个非线性无激活网络，即NAFNet。SOTA结果在各种具有挑战性的基准测试中实现，例如，GoPro上的33.69 dB PSNR（用于图像去模糊），超过了之前的SOTA 0.38 dB，而计算成本仅为其8.4%; SIDD上的40.30 dB PSNR（用于图像去噪），超过了之前的SOTA 0.28 dB，而计算成本不到其一半。

Introduction

包含GELU $15$ 和通道注意模块 $16$ （CA）的基线可以进一步简化：我们揭示基线中的GELU可以被视为门控线性单元 $10$ （GLU）的特例，并且由此我们经验性地证明它可以被简单的门（即，特征映射的元素式乘积）代替。
此外，我们还揭示了CA与GLU在形式上的相似性，并且可以去除CA中的非线性激活函数。

大多数方法可以被视为经典解决方案UNet的变体，这些变体带来了性能增益以及系统复杂度，并且我们将复杂度大致分类为块间复杂度和块内复杂度
与 $30$ 不同的是，我们在GLU中去除了非线性激活函数，而不会降低性能。此外，基于非线性无激活GLU本身包含非线性的事实（因为两个线性变换的乘积提高了非线性），我们的基线可以通过用两个特征图的乘积代替非线性激活函数
据我们所知，它是第一个在没有非线性激活函数的情况下实现SOTA性能的计算机视觉模型。

(a)多级架构 $5，37$ 串联堆叠UNet架构。(b)多尺度融合架构 $25，7$ 融合了不同尺度的特征。（c）UNet架构，一些SOTA方法采用该架构 $39，36$ 。

Baseline

一个普通的块开始图(b)即卷积操作称为PlainNet

为什么使用CNN而不是Transformer

首先，虽然变压器在计算机视觉中表现出良好的性能，但一些作品 $13，23$ 声称它们可能不是实现SOTA结果所必需的。
第二，深度卷积比自我注意机制更简单 $34$ 。
第三，本文不打算讨论transformers和卷积神经网络的优缺点，而只是提供一个简单的基线。

Normalization

虽然 $26$ 放弃了批量规格化 $18$ ，因为小批量可能会带来不稳定的统计数据 $38$ ， $5$ 重新引入实例规格化 $33$ 并避免了小批量问题。随后，在Transformer发展下，层归一化 $3$ 被越来越多的方法使用，包括SOTA方法 $32，39，36，23，22$ 。基于这些事实，我们推测层规范化可能对SOTA恢复器至关重要，因此我们将层规范化添加到上述普通块中。总而言之，我们将层规范化添加到普通块中，因为它可以稳定训练过程。

Activation

在普通块中用GELU代替ReLU，因为它保持了图像去噪的性能，同时为图像去模糊带来了非平凡的增益。

Attention

vanilla自注意机制 $34$ 是由 $12，4$ 采用的，通过所有特征的线性组合来生成目标特征，这些特征通过它们之间的相似性来加权。因此，每个特征都包含全局信息，但其计算复杂度与特征图的大小成二次关系。一些图像恢复任务以高分辨率处理数据，这使得该自注意不实际。
$22，21，36$ 仅在固定大小的局部窗口中应用自注意，以减轻增加的计算的问题。而它缺乏全局信息。我们不采用基于窗口的注意力，因为在平面块中通过深度卷积 $13，23$ 可以很好地捕获局部信息。
$39$ 将空间注意力修改为通道注意力，避免了计算问题，同时保持了每个特征中的全局信息。另外，通道注意的有效性已经在图像恢复任务 $37，8$ 中得到验证，因此我们将通道注意添加到了平坦块中。

到目前为止，我们从头开始构建了一个简单的基线，如表1所示。它可以超越SIDD和GoPro数据集上的SOTA结果，而计算成本只有一小部分。

Nonlinear Activation Free Network

上面描述的基线是简单和有竞争力的，但是在确保简单性的同时，是否有可能进一步提高性能？它能更简单而不损失性能吗？我们试图通过从一些SOTA方法中寻找共性来回答这些问题 $32，39，20，17$ 。我们发现，在这些方法中，采用了门控线性单元 $10$ （GLU）。这意味着GLU可能是有前途的。我们接下来会讨论它。

Gated Linear Units

其中，X表示特征图，f和g是线性变换器，σ是非线性激活函数，例如Sigmoid，并且σ表示逐元素乘法。如上所述，将GLU添加到我们的基线可以提高性能，但块内复杂度也在增加。这不是我们所期望的。为了解决这个问题，我们重新审视了基线中的激活函数，即GELU $15$ ：

其中Φ表示标准正态分布的累积分布函数。根据 $15$ ，GELU可以近似并实现为：

由方程1和方程2，可以注意到GELU是GLU的一个特例，即f，g是单位函数，取σ为Φ。通过这种相似性，我们从另一个角度推测，GLU可以看作是激活函数的推广，有可能取代非线性激活函数。此外，我们注意到GLU本身包含非线性，并且不依赖于σ：即使σ被移除，Gate（X）= f（X）⊙ g（X）也包含非线性。

在此基础上，我们提出了一个简单的GLU变体：在通道维度上直接将特征图分为两部分，然后将它们相乘，如图4c所示，标记为SimpleGate。与等式3中的GELU的复杂实现相比，我们的SimpleGate可以通过元素乘法来实现，仅此而已：

X、Y是相同大小的特征图

结果表明，SimpleGate可以替代GELU，替代后网络中只剩下通道注意力中存在激活函数。作者同样进行了优化和替代。

Simplified Channel Attention

我们将通道注意力 $16$ 引入到我们的块中，因为它捕获全局信息并且计算效率高。如图4a所示：它首先将空间信息压缩到通道中，然后将多层感知应用于它以计算通道注意力，这将用于对特征图进行加权。它可以表示为：

其中X表示特征图，pool表示将空间信息聚合到通道中的全局平均池化操作。σ是一个非线性激活函数Sigmoid，W1，W2是全连接层，两个全连接层之间采用ReLU。如果我们将通道注意力计算视为一个函数，5可以改写为：

可以注意到，Eqn.6与Eqn.1非常相似。这启发我们将通道注意力视为GLU的一种特殊情况，它可以像前一小节中的GLU一样简化。通过保留渠道注意力的两个最重要的角色，即聚合全局信息和渠道信息交互，我们提出了简化的渠道注意力：

简化的通道注意力（方程式7)比原来的简单Eqn.5），如图4a和图4b所示。虽然它更简单，但性能没有损失：SIDD为+0.03 dB（39.93 dB至39.96 dB），GoPro为+0.09 dB（32.76 dB至32.85 dB）。