sun_diffusion.diffusion module

Forward processes and noise schedules for diffusion.

class sun_diffusion.diffusion.VarianceExpandingDiffusion(kappa)

Bases: DiffusionProcess

Variance-expading diffusion process.

Noise schedule is given by \(g(t) = \kappa^t\), which yields a diffusivity of

\[\sigma(t) = \sqrt{\frac{\kappa^{2t} - 1}{2 \log\kappa}}.\]

Parameters:

kappa (float) – Noise scale

noise_coeff(t)

Returns the noise coefficient \(g(t)\) at time t.

Return type:

Tensor

sigma_func(t)

Returns the width of the heat kernel at time t.

Return type:

Tensor

diffuse(x_0, t)

Diffuses input x_0 forward to time t, where

\[x_0 \rightarrow x_t = x_0 + \sigma_t \eta\]

and \(\eta \sim \mathcal{N}(0, \mathbb{I})\).

Parameters:

• x_0 (Tensor) – Input data

• t (Tensor) – Time step to which to diffuse

Return type:

Tensor

class sun_diffusion.diffusion.VarianceExpandingDiffusionSUN(kappa)

Bases: DiffusionSUN

Variance-expanding diffusion on the \({\rm SU}(N)\) group manifold.

Parameters:

kappa (float) – Noise scale

noise_coeff(t)

Returns the noise coefficient \(g(t)\) at time t.

Return type:

Tensor

sigma_func(t)

Returns the width of the heat kernel at time t.

Return type:

Tensor

class sun_diffusion.diffusion.PowerDiffusionSUN(kappa, alpha)

Bases: DiffusionSUN

Power-law diffusion on the \({\rm SU}(N)\) group manifold.

The noise schedule is defined as \(g(t) = \kappa t^\alpha\), which results in a diffusivity of

\[\sigma(t) = \kappa \sqrt{\frac{t^{2\alpha + 1}}{2\alpha + 1}}.\]

Parameters:

• kappa (float) – Noise scale

• alpha (float) – Time exponent

noise_coeff(t)

Returns the noise coefficient \(g(t)\) at time t.

Return type:

Tensor

sigma_func(t)

Returns the width of the heat kernel at time t.

Return type:

Tensor