Stochastic Sampling Variants

While deterministic samplers like DDIM and the higher-order ODE solvers discussed previously offer significant speed improvements by removing the stochastic element inherent in the original DDPM formulation, there are scenarios where reintroducing noise can be beneficial. Stochastic sampling variants use the randomness of the diffusion process, often aligning more closely with the underlying theory and sometimes offering advantages in sample diversity and quality, albeit typically requiring more sampling steps than highly optimized deterministic methods.

The Role of Stochasticity

Recall that the reverse process in DDPM involves predicting the noise added at a timestep $t$ and then sampling $x_{t-1}$ from a Gaussian distribution centered around a function of $x_t$ and the predicted noise, using a learned or fixed variance $\sigma_t^2$. This explicit sampling step introduces randomness.

$$x_{t-1} \sim \mathcal{N}(\mu_\theta(x_t, t), \sigma_t^2 I)$$

Deterministic samplers effectively set $\sigma_t = 0$ (as in DDIM with $\eta=0$) or approximate the process using an Ordinary Differential Equation (ODE) formulation, which describes the expected path of a sample as noise is removed. Stochastic samplers, in contrast, retain a noise term, aiming to solve the corresponding Stochastic Differential Equation (SDE).

The reverse-time SDE often takes a general form like:

$$dx = f(x, t) dt + g(t) d\bar{w}$$

Here, $dx$ represents the change in the sample $x$ over an infinitesimal time step $dt$. The $f(x, t) dt$ term is the drift, guiding the sample towards lower noise levels (this is what ODE solvers approximate). The $g(t) d\bar{w}$ term is the diffusion or stochastic term, where $g(t)$ scales the noise and $d\bar{w}$ represents infinitesimal random noise (from a Wiener process). Stochastic samplers incorporate approximations of this $g(t) d\bar{w}$ term.

Common Stochastic Sampling Approaches

1. DDPM Ancestral Sampler: The original DDPM sampler is inherently stochastic. It directly implements the Markov chain definition of the reverse process, using the model's predicted mean $\mu_\theta(x_t, t)$ and a variance schedule $\sigma_t^2$ (either fixed, like $\beta_t$, or learned) to draw $x_{t-1}$. While foundational, it typically requires many steps (e.g., 1000) for high-quality generation.

2. SDE Solvers (e.g., Euler-Maruyama): These are numerical methods designed to solve SDEs. The Euler-Maruyama method is a straightforward extension of the Euler method for ODEs. An update step looks roughly like:

   $$x_{t-\Delta t} \approx x_t + f(x_t, t) \Delta t + g(t) \sqrt{\Delta t} z_t$$

   where $z_t \sim \mathcal{N}(0, I)$ is standard Gaussian noise. This explicitly adds scaled noise at each step, guided by the SDE formulation. While simple, it might require small $\Delta t$ (many steps) for stability and accuracy.

3. Predictor-Corrector Methods: Samplers inspired by Score-SDE often employ a two-stage process within each time step:

   • Predictor: Estimate $x_{t-\Delta t}$ using an ODE solver step (approximating the drift).
   • Corrector: Refine the predicted $x_{t-\Delta t}$ using score-based methods like Langevin dynamics, which involves taking steps based on the gradient of the data distribution's log-likelihood (the score, $\nabla_x \log p(x)$), inherently adding noise. This helps ensure the sample stays close to the evolving data manifold.

4. Stochastic Variants of Deterministic Solvers: Some methods add noise injection steps to otherwise deterministic solvers (like DDIM with $\eta > 0$). This provides a knob to control the degree of stochasticity, balancing speed and diversity. For example, SDEdit uses a deterministic ODE solver for most of the reverse process but introduces significant noise at an intermediate timestep $t_{edit}$ before continuing the solve, effectively perturbing the sampling path to increase diversity or perform editing tasks.

Comparing Stochastic and Deterministic Samplers

Choosing between stochastic and deterministic samplers involves trade-offs:

• Speed vs. Steps: Highly optimized deterministic ODE solvers (DPM-Solver++, UniPC) are currently the state-of-the-art for generating high-quality samples in extremely few steps (e.g., 5-20). Stochastic samplers generally need more steps (e.g., 20-200 for good SDE solvers, up to 1000 for basic DDPM) to achieve similar quality because the injected noise requires averaging over time.
• Diversity: Stochastic samplers usually produce more diverse outputs for the same initial latent noise $z_T$. The random noise added at each step creates variations in the sampling path, leading to slightly different final samples. Deterministic samplers produce a unique output for a given $z_T$ and step count.
• Sample Quality: The quality comparison is complex. Sometimes, the noise in stochastic methods can slightly improve perceived quality or realism by better matching the noisy training process, potentially avoiding the overly smooth results sometimes seen with few-step deterministic solvers. However, too much noise or insufficient steps can degrade quality.
• Reproducibility: Deterministic samplers are perfectly reproducible given the same initial noise and parameters. Stochastic samplers require seeding the random number generator at each step to be reproducible.
• Theoretical Fidelity: Stochastic samplers adhere more closely to the probabilistic SDE/DDPM formulation from which diffusion models originated.

Starting from the same initial noise at time T, a deterministic sampler follows a single path to the final sample. Stochastic samplers introduce randomness at intermediate steps, leading to different potential paths and final outcomes.

Practical Considerations

When using stochastic samplers:

• Noise Schedule (g(t)): The magnitude of the injected noise throughout the process is significant. It often follows a schedule related to the original diffusion process noise levels. Incorrect scaling can harm sample quality.
• Number of Steps: Finding the right balance between step count, computational cost, and sample quality is necessary. Unlike few-step ODE solvers, stochastic methods often show more gradual quality improvement with increasing steps.
• Use Cases: Stochastic methods are valuable when sample diversity is a high priority, or when trying to exactly replicate results from models trained and evaluated using stochastic sampling (like original DDPM). They can also sometimes help mitigate certain artifacts seen in extremely low-step deterministic sampling.

In summary, stochastic sampling variants provide an alternative approach to generation that embraces randomness. While often slower than state-of-the-art deterministic solvers, they offer benefits in terms of sample diversity and closer alignment with the underlying SDE theory, adding another set of tools for controlling and optimizing the diffusion model generation process.