Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Fast Stochastic Sampling in Diffusion Generative Models

Kushagra Pandey, Maja Rudolph, Stephan Mandt

TL;DR

This work tackles slow inference in diffusion models by introducing Splitting Integrators tailored for augmented diffusion spaces, and shows naive splitting is sub-optimal for fast sampling. By developing Reduced Splitting Integrators that reuse scores, employ timestep embeddings, and regulate noise injection, the method achieves higher sample quality with fewer function evaluations, exemplified by 2.36 FID at 100 NFE on CIFAR-10. The approach yields state-of-the-art or competitive results against leading stochastic samplers across CIFAR-10 and CelebA-64, with practical gains in sampling efficiency. The findings suggest a promising direction for fast, high-fidelity sampling in augmented diffusion models and invite further theoretical and applied exploration, including extensions to position-space-only models.

Abstract

Diffusion models suffer from slow sample generation at inference time. Despite recent efforts, improving the sampling efficiency of stochastic samplers for diffusion models remains a promising direction. We propose Splitting Integrators for fast stochastic sampling in pre-trained diffusion models in augmented spaces. Commonly used in molecular dynamics, splitting-based integrators attempt to improve sampling efficiency by cleverly alternating between numerical updates involving the data, auxiliary, or noise variables. However, we show that a naive application of splitting integrators is sub-optimal for fast sampling. Consequently, we propose several principled modifications to naive splitting samplers for improving sampling efficiency and denote the resulting samplers as Reduced Splitting Integrators. In the context of Phase Space Langevin Diffusion (PSLD) [Pandey \& Mandt, 2023] on CIFAR-10, our stochastic sampler achieves an FID score of 2.36 in only 100 network function evaluations (NFE) as compared to 2.63 for the best baselines.

Towards Fast Stochastic Sampling in Diffusion Generative Models

TL;DR

This work tackles slow inference in diffusion models by introducing Splitting Integrators tailored for augmented diffusion spaces, and shows naive splitting is sub-optimal for fast sampling. By developing Reduced Splitting Integrators that reuse scores, employ timestep embeddings, and regulate noise injection, the method achieves higher sample quality with fewer function evaluations, exemplified by 2.36 FID at 100 NFE on CIFAR-10. The approach yields state-of-the-art or competitive results against leading stochastic samplers across CIFAR-10 and CelebA-64, with practical gains in sampling efficiency. The findings suggest a promising direction for fast, high-fidelity sampling in augmented diffusion models and invite further theoretical and applied exploration, including extensions to position-space-only models.

Abstract

Diffusion models suffer from slow sample generation at inference time. Despite recent efforts, improving the sampling efficiency of stochastic samplers for diffusion models remains a promising direction. We propose Splitting Integrators for fast stochastic sampling in pre-trained diffusion models in augmented spaces. Commonly used in molecular dynamics, splitting-based integrators attempt to improve sampling efficiency by cleverly alternating between numerical updates involving the data, auxiliary, or noise variables. However, we show that a naive application of splitting integrators is sub-optimal for fast sampling. Consequently, we propose several principled modifications to naive splitting samplers for improving sampling efficiency and denote the resulting samplers as Reduced Splitting Integrators. In the context of Phase Space Langevin Diffusion (PSLD) [Pandey \& Mandt, 2023] on CIFAR-10, our stochastic sampler achieves an FID score of 2.36 in only 100 network function evaluations (NFE) as compared to 2.63 for the best baselines.
Paper Structure (20 sections, 32 equations, 2 figures, 3 tables)

This paper contains 20 sections, 32 equations, 2 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Splitting Integrators for Fast Stochastic Sampling
  4. Additional Experimental Results
  5. Conclusion
  6. Splitting Integrators
  7. Introduction to Splitting Integrators
  8. Stochastic Splitting Integrators
  9. Naive Splitting Samplers
  10. Effects of controlling stochasticity
  11. Reduced Splitting Schemes
  12. Local Error Analysis and Justification for design choices
  13. Error Analysis: Naive BAO (NBAO)
  14. Extended Results
  15. Extended Results for Figure \ref{['fig:fig_splitting']}
  16. ...and 5 more sections

Figures (2)

  • Figure 1: Splitting integrators for fast stochastic sampling in diffusion models. a) Sample quality comparisons between naive splitting schemes and the EM Sampler. (b-d) Reduced Splitting schemes outperform their naive counterparts by a large margin. (Left to Right) Comparisons between the naive splitting schemes and their reduced counterparts for the OBA, BAO, and OBAB samplers.
  • Figure 2: Comparison with state-of-the-art methods. a) Reduced OBA outperforms other reduced splitting schemes. (b, c) SPS outperforms competing baseline methods for stochastic sampling for the CIFAR-10 and CelebA-64 datasets, respectively.