fix: do not save simulations in empirical prior (#1700)

Author: gmoss13

sbi/inference/potentials/posterior_based_potential.py

@@ -17,9 +17,7 @@
 )
 from sbi.sbi_types import TorchTransform
 from sbi.utils.sbiutils import (
-    ImproperEmpirical,
     mcmc_transform,
-    warn_empirical_prior_memory_risk,
     within_support,
 )
 from sbi.utils.torchutils import ensure_theta_batched
@@ -102,9 +100,6 @@ def to(self, device: Union[str, torch.device]) -> None:
         self.device = device
         self.posterior_estimator.to(device)
         if self.prior is not None:
-            is_empirical = isinstance(self.prior, ImproperEmpirical)
-            if is_empirical and torch.device(device).type == "cuda":
-                warn_empirical_prior_memory_risk("moving empirical prior to CUDA")
             self.prior.to(device)  # type: ignore
         if self._x_o is not None:
             self._x_o = self._x_o.to(device)

sbi/inference/trainers/npe/npe_base.py

@@ -53,7 +53,6 @@
 from sbi.utils.sbiutils import (
     ImproperEmpirical,
     mask_sims_from_prior,
-    warn_empirical_prior_memory_risk,
 )
 from sbi.utils.torchutils import assert_all_finite
 
