HAT: Hybrid Attention Transformer for Image Restoration

目录

1. 摘要

2. 现阶段问题

3. Motivation

4. 主要贡献

5. 网络架构

6. 不同网络的LAM结果

7. SwinIR与HAT的中间特征可视化

8. 注意力

---

1. 摘要

        通过归因分析attribution analysis method - Local Attribution Map (LAM)，发现目前基于Transformer的方法只能来利用有限的输入空间信息。这意味着 Transformer 的潜力在现有网络中仍未得到充分利用。为了激活更多的输入像素以获得更好的恢复，提出了一种新的混合注意Transformer(HAT)。它结合了通道注意力和基于窗口的自注意力机制，从而利用了它们的互补优势。此外，为了更好地聚合窗口之间信息，引入了一个重叠的交叉注意模块来增强相邻窗口特征之间的交互。在训练阶段，采用了相同的任务预训练策略来进一步利用模型的潜力进行进一步改进。

2. 现阶段问题

• SwinIR 比基于 CNN 的方法（例如 RCAN [10]）在某些场景利用的输入像素更少。

• 在 SwinIR 的中间特征中会出现阻塞伪影。这表明移位窗口机制不能完美地实现跨窗口信息交互。

3. Motivation

        借助LAM诊断工具，现 SwinIR（基于Transformer的方法）在利用的信息范围上并不比基于 CNN 的方法（如RCAN）更广泛。这与常识相矛盾，但也给作者带来了额外的启示。

        首先，这表明 SwinIR 比 CNN 具有更强的映射能力，因此可以利用较少的信息实现更好的性能。其次，由于利用的像素范围有限，SwinIR 可能会恢复错误的纹理，因此如果能够利用更多输入像素，它的性能可能进一步提升。因此，研究人员的目标是设计一种网络，既能够利用类似的自注意力机制，又能够激活更多像素进行重建。他们的 HAT 网络在图中几乎可以看到整个图像，并且能够恢复正确和清晰的纹理。

4. 主要贡献

        提出了一种混合注意力 Transformer，即 HAT。结合通道注意和自注意力机制，以利用前者transformer使用全局信息的能力（激活更多像素）和后者Self-attention强大的representative ability

        引入了一个overlapping的注意力窗口模块来实现相邻窗口特征的更直接的交互，减少阻塞伪影产生。

5. 网络架构

步骤：

HAT模块：

• 总体结构：浅层特征提取——深层特征提取——图像重建（Pixel-Shuffle），L1损失

• 残差混合注意力组（RHAG）深层特征提取的基本单元

• 混合注意力块（HAB）在STL的基础上添加了CAB，修改MSA为(S)W-MSA

• 重叠交叉注意力块（OCAB）在STL的基础上将MSA替换为OCA（基于重叠窗口分区计算，查询更大窗口）

HAT的整体架构： 

        浅层特征提取：使用一个3x3的卷积层HConv(·)从低质量输入图像ILQ中提取浅层特征F0。

        深层特征提取：通过HDF(·)进行深层特征提取，HDF(·)由N1个RHAG和另一个3x3的卷积层HConv(·)组成。每个RHAG逐步处理中间特征，生成更深层次的特征表示。

        图像重建：通过融合浅层特征和深层特征，使用重建模块HRec(·)重建出高质量图像IHQ。

RHAG的结构：
       RHAG是HAT的关键组件，包含多个混合注意力块（HAB）、一个重叠交叉注意力块（OCAB）以及一个带有残差连接的3x3卷积层。RHAG的设计旨在通过混合注意力机制和交叉注意力增强模型对图像特征的理解和表示能力。

HAB的结构：
        HAB结合了通道注意力和基于窗口的自注意力机制，以激活更多的输入像素，从而提高图像重建的质量。HAB的设计采用了类似于标准Swin Transformer块的结构，并保留了基于窗口的自注意力机制。通过引入通道注意力块，HAB能够更好地利用全局信息，增强自注意力的代表性能力。

OCAB (Overlapping Cross-Attention Block)：
        OCAB是HAT网络中用于增强窗口之间信息交互的组件。它通过重叠的窗口分割来计算交叉注意力，这允许网络在不同的特征窗口之间建立联系，从而更好地理解图像的整体结构。OCAB有助于减少由于窗口分割导致的信息隔离问题，使得网络能够更全面地利用输入图像的信息。

