论文阅读笔记：Transcending the Limit of Local Window: Advanced Super-Resolution Transformer with ATD

论文：https://arxiv.org/abs/2401.08209

代码：https://github.com/labshuhanggu/adaptive-token-dictionary

1 背景

1.1 问题

Transformers based SR方法的一个最主要问题是如何平衡令人满意的SR精度和增加的计算复杂度。由于自注意力机制的二次计算复杂度，先前的方法被迫将注意力限制在局部窗口管理计算负载。然而这种基于窗口的方法对感受野施加了约束，影响了性能。虽然最近的研究表明扩大窗口大小改善了感受野，提高了SR性能，但加剧了维数灾难。这个问题强调了需要一种有效的方法来建模长程依赖关系，而不受限于局部窗口，此外传统的图像SR通常采用不考虑图像内容的通用计算，与采用基于矩形局部窗口的划分策略相比，根据图像的具体内容进行划分可能更有利于SR过程。

1.2 提出的方法

本文从超分中经典的字典学习中得到启发，引入令牌字典，增强了图像处理中的交叉和自注意力计算。这种令牌字典提供了三个不同的好处，

• 首先，他能够利用交叉注意力将外部先验融入到图像分析中，这是通过学习概括常见图像结构的辅助令牌来实现的。便于以图像大小成线性比例的复杂度进行高效处理

• 其次，他能够利用全局信息建立长连接，这是通过激活令牌来精炼字典，以通过注意力的反向形式全局地总结图像特定的信息来实现的。

• 最后，他可以利用图像的所有相似部分来增强图像令牌，而不受局部窗口分区的限制。这是通过基于内容的结构划分，根据图像和字典令牌之间的相似性来实现基于类别的自注意力。
  

2 创新点

• 我们引入了令牌字典的思想，它利用一组辅助令牌为每个图像令牌提供先验信息，并以交叉注意力的方式有效地从整幅图像中总结先验信息。

• 我们利用令牌字典将图像令牌分组，并突破局部窗口的边界，以基于类别的自注意力方式利用远程先验。

• 通过结合所提出的令牌字典交叉注意力和基于类别的自注意力，我们的模型可以有效地利用长距离依赖并取得了优于最先进方法的超分结果。

3 方法

通过提出的令牌字典交叉注意力( TDCA )，自适应字典细化( ADR )策略和基于类别的自适应多头自注意力( ACMSA )，来建立本文提出的自适应令牌字典( ATD )网络并用于图像超分辨率。

如图3所示，给定一幅输入的低分辨率图像，首先利用一个3 × 3的卷积层来提取浅层特征。然后将浅层特征经过一系列ATD块，其中每个ATD块包含若干个ATD Transformer层。将令牌字典交叉注意力，基于自适应类别的多头自注意力以及常用的基于移动窗口的多头自注意力结合起来构成Transformer层，这三个注意力模块并行工作，以利用输入特征的外部，全局和局部特征，并通过求和操作对特征进行组合。除了注意力模块，ATD中还是用了Transformer中常用的Layer Norm和FFN层。此外令牌字典开始于每个ATD块内的可学习参数，参与每个Transformer层的令牌字典交叉注意力，并利用自适应字典细化策略使字典适应下一层的输入特征。

在ATD块之后，使用一个额外的卷积和一个PixelShuffle层来生成最终的高分辨率图估计。

4 模块

4.1 预备知识

4.1.1 字典学习

传统的基于字典学习的方法将图像划分成块来建模图像的局部先验，记 $x\in R^d$ 为低分辨率图像中的图像块向量。为了估计相应的高分辨率块 $y\in R^d$，可通过求解系数表示问题来分解信号：

然后用 KaTeX parse error: Undefined control sequence: \* at position 11: D_H\alpha^\̲*̲ 来重构 HR 块。其中 $D_L\in R^{d\times M}$ 和 $D_H\in R^{d\times M}$ 是学习到的 LR 和 HR 字典，M是字典中的基元个数。大多数基于字典学习的SR方法通过学习耦合字典 $D_L$ 和 $D_H$ 来总结来自外部训练数据集的先验信息。为了获得更好的SR结果，还可以根据测试图像进行一些字典优化。

参考论文《Image super-resolution via sparse representation》。

