add predictor corrector sampling

Author: arrjon

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -16,15 +16,12 @@
     integrate_stochastic,
     logging,
     tensor_utils,
-    filter_kwargs,
 )
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .schedules.noise_schedule import NoiseSchedule
 from .dispatch import find_noise_schedule
 
-ArrayLike = int | float | Tensor
-
 
 # disable module check, use potential module after moving from experimental
 @serializable("bayesflow.networks", disable_module_check=True)
@@ -917,27 +914,6 @@ def score_fn(time, xz):
                 seed=self.seed_generator,
                 **integrate_kwargs,
             )
-        elif integrate_kwargs["method"] == "langevin":
-
-            def scores(time, xz):
-                return {
-                    "xz": self.compositional_score(
-                        xz,
-                        time=time,
-                        conditions=conditions,
-                        compute_prior_score=compute_prior_score,
-                        mini_batch_size=mini_batch_size,
-                        training=training,
-                    )
-                }
-
-            state = annealed_langevin(
-                score_fn=scores,
-                noise_schedule=self.noise_schedule,
-                state=state,
-                seed=self.seed_generator,
-                **filter_kwargs(integrate_kwargs, annealed_langevin),
-            )
         else:
 
             def deltas(time, xz):
@@ -957,46 +933,3 @@ def deltas(time, xz):
 
         x = state["xz"]
         return x
-
-
-def annealed_langevin(
-    score_fn: Callable,
-    noise_schedule: Callable,
-    state: dict[str, ArrayLike],
-    steps: int,
-    seed: keras.random.SeedGenerator,
-    start_time: ArrayLike = None,
-    stop_time: ArrayLike = None,
-    langevin_corrector_steps: int = 5,
-    step_size_factor: float = 0.1,
-) -> dict[str, ArrayLike]:
-    """
-    Annealed Langevin dynamics for diffusion sampling.
-
-    for t = T-1,...,1:
-      for s = 1,...,L:
-        eta ~ N(0, I)
-        theta <- theta + (dt[t]/2) * psi(theta, t) + sqrt(dt[t]) * eta
-    """
-    log_snr_t = noise_schedule.get_log_snr(t=start_time, training=False)
-    _, max_sigma_t = noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)
-
-    # main loops
-    for step in range(steps - 1, 0, -1):
-        t = step / steps
-        log_snr_t = noise_schedule.get_log_snr(t=t, training=False)
-        _, sigma_t = noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)
-        annealing_step_size = step_size_factor * keras.ops.square(sigma_t / max_sigma_t)
-
-        sqrt_dt = keras.ops.sqrt(keras.ops.abs(annealing_step_size))
-        for _ in range(langevin_corrector_steps):
-            drift = score_fn(t, **filter_kwargs(state, score_fn))
-            noise = {
-                k: keras.random.normal(keras.ops.shape(v), dtype=keras.ops.dtype(v), seed=seed)
-                for k, v in state.items()
-            }
-
-            # update
-            for k, d in drift.items():
-                state[k] = state[k] + 0.5 * annealing_step_size * d + sqrt_dt * noise[k]
-    return state

bayesflow/utils/integrate.py

@@ -404,7 +404,7 @@ def integrate_stochastic(
     score_fn: Callable = None,
     corrector_steps: int = 0,
     noise_schedule=None,
-    r: float = 0.1,
+    step_size_factor: float = 0.1,
     **kwargs,
 ) -> Union[dict[str, ArrayLike], tuple[dict[str, ArrayLike], dict[str, Sequence[ArrayLike]]]]:
     """
@@ -427,7 +427,7 @@ def integrate_stochastic(
                  Should take (time, **state) and return score dict.
         corrector_steps: Number of corrector steps to take after each predictor step.
         noise_schedule: Noise schedule object for computing lambda_t and alpha_t in corrector.
-        r: Scaling factor for corrector step size.
+        step_size_factor: Scaling factor for corrector step size.
         **kwargs: Additional arguments to pass to the step function.
 
     Returns:
@@ -489,13 +489,13 @@ def body(_loop_var, _loop_state):
                 # where e = 2*alpha_t * (r * ||z|| / ||score||)**2
                 for k in new_state.keys():
                     if k in score:
-                        z_norm = keras.ops.norm(new_state[k], axis=-1, keepdims=True)
+                        z_norm = keras.ops.norm(_corrector_noise[k], axis=-1, keepdims=True)
                         score_norm = keras.ops.norm(score[k], axis=-1, keepdims=True)
 
                         # Prevent division by zero
                         score_norm = keras.ops.maximum(score_norm, 1e-8)
 
-                        e = 2.0 * alpha_t * (r * z_norm / score_norm) ** 2
+                        e = 2.0 * alpha_t * (step_size_factor * z_norm / score_norm) ** 2
                         sqrt_2e = keras.ops.sqrt(2.0 * e)
 
                         new_state[k] = new_state[k] + e * score[k] + sqrt_2e * _corrector_noise[k]