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, \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.
CosineSchedule	Creates a cosine noise schedule.
RectifiedSchedule	Creates a rectified noise schedule.
DecaySchedule	Creates an exponential decay schedule.

Descriptions

class azula.noise.Schedule

Abstract noise schedule.

abstract __call__(t)

Parameters:

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

Returns:

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

Return type:

tuple[Tensor, Tensor]

class azula.noise.VESchedule(sigma_min=0.001, sigma_max=1000.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 \big( t^2 \log \alpha_\min \big) \\ \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[\).

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

Creates a cosine noise schedule.

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

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[\).

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

Creates a rectified noise schedule.

\[\begin{split}\alpha_t & = t \, \alpha_\min + (1 - t) \\ \sigma_t & = t + (1 - t) \, \sigma_\min\end{split}\]

References

Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow (Liu et al. 2022)

https://arxiv.org/abs/2209.03003

Flow Matching for Generative Modeling (Lipman et al. 2023)

https://arxiv.org/abs/2210.02747

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[\).

class azula.noise.DecaySchedule(alpha_min=0.001, sigma_min=0.001, gamma=0.1)

Creates an exponential decay schedule.

\[\begin{split}\alpha_t & = \tau \, \alpha_\min + (1 - \tau) \\ \sigma_t & = \tau + (1 - \tau) \, \sigma_\min\end{split}\]

where

\[\tau = \frac{1 - \gamma^t}{1 - \gamma}\]

The smaller \(\gamma\), the more time is allocated to high signal-to-noise ratios. When \(\gamma\) is close to 1, the schedule becomes equivalent to the RectifiedSchedule.

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[\).

• gamma (float) – The decay factor \(\gamma \in ]0,1[\).