《MB-TaylorFormer V2:基于泰勒公式扩展的改进多分支线性Transformer用于图像修复》学习笔记

paper:2501.04486

GitHub:FVL2020/MB-TaylorFormerV2

第一版:《MB-TaylorFormer:通过泰勒公式扩展的多分支高效Transformer用于图像去雾》学习笔记-CSDN博客 

目录

摘要

1、介绍

2、相关工作

2.1 图像修复

2.2 高效自注意力

2.3 多尺度 Transformer 网络

3、方法

3.1 多分支主干网络

3.2 多尺度补丁嵌入

3.3 泰勒展开的多头自注意力

3.4 聚焦泰勒展开余项(Focused Taylor Expansion Remainder)

3.5 卷积位置编码(Convolutional Positional Encoding, CPE)

计算复杂度

注意力矩阵的秩

4、实验

4.1 实验设置

4.2 图像去雾结果

4.3 图像去雨结果

4.4 图像去雪结果

4.5 图像运动去模糊结果

4.6 图像去噪结果

4.7 消融实验

4.7.1 多尺度补丁嵌入和多分支结构探索

4.7.2 卷积位置编码的有效性

4.7.3 与其他线性自注意力模块的比较

4.7.4 泰勒展开余项分析

4.7.5 偏移的截断范围

4.7.6 聚焦因子'p'的选择


摘要

近年来,由于Transformer网络具有全局感受野和适应输入的能力,它们在图像修复领域表现出了卓越的性能。然而,Softmax注意力机制的二次计算复杂度对其在图像修复任务中的广泛应用,尤其是对于高分辨率图像,构成了显著的限制。为了解决这一挑战,我们提出了一种新的Transformer变体。该变体利用泰勒展开来近似Softmax注意力,并采用范数保持映射的概念来近似一阶泰勒展开的余项,从而实现了线性计算复杂度。此外,我们在提出的Transformer中引入了一个多分支架构,特征多尺度 patch 嵌入,这一架构具有四个明显的优势:

1)各种大小的感受野;

2)多级语义信息;

3)灵活的感受野形状;

4)加速训练和推理速度。(与第一版的区别看摘要是加快了训练速度和推理速度)

因此,所提模型被命名为基于泰勒公式展开的Transformer第二版(简称MB-TaylorFormer V2),其能够同时处理粗到细的特征,捕获长距离的像素交互,并且以有限的计算成本改进泰勒展开余项的近似。实验结果显示,在多种图像修复基准测试中,MB-TaylorFormer V2在图像去雾、去雨、去雪、运动去模糊去噪等多个图像修复任务中,取得了最先进的性能,并且几乎没有额外的计算开销。

第一版主要是去雾,第二版补齐了去雨、去雪、去模糊、去噪的内容。

1、介绍

图像修复技术的发展已经从依赖先验信息的策略[1]转向了基于深度学习的模型。在过去十年里,深度图像修复网络取得了显著进展,这些进展包括多尺度信息融合[2]、精细化卷积变体[4]以及注意力机制[5]等技术的应用,极大地提升了性能。最近,Transformer架构已广泛应用于计算机视觉任务[6][7]。然而,在图像修复任务中应用Transformer时存在两个挑战:

1)Transformer的二次计算复杂度;

2)现有视觉Transformer网络[8][9]通常通过固定卷积核生成固定尺度的tokens。

因此,需要进一步创新来解决这些挑战。

针对第一个挑战,之前的研究通过各种方法减少了Transformer的计算复杂度,如移位窗口[10]、通道自注意力[8]和核函数[11]等。然而,这些方法常常导致一些缺点,如感受野缩小、像素间缺乏交互、值近似不足以及注意力集中问题。因此,我们提出了基于泰勒公式展开的Transformer第二版,命名为TaylorFormer V2。该变体应用了一种新型的注意力机制,称为泰勒展开多头自注意力++(T-MSA++),该机制作用于整个特征图的空间维度。具体来说,T-MSA++由两个部分组成:第一部分是Softmax注意力的一阶泰勒展开,提供其数值的近似;第二部分是对泰勒展开一阶余项的近似,使T-MSA++的注意力函数呈现非线性,从而更集中地关注关键区域。此外,我们利用矩阵乘法的结合律将自注意力的计算复杂度从O(n²)降低到O(n)。这种方法有三个明显的优势:

1)保留了Transformer在数据中建模长程依赖的能力;

2)提供了准确的值近似和更加集中的注意力;

3)将自注意力引导到像素级别的交互,而不是通道级别,从而使特征处理更加细致。

