用泰勒展开优化Transformer,实现高效图像修复
MB-TaylorFormer V2: Improved Multi-branch Linear Transformer Expanded by Taylor Formula for Image Restoration
- 用泰勒展开近似Softmax注意力,将计算复杂度降至线性
- 多分支架构支持多尺度特征处理,提升修复精度与速度
- 适合需要快速高质图像修复的场景,如遥感、医疗影像
近期,由于具备全局感受野和对输入的适应性,Transformer在网络在图像修复领域表现优异。然而,Softmax注意力的二次计算复杂度严重限制了其在高分辨率图像修复任务中的广泛应用。为此,本文提出一种新型Transformer变体,利用泰勒展开近似Softmax注意力,并通过保范映射近似一阶泰勒展开的余项,实现线性计算复杂度。此外,引入多分支架构,结合多尺度补丁嵌入,具有四大优势:1)不同大小的感受野;2)多层次语义信息;3)灵活的感受野形状;4)加速训练与推理。所提模型名为基于泰勒公式扩展的Transformer第二版(MB-TaylorFormer V2),能够同时处理粗到细特征,以有限计算成本捕捉长距离像素关联,并改进泰勒展开余项逼近。在多种图像修复基准测试中,该模型在去雾、去雨、去雪、运动去模糊和去噪等任务上均达到当前最优性能,且计算开销极低。源代码见https://github.com/FVL2020/MB-TaylorFormerV2。
原文摘要 · Abstract (English)
Recently, Transformer networks have demonstrated outstanding performance in the field of image restoration due to the global receptive field and adaptability to input. However, the quadratic computational complexity of Softmax-attention poses a significant limitation on its extensive application in image restoration tasks, particularly for high-resolution images. To tackle this challenge, we propose a novel variant of the Transformer. This variant leverages the Taylor expansion to approximate the Softmax-attention and utilizes the concept of norm-preserving mapping to approximate the remainder of the first-order Taylor expansion, resulting in a linear computational complexity. Moreover, we introduce a multi-branch architecture featuring multi-scale patch embedding into the proposed Transformer, which has four distinct advantages: 1) various sizes of the receptive field; 2) multi-level semantic information; 3) flexible shapes of the receptive field; 4) accelerated training and inference speed. Hence, the proposed model, named the second version of Taylor formula expansion-based Transformer (for short MB-TaylorFormer V2) has the capability to concurrently process coarse-to-fine features, capture long-distance pixel interactions with limited computational cost, and improve the approximation of the Taylor expansion remainder. Experimental results across diverse image restoration benchmarks demonstrate that MB-TaylorFormer V2 achieves state-of-the-art performance in multiple image restoration tasks, such as image dehazing, deraining, desnowing, motion deblurring, and denoising, with very little computational overhead. The source code is available at https://github.com/FVL2020/MB-TaylorFormerV2.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。