PaperlayerView exact revision on arXiv

Source-linked AI summary

Uformer: A General U-Shaped Transformer for Image Restoration

Zhendong Wang, Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, Houqiang Li

arXiv:2106.03106v2cs.CV

TL;DR

Image restoration requires recovering clean images from degraded inputs while efficiently modeling both local context and long-range dependencies. Uformer combines a hierarchical U-shaped Transformer with LeWin window attention and a lightweight multi-scale restoration modulator, achieving state-of-the-art or comparable performance across denoising, deblurring, and deraining tasks. Its broader applicability remains unevaluated beyond the tested restoration settings.

  • Problem

    Image restoration needs models that recover details from degraded images while addressing the limited long-range dependency capture of ConvNets and the quadratic cost of self-attention.

  • Method

    Uformer uses a hierarchical U-shaped encoder-decoder with LeWin non-overlapping window Transformer blocks and a learnable multi-scale restoration modulator in the decoder.

  • Results

    Uformer achieves state-of-the-art or competitive performance across image denoising, motion deblurring, defocus deblurring, and deraining, including a 3.74 dB improvement over the previous best deraining method on SPAD.

  • Takeaways & Limitations

    The architecture captures local and long-range dependencies while adding marginal parameters and computation through its window-based attention and lightweight modulation.

  • Takeaways & Limitations

    Uformer was not evaluated on other vision tasks such as image-to-image translation or image super-resolution.

Abstract

from arXiv · show

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.

1. Introduction

Uformer addresses the difficulty of capturing long-range dependencies efficiently in image restoration by combining a hierarchical Transformer architecture with locally enhanced windows and multi-scale feature modulation. It reports state-of-the-art performance across multiple restoration tasks.

  • Image restoration seeks to recover genuine images from degraded inputs, while ConvNet-based methods are limited in capturing long-range dependencies.
  • Self-attention is often restricted to low-resolution feature maps because its computational complexity is quadratic in feature-map size.
  • Uformer replaces convolution layers with Transformer blocks in a hierarchical U-shaped encoder-decoder while retaining skip-connections.
  • LeWin blocks use non-overlapping window self-attention to reduce high-resolution computation, while hierarchical low-resolution windows capture broader dependencies and depth-wise convolution captures local context.
  • A learnable multi-scale restoration modulator adds spatial biases to decoder features, adapting LeWin blocks for different degradations with a lightweight design.
  • Uformer achieves state-of-the-art or competitive results across denoising, motion deblurring, defocus deblurring, and deraining.It improves over prior methods by 0.14 dB and 0.09 dB on SIDD and DND denoising, by 1.04 dB on defocus deblurring, and by 3.74 dB on SPAD deraining.

2. Related Work

Related work develops image restoration through U-shaped and CNN-based architectures, attention mechanisms, and increasingly efficient vision Transformers. Uformer builds on hierarchical multi-scale features while introducing LeWin blocks for restoration.

  • U-shaped networks with skip-connections are widely used to capture multi-scale information hierarchically for denoising, deblurring, and demoireing.
  • ResNet-based and other CNN-based architectures have been applied broadly to general and task-specific image restoration.
  • Attention mechanisms, including squeeze-and-excitation and non-local networks, have been explored for super-resolution, deraining, and denoising.
  • Uformer combines a hierarchical multi-scale structure with the newly introduced LeWin Transformer block as its basic building component.
  • Vision Transformers capture long-range dependencies through global self-attention, but their quadratic computational cost motivates pyramid and local-window designs.
  • Transformer use in low-level vision remains less explored, with prior restoration work including texture learning for super-resolution and standard Transformer blocks for image restoration.

3. Method

Uformer is a hierarchical U-shaped Transformer for image restoration that combines windowed self-attention with locally enhanced feed-forward processing and decoder-side multi-scale feature modulation.

  • Overall Pipeline: Uformer uses a U-shaped hierarchical encoder-decoder with skip-connections, LeWin Transformer blocks, down-sampling, a bottleneck, and decoder up-sampling stages.The decoder concatenates up-sampled features with corresponding encoder features before restoration processing.
  • LeWin Transformer Block: LeWin blocks replace global self-attention with non-overlapping window attention while retaining long-range dependency modeling through hierarchical low-resolution feature maps.The windowed design reduces computation from O(H^2W^2C) to O(M^2HWC), and low-resolution windows cover larger receptive fields.
  • LeWin Transformer Block: Each LeWin block combines window-based multi-head self-attention with a locally enhanced feed-forward network that uses depth-wise convolution to leverage neighboring pixels.The block applies layer normalization, attention, and LeFF sequentially; LeFF reshapes token features into 2D maps for local processing.
  • Multi-Scale Restoration Modulator: The multi-scale restoration modulator adds learnable M × M × C tensors as shared biases before self-attention in decoder windows.Its window-sized, lightweight addition introduces marginal extra parameters and computational cost while adapting features across decoder stages.

4. Experiments

