pape:2308.14036
GitHub:FVL2020/ICCV-2023-MB-TaylorFormer
目录
摘要
近年来,Transformer 网络由于其全局感受野和对输入的适应性,开始在计算机视觉领域逐渐取代纯卷积神经网络(CNN)。然而,softmax 注意力的二次计算复杂度限制了其在图像去雾任务中的广泛应用,特别是在高分辨率图像上。为了解决这个问题,我们提出了一种新的 Transformer 变体,该变体利用泰勒展开近似 softmax 注意力,并实现了线性计算复杂度。我们还提出了一个多尺度注意力精细化模块,作为补充,以修正泰勒展开的误差。此外,我们为该 Transformer 引入了多分支架构和多尺度 patch 嵌入,通过不同尺度的可变形卷积重叠来嵌入特征。多尺度 patch 嵌入的设计基于三个关键理念:
1)不同大小的感受野;
2)多层次语义信息;
3)灵活的感受野形状。
我们提出的模型,命名为通过泰勒公式扩展的多分支 Transformer(MB-TaylorFormer),可以在patch 嵌入阶段更加灵活地嵌入从粗到细的特征,并以有限的计算成本捕获远距离像素交互。实验结果表明,在多个去雾基准测试中,MB-TaylorFormer 在保持轻量计算负担的同时,达到了最先进的(SOTA)性能。
1、介绍
单张图像去雾是一个图像恢复任务,旨在从雾霾图像中估计出潜在的无雾图像。(任务)
从早期基于卷积神经网络(CNN)的方法[49, 5]及其在去雾任务中的革命性表现开始,去雾任务逐渐从基于先验的方法[89, 20]转向基于深度学习的方法。在过去十年中,深度去雾网络通过多尺度信息融合[49, 47, 37]、卷积的精细变种[74, 36]和注意力机制[45, 86, 21]等先进架构,实现了显著的性能提升。(背景)
近年来,Transformer 被广泛应用于各种计算机视觉任务,并大大推动了高层次视觉任务的进展[39, 71, 61, 6]。在图像去雾领域,直接应用 Transformer 面临几个挑战:
1)Transformer 的计算复杂度与特征图的分辨率成二次关系,这使得其不适合像素对像素的去雾任务。尽管有些方法通过在小空间窗口内应用自注意力[35, 73]来缓解这一问题,但 Transformer 的感受野仍然受到限制;
2)视觉 Transformer 的基本元素通常具有更灵活的尺度[39]。然而,现有的视觉 Transformer 网络[78, 73]通常通过固定的卷积核生成固定尺度的令牌。因此,通过引入灵活的补丁嵌入 Transformer 来解决去雾任务仍有改进空间。(问题)
为了解决第一个挑战,我们提出了一种基于泰勒公式扩展的 Transformer 变体(TaylorFormer),它在整个特征图的空间维度上应用自注意力,并保持线性计算复杂度。具体而言,我们通过对 softmax 进行泰勒展开来计算自注意力的权重,然后通过应用矩阵乘法的结合律将自注意力的计算复杂度从 O(n²) 降低到 O(n)。这一策略带来了三个优势:
1)它保留了 Transformer 在建模数据之间长距离依赖关系的能力,避免了通过窗口划分导致感受野减少的问题[58];
2)它提供了比使用基于核的自注意力公式的方法更强的值逼近,类似于原始的 Transformer[64];
3)它使 Transformer 更关注像素级交互,而非通道级交互[78],这允许对特征进行更细粒度的处理。(方法)
考虑到忽略泰勒公式中 Peano 形式的剩余项[43]所带来的误差,我们引入了一个多尺度注意力精细化(MSAR)模块来改进 TaylorFormer。我们通过对查询和键的局部信息进行卷积来利用图像中的局部相关性,以输出一个带有缩放因子的特征图。特征图的通道数等于多头自注意力(MSA)中的头数,因此每个头都有一个对应的缩放因子。我们的实验表明,所提出的 MSAR 模块在保持微小计算负担的同时,有效提高了模型性能(参见表 3)。(方法)
为了解决第二个挑战,受到 CNN 基础的去雾网络中 inception 模块[60, 59]和可变形卷积[13]成功应用的启发[18, 65, 44, 74],我们为 TaylorFormer 提出了一个基于多尺度 patch 嵌入的多分支编码器-解码器骨干网,命名为 MB-TaylorFormer。多尺度 patch 嵌入具有不同大小的感受野、多层次的语义信息和灵活的感受野形状。考虑到每个 token 的生成应遵循局部相关性先验,我们截断了可变形卷积的偏移量。我们通过深度可分离方法降低了计算复杂度和参数数量。来自不同尺度的token 随后独立地输入到 TaylorFormer 中,最后进行融合。多尺度 patch 嵌入模块能够生成具有不同尺度和维度的令牌,而多分支结构能够同时处理它们,从而捕获更强大的特征。(方法)
总结来说,我们的主要贡献如下:
(1)我们提出了一种基于泰勒展开的线性化 Transformer 新变体,能够在不进行窗口分割的情况下建模像素之间的长距离交互,并引入了一个 MSAR 模块来进一步修正 TaylorFormer 中的自注意力误差;
(2)我们设计了一个具有多尺度 patch 嵌入的多分支架构。通过不同尺寸的感受野、灵活的感受野形状和多层次的语义信息,可以帮助同时生成多尺度的令牌,并捕获更强大的特征;
(3)在公共的合成和真实去雾数据集上的实验结果表明,所提出的 MB-TaylorFormer 在保持较少参数和 MACs 的情况下,达到了最先进的(SOTA)性能。
2、相关工作
2.1 基于 CNN 的图像去
随着 CNN 的发展,图像恢复任务[70, 85, 81, 84, 82, 67, 66]取得了显著进展,其中包括图像去雾[5, 45, 19, 4, 88, 69]。一种方法是基于大气散射模型,如 DehazeNet[5]、DCPDN[80] 和 Aod-Net[31]。然而,当大气散射模型失效时,性能通常无法令人满意。另一种方法是基于图像转换的思想,它不依赖于大气散射模型。这类方法以端到端的方式预测无雾图像,如 EPDN[47]、FFA-Net[45] 和 AECRNet[74]。然而,基于 CNN 的模型很难学习长距离的像素依赖关系。
2.2 高效自注意力
Transformer 的计算复杂度随着特征图空间分辨率的增加而呈二次增长,这使得其对计算资源的需求非常高。一些工作通过滑动窗口[48, 24]或偏移窗口[39, 35, 73, 58]来减少计算负担。然而,这种设计限制了 Transformer 在数据中建模长距离依赖关系的能力。MaxViT[63]通过 Grid 注意力来弥补感受野的减少。然而,Grid 注意力在计算复杂度上并不严格线性,在高分辨率图像上仍然是二次的。另一种方法是修改原始 Transformer 的 softmax 注意力。Restormer[78]在通道之间应用自注意力,并忽略像素之间的全局关系。Performer[12]通过随机投影实现线性复杂度。然而,查询、键和值需要较大的尺寸,这导致计算成本增加。Poly-nl[3]通过高阶多项式桥接了注意力和多项式之间的关系。然而,这在自注意力结构中尚未得到探索。[28, 52, 46, 75, 57]通过核函数分解 softmax,并利用矩阵乘法的结合律实现线性复杂度。这些方法是功能逼近,例如,注意力图中的每个元素都是正的[46]。然而,值的逼近并未考虑在内。我们的方法不使用核函数,而是直接对 softmax 进行泰勒展开,这保证了对 softmax 的功能和数值逼近。
2.3 多尺度 Transformer 网络
在高层视觉领域,[41] 是一个简单的金字塔结构。IFormer[55] 应用 inception 结构来混合高频和低频信息。然而,它没有利用不同的 patch 大小。CrossViT[7] 和 MPViT[30] 通过多个分支处理多尺度补丁,从而获得多尺度感受野。然而,感受野的形状并不灵活。在低层视觉领域,MSP-Former[76] 使用多尺度投影来帮助 Transformer 表现复杂的退化环境。Giqe[54] 通过多分支处理不同大小的特征图。[87] 使用多个子网络表示与任务相关的不同特征。最近的用于恢复任务的 Transformer 网络[78, 73, 58, 26, 41] 构建了一个简单的 U-Net 网络或使用单尺度 patch。然而,这些工作很少进一步探索多尺度 patch 和多分支架构。尽管[29] 在自注意力中使用了可变形卷积,但卷积核的采样点数量是固定的。而我们的多尺度可变形卷积不仅具有灵活的采样点,而且提供了多层次的语义信息。
3、方法
我们的目标是构建一个高效且轻量的基于 Transformer 的去雾网络。为了减少计算复杂度,我们应用了 softmax 自注意力的泰勒展开,以满足结合律,并采用了类似于 Restormer[78] 的 U-Net 结构。为了弥补泰勒展开误差的影响,我们提出了一个 MSAR 模块。在接下来的部分,我们首先描述了 MB-TaylorFormer 的整体架构(图 2a)。然后介绍了三个核心模块:多尺度 patch 嵌入(图 2b)、泰勒展开自注意力(图 2c)和 MSAR 模块(图 2d)。
3.1. 多分支骨干网络
给定一个雾霾图像 ,我们首先通过卷积进行浅层特征提取,生成
。随后,我们采用四阶段的编码器-解码器网络进行深层特征提取。每个阶段包含一个残差块,该块包含多尺度 patch 嵌入和多分支 Transformer 块。我们使用多尺度 patch 嵌入生成多尺度的 token,并将其分别输入多个 Transformer 分支。每个 Transformer 分支包含多个 Transformer 编码器。我们在多分支 Transformer 块的末尾应用 SKFF[79] 模块,以融合由不同分支生成的特征。得益于 TaylorFormer 和多分支设计的优异表现,我们有机会压缩通道数和 Transformer 块的数量。我们在每个阶段应用 pixel-unshuffle 和 pixel-shuffle 操作[53],分别进行下采样和上采样。跳跃连接[51] 用于聚合编码器和解码器的信息,并在第一个阶段之外使用 1×1 卷积层进行维度压缩。我们还在编码器-解码器结构后使用残差块来恢复细节结构和纹理。最后,应用一个 3×3 的卷积层来减少通道并输出残差图像
。因此,我们通过
获得恢复后的图像。为了进一步压缩计算和参数,我们在模型中应用了深度可分离卷积[23,11]。
图像修复经典操作,浅层特征+深层特征。
3.2. 多尺度 patch 嵌入
视觉元素在尺度上可能有很大差异。现有的工作[73,35,78]通过固定卷积核进行 patch 嵌入,这可能导致视觉 token 的尺度单一。为了解决这个问题,我们设计了一种新的多尺度 patch 嵌入,具有三个特性:1)多个感受野的大小,2)多层次的语义信息,3)灵活的感受野形状。具体来说,通过设计多个具有不同卷积核尺度的可变形卷积(DCN)[13],我们使得补丁嵌入能够生成粗细不一的视觉 token,同时具备灵活的变换建模能力。受到堆叠传统层以扩展感受野[56]的操作启发,我们堆叠了多个小核的可变形卷积层,而不是使用单一的大核的可变形卷积层。这不仅增加了网络的深度,从而提供了多层次的语义信息,而且还帮助减少了参数和计算负担。所有可变形卷积层后都跟随 Hardswish[22] 激活函数。
类似于深度可分离卷积[23,11]的策略,我们提出了深度可分离和可变形卷积(DSDCN),它将 DCN 的部分计算解构为深度卷积和逐点卷积,如图 3 所示。
标准 DCN 和 DSDCN 在 图像上的计算成本如下:
其中 M 和 N 分别是输入和输出通道的数量,K 是卷积核的大小。DCN 和 DSDCN 的参数数量分别为:
总的来说,DSDCN 相比于 DCN 大大降低了计算复杂度和参数数量。考虑到图像是局部相关的,而补丁嵌入捕捉到特征图的基本元素,视觉元素(即 token)应更多地集中于局部区域。我们通过截断偏移量来控制 patch 嵌入层的感受野范围,实际选择的偏移量范围为 [−3,3]。如图 4 所示,依据视觉对象的形状,模型能够通过学习选择感受野的大小,其上限为 9×9 ,相当于膨胀卷积[77]的 BF=4 ,下限为 1×1。
当我们并行设置多尺度 patch 嵌入时,不同分支的感受野大小分别为 ,
和
(按升序排列,三个分支)。
3.3 泰勒展开多头自注意力(核心贡献)
设查询(Q)、键(K)和值(V)为维度为 D 的 特征向量序列,其中 h 和 w 分别是图像的高度和宽度。传统 Transformer 的公式如下:
由于 ,
和
,Softmax 导致自注意力的计算复杂度为
,从而带来较高的计算成本。我们旨在将自注意力的计算复杂度从
降低到
。为了实现这一点,我们首先将上式的一般化注意力方程写成如下:
其中带有下标 i 的矩阵是矩阵的第 i 行向量, 表示任意相似度函数。当我们让
时,第(6)式就退化为第(5)式。如果我们对
进行泰勒展开,可以将第(6)式改写为:
接着,我们通过对 和
进行归一化,生成
和
,以近似
。当
和
的范数小于 1 时,我们可以使得注意力图的值全部为正。实际上,我们发现归一化范数到 0.5 时可以取得最佳结果。
如图 5 所示,我们认为 和其一阶泰勒展开近似存在于定义域 [−0.25,0.25]内。因此,我们去除了 Peano 形式的余项[43],并得到了自注意力的泰勒展开表达式:
最后,我们应用矩阵乘法的结合律来处理第(8)式,得到:
补充材料提供了泰勒展开多头自注意力(T-MSA)的伪代码。类似于 MDTA[78],我们采用深度卷积生成 Q,K,V ,以强调局部上下文。我们还采用了多头结构,其中头数从上到下逐步增加。
标准多头自注意力(MSA)模块和 T-MSA 模块对于 图像块的计算复杂度分别为:
通常情况下, 要远大于 D,因此 T -MSA 相较于 MSA 在测试高分辨率图像时提供了更多可能性,并确保其值接近于 MSA。
3.4. 多尺度注意力细化
由于我们在 T-MSA 中对 Softmax 进行了泰勒展开的第一阶近似,并忽略了 Peano 形式的余项,因此不可避免地存在近似误差。对于泰勒展开的 n-阶余项 (其中
),矩阵乘法组合律不能用于将 T-MSA 的计算复杂度降到线性。然而,余项与 Q 和 K 矩阵相关。考虑到图像具有局部相关性,我们学习了 Q 和 K 矩阵的局部信息,以纠正不准确的输出 V′。此外,卷积注意力模块使得 TaylorFormert 更好地处理高频信息[42]。
具体而言,对于多头 和
,我们将它们重塑为
和
,其中 head 表示头的数量。然后,我们沿着通道维度将
和
拼接成一个张量
,该张量随后通过多尺度分组卷积层生成一个门控张量
,其计算公式为:
其中 是 T 的第 head 头,
是带有不同卷积核的卷积操作。由于网络的不同层有不同数量的头,我们为不同数量的头选择相应的多尺度分组卷积。
补充材料提供了 MSAR 模块的详细结构。通过 T-MSA 方法和 MSAR 模块,细化后的 T-MSA 模块的计算公式为:
其中 和
分别表示输入和输出特征图,投影矩阵为
,
这部分提出了通过多尺度分组卷积层进一步优化自注意力机制,减少了计算复杂度,并改善了高频信息的处理。
4、实验
4.1 实验设置
实现细节:我们提供了两种 MB-TaylorFormer 的架构,包括 MB-TaylorFormer-B(基础模型)和 MB-TaylorFormer-L(更大版本)。我们采用随机裁剪和随机翻转作为数据增强手段。初始学习率设置为 ,并通过余弦退火逐渐降低至
。我们仅使用 L1 损失作为损失函数。
数据集:我们在合成数据集(RESIDE)和真实世界数据集(O-HAZE、Dense-Haze)上评估了所提出的 MB-TaylorFormer。RESIDE 的子集,包括 ITS 和 OTS,分别包含 13990 对室内图像和 313950 对室外图像。模型在 SOTS 子集上进行评估。真实世界的数据集 O-HAZE 和 Dense-Haze,分别包含 45 张和 55 张配对图像。我们将每个数据集的最后 5 张图像作为测试集,其余作为训练集。
4.2 合成雾霾图像上的实验
表 1 比较了 MB-TaylorFormer 与现有 SOTA 方法在合成数据集上的表现。我们的基础模型(MB-TaylorFormer-B)在 SOTS-Indoor 上达到了 40.71 dB 的 PSNR 和 0.994 的 SSIM。相比于之前的 SOTA 方法 SGID-PFF [4],在参数数量上减少了 10%,计算成本减少了 14%,但 PSNR 提升了 2.19 dB。此外,我们的更大模型(MB-TaylorFormer-L)在复杂度仅为 SGID-PFF 一半的情况下,获得了 4.12 dB 的增益。我们的方法还超越了第二好的方法 Dehamer [19],在 SOTS-Outdoor 上提高了 2.91 dB。
我们还比较了 MB-TaylorFormer-L 和其他 SOTA 去雾方法的视觉结果。图 6 显示,其他方法生成的图像在阴影和高频区域的恢复效果较差。而我们的 MB-TaylorFormer 生成的无雾霾图像在阴影区域和高频区域恢复的细节更好,且最接近地面真值。
4.3 真实雾霾图像上的实验
我们进一步在真实数据集 O-HAZE [2] 和 Dense-Haze [1] 上评估 MB-TaylorFormer。定量比较见表 1。我们得出以下发现:
1) MB-TaylorFormer 在 O-HAZE 和 Dense-Haze 上分别超越了之前的方法,提升了 0.20 dB 和 0.04 dB 的 PSNR。
2) 通常来说,在小型数据集上训练时,视觉 Transformer 的非凸损失会导致较差的表现 [42]。然而,MB-TaylorFormer 在小型数据集上仍然取得了比其他基于 CNN 的模型更好的 PSNR 和 SSIM 分数。
我们还在视觉质量上与之前的方法进行了比较。图 7 显示,MSBDN [14] 的去雾效果几乎没有去除雾霾,FFA-Net [45] 生成的图像有颗粒感并且细节丢失,SGID-PFF [4] 生成的图像有色彩失真,Dehamer [19] 的结果则存在伪影和纹理丢失。我们的 MB-TaylorFormer 生成的无雾霾图像更清晰。
4.4 消融研究
在本节中,我们对所有 MB-TaylorFormer 模型进行了实验,训练数据为从 SOTS-Indoor 中裁剪的 256×256 像素图像补丁,训练时长为 500 个 epoch。基于 MB-TaylorFormer-B,我们分析了不同模块对框架的影响。我们还使用 MB-TaylorFormer-L 来探索分支维度、通道维度和深度的效果。
多尺度补丁嵌入和多分支结构研究:在表 2 中,我们研究了补丁嵌入和不同分支数目的差异。
具体而言,我们以单分支模型作为基线,采用以下几种方式进行修改:
- 为了研究多分支结构的影响,我们在模型中设计了单尺度和多分支的补丁嵌入(Conv-P)。
- 为了研究不同接收域大小的影响,我们使用并行的扩张卷积层(DF=1, 2)进行补丁嵌入(DilatedConv-P)。
- 为了研究多级语义信息的影响,我们进一步用标准卷积替换扩张卷积,并采用连接两个卷积层的方式(Conv-S)。
- 为了研究灵活接收域形状的影响,我们用 DSDCN 替换标准卷积(DSDCN-S),如图 2(b) 所示。
实验结果表明,性能从最好到最差依次为 DSDCN-S、Conv-S、DilatedConv-S、Conv-P 和 Conv。这表明我们的多尺度补丁嵌入能够灵活地嵌入补丁。如图 8 所示,我们可视化了网络第一阶段的特征图,发现多尺度标记(DSDCN-S)帮助网络获取了比单尺度标记(Conv)更丰富的信息。得益于深度可分离设计,我们的 DSDCN 相较于标准卷积仅增加了少量的参数和计算成本。
多尺度注意力精炼模块的有效性:表 3 显示,MSAR 模块相比于没有 MSAR 模块的对照模型,提供了 1.97 dB 的 PSNR 提升,且仅增加了 0.1M 参数和 0.905G 的 MACs。当我们将门控张量 更改为
时,PSNR 和 SSIM 仅略微提升,但伴随着参数和计算复杂度的大幅增加。综合性能和计算负担考虑,我们最终选择了。
与其他线性自注意力模块的比较:表 4 比较了所提出的 T-MSA 与几种常见的线性自注意力模块。结果显示,TaylorFormer 在性能上显著优于现有的线性自注意力模块,这归功于 T-MSA 能够有效建模远距离像素,并近似 softmax 注意力。
近似误差分析:为了探讨近似误差及其影响,我们研究了不同阶数的泰勒展开对 softmax 注意力的近似效果。考虑到第二阶 T-MSA(T-MSA-2nd)不满足结合律,这会导致显著的计算负担,我们分别对 Swin 进行了第一阶和第二阶泰勒展开。表 5 显示,T-MSA 能有效近似 softmax 注意力,而 T-MSA-2nd 的性能已非常接近 softmax 注意力。然而,T-MSA-2nd 和 softmax 注意力的计算复杂度随着图像分辨率的增加而呈二次增长,因此在实际应用中难以建模远距离像素关系。因此,我们最终选择了第一阶的 T-MSA 作为我们方法的基础。
4.5 泛化能力
尽管本文集中研究了图像去雾任务,我们还评估了所提出的模型在其他任务上的泛化能力,特别是雪和雨去除,使用了 CSD [8] 和 RainCityscapes [25] 数据集。雪去除的训练和测试集设置参考了现有方法 [8, 10],而雨去除的数据集设置则参考了现有方法 [25, 83]。表 6 和表 7 中的结果表明,我们的模型 MB-TaylorFormer 在这些任务上也表现良好,表明其能力不仅限于去雾。补充材料中包括了这些任务的可视化结果。