Diffusion Models

Diffusion Models#

Diffusion Model Types#

class diffuse.diffusion.sde.SDE(beta: Schedule)[source]#

Bases: DiffusionModel

Variance Preserving (VP) SDE for diffusion models.

Implements the forward SDE:

\[dX(t) = -\frac{1}{2}\beta(t) X(t) dt + \sqrt{\beta(t)} dW(t)\]

where \(\beta(t)\) is the noise schedule and \(dW(t)\) is the Wiener process.

This formulation preserves the variance of the data distribution and is the standard choice for diffusion probabilistic models.

Parameters:

beta – Noise schedule (LinearSchedule or CosineSchedule)

noise_level(t: Array) → Array[source]#

Compute noise level \(\sigma(t)\) for diffusion process.

The solution to the VP-SDE is:

\[X(t) = \alpha(t) X_0 + \sigma(t) \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)\]

where \(\alpha(t) = \exp\left(-\frac{1}{2}\int_0^t \beta(s) ds\right)\) and \(\sigma(t) = \sqrt{1 - \alpha^2(t)}\)

Returns:

Noise level \(\sigma(t)\) clipped for numerical stability

sde_coefficients(t: Array) → tuple[Array, Array][source]#

Compute SDE coefficients \(f(t)\) and \(g(t)\).

For the VP-SDE: \(dX(t) = -\frac{1}{2}\beta(t) X(t) dt + \sqrt{\beta(t)} dW(t)\)

Returns:

Tuple of drift coefficient \(f(t) = -\frac{1}{2}\beta(t)\) and diffusion coefficient \(g(t) = \sqrt{\beta(t)}\)

signal_level(t: Array) → Array[source]#

Compute signal level \(\alpha(t) = \exp\left(-\frac{1}{2}\int_0^t \beta(s) ds\right)\).

Returns:

Signal level clipped for numerical stability

class diffuse.diffusion.sde.Flow(tf: float = 1.0)[source]#

Bases: DiffusionModel

Rectified Flow diffusion model with straight-line interpolation paths.

Implements the rectified flow formulation from Liu et al. (2022) with linear schedules:

\[\alpha(t) = 1 - t, \quad \sigma(t) = t\]

This creates straight-line interpolation paths:

\[x_t = (1-t)x_0 + t\varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)\]

which are more amenable to ODE-based sampling with fewer discretization steps.

Parameters:

tf – Final time for the diffusion process (default: 1.0)

References

Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv:2209.03003

noise_level(t: Array) → Array[source]#

Compute noise level \(\sigma(t) = t\).

Returns:

Noise level clipped for numerical stability

sde_coefficients(t: Array) → tuple[Array, Array][source]#

Compute SDE coefficients for rectified flow.

Returns drift \(f(t) = -\frac{1}{1-t}\) and diffusion \(g(t) = \sqrt{\frac{2t}{1-t}}\)

Returns:

Tuple of drift and diffusion coefficients

signal_level(t: Array) → Array[source]#

Compute signal level \(\alpha(t) = 1 - t\).

Returns:

Signal level clipped for numerical stability

class diffuse.diffusion.sde.EDM(tf: float = 1.0)[source]#

Bases: DiffusionModel

Efficient Diffusion Model (EDM) from Karras et al. (2022).

Implements the EDM formulation with constant signal and increasing noise:

\[\alpha(t) = 1, \quad \sigma(t) = t\]

This creates the simple interpolation:

\[x_t = x_0 + t\varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)\]

which simplifies the probability-flow ODE and is particularly effective with Heun’s integration method.

Parameters:

tf – Final time for the diffusion process (default: 1.0)

References

Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the design space of diffusion-based generative models. NeurIPS 35, 26565-26577.

noise_level(t: Array) → Array[source]#

Compute noise level \(\sigma(t) = t\).

Returns:

Noise level clipped for numerical stability

sde_coefficients(t: Array) → tuple[Array, Array][source]#

Compute SDE coefficients for EDM.

Returns drift \(f(t) = 0\) and diffusion \(g(t) = 1\)

Returns:

Tuple of zero drift and unit diffusion coefficients

signal_level(t: Array) → Array[source]#

Compute signal level \(\alpha(t) = 1\).

Returns:

Constant signal level of 1

Noise Schedules#

class diffuse.diffusion.sde.LinearSchedule(b_min: float, b_max: float, t0: float, T: float)[source]#

Bases: object

Linear noise schedule for diffusion processes.

Implements a linear interpolation between minimum and maximum noise levels:

\[\beta(t) = \beta_{\min} + \frac{\beta_{\max} - \beta_{\min}}{T - t_0}(t - t_0)\]
Parameters:

  • b_min – The minimum noise value \(\beta_{\min}\)

  • b_max – The maximum noise value \(\beta_{\max}\)

  • t0 – The starting time \(t_0\)

  • T – The ending time \(T\)

integrate(t: Array, s: Array) → Array[source]#

Compute integral \(\int_s^t \beta(\tau) d\tau\).

Parameters:

  • t – Upper integration bound

  • s – Lower integration bound

Returns:

Integral value

class diffuse.diffusion.sde.CosineSchedule(b_min: float, b_max: float, t0: float, T: float, s: float = 0.008)[source]#

Bases: Schedule

Cosine noise schedule for improved denoising.

Implements the cosine schedule from Nichol & Dhariwal (2021) which provides better signal-to-noise ratio properties than linear schedules. The schedule is based on:

\[\bar{\alpha}(t) = \frac{\cos\left(\frac{t/T + s}{1+s} \cdot \frac{\pi}{2}\right)^2}{\cos\left(\frac{s}{1+s} \cdot \frac{\pi}{2}\right)^2}\]
Parameters:

  • b_min – The minimum beta value \(\beta_{\min}\)

  • b_max – The maximum beta value \(\beta_{\max}\)

  • t0 – The starting time \(t_0\)

  • T – The ending time \(T\)

  • s – Offset parameter for numerical stability (default: 0.008)

References

Nichol, A., & Dhariwal, P. (2021). Improved Denoising Diffusion Probabilistic Models. arXiv:2102.09672

integrate(t: Array, s: Array) → Array[source]#

Compute integral \(\int_s^t \beta(\tau) d\tau\) using \(\bar{\alpha}\) values.

Returns \(\log(\bar{\alpha}(s) / \bar{\alpha}(t))\)

Parameters:

  • t – Upper integration bound

  • s – Lower integration bound

Returns:

Integral value