Rethinking Expressivity and Degradation-Awareness in Attention for All-in-One Blind Image Restoration

ICLR 2026 Conference Submission3343 Authors

Keywords: Transformer, Image Restoration, Representation Learning

Abstract: All-in-one image restoration (IR) aims to recover high-quality images from diverse degradations, which in real-world settings are often mixed and unknown. Unlike single-task IR, this problem requires a model to approximate a family of heterogeneous inverse functions, making it fundamentally more challenging and practically important. Although recent focus has shifted toward large multimodal models, their robustness still depends on faithful low-level inputs, and the principles that govern effective restoration remain underexplored. We revisit attention mechanisms through the lens of all-in-one IR and identify two overlooked bottlenecks in widely adopted Restormer-style backbones: \textit{(i) the value path remains purely linear}, restricting outputs to the span of inputs and weakening expressivity, and \textit{(ii) the absence of an explicit global slot} prevents attention from encoding degradation context. To address these issues, we propose two minimal, backbone-agnostic primitives: a nonlinear value transform that upgrades attention from a selector to a selector–transformer, and a global spatial token that provides an explicit degradation-aware slot. Together, these additions improve restoration across synthetic, mixed, underwater, and medical benchmarks, with negligible overhead and consistent performance gains. Analyses with foundation model embeddings, spectral statistics, and separability measures further clarify their roles, positioning our study as a step toward rethinking attention primitives for robust all-in-one IR.

Primary Area: applications to computer vision, audio, language, and other modalities

Submission Number: 3343