Cat-AIR: Content- & Task-Aware All-in-One Image Restoration

Jiachen Jiang1*, Tianyu Ding2*†, Ke Zhang3, Jinxin Zhou1, Tianyi Chen2, llya Zharkov2, Zhihui Zhu1, Luming Liang2†,
1The Ohio State University, 2Microsoft, 3Johns Hopkins University
*Indicates Equal Contribution, †Indicates Corresponding Author
Paper arXiv Code Demo 🤗 Models 🤗

Denoise

Noisy Clean
Noisy Mask

Derain

Noisy Clean
Noisy Mask

Dehaze

Before After
Before Mask

Deblur

Before After
Before Mask

Light Enhancement

Before After
Before Mask

TL;DR

We present Cat-AIR 🚀, an efficient all-in-one image restoration framework that adaptively balances local and global information via alternating spatial-channel attention, using task- and content-aware routing to allocate convolution or self-attention based on region and task complexity.

🔥🔥 More content comming soon: codes, demos, models.

Overview Image

Illustration of content and task-aware image restoration. Left) Image patches with varying texture complexity and degradation levels. Right) We dynamically apply self-attention for complex regions/tasks, while using convolution for simple regions/tasks.

Demo App 🤗

Video Presentation

Abstract

All-in-one image restoration seeks to recover high-quality images from various types of degradation using a single model, without prior knowledge of the corruption source. However, existing methods often struggle to effectively and efficiently handle multiple degradation types. We present Cat-AIR, a novel Content And Task-aware framework for All-in-one Image Restoration. Cat-AIR incorporates an alternating spatial-channel attention mechanism that adaptively balances the local and global information for different tasks. Specifically, we introduce cross-layer channel attentions and cross-feature spatial attentions that allocate computations based on content and task complexity. Furthermore, we propose a smooth learning strategy that allows for seamless adaptation to new restoration tasks while maintaining performance on existing ones. Extensive experiments demonstrate that Cat-AIR achieves state-of-the-art results across a wide range of restoration tasks, requiring fewer FLOPs than previous methods, establishing new benchmarks for efficient all-in-one image restoration.

FLOPs vs PSNR comparison.

FLOPs vs PSNR
Average performance (PSNR vs. FLOPs) across three degradations (denoising, deraining, dehazing) and five degradations (including deblurring and low-light enhancement). Our methods achieve state-of-the-art results with lower computations.

Cat-AIR architecture

Model Architecture
Overview of Cat-AIR. Our design features alternating channel and spatial attention mechanisms, where channel attention complexity scales across layers and spatial attention complexity adapts based on features. Prompt modules between decoder blocks are inserted to identify degradations.

Method

The core of the method lies in an alternating spatial-channel attention mechanism, which adaptively balances local and global information through cross-layer channel attention and cross-feature spatial attention. These mechanisms dynamically allocate computational resources based on both task and content complexity—using lightweight convolution for simple regions and powerful self-attention for complex ones.

Channel Attention Mechanism
Cross-Layer Channel Attention
Spatial Attention Mechanism
Cross-Feature Spatial attention

Results Visualization

Results Image
Visual comparison on five image restoration tasks.

Mask Visualization

Content-Awareness Visualization
Content-Awareness
Task-Awareness Visualization
Task-Awareness

BibTeX

@misc{jiang2025cataircontenttaskawareallinone,
        title={Cat-AIR: Content and Task-Aware All-in-One Image Restoration}, 
        author={Jiachen Jiang and Tianyu Ding and Ke Zhang and Jinxin Zhou and Tianyi Chen and Ilya Zharkov and Zhihui Zhu and Luming Liang},
        year={2025},
        eprint={2503.17915},
        archivePrefix={arXiv},
        primaryClass={eess.IV},
        url={https://arxiv.org/abs/2503.17915}, 
  }