Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration

May 26, 2026 · View on GitHub

Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration

ICLR 2026 arXiv Code Task Domains

ICLR 2026 Poster · Paper · Pre-trained Models · Results · Citation

DATPRL-IR is the first exploration of Multi-Domain All-in-One Image Restoration (MD-AiOIR). It handles diverse restoration tasks across natural scene, medical imaging, and remote sensing domains with a single unified model by learning domain-aware task prompt representations.

🔥 News

  • 2026.05: We conduct extra experiments on natural-image 5-task, 3-task, and combined-degradation all-in-one settings. Experimental results, model weights, and code are released. Our DATPRL-IR achieves 31.00 dB / 0.919 average PSNR/SSIM on the 5-task setting, 32.95 dB / 0.918 on the 3-task setting, and 30.88 dB / 0.890 on the combined-degradation CDD11 setting.
  • 2026.05: We release the model weights and code for the 3-domain 6-task to 9-task settings reported in the paper.
  • 2026.01: DATPRL-IR is accepted by ICLR 2026.

✨ Highlights

  • First MD-AiOIR framework: one model for multiple image domains and restoration tasks.
  • Dual prompt pools: a task prompt pool and a domain prompt pool encode task-level and domain-level knowledge.
  • Prompt Composition Mechanism (PCM): adaptively composes selected prompts into instance-level representations.
  • Domain-aware restoration: domain priors distilled from MLLMs help the model better understand different image domains.
  • Adaptive Gated Fusion (AGF): dynamically controls the contribution ratio between prompt representations and backbone features at different network layers.
  • Strong scalability: DATPRL-IR remains robust when extending from 6 tasks to 9 tasks.
  • Extra natural-image AiOIR benchmarks: we additionally provide classic natural-image all-in-one settings, including 3-task, 5-task, and combined-degradation settings, together with open-source weights. DATPRL-IR surpasses existing SOTA methods by +0.08 dB, +0.42 dB, and +1.83 dB average PSNR on the 3-task, 5-task, and combined-degradation CDD11 settings, respectively.

🧠 Method Overview

DATPRL-IR adopts a query–retrieval–composition paradigm:

  1. Task Prompt Pool learns task-related knowledge shared by and specific to different restoration tasks.
  2. Domain Prompt Pool learns domain-aware priors for natural, medical, and remote sensing images.
  3. PCM composes the most relevant prompts into instance-level task/domain representations.
  4. Adaptive Gated Fusion (AGF) injects the final domain-aware task prompt representation into the restoration backbone.

📌 Supported Tasks

Multi-domain all-in-one restoration

DomainTaskDataset / Benchmark
Natural Image4× Super-ResolutionDF2K / DIV2K-Val
Natural ImageDerainingRain100L
Natural ImageMotion DeblurringGoPro
Medical ImageMRI Super-ResolutionIXI MRI
Medical ImageCT DenoisingAAPM-Mayo
Medical ImagePET SynthesisPolarStar m660
Remote Sensing Image4× Super-ResolutionUCMerced
Remote Sensing ImageCloud RemovalCUHK-CR1
Remote Sensing ImageDehazingRICE1

Extra natural-image all-in-one settings

SettingTasksDataset / Benchmark
Natural-AiOIR-3TDehazing + Deraining + DenoisingSOTS + Rain100L + WED & BSD400
Natural-AiOIR-5TDehazing + Deraining + Deblurring + Denoising + LowlightSOTS + Rain100L + GoPro + WED & BSD400 + LOL
Natural-AiOIR-CombinedCombined / mixed degradationsCDD11

📊 Results

Main MD-AiOIR results

Qualitative results

Visual inference results for the 3-domain 6-task MD-AiOIR setting are released at 3-domain-6-task-results

Visual inference results for the 3-domain 9-task MD-AiOIR setting are released at 3-domain-9-task-results

Quantitative results