Uformer is evaluated across denoising, deblurring, deraining, and architectural ablations using PSNR and SSIM on multiple datasets. It achieves strong restoration results, with experiments supporting the contributions of its Transformer blocks, hierarchical structure, and multi-scale modulator.

  • Experimental setup: Experiments cover image denoising, motion deblurring, defocus blur removal, deraining, and component ablations across multiple datasets.Evaluation uses PSNR and SSIM, with metrics computed in RGB except deraining, which uses the Y channel in YCbCr.
  • Image denoising: 39.89 dB: Uformer-B achieves this PSNR on SIDD denoising, surpassing other compared methods by at least 0.14 dB.The DND evaluation uses a model trained on SIDD and online-server testing.
  • Motion blur removal: 0.23 dB and 0.36 dB: Uformer outperforms other methods by at least these margins on RealBlur-R and RealBlur-J, respectively.For motion deblurring, the model is trained on GoPro and evaluated on GoPro, HIDE, RealBlur-R, and RealBlur-J.
  • Defocus blur removal: 1.04 dB, 1.15 dB, 1.44 dB, and 1.87 dB: Uformer improves over KPAC, DPDNet, JNB, and DMENet, respectively, on DPD defocus blur removal.Visualizations report recovered images that are sharper and closer to ground truth.
  • Deraining: 3.74 dB: Uformer improves over the previous best work on the real SPAD deraining dataset.Visual results indicate more successful rain removal with fewer artifacts.
  • Ablation study: 0.04 dB, 0.12 dB, and 0.18 dB: Uformer-T, Uformer-S, and Uformer-B outperform matched UNet variants on SIDD denoising.Uformer-T uses fewer parameters and less computation, while Uformer-S uses fewer parameters with slightly higher computation cost.
  • Ablation study: 0.46 dB, 0.03 dB, and 0.41 dB: the multi-scale modulator improves GoPro deblurring, SIDD denoising, and SPAD deraining, respectively.The ablation supports the modulator’s ability to restore more details.

5. Discussion and Conclusion

The discussion presents Uformer as a U-shaped Transformer architecture that combines efficient local processing with long-range dependency modeling and a learnable decoder modulator. Experiments report state-of-the-art performance across several restoration tasks, while the authors note untested applications and broader misuse risks.

  • Discussion and conclusion: Uformer uses LeWin Transformer blocks to handle local context and capture long-range dependencies efficiently.The architecture also inserts a learnable multi-scale restoration modulator into the decoder to address degradations and enhance restoration quality.
  • Discussion and conclusion: Uformer achieves state-of-the-art performance on image denoising, deblurring, and deraining while using less computation and fewer parameters than the UNet family.This conclusion summarizes the reported experiments across several restoration tasks.
  • Limitations and broader impacts: The authors have not evaluated Uformer on image-to-image translation, image super-resolution, and other vision tasks.They identify these applications as directions for future investigation.
  • Limitations and broader impacts: Image restoration techniques may create privacy and authenticity risks when misused in surveillance and image forensics.The authors specifically mention altered camera-identification patterns and multimedia copyright concerns.

A.2. Variants of Skip-Connections

The paper compares three schemes for transferring encoder features to the decoder in Uformer: concatenation, cross-attention, and concatenation-based cross-attention. Their results are similar, with feature concatenation adopted as the default because it performs slightly better.

  • Skip-connection variants: The three skip-connection schemes are Concat-Skip, Cross-Skip, and ConcatCross-Skip.They are designed to deliver learned low-level encoder features to decoder stages.
  • Skip-connection variants: Concat-Skip concatenates encoder and decoder features channel-wise before feeding them to the decoder’s first LeWin block.Cross-Skip instead adds an attention module to that decoder block.
  • Skip-connection variants: The three schemes achieve similar results, while feature concatenation performs slightly better and is selected as Uformer’s default.The comparison includes Concat, Cross, and ConcatCross variants.

B. Additional Experiment for Demoireing

An additional experiment evaluates Uformer for demoireing on the TIP18 dataset. Uformer outperforms four previous methods by margins ranging from 1.53 dB to 3.19 dB.

  • Demoireing: 1.53 dB, 2.29 dB, 3.19 dB, and 2.79 dB: Uformer outperforms MopNet, MSNet, CFNet, and UNet, respectively, on TIP18 demoireing.The experiment is presented as additional evidence of Uformer’s restoration performance.

C. Additional Experimental Settings for Different Tasks

The supplementary settings describe task-specific training and evaluation protocols for denoising and motion deblurring, including crop sizes, training duration, datasets, and test sets.

  • Denoising: Denoising uses randomly cropped 128 × 128 SIDD training samples, 250 epochs, and batch size 32.
  • Denoising: Denoising evaluation uses 256 × 256 SIDD patches and 512 × 512 DND test-image patches, with DND results evaluated online.
  • Motion deblurring: Motion deblurring trains Uformer only on GoPro and evaluates it on GoPro, HIDE, and RealBlur-R/- J test sets.

D. More Visual Comparisons

The supplementary material adds visual comparisons for Uformer and other methods across five image-restoration tasks, with several figures reporting whole-image PSNR alongside visual crops.

  • Additional task comparisons: Figures 9–13 provide additional visual results for denoising, motion deblurring, defocus deblurring, deraining, and demoireing.
  • Image denoising: Figure 9 compares SIDD denoising results, with each patch’s PSNR computed on the corresponding whole image.
  • Deblurring: Figures 10 and 11 report whole-image PSNR for GoPro motion deblurring and DPD defocus deblurring, respectively, while showing zoomed regions for visual comparison in Figure 11.
  • Deraining and demoireing: Figures 12 and 13 present additional visual results for SPAD deraining and TIP18 demoireing.
Loading 2106.03106v2…