An Elementary Approach to Scheduling in Generative Diffusion Models

TL;DR

This work develops a principled, Gaussian-based framework to understand how noise scheduling and time discretization impact diffusion-model sampling. By deriving a closed-form reverse-time trajectory and the KL divergence between the source and reverse-sampled outputs, it connects discretization error to schedule design via Euler–Maclaurin expansions and introduces a modewise tangent-law schedule whose coefficient depends on the source covariance eigenvalues. A globally parameterized tangent law with optimal $\gamma^*$ provides a practical, analytically grounded recipe for noise scheduling, while the KL divergence enables low-cost evaluation of time-discretization strategies for pretrained models. Experiments on CIFAR-10 and FFHQ-64 show the proposed discretization strategy consistently outperforms baselines and competitive search-based methods, especially when the NFEs budget is tight, highlighting the method's potential for efficient, high-quality generative sampling.

Abstract

An elementary approach to characterizing the impact of noise scheduling and time discretization in generative diffusion models is developed. Considering a simplified model where the source distribution is multivariate Gaussian with a given covariance matrix, the explicit closed-form evolution trajectory of the distributions across reverse sampling steps is derived, and consequently, the Kullback-Leibler (KL) divergence between the source distribution and the reverse sampling output is obtained. The effect of the number of time discretization steps on the convergence of this KL divergence is studied via the Euler-Maclaurin expansion. An optimization problem is formulated, and its solution noise schedule is obtained via calculus of variations, shown to follow a tangent law whose coefficient is determined by the eigenvalues of the source covariance matrix. For an alternative scenario, more realistic in practice, where pretrained models have been obtained for some given noise schedules, the KL divergence also provides a measure to compare different time discretization strategies in reverse sampling. Experiments across different datasets and pretrained models demonstrate that the time discretization strategy selected by our approach consistently outperforms baseline and search-based strategies, particularly when the budget on the number of function evaluations is very tight.

Figures (12)

• Figure 1: KL divergence versus steps $N$ for the time discretization strategy \ref{['eq:lam_power_uni']} under (a) VP and (b) VE settings.
• Figure 2: Performance analysis of the parameterized tangent law schedule in Gaussian settings. We report the KL divergence versus the number of time discretization steps $N$. Top Row: Comparison of our parameterized tangent law schedule (using $\gamma^{\star}$) against existing noise schedules ho2020denoisingsong2020scorelipman2023flownichol2021improved. Bottom Row: Ablation study of different $\gamma$ values, verifying the optimality of our derived $\gamma^{\star}$. The left and right columns correspond to low-dimensional ($k=128$) and high-dimensional ($k=1280$) cases, respectively. The results demonstrate the effectiveness of our derived $\gamma^{\star}$ in \ref{['thm:optimal_gamma_thm']}.
• Figure 3: Visualization of the $\eta_{i} = e^{-\lambda_{i}}$ under different time discretization strategies.
• Figure 4: EDM checkpoint with DPM-Solver++ on CIFAR-10 (NFEs=6).
• Figure 5: EDM checkpoint with DPM-Solver++ on CIFAR-10 (NFEs=12).

Theorems & Definitions (13)

• Lemma 1
• proof
• Lemma 2
• proof
• Remark : Convergence and Consistency
• Definition 1
• Theorem 1
• proof
• Theorem 2
• proof