搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（十二）

本文目录

27 low-level多个任务榜首被占领，中科大等联合提出：Uformer
(来自中科院，中科大，刘健庄老师团队)
27.1 Uformer 原理分析
27.1.1 IPT 模型和数据集构建方法概述
27.1.2 Uformer
27.1.3 LeWin Transformer Block
27.1.4 Skip-Connection的几个变种

28 Refiner：改进视觉Transformer的自注意力
(来自 新加坡国立大学，冯佳时老师团队)
28.1 Refiner原理分析
28.1.1 Attention Expansion
28.1.2 Distributed Local Attention (DLA)
28.2 Refiner代码解读

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

本文介绍的第 1 篇文章 Uformer 是视觉Transformer 在底层视觉任务 (Low-Level Tasks) 的新探索，来自中科院刘健庄老师团队。这个 Uformer 适合 4 种底层视觉任务：去噪 (denoising), 去雨 (deraining), 去模糊 (deblurring) 和去摩尔纹 (demoireing)。

Transformer应用于 Low-Level Tasks的上一个有意义的工作是来自北京华为诺亚实验室的 CVPR 2021的工作 IPT: Image Processing Transformer，IPT 适合 3 种底层视觉任务：去噪 (denoising), 去雨 (deraining), 超分 (SISR)。详细内容可以参考下面的链接：

搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（二）

本文介绍的第 2 篇文章 Refiner 是对自注意力的扩展进行了探索，将多头注意力图映射到一个更高的维度空间，以促进其多样性。

27 Low-Level多个任务榜首被占领，中科大等联合提出：Uformer

论文名称：Uformer: A General U-Shaped Transformer for Image Restoration

论文地址：

https://arxiv.org/pdf/2106.03106.pdf

27.1 Uformer原理分析：

要把 Transformer 模型应用在底层视觉任务里面其实是一件很难的事，因为输出要和输入同等大小的尺寸，且要兼顾不同的复原任务 (去雨，去噪，去模糊等等) 。我们首先回顾一下 IPT 是怎么做到的：

27.1.1 IPT 模型和数据集构建方法概述

IPT的整体结构可以分为4个部分：

：从corrupted images中提取特征。由 3层卷积 组成，得到  ，这一步可以表达为：  ，其中  表达第几个任务。

：主要结构，包含Encoder和Deocder，用于恢复丢失的数据。为了区分不同的任务，作者使用了基于特定任务的嵌入向量 Task-specific Embeddings ，这里的  代表 3 种不同的任务。它们都是可学习的向量，目的是学习到不同任务的不同表达。

：最后映射特征。

IPT 的训练数据集构建：

基于ImageNet， 使用几个基于特定任务的不同的操作来生成多个损坏的副本，以服务于不同的任务。原始的ImageNet具有 1M 数据，把它们crop成  的patches。为了得到损坏的副本，作者进行了6种不同的退化(Degradation)操作。

对于超分任务，就对原始ImageNet图片分别进行  下采样。

对于去噪任务，就对原始ImageNet图片分别添加30, 50 noise level Gaussian noise。

对于去雨任务，就对原始ImageNet图片添加雨线。

这样一来，用来训练IPT的整个数据集包含大约1000多万幅图像。

测试数据集：

同样把特定任务对应的测试集crop成  的patches，overlap 10个pixels。为了公平地比较，CNN模型也进行了这样的操作。

27.1.2 Uformer

Uformer 是 2 个针对图像复原任务 (Image Restoration) 的 Transformer 模型，它的特点是：

• 长得很像 U-Net (医学图像分割的经典模型)，有Encoder，有Decoder，且 Encoder 和 Decoder 之间还有 Skip-Connections。
• 在 Transformer 的Block中融合了卷积。

Uformer 不再需要像 IPT 那样通过 ImageNet 构造很大的数据集了，其使用的数据集和 CNN模型一致。

