Attention in Attention for Super-Resolution

点击下方卡片，关注“AIWalker”公众号

重磅干货，第一时间送达

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

code: https://github.com/haoyuc/A2N(尚未开源)

本文是亚马逊&悉尼大学大研究员在注意力机制在图像超分中应用探索。本文从“为什么注意力会起作用，它又是如何起作用”两个问题出发，从理论&实验角度给出了解释，然后针对“图像超分中并非所有注意力都有正向作用”问题提出了注意力丢弃模块并用于构架图像超分。最后，作者通过实验证实：所提方案可以取得比其他轻量化方法更好的性能均衡。

Abstract

已有研究表明：注意力机制对于高性能超分模型非常重要。然而，很少有工作真正讨论“为什么注意力会起作用，它又是如何起作用的”。

我们尝试量化并可视化静态注意力机制并表明：并非所有注意力模块均有益。我们提出了attention in attention network（A2N）用于高精确图像超分。具体来说，A2N由非注意力分支与耦合注意力分支构成。我们提出了Attention dropout module(ADM)为两个分支生成动态注意力权值，它可以用于抑制不重要的注意力调整。这使得注意力模块可以更专注于有益样例而无需其他惩罚，因此能够以少量的额外参数提升注意力模型的容量。

实验结果表明：所提方案可以取得比其他轻量化方法更好的性能均衡。Local attribution maps(LAM)实验同样表明：所提Attention in AttentionA2结构可以从更宽的范围内提取特征。

本文主要贡献包含以下几点：

我们对神经网络不同阶段的注意力层有效性进行了量化分析，提出了一种有效的注意力层简直策略；
我们提出了一种Attention in Attention BlockA2B，它可以为起内部分支动态生成和为1的注意力。由于其中一个分支为注意力分支，故而称所提模块为A2B；
基于A2B提出了A2N，相比类似网络的基线网络，所提方法可以取得更优异的性能。

Motivation

给定输入特征，注意力机制将预测一个注意力图。比如，通道注意力将生成一个1D注意力图；空域注意力将生成一个2D注意力图；通道-空域注意力将生成一个3D注意力图。很自然的，我们提出了两个问题：

图像的每一部分具有更高或者更低的注意力系数呢？
是否注意力机制总是有益于超分模型呢？

Attention Heatmap

LR空间中包含冗余的低频成分以及少量的高频成分。RCAN一文认为：无注意力的模型会对所有特征均等对待，而注意力有助于网络对高频特征赋予更多的注意。然而，很少有工作能够证实上述假设。

为回答上述所提第一个问题，我们通过实验来理解超分中注意力机制的行为。我们构建了一个包含10个注意力模块的网络，每个注意力模块采用通道和空域注意力层，因此每个像素具有独立的注意力系数。

上图给出了某些特征与注意力图的可视化效果，上表给出了注意力图与高通滤波之间的相关系数。尽管这种度量方式无法精确度量注意力响应，但我们的目的是量化不同层之间的相对高通相关性。

从上图&上表可以看到：不同层学习到的注意力变化非常大。比如第一个注意力模块与第十个注意力模块表现处了截然相反的响应，意味着：低层的注意力模块倾向于低频模式，高层的注意力模块倾向于高频模块，中间的模块则具有混合响应。

Attention Dropout

基于上述发现，我们尝试最大限度的减少注意力的使用，同时最小化额外参数量。一个直观的想法：仅在性能关键层保留注意力层。然而，上述分析并不是一种有效的度量注意力层有效性的方案。

为定量度量注意力层的有效性，我们提出了Attention dropout框架。我们通过关闭特定注意力层进行了一系列对比实验，结果见下表。

从上表可以看到：模块深度很大程度影响了注意力模块插入位置。该结果进一步验证了：全网络均匀的设置注意力是次优方案。

Method

已有采用了固定注意力层方案(比如RCAN、PANet)总体同时激活所有注意力图且与图像内容无关。前述实验表明：注意力层的有效性会随位置发生变化。这启发我们构建一个非注意力短连接分支与一个注意力分支并采用动态权值进行混合。

Network Architecture

上图给出了本文所提A2N整体架构示意图，它包含三部分：

浅层特征提取：我们采用单个卷积提取浅层特征，;
堆叠A2B模块：我们采用链式堆叠A2B提取深层特征：
重建模块：完成深度特征提取后，我们通过重建模块进行上采样。在重建模块中，我们首先采用最近邻插值上采样，然后在两个卷积层中间插入一个简化版通道-空域注意力。
全局连接：我们同时还采用了全局连接：对输入通过最近邻插值上采样，然后与上述重建结果相加：

Attention in Attention Block(A2B)

尽管前面我们已经讨论了不同注意力层的动态贡献，但仍然难以手动确定注意力模块的拓扑结构。受启发于动态核，我们提出了可学习ADM以自动丢弃某些不重要的注意力特征，平衡注意力分支于非注意力分支。具体来说，每个ADM采用加权方式控制注意力分支与非注意力分支的动态加权贡献。

如上图c所示，ADM根据输入为不同分支生成了动态加权权值，可以描述如下：

动态加权权值计算方式如下：

上图给出了A2B结构示意图，ADM首先对输入进行均值池化，然后送入两个全连接层+Softmax生成动态加权权值。正如动态卷积一文所提到：约束动态权值可以促进ADM的学习。具体来说，我们采用了sum-to-one约束，即。这种约束方式可以压缩核空间，极大的简化的学习。

Experiments

实验方面采用了常规训练数据DIV2K，测试数据Set5、Set14、B100、Urban100、Manga109等。度量准则采用了PSNR与SSIM。

A2B中的同都数位40，上采样模块中通道数为24。在A2N的实现方面，均包含16个A2B模块；对于A2B模块，我们设计了两种方式：

A2N：非注意力分支采用卷积；
A2N-M：非注意力分支采用卷积。

上表对比了不同A2结构的性能，可以看到：

仅仅保留注意力积分制，模型的性能为28.646dB；采用A2结构后，性能可以提升0.05dB，额外引入参数量为200K。
当我们将通道数降到32时，模型仍有更好的性能，且参数量节省132K。又一次证实：并非所有注意力层都起正向作用。

上表对比了A2结构在多分枝模型中的性能对比，1-4为双分支结构且融合方式为Add；5-8为A2结构。可以看到：

在1-4类型中，非注意力与通道-空域注意力的组合具有最佳的性能；
相比3，7中的注意力与非注意力组合空域带来0.05dB性能提升，仅需额外7K参数。

上表对比了非注意力分支在不同配置下的性能对比，可以看到：

哪怕不进行任务操作，A2结构也可以取得比单个注意力分支更好的性能；
相比单分支注意力结构，A2+组合可以带来0.049dB性能提升，仅需额外的33K参数。这种配置可以取得更好的参数量-性能均衡。
对比no-op与，非注意力分支的卷积对网络有正向贡献，它可以提取注意力分支所无法提取的有效特征。

上表给出了所提方案与SRCNN、FSRCNN、DRRN、VDSR、IMDN、CARN等模型的性能对比，可以看到：

相比其他方案，所提A2N可以取得相当、甚至更好的结果；
相比A2F-M与AWSRN-M，所提A2N-M取得了更佳的性能，同时节省200K参数量。

上图对比了不同方案的在x4任务上的视觉效果对比。可以看到：所提方案取得了更好的视觉效果，可以更精确的重建高频细节。