The Poisson Midpoint Method for Langevin Dynamics: Provably Efficient Discretization for Diffusion Models

TL;DR

The paper introduces the Poisson Midpoint Method, a novel discretization technique for Langevin Dynamics, achieving faster convergence and efficiency in diffusion models for image generation compared to traditional methods.

Contribution

It proposes the Poisson Midpoint Method, a new discretization approach that approximates Langevin Monte Carlo with larger steps, providing provable quadratic speed-up under weak assumptions.

Findings

Maintains image quality with fewer neural network calls.

Outperforms ODE-based methods in diffusion models.

Achieves significant speed-up over traditional LMC.

Abstract

Langevin Dynamics is a Stochastic Differential Equation (SDE) central to sampling and generative modeling and is implemented via time discretization. Langevin Monte Carlo (LMC), based on the Euler-Maruyama discretization, is the simplest and most studied algorithm. LMC can suffer from slow convergence - requiring a large number of steps of small step-size to obtain good quality samples. This becomes stark in the case of diffusion models where a large number of steps gives the best samples, but the quality degrades rapidly with smaller number of steps. Randomized Midpoint Method has been recently proposed as a better discretization of Langevin dynamics for sampling from strongly log-concave distributions. However, important applications such as diffusion models involve non-log concave densities and contain time varying drift. We propose its variant, the Poisson Midpoint Method, which approximates a small step-size LMC with large step-sizes. We prove that this can obtain a quadratic speed up of LMC under very weak assumptions. We apply our method to diffusion models for image generation and show that it maintains the quality of DDPM with 1000 neural network calls with just 50-80 neural network calls and outperforms ODE based methods with similar compute.