• 问答
  • 技术
  • 实践
  • 资源
搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了(二十一)
技术讨论

作者丨科技猛兽
编辑丨极市平台

专栏目录:https://zhuanlan.zhihu.com/p/348593638)

本文目录

40 EDT:用于底层视觉的高效图像处理 Transformer
(来自 港中文,思谋科技)
40.1 EDT 原理分析
40.2 EDT 代码解读

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

这篇文章研究的是用于底层任务的视觉模型的预训练。

40 EDT:用于底层视觉的高效图像处理 Transformer

论文名称:On Efficient Transformer and Image Pre-training for Low-level Vision

论文地址:

https://arxiv.org/abs/2112.10175

40.1 EDT 原理分析:

预训练大模型尽管在 high-level 视觉任务中取得了成功,但是在 low-level 视觉任务方面的研究却很少,本文主要研究这个问题。相关的工作 IPT (一种 Transformer 模型) 借助 ImageNet 进行底层任务预训练,在多个底层任务中取得了强悍的性能。 但是,传统的 Transformer 设计在处理高分辨率输入时面临着巨大的挑战,因为有大量的参数 (例如,IPT为116M) 和密集的计算成本。作者在这个工作里面开发了一种 encoder-decoder-based 的高效 Transformer 模型,发挥 CNN 多尺度信息的优势和 Transformer 建模长距离依赖的优势。对于 ×4 SR Urban100 数据集的超分任务,EDT 只需要 200K 的训练数据,是 IPT 所需训练数据的15.6\%,参数量是 11.6M,仅仅是 IPT 的10.0\%。对于 level-50 的去噪任务,EDT 所需的计算量仅为 38 GFLOPs,远远小于 SwinIR 所需的 451 GFLOPs。作者对 low-level 视觉任务的预训练进行了深入研究,揭示了预训练如何影响模型的内部表征,以及如何进行有效的预训练。

本文的另一个贡献是探索了预训练任务是如何影响 low-level 视觉任务的。 作者发现简单地增加数据规模可能不是最佳选择,相反,多相关任务的预训练更有效,数据使用效率更高。

EDT 模型架构

如下图1所示为 EDT 的模型架构,包含一个轻量的基于卷积的 Encoder-decoder,和一个基于 Transformer 的 Body。 尽管 Transformer 取得了成功,但高额的计算成本使得它们难以处理高分辨率输入。为了提高编码效率,对于高分辨率输入 (如去噪或去雨) 的任务,首先通过 stride convolution 将图像下采样到1/4大小,而对于低分辨率输入 (如 SISR) 的任务,则使用原始大小处理。然后是多个 Transformer 块,以较低的计算成本实现大的感受野。在解码阶段,使用转置卷积将特征上采样到输入大小。此外,在训练过程中引入 skip-connection 以实现快速收敛。特别地,超分辨率任务在输出前还有一个额外的卷积上采样器。

图1:EDT 模型架构

Shifted Crossed Local Attention

为了减少计算成本,Swin Transformer 提出在 Window 内部进行 Self-attention,并通过 Window 的不断移动来增加感受野。但是作者认为这种做法会影响模型的表征能力,因为感受野局限在了 Window 的内部。所以作者设计了一种 "水平,垂直 Window " 的做法,以达到计算复杂度和模型性能的权衡。

Shifted Crossed Local Attention 的过程如下图2所示,给定输入特征 $\mathbf{X} \in {\mathbb R}^{(H \times W) \times C}$ ,把它按照 channel 维度分成均等的2份,分别执行水平和垂直的 Window-based Attention,Window size 分别是 $(h, w)$ 和 $(w, h)$ 。最后把这两部分的结果 concat 在一起。写成公式就是:

$$
\begin{aligned}
&\mathbf{X}=\left[\mathbf{X}{1}, \mathbf{X}{2}\right], \text { where } \mathbf{X}{1}, \mathbf{X}{2} \in \mathbb{R}^{(H \times W) \times C / 2}, \
&\mathbf{X}{1}^{\prime}=\mathrm{H}-\mathrm{MSA}\left(\mathbf{X}{1}\right), \
&\mathbf{X}{2}^{\prime}=\operatorname{V-MSA}\left(\mathbf{X}{2}\right), \
&\text { (S)CL-MSA }(\mathbf{X})=\operatorname{Proj}\left(\left[\mathbf{X}{1}^{\prime}, \mathbf{X}{2}^{\prime}\right]\right),
\end{aligned}
\tag{1}
$$

