2025-05-07 11:47 江苏
关注公众号,发现CV技术之美
本篇分享论文
MB-TaylorFormer V2: Improved Multi-branch Linear Transformer Expanded by Taylor Formula for Image Restoration
已被国际人工智能顶级期刊IEEE TPAMI接收。详细信息如下:作者:金枝,邱钰苇,张凯皓,李宏东,罗文寒
单位:中山大学,哈尔滨工业大学,香港科技大学等
Abstract
本文介绍了一种新的多支路线性Transformer网络,称为MB-TaylorFormer V2,能够有效且高效地进行图像复原任务。
基于之前V1的研究成果,MB-TaylorFormer V2有以下几个贡献:
优化了T-MSA的计算公式,并重新设计了网络结构,引入了一个更集中的版本T-MSA++;
T-MSA++有效地解决了T-MSA在逼近泰勒展开式高阶余项方面的局限性,我们去除了多尺度注意力细化(MSAR)结构,采用卷积位置编码来提供位置信息并增加注意力图的秩;
实现了跨多个分支的并行计算,这一成就促使研究人员考虑使用并行处理技术来加速他们自己的多分支结构。
此外,我们在更广泛的图像复原基准数据集上进行了实验。实验结果表明:MB-TaylorFormer V2在去雾、去雨、去雪、去运动模糊和去噪等多个图像复原任务中优于其它SOTA方案,并且计算开销非常小。
Method
Network Architecture
上图给出了本文所提出的MB-TaylorFormer V2三个核心模块:多分支分层设计(图2.a)多尺度Patch Embedding(图2.b)、Taylor展开的自注意力(图2.c)。
Multi-scale Patch Embedding
相比NLP领域,视觉Token在尺度上非常灵活。现有的工作在补丁嵌入中采用了固定卷积核大小的卷积,这可能导致视觉Token的单一尺度。为了解决这个问题,我们设计了一个新的多尺度Patch Embedding,具有三个特性:
多尺度的感受野
多级语义信息
灵活的感受野形状。
具体来说,通过设计多个并行的具有不同尺度卷积核的可变形卷积(DCN),我们使Patch Embedding能够生成粗糙和精细的视觉标记,以及具备灵活的表达能力,如下图所示。受到堆叠多层3 * 3卷积可获得更大感受野的启发,我们堆叠了几个小卷积核的DCN以获得更丰富的采样点。这不仅增加了网络的深度以提供多级语义信息,还有助于减少参数和计算负担。此外,我们还对DCN加入两个小的改动:
通过对offset截断使得Token更关注局部。
与深度可分离卷积的策略类似,我们提出了深度可分离和可变形卷积(DSDCN),它将 DCN的各部分分解成Depthwis卷积与Pointwise卷积。
Taylor Expanded Multi-head Self-Attention
对于原始的Transformer的Self-attention表达式,如下所示:
我们可以更一般的写成:
其中,当时,上式退化为Softmax-attention。现在对进行泰勒展开,得到:
忽略一阶泰勒展开的余项,并进一步利用矩阵乘法结合律,得到具有线性计算复杂度的Self-attention计算公式:
然而,忽略Softmax-attention的泰勒展开式中的高阶项通常会牺牲注意力图的非线性特性,降低模型对图像中某些重要区域的注意力能力。在下一节中,我们将介绍如何预测Softmax-attention的剩余部分,确保T-MSA++的注意力图在保持线性计算复杂度的同时保留非线性特性。
Focused Taylor Expansion Remainder
从此前的分析中,可以得出结论,余数具有两个性质:1)非负性;2)提供了一个非线性的缩放,以提供更集中的注意力。因此,我们通过引入保范映射近似高阶余项,增强注意力聚焦:
结合一阶项与余项,并引入可学习的调制因子_s_(初始化为 0.5),导出以下公式:
我们将上式命名为Taylor expanded multi-head self-attention ++(T-MSA++),其浮点数计算量和原始Softmax-attention的浮点数计算量如下所示:
Convolutional Positional Encoding
在T-MSA++中,我们采用了Convolutional Positional Encoding(CPE)。这种方法是一种相对位置编码的形式,可以应用于任意分辨率的输入图像。具体地,对于输入_V_,我们利用具有多尺度卷积核的深度卷积(DWC)来执行分组卷积,如下所示:
将CPE的输出与T-MSA++的输出相加,得到最终特征:
对于一般的核模型,其注意力图矩阵的秩存在以下约束:
T-MSA++注意力图的简化计算公式如下:
其中和分别是T-MSA++的简化注意力图和对应于DWC的稀疏注意力图。因此,我们可以推导出以下关系:
理论上,通过学习的全秩,可以提高T-MSA++的注意力图的秩,从而在大多数情况下实现更高的秩,并且T-MSA++表现出更丰富的特征表示。
Experiments
Main Results
上表给出了MB-TaylorFormer V2在去雾任务上和其它SOTA模型的对比,从中可以发现:
在合成去雾数据集ITS上,MB-TaylorFormer-L V2分别比最近的SOTA模型ConIR-B在PSNR上实现了0.12dB改进,同时仅利用了ConIR-B参数数量的84.5%。在OTS上,MB-TaylorFormer L V2实现了第二好的性能,显著优于后续方法C2 PNet和ConvIR-S。
对于真实去雾数据集O-HAZE和NH-Haze,我们的MB-TaylorFormerL V2比之前的SOTA模型ConvIR实现了0.07dB/0.012和0.11dB/0.014的PSNR/SSIM增益,这表明MB-TaylorFormer V2具有很强的图像去雾能力和泛化能力。
与MB-TaylorFormer V1相比,相同尺度的MB-TaylorFormer V2实现了性能的提高,实验结果表明,T-MSA++算法是有效的
上图给出了在合成数据集与真实数据集上的视觉效果对比,可以看出:对比模型的复原结果表此案出明显的伪影和纹理退化,导致不太自然的阴影。相反,我们的复原结果提高了清晰度,并且与真实情况非常相似。
上表给出了MB-TaylorFormer V2在去雨任务上和其它SOTA模型的对比,参照以前的工作,我们在从多个数据集收集的13712个图像对上训练我们的模型,从中可以发现:
与最近的SOTA模型Restormer相比,MB-TaylorFormer-L V2在所有数据集上都实现了最优或次优性能。
在Test1200数据集上,改进可以达到0.12dB,而与Restormer相比,仅利用了62.5%的MAC。
与MB-TaylorFormer-L V1相比,平均PSNR增加了0.34 dB,表明所提出的T-MSA++的有效性。
上图给出了一些具有挑战性的视觉示例,MB-TaylorFormer-L V2可以复原出无雨滴图像,同时保留底层结构内容。
此外,我们还在去雪(表5)、去模糊(表6)、去噪(表7)等图像复原任务上进行了充分的实验。
如上图所示,MB-TaylorFormer-L V2具有最高的视觉质量,特别是在恢复微小文本细节时,清晰度更高。
如上图所示,我们的MB-TaylorFormer-L V2在生成清晰图像的同时保持了精细纹理。
Ablation Studies
Exploration of multi-scale patch embedding and multibranch structures
表8研究了不同的Patch Embedding和不同的支路数的影响,以单支路为baseline,我们可以发现:
在相似的参数量和浮点数计算量下,多支路优于单支路
多尺度感受野(Dilated Conv-P)相比单尺度感受野(Conv-P)带来了+0.12dB的提升
多层次语义信息(Conv-SP)相比不具有多层次语义信息(Dilated Conv-P)带来了+0.26dB的提升
更灵活的感受野形状(DSDCN-SP)相比固定的感受野形状(Conv-SP)带来了+0.75dB的提升
Effectiveness of convolutional positional encoding
表9(a)(b)表明T-MSA++无需MSAR的帮助即可有效地逼近高阶余弦项,并且CPE模块更适合TMSA++。
Comparison with other linear self-attention modules
表9(c)-(i)通过T-MSA++和多种不同的线性Transformer的对比,证明了T-MSA++在图像复原任务上的有效性。
Analysis of approximation errors
如表10所示,为了验证近似误差的影响,我们在Swin的窗口内对softmax-attention进行泰勒展开,我们发现,对softmax-attention进行更高阶的展开能取得更好的性能,这可能时因为更优的数值近似和attention map具有更高的秩。
The truncation range of offsets
表11展示了不同截断范围对模型的影响。我们发现,带有截断偏移的DSDCN比不带截断偏移的DSDCN取得了更好的性能。我们将这种改进归因于我们方法中生成的Token更侧重于特征图的局部区域。我们进一步研究了不同截断范围的影响,最终选择 [-3, 3]作为MBTaylorFormer V2的截断范围。
The choose of focused factor ’p’
表12证明我们的模型性能对_p_的变化具有鲁棒性。具体而言,当_p_在3到8的范围内时,PSNR/SSIM变化不显著。为简单起见,我们对本文中提出的所有模型都选择_p=4_,无需进行额外的调整,以确保可靠的性能,同时最大限度地减少对大量超参数优化的需求。
❝实验室介绍
中山大学智能工程学院的前沿视觉实验室(FVL主页:https://fvl2020.framer.website/ )由学院金枝教授建设并维护,实验室目前聚焦在图像/视频质量增强、视频编解码、3D重建和无接触人体生命体征监测等领域的研究。旨在优化从视频图像的采集、传输到增强以及服务后端应用的完整周期。我们的目标是开发通用的概念和轻量化的方法。为了应对这些挑战,我们将持之以恒地进行相关的研究,并与其他实验室进行合作,希望利用更多关键技术,解决核心问题。长期欢迎有志之士加入我们!
最新 AI 进展报道
请联系:amos@52cv.net
END
欢迎加入「图像复原」交流群👇备注: