Uformer: A General U-Shaped Transformer for Image Restoration

TL;DR

Uformer introduces a hierarchical Transformer architecture with local window self-attention and a multi-scale modulator, achieving state-of-the-art results in various image restoration tasks efficiently.

Contribution

The paper proposes Uformer, a novel Transformer-based model with locally-enhanced window attention and a multi-scale modulator for improved image restoration.

Findings

Uformer outperforms or matches state-of-the-art methods in denoising, deblurring, and deraining.

The locally-enhanced window Transformer reduces computational complexity.

The multi-scale modulator improves detail restoration with minimal extra cost.

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.