@@ -470,18 +469,10 @@ def _get_potential_function(
             to unconstrained space.
         """
 
-        is_empirical = isinstance(prior, ImproperEmpirical)
-        if is_empirical:
-            warn_empirical_prior_memory_risk(
-                "disabling parameter transforms for empirical prior"
-            )
-
         potential_fn, theta_transform = posterior_estimator_based_potential(
             posterior_estimator=estimator,
             prior=prior,
             x_o=None,
-            # Disable transforms if prior is empirical to avoid sampling issues.
-            enable_transform=not is_empirical,
         )
         return potential_fn, theta_transform
 

sbi/utils/sbiutils.py

@@ -35,7 +35,6 @@
 from torch.optim.adam import Adam
 
 from sbi.sbi_types import TorchTransform
-from sbi.utils.torchutils import atleast_2d
 
 
 def warn_if_zscoring_changes_data(x: Tensor, duplicate_tolerance: float = 0.1) -> None:
@@ -242,21 +241,6 @@ def biject_transform_zuko(
     )
 
 
-def warn_empirical_prior_memory_risk(context: Optional[str] = None) -> None:
-    """Emit a standardized warning about empirical-prior memory/VRAM risks.
-
-    Args:
-        context: Optional context string to append to the warning.
-    """
-    base = (
-        "Empirical prior memory/VRAM risk: empirical priors retain all simulations "
-        "as support and may trigger operations over large supports. This can "
-        "significantly increase memory usage and cause out-of-memory (OOM) errors."
-    )
-    message = f"{base} Context: {context}" if context else base
-    warnings.warn(message, stacklevel=2)
-
-
 def z_standardization(
     batch_t: Tensor,
     structured_dims: bool = False,
@@ -752,13 +736,6 @@ def mcmc_transform(
         (or z-scored) to constrained (or non-z-scored) space.
     """
     if enable_transform:
-        if isinstance(prior, (ImproperEmpirical, Empirical)):
-            warn_empirical_prior_memory_risk(
-                "disabled parameter transforms to avoid sampling-based moments"
-            )
-            return torch_tf.IndependentTransform(
-                torch_tf.identity_transform, reinterpreted_batch_ndims=1
-            )
 
         def prior_mean_std_transform(prior, device):
             try:
@@ -850,7 +827,12 @@ def check_transform(
 ) -> None:
     """Check validity of transformed and re-transformed samples."""
 
-    theta = prior.sample(torch.Size((2,)))
+    # check transform with prior samples
+    try:
+        theta = prior.sample(torch.Size((2,)))
+    except NotImplementedError:
+        # Prior has no sampling method, use the prior mean instead
+        theta = prior.mean.repeat(2, *[1] * prior.mean.dim())
 
     theta_unconstrained = transform.inv(theta)
     assert (
@@ -881,9 +863,15 @@ class ImproperEmpirical(Empirical):
     def __init__(self, values: Tensor, log_weights: Optional[Tensor] = None):
         super().__init__(values, log_weights=log_weights)
         # Warn if extremely large to inform about memory/serialization cost.
-        support_size = values.shape[0]
-        if support_size > 10_000_000:  # 10M still works well on modern hardware.
-            warn_empirical_prior_memory_risk(f">10M support size (size={support_size})")
+        self._mean = self._compute_mean(values, log_weights)
+        self._variance = self._compute_variance(values, log_weights)
+
+    def sample(self, sample_shape=torch.Size()):
+        raise NotImplementedError(
+            "Sampling from ImproperEmpirical is not supported. If you are using "
+            "likelihood or ratio estimation, or multi-round inference, you need to "
+            "define a prior distribution."
+        )
 
     def log_prob(self, value: Tensor) -> Tensor:
         """
@@ -895,8 +883,79 @@ def log_prob(self, value: Tensor) -> Tensor:
         Returns:
             Tensor of as many ones as there were parameter sets.
         """
-        value = atleast_2d(value)
-        return zeros(value.shape[0])
+        raise NotImplementedError(
+            "Evaluating log_prob from ImproperEmpirical is not supported. If you are "
+            "using likelihood or ratio estimation, or multi-round inference, you need "
+            "to define a prior distribution."
+        )
+
+    def _compute_mean(self, values: Tensor, weights: Optional[Tensor] = None) -> Tensor:
+        """
+        Return the mean of the empirical distribution.
+
+        Args:
+            values: The empirical samples.
+            weights: Optional weights for the samples.
+
+        Returns:
+            The mean of the empirical distribution.
+        """
+        if weights is None:
+            return torch.mean(values, dim=0)
+        else:
+            normalized_weights = torch.nn.functional.softmax(weights, dim=0)
+            return torch.sum(normalized_weights.unsqueeze(-1) * values, dim=0)
+
+    def _compute_variance(
+        self, values: Tensor, weights: Optional[Tensor] = None
+    ) -> Tensor:
+        """
+        Return the standard deviation of the empirical distribution.
+
+        Args:
+            values: The empirical samples.
+            weights: Optional weights for the samples.
+
+        Returns:
+            The standard deviation of the empirical distribution.
+        """
+        if weights is None:
+            variance = torch.var(values, dim=0)
+        else:
+            normalized_weights = torch.nn.functional.softmax(weights, dim=0)
+            variance = torch.sum(
+                normalized_weights.unsqueeze(-1) * (values - self._mean) ** 2,
+                dim=0,
+            )
+            # bias correction
+            variance = variance / (1 - torch.sum(normalized_weights**2))
+        return variance
+
+    @property
+    def mean(self) -> Tensor:
+        return self._mean
+
+    @property
+    def variance(self) -> Tensor:
+        return self._variance
+
+    @property
+    def stddev(self) -> Tensor:
+        return torch.sqrt(self._variance)
+
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move the distribution to a different device.
+
+        Args:
+            device: The device to move the distribution to.
+
+        Returns:
+            The distribution on the specified device.
+        """
+        self._mean = self._mean.to(device)
+        self._variance = self._variance.to(device)
+        super().to(device)
 
 
 def mog_log_prob(

tests/sbiutils_test.py

@@ -1,7 +1,6 @@
 # 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/>
 
-import warnings
 from typing import Tuple
 
 import matplotlib.pyplot as plt
@@ -21,7 +20,7 @@
 from sbi.inference.trainers.npe.npe_a import NPE_A_MDN
 from sbi.neural_nets import classifier_nn, likelihood_nn, posterior_nn
 from sbi.utils import BoxUniform, get_kde
-from sbi.utils.sbiutils import ImproperEmpirical, mcmc_transform, z_score_parser
+from sbi.utils.sbiutils import z_score_parser
 
 
 def test_conditional_density_1d():
@@ -555,15 +554,3 @@ def test_z_scoring_structured(z_x, z_theta, builder):
     # plt.plot(x_zstructured.T)
     # plt.title('z-scored: structured dims');
     # plt.show()
-
-
-def test_mcmc_transform_emits_warning_for_improper_empirical():
-    values = torch.randn(100, 3)
-    logw = torch.zeros(values.shape[0])
-    prior = ImproperEmpirical(values, log_weights=logw)
-    with warnings.catch_warnings(record=True) as w:
-        warnings.simplefilter("always")
-        _ = mcmc_transform(prior, enable_transform=True)
-        assert any("Empirical prior memory/VRAM risk" in str(ww.message) for ww in w), (
-            "Expected generic empirical prior memory/VRAM risk warning."
-        )