稀疏表示是指，使用过完备字典中少量向量的线性组合来表示某个元素。过完备字典是一个列数大于行数的行满秩矩阵，也就是说，它的列向量有无数种线性组合来表达列向量空间中的任意点。由于它的列数通常远大于行数，可以使用占比很小的列向量来表示特定的向量，我们称这种表示为稀疏表示。

训练过完备字典的过程称为稀疏编码。设训练数据集为 $X=(x_1,x_2,\dots x_n)\in R^{m\times n}$，待训练矩阵为 $A\in R^{m\times K}$，矩阵对每一行数据的表示权重 $\alpha =({\alpha }_1,{\alpha }_2,\dots {\alpha }_n)\in R^{K\times n}$。则进行如下优化：

即约束重建向量和原始向量之间的差距，优化表示权重 $\alpha$ 的稀疏性。通过优化获得所需的过完备字典后，就可以用它来表示新的数据，对于某个向量 $y$，可以进行类似的优化来获得它在字典中的稀疏表示：

为了方便优化，对上式进行调整：

在训练阶段，同时对低分辨率数据集 $Y=(y_1,y_2,\dots y_n)$ 和高分辨率数据集 $X=(x_1,x_2,\dots x_n)$ 训练完备字典 $D_L,D_H$，使得LR数据 $y_i$ 和HR数据 $x_i$ 能够以相同的稀疏编码 ${\alpha }_i$ 分别被 $D_L$ 和 $D_H$ 表示，即 $D_L{\alpha }_i=y_i,D_H{\alpha }_i=x_i$。

在测试阶段，已经有了训练好的D_L和D_H。对于测试图像 $y_t$，首先通过优化 $y_t$ 在 $D_L$ 中的稀疏表示 ${\alpha }_t$，此时有 $D_L{\alpha }_t\approx y_t$。然后用这个表示通过 $D_H$ 来映射出对应的SR图像，即 $\widehat{x_t}=D_H{\alpha }_t$。

4.1.2 ISR中的VIT

Transformer中的核心是自注意力操作，利用令牌之间的相似性作为权重来相互增强图像特征：

然而自注意力计算的复杂程度随着输入令牌数呈二次方增长，不同的策略包括移位窗口，锚点注意力和移位交叉注意力被提出来缓解窗口大小受限问题。

4.1.3 使用令牌字典的交叉/自注意力

作者发现基于字典学习的图像SR的分解和重建思想与自注意力计算的过程类似。具体来说，式（1）在求解稀疏表示模型，寻找相似的LR字典基元，并与对应的HR字典基元重构HR信号，而基于注意力的方法则使用归一化的点积运算来确定注意力权重并组合value令牌。

基于上述观察，字典学习的思想可以很容易的融入到Transformer框架中，打破局部窗口的限制。

4.2 令牌字典交叉注意力

现有的多头自注意力通过输入的特征自身来生成query，key和value。作者旨在引入一个额外的字典 $D\in R^{M\times d}$，并将其初始化为网络参数，以便在训练阶段总结外部先验。利用学习到的令牌字典 $D$ 生成 Key 字典 $K_D$ 和 Value 字典 $V_D$，利用输入特征 $X\in R^{N\times d}$ 生成 Query:

其中 $W^Q\in R^{d\times d/r},W^K\in R^{d\times d/r},W^V\in R^{d\times d}$是query令牌，key字典令牌和value字典令牌的线性变换。令 $M<<N$ 来保持较低的计算成本。同时，query令牌和key字典令牌的特征维度降为 $1/r$，以减小模型规模和复杂度，其中 $r$ 为降低比例。然后使用key字典和value字典通过交叉注意力计算来增强查询令牌：

式（4）中，$\tau$ 为可学习参数，用于调整相似度取值范围，$Si{m}_{cos}$ 表示计算两个令牌之间的余弦相似度。$S=Si{m}_{cos}(Q_X,K_D)\in R^{N\times M}$是query图像令牌和key字典令牌之间的相似性图。使用归一化余弦距离代替MSA中的点积运算，因为作者希望字典中的每个令牌都有平等被选中的机会，然后使用SoftMax函数将相似度图$S$转换成注意力图 $A$ 用于后续计算。

4.3 自适应字典细化

由于每一层的图像特征被多层感知机MLPs投影到不同的特征空间，因此需要为每一层不同的令牌字典，以提供每个特定特征空间中的外部先验。这将导致大量的附加参数，作者引入了一种自适应的细化策略，该策略基于相似度图对前一层的令牌字典进行细化，并以反向注意力的形式更新特征。

