fix mapping min/max snr to t_min/max

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -146,14 +146,15 @@ def __init__(self, min_log_snr: float = -15, max_log_snr: float = 15):
         self._log_snr_min = min_log_snr
         self._log_snr_max = max_log_snr
 
-        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
-        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
     def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
         t_trunc = self._t_min + (self._t_max - self._t_min) * t
         # SNR = -log(exp(t^2) - 1)
-        return -ops.log(ops.exp(ops.square(t_trunc)) - 1)
+        # equivalent, but more stable: -t^2 - log(1 - exp(-t^2))
+        return -ops.square(t_trunc) - ops.log(1 - ops.exp(-ops.square(t_trunc)))
 
     def get_t_from_log_snr(self, log_snr_t: Union[float, Tensor], training: bool) -> Tensor:
         """Get the diffusion time (t) from the log signal-to-noise ratio (lambda)."""
@@ -205,8 +206,8 @@ def __init__(self, min_log_snr: float = -15, max_log_snr: float = 15, s_shift_co
         self._log_snr_max = max_log_snr
         self._s_shift_cosine = s_shift_cosine
 
-        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
-        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
     def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
@@ -266,8 +267,8 @@ def __init__(self, sigma_data: float = 0.5, sigma_min: float = 0.002, sigma_max:
         # convert EDM parameters to signal-to-noise ratio formulation
         self._log_snr_min = -2 * ops.log(sigma_max)
         self._log_snr_max = -2 * ops.log(sigma_min)
-        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
-        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_min = self.get_t_from_log_snr(log_snr_t=self._log_snr_max, training=True)
+        self._t_max = self.get_t_from_log_snr(log_snr_t=self._log_snr_min, training=True)
 
     def get_log_snr(self, t: Tensor, training: bool) -> Tensor:
         """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""