CosineNoiseSchedule#

class bayesflow.experimental.diffusion_model.CosineNoiseSchedule(min_log_snr: float = -15, max_log_snr: float = 15, shift: float = 0.0, weighting: Literal['sigmoid', 'likelihood_weighting'] = 'sigmoid')[source]#

Bases: NoiseSchedule

Cosine noise schedule for diffusion models. This schedule is based on the cosine schedule from [1].

A cosine schedule is a popular technique for controlling how the variance (noise level) or learning rate evolves during the training of diffusion models. It was proposed as an improvement over the original linear beta schedule in [2]

[1] Dhariwal, P., & Nichol, A. (2021). Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34, 8780-8794. [2] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840-6851.

Initialize the cosine noise schedule.

Parameters:

min_log_snrfloat, optional: The minimum log signal-to-noise ratio (lambda). Default is -15.
max_log_snrfloat, optional: The maximum log signal-to-noise ratio (lambda). Default is 15.
shiftfloat, optional: Shift the log signal-to-noise ratio (lambda) by this amount. Default is 0.0. For images, use shift = log(base_resolution / d), where d is the used resolution of the image.
weightingLiteral[“sigmoid”, “likelihood_weighting”], optional: The type of weighting function to use for the noise schedule. Default is “sigmoid”.

get_log_snr(t: Tensor | float, training: bool) → Tensor[source]#: Get the log signal-to-noise ratio (lambda) for a given diffusion time.

get_t_from_log_snr(log_snr_t: Tensor | float, training: bool) → Tensor[source]#: Get the diffusion time (t) from the log signal-to-noise ratio (lambda).

derivative_log_snr(log_snr_t: Tensor, training: bool) → Tensor[source]#: Compute d/dt log(1 + e^(-snr(t))), which is used for the reverse SDE.

get_config()[source]#

classmethod from_config(config, custom_objects=None)[source]#

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_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_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.

validate()#: Validate the noise schedule.