[Paper Review] Uformer: A General U-Shaped Transformer for Image Restoration

Zhendong Wang, Xiaodong Cun+4 more|arXiv (Cornell University)|Jun 6, 2021

Advanced Image Processing Techniques76 references109 citations

TL;DR

Uformer introduces a U-shaped Transformer with Locally-Enhanced Window (LeWin) blocks and a lightweight multi-scale restoration modulator to achieve state-of-the-art results on denoising, deblurring, defocus deblurring, and deraining with efficient computation.

ABSTRACT

In this paper, we present Uformer, an effective and efficient Transformer-based architecture for image restoration, in which we build a hierarchical encoder-decoder network using the Transformer block. In Uformer, there are two core designs. First, we introduce a novel locally-enhanced window (LeWin) Transformer block, which performs nonoverlapping window-based self-attention instead of global self-attention. It significantly reduces the computational complexity on high resolution feature map while capturing local context. Second, we propose a learnable multi-scale restoration modulator in the form of a multi-scale spatial bias to adjust features in multiple layers of the Uformer decoder. Our modulator demonstrates superior capability for restoring details for various image restoration tasks while introducing marginal extra parameters and computational cost. Powered by these two designs, Uformer enjoys a high capability for capturing both local and global dependencies for image restoration. To evaluate our approach, extensive experiments are conducted on several image restoration tasks, including image denoising, motion deblurring, defocus deblurring and deraining. Without bells and whistles, our Uformer achieves superior or comparable performance compared with the state-of-the-art algorithms. The code and models are available at https://github.com/ZhendongWang6/Uformer.

Motivation & Objective

• Motivate the need for effective long-range dependency modeling in image restoration beyond traditional ConvNets.
• Propose a general U-shaped Transformer architecture suitable for multi-scale image restoration tasks.
• Develop efficient LeWin Transformer blocks to balance local detail and global context.
• Introduce a lightweight multi-scale restoration modulator to enhance detail recovery across scales.
• Demonstrate state-of-the-art or competitive performance across denoising, deblurring, defocus blur, and deraining datasets.

Proposed method

• Propose a hierarchical UNet-like encoder–decoder with skip connections where Convolutions are replaced by LeWin Transformer blocks.
• Introduce Locally-Enhanced Window (LeWin) Transformer block combining non-overlapping window-based self-attention (W-MSA) and a Locally-Enhanced Feed-Forward Network (LeFF) with a depthwise convolution.
• Use non-overlapping MxM windows for self-attention to reduce complexity from O(H^2W^2C) to O(M^2HW C).
• Incorporate a multi-scale restoration modulator as learnable window-based biases added to decoder features to adapt representations for restoration across scales.

Experimental results

Research questions

• RQ1 Can a Transformer-based U-shaped architecture with local windowed self-attention and local context FFN capture both local details and long-range dependencies effectively for image restoration?
• RQ2 Does a lightweight multi-scale restoration modulator improve restoration quality across diverse degradation types without substantial computational overhead?
• RQ3 What are the performance and efficiency trade-offs of LeWin blocks versus traditional CNN or global-attention Transformers in denoising, deblurring, and deraining tasks?

Key findings

• Uformer-B achieves 39.89 dB PSNR on SIDD and 39.98 dB on DND, surpassing prior state-of-the-art on these real-noise datasets.
• On motion deblurring, Uformer achieves state-of-the-art or competitive results across GoPro, RealBlur-R/J, and HIDE datasets.
• On defocus blur, Uformer outperforms prior methods by up to 1.87 dB PSNR on DPD with improved SSIM.
• On real rain removal (SPAD), Uformer-B attains 47.84 dB PSNR and 0.9925 SSIM, improving over prior best by 3.74 dB in PSNR.
• Ablations show LeWin blocks outperform vanilla UNet variants, and locality enhanced FFN contributes to performance gains, with the modulator providing additional gains (notably in SPAD).
• The proposed Modulator yields notable improvements in deblurring (0.46 dB) and gains in denoising and deraining tasks.”