针对第二个挑战,MPViT[12]通过并行卷积分支使用多尺度 patch 来解决这一问题。然而,我们发现它的灵活性仍有提升空间。受CNN基于修复网络中的DCN[13]和Inception模块[14]成功应用的启发,我们将多分支编码器-解码器骨干网络引入当前的TaylorFormer V2,形成了MB-TaylorFormer V2,它基于多尺度 patch 嵌入。这种嵌入提供了多样的感受野大小、多级语义信息和灵活的感受野形状。此外,由于Transformer的计算复杂度在通道维度上是二次的,多分支设计允许使用较少的通道,从而进一步减少计算成本。多尺度 patch 嵌入生成了具有不同尺度和维度的tokens,这些来自不同尺度的tokens被同时输入到不同的分支中,最后进行融合。

多分支编码器-解码器骨干网络通过多路径特征提取与融合,解决了传统单分支模型在复杂任务中特征表达能力不足的问题。

总之,我们的主要贡献如下:

  1. 我们使用泰勒公式对Softmax注意力进行一阶泰勒展开,使其满足矩阵乘法的结合律,从而在线性计算复杂度下建模像素之间的长距离交互;
  2. 基于范数保持映射,我们用线性计算复杂度近似泰勒展开的高阶项,解决了Softmax注意力一阶展开的注意力不集中的问题;
  3. 我们设计了一个多分支架构,结合多尺度 patch 嵌入。该设计具备多个感受野大小、灵活的感受野形状和多级语义信息,能够同时处理不同尺度的tokens;
  4. 在图像去雾、去雨、去雪、运动模糊和去噪等任务中的实验结果表明,所提MB-TaylorFormer V2在较少的计算复杂度和较小的参数量下实现了最先进的性能。

这项工作是我们在2023年ICCV会议上发表的论文的扩展[15]。与之前的工作相比,我们加入了大量新的内容和额外的实验。

(1) 我们通过泰勒展开对Softmax注意力进行了解构,并基于研究结果优化了T-MSA公式,重新设计了网络结构,引入了更加集中的版本——T-MSA++。

(2) 由于T-MSA++有效解决了T-MSA在近似泰勒展开高阶余项时的局限性,我们去除了多尺度注意力精化(MSAR)结构,并采用卷积位置编码来提供位置信息并提高注意力图的秩。

(3) 我们在多个分支之间实现了并行计算,从而在硬件上提高了推理速度。这一成就促使研究人员考虑使用并行处理技术加速他们自己的多分支结构。

(4) 如图1所示,我们验证了MB-TaylorFormer V2在更广泛的图像修复任务中的泛化能力。

2、相关工作

2.1 图像修复

近年来,图像修复领域出现了向基于数据驱动的卷积神经网络(CNN)架构[16][17]转变的显著趋势,这些方法相较于传统的修复方法[1]展现了优越的性能。在各种CNN设计中,基于编码器-解码器架构的U-Net[18]在修复任务中受到了广泛关注。其偏好原因在于U-Net架构能够提供层级的多尺度表示,能够有效地捕捉复杂特征,同时保持计算效率。此外,涉及注意力机制的策略已经成为图像修复的一个重要方向,强调对不同类型退化区域的自适应聚焦[19]。生成对抗网络(GAN)的融合也变得越来越流行,它使得在传统的像素级误差度量下,能够恢复图像的细节[20]。基于物理先验的一些方法[21][22]也受到了关注。例如,Dutta等人[22]提出了一种名为DIVA的深度神经网络,展现了基于量子多体物理理论的自适应去噪算法(De-QuIP),并在多个图像修复任务中取得了最先进的性能。此外,作为一种新颖且有效的深度学习架构,Transformer在图像修复领域受到了广泛关注。Yang等人[23]在图像超分辨率任务中引入了Transformer架构,通过重建高分辨率和低分辨率图像纹理细节之间的自注意力关系,提升了超分辨率图像的细节信息。Chen等人[7]提出了一种基于Transformer的通用图像修复方法,通过利用预训练的 IPT 模型权重并在特定任务上进行微调,有效提高了性能。更多改进版的Transformer[8][9][24]也相继提出。有关图像修复中主要设计选择的全面概述,我们推荐参考NTIRE挑战赛的报告[25][26][27]以及近期的文献综述[28][29]。

2.2 高效自注意力

Transformer的计算复杂度随着特征图空间分辨率的增加而呈二次增长,这对计算资源提出了较高的要求。一些方法通过采用滑动窗口[30]或移位窗口[9]等技术缓解了这一负担。

例如 Swin-Transformer。