为了引入所提出的自适应细化策略，作者为输入特征和令牌字典建立层索引 $(l)$ ，即 $X^{(l)}$ 和 $D^{(l)}$ 分别表示第 $l$ 层的输入特征和标记字典，只对初始层 $D^{(1)}$ 建立令牌字典作为网络参数来融入外部先验知识。在接下来的层中，每个字典 $D^{(l)}$ 都是在上一层 $D^{(l-1)}$ 的基础上改进的。对于第 $l$ 层的字典 { d i ( l ) } i = 1 , … , M \{d_i^{(l)}\}_{i=1,…,M} {di(l)​}i=1,…,M​ 中的每一个令牌，选择增强特征 $X^{(l+1)}$中相应的相似令牌，即用第 $l$ 层的输出对其进行细化。具体的，用 $a_i^{(l)}$ 表示注意力图 $A^{(l)}$ 的第 $i$ 列，它包含了 $d_i^{(l)}$ 与所有的 N 个query令牌 $X^{(l)}$ 之间的相似注意力。因此，基于每个 $a_i^{(l)}$，可以选择相应的增强令牌 $X^{(l+1)}$ 来重构新的令牌字典元素 $d_i^{(l)}$，并将他们组合起来形成 $D^{(l)}$：

相当于把 $l$ 层的字典令牌作为 query，$l$ 层的输入特征作为 key，$l$ 层的输出特征（也是 $l+1$ 层的输入特征）作为 value 来更新字典令牌。

其中 $Norm$ 为归一化层，用于调整注意力图的范围。这种细化可以看做是式（4）的反向形式。然后基于一个可学习的参数 $\sigma$，自适应地组合 ${\hat{D}}^{(l-1)}$ 和 $D^{(l-1)}$ 得到 $D^{(l)}$。通过这种方式，细化的令牌字典能够同时整合输入图像的外部先验和特点的内部先验。

4.4 基于类别的自适应注意力

为了更好的利用自注意力，自适应特征划分可能是一个更合适的选择，由于TDCA得到的输入特征与令牌字典之间的注意力图隐含了每一个像素的类别信息，因此可以对输入特征进行分类，根据字典令牌相似度图将每个像素分类不同的类别 ${\theta }^1,{\theta }^2\dots {\theta }^M$：

式中 $A\in R^{N\times M}$ 为式（4）中计算出的相似度图。如果像素 $A_{ji}$ 在 $A_{j1},A_{j2},...A_{jM}$中最高，这表明像素 $x_j$ 更可能与字典中的第 $i$ 个令牌 $d_i$ 属于同一类别，则将像素 $x_j$ 归为 ${\theta }^i$ 类。因此，每个类别都可以被看作是一个不规则形状的窗口，其中包含了相同类别的令牌。然而每个类别的令牌数量可能不同，这导致并行效率低和计算负担显著。为了解决不平衡分类问题，作者将 $\theta$ 进一步划分为子类 $\varphi$：

其中，类别 ${\theta }^i$ 包含了 $n_i$ 个令牌。将每个类别扁平化并串联成 $\Phi$，再划分子类别 $[{\varphi }^1,{\varphi }^2,...{\varphi }^j,...]$。划分后，所有子类具有相同的固定大小 $n_s$。如下图B.2。

AC-MSA可以被表示为：

其中 $Categorize$ 运算为式（6）和式（7）的组合。将每个子类别视为一个注意力组，并在每个组内执行多头自注意力。最后用 $UnCategorize$ 反分类操作将每个像素返回他原来的位置，如图。

虽然子类别划分限制了每个注意力组的大小，但只对感受野有轻微影响，这是因为式（7）中划分子集时的排序操作产生了随机打乱的效果，因此每个子类别都可以看作该类别的随机样本，某个子类别的令牌仍然可以遍布整个特征图，保持全局感受野。

5 效果

5.1 消融实验

各组件的消融实验。

基于类别的自注意的消融实验。使用随机令牌字典分类效果比基线好，使用学习到的自适应令牌字典可以更准确执行分类。

子类别尺寸ns大小和字典大小的消融实验。

5.2 和SOTA方法对比

和SOTA方法在各数据集上的指标对比。

和轻量级SOTA方法在各数据集上的指标对比。

和最近SOTA方法的性能和计算量对比。

和SOTA方法的可视化效果对比。

基于类别的自适应多头自注意力中每个类别所在位置。