Entropy-Based Dimension-Free Convergence and Loss-Adaptive Schedules for Diffusion Models
Ahmad Aghapour, Erhan Bayraktar, Ziqing Zhang
TL;DR
The paper introduces a dimension-free framework for diffusion-model discretization analysis based on Shannon entropy, deriving a bound of $O(H^2/K)$ on KL divergence between target and generated distributions. It recasts discretization error as an MMSE functional and proves an entropy-based bound on the MMSE derivative, enabling a geometrically spaced (log-SNR) grid and a two-term KL bound separating entropy-driven discretization from model-parameter errors. Building on this theory, the authors propose LAS (Loss-Adaptive Schedule), a training-loss-informed discretization strategy that can be computed post-training via a dynamic-programming optimization, and show substantial sampling improvements over common schedules. The approach is demonstrated on toy GMMs and a large-scale ImageNet latent-diffusion model, highlighting practical gains in sampling efficiency and quality with fewer function evaluations. Together, the results provide both a theoretical dimension-free convergence guarantee and a practical, loss-based scheduling method that improves diffusion-model sampling without heavy post-training cost.
Abstract
Diffusion generative models synthesize samples by discretizing reverse-time dynamics driven by a learned score (or denoiser). Existing convergence analyses of diffusion models typically scale at least linearly with the ambient dimension, and sharper rates often depend on intrinsic-dimension assumptions or other geometric restrictions on the target distribution. We develop an alternative, information-theoretic approach to dimension-free convergence that avoids any geometric assumptions. Under mild assumptions on the target distribution, we bound KL divergence between the target and generated distributions by $O(H^2/K)$ (up to endpoint factors), where $H$ is the Shannon entropy and $K$ is the number of sampling steps. Moreover, using a reformulation of the KL divergence, we propose a Loss-Adaptive Schedule (LAS) for efficient discretization of reverse SDE which is lightweight and relies only on the training loss, requiring no post-training heavy computation. Empirically, LAS improves sampling quality over common heuristic schedules.