然而,这些设计对Transformer捕捉数据中的长程依赖能力造成了限制。MaxViT[31]通过Grid attention解决了感受野减小的问题,但Grid attention在高分辨率图像上的计算复杂度依然是二次的。另一种策略是修改标准Transformer的注意力机制。Restormer[8]在通道之间引入了自注意力,但忽略了像素之间的全局交互。Performer[32]通过随机投影实现了线性复杂度,但查询、键和值的大小需要很大,从而增加了计算开销。Poly-nl[33]建立了注意力与高阶多项式之间的联系,但这种方法在自注意力结构中尚未得到探索。其他模型[11][34][35]通过核函数分解Softmax,并利用矩阵乘法的结合律实现了线性复杂度。然而,这些模型需要构建特殊的核函数来近似Softmax注意力的功能。例如,[11]要求注意力图的每个元素为非负值;[35]要求注意力图表现出局部相关性;[34]要求在相关区域上给予更多的关注。然而,它们都忽视了数值近似的问题。

2.3 多尺度 Transformer 网络

在高层视觉领域,除了简单的金字塔网络[36],IFormer[37]通过整合 inception 结构来融合高频和低频信息。然而,它忽视了不同尺寸 patch 的利用。CrossViT[38] 和 MPViT[12] 通过多个分支处理多尺度 patch ,旨在实现多样化的感受野。然而,由于固定形状的卷积核,感受野形状的灵活性受到限制。在低层视觉领域,MSP-Former[39] 使用多尺度投影来辅助 Transformer 捕捉复杂的退化环境。Giqe[40] 采用多分支方法处理不同尺寸的特征图。[41] 采用多个子网络捕捉与任务相关的不同特征。GridFormer[42] 设计了一种网格结构,使用残差密集 Transformer 块来捕捉多尺度信息。最近为修复任务设计的 Transformer 网络[8][9] 构建了简单的 U-Net 架构,采用单尺度补丁。然而,这些努力很少深入探索多尺度 patch 和多分支架构的使用。虽然[43] 在自注意力中利用了可变形卷积,但值得注意的是卷积核中的采样点数量仍然是固定的。相比之下,我们的多尺度可变形卷积不仅拥有灵活的采样点,还能提供多层次的语义信息。

inception 结构:通过并行多尺度特征提取与高效降维设计,解决传统卷积神经网络参数过多、计算复杂度过高的问题。

3、方法

框架,与一代版本对比。

二代

一代 

3.1 多分支主干网络

给定一张退化的图像 I \in \mathbb{R}^{3 \times h \times w} ,我们首先进行卷积以提取浅层特征,生成 F_o \in \mathbb{R}^{c \times h \times w} 。接下来,使用四阶段的编码器-解码器网络进行深层特征提取。在每个阶段中,我们结合了残差块,包括多尺度 patch 嵌入多分支 Transformer 块。我们将 Transformer 中的 FFN 层替换为 SKFF [8],它可以自适应地选择和融合来自不同卷积核尺寸的特征,从而使网络能够有效地捕获多尺度信息。利用多尺度 patch 嵌入,我们生成不同尺度的 tokens,然后将它们同时输入到多个 Transformer 分支中。每个 Transformer 分支由多个 Transformer 编码器组成,不同分支并行计算。在多分支 Transformer 块的末尾,我们使用 SKFF 模块[44] 来融合不同分支生成的特征,通过注意力机制选择性地融合互补特征。得益于这种设计,我们可以将通道数分配到多个分支中。一般而言,T-MSA++ 的计算复杂度随着通道数的增长而呈二次增长,并且通道数远小于 tokens 的数量。此外,通过将通道分解到多个分支中,分而治之的方式进一步减少了整体计算成本。我们在每个阶段使用 pixel unshuffle 和 pixel shuffle 操作[45],分别进行下采样和上采样特征。跳跃连接[46] 被用来整合来自编码器和解码器的信息,并且使用 1 × 1 卷积层进行维度缩减(除第一阶段外)。在编码器-解码器结构后还应用了残差块,以便精细化结构和纹理细节。最后,我们使用 3 × 3 卷积层减少通道数,并生成恢复图像 R \in \mathbb{R}^{3 \times h \times w} 。恢复后的图像为 I' = I + R。为了进一步减少计算成本,我们在模型中加入了深度可分离卷积[47]。

I' = I + R:退化图像+恢复图像

3.2 多尺度补丁嵌入

视觉 tokens 在尺度上展现出显著的差异。以往的方法[8][9][24] 常常使用固定卷积核的卷积来进行补丁嵌入,这可能导致生成单一尺度的视觉 tokens。为了解决这一局限性,我们提出了一种新颖的多尺度 patch 嵌入方法,具有三个关键特性:

1)不同尺寸的感受野;

2)灵活的感受野形状;

3)多层次的语义信息。

具体而言,我们采用了多个不同尺度卷积核的 DCN [13] 层。这使得 patch 嵌入能够生成不同粗细程度的视觉 tokens,同时促进灵活的变换建模。受到堆叠常规层来扩展感受野的概念启发[48],我们堆叠了多个小卷积核的 DCN 层,而不是使用单一大卷积核的 DCN 层。这种方法不仅增加了网络的深度,提供了多层次的语义信息,而且有助于减少参数数量和计算负担。所有 DCN 层后面都跟随 Hardswish [49] 激活函数。

