【ICCV2023】MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing

论文摘要：

🧐

本文介绍了一种新的多支路线性Transformer网络，称为MB-TaylorFormer，能够有效且高效的进行图像去雾任务。

🔥

本文提出了一种新的 Transformer 变体，通过应用泰勒展开来近似softmax-attention，实现了线性的计算复杂度。该模型还引入了多尺度注意力细化模块和多分支架构，通过多尺度块嵌入和可变形卷积来更灵活地嵌入特征。

动机：

1、传统 Transformer 计算量大

Transformer 的计算复杂度与特征图的分辨率呈二次关系，这导致了在像素到像素的去雾任务中表现不佳。一些工作将自注意力应用于较小的空间窗口，以缓解计算复杂度的问题。然而，由于 Transformer 的局部自注意力操作，其感受野受到了限制。

2、固定卷积核导致视觉标记尺度单一

目前，视觉 Transformer 网络通常会将具有固定的卷积核用于生成 token，这可能会导致视觉标记的单一尺度。但是，基于最新研究结果 visual Transformer 的基本元素通常具有更灵活的规模。因此，将灵活的 Patch Embedding 应用于 Transformer 网络中，可以为去雾任务带来更多的改进空间。

本文方法：

整体框架（MB-TaylorFormer）

本文构建一个高效轻量的基于Transformer的去雾网络，为了降低计算复杂性，使用结合律来处理 softmax-attention 的泰勒展开式，并采用了与 Restormer 类似的 U-net 结构。同时，为了弥补泰勒展开误差的影响，本文还提出了一个MSAR模块。

对于给定模糊图像，本文首先通过卷积进行浅层特征提取，生成对应特征图。接着，本文将生成的特征图输入到类似 Unet 的编码器-解码器网络结构中进行深层特征提取。在每个阶段中，使用一个包含多尺度图像块嵌入模块（Multi-scale Patch Embedding）和多支路 Transformer 块（Multi-branch Transformer Block）的残差结构块。

多尺度图像块嵌入模块用来生成多尺度的 token 标记，然后分别将它们输入到多个 Transformer 支路中。每个 Transformer 支路包含多个Transformer编码器，并在多支路 Transformer 块的末尾应用SKFF（Split-Attention Geographical Gated Fusion）模块，用于融合不同支路生成的特征。

在每个阶段中，本文使用跳跃连接来聚合编码器和解码器的信息，并在处理过程中使用一个卷积层进行降维（除了第一个阶段）。在编码器-解码器结构的最后，还应用了一个残差块来恢复图像的细节。最后，通过应用一个卷积层来降低通道数，并输出一个残差图像。因此，通过来获得恢复后的图像。

多尺度图像块嵌入模块（Multi-scale Patch Embedding）

为了解决固定的卷积核导致视觉特征只有单一尺度的问题，本文设计了一种新的多尺度图像块嵌入，该方法具有以下三个特点：1）支持多种接收域大小，2）提供多层次的语义信息，3）具备接收域形状的灵活性。

具体而言，通过设计不同尺度的可变形卷积（DCN）并行进行，我们使图像块嵌入能够生成粗细不同的视觉特征，并具备灵活的变换建模能力。本文堆叠了几个带有小卷积核的可变形卷积层，而不是使用单个带有大卷积核的可变形卷积层。这不仅增加了网络的深度，同时还提供了多层次的语义信息，并帮助减少参数和计算负担。所有的可变形卷积层后面都跟着 Hardswish 激活函数。

类似于深度可分离卷积的设计思路，本文提出了深度可分离可变形卷积（DSDCN），它将可变形卷积的一部分用深度卷积和逐点卷积来替代，具体结构如上图所示。

对于一个大小为的图像，标准的可变形卷积（DCN）和深度可分离可变形卷积（DSDCN）的计算复杂度分别如下所示：

其中，和分别表示输入和输出通道数，为卷积核大小。而标准的可变形卷积（DCN）和深度可分离可变形卷积（DSDCN）的参数数量分别为：

综上所述，相较于标准的可变形卷积（DCN），DSDCN大大降低了计算复杂度和参数数量。

考虑到图像的局部相关性和图像块嵌入可以捕捉到特征图的基本元素，为了使视觉元素能够更加关注局部区域，本文通过限制图像块嵌入层的接收域范围来实现这个目的。具体而言，作者选择了 [-3, 3] 作为偏移量的截断范围。如下图所示，根据视觉对象的形状，模型可以学习选择不同大小的接收域，接收域的上限为9×9，相当于BF=4的膨胀卷积；下限为1×1。

泰勒展开多头自注意力模块（Taylor Expanded Multi-head Self-Attention）

假设 Q、K 和 V 分别是一个具有维度 D 的 h×w 特征向量序列，其中 h 和 w 分别是图像的高度和宽度。Vanilla Transformer 的公式如下所示：

在这个公式中，，和。由于 Softmax 函数导致自注意力的计算复杂度为，这会导致计算成本非常高。我们的目标是将自注意力的计算复杂度从降低到。为了实现这一目标，作者首先进行广义注意力的建模，如下所示：

其中，下标表示矩阵的第行的向量，表示任意相似性函数。当使得时，并在公式中应用泰勒展开式对进行一阶泰勒展开，得到以下形式：

进一步地，通过将向量和的归一化生成和来近似。当和的范数小于1时，可以使注意力图的值全为正数。作者认为存在一个的近似值，其一阶Taylor展开在定义域 [-0.25, 0.25] 内。因此，消除Peano余式，得到自注意力的泰勒展开表达式如下：

最后，将矩阵乘法结合律应用于上述公式可以得到：

作者使用深度卷积方法对，和进行局部上下文的强调，同时，还使用了多头结构。对于大小为的图像，标准多头自注意力（MSA）模块和 T-MSA 模块的计算复杂性如下：

通常情况下，远大于D，因此 T-MSA 比 MSA 更适用于测试高分辨率图像，并且保证了值与 MSA 接近。

多尺度注意力优化（Multi-scale Attention Refinement）

首先，为了降低计算复杂度，我们使用一阶泰勒展开对 softmax 进行近似，并忽略了Peano剩余项的形式，这样就产生了必然的近似误差。然而，对于多头自注意力机制中的计算，这样的近似并不能保持线性计算复杂度，因为剩余项与输入矩阵、相关。为了解决这个问题，考虑到图像的局部相关性，我们学习了局部信息并用它们来纠正不精确的输出。

具体而言，对于多头注意力机制中的和，我们将它们重塑为和，并将和沿着通道维度连接，生成一个张量。

然后，张量通过一个多尺度分组卷积层，生成一个门控张量，表达式如下所示：

其中，是张量中的第个头，表示使用不同卷积核的卷积。由于网络的不同层级具有不同数量的头，我们为不同头的数量选择相应的多尺度分组卷积操作。

通过T-MSA方法生成的特征和MSAR模块输出的特征经过计算，得到了改进的T-MSA模块，相关计算公式如下：

其中和分别表示输入和输出特征图。投影操作由参数矩阵，，，以及完成。