from abc import ABC, abstractmethod
from typing import Literal
from keras import ops
from bayesflow.types import Tensor
from bayesflow.utils import logging
from bayesflow.utils.serialization import deserialize, serializable
[docs]
@serializable("bayesflow.networks")
class NoiseSchedule(ABC):
r"""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)
"""
def __init__(
self,
name: str,
variance_type: Literal["preserving", "exploding"],
weighting: Literal["likelihood_weighting"] = None,
):
"""
Initialize the noise schedule with given variance and weighting strategy.
Parameters
----------
name : str
The name of the noise schedule.
variance_type : Literal["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".
weighting : Literal["likelihood_weighting"], optional
The type of weighting function to use for the noise schedule.
Default is None, which means no weighting is applied.
"""
self.name = name
self._variance_type = variance_type
self.log_snr_min = None # should be set in the subclasses
self.log_snr_max = None # should be set in the subclasses
self._weighting = weighting
[docs]
@abstractmethod
def log_snr(self, t: float | Tensor, training: bool) -> Tensor:
"""Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
pass
[docs]
@abstractmethod
def time_from_log_snr(self, log_snr_t: float | Tensor, training: bool) -> Tensor:
"""Get the diffusion time (t) from the log signal-to-noise ratio (lambda)."""
pass
[docs]
@abstractmethod
def derivative_log_snr(self, log_snr_t: float | Tensor, training: bool) -> Tensor:
r"""Compute \beta(t) = d/dt log(1 + e^(-snr(t))). This is usually used for the reverse SDE."""
pass
[docs]
def get_drift(self, log_snr_t: Tensor, x: Tensor = None, training: bool = False) -> Tensor | tuple[Tensor, Tensor]:
r"""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.
"""
beta = self.derivative_log_snr(log_snr_t=log_snr_t, training=training)
if self._variance_type == "preserving":
g_squared = beta
elif self._variance_type == "exploding":
# [1], Eq. 103, reuses the VP g^2 here, which is wrong (Eq. 105)
g_squared = beta * (1 + ops.exp(-log_snr_t))
else:
raise ValueError(f"Unknown variance type: {self._variance_type}")
if x is None: # return g^2 only
return g_squared
if self._variance_type == "preserving":
f = -0.5 * beta * x
elif self._variance_type == "exploding":
f = ops.zeros_like(beta)
else:
raise ValueError(f"Unknown variance type: {self._variance_type}")
return f, g_squared
[docs]
def get_alpha_sigma(self, 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)
"""
if self._variance_type == "preserving":
# variance preserving schedule
alpha_t = ops.sqrt(ops.sigmoid(log_snr_t))
sigma_t = ops.sqrt(ops.sigmoid(-log_snr_t))
elif self._variance_type == "exploding":
# variance exploding schedule
alpha_t = ops.ones_like(log_snr_t)
sigma_t = ops.sqrt(ops.exp(-log_snr_t))
else:
raise TypeError(f"Unknown variance type: {self._variance_type}")
return alpha_t, sigma_t
[docs]
def weights_for_snr(self, 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_t : Tensor
A tensor containing the log signal-to-noise ratio values.
loss_type : str
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.
"""
if self._weighting is None:
return ops.ones_like(log_snr_t)
elif self._weighting == "likelihood_weighting":
# likelihood weighting based on Song et al. (2021); g^2 / sigma^2 = -dlambda/dt of the training schedule
g_squared = self.get_drift(log_snr_t, training=True)
_, sigma_t = self.get_alpha_sigma(log_snr_t)
return g_squared / ops.square(sigma_t)
else:
raise TypeError(f"Unknown weighting type: {self._weighting}")
[docs]
def get_config(self):
return {"name": self.name, "variance_type": self._variance_type, "weighting": self._weighting}
[docs]
@classmethod
def from_config(cls, config, custom_objects=None):
return cls(**deserialize(config, custom_objects=custom_objects))
[docs]
def validate(self, loss_type: str = "noise"):
"""Validate the noise schedule and warn if its weighting is not defined for `loss_type`."""
if self.log_snr_min >= self.log_snr_max:
raise ValueError("min_log_snr must be less than max_log_snr.")
# Validate log SNR values and corresponding time mappings for both training and inference
for training in (True, False):
if not ops.isfinite(self.log_snr(0.0, training=training)):
raise ValueError(f"log_snr(0.0) must be finite (training={training})")
if not ops.isfinite(self.log_snr(1.0, training=training)):
raise ValueError(f"log_snr(1.0) must be finite (training={training})")
if not ops.isfinite(self.time_from_log_snr(self.log_snr_max, training=training)):
raise ValueError(f"t(log_snr_max) must be finite (training={training})")
if not ops.isfinite(self.time_from_log_snr(self.log_snr_min, training=training)):
raise ValueError(f"t(log_snr_min) must be finite (training={training})")
# Validate log SNR derivatives at the boundaries
for boundary, name in [(self.log_snr_max, "log_snr_max (t=0)"), (self.log_snr_min, "log_snr_min (t=1)")]:
if not ops.isfinite(self.derivative_log_snr(boundary, training=False)):
raise ValueError(f"derivative_log_snr at {name} must be finite.")
if self._weighting is not None and loss_type != "noise":
logging.warning(
f"The '{self._weighting}' weighting is defined for the noise loss, "
f"but loss_type='{loss_type}' is used. You might want to adjust the weighting."
)
