SASFormer: Stage-Aware Selective Attention Transformer for All-in-One Blind Image Restoration
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.