CoTIR: Universal Image Restoration

• CoTIR is a universal image restoration framework that integrates chain-of-thought reasoning into a single model to tackle diverse, complex degradations.
• It reframes restoration as a specialized image editing task by fine-tuning a large-scale pre-trained editing model, improving generalization on composite degradations.
• CoTIR employs differentiable Lagrangian optimization with adaptive constraint enforcement, achieving efficient runtime and state-of-the-art perceptual quality.

Searching arXiv for CoTIR and related restoration papers. CoTIR is a universal image restoration framework that internalizes Chain-of-Thought reasoning within a single model for recovering high-quality images from degraded inputs under complex, mixed degradations. It is introduced in "Universal Image Restoration via Internalized Chain-of-Thought Reasoning" (Guo et al., 16 Jun 2026), which frames restoration as a specialized subtask of image editing, fine-tunes a large-scale pre-trained editing model for restoration, and encodes structured CoT-style reasoning into the learning objective through a differentiable formulation inspired by Lagrangian optimization. The framework is accompanied by CoTIR-Bench, a benchmark of 5.2 million samples with CoT-style reasoning traces, and is positioned against both all-in-one restoration models and multi-round CoT-based restoration pipelines (Guo et al., 16 Jun 2026).

1. Definition and problem setting

CoTIR addresses image restoration in settings where degradation is complex, mixed, and highly ill-posed. The motivating observation is that unified all-in-one models decline in performance as degradation complexity increases, while recent approaches that adopt Chain-of-Thought reasoning for multi-round restoration rely on sequential inference through specialized modules, which increases computational cost and weakly models interactions between degradations (Guo et al., 16 Jun 2026).

Within this formulation, CoTIR replaces explicit multi-step tool chaining with an internalized reasoning process learned inside a single model. The paper identifies two limitations of recent CoT-based restoration systems: increased computational cost due to multi-step processing, and weak modeling of interactions between degradations during stepwise inference. CoTIR is designed to address both by performing holistic restoration without chaining specialized restorers (Guo et al., 16 Jun 2026).

This distinguishes CoTIR from a broader set of CoT-related methods in other domains. For example, Chain-of-Thought has been used in composed image retrieval through caption generation, ranking, or explicit retrieval-oriented reasoning in systems such as CSMCIR (Qian et al., 7 Jan 2026), SDR-CIR (Sun et al., 4 Feb 2026), and CIR-CoT (Lin et al., 9 Oct 2025). A plausible implication is that CoTIR belongs to a separate line of work in which CoT is not exposed as an external reasoning trace for retrieval or ranking, but is instead encoded as an internal restoration prior (Guo et al., 16 Jun 2026).

2. Framing restoration as image editing

A central premise of CoTIR is that image restoration can be viewed as a specialized image editing problem. On this basis, the framework uses a large-scale pre-trained editing model, specifically the FLUX model, as its initialization. The paper states that such models already exhibit strong zero-shot capability in restoration, though fine-tuning remains necessary for optimal performance (Guo et al., 16 Jun 2026).

The argument is formalized through an initialization-gap perspective. Let $\theta_{\text{edit}}$ denote the parameters of a pre-trained editing model, $\theta_{\text{ir}}$ those of a restoration-specific model, and $D_{\mathcal{R}}$ the distribution of universal or degraded images. The gap from an initialization $\theta_0$ to an optimal restoration model is written as

$$\Delta_{D_{\mathcal{R}}}(\theta_0)=\mathbb{E}_{(x,y)\sim D_{\mathcal{R}}}\big[\log p_{\theta_{\mathcal{R}}^*}(y|x) - \log p_{\theta_0}(y|x)\big].$$

The paper then states the ordering

$$\Delta_{D_{\mathcal{R}}}(\theta_{\text{edit}})\leq \Delta_{D_{\mathcal{R}}}(\theta_{\text{ir}})\leq \Delta_{D_{\mathcal{R}}}(\theta_{\text{rand}}),$$

which is used to motivate editing-model initialization both empirically and theoretically (Guo et al., 16 Jun 2026).

The significance of this design choice is not merely architectural. It suggests that restoration under diverse composite degradations benefits from priors learned for generic image transformation rather than from narrowly restoration-specific initialization. The paper explicitly reports that using the editing model, rather than a restoration-specific or random initializer, improves generalization to novel degradations (Guo et al., 16 Jun 2026).