SettingDomainsTasksAvg. PSNR ↑Avg. SSIM ↑
DATPRL-IR-6T3630.770.8653
DATPRL-IR-7T3730.740.8643
DATPRL-IR-8T3830.740.8647
DATPRL-IR-9T3930.780.8645

Natural-image all-in-one results

Qualitative results

Visual inference results for the natural-image 5-task setting are released at 5-task-results

Visual inference results for the natural-image 3-task setting are released at 3-task-results

Visual inference results for the natural-image mixed-degradation CDD11 setting are released at CDD11-results

Quantitative results

Quantitative results on 5-Task natural image AiOIR.

Quantitative results on 3-Task natural image AiOIR.

Quantitative results on mixed-degradation natural image AiOIR with CDD11 dataset.

🧩 Repository Structure

ICLR/
├── All_in_One/                     # Extra natural-image all-in-one experiments
│   ├── 5task_test/                 # Test logs for the natural-image 5-task setting
│   ├── net/                        # Network architecture
│   │   ├── promptpoolnafnetv3_arch.py
│   │   ├── promptrestormerv3_arch.py
│   │   └── restormer_arch.py
│   ├── utils/                      
│   ├── options.py                  
│   ├── test_v2.py                  # Testing 
│   ├── train_3task.py              # Training for the 3-task setting
│   └── train_5task.py              # Training for the 5-task setting
├── MD_All_in_One/                  # Main multi-domain all-in-one restoration code
│   ├── experiments/                # Experiment records / checkpoints
│   ├── options/                    # YAML configuration files
│   │   ├── train/
│   │   └── test/
│   ├── scripts/                    # Helper scripts
│   ├── requirements.txt            # Python dependencies
│   ├── setup.cfg
│   └── setup.py
├── figure/                         # Figures used in README
└── README.md

🛠️ Installation

cd MD_All_in_One

pip install -r requirements.txt
python setup.py develop

The codebase follows the BasicSR-style training / testing pipeline. Please adjust the PyTorch and CUDA versions according to your own environment.

🚀 Inference

Test MD-AiOIR models

cd MD_All_in_One
python basicsr/test.py -opt options/test/test_allinone_6task_ours.yml

Test natural-image all-in-one models

# 5-task natural-image all-in-one model
cd All_in_One
python test_v2.py

# Combined-degradation natural-image all-in-one model
cd MD_All_in_One
python basicsr/test.py -opt options/test/test_cdd11.yml

📦 Pre-trained Models

ModelDomainsTasksCheckpoint
DATPRL-IR-6TNatural + Medical + Remote Sensing63Domain-6Task
DATPRL-IR-7TNatural + Medical + Remote Sensing73Domain-7Task
DATPRL-IR-8TNatural + Medical + Remote Sensing83Domain-8Task
DATPRL-IR-9TNatural + Medical + Remote Sensing93Domain-9Task
Natural-AiOIR-3TNatural3Natural-3Task
Natural-AiOIR-5TNatural5Natural-5Task
Natural-AiOIR-CombinedNaturalCombined degradationNatural-CDD11

🙏 Acknowledgements

This project is built upon and inspired by several excellent open-source projects and papers:

We sincerely thank the authors for their contributions to the community.

📄 Citation

If this work is useful for your research, please consider citing:

@inproceedings{dong2026datprlir,
  title     = {Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration},
  author    = {Dong, Guanglu and Li, Chunlei and Ren, Chao and Hu, Jingliang and Shi, Yilei and Zhu, Xiao Xiang and Mou, Lichao},
  booktitle = {International Conference on Learning Representations},
  year      = {2026}
}

@article{dong2026learning,
  title   = {Learning Domain-Aware Task Prompt Representations for Multi-Domain All-in-One Image Restoration},
  author  = {Dong, Guanglu and Li, Chunlei and Ren, Chao and Hu, Jingliang and Shi, Yilei and Zhu, Xiao Xiang and Mou, Lichao},
  journal = {arXiv preprint arXiv:2603.01725},
  year    = {2026}
}

📬 Contact

For questions, please open an issue in this repository or contact the authors listed in the paper.