azula.noise

Noise schedules.

A noise schedule is a mapping from a time \(t \in [0, 1]\) to the signal scale \(\alpha_t \in \mathbb{R}_+\) and the noise scale \(\sigma_t \in \mathbb{R}_+\) in a perturbation kernel

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

from a “clean” random variable \(X\) to a “noisy” random variable \(X_t\). The only constraint is for the signal-to-noise (SNR) ratio \(\frac{\alpha_t}{\sigma_t}\) to be monotonically decreasing with respect to the time \(t\). Typically, the initial signal scale \(\alpha_0\) is set to 1 and the initial noise is small enough (\(0 < \sigma_0 \ll 1\)) that \(X_0\) is almost equivalent to \(X\).

Note that the relation between \(X_s\) and \(X_t\) (\(0 \leq s \leq t\)) is not enforced by the noise schedule. For example,

\[\begin{split}Z & \sim \mathcal{N}(0, I) \\ X_s & = \alpha_s X + \sigma_s Z \\ X_t & = \alpha_t X + \sigma_t Z\end{split}\]

and

\[\begin{split}Z_1, Z_2 & \sim \mathcal{N}(0, I) \\ X_s & = \alpha_s X + \sigma_s Z_1 \\ X_t & = \frac{\alpha_t}{\alpha_s} X_s + \sqrt{\sigma_t^2 - \frac{\alpha_t^2}{\alpha_s^2} \sigma_s^2} \, Z_2\end{split}\]

are both compatible with the perturbation kernel \(p(X_t \mid X)\).

Classes

Schedule	Abstract noise schedule.
VESchedule	Creates a variance exploding (VE) noise schedule.
VPSchedule	Creates a variance preserving (VP) noise schedule.

Descriptions

class azula.noise.Schedule(*args, **kwargs)

Abstract noise schedule.

abstract forward(t)

Parameters:

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

Returns:

The signal and noise scales \(\alpha_t\) and \(\sigma_t\), with shape \((*, 1)\).

Return type:

Tuple[Tensor, Tensor]

class azula.noise.VESchedule(sigma_min=0.001, sigma_max=100.0)

Creates a variance exploding (VE) noise schedule.

\[\begin{split}\alpha_t & = 1 \\ \sigma_t & = \exp \big( (1 - t) \log \sigma_\min + t \log \sigma_\max \big)\end{split}\]

References

Generative Modeling by Estimating Gradients of the Data Distribution (Song et al., 2019)

https://arxiv.org/pdf/1907.05600

Score-Based Generative Modeling through Stochastic Differential Equations (Song et al., 2021)

https://arxiv.org/abs/2011.13456

Parameters:

• sigma_min (float) – The initial noise scale \(\sigma_\min \in \mathbb{R}_+\).

• sigma_max (float) – The final noise scale \(\sigma_\max \in \mathbb{R}_+\).

class azula.noise.VPSchedule(alpha_min=0.001, sigma_min=0.001)

Creates a variance preserving (VP) noise schedule.

\[\begin{split}\alpha_t & = \exp(t^2 \log \alpha_\min) \\ \sigma_t & = \sqrt{ 1 - \alpha_t^2 + \sigma_\min^2 }\end{split}\]

References

Denoising Diffusion Probabilistic Models (Ho et al. 2020)

https://arxiv.org/abs/2006.11239

Score-Based Generative Modeling through Stochastic Differential Equations (Song et al., 2021)

https://arxiv.org/abs/2011.13456

Parameters:

• alpha_min (float) – The final signal scale \(\alpha_\min \in [0, 1]\).

• sigma_min (float) – The initial noise scale \(\sigma_\min \in [0, 1]\).