1 背景
1.1 问题
虽然Transformer based的超分模型效果比CNN好,但是原因还是不得而知。一个直观的解释是,这种网络可以受益于自注意力机制并利用远距离信息。
因此作者采用归因分析方法LAM来考察SwinIR中用于重建的信息所涉及的范围,发现SwinIR并没有比基于CNN的方法在超分辨率上利用更多的输入像素,如图2所示。
此外,虽然SwinIR获得了更高的量化性能,但由于使用的信息范围有限,在一些样本中产生了不如RCAN的结果,这些现象说明Transformer对局部信息的建模能力更强,但其利用信息的范围还有待扩大。作者还发现在SwinIR的中间特征会出现块状伪影,如图3所示。论证了平移窗口机制无法完美实现跨窗信息交互。
1.2 提出的方法
为了解决上述问题,进一步挖掘Transformer在超分辨率重建中的潜力,本文提出了一种混合注意力Transformer,即 HAT。
HAT结合了通道注意力和自注意力机制,以利用前者获取全局信息的能力和后者强大的表征能力。
此外,引入重叠交叉注意力模块,以实现相邻窗口特征更直接的交互,受益于这些设计,模块可以激活更多的像素重建,从而获得更显著的性能提升。
2 创新点
设计了一种新颖的混合注意力Transformer( Hybrid Attention Transformer,HAT ),将自注意力、通道注意力和一种新的重叠交叉注意力相结合,以激活更多的像素,从而更好地进行重建。
提出了一种有效的同任务预训练策略来进一步挖掘SR Transformer的潜力,并表明了大规模数据预训练对于该任务的重要性。
方法达到了最先进的性能。通过进一步扩展HAT构建大模型,极大地扩展了SR任务的性能上界。
3 方法
整体网络由3部分组成,包括浅层特征提取、深层特征提取和图像重建。对于给定的低分辨率输入 I L R ∈ R H × W × C i n I_{LR}∈R^{H×W×C_{in}} ILR​∈RH×W×Cin​,首先利用一个卷积层来提取浅层特征 F 0 ∈ R H × W × C F_0∈R^{H×W×C} F0​∈RH×W×C,其中 C i n C_{in} Cin​ 和 C C C 是输入图像和中间特征的通道数。然后利用一系列残差混合注意力组 RHAG 和一个 3×3 卷积层 H C o n v ( ⋅ ) H_{Conv}(·) HConv​(⋅) 进行深度特征提取。之后,使用一个全局残差来融合浅层特征 F 0 F_0 F0​ 和深层特征 F D ∈ R H × W × C F_D∈R^{H×W×C} FD​∈RH×W×C,在最后通过重建模块对高分辨率结果进行重建,如图4。
4 模块
4.1 混合注意力模块(HAB)
如图2中所示,当采用通道注意力时,更多的像素被激活,因为全局信息参与计算通道注意力权重。此外,还有很多工作表明卷积可以帮助Transformer获得更好的视觉表示或实现更简单的优化。因此,作者在标准Transformer块中融入基于通道注意力的卷积块来增强网络的表达能力。
如图4所示,在第一个Layer Norm层后的标准Swin Transformer块中并联一个通道注意力模块CAB,该模块与基于窗口的多头自注意力模块W-MSA并联,在连续的HAB块中,间隔使用基于移动窗口的自注意力模块SW-MSA。为了避免CAB和MSA在优化和可视化表达上可能存在冲突,在CAB的输出上乘以一个较小的常数
α
\alpha
α。对于给定的输入特征
X
X
X,整个HAB的计算过程如下:
其中
X
N
X_N
XN​ 和
X
M
X_M
XM​ 表示中间特征。
Y
Y
Y 表示HAB的输出。将每一个像素看作是一个块的嵌入的token。给定输入特征
H
×
W
×
C
H×W×C
H×W×C,现将其划分成
H
W
M
2
\frac{HW}{M^2}
M2HW​ 个尺寸为
M
×
M
M×M
M×M 的局部窗口,然后在每个窗口内部计算自注意力:
其中 d d d 表示 query 和 key 的维度。 B B B 表示相对位置编码。作者发现扩大窗口大小可以显著的扩大使用像素的范围,于是使用了一个大的窗口大小来计算自注意力。同时为了建立相邻非重叠窗口之间的联系,作者还是用窗口移动划分方法,并将窗口移动大小设置为窗口大小的一半。
CAB由两个带有GELU激活的标准卷积层和一个通道注意力模块组成,如图4所示。由于基于Transformer的结构往往需要较多的嵌入,直接使用宽度不变的卷积会产生较大的计算开销,所以作者使用一个常数 β \beta β 来压缩两个卷积层将通道数。对于一个具有 C C C 个通道的输入特征,将第一个卷积层后的输出特征的通道数压缩为 C β \frac{C}{\beta} βC​。然后通过第二层将输出通道扩展为 C C C。其次,利用一个标准的CA模块自适应的缩放通道特征。
4.2 重叠交叉注意力模块(OCAB)
作者引入OCAB直接建立跨窗口连接,增强窗口自注意力的表示能力。OCAB类似于一个标准的Swin Transformer,由一个重叠交叉注意力层OCA和MLP层组成。但是对于OCA,如图5所示,作者使用不同的窗口大小对投影后的特征进行划分。具体来说,对于输入特征
X
X
X 的
X
Q
,
X
K
,
X
V
∈
R
H
×
W
×
C
X_Q,X_K,X_V∈R^{H×W×C}
XQ​,XK​,XV​∈RH×W×C,
X
Q
X_Q
XQ​ 被划分成大小为
M
×
M
M×M
M×M 的
H
W
M
2
\frac{HW}{M^2}
M2HW​ 个非重叠区域,而
X
K
,
X
V
X_K,X_V
XK​,XV​ 被划分成大小为
M
o
×
M
o
M_o×M_o
Mo​×Mo​ 的
H
W
M
2
\frac{HW}{M^2}
M2HW​ 个重叠窗口,其计算公式为:
其中 γ \gamma γ 是控制重叠尺寸的常数。为了更好的理解这个操作,标准的窗口划分可以为认为是一个滑动划分,其核大小和步长都等于窗口大小 M M M,相比之下,重叠窗口划分可以看做是一个滑动划分,其核大小等于等于 M o M_o Mo​,而步长等于 M M M。窗口超出的尺寸采用补零的方式来填充。计算注意力矩阵的方式同式2,同样采样相对位置偏差 B 属于 R M × M o B属于R^{M×M_o} B属于RM×Mo​。与WSA不同的是,WSA的query,key和value都是通过相同窗口特征的计算得到,而OCA从更大的视野中计算key和value,query可以利用更多有用的信息。
4.3 同任务预训练
预训练在许多高级视觉任务上被证明是有效的。最近的工作也证明了预训练对低级视觉任务是有益的。IPT 强调使用各种低级任务,如去噪、去雨、超分辨率等,而EDT则利用特定任务的不同退化程度进行预训练。这些工作集中于考察针对某一目标任务的多任务预训练的效果。相比之下,作者基于同样的任务直接在更大规模的数据集(即ImageNet )上进行预训练,表明预训练的有效性更依赖于数据的规模和多样性。例如,当我们想要训练一个× 4 SR的模型时,我们首先在ImageNet上训练一个× 4 SR的模型,然后再对其进行微调,如在DF2K上。所提出的策略,即同任务预训练,更简单的同时带来更多的性能提升。值得一提的是,足够的训练迭代次数进行预训练和合适的小学习率进行微调对于预训练策略的有效性非常重要。我们认为这是由于Transformer需要更多的数据和迭代次数来学习任务的一般性知识,但需要较小的学习率进行微调,以避免对特定数据集的过拟合。
5 效果
5.1 消融实验
不同窗口尺寸的消融实验和可视化效果对比。
提出的OCAB和CAB模块的消融实验和可视化效果。
同任务预训练策略消融实验的效果。
5.2 和SOTA方法对比
和SOTA方法的指标对比。
和SOTA方法的可视化效果对比。