Unifying Design Space of Diffusion Models

• The paper introduces arbitrary, task-specific noise patterns via user-selectable basis functions, expanding the EDM framework for tailored image restoration.
• It demonstrates that integrating structured noise with deterministic reverse sampling maintains computational efficiency while significantly boosting restoration performance.
• The methodology unifies forward and reverse diffusion processes, proving that classical EDMs are a special case within a broader, flexible design space.

Diffusion-based generative models (EDMs) employ a stochastic process that gradually corrupts data by adding noise, and then learn to reverse this process to generate or restore samples. The unification of EDMs provided a systematic understanding of how forward and reverse diffusion, network parameterization, loss weighting, and sampling discretizations influence performance and efficiency. Recent developments have extended the EDM framework by generalizing the structure of noise in the forward process, thereby expanding its applicability to advanced restoration tasks beyond the limitations of pure Gaussian noise. "Elucidating the Design Space of Arbitrary-Noise-Based Diffusion Models" (EDA) (Qiu et al., 24 Jul 2025) demonstrates how arbitrary, task-specific noise patterns can be incorporated into the fundamental architecture of EDMs, leading to both superior practical results and a deeper, more flexible theoretical foundation.

1. Expansion of the Diffusion Model Design Space

Standard EDMs unify the training and sampling of diffusion models via pixel-wise independent Gaussian noise, which serves as the default corruption mechanism in the forward process. While the module design affords considerable flexibility (e.g., in discretization, parameterization, and schedules), the class of allowable noise processes remains restricted to isotropic Gaussian patterns. EDA generalizes this by allowing the forward noise to be an arbitrary linear combination of user-selectable basis functions, yielding a noise process with arbitrary covariance tailored to the restoration task:

$$N = \sum_{m=1}^M \frac{\eta + \epsilon_m}{\eta + 1} \cdot h_{m,x_0}$$

Here, $\epsilon_m \sim \mathcal{N}(0,1)$ are independent, $h_{m,x_0}$ are a collection of basis vectors (which can represent, for example, global smoothness, sharp edges, or local boundary-aware structures), and $\eta$ modulates stochasticity. The resulting noise model allows the forward process to inject noise with spatially correlated, smooth, or heterogeneous properties, bypassing the limitations of pixelwise Gaussianity. This generalization is mathematically encoded as a multivariate Gaussian with structured covariance, $\Sigma_{x_0} = H_{x_0} H_{x_0}^\top$, and the forward SDE is

$$dx = [f(t)x + \varphi_{x_0}(t)]dt + g(t)\sum_{m=1}^M h_{m,x_0} dw_t^m$$

This constructs the foundation for the forward process's arbitrary-noise capability while preserving the core module separation and training objectives of EDM.

2. Roles of Arbitrary Noise Patterns

By selecting appropriate bases $H_{x_0}$, EDA adapts to the structure of degradation in different restoration tasks:

• MRI Bias Field Correction: Noise is modeled as globally smooth, e.g., via low-order polynomials or trigonometric functions, reflecting the smooth intensity bias fields in MR imaging.
• CT Metal Artifact Reduction: Noise is configured to be globally sharp, akin to the streaking artifacts from metal implants; bases can be constructed from the difference between artifact-affected and artifact-free images.
• Natural Image Shadow Removal: The noise pattern is designed to be local and boundary-aware, with basis functions focused around shadow edges, so that the restoration process is attentive to subtle local structure.

This flexibility allows EDA to match the corruption patterns present in real-world data, reducing the transformation distance from degraded to clean images and improving restoration quality while avoiding the information loss associated with Gaussian noise injection in standard EDMs.

3. Computational Efficiency and Theoretical Properties

Despite the increase in noise model complexity, EDA maintains computational cost equivalent to EDM. This is established by the derivation of the deterministic sampling update:

$$\frac{dx}{dt} = \frac{s'(t)}{s(t)}x + \frac{\sigma'(t)}{\sigma(t)}x - \frac{\sigma'(t)s(t)}{\sigma(t)}D(x;\sigma)$$

where $D(x;\sigma)$ is the denoiser network and the update's analytic form does not depend on the complexity of the noise covariance. During the reverse process, the arbitrary noise expansion terms cancel upon substitution of the optimal score function, leaving the sampling algorithm (Euler’s method or higher-order integrators) unchanged in computational structure from EDM. Both training and sampling thus proceed at identical speed and memory cost as classic Gaussian-noise EDMs. The paper rigorously demonstrates that EDM is a special case of EDA when the basis is chosen as the standard one-hot pixel basis and $\eta = 0$.

4. Empirical Validation on Restoration Tasks

EDA's practical impact is validated on three demanding image restoration tasks by selecting task-appropriate noise models:

Restoration Task	Noise Structure	Specific Basis $H_{x_0}$	Performance Highlights
MRI Bias Field Correction	Global smooth	Low-order polynomials/trig functions	Outperforms N4, MICO, ABCNet in SSIM, PSNR, and COCO
CT Metal Artifact Reduction	Global sharp	Artifact-based residuals	Competes with and sometimes exceeds dual-domain approaches
Natural Image Shadow Removal	Local (boundary-aware)	Edge-focused basis	Exceeds baselines in PSNR, SSIM, and LAB-RMSE on ISTD

All EDA evaluations use only five sampling steps, yet often surpass prior diffusion-based approaches (e.g., "Refusion" with 100 steps) and classic non-diffusion restoration pipelines in both fidelity and speed. For example, the MRI correction benchmark sees per-slice processing times reduced from nearly 10 seconds to fractions of a second, highlighting the practical advantages of the method. Shadow removal and metal artifact reduction also achieve leading quantitative and qualitative results due to the tuned basis adaptation.

5. Sampling Steps and Restoration Complexity

A crucial implication of EDA's generalization is the ability to minimize the required number of sampling steps for effective restoration. Because arbitrary noise patterns can be matched to the corruption structure (rather than artificially adding white noise), the diffusion process operates over a much shorter, more meaningful transformation path from input degradation to clean target. The deterministic reverse process begins directly at the degraded image, so every integration step during denoising addresses the actual corruption present, rather than restoring artificially induced Gaussian blur. This results in exceptional practical speedup and is especially significant for clinical or real-time scenarios.

6. Theoretical and Practical Significance

EDA demonstrates that enriching the design space with arbitrary noise models—without altering the overall EDM architecture—permits restoration in scenarios where classic models fail or incur excess computational cost. Theoretical analysis proves that this is achievable without additional complexity in training or inference. By bridging the gap between tailored, data-aware restoration and the universal generative modeling framework of EDM, EDA positions itself as both a generalization and a strict superset of prior approaches. This enables systematic advances in both mathematical modeling and high-throughput, high-fidelity image restoration.

Key equations for the arbitrary-noise framework include:

• Generalized noise: $$N = \sum_{m=1}^M \frac{\eta+\epsilon_m}{\eta+1} h_{m,x_0}$$
• Forward process: $$x_t = s(t)x_0 + s(t)\sigma(t) \sum_m \frac{\eta+\epsilon_m}{\eta+1} h_{m,x_0}$$
• Reverse SDE/PFODE: $$\frac{dx}{dt} = \frac{s'(t)}{s(t)}x + \frac{\sigma'(t)}{\sigma(t)}x - \frac{\sigma'(t)s(t)}{\sigma(t)}D(x;\sigma)$$

In conclusion, EDA (Qiu et al., 24 Jul 2025) extends the EDM framework by unifying arbitrary task-driven noise modeling with the established strengths of modular forward/reverse design, achieving state-of-the-art results and offering a powerful new dimension in the design of diffusion-based restoration systems.