下图2是 Uformer 的架构示意，作者已经把张量的维度标注在了图里面。可以看到这个架构整体长得很像 U-Net，有Encoder，有Decoder，且 Encoder 和 Decoder 之间还有 Skip-Connections。

1) 卷积提特征：

输入一个Degraded Image  ，Uformer 首先通过  卷积和 LeakyReLU激活函数来提取底层特征，得到  。

2) Encoder：

接着  进入U形结构的Encoder部分，一共是  个阶段，每个阶段由2个 LeWin Transformer Block 和 1个下采样 Down-sampling 层组成。LeWin Transformer Block 首先把张量 reshape成1D的序列特征  ，它可以借助 Self-attention 机制的建模长序列信息的能力并通过 non-overlapping windows的办法节约计算量。下采样 Down-sampling 层首先把张量 reshape成2D的空间特征，再通过一个  的卷积操作进行下采样。经过这2步以后，第  个stage输出的张量维度是  ，第  个stage输出的张量维度是  。

3) Bottleneck Stage：

接着模型到了Bottleneck Stage，本质上是通过2个LeWin Transformer Block。

4) Decoder：

Decoder 用于特征的重建，一共是  个阶段，每个阶段由 1个上采样 Up-sampling 层 2个 LeWin Transformer Block 组成。上采样层也是先把张量 reshape成2D的空间特征，再通过  的 Transposed Convolution 执行，这个层可以减少特征的一半的channel数以及使得空间维度增加2倍。接着，这个上采样后的特征和来自Skip-Connection的 Encoder的特征一起输入 LeWin Transformer Block，通过它来恢复丢失的数据。  个阶段完成以后，通过一个3×3卷积得到 Residual Image  。

实验中设置  ，就是 Encoder和 Decoder都是4个阶段，损失函数是Charbonnier loss：

式中  是真值，  是个微小数值。

27.1.3 LeWin Transformer Block

接下来介绍下这个 LeWin Transformer Block 的机制。传统的 Transformer在计算 self-attention 时会计算全部 tokens 之间的attention，假设序列长度是  ，即有  个 tokens，Embedding dimension为  ，则计算复杂度是  。因为分类任务的图片大小一般是224×224，当我们把图片分块成16×16大小的块时，可以分成196个patches，token数就是196。

但是对于 Image Restoration 的任务一般都是高分辨率的图片，比如1048×1048的图片，此时得到的patches数量会很多，如果还计算全部 tokens 之间的attention的话，计算量开销会变得很大，显然不合适。所以第1个结论是：计算全部 tokens 之间的 attention 是不合适的。

另外，对于 Image Restoration 的任务local的上下文信息很重要。那么什么是local的上下文信息呢？就是一个patch的上下左右patch的信息。那么为什么local的上下文信息信息很重要呢？因为一个 degraded pixel 的邻近 pixel 可以帮助这个 pixel 完成重建。所以第2个结论是：每个 patch 的 local 的上下文信息很重要。

为了解决第1个问题，作者提出了一种 locally-enhanced window (LeWin) 的 Transformer Block，如上图2(b)所示。它既得益于 Self-attention 的建模长距离信息的能力，也得益于卷积网络的捕获局部信息的本领。

Window-based Multi-head Self-Attention (W-MSA)

这个模块是LeWin Transformer Block 的第1步。具体而言，它不使用全部的 tokens 来做 Self-attention，而是在不重叠的局部窗口 (non-overlapping local windows) 里面去做attention，以减少计算量。

输入2D的特征图  ，宽为  高为  ，我们把这个特征图分成大小为  的相互不重叠的 windows，如下图 3 我们有16个 windows。这样以后把每个 windows 的张量展平后得到的维度都是  。再对每个 window 执行 Self-attention。假设 head 的数量是  ，每个 head 的维度是  。则计算某个 window 中的第  个 self-attention 操作可以写为：

式中，  代表第  个 window，  代表第  个 window 的第  个 head 的结果。  代表第  个 head 的所有 windows 综合起来的结果。

第  个head 的queries, keys, 和 values 的映射矩阵。

