2.2. Samplers based on Ordinary Differential Equations (ODEs)
Another approach to sampling using Diffusion Models (DMs) is to use the Probability Flow Ordinary Differential Equation (ODE), given by: \[ d\boldsymbol{x} = \left[\boldsymbol{f}(\boldsymbol{x}, t) - \frac{1}{2}g(t)^2 \nabla_{\boldsymbol{x}} \log p_t(\boldsymbol{x})\right] dt, \tag{1}\] where \(\boldsymbol{f}(\boldsymbol{x}, t) : \mathbb R^{n \times 1} \to \mathbb R^n\) is the drift coefficient, \(g(t): \mathbb R\to \mathbb R\) is the diffusion coefficient and \(\nabla_{\boldsymbol{x}} \log p_t(\boldsymbol{x}) : \mathbb R^{n \times 1} \to \mathbb R^n\) is the score function. As mentioned in the Stochastic Differential Equations (SDEs) formulation, the Probability Flow ODE is a deterministic process that shares the same marginal distributions as the forward SDE. This allows for sampling from the data distribution by solving an ODE instead of simulating a stochastic process, while also making the process deterministic.
Denoising Diffusion Implicit Models (DDIMs)
After Denoising Diffusion Probabilistic Models (DDPMs), it was observed that the denoising loss function depended only on the marginal distribution \(q(\boldsymbol{x}_t \mid \boldsymbol{x}_0)\), and not on the entire Markov chain. That raised the idea of formulating different inference distributions that matched the mentioned posterior. With that, Denoising Diffusion Implicit Models (DDIMs) (J. Song, Meng, and Ermon 2021) were introduced, with the goal of decreasing the number of inference iterations. Concurrently, it was observed that the DDIM inference scheme could be interpreted as sampling from the DDPM Probability FLow ODE (Y. Song et al. 2021).
Let the DDPM Probability Flow ODE be given by choosing \(\boldsymbol{f}(\boldsymbol{x}, t) = -\frac{1}{2} \beta(t) \boldsymbol{x}\) and \(g(t)^2 = \beta(t)\): \[ d\boldsymbol{x} = \left[-\frac{1}{2} \beta(t) \boldsymbol{x} - \frac{1}{2}\beta(t) \nabla_{\boldsymbol{x}} \log p_t(\boldsymbol{x})\right] dt, \tag{2}\] where \(\beta(t): \mathbb R \to \mathbb R\) is the noise schedule. DDIM introduced a non-Markovian inference distribution \(q_{\boldsymbol{\sigma}}(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0)\) that matched the DDPM posterior \(q(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0)\). Here, \(\boldsymbol{\sigma} \in \mathbb R^T\) represents the noise level. Therefore, the DDIM sampling procedure could be interpreted as solving the Probability Flow ODE, with its inference scheme being given by: \[ \boldsymbol{x}_{t-1} = \sqrt{\alpha_{t-1}}\left(\frac{\boldsymbol{x}_t - \sqrt{1 - \alpha_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)}{\sqrt{\alpha_t}}\right) + \sqrt{1 - \alpha_{t-1} - \sigma_t^2} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) + \sigma_t\boldsymbol{\epsilon}, \tag{3}\] where \(\boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\) is a Gaussian noise vector. When \(\sigma_t = 0\), the DDIM inference scheme becomes deterministic, while non-zero values control the stochasticity of the sampling process. The DDIM sampler is presented in Algorithm 1.
Stochastic Churn Sampler
The Stochastic Churn sampler was introduced in Elucidated Diffusion Models (EDMs) (Karras et al. 2022) as a stochastic extension of the deterministic Heun solver. EDM controlled the stochasticity through a hyperparameter \(\gamma_i \in \mathbb{R}_{\geq 0}\) that governed the amount of noise added to \(\boldsymbol{x}_i\) before the predictor step.
Let \(\boldsymbol{D}_{\boldsymbol{\theta}}: \mathbb R^{n \times 1} \to \mathbb R^n\) be the trained denoiser network, and let \(\{t_i\}_{i=0}^{T}\) denote a decreasing noise schedule, with \(t_0 > t_1 > \cdots > t_T\). The idea of the Stochastic Churn sampler is to first increase the noise level from \(t_i\) to \(\hat{t}_i\) according to: \[ \hat{t}_i = (1+\gamma_i)t_i, \tag{4}\] and then perturb the current sample according to: \[ \boldsymbol{\hat{x}}_i = \boldsymbol{x}_i + \sqrt{\hat{t}_i^2-t_i^2}\boldsymbol{\epsilon}_i, \tag{5}\] where \(\boldsymbol{\epsilon}_i \sim \mathcal{N}(\boldsymbol{0}, \sigma_{\text{Std}}^2\boldsymbol{I})\). Thus, the current sample is moved to a higher noise level before the deterministic denoising step. The sampler then applies Heun’s method to integrate the denoising ODE from \(\hat{t}_i\) to \(t_{i+1}\). In particular, the first evaluation defines the direction: \[ \boldsymbol{d}_i = \frac{\boldsymbol{\hat{x}}_i - \boldsymbol{D}_{\boldsymbol{\theta}}(\boldsymbol{\hat{x}}_i,\hat{t}_i)}{\hat{t}_i}, \tag{6}\] which is used to obtain a first-order prediction at \(t_{i+1}\). A second evaluation of the denoiser is then used to correct this prediction through the average of the two directions.
The stochastic perturbation allows the sampler to explore different denoising trajectories instead of following a single deterministic trajectory. With \(\gamma_i = 0\), \(\hat{t}_i=t_i\) and the stochastic perturbation does not occur, recovering the deterministic Heun solver. Larger values of \(\gamma_i\) increase the temporary noise level and consequently introduce greater stochasticity into the generated trajectory. The full algorithm is presented in (alg:edm_churn?).
```pseudocode #| label: alg:edm_churn #| html-indent-size: “1.2em” #| html-line-number: true