Chain-of-Restoration: Multi-Task Image Restoration Models are Zero-Shot Step-by-Step Universal Image Restorers

TL;DR

This work introduces Universal Image Restoration (UIR), a zero-shot framework that handles both isolated and composite degradations by training on a base set of degradations rather than all combinations. The Chain-of-Restoration (CoR) mechanism augments pre-trained multi-task image restoration models with a Degradation Discriminator to iteratively remove degradation bases, enabling step-by-step restoration of unknown composites. The authors also provide the first UIR dataset UIRD-12 and demonstrate through extensive experiments that CoR substantially improves composite-degradation removal, often matching or surpassing state-of-the-art methods trained on full degradation sets while reducing training burden. The approach highlights practical gains in efficiency and generalization for universal restoration, with insights into degradation coupling, base selection, and sequence control for future improvements.

Abstract

Despite previous image restoration (IR) methods have often concentrated on isolated degradations, recent research has increasingly focused on addressing composite degradations involving a complex combination of multiple isolated degradations. However, current IR methods for composite degradations require building training data that contain an exponential number of possible degradation combinations, which brings in a significant burden. To alleviate this issue, this paper proposes a new task setting, i.e. Universal Image Restoration (UIR). Specifically, UIR doesn't require training on all the degradation combinations but only on a set of degradation bases and then removing any degradation that these bases can potentially compose in a zero-shot manner. Inspired by the Chain-of-Thought that prompts large language models (LLMs) to address problems step-by-step, we propose Chain-of-Restoration (CoR) mechanism, which instructs models to remove unknown composite degradations step-by-step. By integrating a simple Degradation Discriminator into pre-trained multi-task models, CoR facilitates the process where models remove one degradation basis per step, continuing this process until the image is fully restored from the unknown composite degradation. Extensive experiments show that CoR can significantly improve model performance in removing composite degradations, achieving comparable or better results than those state-of-the-art (SoTA) methods trained on all degradations.

Figures (14)

• Figure 1: Comparison of CoT and CoR. (a) Zero-Shot CoT for LLMs. The core idea of CoT is to ask the LLM to response to break down the question into smaller components and solve them step by step. (b) Our proposed CoR. We take the composite degradation as combination of multiple degradation bases. Using multi-task models that are trained on these bases, we ask the model to remove the composite degradation step by step. In this paper, "multi-task model" refers to any image restoration model trained on more than one degradations.
• Figure 2: Comparison of task settings and classification of previous image restoration models. (I) One-to-One: In this setting, models are trained on an isolated degradation and tested on it. (II) One-to-Many: In this setting, models are trained on multiple isolated degradations simultaneously and tested on them simultaneously. (III) One-to-Composite: In this setting, models are trained on multiple isolated degradations and composite degradations simultaneously and tested on them simultaneously. (IV) One-to-Universal: In this setting, models are trained on a set of base degradations simultaneously and tested on combinations of these base degradations.
• Figure 3: (a) To master long jump, the only requirements are knowing how to run and jump, executed step by step. (b) For an image degraded by both rain and haze, restoring it requires only a model that can dehaze and derain, without any additional training.
• Figure 4: Visualization of $TR_n(k)$ and $IR_n(k)$ when $n=20$.
• Figure 5: The visualization of the step-by-step process of CoR with different methods on UIRD-12 and CDD-11.