feat: score-based iid sampling (#1381)

* npse MAP

* set default enable_Transform to True

* ruff formatting version change

* sampling via diffusion twice

* batched sampling for score-based posteriors

* iid api integration

* new ruff

* adding corrector back

* adding untesed GAUSS

* All other methods

* reformat

* messy version of simple gauss, API to sample method

* jac method (still needs feasible Lambda projection to work)

* Only jac left

* Formating and so on

* Adding correct tests

* Add empirical support - but this doesnt work that well

* A bunch of auto marginalize and denois methods

* general prior with GMM approx

* Bunch of reffactorings and customizability

* Update API docstirngs

* New tests, passes all now

* Formating linting, type error form other PR

* Remove assert that IID data is not supported

* Ruff linting

* Fixing comments

* Rearangements

* Missing types

* Improving tests doc

* Refactored tests with fixtures!

* refactored the tests

* Format tests

* more documentaiton

* Fix FNPE intitialization - not that nice

* Last issues resolved

* improve docstrings

* Raise error on potential. Make BaseClass abstract and remove redundant arguments.

* cosmetics

* update changelog

---------

Co-authored-by: Guy Moss <guy.moss13@gmail.com>
Co-authored-by: Jan <jan.boelts@mailbox.org>

Author: 3 people

CHANGELOG.md

@@ -4,8 +4,9 @@
 
 ### ✨ Highlights
 
-- Add `CategoricalMADE` by @jnsbck in https://github.com/sbi-dev/sbi/pull/1269  **(Major New Feature)**
+- feat: add `CategoricalMADE` by @jnsbck in https://github.com/sbi-dev/sbi/pull/1269  **(Major New Feature)**
 - tests: `mini-sbibm` by @manuelgloeckler in https://github.com/sbi-dev/sbi/pull/1335 **(Major New Feature)**
+- feat: Score-based iid sampling by @manuelgloeckler in https://github.com/sbi-dev/sbi/pull/1381 **(Major New Feature)**
 - Drop python3.9 support, fix ci by @janfb in https://github.com/sbi-dev/sbi/pull/1412 **(Python Version Support Change)**
 - additional features for NPSE by @gmoss13 in https://github.com/sbi-dev/sbi/pull/1370 **(Enhancement)**
 
@@ -21,6 +22,7 @@
 - fix mnle tests by @janfb in https://github.com/sbi-dev/sbi/pull/1415
 - fix: protocol and refactor for custom potential by @janfb in https://github.com/sbi-dev/sbi/pull/1409
 - fix docs workflow by @janfb in https://github.com/sbi-dev/sbi/pull/1419
+- fix: gpu-handling for CategoricalMADE by @janfb in https://github.com/sbi-dev/sbi/pull/1448
 
 ### 🛠️ Maintenance & Improvements
 

sbi/inference/posteriors/score_posterior.py

@@ -105,6 +105,8 @@ def sample(
         corrector_params: Optional[Dict] = None,
         steps: int = 500,
         ts: Optional[Tensor] = None,
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[Dict] = None,
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
@@ -123,6 +125,19 @@ def sample(
             steps: Number of steps to take for the Euler-Maruyama method.
             ts: Time points at which to evaluate the diffusion process. If None, a
                 linear grid between t_max and t_min is used.
+            iid_method: Which method to use for computing the score in the iid setting.
+                We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss". The
+                fnpe method is simple and generally applicable. However, it can become
+                inaccurate already for quite a few iid samples (as it based on heuristic
+                approximations), and should be used at best only with a `corrector`. The
+                "gauss" methods are more accurate, by aiming for an efficient
+                approximation of the correct marginal score in the iid case. This
+                however requires estimating some hyperparamters, which is done in a
+                systematic way in the "auto_gauss" (initial overhead) and "jac_gauss"
+                (iterative jacobian computations are expensive). We default to
+                "auto_gauss" for these reasons.
+            iid_params: Additional parameters passed to the iid method. See the specific
+                `IIDScoreFunction` child class for details.
             max_sampling_batch_size: Maximum batch size for sampling.
             sample_with: Deprecated - use `.build_posterior(sample_with=...)` prior to
                 `.sample()`.
@@ -138,7 +153,10 @@ def sample(
 
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.score_estimator.condition_shape)
-        self.potential_fn.set_x(x, x_is_iid=True)
+        is_iid = x.shape[0] > 1
+        self.potential_fn.set_x(
+            x, x_is_iid=is_iid, iid_method=iid_method, iid_params=iid_params
+        )
 
         num_samples = torch.Size(sample_shape).numel()
 

sbi/inference/potentials/score_based_potential.py

@@ -1,7 +1,7 @@
 # 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 typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -10,6 +10,7 @@
 from zuko.transforms import FreeFormJacobianTransform
 
 from sbi.inference.potentials.base_potential import BasePotential
+from sbi.inference.potentials.score_fn_iid import get_iid_method
 from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
@@ -56,7 +57,8 @@ def __init__(
         score_estimator: ConditionalScoreEstimator,
         prior: Optional[Distribution],
         x_o: Optional[Tensor] = None,
-        iid_method: str = "iid_bridge",
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[Dict[str, Any]] = None,
         device: str = "cpu",
     ):
         r"""Returns the score function for score-based methods.
@@ -65,19 +67,24 @@ def __init__(
             score_estimator: The neural network modelling the score.
             prior: The prior distribution.
             x_o: The observed data at which to evaluate the posterior.
-            iid_method: Which method to use for computing the score. Currently, only
-                `iid_bridge` as proposed in Geffner et al. is implemented.
+            iid_method: Which method to use for computing the score in the iid setting.
+                We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss".
+            iid_params: Parameters for the iid method, for arguments see
+                `IIDScoreFunction`.
             device: The device on which to evaluate the potential.
         """
         self.score_estimator = score_estimator
         self.score_estimator.eval()
         self.iid_method = iid_method
+        self.iid_params = iid_params
         super().__init__(prior, x_o, device=device)
 
     def set_x(
         self,
         x_o: Optional[Tensor],
         x_is_iid: Optional[bool] = False,
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[Dict[str, Any]] = None,
         atol: float = 1e-5,
         rtol: float = 1e-6,
         exact: bool = True,
@@ -90,12 +97,20 @@ def set_x(
         Args:
             x_o: The observed data.
             x_is_iid: Whether the observed data is IID (if batch_dim>1).
+            iid_method: Which method to use for computing the score in the iid setting.
+                We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss".
+            iid_params: Parameters for the iid method, for arguments see
+                `IIDScoreFunction`.
             atol: Absolute tolerance for the ODE solver.
             rtol: Relative tolerance for the ODE solver.
             exact: Whether to use the exact ODE solver.
         """
         super().set_x(x_o, x_is_iid)
-        if self._x_o is not None:
+        self.iid_method = iid_method
+        self.iid_params = iid_params
+        # NOTE: Once IID potential evaluation is supported. This needs to be adapted.
+        # See #1450.
+        if not x_is_iid and (self._x_o is not None):
             self.flow = self.rebuild_flow(atol=atol, rtol=rtol, exact=exact)
 
     def __call__(
@@ -112,6 +127,15 @@ def __call__(
         Returns:
             The potential function, i.e., the log probability of the posterior.
         """
+
+        if self.x_is_iid:
+            raise NotImplementedError(
+                "Potential function evaluation in the IID setting is not yet supported"
+                " for score-based methods. Sampling does however work via `.sample`. "
+                "If you intended to evaluate the posterior given a batch of (non-iid) "
+                "x use `log_prob_batched`."
+            )
+
         theta = ensure_theta_batched(torch.as_tensor(theta))
         theta_density_estimator = reshape_to_sample_batch_event(
             theta, theta.shape[1:], leading_is_sample=True
@@ -160,10 +184,15 @@ def gradient(
                     input=theta, condition=self.x_o, time=time
                 )
             else:
-                raise NotImplementedError(
-                    "Score accumulation for IID data is not yet implemented."
+                assert self.prior is not None, "Prior is required for iid methods."
+
+                iid_method = get_iid_method(self.iid_method)
+                score_fn_iid = iid_method(
+                    self.score_estimator, self.prior, **(self.iid_params or {})
                 )
 
+                score = score_fn_iid(theta, self.x_o, time)
+
         return score
 
     def get_continuous_normalizing_flow(
@@ -205,9 +234,6 @@ def rebuild_flow(
         x_density_estimator = reshape_to_batch_event(
             self.x_o, event_shape=self.score_estimator.condition_shape
         )
-        assert x_density_estimator.shape[0] == 1 or not self.x_is_iid, (
-            "PosteriorScoreBasedPotential supports only x batchsize of 1`."
-        )
 
         flow = self.get_continuous_normalizing_flow(
             condition=x_density_estimator, atol=atol, rtol=rtol, exact=exact