Source-linked AI summary
MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing
Yuwei Qiu, Kaihao Zhang, Chenxi Wang, Wenhan Luo, Hongdong Li, Zhi Jin
TL;DR
Image dehazing requires restoring haze-free images while avoiding the quadratic cost of softmax attention and the rigidity of fixed-scale tokenization. MB-TaylorFormer uses Taylorized full-map attention, MSAR refinement, and multi-scale deformable multi-branch embedding. The paper reports SOTA performance with few parameters and MACs across synthetic and real dehazing datasets.
Problem
Transformers offer global receptive fields but their quadratic attention cost and fixed-scale tokenization limit efficient, flexible application to image dehazing.
Method
MB-TaylorFormer combines Taylor-expanded softmax attention, MSAR error refinement, and a multi-branch architecture with multi-scale deformable patch embedding.
Results
MB-TaylorFormer achieves state-of-the-art performance on public synthetic and real dehazing datasets with few parameters and MACs.
Takeaways & Limitations
The model supports flexible multi-scale token embedding and long-distance pixel interactions while retaining a light computational burden.
Takeaways & Limitations
First-order Taylor attention introduces approximation error because the remainder term is ignored.
Abstract
from arXiv · showhide
In recent years, Transformer networks are beginning to replace pure convolutional neural networks (CNNs) in the field of computer vision due to their global receptive field and adaptability to input. However, the quadratic computational complexity of softmax-attention limits the wide application in image dehazing task, especially for high-resolution images. To address this issue, we propose a new Transformer variant, which applies the Taylor expansion to approximate the softmax-attention and achieves linear computational complexity. A multi-scale attention refinement module is proposed as a complement to correct the error of the Taylor expansion. Furthermore, we introduce a multi-branch architecture with multi-scale patch embedding to the proposed Transformer, which embeds features by overlapping deformable convolution of different scales. The design of multi-scale patch embedding is based on three key ideas: 1) various sizes of the receptive field; 2) multi-level semantic information; 3) flexible shapes of the receptive field. Our model, named Multi-branch Transformer expanded by Taylor formula (MB-TaylorFormer), can embed coarse to fine features more flexibly at the patch embedding stage and capture long-distance pixel interactions with limited computational cost. Experimental results on several dehazing benchmarks show that MB-TaylorFormer achieves state-of-the-art (SOTA) performance with a light computational burden. The source code and pre-trained models are available at https://github.com/FVL2020/ICCV-2023-MB-TaylorFormer.
1. Introduction
Image dehazing seeks haze-free images, but Transformers face quadratic attention cost and fixed-scale tokenization. MB-TaylorFormer addresses these challenges with linearized attention, refinement, and multi-scale multi-branch embedding.
- Single image dehazing estimates latent haze-free images from hazy inputs.
- Transformer attention scales quadratically with feature-map resolution, while windowed attention restricts the receptive field for pixel-to-pixel dehazing.
- Existing visual Transformers generally produce fixed-scale tokens with fixed convolution kernels, motivating more flexible patch embedding for dehazing.
- TaylorFormer approximates softmax attention with Taylor expansion and associative matrix multiplication, reducing self-attention complexity from O(n^2) to O(n).
- MSAR refines TaylorFormer by using local query-key correlations to generate per-head scaling factors that compensate for Taylor-expansion error.
- MB-TaylorFormer combines Taylorized attention with a multi-branch architecture whose deformable multi-scale patch embedding provides varied receptive fields, semantic levels, and shapes.
2. Related Works
Related work spans CNN-based dehazing and multi-scale Transformer designs. Prior Transformers explore multi-scale processing, but the cited works generally lack the combination of patch-size diversity, flexible receptive-field shapes, and multi-branch recovery architecture emphasized here.
- CNN-based dehazing methods use atmospheric-scattering models or end-to-end image conversion to predict haze-free images.
- TaylorFormer performs Taylor expansion directly on softmax to provide a functional and numerical approximation.
- Multi-scale Transformer Networks: Prior multi-scale Transformers use pyramid, inception, projection, or multi-branch designs, but CrossViT and MPViT lack flexible receptive-field shapes.
- Multi-scale Transformer Networks: Recovery Transformers commonly use simple U-net structures or single-scale patches, while the proposed design combines multi-scale deformable convolution with multi-branch processing.
3. MB-TaylorFormer
MB-TaylorFormer combines a four-stage encoder-decoder with multi-scale deformable patch embedding and Taylor-expanded self-attention. The design targets long-distance pixel interactions and high-resolution dehazing with reduced computational complexity.
- Multi-branch Backbone: The backbone uses four encoder-decoder stages, each containing multi-scale patch embedding, multi-branch Transformer blocks, and feature fusion with SKFF.Pixel-unshuffle and pixel-shuffle operations down-sample and up-sample features across stages.
- Multi-scale Patch Embedding: Parallel deformable convolutions with different kernel scales generate coarse and fine tokens with multiple receptive-field sizes, multi-level semantics, and flexible shapes.Stacking small deformable kernels increases network depth while reducing parameters and computational burden relative to a single large kernel.
- Multi-scale Patch Embedding: DSDCN decomposes deformable convolution into depthwise and pointwise convolutions, greatly reducing computational complexity and parameter count compared with DCN.Its offsets use depthwise and pointwise convolutions, while outputs use depthwise deformable and pointwise convolutions.
- Multi-scale Patch Embedding: Truncating deformable offsets to [−3, 3] constrains each learned receptive field between 1 × 1 and 9 × 9.For three parallel branches, receptive-field ranges are x ∈[1, 9], y ∈[x, x + 8], and z ∈[y, y + 8].
- Taylor Expanded Multi-head Self-Attention: TaylorFormer applies a first-order Taylor expansion to softmax-attention and uses associative matrix multiplication to reduce complexity from O(n^2) to O(n).The approximation operates across the entire feature map rather than splitting attention into spatial windows.
- Taylor Expanded Multi-head Self-Attention: T-MSA has complexity 18hwD + 7hwD^2, and the paper reports that its values remain close to MSA while enabling high-resolution testing.The comparison assumes h×w is much larger than D.
4. Experiments
Experiments evaluate MB-TaylorFormer on synthetic and real dehazing datasets, ablate its modules, compare linear attention variants, and test snow- and rain-removal generalization. The model achieves strong restoration quality while reducing computational cost.
- Experiments on Synthetic Hazy Images: 40.71dB PSNR and 0.994 SSIM on SOTS-Indoor are achieved by MB-TaylorFormer-B, with 10% of SGID-PFF’s parameters and 14% of its computational cost.MB-TaylorFormer-L gains 4.12dB over SGID-PFF with approximately half its complexity.
- Experiments on Real Hazy Images: Up to 0.20dB and 0.04dB PSNR improvements over previous methods are reported on O-HAZE and Dense-Haze, respectively.MB-TaylorFormer achieves the best PSNR and SSIM scores among the compared CNN-based models on these small datasets.
- Ablation Studies: DSDCN-S performs best among the tested patch-embedding variants, while multi-scale tokens provide richer feature information than single-scale tokens.The depth-separable DSDCN introduces only a tiny increase in parameters and computational cost compared with depth-separable convolution.
- Ablation Studies: 1.97dB PSNR gain is obtained from MSAR, with increases of only 0.1M parameters and 0.905G MACs.Expanding the gating tensor yields only a tiny metric improvement but substantially increases parameters and computation.
- Ablation Studies: TaylorFormer outperforms the compared linear self-attention modules by modeling long-distance pixels while approximating softmax-attention.The first-order Taylor variant is selected because second-order and softmax attention retain quadratic complexity with image resolution.
- Generalization: MB-TaylorFormer performs well on snow and rain removal, indicating that its capabilities are not limited to dehazing.These evaluations use the CSD and RainCityscapes datasets.
5. Conclusion
MB-TaylorFormer combines multi-scale patch embedding with Taylor-expanded self-attention in a multi-branch linearized Transformer. The design supports flexible visual-token embedding, long-range attention, local correction, and effectiveness with lightness and generalization across datasets.
- Conclusion: MB-TaylorFormer combines flexible multi-scale patch embedding with Taylor-expanded self-attention in a multi-branch linearized Transformer.The patch embedding uses flexible receptive-field shapes, multiple scales, and multi-level semantic information.
- Conclusion: Taylor expansion reduces self-attention complexity, while gating attention corrects its output to support long-range attention and local corrections.The conclusion presents these components as the basis for the model’s efficiency and restoration capability.
- Conclusion: Experiments across various datasets demonstrate the effectiveness, lightness, and generalization of MB-TaylorFormer.