RUN: Rethinking the UNet Architecture for Efficient Image Restoration

Zhijian Wu, Jun Li, Chang Xu, Dingjiang Huang, Steven C. H. Hoi

Abstract: Recent advanced image restoration (IR) methods typically stack homogeneous operators hierarchically in the UNet architecture. To achieve higher accuracy, these models are now going deeper and more complex, making them resource-intensive. After comprehensively reviewing different operators within modern networks, we provide an in-depth analysis of their individual favorable properties and invent a novel efficient IR network by redesigning the UNet architecture (RUN) with heterogeneous operators. Specifically, we propose three heterogeneous operators for different relational interactions concerning the specificity of different hierarchical features of the UNet architecture. First, the spatial self-attention block (SSA Block) processes high-resolution top-level features by modeling pixel interactions from the spatial dimension. Second, the channel self-attention block (CSA Block) performs channel recalibration and information transmission for the bottom-level features with rich channels. Finally, a simple and efficient convolution block (Conv Block) is used to facilitate middle-order information propagation, which complements the self-attention mechanism to achieve local-global coupling. Based on these designs, our RUN enables more comprehensive information dissemination and interaction regardless of topological distance, thus achieving superior performance while maintaining desirable computational budgets. Extensive experiments show that our RUN achieves state-of-the-art results for a variety of IR tasks, including image deblurring, image denoising, image deraining, and low-light image enhancement.