sbi-dev
diff --git a/‎docs/advanced_tutorials/19_flowmatching_and_scorematching.ipynb‎
Lines changed: 285 additions & 15 deletions b/‎docs/advanced_tutorials/19_flowmatching_and_scorematching.ipynb‎
Lines changed: 285 additions & 15 deletions
diff --git a/‎docs/advanced_tutorials/20_score_based_methods_new_features.ipynb‎
Lines changed: 98 additions & 23 deletions b/‎docs/advanced_tutorials/20_score_based_methods_new_features.ipynb‎
Lines changed: 98 additions & 23 deletions
diff --git a/‎sbi/inference/__init__.py‎
Lines changed: 1 addition & 2 deletions b/‎sbi/inference/__init__.py‎
Lines changed: 1 addition & 2 deletions
diff --git a/‎sbi/inference/trainers/base.py‎
Lines changed: 12 additions & 6 deletions b/‎sbi/inference/trainers/base.py‎
Lines changed: 12 additions & 6 deletions
diff --git a/‎sbi/inference/trainers/npse/__init__.py‎
Lines changed: 0 additions & 4 deletions b/‎sbi/inference/trainers/npse/__init__.py‎
Lines changed: 0 additions & 4 deletions
diff --git a/‎sbi/inference/trainers/fmpe/__init__.py‎ ‎sbi/inference/trainers/vfpe/__init__.py‎sbi/inference/trainers/fmpe/__init__.py renamed to sbi/inference/trainers/vfpe/__init__.py
Lines changed: 2 additions & 2 deletions b/‎sbi/inference/trainers/fmpe/__init__.py‎ ‎sbi/inference/trainers/vfpe/__init__.py‎sbi/inference/trainers/fmpe/__init__.py renamed to sbi/inference/trainers/vfpe/__init__.py
Lines changed: 2 additions & 2 deletions
diff --git a/‎…/trainers/npse/vector_field_inference.py‎ ‎…rence/trainers/vfpe/base_vf_inference.py‎sbi/inference/trainers/npse/vector_field_inference.py renamed to sbi/inference/trainers/vfpe/base_vf_inference.py
Lines changed: 93 additions & 8 deletions b/‎…/trainers/npse/vector_field_inference.py‎ ‎…rence/trainers/vfpe/base_vf_inference.py‎sbi/inference/trainers/npse/vector_field_inference.py renamed to sbi/inference/trainers/vfpe/base_vf_inference.py
Lines changed: 93 additions & 8 deletions
diff --git a/‎sbi/inference/trainers/fmpe/fmpe.py‎ ‎sbi/inference/trainers/vfpe/fmpe.py‎sbi/inference/trainers/fmpe/fmpe.py renamed to sbi/inference/trainers/vfpe/fmpe.py
Lines changed: 32 additions & 14 deletions b/‎sbi/inference/trainers/fmpe/fmpe.py‎ ‎sbi/inference/trainers/vfpe/fmpe.py‎sbi/inference/trainers/fmpe/fmpe.py renamed to sbi/inference/trainers/vfpe/fmpe.py
Lines changed: 32 additions & 14 deletions
@@ -7,12 +7,11 @@
     check_if_proposal_has_default_x,
     infer,
 )
-from sbi.inference.trainers.fmpe import FMPE
 from sbi.inference.trainers.marginal import MarginalTrainer
 from sbi.inference.trainers.nle import MNLE, NLE_A
 from sbi.inference.trainers.npe import MNPE, NPE_A, NPE_B, NPE_C  # noqa: F401
-from sbi.inference.trainers.npse import NPSE
 from sbi.inference.trainers.nre import BNRE, NRE_A, NRE_B, NRE_C  # noqa: F401
+from sbi.inference.trainers.vfpe import FMPE, NPSE
 
 SNL = SNLE = SNLE_A = NLE = NLE_A
 _nle_family = ["NLE"]
 
