【论文解读】NBNet：抛开复杂的网络结构设计，旷世&快手提出子空间注意力模块用于图像降噪

作者丨小小刘

审稿丨邓富城

编辑丨极市平台

导读

本文对旷视科技2021年关于图像去噪的新作"NBNet进行解读，该工作抛开复杂的网络结构设计和精确的图像噪声建模，创新性的提出子空间基向量生成和投影操作。

论文地址：https://arxiv.org/abs/2012.15028

开源代码：https://github.com/megvii-research/NBNet

NBNet创新简述

本文从图像自适应投影这一新视角出发，提出一种新的图像去噪网络------NBNet，通过在特征空间中学习一系列重建基底使网络能分离信号和噪声。选择信号子空间相应的基底并将输入信号投影到信号子空间中，以实现图像去噪。核心是希望投影能保持输入信号的局部结构，尤其是低亮度和弱纹理的区域。为此，作者提出子空间注意力SSA，一个non-local注意力模块来显式地学习基底生成和子空间投影，最终设计一个UNet结构的网络完成端到端的图像去噪。大量实验表明提出的NBNet在PSNR和SSIM指标上都取得了最好的性能。

NBNet提出动机

图像去噪是图像处理和计算机视觉领域的基本任务，一个典型的加性噪声模型可用下式表示：y=x n，想要恢复出干净的图像x其实是一个病态问题。传统方法是利用图像先验和噪声模型估计图像或是噪声，例如NLM和BM3D利用图像局部相似性和噪声独立性，小波去噪方法利用图像在变换域上的稀疏性。近年来，基于深度学习的去噪方法隐式地利用图像先验或从大规模的成对数据集中学习噪声分布。尽管基于深度学习的方法取得了很大成功，但在弱纹理或高频细节这样的困难场景下，恢复高质量图像仍然是一项挑战。因为卷积网络通常利用局部滤波器分离噪声和信号，但在低信噪比下，如果没有全局结构信息局部响应就会很容易失败。

为解决上述问题，通过投影来利用非局部图像信息从而设计图像去噪网络。如上图是图像投影和基底示意图，从输入图像生成一系列的图像基底向量，然后从这些基底向量构成的子空间重建出去噪图像。因为自然图像通常处于低秩的信号空间，因此通过准确地学习和生成这些基底向量，重建图像能最大程度保留原始信息并抑制噪声。

NBNet的具体结构

NBNet整体是UNet形式的网络，其中关键的是子空间注意力模块SSA，其学习子空间基底向量，如图1是NBNet的整体结构。NBNet的创新在于子空间投影，包括两个主要的步骤：基底向量生成和投影。

1 、基底向量生成，从输入图像特征图谱生成子空间基底向量。基底生成函数用

f_{theta}

表示, 以编码阶段的中间特征图谱为输入, 则生成的基底矩阵为

boldsymbol{V}=f_{theta}left(boldsymbol{X}_{1}, boldsymbol{X}_{2}right), boldsymbol{V}=left[v_{1}, v_{2}, cdots, v_{k}right]

。基底生成函数用

f_{theta}

可以一个残差卷积 block 实现，输出通道数为

, 然后将输出变形为

H W

times K

。下图中(b) 为残差卷积 block;(c) 为子空间注意力模块 SSA，其中 basis vectors 即为生成的基底向量。

2、投影，如上图中(c)的 projection, 通过正交线性投影将图像特征

boldsymbol{X}_{1}

投影到子空间

mathrm{V} in mathrm{R}^{N}

中。用

boldsymbol{P}: mathrm{R}^{N} rightarrow mathrm{V}

表示信号子空间的正交投影矩阵，则

boldsymbol{P}

可由

boldsymbol{V}

计算得到，如下式:

boldsymbol{P}=boldsymbol{V}left(boldsymbol{V}^{mathrm{T}} boldsymbol{V}right)^{-1} boldsymbol{V}^{mathrm{T}}

其中标准化项

left(boldsymbol{V}^{mathrm{T}} boldsymbol{V}right)^{-1}

是必要的，因为基底生成过程不能保证基向量间是相互正交的。最终图像特征

boldsymbol{X}_{1}

在信号子空间重建为

boldsymbol{Y}

, 即

boldsymbol{Y}=boldsymbol{P X}_{1}

。投影运算其实仅包含线性矩阵运算，经过适当的变形，投影运算是完全可微，在现代神经网络框架中易于实现。

图 1 NBNet网络整体结构

3、NBNet具体实现和损失函数。NBNet整体结构是典型的UNet形式，如上图1，包括4个编码阶段和4个相应的解码阶段。下采样使用k4s2的卷积操作，上采样使用反卷积。每个卷积层激活函数都使用LeakyReLU。

提出的SSA模块放置在跳转连接中。因为大尺寸低层特征图谱包含更多原始图像的细节信息，因此以编码阶段低层特征图谱为，解码阶段高层特征图谱为。即利用上采样的高层特征指导来自跳转连接的低层特征投影到信号子空间中。经投影变换后的特征与原始高层特征concat，再进行解码。与传统UNet形式结构相比，NBNet中的低层特征在与高层特征融合前先由SSA模块进行投影变换。最后一个解码器的输出通过一个线性3×3卷积层输出全局残差（针对有噪声的输入而言），再加上输入后得到最终的去噪结果。

以L1距离为网络训练的损失函数：

mathrm{L}(G, boldsymbol{x}, boldsymbol{y})=|boldsymbol{x}-G(boldsymbol{y})|_{1}