在attention module里面仍然采用的是相对位置编码 (relative position encoding)，所以attention的表达式为：

式中，  是相对位置编码，其值来自于可学习的  矩阵。

问：Window-based Multi-head Self-Attention (W-MSA)能够节约多少计算量？

答： 假设特征图为  ，如果使用全局 Self-attention，则序列长度  ，计算量为  。如果使用全局 W-MSA，计算量为  ，注意  是每个windows的计算量，一共是  个windows。

Locally-enhanced Feed-Forward Network (LeFF)

在之前的 CvT 和 CeiT (链接如下) 中已经指出，Feed-Forward Network (FFN) 利用局部上下文信息的能力有限。但是相邻像素的信息却对图像复原的任务来讲很重要，因此作者仿照之前的做法在FFN中添加了 Depth-wise Convolution，如下图4所示。

搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（六）

首先特征通过一个 FC 层来增大特征的维度 (feature dimension)，再把序列化的1D特征 reshape 成 2D的特征，并通过 3×3 的 Depth-wise Convolution来建模 local 的信息。

再把2D的特征 reshape 成序列化的1D特征，并通过一个 FC 层来减少特征的维度到原来的值，每一个 FC 层或 Depth-wise 卷积层后使用 GeLU 激活函数。

综上，locally-enhanced window (LeWin) 的过程可以用下式表达，如图 2 (b) 所示。

• 27.1.4 Skip-Connection 的几个变种

为了研究如何把 Encoder 里面的 low-level 的特征更好地传递给 Decoder，作者探索了3种Skip-Connection的方式，分别是：

• Concatenation-based Skip-connection (Concat-Skip).
• Cross-attention as Skip-connection (Cross-Skip)
• Concatenation-based Cross-attention as Skip-connection (Concat Cross-Skip)

名字听上去非常花里胡哨，但其实很直觉，如下图5所示。

首先 Decoder 中的每个 LeWin Transformer Block 要综合2方面的信息，其一是来自 Encoder的输出  ，其二是来自 Decoder的输出  。

第1种 Concat-Skip 的做法就是直接把二者在 channel 这个维度给 concatenate 在一起，得到的张量的维度是  ，把它输入给 W-MSA 模块，如图 5 (a) 所示。

第2种 Cross-Skip 的做法就是给Decoder阶段的每个 LeWin Transformer block 里面加一个额外的attention module，这个 attention module 的 query 是第1个W-MAS的输出，key 和 value 都是  ，如图 5 (b) 所示。

第3种 Concat Cross-Skip 的做法就是把二者在 channel 这个维度给 concatenate 在一起，且W-MSA 的query 是  ，key 和 value 都是 concat 后的结果，如图 5 (c) 所示。

Experiments：

实验设置：

• 优化器：AdamW
• momentum：0.9,0。999
• weight decay：0.02
• 数据增强：horizontal flipping，旋转90°，180°，270°。
• 学习率变化策略：cosine decay strategy
• 初始学习率：2e-4，变化到1e-6。
• window size：8×8，则一个 window 是8×8的，一个wendow的序列长度是64。
• Evaluation metrics：PSNR 和 SSIM。

对比实验：

作者在针对图像去噪任务的数据集 SIDD 上面使用2个 NVIDIA V100 GPUs 训练了250 epochs，图片大小128 × 128，batch size=32，实验结果如下图6 (a-e)所示。

(a) [LeWin Transformer Block的必要性]： 把所有的 LeWin Transformer Block 换成是普通的基于卷积的 ResBlocks，得到所谓的 U-Net架构，使用 Concat-Skip 作为 Skip-Connection 的连接方法，结果如下。可以发现Uformer 在参数量更小的情况下超越了UNet  0.04dB。Uformer 在参数量更小的情况下超越了UNet  0.12dB。下标代表第1个卷积层的输出channel数。证明了 LeWin Transformer Block 的必要性。

