Universal Image Restoration via Internalized Chain-of-Thought Reasoning

2026-06-16

AI Summary

Image restoration under complex mixed degradations is highly ill-posed; existing unified models suffer performance degradation as degradation complexity increases, while multi-stage approaches incur substantial computational overhead and inadequately model degradation interactions. This work reframes restoration as an image editing subtask, fine-tuning a large-scale pretrained editing model and introducing, for the first time, a differentiable chain-of-thought reasoning mechanism. This mechanism is internalized within a single end-to-end model via a Lagrangian-optimization-inspired objective function. The authors also construct CoTIR-Bench, the first large-scale benchmark incorporating explicit reasoning trajectories. Experiments demonstrate that the proposed method significantly outperforms both existing unified models and multi-stage approaches in terms of perceptual quality and fidelity on CoTIR-Bench and real-world composite degradation scenarios.

Abstract

Image restoration seeks to recover high-quality images from degraded inputs but becomes highly ill-posed under complex, mixed degradations. While unified all-in-one models are common, their performance declines as degradation complexity increases. Recent works adopt Chain-of-Thought (CoT) reasoning for multi-round restoration using specialized modules. However, this approach faces two key limitations: (i) increased computational cost due to multi-step processing, and (ii) weak modeling of interactions between degradations during stepwise inference. We introduce CoTIR, a universal image restoration framework that internalizes CoT reasoning within a single model. Concretely, we view image restoration as a specialized subtask of image editing, which implies that a large-scale pre-trained editing model provides a more favorable optimization starting point. Building on this, we fine-tune the model for restoration and further encode structured CoT-style reasoning into the learning objective via a differentiable formulation inspired by Lagrangian optimization, enabling holistic restoration without chaining specialized restorers. To facilitate training and evaluation, we further present CoTIR-Bench, a large-scale benchmark comprising 5.2 million samples with CoT-style reasoning traces. Extensive experiments on CoTIR-Bench and broad real composite degradation scenes show that CoTIR achieves stronger perceptual quality and more competitive fidelity than both all-in-one models and multi-round restoration methods. The source code is available at https://github.com/gy65896/CoTIR.

Problem

• image restoration
• mixed degradations
• Chain-of-Thought reasoning
• computational cost
• degradation interaction

Innovation

• Chain-of-Thought Reasoning
• Universal Image Restoration
• Internalized Reasoning
• Differentiable Optimization
• Image Editing Pretraining