CosineNoiseSchedule

class bayesflow.networks.inference.diffusion.CosineNoiseSchedule(min_log_snr: float = -12, max_log_snr: float = 12, shift: float = 0.0, weighting: Literal['likelihood_weighting'] | None = None)

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

The default weighting=None corresponds to the settings used in the benchmark of [3].

[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. [3] Arruda, J., Bracher, N., Köthe, U., Hasenauer, J., & Radev, S. T. (2026). Diffusion Models in Simulation-Based Inference: A Tutorial Review. https://openreview.net/forum?id=86Q261z005

Parameters:

min_log_snrfloat, optional

The minimum log signal-to-noise ratio (lambda). Default is -12.

max_log_snrfloat, optional

The maximum log signal-to-noise ratio (lambda). Default is 12.

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[“likelihood_weighting”] | None, optional

The type of weighting function to use for the noise schedule. Default is None.

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[“likelihood_weighting”], optional

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

log_snr(t: Tensor | float, training: bool) → Tensor

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

time_from_log_snr(log_snr_t: Tensor | float, training: bool) → Tensor

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

derivative_log_snr(log_snr_t: Tensor, training: bool) → Tensor

Compute d/dt log(1 + e^(-snr(t))), which is used for the reverse SDE.

get_config()

classmethod from_config(config, custom_objects=None)

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

validate(loss_type: str = 'noise')

Validate the noise schedule and warn if its weighting is not defined for loss_type.

weights_for_snr(log_snr_t: Tensor, loss_type: str = 'noise') → 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. The weightings are defined for the noise prediction loss.

Supported weighting strategies: - “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.

loss_typestr

The loss the weights are applied to. Default is “noise”.

Returns:

Tensor

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

Raises:

TypeError

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