(b) [分层结构还是单尺度架构]： 作者又构造了一个ViT-based的架构，只有单尺度，而不是像U-Net架构这样有不同大小的尺度的特征。ViT架构只有单尺度，就是所有特征的维度都是一致的。那这玩意用在 Image Restoration这种复杂的底层任务好使吗？

作者搞了个这样的架构出来：包含12个 Transformer Block，输入输出维度一模一样。头部 head部分是2层卷积用于提取特征，尾部是2层卷积用于重建恢复丢失的数据。图片分块是16×16，Embedding dimension是256。结果如下表所示，一顿操作猛如虎，结果一看二百五，PSNR比Uformer 低了0.15 dB。

(c) [Skip-Connection 的选择]： 如下图所示，毫无疑问选择第1种 Concat方法直接把二者在 channel 这个维度给 concatenate 在一起。

(d) [增强 Locality 的位置的选择]： 说白了就是使用 Depth-wise Convolution的位置，如下图所示。结果发现还是在 FFN 里面使用 Depth-wise Convolution，即使用 Locally-enhanced Feed-Forward Network (LeFF) 的效果最好。

去噪实验结果：

数据集：SIDD, DND。

去噪实验结果如下图7所示，作者对比了8个主要的去噪模型，它们分别是：BM3D, RIDNet, VDN, CycleISP, NBNet, DANet, MIRNet, MPRNet。

Uformer  达到了39.77dB的PSNR，超过了其他的模型。对于DND数据集来讲，Uformer 比最好的方法NBNet高了0.07dB。PSNR与计算量的帕累托曲线呈现在图8中。Uformer 可以获得比大多数模型更好的性能，但以最少的计算成本。这证明了Uformer的效率和有效性。可视化的结果如图9所示。

去雨实验结果：

数据集：SPAD。

去雨实验结果如下图10和11所示。作者对比了7个主要的去雨模型，它们分别是：GMM, DDN, RESCAN, PReNet, JORDER-E, SPANet, RCDNet。Uformer表现出了明显更好的性能，实现了5dB的提升。

去模糊和去摩尔纹实验结果：

Debluring 数据集：dual-pixel dataset DPD。

Moire Pattern数据集：TIP18。

去模糊实验结果如下图12和13所示。作者对比了3个主要的去噪模型，它们分别是：DMENet, JNB, DPDNet。Uformer分别有 +2.24 dB, +1.81 dB, 和+1.55 dB的提升。

作者对比了3个主要的去摩尔纹模型，它们分别是：UNet, CFNet, MS-Net。Uformer比之前的方法 CFNet 和 MS-Net 分别提高了 1.21 dB和 0.31 dB。

小结

Uformer 是 2 个针对图像复原任务 (Image Restoration) 的 Transformer 模型，它的特点是长得很像 U-Net (医学图像分割的经典模型)，有Encoder，有Decoder，且 Encoder 和 Decoder 之间还有 Skip-Connections。另外，在 Transformer 的 LeWin Transformer Block 中包含2个结构，一个是W-MSA，就是基于windows的attention模块，另一个是LeFF，就是融入了卷积的FFN。

28 Refiner：改进视觉Transformer的自注意力

论文名称：Refiner: Refining Self-attention for Vision Transformers

论文地址：

https://arxiv.org/pdf/2106.03714.pdf

28.1 Refiner 原理分析：

Refiner 是一个针对 Self-attention 机制本身做改进的工作。改进可以用2点来概括：

• Attention Expansion：通过增大channel数来增加 head 的数量。

• Distributed Local Attention：通过融入卷积来增强 attention 局部信息的交互。

Motivation

Self-attention (SA) 机制建模的是全局信息，这样的能力比 CNN 建模局部信息的表达能力强，但也需要更多的训练样本和更长的训练时间来完成训练。

DeiT，CaiT，T2T-ViT等都是减少使用训练数据集的有效手段，但是它们几个都是针对 Transformer 整体架构做改进的工作，无一针对 Self-attention 机制本身做改进。那么问题是：为 NLP 任务设计的 Self-attention 机制是真的适合直接迁移到 CV 任务吗？