紧接着下一层,作者按照 Swin Transformer 将水平和垂直的 Window 分别移位 $(\lfloor \frac{h}{2} \rfloor, \lfloor \frac{w}{2} \rfloor)$ 和 $(\lfloor \frac{w}{2} \rfloor, \lfloor \frac{h}{2} \rfloor)$ 。

图2:Shifted Crossed Local Attention

Shifted Crossed Local Attention 的计算复杂度为:

$$
\Omega((\mathrm{S}) \mathrm{CW}-\mathrm{MSA})=4 H W C^{2}+2 h w H W C
\tag{2}
$$

关于 Window Size 的设置,作者做了一组消融实验,任务是 ×2 SR,结果如下图3所示,发现大的 Window Size 比小的好,证明感受野对性能的积极影响。另外,Window Size 设置为 (6, 24) 时效果最好,证明了本文提出的水平和垂直的 Window-based Attention 的作用。

图3: Window Size 的设置消融实验

Anti-Blocking FFN

为了消除 window partition 可能造成的阻塞效应,作者设计了一种抗阻塞前馈网络 (Anti-FFN) ,其表达式为

$$
\begin{aligned}
&\mathbf{X}^{\prime}=\operatorname{Act}(\operatorname{Linear}(\mathbf{X})) \
&\mathbf{X}=\operatorname{Linear}\left(\operatorname{Act}\left(\operatorname{Anti-Block}\left(\mathbf{X}^{\prime}\right)\right)\right)
\end{aligned}
\tag{3}$$

Anti-blocking 操作是一个 5×5 的 Depthwise Convolution。

模型变体一共是下面4种:EDT-T,EDT-S,EDT-B,EDT-L。

图4:模型变体

Centered Kernel Alignment 研究隐藏层的表示相似性

衡量特征之间的相似性,一般通过衡量特征中样本对之间的相似性。那么衡量样本对之间的相似性的指标是相似性度量矩阵 (Representational Similarity Matrices) 。给定 $m$ 个数据点,得到两个层的输出 $X \in \mathbb{R}^{m \times p{1}}$ 和 $Y \in \mathbb{R}^{m \times p{2}}$ ,其中 $p_1$ 和 $p_2$ 分别代表了输出特征的神经元数量。相似性度量矩阵形式不一,比较常见的是以下几种:

  • 基于点积的相似性度量 (Dot Product-Based Similarity):

$$
\left\langle\operatorname{vec}\left(X X^{\mathrm{T}}\right), \operatorname{vec}\left(Y Y^{\mathrm{T}}\right)\right\rangle=\operatorname{tr}\left(X X^{\mathrm{T}} Y Y^{\mathrm{T}}\right)=\left|Y^{\mathrm{T}} X\right|_{\mathrm{F}}^{2}
\tag{S.1}
$$

其中, $XX^\text{T}$ 表示从特征 $X$ 的视角来衡量样本之间的关系,$YY^\text{T}$ 表示从特征 $Y$ 的视角来衡量样本之间的关系。这两个关系的内积即为特征 $X$ 和特征 $Y$ 之间的相似性。

  • 基于 Hilbert-Schmidt 独立性的度量 (Hilbert-Schmidt Independence Criterion):

对于具有0均值的 $X$ 和 $Y$ 而言,式 S.1 变成了:

$$
\frac{1}{(m-1)^{2}} \operatorname{tr}\left(X X^{\mathrm{T}} Y Y^{\mathrm{T}}\right)=\left|\operatorname{cov}\left(X^{\mathrm{T}}, Y^{\mathrm{T}}\right)\right|_{\mathrm{F}}^{2}
\tag{S.2}
$$

式 S.2 成立是因为 (这里补一个推导过程,方便理解: )\~):

