The Effect of Stochasticity in Score-Based Diffusion Sampling: a KL Divergence Analysis

TL;DR

This paper analyzes how stochasticity affects score-based diffusion sampling by deriving KL divergence bounds, revealing that stochasticity can both correct errors or amplify them depending on score accuracy and structure.

Contribution

The work provides a theoretical KL divergence analysis of stochasticity in diffusion sampling, including bounds and insights for both exact and approximate score functions.

Findings

Stochasticity can decrease KL divergence with exact scores.

Trade-off exists between error correction and amplification for approximate scores.

Numerical and analytical examples illustrate the theoretical bounds.

Abstract

Sampling in score-based diffusion models can be performed by solving either a reverse-time stochastic differential equation (SDE) parameterized by an arbitrary time-dependent stochasticity parameter or a probability flow ODE, corresponding to the stochasticity parameter set to zero. In this work, we study the effect of this stochasticity on the generation process through bounds on the Kullback-Leibler (KL) divergence, complementing the analysis with numerical and analytical examples. Our main results apply to linear forward SDEs with additive noise and Lipschitz-continuous score functions, and quantify how errors from the prior distribution and score approximation propagate under different choices of the stochasticity parameter. The theoretical bounds are derived using log-Sobolev inequalities for the marginals of the forward process, which enable a more effective control of the KL divergence decay along sampling. For exact score functions, we find that stochasticity acts as an error-correcting mechanism, decreasing KL divergence along the sampling trajectory. For an approximate score function, there is a trade-off between error correction and score error amplification, so that stochasticity can either improve or worsen the performance, depending on the structure of the score error. Numerical experiments on simple datasets and a fully analytical example are included to illustrate and enlighten the theoretical results.