Reverse Transition Kernel: A Flexible Framework to Accelerate Diffusion Inference

TL;DR

This paper introduces a flexible reverse transition kernel framework for diffusion models, enabling efficient inference with fewer subproblems and providing strong theoretical convergence guarantees.

Contribution

The authors propose a novel RTK framework that reduces the number of subproblems to nearly one, and develop RTK-MALA and RTK-ULD algorithms with improved convergence guarantees for diffusion inference.

Findings

RTK-ULD achieves $ ilde O(d^{1/2} ext{epsilon}^{-1})$ convergence in TV distance.

RTK-MALA attains $ ext{O}(d^{2} ext{log}(d/ ext{epsilon}))$ convergence rate.

Numerical experiments support the theoretical convergence improvements.

Abstract

To generate data from trained diffusion models, most inference algorithms, such as DDPM, DDIM, and other variants, rely on discretizing the reverse SDEs or their equivalent ODEs. In this paper, we view such approaches as decomposing the entire denoising diffusion process into several segments, each corresponding to a reverse transition kernel (RTK) sampling subproblem. Specifically, DDPM uses a Gaussian approximation for the RTK, resulting in low per-subproblem complexity but requiring a large number of segments (i.e., subproblems), which is conjectured to be inefficient. To address this, we develop a general RTK framework that enables a more balanced subproblem decomposition, resulting in $\tilde O(1)$ subproblems, each with strongly log-concave targets. We then propose leveraging two fast sampling algorithms, the Metropolis-Adjusted Langevin Algorithm (MALA) and Underdamped Langevin Dynamics (ULD), for solving these strongly log-concave subproblems. This gives rise to the RTK-MALA and RTK-ULD algorithms for diffusion inference. In theory, we further develop the convergence guarantees for RTK-MALA and RTK-ULD in total variation (TV) distance: RTK-ULD can achieve $\epsilon$ target error within $\tilde{\mathcal O}(d^{1/2}\epsilon^{-1})$ under mild conditions, and RTK-MALA enjoys a $\mathcal{O}(d^{2}\log(d/\epsilon))$ convergence rate under slightly stricter conditions. These theoretical results surpass the state-of-the-art convergence rates for diffusion inference and are well supported by numerical experiments.