Source-linked AI summary
MB-TaylorFormer V2: Improved Multi-branch Linear Transformer Expanded by Taylor Formula for Image Restoration
Zhi Jin, Yuwei Qiu, Kaihao Zhang, Hongdong Li, Wenhan Luo
TL;DR
High-resolution image restoration is constrained by Softmax-attention’s quadratic complexity and fixed-scale tokenization. MB-TaylorFormer V2 uses Taylor-expanded attention with a norm-preserving remainder approximation and multi-branch multi-scale patch embedding, achieving state-of-the-art results across several restoration tasks with limited overhead.
Problem
Image-restoration Transformers are limited by quadratic Softmax-attention complexity and fixed-scale tokens, especially for high-resolution inputs.
Method
MB-TaylorFormer V2 combines T-MSA++, which linearly approximates Softmax-attention and its higher-order remainder, with parallel branches using multi-scale patch embedding.
Results
The model achieves state-of-the-art performance across dehazing, deraining, desnowing, motion deblurring, and denoising with lower computational complexity and fewer parameters.
Takeaways & Limitations
The design supports long-distance pixel interactions and coarse-to-fine feature processing while maintaining linear attention complexity and improving inference speed.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper addresses Transformer-based image restoration’s quadratic attention cost and fixed-scale tokens by combining Taylor-expanded attention with a multi-branch, multi-scale architecture.
- Existing image-restoration Transformers face quadratic computational complexity and generally generate tokens at a fixed scale.
- T-MSA++ applies a first-order Taylor expansion to Softmax-attention and approximates its first-order remainder, preserving long-range pixel interactions with O(n) complexity.
- The proposed multi-scale patch embedding creates tokens with varying scales and dimensions, which are processed simultaneously by parallel branches and then fused.
- The multi-branch design provides diverse receptive-field sizes, multi-level semantic information, and flexible receptive-field shapes while reducing channel-related computational cost.
- The paper reports state-of-the-art performance with lower computational complexity and fewer parameters across dehazing, deraining, desnowing, motion deblurring, and denoising.
- The V2 revision redesigns T-MSA, removes MSAR in favor of convolutional positional encoding, and enables parallel branch computation for higher inference speed.
2 RELATED WORKS
Related work highlights the tension between efficient attention and global receptive fields, while multi-scale Transformer designs still constrain receptive-field shape and flexibility.
- Encoder-decoder CNNs remain prominent in restoration because their hierarchical multi-scale representations capture intricate features.
- Window-based attention reduces computational burden but limits long-range dependency modeling, while Grid attention remains quadratic on high-resolution images.
- Channel self-attention overlooks global pixel interactions, and random-projection approaches can incur extra cost because queries, keys, and values remain large.
- Several efficient attention methods impose constraints such as non-negative or locally correlated attention maps, while overlooking numerical approximations.
- Multi-branch methods such as CrossViT and MPViT provide diverse receptive fields, but fixed-shape convolutional kernels constrain receptive-field flexibility.
- DSDCN uses depthwise and pointwise convolutions to generate offsets and outputs, while deformable sampling provides flexible receptive fields and multi-level semantic information.
3 METHOD
MB-TaylorFormer V2 combines Taylor-expanded linear attention with a multi-branch, multi-scale encoder-decoder to reduce computational cost while preserving focused long-range interactions and diverse receptive fields.
- Architecture: The model uses a four-stage encoder-decoder with residual blocks containing multi-scale patch embedding and parallel Transformer branches.Different-scale tokens are processed simultaneously, and SKFF merges branch features adaptively.
- Multi-scale Patch Embedding: Multi-scale patch embedding provides varied receptive-field sizes, flexible shapes, and multi-level semantic information.Stacked deformable convolutions with small kernels increase depth while reducing parameters and computational load compared with a single large kernel.
- Taylor Expanded Self-Attention: Taylor expansion approximates Softmax-attention so matrix associativity reduces self-attention complexity from quadratic in spatial resolution to linear complexity.The first-order expansion preserves an efficient computation form for long-distance pixel interactions.
- Focused Taylor Expansion Remainder: Ignoring higher-order Taylor terms weakens attention nonlinearity and focus, so T-MSA++ estimates the remainder with a norm-preserving mapping.The mapping maintains feature norms, adjusts directions, and uses non-negativity and nonlinear scaling to emphasize more similar query-key pairs.
- Focused Taylor Expansion Remainder: The mapping increases attention weights for more similar query-key vectors, producing a closer approximation to Softmax-attention.The mapping reduces cosine distance for initially similar vectors and increases it for initially dissimilar vectors.
- Attention and Positional Encoding: T-MSA++ combines the first-order Taylor term with a mapped remainder and uses convolutional positional encoding based on multi-scale depthwise convolutions.The positional encoding can be applied to input images of arbitrary resolutions, while the attention map can more readily achieve full rank.
4 EXPERIMENTS
Experiments evaluate MB-TaylorFormer V2 across five image restoration tasks and ablate its attention, positional encoding, receptive-field, and focused-factor designs. The model reports strong restoration performance, improved attention focus, and favorable computational efficiency across these evaluations.
- Benchmark evaluations: MB-TaylorFormer V2 is evaluated on image dehazing, deraining, desnowing, motion deblurring, and denoising benchmarks.Three model variants are described: MB-TaylorFormer-B V2, MB-TaylorFormer-L V2, and MB-TaylorFormer-XL V2.
- Image dehazing results: 0.12dB and 0.77dB PSNR improvements over ConIR-B are reported on ITS and Haze4K, respectively, with MB-TaylorFormer-L V2 using 84.5% of ConIR-B’s parameters.On O-HAZE and NH-HAZE, the reported PSNR/SSIM gains over ConvIR are 0.07dB/0.012 and 0.11dB/0.014, respectively.
- Image deraining results: 0.12dB is the reported maximum PSNR improvement over Restormer on Test1200, while MB-TaylorFormer-L V2 uses 62.5% of Restormer’s MACs.The model is reported to achieve optimal or suboptimal performance across five deraining datasets.
- Image desnowing results: 0.09dB and 0.16dB PSNR gains over ConvIR-B are reported on Snow100K and SRRS, respectively, with 84.5% of ConvIR-B’s parameters.The text attributes desnowing performance to long-range interactions from the Transformer’s larger receptive field.
- Image motion deblurring results: 0.32dB and 0.21dB PSNR gains over Restormer and NAFNet are reported on GoPro, while gains over DiffIR are 0.04dB on GoPro and 0.11dB on HIDE.MB-TaylorFormer-XL V2 also surpasses MB-TaylorFormer-XL V1 by 0.29dB on GoPro and 0.33dB on HIDE.
- Image denoising results: 0.4dB PSNR improvement over MPRNet is reported on SIDD, while the gain over MB-TaylorFormer-L V1 reaches 0.13dB.The qualitative results report clean images with preserved fine textures.
- Ablation study: DSDCN-SP, Conv-SP, Dilated Conv-P, Conv-P, and Conv rank from best to worst in the multi-scale patch-embedding ablation.The configurations probe multi-branch structure, receptive-field sizes, multi-level semantics, and flexible receptive-field shapes.
5 CONCLUSION
The work addresses inaccurate Transformer approximations and inflexible token scales in image restoration. MB-TaylorFormer V2 improves Taylor expansion approximation, adopts multi-scale patch embedding and parallel branches, and reports state-of-the-art results across several restoration tasks.
- Existing image restoration methods suffer from inaccurate approximations and inflexible tokens.
- MB-TaylorFormer V2 improves Linear Transformer approximation by addressing insufficient receptive fields, missing pixel self-attention, and neglected value approximations.
- Multi-scale patch embedding enhances token-scale flexibility, while improved remainder approximation and parallel multi-branch processing target crucial image areas.
- MB-TaylorFormer V2 achieves state-of-the-art performance in dehazing, deraining, motion deblurring, and denoising.
- The authors identify hardware support optimization as future work to make TaylorFormer more hardware-friendly.