[Paper Review] Towards Effective Multiple-in-One Image Restoration: A Sequential and Prompt Learning Strategy

Xiangtao Kong, Chao Dong+1 more|arXiv (Cornell University)|Jan 7, 2024

Photoacoustic and Ultrasonic Imaging4 citations

TL;DR

This paper proposes a sequential and prompt learning strategy for Multiple-in-One Image Restoration (MiO IR), enabling a single model to effectively handle seven diverse image restoration tasks. By sequentially training on tasks and using explicit or adaptive prompts for task adaptation, the method achieves state-of-the-art performance across CNN and Transformer backbones, with up to +2.09 dB PSNR gain on out-of-distribution sets.

ABSTRACT

While single task image restoration (IR) has achieved significant successes, it remains a challenging issue to train a single model which can tackle multiple IR tasks. In this work, we investigate in-depth the multiple-in-one (MiO) IR problem, which comprises seven popular IR tasks. We point out that MiO IR faces two pivotal challenges: the optimization of diverse objectives and the adaptation to multiple tasks. To tackle these challenges, we present two simple yet effective strategies. The first strategy, referred to as sequential learning, attempts to address how to optimize the diverse objectives, which guides the network to incrementally learn individual IR tasks in a sequential manner rather than mixing them together. The second strategy, i.e., prompt learning, attempts to address how to adapt to the different IR tasks, which assists the network to understand the specific task and improves the generalization ability. By evaluating on 19 test sets, we demonstrate that the sequential and prompt learning strategies can significantly enhance the MiO performance of commonly used CNN and Transformer backbones. Our experiments also reveal that the two strategies can supplement each other to learn better degradation representations and enhance the model robustness. It is expected that our proposed MiO IR formulation and strategies could facilitate the research on how to train IR models with higher generalization capabilities.

Motivation & Objective

• Address the challenge of training a single model to handle multiple, diverse image restoration tasks simultaneously.
• Overcome optimization instability and task interference when training on mixed degradation types.
• Improve model generalization and robustness across in- and out-of-distribution test sets.
• Develop a unified framework that supports high-quality supervision using clean, artifact-free ground-truth images.
• Investigate complementary strategies for stable training and effective task-specific adaptation.

Proposed method

• Propose sequential learning: train the model on one image restoration task at a time in a carefully designed order, rather than mixing all tasks during training.
• Introduce prompt learning to adapt the model to different tasks: use either explicit prompts (e.g., task embeddings) or adaptive visual prompts extracted from input images.
• Apply explicit prompt learning by injecting task-specific conditioning signals into the network via learnable embeddings.
• Implement adaptive prompt learning by dynamically generating task-aware features from the input image itself.
• Combine sequential and prompt learning strategies to improve degradation representation learning and model robustness.
• Use high-quality, artifact-free ground-truth images from curated datasets to avoid low-quality supervision signals.

Experimental results

Research questions

• RQ1 How can we stabilize training when optimizing for multiple, conflicting image restoration objectives simultaneously?
• RQ2 What training strategy enables better generalization across diverse image restoration tasks without performance degradation?
• RQ3 Can prompt learning effectively guide a single model to adapt to different restoration tasks without fine-tuning?
• RQ4 How do sequential and prompt learning strategies complement each other in improving model performance and robustness?
• RQ5 To what extent can a unified model trained with these strategies outperform task-specific models on both in- and out-of-distribution test sets?

Key findings

• Sequential learning improves optimization stability, yielding an average PSNR gain of 0.29 dB for SRResNet and 0.85 dB for SwinIR across all seven tasks.
• The combination of sequential learning and explicit prompt learning achieves a +2.09 dB PSNR improvement on out-of-distribution test sets for SwinIR.
• Adaptive prompt learning improves generalization, especially on unseen degradation types, with a +1.06 dB gain on the Unknown test set for SRResNet.
• The proposed MiO IR formulation with high-quality ground-truth data avoids performance degradation from low-quality supervision signals.
• The sequential and prompt learning strategies enhance degradation representation learning, improving robustness on wild and challenging test sets.
• The method achieves state-of-the-art performance across multiple backbones (SRResNet, SwinIR, Restormer, Uformer), demonstrating broad applicability.