IR-Flow: Bridging Discriminative and Generative Image Restoration via Rectified Flow

1. Problem Statement

Image restoration aims to recover clean images from degraded inputs (e.g., rain, noise, blur). The field currently faces a fundamental trade-off between discriminative and generative paradigms:

• Discriminative Methods: Typically use supervised learning to minimize pixel-wise distortion (e.g., $\ell_2$ loss). While efficient (single-step inference), they often produce overly smooth outputs that lack high-frequency details and perceptual realism because they learn the conditional expectation $E[x|y]$.
• Generative Methods (Diffusion/SDE): Model the full data distribution to capture textures and details. However, they suffer from inefficient multi-step sampling (often requiring 50–100 steps) and complex noise-residual coupling, where the model must simultaneously learn to remove noise and reconstruct the residual, leading to difficult optimization and slow inference.

The core challenge is to create a unified framework that achieves the efficiency of discriminative models (few-step inference) while maintaining the perceptual quality of generative models, without the computational burden of standard diffusion processes.

2. Methodology: IR-Flow

The authors propose IR-Flow, a framework based on Rectified Flow that establishes a direct, linear transport map between the degraded image domain and the clean image domain.

A. Multilevel Data Distribution Flows

Instead of treating degradation as a single state, the authors construct multilevel data distribution flows.

• They artificially create intermediate states with controlled degradation levels between the clean image ($X_0$) and the degraded image ($X_1$).
• This expands the model's ability to learn from and adapt to varying degrees of degradation, reinforcing the local structure of the distribution and preventing coarse estimates in sparse regions.
• Mathematically, this creates a continuous path $X_t = tX_1 + (1-t)X_0$ for $t \in [0, 1]$.

B. Cumulative Velocity Field (CVF)

Standard Rectified Flow learns an instantaneous velocity field $v(X_t, t)$ representing the tangent direction of the trajectory. The authors argue this requires small step sizes for accurate integration.

• Innovation: They propose a Cumulative Velocity Field defined as $\hat{v}(X_t, t) = E[X_t - X_0 | X_t]$.
• Mechanism: Instead of pointing along the tangent, $\hat{v}$ points directly from the current intermediate state $X_t$ toward the target clean state $X_0$.
• Benefit: This reduces the "transport energy" by a factor of 3 compared to standard velocity fields ($A[\hat{v}] = \frac{1}{3}A[v]$), leading to faster convergence, lower modeling complexity, and more direct trajectories.

C. Multi-step Consistency Training (MCT)

To ensure the learned trajectories remain linear and coherent during inference, a Multi-step Consistency Loss is introduced.

• Objective: The model is trained not only to match the velocity field but also to ensure that if the ODE solver takes multiple steps from a noisy state, the final reconstruction aligns with the ground truth.
• Loss Function: $L = L_S + \lambda_{MCT} L_{MCT}$, where $L_{MCT}$ measures the distance between the ODE-sampled result and the clean target.
• Result: This enforces trajectory coherence, allowing for high-quality restoration with very few sampling steps (1–4 steps).

3. Key Contributions

1. Unified Framework: IR-Flow bridges the gap between discriminative and generative paradigms by constructing a direct linear transport flow between degraded and clean images, eliminating the need for complex noise-residual coupling.
2. Cumulative Velocity Field (CVF): A novel formulation that captures the directional flow toward the target state, significantly accelerating training convergence and reducing transport energy compared to standard Rectified Flow.
3. Few-Step Consistency: The introduction of a trajectory consistency constraint allows the model to achieve state-of-the-art performance with only 1 to 4 sampling steps, drastically reducing inference time.
4. Robustness: The multilevel flow construction enables strong generalization to Out-of-Distribution (OOD) degradations and varying noise levels without retraining.

4. Experimental Results

The method was evaluated on Image Deraining, Denoising, Raindrop Removal, Super-Resolution, Deblurring, and Dehazing.

• Quantitative Performance:
  • Deraining (Rain100H/L): Achieved SOTA PSNR (32.69 dB) and SSIM with only 1-2 NFE (Number of Function Evaluations), outperforming IR-SDE (which requires 100 NFEs) and supervised methods like Restormer.
  • Denoising: On the McMaster dataset with $\sigma=25$, IR-Flow achieved 33.15 PSNR in 1 step, surpassing IR-SDE (32.39 PSNR in 22 steps) and DnCNN.
  • Efficiency: IR-Flow is 100x faster than IR-SDE and 7x faster than RDDM at inference time, while maintaining comparable or better perceptual quality (LPIPS).
• Qualitative Performance: Visual results show that IR-Flow recovers fine textures and details (e.g., fur, rain streaks) better than discriminative models, which tend to be blurry, and matches the quality of multi-step diffusion models with a fraction of the computation.
• Generalization: The model demonstrated strong robustness on real-world internet images and varying noise levels (unified training on $\sigma \in [0, 50]$), outperforming methods trained on specific noise levels.

5. Significance

IR-Flow represents a significant shift in image restoration by demonstrating that generative capabilities do not require expensive multi-step sampling.

• Efficiency: It solves the "speed vs. quality" dilemma, offering diffusion-like perceptual quality with the speed of a single-pass discriminative network.
• Physical Interpretability: By modeling a linear transport flow with cumulative velocity, the method provides a more physically interpretable and mathematically grounded approach than noise-perturbation-based diffusion.
• Practicality: The ability to handle diverse degradation levels with a single unified model makes it highly suitable for real-world deployment where degradation types and intensities vary unpredictably.

In summary, IR-Flow successfully unifies the best of both worlds: the distortion fidelity of supervised learning and the perceptual realism of generative modeling, achieved through an efficient, few-step ODE-based framework.