Transformer improves image restoration by using stage aware attention
SASFormer: Stage-Aware Selective Attention Transformer for All-in-One Blind Image Restoration
Computer Vision and Pattern Recognition
Summary
Image restoration tries to fix pictures that have been damaged in different ways using a single computer model. The authors found that treating different parts of the model, called encoder and decoder, differently helps improve the fixing process. They made a new model called SASFormer that uses special attention in each part: the encoder focuses on preserving content and spotting damage, while the decoder focuses on reconstructing the image with flexible attention. Their tests show this method restores images better and with fewer parameters than others.
What this means in practice
- •For mobile app developers: Integrate SASFormer to enhance photo repair features in apps that fix multiple types of image damage using a single model.
- •For digital content editors: Use SASFormer to improve image quality in workflows involving restoration of photos degraded by various conditions without needing specialized models.
Authors
Xiangpeng Tian, Mingyang Li, Fajing Liu, Xiaofeng Wang, Xiaohai He, Chao Ren
Abstract
All-in-one image restoration aims to recover high-quality images from multiple degradation types using a single unified model. Recent methods improve unified restoration with hybrid attention, prompt learning, or routing mechanisms, but generally focus on enhancing degradation adaptivity while employing similar feature modeling strategies across network stages. Consequently, the distinct representation requirements of encoding and decoding under diverse degradations remain insufficiently explored. In this paper, we propose SASFormer, a Stage-Aware Selective Attention Transformer for all-in-one blind image restoration. SASFormer follows a stage-aware design principle: the encoder uses established hybrid channel-spatial attention to build robust content-preserving and degradation-aware representations, while the decoder employs adaptive selective attention for restoration-oriented reconstruction. To address the stage-specific requirements of decoder reconstruction, we develop Adaptive Selective Channel Self-Attention (ASCSA) and Adaptive Selective Spatial Self-Attention (ASSSA), enabling the decoder to dynamically identify restoration-relevant channel responses and spatial contexts. Specifically, ASCSA performs adaptive query-channel selection and attention masking, while ASSSA adaptively integrates spatial attention across multiple window sizes. Compact implicit prompts further modulate the selection process without requiring explicit degradation labels. Extensive experiments on 3-, 5-, and 11-degradation settings demonstrate that SASFormer achieves strong restoration performance with favorable parameter efficiency. Moreover, stage-wise representation analysis and ablation studies provide empirical support for the proposed stage-aware design and the effectiveness of decoder-side adaptive selection.