Elucidating the Design Space of Arbitrary-Noise-Based Diffusion Models

TL;DR

This paper introduces EDA, a flexible diffusion model framework that handles arbitrary noise patterns, improving image restoration tasks with fewer sampling steps and maintaining computational efficiency.

Contribution

EDA extends the diffusion model design space to include arbitrary noise, offering a unified, modular approach with theoretical guarantees and practical validation on diverse image restoration tasks.

Findings

EDA achieves competitive results with only 5 sampling steps.

It effectively handles various noise types in medical and natural images.

EDA maintains computational efficiency despite increased noise complexity.

Abstract

Although EDM aims to unify the design space of diffusion models, its reliance on fixed Gaussian noise prevents it from explaining emerging flow-based methods that diffuse arbitrary noise. Moreover, our study reveals that EDM's forcible injection of Gaussian noise has adverse effects on image restoration task, as it corrupts the degraded images, overextends the restoration distance, and increases the task's complexity. To interpret diverse methods for handling distinct noise patterns within a unified theoretical framework and to minimize the restoration distance, we propose EDA, which Elucidates the Design space of Arbitrary-noise diffusion models. Theoretically, EDA expands noise pattern flexibility while preserving EDM's modularity, with rigorous proof that increased noise complexity introduces no additional computational overhead during restoration. EDA is validated on three representative medical image denoising and natural image restoration tasks: MRI bias field correction (global smooth noise), CT metal artifact removal (global sharp noise) and natural image shadow removal (local boundary-aware noise). With only 5 sampling steps, competitive results against specialized methods across medical and natural tasks demonstrate EDA's strong generalization capability for image restoration. Code is available at: https://github.com/PerceptionComputingLab/EDA.