CAB (Channel Attention Block)：
        CAB是HAT网络中用于动态调整特征通道重要性的模块。它通过全局信息来学习每个通道的权重，从而自适应地调整通道特征的响应。CAB有助于网络集中注意力于更重要的特征，并抑制不那么重要的信息，这对于提高图像恢复的质量至关重要。
        CAB由两个标准的卷积层和一个GELU（Gaussian Error Linear Units）激活函数组成，后面跟着一个通道注意力（Channel Attention, CA）模块。GELU激活函数是一种在深度学习中常用的非线性激活函数，它可以增加模型的表达能力。通道注意力模块的作用是动态调整不同通道的特征响应，使模型能够关注更重要的特征。
 

OCA：

OCA的重叠窗口划分：
        OCA是为了直接建立相邻窗口间的连接并增强窗口自注意力的代表性能力而设计的。在OCA中，输入特征X被划分为不同的窗口，其中查询（XQ）、键（XK）和值（XV）被分别划分为不重叠和重叠的窗口。
        XQ被划分为大小为MxM的非重叠窗口，每个窗口包含MxM个像素。XK和XV被划分为大小为Mo x Mo的重叠窗口，其中Mo = (1 + γ) x M，γ是一个控制重叠大小的常数。这意味着每个重叠窗口会覆盖更多的像素，从而在计算自注意力时可以利用更多的上下文信息。通过零填充（zero-padding）和相对位置偏置（relative position bias）来确保窗口大小的一致性，并计算注意力矩阵。
        OCA的设计允许模型在更大的感受野内捕捉信息，从而提高图像重建任务的性能。通过这种方式，HAT能够有效地激活更多的输入像素，实现更高质量的图像重建。
重叠：
        与非重叠窗口分割不同，重叠窗口分割允许相邻的窗口之间有重叠的部分。这意味着每个窗口并不完全独立，它们在边缘处共享一些像素。这种设计允许模型在处理一个窗口时，同时考虑到相邻窗口的信息，从而提供更连贯和全局的特征表示。

计算注意力：
        在OCA中，模型将计算每个窗口内部的自注意力，同时也会计算跨越窗口边界的交叉注意力。交叉注意力允许模型在不同的窗口之间建立联系，这对于理解图像的整体结构和上下文非常重要。

6. 不同网络的LAM结果

        对于基于CNN的方法（EDSR,RCAN），利用更多的输入像素，其PSNR越高，重建效果较好。 但对于基于transformer的SwinIR方法，并没有利用比RCAN更广的范围，但是其PSNR效果更好。

7. SwinIR与HAT的中间特征可视化

可以观察到SwinIR中间特征中明显的阻塞伪影，如图3所示。

这些伪影是由窗口划分机制引起的，这表明移位窗口机制在构建跨窗口连接是低效的。

这也可能是 SwinIR 在图 2 中 不使用更多像素进行重建的原因。

（换句话说，由于窗口之间的连接不够充分，模型可能无法有效地利用整个图像的信息来进行准确的图像重建。）

8. 注意力

通道注意力（Channel Attention）

• 通道注意力机制主要关注的是输入特征的通道维度。在深度学习模型中，每个像素点通常会被表示为一个多通道的特征向量，其中每个通道可能捕捉到不同的空间信息。

• 通道注意力的目标是识别出对于当前任务最重要的通道，并给予它们更多的权重，同时降低不重要通道的影响。这样做可以增强模型对关键特征的敏感性，从而提升特征的表达能力。

• 在HAT中，通道注意力通过全局统计信息来实现，这意味着它会考虑到整个图像范围内的特征分布，而不是局限于局部区域。这有助于模型捕捉到全局上下文信息，对于图像恢复尤其重要，因为全局信息可以帮助生成更加连贯和准确的输出。

基于窗口的自注意力（Window-Based Self-Attention）

• 自注意力机制允许模型在序列的不同位置之间建立直接的联系，而窗口化的自注意力则是将输入图像分割成较小的局部窗口，并在每个窗口内独立地计算自注意力。

• 这种方法可以减少计算量，因为不需要考虑图像中所有像素之间的相互作用，只需考虑局部窗口内的相互作用即可。这对于大型图像或高分辨率图像尤其有效，因为它减少了模型的计算负担。

• 窗口化的自注意力有助于模型捕捉局部细节和纹理信息，这对于图像恢复任务（如去噪、超分辨率）来说是非常重要的。通过在局部窗口内计算注意力权重，模型可以更好地理解图像的局部结构，并在恢复过程中保留这些细节。

        在HAT模型中，这两种机制被结合起来，形成了一种混合注意力（Hybrid Attention）框架。通道注意力提供了全局信息和上下文感知，而窗口化的自注意力则专注于捕捉和恢复图像的局部特征。这种混合方法使得HAT能够有效地平衡全局和局部的视觉信息，从而在图像恢复任务中取得了优异的性能。