$$
\begin{aligned}
\frac{1}{(m-1)^{2}} \operatorname{tr}\left(X X^{\mathrm{T}} Y Y^{\mathrm{T}}\right) &=\frac{1}{(m-1)^{2}} \operatorname{tr}\left(\left(Y^{\mathrm{T}} X\right)^{T} Y^{\mathrm{T}} X\right) \
&=\frac{1}{(m-1)^{2}} \sum{i=1}^{p} \sum{j=1}^{p}\left(Y^{\mathrm{T}} X{i, j}\right)\left(Y^{\mathrm{T}} X{i, j}\right) \
&=\frac{1}{(m-1)^{2}} \sum{i=1}^{p} \sum{j=1}^{p}\left(Y^{\mathrm{T}} X{i, j}\right)^{2} \
&=\frac{1}{(m-1)^{2}} \operatorname{cov}\left(X^{\mathrm{T}}, Y^{\mathrm{T}}\right)
{i, j}^{2} \
&=\left|\operatorname{cov}\left(X^{\mathrm{T}}, Y^{\mathrm{T}}\right)\right|_{\mathrm{F}}^{2}
\end{aligned}
\tag{S.3}
$$

基于 Hilbert-Schmidt 独立性的度量把 S.2 式的协方差的 Frobenius 范数的平方改成了 Hilbert-Schmidt 范数的平方。定义两个核函数 $K_{ij} = k(\mathbf{x}_i, \mathbf{x}j)$ , $L{ij} = l(\mathbf{y}_i, \mathbf{y}_j)$ ,HSIC 度量可以写成:

$$
\operatorname{HSIC}(K, L)=\frac{1}{(m-1)^{2}} \operatorname{tr}(K H L H)
\tag{S.4}
$$

式中, $H$ 是类似于 S.2 式中 $\frac{1}{(m-1)^2}$ 的中心矩阵, $H_n = I_n - \frac{1}{n} \mathbf{1}\mathbf{1}^\text{T}$ 。

式 S.4 和式 S.2 的联系是:当使用线性核函数 $k(\mathbf{x}, \mathbf{y}) = l(\mathbf{x}, \mathbf{y}) = \mathbf{x}^\text{T}\mathbf{y}$ 时,HSIC 度量退化为点积相似性度量。

  • CKA 度量 (Centered Kernel Alignment):

$$
\operatorname{CKA}(K, L)=\frac{\operatorname{HSIC}(K, L)}{\sqrt{\operatorname{HSIC}(K, K) \operatorname{HSIC}(L, L)}} .
\tag{5}
$$

HSIC 对各向同性标度是变化的,但可以通过归一化使其不变。给定 $m$ 个数据点,得到两个层的输出 $X \in \mathbb{R}^{m \times p{1}}$ 和 $Y\in \mathbb{R}^{m \times p{2}}$ ,其中 $p_1$ 和 $p_2$ 分别代表了输出特征的神经元数量。

EDT 通过格莱姆矩阵 $K=XX^\top$ 和 $L=YY^\top$ 来计算 CKA。

其实 $K,L$ 还有许多种计算的方法,如 RBF kernel: $k(\textbf{x}_i, \textbf{x}_j) = \exp(-||\textbf{x}_i - \textbf{x}_j||_2^2/(2\sigma^2))$ 。

作者通过 CKA 方法研究 EDT 模型的层与层之间的相似性。对于下图4 (a)(b) 的超分任务,Transformer Body 具有很强的相似性。 超分任务时,EDT 大概可以分成3块:第一块对应于模型头部,它将输入图像转换为Transformer 的 tokens。第2和第3块结构是 Transformer 主体。

图5:\(a\)-\(c\):×2EDT-S SR 模型×2EDT-B SR 模型,level-15EDT-B 去噪模型单任务训练时层与层之间的 CKA 相似性。\(d\):EDT-B SR模型与EDT-B去噪模型的交叉模型比较。\(e\)-\(g\):经过/不经过预训练时得到的模型的CKA相似性

图6:\(a\)-\(d\):×2EDT-S SR 模型在不进行预训练,单任务 \(×2 SR\) 预训练,和相关任务 \(×2,×3,×4 SR\) 预训练,不相关任务 \(×2,×3SR, level-15 denoising\) 预训练时的 CKA 相似性。\(e\)-\(h\):与之相对应的 Transformer 块的 attention head 的平均距离

发现:

  • 超分模型表现出明显的阶段性特征,各阶段的比例随模型尺寸的变化而变化,去噪模型呈现相对均匀的结构。
  • 去噪模型层与超分模型的低层更加相似,包含更多的局部信息。
  • 单任务预训练主要影响超分模型的高层,但对去噪模型影响有限。
  • SR 模型表示在低层中包含更多的局部信息,与之相对应的 Transformer 块的 attention head 的平均距离较小;在高层中包含更多的全局信息,与之相对应的 Transformer 块的 attention head 的平均距离较大。
  • 与单任务预训练相比,相关多任务预训练 (图(g)) 将更多的第3个块的全局表征转化成了局部表征,增加了第2块的范围。不相关多任务预训练 (图(h)) 的这种现象的表现偏弱。
  • 3种预训练方法都可以通过将不同程度的局部信息作为一种归纳偏差引入模型的中间层,从而大大提高模型的性能,其中相关多任务的预训练效果最好。