类似于深度可分离卷积[47][50] 的方法,我们引入了一种名为深度可分离可变形卷积(DSDCN)的新技术。该方法将 DCN 的组件分解为深度卷积和逐点卷积,如图 3 所示。

标准 DCN 和 DSDCN 对于分辨率为 h \times w 图像的计算复杂度分别为:

\Omega(\text{DCN}) = 2DK^4hw + D^2K^2hw + 4DK^2hw

\Omega(\text{DSDCN}) = 8DK^2hw + D^2hw

其中,D 是特征图的通道数,K 是卷积核的大小。与 DCN 相比,DSDCN 大大降低了计算复杂度。

鉴于图像通常表现出局部相关性,而 patch 嵌入捕捉了特征图的基本元素,视觉元素(即 tokens)应更加关注局部区域。为了控制 patch 嵌入层的感受野范围,我们截断了偏移量,实际上选择其范围为 [−3,3]。如图 4 所示,根据视觉对象的形状,模型可以通过学习自动选择感受野的大小。

这个选择过程的上界为 9×9,相当于一个膨胀卷积[51],膨胀因子为 4,下界为 1×1。在设置并行的多尺度 patch 嵌入时,不同分支的感受野尺寸为x_1 \in [1, 9] ,x_2 \in [x_1, x_1 + 8],和 x_3 \in [x_2, x_2 + 8](用于三个分支)。在表 11 中的实验结果表明,适当限制每个 token 的感受野能够增强性能。

3.3 泰勒展开的多头自注意力

设查询(Q)、键(K)和值(V)分别表示具有维度 D 的 h \times w 特征向量序列。原始 Transformer 的公式[52] 如下所示:

V' = \text{Softmax}\left(\frac{Q K^T}{\sqrt{D}}\right) V

假设Q \in \mathbb{R}^{hw \times D}K \in \mathbb{R}^{hw \times D},以及 V \in \mathbb{R}^{hw \times D},应用 Softmax 结果将导致自注意力的计算复杂度为 O(h^2w^2),从而产生高计算成本。

为了将自注意力的计算复杂度从 O(h^2w^2) 降低到 O(hw),我们首先将方程(3)的广义注意力表达式表示为如下形式:

V'_i = \frac{\sum_{j=1}^N f(Q_i, K_j) V_j}{\sum_{j=1}^N f(Q_i, K_j)}

其中,带有下标 i 和 j 的矩阵分别是矩阵的第 i 行和第 j 行的向量。函数 f(\cdot) 表示任意的映射函数。当我们设定 f(Q_i, K_j) = \exp\left(\frac{Q_i K_j^T}{\sqrt{D}}\right) 时,方程(4)便变为方程(3)。

如果我们对 \exp\left(\frac{Q_i K_j^T}{\sqrt{D}}\right) 进行泰勒展开的一级近似,可以将方程(4)重写为:

V'_i = \frac{\sum_{j=1}^N \left( 1 + Q_i K_j^T + o(Q_i K_j^T) \right) V_j}{\sum_{j=1}^N \left( 1 + Q_i K_j^T + o(Q_i K_j^T) \right)}

为了近似 \exp\left(\frac{Q_i K_j^T}{\sqrt{D}}\right) ,并确保注意力图中的权重始终大于零,我们将 Q_i 和 K_j 的幅度归一化为 1,生成 \tilde{Q_i} 和 \tilde{K_j} ​。由此,我们得到自注意力的泰勒展开表达式:

V'_i = \frac{ \sum_{j=1}^N \left( 1 + \tilde{Q_i} \tilde{K_j}^T + o(\tilde{Q_i} \tilde{K_j}^T) \right) V_j}{\sum_{j=1}^N \left( 1 + \tilde{Q_i} \tilde{K_j}^T + o(\tilde{Q_i} \tilde{K_j}^T) \right)}

如果忽略泰勒展开中的高阶项,我们可以简化方程(6),并利用矩阵乘法的结合律来降低计算复杂度,如下所示:

V'_i = \frac{ \sum_{j=1}^N V_j + \tilde{Q_i} \sum_{j=1}^N \tilde{K_j}^T V_j}{N+ \tilde{Q_i} \sum_{j=1}^N \tilde{K_j}^T}

然而,忽略 Softmax 注意力泰勒展开中的高阶项通常会牺牲注意力图的非线性特征,减少模型对图像中重要区域的关注能力。为了在保持线性计算复杂度的同时确保 T-MSA++ 的注意力图保持非线性特性,我们将在下一节介绍如何预测 Softmax 注意力的余项。

3.4 聚焦泰勒展开余项(Focused Taylor Expansion Remainder)

