cs.CVOct 31, 2026

SASFormer: Stage-Aware Selective Attention Transformer for All-in-One Blind Image Restoration

Authors: Xiangpeng Tian, Mingyang Li, Fajing Liu, Xiaofeng Wang, Xiaohai He, Chao Ren

Organizations: College of Electronics and Information Engineering, Sichuan University, Chengdu, 610065, China · College of Intelligent Systems Science and Engineering, Hubei Minzu University, Enshi 445000, China

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.

Figures & tables

Explore similar work

Mar 16, 2026

M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts

While Transformer-based architectures have dominated recent advances in all-in-one image restoration, they remain fundamentally reactive: propagating degradations rather than proactively suppressing them. In the absence of explicit suppression mechanisms, degraded signals interfere with feature learning, compelling the decoder to balance artifact removal and detail preservation, thereby increasing model complexity and limiting adaptability. To address these challenges, we propose M2IR, a novel restoration framework that proactively regulates degradation propagation during the encoding stage and efficiently eliminates residual degradations during decoding. Specifically, the Mamba-Style Transformer (MST) block performs pixel-wise selective state modulation to mitigate degradations while preserving structural integrity. In parallel, the Adaptive Degradation Expert Collaboration (ADEC) module utilizes degradation-specific experts guided by a DA-CLIP-driven router and complemented by a shared expert to eliminate residual degradations through targeted and cooperative restoration. By integrating the MST block and ADEC module, M2IR transitions from passive reaction to active degradation control, effectively harnessing learned representations to achieve superior generalization, enhanced adaptability, and refined recovery of fine-grained details across diverse all-in-one image restoration benchmarks. Our source codes are available at https://github.com/Im34v/M2IR.
May 2, 2026

Degradation-Aware Adaptive Context Gating for Unified Image Restoration

Unified image restoration using a single model often faces task interference due to diverse degradations. To address this, we propose DACG-IR (Degradation-Aware Adaptive Context Gating), which enables explicit perception of degradation characteristics to dynamically modulate feature representations. Our method constructs degradation-aware contextual representations from the input to modulate attention distribution, frequency-domain features, and feature aggregation. Specifically, a lightweight multi-scale degradation-aware module extracts coarse degradation information and generates layer-wise prompts. These prompts guide attention temperature and output gating in encoder and decoder blocks for adaptive feature extraction. Additionally, a spatial-channel dual-gated adaptive fusion mechanism refines encoder features, suppressing noise propagation from shallow to deep layers. This design effectively suppresses degradation-induced noise while preserving informative structures. Experiments show DACG-IR outperforms state-of-the-art methods in single-task, all-in-one, adverse weather removal, and composite degradation settings. Code: https://github.com/HlHomes/DACG-IR-code
Aug 7, 2026

Bend the Basics: Degradation-Aware Deformable Tokenization for All-in-One Image Restoration

All-in-one image restoration seeks a single model that can recover images degraded by diverse and spatially non-uniform corruptions. However, many unified Transformers rely on fixed patch partitioning: task/degradation condition is injected only into the backbone blocks after tokenization, leaving the embedding and reconstruction stages insensitive to local degradation variations. In contrast to previous approaches, we present Flexible Image Transformer (FIT) that explicitly models degradation awareness across the entire pipeline, from patch sampling to pixel reconstruction. Specifically, FIT employs a lightweight Degradation Encoder to predict a global degradation vector g\mathbf{g} and a spatial degradation map M\mathbf{M} from local degradation severity, which jointly condition the patch embedding and unembedding through adaptive deformation. Moreover, to improve robustness across degradation types, we introduce a task-token dropout strategy that regularizes task conditioning during training. On five standard benchmarks (BSD68, Rain100L, SOTS, GoPro, and LOLv1), FIT achieves state-of-the-art performance with 30.72 dB average PSNR on the five-degradation setting and 32.83 dB on the three-degradation setting, outperforming recent unified restoration methods by +0.5∼\sim1.1 dB. Moreover, the learned offsets provide a direct handle for visualizing degradation-aware spatial adaptation.