实验结果

数据集: 作者按照 IPT 的做法使用 ImageNet 做预训练,不同的是只使用 200k 张,占原数据集的15.6\%。

任务: 超分 (SR),去噪 (denoising),去雨 (deraining)

去噪任务直接在 clean image 上面加 Gaussian noises,去雨任务直接在 clean image 上面加雨线。

实验设置: SR:×2,×3,×4,去噪:15/25/50,去雨:light/heavy。

预训练方法: single task,related tasks 和 unrelated tasks

  • single task:指针对某一特定任务 (如×2 SR) 进行单一模型的训练。
  • related tasks:对高度相关的任务 (如×2,×3,×4 SR) 训练单一模型。
  • unrelated tasks:对不相关的任务(如×2,×3SR, level-15 去噪) 训练单一模型。

其中,unrelated tasks 是相关工作 IPT 采用的方式。

Classical SR 实验结果

作者将 EDT 与基于 CNN 的方法以及基于 Transformer 的方法进行比较。如下图7所示,EDT 在×2、×3和×4尺度上优于许多算法。即使没有经过预训练,EDT-B 的性能仍然很具有竞争力。

图7:Classical SR 实验结果

Lightweight SR 实验结果

在轻量级SR算法中,EDT 在所有基准数据集上都取得了最好的结果。尽管 SwinIR 使用更大的图片块 (64×64),并采用经过良好训练的 ×2 模型作为 ×3 和 ×4 模型的初始化,但是 EDT-T (无预训练) 仍然在 Urban100和Manga109 上有接近 0.2dB 和 0.4dB 的提升。

图8:Lightweight SR 实验结果

Denoising 实验结果

在下图9中,除了一系列之前的去噪方法外,作者还展示了的三个模型:(1) 无预训练的EDT-B。(2) EDT-B,带预培训;(3) EDT-B无下采样和预训练。值得注意的是,与从预训练中获益的 SR 模型不同,去噪模型通过预训练任务仅有 0.02-0.11dB 的增益。 一个可能的解释是,我们在去噪任务中使用了一个大型的训练数据集,它已经提供了足够的数据,满足了我们模型的需要。 此外,我们发现基于 Encoder-Decoder 的框架在高噪声水平 (例如 $σ = 50$ ) 下表现良好,而在低噪声水平 (例如 $σ = 15$ ) 下表现略差。这可能是 EDT 的下采样操作造成的。为了验证这个假设,我们训练了另一个 EDT-B 模型,没有下采样。如图9中的 EDT-B $^*$ 所示,它在低水平的噪声下确实获得了更好的性能。EDT 模型是去噪任务的一个很好的选择,因为它在性能和计算复杂度之间达到了一个最佳点。例如,EDT-B (38G)的 FLOPs 仅为 SwinIR (451G) 的8.4\%。

图9:Denoising 实验结果

Deraining 实验结果

作者还评估了 EDT 模型在基准数据集 Rain100L (小的雨线) 和 Rain100H (大的雨线) 的表现,结果如下图10所示。在去雨任务上,虽然 EDT-B (11.5 M) 的模型尺寸远小于 IPT (116 M) ,但在小雨环境下,它的性能仍比 IPT 高出近 0.9dB。与此同时,EDT 在大雨环境下比第二好的 RCDNet 获得了2.97 dB 的显著优势,说明 EDT 在严重退化的复原任务中表现良好。

图10:Deraining 实验结果

总结

作者在本文中提出了一种高效率的基于 Encoder-Decoder 的图像复原 Transformer,并对底层视觉的图像预训练任务进行了深入的分析。 作者发现,预训练通过将更多的全局表征转化成了局部表征,从而大大提高模型的性能。 此外,作者发现预训练的效果是特定于任务的,可以使 SR 任务性能有显著提高,而去噪任务效果只是略有提升。 最后,作者认为多任务预训练在挖掘图像先验方面表现出巨大的潜力,远比使用更大的预训练数据集更有效。

  • 0
  • 0
  • 235
收藏
暂无评论