从图 5 的分析可以得出结论,余项 o(\tilde{Q_i} \tilde{K_j}^T) 具有两个特性:

  1. 非负性;
  2. 提供了 \tilde{Q_i} \tilde{K_j}^T 的非线性缩放,以提供更聚焦的注意力。

因此,我们建立了一个新的映射函数如下所示:

o(\tilde{Q_i} \tilde{K_j}^T) = \varphi_p(\tilde{Q_i}) \varphi_p^T(\tilde{K_j})

其中,\varphi_p(x) = \frac{\| \text{ReLU}(x) \|}{\| \text{ReLU}(x^p) \| } \text{ReLU}(x^p)​,
其中 x^p 表示 x 的逐元素幂运算。

我们采用与之前线性注意力模块相似的 ReLU 函数,以确保输入的非负性,并保证公式(8)中分母的有效性。直接观察可以发现,特征的范数在映射后保持不变,即:

\| x \| = \| \varphi_p(x) \|

这意味着只有特征的方向会被调整。图 6(a) 展示了我们映射函数的原理。该映射函数在 Q_i 和 K_j 初始距离较小时,减小了它们之间的余弦距离;反之,当它们的初始距离较大时,增大了它们之间的余弦距离。该映射函数的独特性质使得 T-MSA++ 能够在注意力图中对具有更高相似度的 Q_iK_j 向量赋予更大的权重。因此,T-MSA++ 能够更接近地逼近 Softmax 注意力,如图 6(b) 所示。

显然,通过该映射函数,o(\tilde{Q_i} \tilde{K_j}^T) 满足以下特性:

  1. 非负性;
  2. 当 p>1 时,对于较大的 \tilde{Q_i} \tilde{K_j}^T ,满足以下关系:

\varphi_p(Q_i) \varphi_p^T(K_j) > Q_i K_j^T

对于较小的 \tilde{Q_i} \tilde{K_j}^T ,满足以下关系:

\varphi_p(Q_i) \varphi_p^T(K_j) < Q_i K_j^T

为了防止聚焦泰勒展开余项过度干扰 T-MSA++ 对 Softmax 注意力的数值近似,我们在 \varphi_p(\tilde{Q_i}) \varphi_p^T(\tilde{K_j}) 前引入一个可学习的调节因子 s,该因子初始化为 0.5,并且在模型训练过程中可以学习。

此外,从方程(6)可以推导出以下公式:

V'_i = \sum_{j=1}^N f_1(Q_i, K_j) V_j + \sum_{j=1}^N f_r(Q_i, K_j) V_j

= \frac{\sum_{j=1}^N V_j + \tilde{Q_i} \sum_{j=1}^N \tilde{K_j}^T V_j}{N++ \tilde{Q_i} \sum_{j=1}^N \tilde{K_j}^T + s \cdot \varphi_p(\tilde{Q_i}) \sum_{j=1}^N \varphi_p^T(\tilde{K_j})}+ \frac{s \cdot \varphi_p(\tilde{Q_i}) \sum_{j=1}^N \varphi_p^T(\tilde{K_j}) V_j }{N+ \tilde{Q_i} \sum_{j=1}^N \tilde{K_j}^T + s \cdot \varphi_p(\tilde{Q_i}) \sum_{j=1}^N \varphi_p^T(\tilde{K_j}).}

3.5 卷积位置编码(Convolutional Positional Encoding, CPE)

自注意力机制对位置是不可知的,尽管一些位置编码方法 [10], [52] 通过加入位置信息来解决这个问题。然而,这些方法通常要求固定的窗口或输入尺寸。在 T-MSA++ 中,我们采用了一种简单的方法,称为卷积位置编码(CPE)。这种方法是一种相对位置编码方法,可以应用于任意分辨率的输入图像。

具体来说,对于输入 V,我们使用深度卷积(DWC)和多尺度卷积核来执行分组卷积,如下所示:

V_I, V_{II}, \dots = \text{Split}(V)

\text{CPE}(V) = \text{Cat}(\text{DWC}_{3 \times 3}(V_I), \text{DWC}_{5 \times 5}(V_{II}), \dots)

然后我们将其添加到前一部分中获得的输出 V′ ,得到最后的公式:

\text{T-MSA++}(Q, K, V) = V' + \text{CPE}(V)

计算复杂度

在输入特征分辨率为 h \times w 的情况下,Softmax-attention 和 T-MSA++ 的计算复杂度分别为:

Omega(\text{Softmax-attention}) = 2(hw)^2 D + 4hwD^2

\Omega(\text{T-MSA++}) = 8hwD^2 + 4K^2hwD

其中,D 是输入通道的数量。

注意力矩阵的秩

对于一般的卷积核模型 [11],其注意力矩阵的秩有以下限制:

