azula.sample

Reverse diffusion samplers.

For a distribution \(p(X)\) over \(\mathbb{R}^D\), the perturbation kernel (see azula.noise)

\[p(X_t \mid X) = \mathcal{N}(X_t \mid \alpha_t X_t, \sigma_t^2 I)\]

defines a series of marginal distributions

\[p(X_t) = \int p(X_t \mid x) \, p(x) \operatorname{d}\!x \, .\]

The goal of diffusion models is to generate samples from \(p(X_0)\). To this end, reverse transition kernels \(q(X_s \mid X_t)\) from \(t\) to \(s < t\) are chosen. Then, starting from \(x_1 \sim p(X_1)\), \(T\) transitions \(x_{t_{i-1}} \sim q(X_{t_{i-1}} \mid x_{t_i})\) are simulated from \(t_T = 1\) to \(t_0 = 0\). If the kernels are consistent with the marginals \(p(X_t)\), that is if

\[p(X_{t_{i-1}}) \approx \int q(X_{t_{i-1}} \mid x_{t_i}) \, p(x_{t_i}) \, \operatorname{d}\!x_{t_i} \, ,\]

for all \(i = 1, \dots, T\), the vectors \(x_{t_i}\) are distributed according to \(p(X_{t_i})\), including the last one \(x_{t_0} = x_0\).

Classes

Sampler

Abstract reverse diffusion sampler.

DDPMSampler

Creates an DDPM sampler.

DDIMSampler

Creates a DDIM sampler.

Descriptions

class azula.sample.Sampler(steps)

Abstract reverse diffusion sampler.

Parameters:

steps (int) – The number of discretization steps \(T\). By default, the step size \(t - s\) is constant.

forward(x_1, **kwargs)

Simulates the reverse process from \(t = 1\) to \(0\).

Parameters:

  • x_1 (Tensor) – The noisy vector \(x_1\), with shape \((*, D)\).

  • kwargs – Optional keyword arguments.

Returns:

The clean vector \(x_0\), with shape \((*, D)\).

Return type:

Tensor

abstract step(x_t, t, s, **kwargs)

Simulates the reverse process from \(t\) to \(s \leq t\).

Parameters:

  • x_t (Tensor) – The current vector \(x_t\), with shape \((*, D)\).

  • t (Tensor) – The current time \(t\), with shape \((*)\).

  • s (Tensor) – The target time \(s\), with shape \((*)\).

  • kwargs – Optional keyword arguments.

Returns:

The new vector \(x_s \sim q(X_s \mid x_t)\), with shape \((*, D)\).

Return type:

Tensor

class azula.sample.DDPMSampler(denoiser, **kwargs)

Creates an DDPM sampler.

\[x_s = \alpha_s \mu_\phi(x_t) + \sigma_s \, \sqrt{1 - \tau} \, \frac{x_t - \alpha_t \mu_\phi(x_t)}{\sigma_t} + \sigma_s \, \sqrt{\tau} \, \epsilon\]

where \(\epsilon \sim \mathcal{N}(0, I)\) and

\[\tau = 1 - \frac{\alpha_t^2}{\alpha_s^2} \frac{\sigma_s^2}{\sigma_t^2} \, .\]

References

Denoising Diffusion Probabilistic Models (Ho et al., 2020)
https://arxiv.org/abs/2006.11239
Parameters:
class azula.sample.DDIMSampler(denoiser, eta=None, **kwargs)

Creates a DDIM sampler.

\[x_s = \alpha_s \mu_\phi(x_t) + \sigma_s \, \sqrt{1 - \eta \, \tau} \, \frac{x_t - \alpha_t \mu_\phi(x_t)}{\sigma_t} + \sigma_s \, \sqrt{\eta \, \tau} \, \epsilon\]

where \(\epsilon \sim \mathcal{N}(0, I)\) and

\[\tau = 1 - \frac{\alpha_t^2}{\alpha_s^2} \frac{\sigma_s^2}{\sigma_t^2} \, .\]

References

Denoising Diffusion Implicit Models (Song et al., 2021)
https://arxiv.org/abs/2010.02502
Parameters:

  • denoiser (GaussianDenoiser) – A Gaussian denoiser.

  • eta (float | None) – The stochasticity hyperparameter \(\eta \in \mathbb{R}_+\). If \(\eta = 1\), DDIMSampler is equivalent to DDPMSampler. If eta is None, \(\eta = 0\) and the sampler is deterministic.

  • kwargs – Keyword arguments passed to Sampler.