Self-attention 机制如下图14所示，图片分成一个个patches之后，每个patch可以看做是一个token。对于每个token来讲，Self-attention 机制会建模它与其他所有tokens之间的relationship。这样，一个层的 tokens 之间就能够充分地相互交换信息，从而给模型提供了很好的表达能力。但是问题是：不同的 tokens 会变得越来越相似，尤其是当模型加深时。

这种现象也称为过度平滑 (over-smoothing)，已经在下面几个工作中有过探索，这种现象会严重影响 ViT 模型的性能。

Deep-ViT，CaiT 链接：

搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（八）

在这个工作中，作者通过直接 Refine attention map来解决这个问题。在CaiT这个工作里面作者给出结论：增加 head 的数量可以有效地提高模型的性能，受此启发，本文作者研究了如何使用类似的策略来促进 attention map 的多样性。

MHSA 的操作可以用下式表示：

式中，  ，  ，  ，  ，  是 head 的数量。

作者认为过度平滑 (over-smoothing) 问题会导致特征在通过模型的不同 Block 时变化会比较缓慢，就是前一个Block和后一个Block之间的特征差异没那么大了。那么如何去衡量这个特征变化速度这个指标呢？作者使用的是：CKA 相似度 (Centered Kernel Alignment similarity)， 衡量每个Block输出的中间 token 和最后1个 Block 的相似性，并根据第1个和最后1个 Block 的相似性作归一化。这样的度量捕获了2个属性：中间特征收敛到最终特征的速度有多快，以及最终特征与第一层特征有多大的不同。

如下图15所示是ViT，DeiT，ResNet 3种模型的 CKA 相似度，ViT模型随着层数的加深，中间层与最后一层特征的相似性增加缓慢，但是DeiT 和 ResNet 的这种相似性增长迅速 (evolving faster)，并且最后1层和第1层的差异比 ViT 大。

• 28.1.1 Attention Expansion

增加 head 的数量可以有效地提高模型的性能，但是这就会使得注意力图不够全面和准确：假如模型的 Embedding dimension 是 192，本来只有 3 个heads，那么每个head的dimension是 64。如果去直接增加 head 的数量，比如增加到 6 个 heads，那么每个head的dimension只有 32，削弱了注意力图的表达能力。

为了解决这一难题，作者探索了注意力扩展，将 multi-head self-attention 线性映射到高维空间中。因此，在增加 head 的数量时，也能够保证每个head的dimension是不变的。使模型既能享受到更多 head 的好处，又能享受到高嵌入维度的好处。

具体的做法是：使用线性变换矩阵  把 attention maps  映射到  ，并满足：

假设还是上面的例子，这个新的 attention map，原来模型的 Embedding dimension 是 192，本来只有 3 个heads，那么每个head的dimension是 64。现在模型的 Embedding dimension 扩展到了 384，以及 6 个heads，那么每个head的dimension仍然是 64。这样既能享受到更多 head 的好处，又能享受到高嵌入维度的好处。

• 28.1.2 Distributed Local Attention (DLA)

作者认为 过度平滑 (over-smoothing) 的第2个原因是ViT模型忽略了 tokens 之间的 local relationship，局部性 (locality) 和空间不变性 (权重共享) 已被证明是 CNN 成功的关键，因此作者想把卷积融入 attention 里面，来同时利用 attention 机制的全局建模能力 (global-range modeling) 和卷积操作的局部建模能力 (reinforcing local patterns)。

具体的做法是：比如说attention map的第  个head  ，对其使用卷积核  进行一步卷积操作：

接下来使用这个新的 attention map 与 value 作 feature aggregation 操作：

以上 8,9 两式就是本文采取的办法，接着作者关于加入卷积的合理性给出了一点分析：

假设我们在attention之后使用1D卷积，卷积核大小为  ，那么这种 Convolution-Augmented SA可以写作：