3. Internalized Chain-of-Thought reasoning

CoTIR encodes a "Thinking $\to$ Planning $\to$ Action" process inside the model. The reasoning process is mapped to three structured intermediate representations: $c_s$ for sharp feature description, $c_d$ for degradation pattern description, and $\theta_{\text{ir}}$0 for restoration plan (Guo et al., 16 Jun 2026).

The three phases are defined as follows. "Thinking" disentangles features into scene or clean content and degradation patterns. "Planning" analyzes interactions between these representations to formulate a restoration strategy. "Action" restores the corrupted image by holistically applying the learned plan. Rather than producing a discrete multi-round trajectory at inference time, CoTIR supervises these intermediate reasoning outputs during training and uses them as explicit constraints (Guo et al., 16 Jun 2026).

This internalization is the paper’s central conceptual move. Prior CoT-based restoration pipelines, including Q-Agent, Chain-of-Restoration, and RestoreAgent, are described as sequentially applying specialized restoration modules, each attending to a specific degradation. CoTIR instead learns a unified latent process that represents scene content, degradation structure, and restoration strategy jointly (Guo et al., 16 Jun 2026). This suggests that the method treats degradation composition as an interacting system rather than as a sequence of independent corrections.

The paper further reports that split supervision for each reasoning phase is more effective than a holistic union-constraint formulation or a no-constraint baseline. This is presented in the ablations as evidence that the structured decomposition into $\theta_{\text{ir}}$1, $\theta_{\text{ir}}$2, and $\theta_{\text{ir}}$3 is not merely interpretive but functionally important for optimization and performance (Guo et al., 16 Jun 2026).

4. Optimization and architectural formulation

CoTIR trains the model to predict both the restored output image $\theta_{\text{ir}}$4 and the intermediate reasoning outputs $\theta_{\text{ir}}$5. The main objective is a Flow Matching loss in latent space, subject to soft constraints that require the predicted reasoning variables to remain close to their targets within tolerances $\theta_{\text{ir}}$6 (Guo et al., 16 Jun 2026):

$\theta_{\text{ir}}$7

The constrained problem is converted into a Lagrangian:

$\theta_{\text{ir}}$8

with learnable Lagrange multipliers $\theta_{\text{ir}}$9 (Guo et al., 16 Jun 2026).

Training uses a primal-dual procedure to solve

$D_{\mathcal{R}}$0

The primary optimizer updates model parameters $D_{\mathcal{R}}$1, while a secondary optimizer updates $D_{\mathcal{R}}$2, increasing weights when constraints are violated and relaxing them otherwise (Guo et al., 16 Jun 2026).

The ablation studies attribute better optimization dynamics and better performance to adaptive, learnable Lagrange multipliers than to fixed weights. The same section reports that, in the CoT Adapter, Deep Gated-Attention outperforms Cross-Attention. Together, these findings characterize CoTIR not only as a conceptual framework for internalized reasoning, but also as a specific constrained optimization design whose behavior depends on adaptive constraint enforcement and the chosen adapter architecture (Guo et al., 16 Jun 2026).

5. CoTIR-Bench

CoTIR-Bench is introduced as a large-scale benchmark for training and evaluating universal image restoration with explicit CoT-style reasoning traces. The benchmark contains 5.2 million samples and is built by aggregating more than 60 existing datasets (Guo et al., 16 Jun 2026).

Its coverage includes multiple degradations such as rain, haze, low-light, noise, blur, compression, underwater, and shadows, with both single and composite cases. For each low-quality and high-quality image pair, the benchmark provides a triple consisting of scene description, degradation type or types, and restoration plan. These reasoning traces are generated using advanced vision-LLMs, including Qwen2.5-VL, under a designed instruction template that takes degradation types and low- and high-quality images as input (Guo et al., 16 Jun 2026).

The dataset construction process also includes filtering and validation, including CLIP consistency checks, and the paper notes that a standardized test set is curated for unified evaluation (Guo et al., 16 Jun 2026). This gives CoTIR-Bench a dual role. It is simultaneously a supervision source for the structured reasoning variables and an evaluation substrate for broad comparison across restoration regimes.

