Preconditioned Score-based Generative Models

TL;DR

This work addresses the bottleneck of slow sampling in score-based generative models by diagnosing ill-conditioning in Langevin dynamics and reverse diffusion. It introduces preconditioned diffusion sampling (PDS), a training-free method that applies a data-driven linear preconditioner M, constructed from frequency and pixel statistics and implemented via FFT, to balance coordinate scales with marginal computational overhead. Theoretical analysis shows that PDS preserves the steady-state distribution of Langevin dynamics and the final-state distribution of reverse diffusion, ensuring unbiased generation while enabling larger step sizes. Empirically, PDS yields up to 28× faster high-resolution image generation (e.g., 1024×1024) with competitive or superior quality compared with state-of-the-art accelerations across CIFAR-10, CelebA, FFHQ, ImageNet, and COCO. The approach offers a scalable, model-agnostic path to practical, fast SGMs without retraining or additional learned components.

Abstract

Score-based generative models (SGMs) have recently emerged as a promising class of generative models. However, a fundamental limitation is that their sampling process is slow due to a need for many (e.g., 2000) iterations of sequential computations. An intuitive acceleration method is to reduce the sampling iterations which however causes severe performance degradation. We assault this problem to the ill-conditioned issues of the Langevin dynamics and reverse diffusion in the sampling process. Under this insight, we propose a novel preconditioned diffusion sampling (PDS) method that leverages matrix preconditioning to alleviate the aforementioned problem. PDS alters the sampling process of a vanilla SGM at marginal extra computation cost and without model retraining. Theoretically, we prove that PDS preserves the output distribution of the SGM, with no risk of inducing systematical bias to the original sampling process. We further theoretically reveal a relation between the parameter of PDS and the sampling iterations, easing the parameter estimation under varying sampling iterations. Extensive experiments on various image datasets with a variety of resolutions and diversity validate that our PDS consistently accelerates off-the-shelf SGMs whilst maintaining the synthesis quality. In particular, PDS can accelerate by up to 28x on more challenging high-resolution (1024x1024) image generation. Compared with the latest generative models (e.g., CLD-SGM and Analytic-DDIM), PDS can achieve the best sampling quality on CIFAR-10 at an FID score of 1.99. Our code is publicly available to foster any further research https://github.com/fudan-zvg/PDS.

Figures (22)

• Figure 1: Facial images at a resolution of $1024\times 1024$ generated by NCSN++ song2020score under a variety of sampling iterations (top) without and (bottom) with our PDS. NCSN++ decades quickly with increasingly reduced sampling iterations, which can be well solved with PDS. For generating a batch of 8 images, PDS reduces the time cost of generating $8$ human facial samples from $1920$ sec (more than half an hours) to $68$ sec (about one minute) on one NVIDIA RTX 3090 GPU. Dataset: FFHQ karras2019style. More samples in Appendix.
• Figure 2: Illustration of our preconditioning method for accelerating sampling process.
• Figure 3: Sampling on a two-dimensional Gaussian distribution using Langevin dynamics with step size $0.1$ (left), Langevin dynamics with step $1$ (middle) and preconditioned Langevin dynamics with step $1$ (right).
• Figure 4: Sampling on LSUN (church and bedroom) ($256 \times 256$). Top: NCSN++ song2020scoreBottom: NCSN++ with PDS. Both use $80$ iterations and $\alpha=2.2$. More samples in Appendix.
• Figure 5: Examples of frequency preconditioning $R_f$ with $(r,\lambda) = (0.1H,1.6)$ in Eq. \ref{['eq: freq_mask_a2']}.

Theorems & Definitions (6)

• theorem 1
• proof
• theorem 2
• proof
• theorem 3
• proof