Technical Summary: Diffusion Model-Based Image Restoration via Interactive Learning of Determinism and Stochasticity

Problem Statement

Image restoration (IR) aims to reconstruct latent visual information from degraded observations, a task complicated by the significant distributional discrepancies among different types of degradation (e.g., noise, rain, shadows, low-light) in feature space. Existing deep learning methods, including CNNs, Transformers, and GANs, often struggle to generalize across diverse IR tasks without task-specific redesigns. While diffusion models offer high fidelity and interpretability, current approaches face two primary limitations:

1. Inefficiency in IR: Standard diffusion models start the reverse process from pure noise, which is inefficient for IR tasks where the input already contains structural information.
2. Static Modeling of Determinism and Stochasticity: Previous attempts to unify deterministic (residual-based) and stochastic (noise-based) diffusion, such as the Residual Denoising Diffusion Model (RDDM), employ independent U-Nets for each path. This static separation fails to capture the distinct probability distributions of different degradation types and may introduce excessive diversity that hinders precise reconstruction.

Methodology

The authors propose a novel diffusion framework that employs Interactive Learning of Determinism and Stochasticity (ILDS) to address these limitations. The core architecture utilizes a dual-path network based on two simple U-Nets.

1. Forward Diffusion Process Construction

Unlike standard diffusion that adds only noise, the proposed forward process incorporates both residuals (determinism) and noise (stochasticity). The transition is defined as:
$$x_t = x_{t-1} + \alpha_t x_{res} + \beta_t \epsilon_{t-1}$$
where $x_{res} = x_{in} - x_0$ represents the degradation residual, and $\epsilon$ is Gaussian noise. This allows the forward trajectory to preserve degradation-related structures, enabling the reverse process to start from states closer to the target than pure noise, thereby reducing the required sampling steps.

2. Interactive Learning Mechanism

To overcome the limitations of independent modeling in RDDM, the authors introduce an interactive learning strategy within the U-Net architecture:

• Dual-Path Prediction: Two U-Nets are used to predict the noise ($\epsilon_\theta$) and the residual ($x_{res}^\theta$) respectively.
• Feature Fusion: Instead of operating in isolation, the encoders of the two U-Nets share features. Specifically, after the first downsampling block, features from both paths are concatenated and processed through $1 \times 1$ convolutions ($\tilde{h}_{res}$ and $\tilde{h}_\epsilon$).
• Dynamic Time-Dependent Modeling: The timestep inputs for the two paths are dynamically modulated by parameters controlling temporal accumulation ($\bar{\alpha}_t$) and noise intensity ($\bar{\beta}_t$). This allows the model to adaptively adjust the contribution of deterministic and stochastic components based on the specific feature probability distribution of the degradation type.
• Shared Context: The interacted features provide shared cross-path contextual information during the decoding phase, ensuring that the deterministic guidance and stochastic diversity influence each other.

3. Training and Generation

The model is trained to minimize the Kullback-Leibler (KL) divergence between the true posterior and the learned reverse transition. The final loss function is the sum of squared errors for both noise and residual predictions:
$$L = L_{res}(\theta) + L_{\epsilon}(\theta)$$
Notably, these components are unweighted, as the interactive mechanism facilitates a self-adaptive equilibrium during training. During inference, the model generates high-quality images by iteratively sampling from the conditional distribution, with the process becoming fully deterministic when the noise scale factor $\eta = 0$.

Key Contributions

1. Novel Diffusion Framework: The design of a framework that interactively models deterministic and stochastic aspects of image degradation, allowing the diffusion process to adapt to tasks requiring either high fidelity (precision) or diversity.
2. Dynamic Time-Dependent Modeling: The incorporation of temporal accumulation and noise intensity parameters to dynamically modulate timestep inputs for dual diffusion paths, enhancing generalization across diverse IR scenarios.
3. Efficient Dual-Path Architecture: The implementation of the dual-path diffusion using only two simple U-Nets with shared encoder features. This design achieves competitive performance with minimal parameter overhead and high scalability.

Experimental Results

The proposed method was evaluated on four representative IR tasks: Image Denoising, Shadow Removal, Low-Light Enhancement, and Image Deraining.

• Performance: The method achieved superior or competitive performance across all tasks.
  • Denoising (SIDD): Achieved the best PSNR (41.67 dB) and SSIM (0.9385), outperforming RDDM and other baselines.
  • Shadow Removal (ISTD): Secured the best PSNR (29.29 dB) and MAE (8.776), with the second-best SSIM.
  • Deraining (Raindrop): Achieved the best PSNR (31.30 dB) and SSIM (0.9460), successfully removing rain streaks where other methods left artifacts.
  • Low-Light Enhancement (LOL): Achieved the second-best performance (PSNR 24.04, SSIM 0.9118), slightly trailing LLFormer in metrics but maintaining strong structural preservation.
• Efficiency: The interactive learning significantly improved sampling efficiency. The model required only 1–4 sampling steps to achieve high-quality results, compared to hundreds or thousands for other diffusion-based methods (e.g., DiffusionVI required 1000 steps).
• Ablation Studies:
  • Removing the interaction module resulted in performance drops across all tasks, confirming the benefit of feature fusion.
  • Using dynamic time-dependent modeling ($\bar{\alpha}_t, \bar{\beta}_t$) outperformed static timestep modeling ($t$).
  • Optimal sampling steps varied by task (1 for low-light, 2 for denoising, 3 for shadow removal, 4 for deraining).

Significance and Claims

The paper claims that the proposed framework offers a unified solution for multi-task image restoration that balances reconstruction accuracy and diversity. By integrating deterministic and stochastic modeling through interactive learning, the approach provides stronger interpretability than GAN-based methods and better generalization than task-specific CNNs or Transformers.

The authors acknowledge limitations, particularly in low-light enhancement tasks with large illumination gradients where brightness restoration may be insufficient compared to specialized baselines. They also note the current reliance on paired clean-degraded data for supervised training. Future work is directed toward optimizing diffusion equations for illumination-sensitive tasks and developing unsupervised training strategies to reduce data dependency.

Overall, the study demonstrates that a simple, interactive dual-path diffusion architecture can effectively handle the heterogeneity of image degradation processes, offering a scalable and efficient alternative for diverse image restoration challenges.