Your Pre-trained Diffusion Model Secretly Knows Restoration

1. Problem Statement

All-in-One Restoration (AiOR) aims to restore images or videos corrupted by diverse degradations (e.g., haze, rain, snow, motion blur, low light) using a single unified model. While recent approaches leverage large pre-trained Text-to-Image/Video (T2IV) Latent Diffusion Models (LDMs) like Stable Diffusion, FLUX, and WAN due to their strong generative priors, existing methods face significant limitations:

• Fine-tuning/ControlNet: Most methods require fine-tuning the backbone or adding heavy ControlNet-style modules. This is computationally expensive, risks overfitting to training distributions, and can weaken the model's original generalization priors.
• Ineffective Prompting: Directly using natural language prompts (e.g., "remove haze") or optimizing text tokens (Token-space optimization) fails to elicit restoration behavior. Instead, these methods tend to suppress noise while retaining the degraded structure, effectively "denoising" the degradation rather than removing it.
• Trajectory Mismatch: Naive prompt learning suffers from a disconnect between training and inference. Training often uses a forward noising process anchored at the degraded image ($z_{deg}$), while inference requires a reverse trajectory moving from degraded to clean ($z_{clean}$). This mismatch causes instability and artifacts.

2. Methodology

The authors propose a framework that unlocks the inherent restoration capabilities of frozen pre-trained diffusion models by learning lightweight prompts directly in the embedding space of the text encoder, rather than the token space.

A. Embedding-Space Prompt Optimization

Instead of optimizing discrete text tokens or using natural language, the method optimizes a continuous context embedding vector ($p$) at the output of the frozen text encoder.

• Observation: While token-space prompts fail, directly optimizing the embedding $p$ allows the frozen backbone to map noisy degraded latents ($z_t^{deg}$) toward clean targets ($z_{clean}$).
• Mechanism: The model minimizes the reconstruction error of the clean latent given the noisy input and the learned prompt $p$.

B. Addressing Trajectory Mismatch via Diffusion Bridges

The authors identify that naive training (using a forward process $z_t = \alpha(t)z_{deg} + \sigma(t)\epsilon$) creates a distribution mismatch during inference. To resolve this, they introduce a Diffusion Bridge formulation:

1. The Problem: Naive training exposes the prompt only to states anchored at $z_{deg}$. Inference, however, requires the model to traverse a path where the signal component transitions from $z_{deg}$ to $z_{clean}$.
2. The Solution: They define a training trajectory that smoothly interpolates between the degraded and clean endpoints while maintaining a noise schedule compatible with the frozen backbone.
3. EBR Bridge: They adopt an Energy-Oriented Diffusion Bridge (EBR) formulation. The intermediate states are constructed as:
   $$z_t^{ebr} = (1-t) \left[ (1-\frac{t}{T_0})z_{clean} + (\frac{t}{T_0})z_{deg} \right] + t\epsilon$$
   This ensures the prompt is trained on a monotone path from noisy-degraded to clean, aligning the training dynamics with the inference sampling trajectory.

C. Implementation Details

• Architecture: The diffusion backbone (WAN for video, FLUX for images) and the text encoder are kept frozen.
• Prompt Design: A separate prompt $p^{(k)}$ is learned for each degradation type $k$.
• Residual Injection: To improve parameter efficiency, the method learns a residual prompt added to the null-text context at the conditioning interface (e.g., cross-attention input) rather than optimizing the full embedding from scratch.
• Mixed Degradations: For mixed degradations, prompts are combined by averaging their predicted velocities during sampling.

3. Key Contributions

1. Discovery of Hidden Restoration Priors: The paper demonstrates that pre-trained diffusion models inherently possess restoration knowledge that is inaccessible via natural language or token-space optimization but can be unlocked via embedding-space conditioning.
2. Trajectory Alignment: The authors identify the train-test trajectory mismatch as a primary cause of instability in naive prompt learning and propose a Diffusion Bridge (specifically EBR-style) to align training and inference dynamics.
3. Lightweight, Generalizable Framework: They instantiate this approach on the WAN (1.3B) video model and FLUX (12B) image model. The method achieves state-of-the-art performance with no backbone fine-tuning and only a small number of learnable parameters (e.g., ~300K for WAN vs. billions for the backbone).

4. Experimental Results

The method was evaluated on diverse datasets for both image and video restoration, including in-distribution (ID), out-of-distribution (OOD), and unseen/mixed degradations.

• Performance:
  • Image Restoration (FLUX): Outperformed or matched state-of-the-art (SOTA) methods like AutoDIR, PixWizard, and FoundIR across metrics (PSNR, SSIM, LPIPS, CLIPIQA, MUSIQ). Notably, it showed superior generalization on OOD datasets (e.g., HazeRD, WeatherBench) and unseen degradations (TOLED, POLED).
  • Video Restoration (WAN): Achieved competitive results against specialized video restoration models (ViWS-Net, AverNet) and significantly outperformed image-based methods applied to video frames. It excelled in temporal consistency metrics (DOVER) due to the strong temporal priors of the pre-trained WAN model.
• Ablation Studies:
  • Bridge vs. Naive: Using the EBR bridge significantly outperformed naive training and DDBM-style bridges, confirming the necessity of trajectory alignment.
  • Embedding vs. Token: Embedding-space optimization was proven essential; token-space optimization failed to restore images.
  • Residual Prompt: Learning a residual prompt was more parameter-efficient and effective than optimizing the full context embedding.

5. Significance

This work fundamentally shifts the paradigm for leveraging pre-trained diffusion models in low-level vision tasks:

• Efficiency: It eliminates the need for expensive fine-tuning or heavy ControlNet modules, making high-quality restoration accessible with minimal computational overhead.
• Generalization: By preserving the frozen backbone's original priors, the method exhibits superior robustness to real-world, complex, and unseen degradations compared to methods that overfit to specific training distributions.
• Insight: It reveals a critical insight into the "black box" of diffusion models: restoration capabilities exist within the model's prior but are gated by the conditioning mechanism, requiring specific embedding-space interventions rather than semantic text prompts.

In summary, the paper proves that pre-trained diffusion models are "secretly" capable of restoration, provided the conditioning mechanism is optimized correctly to align with the model's internal generative dynamics.