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.
`LMSSampler`	Creates a linear multi-step (LMS) 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

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:

denoiser (GaussianDenoiser) – A Gaussian denoiser.
kwargs – Keyword arguments passed to Sampler.

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:

denoiser (GaussianDenoiser) – A Gaussian denoiser.
eta (float) – The stochasticity hyperparameter \(\eta \in \mathbb{R}_+\). If \(\eta = 1\), DDIMSampler is equivalent to DDPMSampler. If \(\eta = 0\), DDIMSampler is equivalent to EulerSampler.
kwargs – Keyword arguments passed to Sampler.