NoiseSchedule

class bayesflow.experimental.diffusion_model.NoiseSchedule(name: str, variance_type: Literal['preserving', 'exploding'], weighting: Literal['sigmoid', 'likelihood_weighting'] = None)

Bases: ABC

Noise schedule for diffusion models. We follow the notation from [1].

The diffusion process is defined by a noise schedule, which determines how the noise level changes over time. We define the noise schedule as a function of the log signal-to-noise ratio (lambda), which can be interchangeably used with the diffusion time (t).

The noise process is defined as: z = alpha(t) * x + sigma(t) * e, where e ~ N(0, I). The schedule is defined as: lambda(t) = log sigma^2(t) - log alpha^2(t).

We can also define a weighting function for each noise level for the loss function. Often the noise schedule is the same for the forward and reverse process, but this is not necessary and can be changed via the training flag.

[1] Variational Diffusion Models 2.0: Understanding Diffusion Model Objectives as the ELBO with Simple Data Augmentation: Kingma et al. (2023)

Initialize the noise schedule with given variance and weighting strategy.

Parameters:

namestr

The name of the noise schedule.

variance_typeLiteral[“preserving”, “exploding”]

If the variance of noise added to the data should be preserved over time, use “preserving”. If the variance of noise added to the data should increase over time, use “exploding”. Default is “preserving”.

weightingLiteral[“sigmoid”, “likelihood_weighting”], optional

The type of weighting function to use for the noise schedule. Default is None, which means no weighting is applied.

abstractmethod get_log_snr(t: float | Tensor, training: bool) → Tensor

Get the log signal-to-noise ratio (lambda) for a given diffusion time.

abstractmethod get_t_from_log_snr(log_snr_t: float | Tensor, training: bool) → Tensor

Get the diffusion time (t) from the log signal-to-noise ratio (lambda).

abstractmethod derivative_log_snr(log_snr_t: float | Tensor, training: bool) → Tensor

Compute beta(t) = d/dt log(1 + e^(-snr(t))). This is usually used for the reverse SDE.

get_drift_diffusion(log_snr_t: Tensor, x: Tensor = None, training: bool = False) → Tensor | tuple[Tensor, Tensor]

Compute the drift and optionally the squared diffusion term for the reverse SDE. It can be derived from the derivative of the schedule:

math::

beta(t) = d/dt log(1 + e^{-snr(t)})

f(z, t) = -0.5 * beta(t) * z

g(t)^2 = beta(t)

The corresponding differential equations are:

SDE: d(z) = [ f(z, t) - g(t)^2 * score(z, lambda) ] dt + g(t) dW
ODE: dz = [ f(z, t) - 0.5 * g(t)^2 * score(z, lambda) ] dt

For a variance exploding schedule, one should set f(z, t) = 0.

get_alpha_sigma(log_snr_t: Tensor) → tuple[Tensor, Tensor]

Get alpha and sigma for a given log signal-to-noise ratio (lambda).

Default is a variance preserving schedule:

alpha(t) = sqrt(sigmoid(log_snr_t)) sigma(t) = sqrt(sigmoid(-log_snr_t))

For a variance exploding schedule, one should set alpha^2 = 1 and sigma^2 = exp(-lambda)

get_weights_for_snr(log_snr_t: Tensor) → Tensor

Compute loss weights based on log signal-to-noise ratio (log-SNR).

This method returns a tensor of weights used for loss re-weighting in diffusion models, depending on the selected strategy. If no weighting is specified, uniform weights (ones) are returned.

Supported weighting strategies: - “sigmoid”: Based on Kingma et al. (2023), uses a sigmoid of shifted log-SNR. - “likelihood_weighting”: Based on Song et al. (2021), uses ratio of diffusion drift

to squared noise scale.

Parameters:

log_snr_tTensor

A tensor containing the log signal-to-noise ratio values.

Returns:

Tensor

A tensor of weights corresponding to each log-SNR value.

Raises:

TypeError

If the weighting strategy specified in self._weighting is unknown.

get_config()

classmethod from_config(config, custom_objects=None)

validate()

Validate the noise schedule.