[Paper Review] Non-Uniform Diffusion Models

Georgios Batzolis, Jan Stanczuk+2 more|arXiv (Cornell University)|Jul 20, 2022

Advanced Neuroimaging Techniques and Applications4 citations

TL;DR

This paper introduces non-uniform diffusion models, where different regions of an image diffuse at varying speeds via distinct stochastic differential equations (SDEs), enabling multi-scale generative modeling. The proposed method achieves a 4.4× faster sampling speed at 128×128 resolution with improved FID scores compared to uniform diffusion, while also introducing a novel conditional score estimator (CMDE) that outperforms existing methods in conditional image generation tasks.

ABSTRACT

Diffusion models have emerged as one of the most promising frameworks for deep generative modeling. In this work, we explore the potential of non-uniform diffusion models. We show that non-uniform diffusion leads to multi-scale diffusion models which have similar structure to this of multi-scale normalizing flows. We experimentally find that in the same or less training time, the multi-scale diffusion model achieves better FID score than the standard uniform diffusion model. More importantly, it generates samples $4.4$ times faster in $128 imes 128$ resolution. The speed-up is expected to be higher in higher resolutions where more scales are used. Moreover, we show that non-uniform diffusion leads to a novel estimator for the conditional score function which achieves on par performance with the state-of-the-art conditional denoising estimator. Our theoretical and experimental findings are accompanied by an open source library MSDiff which can facilitate further research of non-uniform diffusion models.

Motivation & Objective

• To address the high inference time of standard diffusion models, especially in high-resolution image generation.
• To explore the potential of non-uniform diffusion—where different image regions diffuse at different rates—as a means to improve efficiency and performance.
• To develop a new conditional score estimator, CMDE, that unifies and improves upon existing conditional denoising and diffusive estimators.
• To provide a theoretical foundation for the conditional denoising estimator (CDE), which had previously been used without formal justification.
• To empirically evaluate and compare conditional score estimation methods across super-resolution, inpainting, and edge-to-image translation tasks.

Proposed method

• The paper proposes a non-uniform diffusion framework in which different spatial regions of an image are governed by distinct SDEs with varying diffusion speeds, leading to multi-scale diffusion dynamics.
• It introduces the conditional multi-speed diffusive estimator (CMDE), a novel estimator for the conditional score function that interpolates between conditional denoising and conditional diffusive estimation.
• The method uses a learnable diffusion speed function σ^y(t) that adapts to the conditional observation y, allowing region-specific diffusion rates.
• A theoretical proof is provided for the consistency of the conditional denoising estimator (CDE), validating its use in practice.
• The framework is implemented in an open-source library, MSDiff, to support reproducibility and further research.
• The approach is evaluated on three conditional image generation tasks: super-resolution, inpainting, and edge-to-image translation, using FID, LPIPS, and diversity metrics.

Experimental results

Research questions

• RQ1 Can non-uniform diffusion, where different image regions diffuse at different speeds, lead to faster and more accurate generative modeling?
• RQ2 How does the proposed CMDE estimator compare to existing conditional score estimators (CDE and CDiffE) in terms of sample quality and training stability?
• RQ3 Does non-uniform diffusion improve the optimization landscape and posterior approximation in conditional image generation?
• RQ4 Is the conditional denoising estimator (CDE) theoretically consistent, and can this be formally proven?
• RQ5 Can the multi-scale diffusion framework achieve better FID scores with reduced training time and significantly faster sampling than standard uniform diffusion models?

Key findings

• The multi-scale diffusion model achieves a 4.4× speed-up in sampling at 128×128 resolution compared to standard uniform diffusion, with no loss in sample quality.
• In less training time, the multi-scale model reaches better FID scores than the uniform baseline, demonstrating improved training efficiency.
• The CMDE estimator outperforms both CDE and CDiffE in FID scores on inpainting and super-resolution tasks, indicating improved posterior approximation.
• On edge-to-image translation, CDiffE fails to produce consistent samples, while CMDE and CDE perform significantly better, with CMDE achieving the best FID.
• The paper provides a formal proof of consistency for the conditional denoising estimator (CDE), offering a theoretical justification for its widespread use.
• The open-source library MSDiff is released to support further research in non-uniform and conditional diffusion models.