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.

EulerSampler

Creates a deterministic Euler (1st order) sampler.

HeunSampler

Creates a deterministic Heun (2nd order) sampler.

PCSampler

Creates a predictor-corrector (PC) 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.

init(shape, mean=None, var=None, **kwargs)

Draws an initial noisy vector \(x_1\).

\[x_1 \sim \mathcal{N}(\alpha_1 \mathbb{E}[X], \alpha_1^2 \mathbb{V}[X] + \sigma_1^2 I)\]
Parameters:

  • shape (Sequence[int]) – The shape \((*, D)\) of the vector.

  • mean (Tensor | None) – The mean \(\mathbb{E}[X]\) of \(p(X)\), with shape \(()\) or \((*, D)\). If None, use 0 instead.

  • var (Tensor | None) – The variance \(\mathbb{V}[X]\) of \(p(X)\), with shape \(()\) or \((*, D)\). If None, use 1 instead.

  • kwargs – Keyword arguments passed to torch.randn.

Returns:

A noisy vector \(x_1\), with shape \((*, D)\).

Return type:

Tensor

forward(x_1, **kwargs)

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

Parameters:

  • x_1 (Tensor) – A 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=0.0, **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:
class azula.sample.EulerSampler(denoiser, **kwargs)

Creates a deterministic Euler (1st order) sampler.

The integration is carried with respect to the noise-to-signal ratio \(\frac{\sigma_t}{\alpha_t}\) rather than the time \(t\).

Wikipedia

https://wikipedia.org/wiki/Euler_method

Parameters:
class azula.sample.HeunSampler(denoiser, **kwargs)

Creates a deterministic Heun (2nd order) sampler.

The integration is carried with respect to the noise-to-signal ratio \(\frac{\sigma_t}{\alpha_t}\) rather than the time \(t\).

Wikipedia

https://wikipedia.org/wiki/Heun%27s_method

Parameters:
class azula.sample.PCSampler(denoiser, corrections=1, delta=0.01, **kwargs)

Creates a predictor-corrector (PC) sampler.

References

Score-Based Generative Modeling through Stochastic Differential Equations (Song et al., 2021)
https://arxiv.org/abs/2011.13456
Parameters:

  • denoiser (GaussianDenoiser) – A Gaussian denoiser.

  • corrections (int) – The number of Langevin corrections per step.

  • delta (float) – The amplitude of Langevin corrections.

  • kwargs – Keyword arguments passed to Sampler.