The Missing U for Efficient Diffusion Models

TL;DR

This work tackles the efficiency bottleneck in diffusion probabilistic models by replacing the standard U-Net denoiser with a continuous U-Net designed around second-order Neural Ordinary Differential Equations and time embeddings. The approach yields substantial parameter and FLOP reductions, faster reverse-step convergence, and competitive image synthesis quality, while delivering improved denoising performance in perceptual metrics. The authors provide a theoretical basis comparing the Probability Flow ODE to SDEs and demonstrate memory-efficient backpropagation via adjoint methods, making the method attractive for deployment in resource-constrained settings. Overall, the paper presents a modular, orthogonal advancement that can augment existing diffusion-model improvements and extend efficiently to downstream tasks.

Abstract

Diffusion Probabilistic Models stand as a critical tool in generative modelling, enabling the generation of complex data distributions. This family of generative models yields record-breaking performance in tasks such as image synthesis, video generation, and molecule design. Despite their capabilities, their efficiency, especially in the reverse process, remains a challenge due to slow convergence rates and high computational costs. In this paper, we introduce an approach that leverages continuous dynamical systems to design a novel denoising network for diffusion models that is more parameter-efficient, exhibits faster convergence, and demonstrates increased noise robustness. Experimenting with Denoising Diffusion Probabilistic Models (DDPMs), our framework operates with approximately a quarter of the parameters, and $\sim$ 30\% of the Floating Point Operations (FLOPs) compared to standard U-Nets in DDPMs. Furthermore, our model is notably faster in inference than the baseline when measured in fair and equal conditions. We also provide a mathematical intuition as to why our proposed reverse process is faster as well as a mathematical discussion of the empirical tradeoffs in the denoising downstream task. Finally, we argue that our method is compatible with existing performance enhancement techniques, enabling further improvements in efficiency, quality, and speed.

Figures (14)

• Figure 1: Visual representation of our framework featuring implicit deep layers tailored for denoising in the reverse process of a DDPM, enabling the reconstruction of the original data from a noise-corrupted version.
• Figure 2: Architectural components of the continuous U-Net. On the left, the Dynamic ODE Block represents the core unit of our continuous model, detailing the integration of the ODE solver and function approximator within the network structure. The right panel expands on the ODE Function Approximator, highlighting the convolutional layers, group normalisation, time embeddings, and the incorporation of scale, shift operations, and residual connections to accurately adapt the network's dynamics during the diffusion process.
• Figure 3: Randomly selected generated samples by our model (right) and the baseline U-Net-based DDPM (left) trained on CelebA and LSUN Church.
• Figure 4: Visualisation of noise accumulation in images over increasing timesteps. As timesteps advance, the images exhibit higher levels of noise, showcasing the correlation between timesteps and noise intensity. The progression highlights the effectiveness of time embeddings in predicting noise magnitude at specific stages of the diffusion process.
• Figure 5: Original image (left), with Gaussian noise (second), and denoised using our continuous U-Net (third and fourth). As noise increases, U-Net struggles to recover the fine-grained details such as the glasses.

Theorems & Definitions (1)

• Proposition 4.1