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) \, dx \, .\]

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}) \, dx_{t_i} \, ,\]

for all \(i = 1, \dots, T\), the tensors \(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 explicit Euler (1st order) sampler.
HeunSampler	Creates a explicit Heun (2nd order) sampler.
ABSampler	Creates an Adams-Bashforth (AB) multi-step sampler.
EABSampler	Creates an exponential Adams-Bashforth (EAB) multi-step sampler.
PCSampler	Creates a predictor-corrector (PC) sampler.

Descriptions

class azula.sample.Sampler(start=1.0, stop=0.0, steps=64)

Abstract reverse diffusion sampler.

Parameters:

• start (float) – The starting time \(t_T\).

• stop (float) – The stopping time \(t_0\).

• steps (int) – The number of discretization steps \(T\). By default, the step size \(t_{i} - t_{i-1}\) is constant.

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

Draws an initial noisy tensor \(x_{t_T}\).

\[x_{t_T} \sim \mathcal{N}(\alpha_{t_T} \mathbb{E}[X], \alpha_{t_T}^2 \mathbb{V}[X] + \sigma_{t_T}^2 I)\]

Parameters:

• shape (Sequence[int]) – The shape \((*)\) of the tensor.

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

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

• kwargs – Keyword arguments passed to torch.randn.

Returns:

A noisy tensor \(x_1\), with shape \((*)\).

Return type:

Tensor

forward(x, **kwargs)

Simulates the reverse process from \(t_T\) to \(t_0\).

Parameters:

• x (Tensor) – A noisy tensor \(x_{t_T}\), with shape \((*, D)\).

• kwargs – Optional keyword arguments.

Returns:

The clean(er) tensor \(x_{t_0}\), with shape \((*, D)\).

Return type:

Tensor

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

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

Parameters:

• x_t (Tensor) – The current tensor \(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 tensor \(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 \gets \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 \gets \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.

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

Creates a explicit Euler (1st order) sampler.

Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = t\) such that

\[x_s = x_t + \int_t^s z(x_u) \, du\]

where \(z(x_t) = \frac{x_t - \mathbb{E}[X \mid x_t]}{t}\). The explicit Euler step for this integral is

\[x_s \gets x_t + (s - t) \, z(x_t)\]

Wikipedia

https://wikipedia.org/wiki/Euler_method

Parameters:

• denoiser (GaussianDenoiser) – A Gaussian denoiser.

• kwargs – Keyword arguments passed to Sampler.

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

Creates a explicit Heun (2nd order) sampler.

Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = t\) such that

\[x_s = x_t + \int_t^s z(x_u) \, du\]

where \(z(x_t) = \frac{x_t - \mathbb{E}[X \mid x_t]}{t}\). The explicit Heun step for this integral is

\[\begin{split}x_s & \gets x_t + (s - t) \, z(x_t) \\ x_s & \gets x_t + (s - t) \frac{z(x_t) + z(x_s)}{2}\end{split}\]

Wikipedia

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

Parameters:

• denoiser (GaussianDenoiser) – A Gaussian denoiser.

• kwargs – Keyword arguments passed to Sampler.

class azula.sample.ABSampler(denoiser, order=3, **kwargs)

Creates an Adams-Bashforth (AB) multi-step sampler.

Note

This sampler is equivalent to the linear multi-step (LMS) sampler from Katherine Crowson’s k-diffusion.

Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = t\) such that

\[x_s = x_t + \int_t^s z(x_u) \, du\]

where \(z(x_t) = \frac{x_t - \mathbb{E}[X \mid x_t]}{t}\). The \(n\)-th order Adams-Bashforth step for this integral is

\[x_s \gets x_t + \sum_{i=0}^{n-1} z(x_{t_i}) \int_t^s \ell_i(u) \, du\]

where \(t_i\) are previous time steps and the polynomials \(\ell_i(t)\) form their Lagrange basis.

\[\ell_i(t_j) = \sum_{k=0}^{n-1} a_{ik} \, t_j^k = \delta_{ij}\]

Therefore, the Adams-Bashforth coefficients are

\[\begin{split}\int_t^s \ell_i(u) \, du & = \sum_{k=0}^{n-1} a_{ik} \int_t^s u^k \, du \\ & = \sum_{k=0}^{n-1} a_{ik} \left[ \frac{u^{k+1}}{k+1} \right]_t^s\end{split}\]

Wikipedia

https://wikipedia.org/wiki/Linear_multistep_method

Parameters:

• denoiser (GaussianDenoiser) – A Gaussian denoiser.

• order (int) – The order \(n\) of the multi-step method.

• kwargs – Keyword arguments passed to Sampler.

static adams_bashforth(t, n=3)

Returns the coefficients of the \(n\)-th order Adams-Bashforth method.

Parameters:

• t (Tensor) – The integration variable, with shape \((m + 1)\).

• n (int) – The method order \(n \leq m\).

Returns:

The coefficients, with shape \((n)\).

Return type:

Tensor

class azula.sample.EABSampler(denoiser, order=3, **kwargs)

Creates an exponential Adams-Bashforth (EAB) multi-step sampler.

Note

This sampler is a multi-step generalization of the DPM-Solver sampler from Cheng Lu’s dpm-solver.

Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = e^t\) such that

\[x_s = x_t + \int_t^s e^u \, z(x_u) \, du\]

where \(z(x_t) = \frac{x_t - \mathbb{E}[X \mid x_t]}{e^t}\). The \(n\)-th order exponential Adams-Bashforth step for this integral is

\[x_s \gets x_t + \sum_{i=0}^{n-1} z(x_{t_i}) \int_t^s e^u \, \ell_i(u) \, du\]

where \(t_i\) are previous time steps and the polynomials \(\ell_i(t)\) form their Lagrange basis.

\[\ell_i(t_j) = \sum_{k=0}^{n-1} a_{ik} \, t_j^k = \delta_{ij}\]

Therefore, the exponential Adams-Bashforth coefficients are

\[\begin{split}\int_t^s e^u \, \ell_i(u) \, du & = \sum_{k=0}^{n-1} a_{ik} \int_t^s e^u \, u^k \, du \\ & = \sum_{k=0}^{n-1} a_{ik} \left[ (-1)^k \, k! \, e^u \sum_{j=0}^k \frac{(-u)^j}{j!} \right]_t^s\end{split}\]

References

Exponential Adams Bashforth ODE solver for stiff problems (Coudière et al., 2018)

https://arxiv.org/abs/1804.09927

DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps (Lu et al., 2022)

https://arxiv.org/abs/2206.00927

Parameters:

• denoiser (GaussianDenoiser) – A Gaussian denoiser.

• order (int) – The order \(n\) of the multi-step method.

• kwargs – Keyword arguments passed to Sampler.

static exponential_adams_bashforth(t, n=3)

Returns the coefficients of the \(n\)-th order exponential Adams-Bashforth method.

Parameters:

• t (Tensor) – The integration variable, with shape \((m + 1)\).

• n (int) – The method order \(n \leq m\).

Returns:

The coefficients, with shape \((n)\).

Return type:

Tensor

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.