\text{Rank}(\varphi(Q)\varphi^T(K)) \leq \min(\text{Rank}(\varphi(Q)), \text{Rank}(\varphi^T(K))) \leq \min(hw, D)

其中 D 通常远小于 hw ,尤其是在图像恢复任务中,因此很难实现全秩。相反,T-MSA++ 的注意力矩阵更有可能实现全秩。为了更好的说明,我们忽略了 T-MSA++ 中分母的归一化效应。由于分母的归一化涉及到每一行所有元素的按比例缩放,因此它不会影响注意力矩阵的秩。T-MSA++ 注意力矩阵的简化公式如下:

\text{Matt} = 1 + QK^T + \varphi_p(Q) \varphi_p^T(K) + M_{\text{DWC}}

其中 \text{Matt}  和 M_{\text{DWC}}​ 分别表示 T-MSA++ 的简化注意力矩阵和与 DWC 相对应的稀疏矩阵。因此,我们可以得到以下关系:

\text{Rank}(\text{Matt}) \leq \min \left( 1 + \min(\text{Rank}(Q), \text{Rank}(K^T)) + \min(\text{Rank}(\varphi(Q)), \text{Rank}(\varphi^T(K))) + \text{Rank}(M_{\text{DWC}}), hw \right)

理论上,通过学习参数 M_{\text{DWC}}​,可以实现全秩,从而使 T-MSA++ 的注意力矩阵具有更高的秩。因此,在大多数情况下,T-MSA++ 展示了更丰富的特征表示。

4、实验

4.1 实验设置

我们评估了所提出的 MB-TaylorFormer V2 在多个基准数据集上的表现,涵盖了五个不同的图像恢复任务:(a) 图像去雾 (b) 图像去雨 (c) 图像去雪 (d) 图像运动去模糊 (f) 图像去噪。

实现细节:我们提出了三种 MB-TaylorFormer V2 的变体,分别为 MB-TaylorFormer-B V2(基础模型)、MB-TaylorFormer-L V2(大规模变体)和 MB-TaylorFormer-XL V2(超大规模变体),其详细结构如表1所示。数据增强通过随机裁剪和翻转进行。初始学习率设置为 3 \times 10^{-4} ,并通过余弦退火方法逐步减小至 1 \times 10^{-6} 。损失函数包括 L1 损失和 FFT 损失 [54]。所有对比方法均在相同的训练数据集上训练,并在相同的测试数据集上评估。

4.2 图像去雾结果

我们在合成数据集(ITS [63],OTS [63] 和 HAZE4K [68])和真实世界数据集(O-HAZE [89],Dense-Haze [64] 和 NH-HAZE [90])上,采用与 [15] 相同的设置逐步训练 MB-TaylorFormer V2。表2和图1(a)中的定量结果表明,我们的模型显著优于其他模型。

 具体而言,MB-TaylorFormer-L V2 在合成数据集 ITS 和 Haze4K 上,分别比最新的 SOTA 模型 ConIR-B [62] 提高了 0.12dB 和 0.77dB 的 PSNR,同时参数数量仅为 ConIR-B 的 84.5%。在户外合成去雾数据集 OTS 上,MB-TaylorFormer-L V2 取得了第二好的性能,显著优于后续的 C2PNet [61] 和 ConvIR-S [62] 方法。对于小规模的真实世界数据集 O-HAZE 和 NH-Haze,MB-TaylorFormer-L V2 分别在 PSNR/SSIM 上较之前表现最好的模型 ConvIR 提高了 0.07dB/0.012 和 0.11dB/0.014。这表明 MB-TaylorFormer V2 在图像去雾任务中具有强大的能力和良好的泛化能力。与 MB-TaylorFormer V1 [15] 相比,MB-TaylorFormer V2 在相同规模下取得了更好的性能,表明提出的 T-MSA++ 方法有效,并且更加关注关键区域。我们还展示了 MB-TaylorFormer-L V2 与其他 SOTA 去雾模型的视觉对比结果。如图7和图8所示,比较突出了我们方法和其他模型生成的阴影之间的巨大差异。值得注意的是,其他模型生成的图像存在明显的伪影和纹理退化,导致阴影效果不自然。而我们的方式生成的去雾图像具有更高的清晰度、更干净,并且与真实图像更为相似。

4.3 图像去雨结果

我们按照之前的工作 [8],在从多个数据集收集的 13,712 对干净-雨图像对上训练模型 [16],[72],[73],[74],[75],[76]。使用这个训练好的模型,我们在多个测试集上进行评估,包括 Rain100H [70]、Rain100L [70]、Test100 [69]、Test2800 [72] 和 Test1200 [16]。我们使用 YCbCr 色彩空间中的 Y 通道来计算 PSNR(dB)/SSIM 分数。表4中的结果显示,MB-TaylorFormer-L V2 在所有五个数据集上都表现出一致且可比较的性能提升。

