Unify flow matching and score-based models

sbi/diagnostics/sbc.py

@@ -10,7 +10,7 @@
 from torch.distributions import Uniform
 from tqdm.auto import tqdm
 
-from sbi.inference import DirectPosterior, ScorePosterior
+from sbi.inference import DirectPosterior, VectorFieldPosterior
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.posteriors.vi_posterior import VIPosterior
 from sbi.utils.diagnostics_utils import (
@@ -186,7 +186,9 @@ def get_nltp(thetas: Tensor, xs: Tensor, posterior: NeuralPosterior) -> Tensor:
         nltp: negative log probs of true parameters under approximate posteriors.
     """
     nltp = torch.zeros(thetas.shape[0])
-    unnormalized_log_prob = not isinstance(posterior, (DirectPosterior, ScorePosterior))
+    unnormalized_log_prob = not isinstance(
+        posterior, (DirectPosterior, VectorFieldPosterior)
+    )
 
     for idx, (tho, xo) in enumerate(zip(thetas, xs, strict=False)):
         # Log prob of true params under posterior.

sbi/inference/__init__.py

@@ -38,13 +38,14 @@
     ImportanceSamplingPosterior,
     MCMCPosterior,
     RejectionPosterior,
-    ScorePosterior,
     VIPosterior,
+    VectorFieldPosterior,
 )
 from sbi.inference.potentials import (
     likelihood_estimator_based_potential,
     mixed_likelihood_estimator_based_potential,
     posterior_estimator_based_potential,
     ratio_estimator_based_potential,
+    vector_field_estimator_based_potential,
 )
 from sbi.utils.simulation_utils import simulate_for_sbi

sbi/inference/posteriors/__init__.py

@@ -3,5 +3,5 @@
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.posteriors.mcmc_posterior import MCMCPosterior
 from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
-from sbi.inference.posteriors.score_posterior import ScorePosterior
+from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
 from sbi.inference.posteriors.vi_posterior import VIPosterior

sbi/inference/potentials/__init__.py

@@ -8,3 +8,6 @@
 from sbi.inference.potentials.ratio_based_potential import (
     ratio_estimator_based_potential,
 )
+from sbi.inference.potentials.vector_field_potential import (
+    vector_field_estimator_based_potential,
+)

sbi/inference/potentials/score_fn_iid.py

@@ -7,7 +7,7 @@
 from torch import Tensor
 from torch.distributions import Distribution
 
-from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
+from sbi.neural_nets.estimators import ConditionalVectorFieldEstimator
 from sbi.utils.score_utils import (
     add_diag_or_dense,
     denoise,
@@ -58,7 +58,7 @@ def decorator(cls: Type["IIDScoreFunction"]) -> Type["IIDScoreFunction"]:
 class IIDScoreFunction(ABC):
     def __init__(
         self,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         device: str = "cpu",
     ) -> None:
@@ -76,12 +76,34 @@ def __init__(
         this requires some adjustments.
 
         Args:
-            score_estimator: The neural network modeling the score.
+            vector_field_estimator: The neural network modeling the score.
             prior: The prior distribution.
             device: The device on which to evaluate the potential. Defaults to "cpu".
         """
 
-        self.score_estimator = score_estimator.to(device).eval()
+        if not vector_field_estimator.SCORE_DEFINED:
+            raise ValueError(
+                "Score is not defined for this vector field estimator. "
+                "You are probably using a vector field estimator that does not "
+                "implement the score function, e.g. an optimal transport-based "
+                "flow matching. IID methods require the score function to be defined "
+                "so they are not applicable to this vector field estimator. "
+                "If you have implemented a custom vector field "
+                "estimator with score defined, set SCORE_DEFINED to True."
+            )
+        if not vector_field_estimator.MARGINALS_DEFINED:
+            raise ValueError(
+                "Marginals are not defined for this vector field estimator. "
+                "You are probably using a vector field estimator that does not "
+                "implement the marginals mean_t_fn and std_fn, e.g. an "
+                "optimal transport-based flow matching. IID methods require the "
+                "marginals to be defined so they are not applicable to this "
+                "vector field estimator. "
+                "If you have implemented a custom vector field "
+                "estimator with marginals defined, set MARGINALS_DEFINED to True."
+            )
+
+        self.vector_field_estimator = vector_field_estimator.to(device).eval()
         self.prior = prior
 
     @abstractmethod
@@ -106,17 +128,17 @@ def __call__(
 
 
 @register_iid_method("fnpe")
-class FNPEScoreFunction(IIDScoreFunction):
+class FactorizedNPEScoreFunction(IIDScoreFunction):
     def __init__(
         self,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         device: str = "cpu",
         prior_score_weight: Optional[Callable[[Tensor], Tensor]] = None,
     ) -> None:
         r"""
-        The FNPEScoreFunction implments the "Factorized Neural Posterior Estimation"
-        method for score-based models [1].
+        The FactorizedNPEScoreFunction implments the
+        "Factorized Neural Posterior Estimation" method for score-based models [1].
 
         This method does not apply the necessary corrections for the score function, but
         instead uses a simple weighting of the prior score. This is generally applicable
@@ -131,15 +153,15 @@ def __init__(
             (https://arxiv.org/abs/2209.14249)
 
         Args:
-            score_estimator: The neural network modeling the score.
+            vector_field_estimator: The neural network modeling the score.
             prior: The prior distribution.
             device: The device on which to evaluate the potential. Defaults to "cpu".
             prior_score_weight: A function to weight the prior score. Defaults to the
                 linear interpolation between zero (at t=0) and one (at t=t_max).
         """
-        super().__init__(score_estimator, prior, device)
+        super().__init__(vector_field_estimator, prior, device)
         if prior_score_weight is None:
-            t_max = score_estimator.t_max
+            t_max = vector_field_estimator.t_max
 
             def prior_score_weight_fn(t):
                 return (t_max - t) / t_max
@@ -168,13 +190,13 @@ def __call__(
         assert conditions.ndim == 2, "Conditions must have shape [iid,...]"
 
         if time is None:
-            time = torch.tensor([self.score_estimator.t_min])
+            time = torch.tensor([self.vector_field_estimator.t_min])
 
         N = conditions.shape[0]
 
         # Compute the per-sample score
         inputs = ensure_theta_batched(inputs)
-        base_score = self.score_estimator(inputs, conditions, time)
+        base_score = self.vector_field_estimator.score(inputs, conditions, time)
 
         # Compute the prior score
         prior_score = self.prior_score_weight_fn(time) * self.prior_score_fn(inputs)
@@ -214,7 +236,7 @@ def prior_score_fn(self, theta: Tensor) -> Tensor:
 class BaseGaussCorrectedScoreFunction(IIDScoreFunction):
     def __init__(
         self,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         ensure_lam_psd: bool = True,
         lam_psd_nugget: float = 0.01,
@@ -244,13 +266,13 @@ def __init__(
             (https://arxiv.org/abs/2411.02728)
 
         Args:
-            score_estimator: The neural network modelling the score.
+            vector_field_estimator: The neural network modelling the score.
             prior: The prior distribution.
             ensure_lam_psd: Whether to ensure the precision matrix is positive definite.
             lam_psd_nugget: The nugget value to ensure positive definiteness.
             device: The device on which to evaluate the potential. Defaults to "cpu".
         """
-        super().__init__(score_estimator, prior, device)
+        super().__init__(vector_field_estimator, prior, device)
         self.ensure_lam_psd = ensure_lam_psd
         self.lam_psd_nugget = lam_psd_nugget
 
@@ -289,8 +311,8 @@ def marginal_denoising_posterior_precision_est_fn(
         precisions_posteriors = self.posterior_precision_est_fn(conditions)
 
         # Denoising posterior via Bayes rule
-        m = self.score_estimator.mean_t_fn(time)
-        std = self.score_estimator.std_fn(time)
+        m = self.vector_field_estimator.mean_t_fn(time)
+        std = self.vector_field_estimator.std_fn(time)
 
         if precisions_posteriors.ndim == 4:
             Ident = torch.eye(precisions_posteriors.shape[-1])
@@ -315,8 +337,8 @@ def marginal_prior_score_fn(self, time: Tensor, inputs: Tensor) -> Tensor:
         try:
             with torch.enable_grad():
                 inputs = inputs.clone().detach().requires_grad_(True)
-                m = self.score_estimator.mean_t_fn(time)
-                std = self.score_estimator.std_fn(time)
+                m = self.vector_field_estimator.mean_t_fn(time)
+                std = self.vector_field_estimator.std_fn(time)
                 p = marginalize(self.prior, m, std)
                 log_p = p.log_prob(inputs)
                 prior_score = torch.autograd.grad(
@@ -342,8 +364,8 @@ def marginal_denoising_prior_precision_fn(
         Returns:
             Marginal denoising prior precision.
         """
-        m = self.score_estimator.mean_t_fn(time)
-        std = self.score_estimator.std_fn(time)
+        m = self.vector_field_estimator.mean_t_fn(time)
+        std = self.vector_field_estimator.std_fn(time)
 
         p_denoise = denoise(self.prior, m, std, inputs)
 
@@ -388,9 +410,11 @@ def __call__(
         N, *_ = conditions.shape
 
         if time is None:
-            time = torch.tensor([self.score_estimator.t_min])
+            time = torch.tensor([self.vector_field_estimator.t_min])
 
-        base_score = self.score_estimator(inputs, conditions, time, **kwargs)
+        base_score = self.vector_field_estimator.score(
+            inputs, conditions, time, **kwargs
+        )
         prior_score = self.marginal_prior_score_fn(time, inputs)
 
         # Marginal prior precision
@@ -438,7 +462,7 @@ def __call__(
 class GaussCorrectedScoreFn(BaseGaussCorrectedScoreFunction):
     def __init__(
         self,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         posterior_precision: Optional[Tensor] = None,
         scale_from_prior_precision: float = 2.0,
@@ -453,7 +477,7 @@ def __init__(
         precision. This method simply scales the prior precision by a factor.
 
         Args:
-            score_estimator: The neural network modeling the score.
+            vector_field_estimator: The neural network modeling the score.
             prior: The prior distribution.
             posterior_precision: Optional preset posterior precision.
             scale_from_prior_precision: If not provided it simply increases the prior
@@ -463,7 +487,7 @@ def __init__(
             device: The device on which to evaluate the potential. Defaults to "cpu".
         """
         super().__init__(
-            score_estimator, prior, enable_lam_psd, lam_psd_nugget, device=device
+            vector_field_estimator, prior, enable_lam_psd, lam_psd_nugget, device=device
         )
 
         if posterior_precision is None:
@@ -511,7 +535,7 @@ def estimate_prior_precision(
         except Exception:
             d = prior.event_shape[0]
             num_samples = int(math.sqrt(d) * 1000)
-            prior_samples = prior.sample((num_samples,))
+            prior_samples = prior.sample(torch.Size((num_samples,)))
             prior_precision_estimate = 1 / torch.var(prior_samples, dim=0)
             posterior_precision = scale_from_prior_precision * prior_precision_estimate
         return posterior_precision
@@ -521,7 +545,7 @@ def estimate_prior_precision(
 class AutoGaussCorrectedScoreFn(BaseGaussCorrectedScoreFunction):
     def __init__(
         self,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         enable_lam_psd: bool = True,
         lam_psd_nugget: float = 0.01,
@@ -533,11 +557,11 @@ def __init__(
         r"""
         This method extends the by estimating the posterior precision using
         approximate posterior samples obtained from a diffusion model (using the
-        score_estimator) [1]. This method has a slight initialization overhead, but
-        generally provides more accurate results than simple heuristics.
+        vector_field_estimator) [1]. This method has a slight initialization overhead,
+        but generally provides more accurate results than simple heuristics.
 
         Args:
-            score_estimator: The neural network modeling the score.
+            vector_field_estimator: The neural network modeling the score.
             prior: The prior distribution.
             enable_lam_psd: Whether to ensure the precision matrix is positive definite.
             lam_psd_nugget: The nugget value to ensure positive definiteness.
@@ -548,7 +572,7 @@ def __init__(
             device: The device on which to evaluate the potential. Defaults to "cpu".
         """
         super().__init__(
-            score_estimator, prior, enable_lam_psd, lam_psd_nugget, device=device
+            vector_field_estimator, prior, enable_lam_psd, lam_psd_nugget, device=device
         )
         self.precision_est_only_diag = precision_est_only_diag
         self.precision_est_budget = precision_est_budget
@@ -565,7 +589,7 @@ def posterior_precision_est_fn(self, conditions: Tensor) -> Tensor:
             Estimated posterior precision.
         """
         return self.estimate_posterior_precision(
-            self.score_estimator,
+            self.vector_field_estimator,
             self.prior,
             conditions,
             precision_est_only_diag=self.precision_est_only_diag,
@@ -577,7 +601,7 @@ def posterior_precision_est_fn(self, conditions: Tensor) -> Tensor:
     @functools.lru_cache()
     def estimate_posterior_precision(
         cls,
-        score_estimator: "ConditionalScoreEstimator",
+        vector_field_estimator: ConditionalVectorFieldEstimator,
         prior: Distribution,
         conditions: Tensor,
         precision_est_only_diag: bool = False,
@@ -589,7 +613,7 @@ def estimate_posterior_precision(
         for each observation, then empirically estimating the precision matrix.
 
         Args:
-            score_estimator: The score estimator.
+            vector_field_estimator: The score estimator.
             prior: The prior distribution.
             conditions: Observed data.
             precision_est_only_diag: Whether to estimate only the diagonal of the
@@ -606,9 +630,9 @@ def estimate_posterior_precision(
         # the results efficiently.
 
         # NOTE: To avoid circular imports :(
-        from sbi.inference.posteriors.score_posterior import ScorePosterior
+        from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
 
-        posterior = ScorePosterior(score_estimator, prior)
+        posterior = VectorFieldPosterior(vector_field_estimator, prior)
 
         if precision_est_budget is None:
             if precision_est_only_diag:
@@ -674,16 +698,16 @@ def marginal_denoising_posterior_precision_est_fn(
         with torch.enable_grad():
             # NOTE: torch.func can be realtively unstable...
             jac_fn = torch.func.jacrev(  # type: ignore
-                lambda x: self.score_estimator(x, conditions, time)
+                lambda x: self.vector_field_estimator.score(x, conditions, time)
             )
             jac_fn = torch.func.vmap(torch.func.vmap(jac_fn))  # type: ignore
             jac = jac_fn(inputs).squeeze(1)
 
         # Must be symmetrical
         jac = 0.5 * (jac + jac.transpose(-1, -2))
 
-        m = self.score_estimator.mean_t_fn(time)
-        std = self.score_estimator.std_fn(time)
+        m = self.vector_field_estimator.mean_t_fn(time)
+        std = self.vector_field_estimator.std_fn(time)
         cov0 = std**2 * jac + torch.eye(d)[None, None, :, :]
 
         denoising_posterior_precision = m**2 / std**2 + torch.inverse(cov0)