A plausible implication is that CoTIR-Bench operationalizes internalized CoT by supplying supervision at the level of restoration cognition rather than only at the level of input-output image pairs. In that respect, it resembles reasoning-annotated datasets developed in other CoT settings, although its target domain is image restoration rather than language reasoning or retrieval (Guo et al., 16 Jun 2026).

6. Empirical performance and comparative findings

The empirical comparison in CoTIR spans all-in-one restoration models and multi-round CoT-based methods. The paper reports that CoTIR consistently achieves state-of-the-art or near-best performance on no-reference perceptual metrics including CLIP-IQA+, Q-Align, LIQE, and MACLIP, especially on composite degradation scenarios where degradation types interact (Guo et al., 16 Jun 2026).

On full-reference metrics such as PSNR and SSIM, CoTIR remains highly competitive, although the paper notes that some CNN- or Transformer-based methods may be slightly better in strictly pixel-wise scores because they do not aggressively remove mild residual degradations (Guo et al., 16 Jun 2026). This distinction is important for interpreting the method’s strengths: the reported advantage is strongest in perceptual quality and restoration realism under degradation composition, rather than exclusively in distortion-minimization metrics.

The framework is also reported to outperform both all-in-one and multi-round baselines on complex, real-world compositional degradation, including difficult unseen degradation combinations. Qualitative and quantitative results are described as showing stronger generalization than even dedicated multi-round agentic pipelines in such settings (Guo et al., 16 Jun 2026).

Efficiency is another explicit outcome. Using a Flux-2-4B backbone with 5 sampling steps on 512$D_{\mathcal{R}}$3512 images, the paper gives an example runtime of 1.84s for CoTIR and characterizes it as much lower than DA-CLIP, LucidFlux, and agentic pipelines (Guo et al., 16 Jun 2026). The same section states that user studies with 30 volunteers and 200 scores each rate CoTIR’s outputs most favorably for perceptual quality, and that real-scene experiments on both seen and unseen scenarios show clearer, more realistic restoration than the baselines (Guo et al., 16 Jun 2026).

7. Position within CoT research and limitations of interpretation

CoTIR occupies a specific position within the broader Chain-of-Thought literature. In language modeling, recent work has focused on improving CoT efficiency, such as Constraint-Rectified Training, which minimizes reasoning length while guarding accuracy through constrained optimization (Wu et al., 13 Feb 2026), or on converting explicit rationales into reusable latent functional units, as in CIRF (Lee et al., 27 May 2026). Other work examines whether CoT traces faithfully reflect internal computation, as in CIE-Scorer (Shen et al., 25 May 2026). CoTIR shares with these papers an interest in internalizing or constraining reasoning, but applies that logic to image restoration rather than to textual reasoning tasks.

Within vision-language systems, CoT has also been used for composed image retrieval. CSMCIR employs Multi-level Chain-of-Thought prompting to generate target captions and support symmetric query-target alignment (Qian et al., 7 Jan 2026); SDR-CIR uses Selective CoT and Semantic Debias Ranking for training-free zero-shot retrieval (Sun et al., 4 Feb 2026); CIR-CoT introduces explicit end-to-end reasoning for interpretable composed image retrieval (Lin et al., 9 Oct 2025). CoTIR differs in that its CoT is not a textual rationale used directly for retrieval or ranking, but a structured supervisory signal internalized into a restoration model (Guo et al., 16 Jun 2026).

The paper makes several strong claims about novelty, including that CoTIR is the first to internalize CoT reasoning as part of the model’s own latent representations and training constraints, and that CoTIR-Bench is the first 5M+ sample benchmark for universal restoration with explicit CoT annotations (Guo et al., 16 Jun 2026). These claims define CoTIR’s place in the literature as an attempt to recast CoT from an inference-time procedural scaffold into a training-time structural prior.

At the same time, the article’s evidence base is centered on the reported benchmark results, ablations, runtime example, and user study in the original paper (Guo et al., 16 Jun 2026). This suggests that further evaluation across additional restoration backbones, alternative editing initializations, and different forms of reasoning annotation would be relevant directions for confirming the generality of the framework, although such extensions are not claimed in the paper itself.