与最近的 SOTA 模型 Restormer [8] 相比,MB-TaylorFormer-L V2 在所有数据集上都达到了最佳或次优的性能。在特定数据集(如 Test1200 [16])上,性能提升可高达 0.12dB,同时相比 Restormer,MB-TaylorFormer-L V2 仅使用了 62.5% 的 MACs。此外,与 MB-TaylorFormer-L V1 [15] 相比,平均 PSNR 提升了 0.34dB,表明提出的 T-MSA++ 是有效的。图9和图1(b)中展示了具有挑战性的视觉示例,其中 MB-TaylorFormer-L V2 能够去除雨滴,同时保留图像的结构内容。

4.4 图像去雪结果

我们在一个合成数据集上进行去雪实验。具体来说,我们在 Snow100K [80] 和 SRRS [81] 数据集上进行训练和测试。我们将 MB-TaylorFormer-L V2 与其他方法进行比较,并在表5中报告了结果。在 Snow100K 数据集上,MB-TaylorFormer-L V2 在 PSNR 上超越了之前最佳方法 ConvIR-B 0.09dB,并超过了 IRNeXt 方法 0.4dB。在 SRRS 数据集上,MB-TaylorFormer-L V2 在 PSNR 上超越了之前最佳方法 ConvIR-B 0.16dB。此外,我们的模型的参数数量仅为 ConvIR-B 的 84.5%。图10展示了视觉比较。我们的方法有效避免了伪影,并在去除雪的效果上表现更好。这归功于 Transformer 更大的感受野,它能够进行长距离的交互,从而聚集来自远处区域的信息进行图像恢复。相比之下,卷积方法在处理大面积退化时存在困难。

4.5 图像运动去模糊结果

MB-TaylorFormer-XL V2在GoPro数据集上进行训练,任务为图像运动去模糊。随后,MB-TaylorFormer-XL V2在两个已建立的数据集GoPro和HIDE上进行了评估。我们将MB-TaylorFormer-XL V2与当前最先进的图像运动去模糊模型进行比较,包括Restormer [8]、NAFNet [96]和DiffIR [97]。定量结果(包括PSNR和SSIM指标)见表6。

值得注意的是,MB-TaylorFormer-XL V2在性能上优于其他运动去模糊模型。具体来说,在被认为是困难数据集的GoPro数据集上,MB-TaylorFormer-XL V2比Restormer [8]和NAFNet [96]分别提高了0.32dB和0.21dB。此外,与DiffIR相比,MB-TaylorFormer-XL V2在GoPro和HIDE数据集上的表现也有所提升,分别提高了0.04dB和0.11dB。这些结果突显了MB-TaylorFormer V2在实现当前最先进的运动去模糊性能方面的有效性。此外,虽然MB-TaylorFormer-XL V1在GoPro上的PSNR值为32.95dB,在HIDE上为31.33dB,但我们的MB-TaylorFormer-XL V2展示了更好的性能,分别比MB-TaylorFormer-XL V1提高了0.29dB和0.33dB。定性结果如图11和图1(c)所示。值得注意的是,MB-TaylorFormer-XL V2在恢复细小文本细节方面表现出色,具有更高的清晰度。这一定性评估与定量研究结果一致,进一步确认了我们的MB-TaylorFormer-XL V2的卓越性能。

4.6 图像去噪结果

我们在一个实际数据集上进行去噪实验。具体来说,我们在SIDD数据集上进行了训练和测试。与之前的研究 [8] 一致,去噪使用了无偏的MB-TaylorFormer-LV2模型,这使其能够适应各种噪声水平。表7显示了我们模型的优越性能。特别是在SIDD数据集上,MB-TaylorFormer-LV2在PSNR上取得了显著的提升,超过了之前领先的CNN模型MPRNet [96],提高了0.4dB。此外,与MB-TaylorFormer-LV1 [15]相比,MB-TaylorFormer-LV2的PSNR提升高达0.13dB。图12中的视觉结果表明,我们的MB-TaylorFormer-LV2在生成干净的图像的同时,能够保持精细的纹理,这表明与CNN方法相比,MB-TaylorFormer-LV2能够有效利用Transformer的低通特性去除高频噪声。此外,与一些Transformer方法相比,MB-TaylorFormer-LV2能够利用其更精确的全局建模能力,进一步提升去噪效果。

4.7 消融实验

在本节中,我们使用去雾ITS [63] 数据集对MB-TaylorFormer-BV2模型进行消融实验,以评估和理解模型各个模块的鲁棒性和有效性。

4.7.1 多尺度补丁嵌入和多分支结构探索

在表8中,我们研究了补丁嵌入的变化以及采用不同分支数量的影响。我们的基准模型是一个基于标准单尺度卷积的单分支配置,如图13(a)所示。

