azula.sample¶
Reverse diffusion samplers.
For a distribution \(p(X)\) over \(\mathbb{R}^D\), the perturbation kernel (see
azula.noise)
defines a series of marginal distributions
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
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¶
Abstract reverse diffusion sampler. |
|
Creates an DDPM sampler. |
|
Creates a DDIM sampler. |
|
Creates a explicit Euler (1st order) sampler. |
|
Creates a explicit Heun (2nd order) sampler. |
|
Creates an Itô SDE sampler. |
|
Creates an Adams-Bashforth (AB) multi-step sampler with noise (\(z\)) prediction. |
|
Creates an Adams-Bashforth (AB) multi-step sampler with velocity (\(v\)) prediction. |
|
Creates an exponential Adams-Bashforth (EAB) multi-step sampler with noise (\(z\)) prediction. |
|
Creates an exponential Adams-Bashforth (EAB) multi-step sampler with data (\(x\)) prediction. |
|
Creates a Rosenbrock-type exponential Adams-Bashforth (REAB) multi-step sampler. |
|
Creates a predictor-corrector (PC) sampler. |
Descriptions¶
- class azula.sample.Sampler(start=1.0, stop=0.0, steps=64, silent=False, dtype=None, device=None)¶[source]
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.
silent (bool) – Whether to hide the sampling progress bar or not.
dtype (dtype | None) – The time data type.
device (device | None) – The time device.
- init(shape, mean=0.0, var=1.0, **kwargs)¶[source]
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:
- Returns:
A noisy tensor \(x_{t_T}\), with shape \((*)\).
- Return type:
- class azula.sample.DDPMSampler(denoiser, **kwargs)¶[source]
Creates an DDPM sampler.
\[x_s \gets \alpha_s \mathbb{E}[X \mid x_t] + \sigma_s \, \sqrt{1 - \tau} \, \frac{x_t - \alpha_t \mathbb{E}[X \mid x_t]}{\sigma_t} + \sigma_s \, \sqrt{\tau} \, \varepsilon\]where \(\varepsilon \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)
- class azula.sample.DDIMSampler(denoiser, eta=0.0, **kwargs)¶[source]
Creates a DDIM sampler.
\[x_s \gets \alpha_s \mathbb{E}[X \mid x_t] + \sigma_s \, \sqrt{1 - \eta \, \tau} \, \frac{x_t - \alpha_t \mathbb{E}[X \mid x_t]}{\sigma_t} + \sigma_s \, \sqrt{\eta \, \tau} \, \varepsilon\]where \(\varepsilon \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)- Parameters:
denoiser (Denoiser) – A denoiser \(q_\phi(X \mid X_t)\).
eta (float) – The stochasticity parameter \(\eta \in \mathbb{R}_+\). If \(\eta = 1\),
DDIMSampleris equivalent toDDPMSampler. If \(\eta = 0\),DDIMSampleris equivalent toEulerSampler.kwargs – Keyword arguments passed to
Sampler.
- class azula.sample.EulerSampler(denoiser, **kwargs)¶[source]
Creates a explicit Euler (1st order) sampler.
Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = \mathbb{E}[Z \mid x_t] = z(x_t)\end{split}\]and
\[x_s = x_t + \int_t^s z(x_u) \, du \, .\]The explicit Euler step for this integral is
\[x_s \gets x_t + (s - t) \, z(x_t)\]Wikipedia
- class azula.sample.HeunSampler(denoiser, **kwargs)¶[source]
Creates a explicit Heun (2nd order) sampler.
Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = \mathbb{E}[Z \mid x_t] = z(x_t)\end{split}\]and
\[x_s = x_t + \int_t^s z(x_u) \, du \, .\]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
- class azula.sample.ItoSampler(denoiser, eta=1.0, temperature=1.0, **kwargs)¶[source]
Creates an Itô SDE sampler.
Let’s consider the Itô SDE
\[dx_t = \left[ f_t \, x_t - \frac{1 + \eta^2}{2 \tau} g_t^2 \, \nabla_{x_t} \log p(x_t) \right] dt + \eta \, g_t \, dw_t\]where \(\eta \geq 0\) controls stochasticity, \(\tau \geq 0\) controls temperature, and
\[\begin{split}f_t & = \partial_t \log \alpha_t \\ g_t^2 & = \alpha_t^2 \, \partial_t \frac{\sigma_t^2}{\alpha_t^2} \, .\end{split}\]The integral form of this semi-linear Itô SDE is
\[x_s = \Psi(t, s) \, x_t + \frac{1 + \eta^2}{2} \int_t^s \Psi(u, s) \, g_u^2 \, \nabla_{x_u} \log p(x_u) \, du + \eta \int_t^s \Psi(u, s) \, g_u \, dw_u\]where
\[\Psi(t, s) = \exp \left( \int_t^s f_u \, du \right) = \frac{\alpha_s}{\alpha_t} \, .\]By Tweedie’s formula, we have
\[\begin{split}\int_t^s \Psi(u, s) \, g_u^2 \, \nabla_{x_u} \log p(x_u) \, du & = \int_t^s \Psi(u, s) \, g_u^2 \frac{\alpha_u \mathbb{E}[x_0 | x_u] - x_u}{\sigma_u^2} du \\ & = 2 \alpha_s \int_t^s \partial_u \frac{\sigma_u}{\alpha_u} \frac{\alpha_u \mathbb{E}[x_0 | x_u] - x_u}{\sigma_u} du \\ & \approx 2 \alpha_s \frac{\alpha_t \mathbb{E}[x_0 | x_t] - x_t}{\sigma_t} \int_t^s \partial_u \frac{\sigma_u}{\alpha_u} du \\ & \approx 2 \left( \frac{\sigma_s}{\sigma_t} - \frac{\alpha_s}{\alpha_t} \right) (\alpha_t \mathbb{E}[x_0 | x_t] - x_t) \, .\end{split}\]Finally, by Itô isometry, we have
\[\begin{split}\int_t^s \Psi(u, s) \, g_u \, dw_u & = \int_s^t \Psi(u, s) \, g_u \, dw_u \\ & = \varepsilon \, \sqrt{\int_s^t \Psi(u, s)^2 \, g_u^2 \, du} \\ & = \varepsilon \, \sqrt{\int_s^t \alpha_s^2 \, \partial_u \frac{\sigma_u^2}{\alpha_u^2} du} \\ & = \varepsilon \, \alpha_s \sqrt{\frac{\sigma_t^2}{\alpha_t^2} - \frac{\sigma_s^2}{\alpha_s^2}}\end{split}\]where \(\varepsilon \sim \mathcal{N}(0, 1)\).
- class azula.sample.zABSampler(denoiser, order=2, **kwargs)¶[source]
Creates an Adams-Bashforth (AB) multi-step sampler with noise (\(z\)) prediction.
Note
This sampler is equivalent to the \(\rho\mathrm{AB}\) sampler from Zhang et al. (2023) and 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
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = \mathbb{E}[Z \mid x_t] = z(x_t)\end{split}\]and
\[x_s = x_t + \int_t^s z(x_u) \, du \, .\]The \(n\)-th order Adams-Bashforth step for this integral is
\[x_s \gets x_t + \sum_{i=1}^{n} z(x_{t_i}) \int_t^s \ell_i(u) \, du\]where \(t_i\) are previous time steps and the polynomials \(\ell_i(u) = \sum_{k=1}^{n} a_{ik} \, u^k\) form their Lagrange basis. The coefficients \(a_{ik}\) are determined by solving the system of linear equations \(\ell_i(t_j) = \delta_{ij}\). Then, the Adams-Bashforth coefficients are
\[\begin{split}\int_t^s \ell_i(u) \, du & = \sum_{k=1}^{n} a_{ik} \int_t^s u^k \, du \\ & = \sum_{k=1}^{n} a_{ik} \left[ \frac{u^{k+1}}{k+1} \right]_t^s\end{split}\]Wikipedia
https://wikipedia.org/wiki/Linear_multistep_method
References
Fast Sampling of Diffusion Models with Exponential Integrator (Zhang et al., 2023)
- class azula.sample.vABSampler(denoiser, order=2, **kwargs)¶[source]
Creates an Adams-Bashforth (AB) multi-step sampler with velocity (\(v\)) prediction.
Without loss of generality, let’s assume \(\alpha_t = 1 - t\) and \(\sigma_t = t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = -\mathbb{E}[X \mid x_t] + \mathbb{E}[Z \mid x_t] = v(x_t)\end{split}\]and
\[x_s = x_t + \int_t^s v(x_u) \, du \, .\]The \(n\)-th order Adams-Bashforth step for this integral is
\[x_s \gets x_t + \sum_{i=1}^{n} v(x_{t_i}) \int_t^s \ell_i(u) \, du\]where \(t_i\) are previous time steps and the polynomials \(\ell_i(u)\) form their Lagrange basis.
See also
- class azula.sample.zEABSampler(denoiser, order=2, **kwargs)¶[source]
Creates an exponential Adams-Bashforth (EAB) multi-step sampler with noise (\(z\)) prediction.
Note
This sampler is a multi-step generalization of the DPM-Solver sampler from Lu et al. (2022).
Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = e^t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = e^t \, \mathbb{E}[Z \mid x_t] = e^t \, z(x_t)\end{split}\]and
\[x_s = x_t + \int_t^s e^u \, z(x_u) \, du \, .\]The \(n\)-th order exponential Adams-Bashforth step for this integral is
\[x_s \gets x_t + \sum_{i=1}^{n} 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(u) = \sum_{k=1}^{n} a_{ik} \, u^k\) form their Lagrange basis. The coefficients \(a_{ik}\) are determined by solving the system of linear equations \(\ell_i(t_j) = \delta_{ij}\). Then, the exponential Adams-Bashforth coefficients are
\[\begin{split}\int_t^s e^u \, \ell_i(u) \, du & = \sum_{k=1}^{n} a_{ik} \int_t^s e^u \, u^k \, du \\ & = \sum_{k=1}^{n} a_{ik} \left[ (-1)^k \, k! \, e^u \sum_{j=0}^k \frac{(-u)^j}{j!} \right]_t^s\end{split}\]References
DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps (Lu et al., 2022)Exponential Adams Bashforth ODE solver for stiff problems (Coudière et al., 2018)
- class azula.sample.xEABSampler(denoiser, order=2, **kwargs)¶[source]
Creates an exponential Adams-Bashforth (EAB) multi-step sampler with data (\(x\)) prediction.
Note
This sampler is a multi-step generalization of the DPM-Solver++ sampler from Lu et al. (2022).
Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = e^t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = e^t \, \mathbb{E}[Z \mid x_t] = x_t - \mathbb{E}[X \mid x_t] = x_t - x(x_t)\end{split}\]and
\[x_s = e^{s-t} \, x_t - \int_t^s e^{s-u} \, x(x_u) \, du \, .\]The \(n\)-th order exponential Adams-Bashforth step for this integral is
\[x_s \gets e^{s-t} \, x_t - e^s \sum_{i=1}^{n} x(x_{t_i}) \int_t^s e^{-u} \, \ell_i(u) \, du\]where \(t_i\) are previous time steps and the polynomials \(\ell_i(u) = \sum_{k=1}^{n} a_{ik} \, u^k\) form their Lagrange basis. The coefficients \(a_{ik}\) are determined by solving the system of linear equations \(\ell_i(t_j) = \delta_{ij}\). Then, the exponential Adams-Bashforth coefficients are
\[\begin{split}\int_t^s e^{-u} \, \ell_i(u) \, du & = \sum_{k=1}^{n} a_{ik} \int_t^s e^{-u} \, u^k \, du \\ & = \sum_{k=1}^{n} a_{ik} \left[ -k! \, e^{-u} \sum_{j=0}^k \frac{u^j}{j!} \right]_t^s\end{split}\]References
DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models (Lu et al., 2022)
- class azula.sample.REABSampler(denoiser, order=2, **kwargs)¶[source]
Creates a Rosenbrock-type exponential Adams-Bashforth (REAB) multi-step sampler.
Note
This sampler is a multi-step generalization of the DPM-Solver-v3 sampler from Zheng et al. (2023).
Without loss of generality, let’s assume \(\alpha_t = 1\) and \(\sigma_t = e^t\) such that
\[\begin{split}\frac{d x_t}{dt} & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\ & = e^t \, \mathbb{E}[Z \mid x_t] = x_t - \mathbb{E}[X \mid x_t] \\ & = a_t \, x_t + b_t \frac{(1 - a_t) \, x_t - \mathbb{E}[X \mid x_t]}{b_t} = a_t \, x_t + b_t \, f(x_t)\end{split}\]and
\[x_s = \Psi(t, s) \, x_t + \int_t^s \Psi(u, s) \, b_u \, f(x_u) \, du\]where \(a_t = \frac{e^{2t}}{1 + e^{2t}}\), \(b_t = \sqrt{a_t}\) and
\[\Psi(t, s) = \exp \left( \int_t^s a_u \, du \right) = \sqrt{\frac{1 + e^{2s}}{1 + e^{2t}}} \, .\]The \(n\)-th order exponential Adams-Bashforth step for this integral is
\[x_s \gets \sqrt{\frac{1 + e^{2s}}{1 + e^{2t}}} \, x_t + \sqrt{1 + e^{2s}} \sum_{i=1}^{n} f(x_{t_i}) \int_t^s \frac{e^u}{1 + e^{2u}} \ell_i(u) \, du\]where \(t_i\) are previous time steps and the polynomials \(\ell_i(u) = \sum_{k=1}^{n} a_{ik} \, u^k\) form their Lagrange basis. The coefficients \(a_{ik}\) are determined by solving the system of linear equations \(\ell_i(t_j) = \delta_{ij}\). Then, the exponential Adams-Bashforth coefficients are
\[\int_t^s \frac{e^u}{1 + e^{2u}} \ell_i(u) \, du = \sum_{k=1}^{n} a_{ik} \int_t^s \frac{e^u \, u^k}{1 + e^{2u}} du\]where the last integral is estimated by numerical integration.
References
DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics (Zheng et al., 2023)Exponential Rosenbrock-Type Methods (Hochbruck et al., 2009)