Pyramid attention networks for image restoration

paper:https://arxiv.org/abs/2004.13824

code: https://github.com/SHI-Labs/Pyramid-Attention-Networks

1. 基本思想

作者指出，当前基于深度学习的方法只是在单个尺度上利用了self-similarity（do not take full advantage of self-similarities by relying on self-attention neural modules that only process information at the same scale）。在该论文里，作者提出了 pyramid attention module 用于图像修复，可以利用多尺度特征金字塔捕获 long-range 的特征关联。

在传统图像去噪方法中，NLmean、BM3D都是采用自相似性进行图像去噪的佼佼者，图像的自相似性非常重要。同时，作者指出，基于深度学习的方法大多使用non-local方法来获取 long-range 的信息，这样的模块早期被用于识别和分类任务，但是在 low-level 的视觉任务中也被证明非常有效。

但是，直接将 self-attention 应用于图像修复有一定的局限性：

multi-scale non-local prior has not been explorered.
pixel-wise matching used in the self-attention module is usually noisy for image restoration.

基于上述分析，作者提出了non-local pyramid attention 模块来捕获long-range信息。

如图所示，该方法可以挖掘 multi-scale self-examplars 来提高重建性能，主要有三个优点：

将现有non-local方法的扩展到多尺度空间（现有non-local方法的搜索空间仅限于金字塔的最底端）
不同大小特征之间的 long-range 依赖被较好的建模
可以降低特征中的噪声，进而提升性能。

2. 方法

如图2a所示，non-local attention 在全图进行响应特征平均，可以这样描述：

\[y^i=\frac{1}{\sigma(x)}\sum\phi(x^i, x^j)\theta(x^j)
\]

其中\(\phi\) 计算相似性，\(\theta\)表示特征变换。上述公式仅在统一尺度空间进行处理，难以利用跨尺度的特征。

图2b为 scale agnostic attention，是 2a 的一种扩展。给定一个尺度\(s\)，如何评价\(x^j\)和\(x^j_{\delta(s)}\)之间的相似性，以及如何把\(x^j_{\delta(s)}\)聚合到\(y^i\)是两个关键问题。为了解决这两个问题，作者将\(x^j_{\delta(s)}\)下采样为像素特征\(z^j\)，由于需要在整个特征图上搜索，作者对原始输入\(x\)进行下采样得到特征图\(z\)，然后相似性可以这样计算：

\[y^i=\frac{1}{\sigma(x,z)}\sum_j\phi(x^i,z^i)\theta(z^i)
\]

同时，作者将该计算方式扩展到金字塔形式，金字塔注意力具有跨尺度特征自相似建模能力。

在具体实现时，相似性函数\(\phi\)使用了embeded gaussian。在特征变换函数\(\theta\)方面，使用简单的线性变换。同时，为了改进匹配的鲁棒性，添加了额外的近邻相似约束，类似于经典的非局部均值滤波，添加了一个强约束信息：近邻且相似。这样有助于学习相关特征，同时抑制不相关特征。

图3 为 Pyramid attention 模块的一个示意图，在损失函数方面，作者选择了L1损失。

3. 实验与结论

作者使用了一个5层的金字塔，网络包含80个残差块，提出的PA模块插入在第40个残差块之后，第41个残差块之前。该方法在图像降噪、去马赛克、去压缩伪影、超分辨率方面都进行了实验，效果都比较好。具体细节和ablation study 可以参考论文，不再细说。

该论文使用了一个 simple backbone，然后只加了一个 pyramid attention 模块，就取得了非常好的效果。作者认为，该模块可以在以后的图像修复网络中，作为基础模块来使用。