接着,我们做出以下修改:

  1. 多分支结构:为了评估多分支结构的影响,我们设计了一个补丁嵌入模型,使用单尺度卷积和多分支并行配置(Conv-P),如图13(b)所示。
  2. 多种感受野尺寸:为了探究多种感受野尺寸的效果,我们在补丁嵌入中引入了并行扩张卷积层(DF=1, 2)(DilatedConv-P),如图13(c)所示。
  3. 多级语义信息研究:为了深入研究多级语义信息的影响,我们将扩张卷积替换为标准卷积,采用两层卷积层的串联方式进行补丁嵌入(Conv-SP),如图13(d)所示。
  4. 灵活感受野形状检验:为了评估灵活感受野形状的影响,我们将标准卷积替换为DSDCN(DSDCN-SP),如图13(e)所示。
    实验结果表明,性能从最好到最差的排序为:DSDCN-SP、Conv-SP、DilatedConv-P、Conv-P和Conv。这表明我们的多尺度补丁嵌入方法在补丁表示方面提供了灵活性。

4.7.2 卷积位置编码的有效性

表9(a)展示了我们的卷积位置编码(CPE)模块相较于没有CPE模块的对比模型,在PSNR上获得了1.24dB的显著提升,并且仅增加了0.05M的参数和0.45G的MACs,这得益于CPE模块提供了更高排名的注意力图以及token的相对位置信息。此外,表9(b)表明,与我们之前采用的MSAR模块[15](用于提供局部误差校正和位置编码)相比,CPE在PSNR上获得了0.27dB的提升。这表明T-MSA++能够有效地逼近更高阶的余弦项,而不依赖MSAR,并且CPE模块更适合T-MSA++。

4.7.3 与其他线性自注意力模块的比较

表9(c)-(i)展示了我们提出的T-MSA++与几种常见线性自注意力模块的比较。结果表明,TaylorFormer在现有的线性自注意力模块中具有显著优势。

这归因于几个原因:

  1. T-MSA++相比于MDTA [8]具有更细粒度的自注意力能力;
  2. 我们的模型在建模长距离像素时优于Cswin [106]和Swin [10];
  3. PVTv2 [105]中的池化机制导致信息丢失;
  4. LinFormer [107]依赖于构建一个可学习的低秩矩阵,导致参数数量过多,同时限制了注意力图的秩。

4.7.4 泰勒展开余项分析

为了探索余项及其影响,我们研究了不同阶数的泰勒展开对Softmax-注意力的影响。考虑到在T-MSA++或T-MSA [15]的二阶泰勒展开(T-MSA-2nd)中关联律不可应用,这导致了显著的计算负担,我们为Swin进行了第一次和第二次泰勒展开。表10显示,T-MSA-1st能够逼近Softmax-注意力的性能,而更高阶的泰勒展开,如T-MSA-2nd,能更好地逼近Softmax-注意力。然而,T-MSA-2nd的计算复杂度随图像分辨率的增加而平方增长,因此在实际应用中难以建模长距离依赖。与T-MSA-2nd的二次计算复杂度不同,T-MSA++是一个具有线性计算复杂度的算法。图14可视化了模型第一层的注意力图。我们观察到,对于椅子前面的点(绿色点),注意力图中对应椅子前面的点具有更高的权重;而对于椅子侧面的点(蓝色点),注意力图中对应的点也具有更高的权重。这表明T-MSA++能够集中注意力于更重要的区域。可以得出结论,T-MSA++具备类似于Softmax-注意力的关注能力。

4.7.5 偏移的截断范围

表11展示了不同截断范围对模型的影响。我们发现,带有截断偏移的DSDCN比没有截断偏移的DSDCN表现更好。我们将这种改进归因于生成的tokens在我们的模型中更集中于特征图的局部区域。我们进一步研究了不同截断范围的效果,并最终选择[-3,3]作为MB-TaylorFormerV2的截断范围。

4.7.6 聚焦因子'p'的选择

我们的模型对'p'的变化具有鲁棒性。具体来说,当'p'在3到8之间时,PSNR/SSIM变化不大(见表12)。这表明模型对该超参数不敏感。网络第一层的注意力图如图15所示,随着'p'的增加,模型对关键区域的关注显著增强。这使得模型能够准确捕捉图像中的重要特征,从而提高整体性能。然而,当'p'超过某个阈值时,这种注意力增强的效果趋于平稳,可能会导致对局部细节的过度关注,而忽略全局信息。根据图15,p=4是一个平衡的选择,因为它能够平滑地关注重要区域。因此,出于简化考虑,我们选择p=4作为本文中所有模型的参数,且无需额外调整,以确保可靠的性能,同时减少了广泛超参数优化的需求。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值