@@ -673,13 +673,19 @@ def _raise_deprecation_warning(
         """
 
         deprecated_params = deprecated_params.copy()
-
-        is_default_mcmc_method = kwargs.get("mcmc_method") == "slice_np_vectorized"
-        is_default_vi_method = kwargs.get("vi_method") == "rKL"
-
-        if not is_default_mcmc_method:
+        default_mcmc_method = "slice_np_vectorized"
+        default_vi_method = "rKL"
+
+        # Check if deprecated parameters are used
+        if (
+            kwargs.get("mcmc_method") == default_mcmc_method
+            or kwargs.get("mcmc_method") is None
+        ):
             deprecated_params.append("mcmc_method")
-        if not is_default_vi_method:
+        if (
+            kwargs.get("vi_method") == default_vi_method
+            or kwargs.get("vi_method") is None
+        ):
             deprecated_params.append("vi_method")
 
         if deprecated_params:
 
@@ -1,4 +1,4 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
-
-from sbi.inference.trainers.fmpe.fmpe import FMPE
+from sbi.inference.trainers.vfpe.fmpe import FMPE
+from sbi.inference.trainers.vfpe.npse import NPSE
@@ -4,7 +4,7 @@
 import time
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Any, Callable, Optional, Protocol, Tuple, Union
+from typing import Any, Callable, Literal, Optional, Protocol, Tuple, Union
 
 import torch
 from torch import Tensor, ones
@@ -60,7 +60,10 @@ class VectorFieldTrainer(NeuralInference, ABC):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        vector_field_estimator_builder: Union[str, VectorFieldEstimatorBuilder] = "mlp",
+        vector_field_estimator_builder: Union[
+            Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
+            VectorFieldEstimatorBuilder,
+        ] = "mlp",
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
         summary_writer: Optional[SummaryWriter] = None,
@@ -106,15 +109,19 @@ def __init__(
         check_estimator_arg(vector_field_estimator_builder)
         if isinstance(vector_field_estimator_builder, str):
             self._build_neural_net = self._build_default_nn_fn(
-                vector_field_estimator_builder=vector_field_estimator_builder, **kwargs
+                model=vector_field_estimator_builder, **kwargs
             )
         else:
             self._build_neural_net = vector_field_estimator_builder
 
         self._proposal_roundwise = []
 
     @abstractmethod
-    def _build_default_nn_fn(self, **kwargs) -> VectorFieldEstimatorBuilder:
+    def _build_default_nn_fn(
+        self,
+        model: Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
+        **kwargs,
+    ) -> VectorFieldEstimatorBuilder:
         pass
 
     def append_simulations(
@@ -209,12 +216,13 @@ def train(
         training_batch_size: int = 200,
         learning_rate: float = 5e-4,
         validation_fraction: float = 0.1,
-        stop_after_epochs: int = 50,
-        max_num_epochs: int = 500,
+        stop_after_epochs: int = 20,
+        max_num_epochs: int = 2**31 - 1,
         clip_max_norm: Optional[float] = 5.0,
         calibration_kernel: Optional[Callable] = None,
         ema_loss_decay: float = 0.1,
-        validation_times: Union[Tensor, int] = 20,
+        validation_times: Union[Tensor, int] = 10,
+        validation_times_nugget: float = 0.05,
         resume_training: bool = False,
         force_first_round_loss: bool = False,
         discard_prior_samples: bool = False,
@@ -253,6 +261,9 @@ def train(
                 training and validation losses.
             validation_times: Diffusion times at which to evaluate the validation loss
                 to reduce variance of validation loss.
+            validation_times_nugget: As both diffusion and flow matching losses often
+                have high variance losses at the end, we add a small nugget to compute
+                the validation loss. Default is 0.05 i.e. t_min + 0.05 or t_max - 0.5.
             resume_training: Can be used in case training time is limited, e.g. on a
                 cluster. If `True`, the split between train and validation set, the
                 optimizer, the number of epochs, and the best validation log-prob will
@@ -341,7 +352,9 @@ def default_calibration_kernel(x):
 
         if isinstance(validation_times, int):
             validation_times = torch.linspace(
-                self._neural_net.t_min, self._neural_net.t_max, validation_times
+                self._neural_net.t_min + validation_times_nugget,
+                self._neural_net.t_max - validation_times_nugget,
+                validation_times,
             )
         assert isinstance(
             validation_times, Tensor
@@ -431,6 +444,8 @@ def default_calibration_kernel(x):
                         times_batch, *([1] * (masks_batch.ndim - 1))
                     )
 
+                    # This will repeat the validation times for each batch in the
+                    # validation set.
                     validation_times_rep = validation_times.repeat_interleave(
                         val_batch_size, dim=0
                     )
@@ -486,6 +501,76 @@ def default_calibration_kernel(x):
 
         return deepcopy(self._neural_net)
 
+    def _converged(self, epoch: int, stop_after_epochs: int) -> bool:
+        """Return whether the training converged yet and save best model state so far.
+
+        Diffusion or flow matching objectives are inherently more stochastic than MLE
+        for e.g. NPE because they additionally add "noise" by construction. We hence
+        use a statistical approach to detect convergence by tracking standard deviation
+        of validation losses. Training is considered converged when the current loss is
+        significantly worse than the best loss for a sustained period (more than 2 std
+        deviations above best).
+
+        NOTE: The standard deviation of the `validation_loss `is computed in a running
+            fashion over the most recent 2 × stop_after_epochs loss values.
+
+        Args:
+            epoch: Current epoch in training.
+            stop_after_epochs: How many fruitless epochs to let pass before stopping.
+
+        Returns:
+            Whether the training has stopped improving, i.e. has converged.
+        """
+        converged = False
+
+        assert self._neural_net is not None
+        neural_net = self._neural_net
+
+        # Initialize tracking variables if not exists
+        if not hasattr(self, '_best_val_loss'):
+            self._best_val_loss = float('inf')
+            self._epochs_since_last_improvement = 0
+            self._best_model_state_dict = None
+
+        # Check if we have a new best loss
+        if self._val_loss < self._best_val_loss:
+            self._best_val_loss = self._val_loss
+            self._epochs_since_last_improvement = 0
+            self._best_model_state_dict = deepcopy(neural_net.state_dict())
+        else:
+            # Only start statistical analysis after we have enough data
+            if len(self._summary["validation_loss"]) >= stop_after_epochs:
+                # Calculate running statistics of recent losses
+                recent_losses = torch.tensor(
+                    self._summary["validation_loss"][-stop_after_epochs * 2 :]
+                )
+                loss_std = recent_losses.std().item()
+
+                # Calculate how many standard deviations the current loss is from the
+                # best
+                diff_to_best_normalized = (
+                    self._val_loss - self._best_val_loss
+                ) / loss_std
+                # Consider it "no improvement" if current loss is significantly
+                # worse than the best loss (more than 2 std deviations above best)
+                # This accounts for natural fluctuations while being sensitive to
+                # real degradation
+                if diff_to_best_normalized > 2.0:
+                    self._epochs_since_last_improvement += 1
+                else:
+                    # Reset counter if loss is within acceptable range
+                    self._epochs_since_last_improvement = 0
+            else:
+                return False
+
+        # If no validation improvement over many epochs, stop training.
+        if self._epochs_since_last_improvement > stop_after_epochs - 1:
+            if self._best_model_state_dict is not None:
+                neural_net.load_state_dict(self._best_model_state_dict)
+            converged = True
+
+        return converged
+
     def _get_potential_function(
         self, prior: Distribution, estimator: ConditionalVectorFieldEstimator
     ) -> Tuple[VectorFieldBasedPotential, TorchTransform]:
 
@@ -2,6 +2,7 @@
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
 
+import warnings
 from typing import Any, Dict, Literal, Optional, Union
 
 from torch.distributions import Distribution
@@ -10,12 +11,12 @@
 from sbi import utils as utils
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.posteriors.posterior_parameters import VectorFieldPosteriorParameters
-from sbi.inference.trainers.npse.vector_field_inference import (
+from sbi.inference.trainers.vfpe.base_vf_inference import (
     VectorFieldEstimatorBuilder,
     VectorFieldTrainer,
 )
-from sbi.neural_nets import flowmatching_nn
-from sbi.neural_nets.estimators import ConditionalVectorFieldEstimator
+from sbi.neural_nets.estimators.base import ConditionalVectorFieldEstimator
+from sbi.neural_nets.factory import posterior_flow_nn
 
 
 class FMPE(VectorFieldTrainer):
@@ -24,7 +25,11 @@ class FMPE(VectorFieldTrainer):
     def __init__(
         self,
         prior: Optional[Distribution],
-        density_estimator: Union[str, VectorFieldEstimatorBuilder] = "mlp",
+        vf_estimator: Union[
+            Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
+            VectorFieldEstimatorBuilder,
+        ] = "mlp",
+        density_estimator: Optional[VectorFieldEstimatorBuilder] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "WARNING",
         summary_writer: Optional[SummaryWriter] = None,
@@ -35,11 +40,14 @@ def __init__(
 
         Args:
             prior: Prior distribution.
-            density_estimator: Neural network architecture used to learn the
-                vector field estimator. Can be a string (e.g. 'mlp' or 'ada_mlp') or a
-                callable that implements the `VectorFieldEstimatorBuilder` protocol
-                with `__call__` that receives `theta` and `x` and returns a
+            vf_estimator: Neural network architecture used to learn the
+                vector field estimator. Can be a string (e.g. 'mlp', 'ada_mlp',
+                'transformer' or 'transformer_cross_attn') or a callable that
+                implements the `VectorFieldEstimatorBuilder` protocol with
+                `__call__` that receives `theta` and `x` and returns a
                 `ConditionalVectorFieldEstimator`.
+            density_estimator: Deprecated. Use `vf_estimator` instead. When passed, a
+                warning is raised and the `vf_estimator="mlp"` default is used.
             device: Device to use for training.
             logging_level: Logging level.
             summary_writer: Summary writer for tensorboard.
@@ -48,17 +56,24 @@ def __init__(
                 `density_estimator` is a string.
         """
 
+        if density_estimator is not None:
+            warnings.warn(
+                "`density_estimator` is deprecated and will be removed in a future "
+                "release. Use `vf_estimator` instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            vf_estimator = density_estimator
+
         super().__init__(
             prior=prior,
             device=device,
             logging_level=logging_level,
             summary_writer=summary_writer,
             show_progress_bars=show_progress_bars,
-            vector_field_estimator_builder=density_estimator,
+            vector_field_estimator_builder=vf_estimator,
             **kwargs,
         )
-        # density_estimator name is kept since it is public API, but it is
-        # actually misleading since it is a builder for an estimator.
 
     def build_posterior(
         self,
@@ -106,6 +121,9 @@ def build_posterior(
             posterior_parameters=posterior_parameters,
         )
 
-    def _build_default_nn_fn(self, **kwargs) -> VectorFieldEstimatorBuilder:
-        model = kwargs.pop("vector_field_estimator_builder", "mlp")
-        return flowmatching_nn(model=model, **kwargs)
+    def _build_default_nn_fn(
+        self,
+        model: Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
+        **kwargs,
+    ) -> VectorFieldEstimatorBuilder:
+        return posterior_flow_nn(model=model, **kwargs)