feat: Add AmortizedVIPosterior for amortized variational inference

sbi/inference/__init__.py

@@ -33,8 +33,6 @@
 _abc_family = ["ABC", "MCABC", "SMC", "SMCABC"]
 
 
-__all__ = _npe_family + _nre_family + _nle_family + _abc_family + ["FMPE", "NPSE"]
-
 from sbi.inference.posteriors import (
     DirectPosterior,
     EnsemblePosterior,
@@ -53,4 +51,27 @@
 )
 from sbi.utils.simulation_utils import simulate_for_sbi
 
-__all__ = ["FMPE", "MarginalTrainer", "NLE", "NPE", "NPSE", "NRE", "simulate_for_sbi"]
+__all__ = (
+    _npe_family
+    + _nre_family
+    + _nle_family
+    + _abc_family
+    + [
+        "FMPE",
+        "MarginalTrainer",
+        "NPSE",
+        "DirectPosterior",
+        "EnsemblePosterior",
+        "ImportanceSamplingPosterior",
+        "MCMCPosterior",
+        "RejectionPosterior",
+        "VIPosterior",
+        "VectorFieldPosterior",
+        "simulate_for_sbi",
+        "likelihood_estimator_based_potential",
+        "mixed_likelihood_estimator_based_potential",
+        "posterior_estimator_based_potential",
+        "ratio_estimator_based_potential",
+        "vector_field_estimator_based_potential",
+    ]
+)

sbi/inference/posteriors/posterior_parameters.py

@@ -7,7 +7,6 @@
     Any,
     Callable,
     Dict,
-    Iterable,
     Literal,
     Optional,
     Union,
@@ -17,7 +16,7 @@
 )
 
 from sbi.inference.posteriors.vi_posterior import VIPosterior
-from sbi.sbi_types import PyroTransformedDistribution, TorchTransform
+from sbi.sbi_types import TorchTransform, VariationalDistribution
 from sbi.utils.typechecks import (
     is_nonnegative_int,
     is_positive_float,
@@ -334,61 +333,73 @@ def validate(self):
 @dataclass(frozen=True)
 class VIPosteriorParameters(PosteriorParameters):
     """
-    Parameters for initializing VIPosterior.
+    Parameters for VIPosterior, supporting both single-x and amortized VI.
 
     Fields:
-        q: Variational distribution, either string, `TransformedDistribution`, or a
-            `VIPosterior` object. This specifies a parametric class of distribution
-            over which the best possible posterior approximation is searched. For
-            string input, we currently support [nsf, scf, maf, mcf, gaussian,
-            gaussian_diag]. You can also specify your own variational family by
-            passing a pyro `TransformedDistribution`.
-            Additionally, we allow a `Callable`, which allows you the pass a
-            `builder` function, which if called returns a distribution. This may be
-            useful for setting the hyperparameters e.g. `num_transfroms` within the
-            `get_flow_builder` method specifying the number of transformations
-            within a normalizing flow. If q is already a `VIPosterior`, then the
-            arguments will be copied from it (relevant for multi-round training).
-        vi_method: This specifies the variational methods which are used to fit q to
-            the posterior. We currently support [rKL, fKL, IW, alpha]. Note that
-            some of the divergences are `mode seeking` i.e. they underestimate
-            variance and collapse on multimodal targets (`rKL`, `alpha` for alpha >
-            1) and some are `mass covering` i.e. they overestimate variance but
-            typically cover all modes (`fKL`, `IW`, `alpha` for alpha < 1).
-        parameters: List of parameters of the variational posterior. This is only
-            required for user-defined q i.e. if q does not have a `parameters`
-            attribute.
-        modules: List of modules of the variational posterior. This is only
-            required for user-defined q i.e. if q does not have a `modules`
-            attribute.
+        q: Variational distribution. Either a string specifying the flow type
+            [nsf, maf, naf, unaf, nice, sospf, gaussian, gaussian_diag], a
+            `TransformedDistribution`, a `VIPosterior` object, or a `Callable`
+            builder function. For amortized VI, use string flow types only.
+            If q is already a `VIPosterior`, arguments are copied from it
+            (relevant for multi-round training).
+        vi_method: Variational method for fitting q to the posterior. Options:
+            [rKL, fKL, IW, alpha]. Some are "mode seeking" (rKL, alpha > 1) and
+            some are "mass covering" (fKL, IW, alpha < 1). Currently only used
+            for single-x VI; amortized VI uses ELBO (rKL).
+        num_transforms: Number of transforms in the normalizing flow.
+        hidden_features: Hidden layer size in the flow networks.
+        z_score_theta: Method for z-scoring θ (the parameters being modeled).
+            One of "none", "independent", "structured". Use "structured" for
+            parameters with correlations.
+        z_score_x: Method for z-scoring x (the conditioning variable, amortized
+            VI only). One of "none", "independent", "structured". Use
+            "structured" for structured data like images.
+
+    Note:
+        For custom distributions that lack `parameters()` and `modules()` methods,
+        pass these via `VIPosterior.set_q(q, parameters=..., modules=...)` instead.
     """
 
     q: Union[
-        Literal["nsf", "scf", "maf", "mcf", "gaussian", "gaussian_diag"],
-        PyroTransformedDistribution,
+        Literal[
+            "nsf", "maf", "naf", "unaf", "nice", "sospf", "gaussian", "gaussian_diag"
+        ],
+        VariationalDistribution,
         "VIPosterior",
         Callable,
     ] = "maf"
     vi_method: Literal["rKL", "fKL", "IW", "alpha"] = "rKL"
-    parameters: Optional[Iterable] = None
-    modules: Optional[Iterable] = None
+    num_transforms: int = 5
+    hidden_features: int = 50
+    z_score_theta: Literal["none", "independent", "structured"] = "independent"
+    z_score_x: Literal["none", "independent", "structured"] = "independent"
 
     def validate(self):
         """Validate VIPosteriorParameters fields."""
-
-        valid_q = {"nsf", "scf", "maf", "mcf", "gaussian", "gaussian_diag"}
+        valid_q = {
+            "nsf",
+            "maf",
+            "naf",
+            "unaf",
+            "nice",
+            "sospf",
+            "gaussian",
+            "gaussian_diag",
+        }
 
         if isinstance(self.q, str) and self.q not in valid_q:
             raise ValueError(f"If `q` is a string, it must be one of {valid_q}")
         elif not isinstance(
-            self.q, (PyroTransformedDistribution, VIPosterior, Callable, str)
+            self.q, (VariationalDistribution, VIPosterior, Callable, str)
         ):
             raise TypeError(
-                "q must be either of typr PyroTransformedDistribution,"
-                " VIPosterioror or Callable"
+                "q must be either of type VariationalDistribution,"
+                " VIPosterior or Callable"
             )
 
-        if self.parameters is not None and not isinstance(self.parameters, Iterable):
-            raise TypeError("parameters must be iterable or None.")
-        if self.modules is not None and not isinstance(self.modules, Iterable):
-            raise TypeError("modules must be iterable or None.")
+        if self.num_transforms < 1:
+            raise ValueError(f"num_transforms must be >= 1, got {self.num_transforms}")
+        if self.hidden_features < 1:
+            raise ValueError(
+                f"hidden_features must be >= 1, got {self.hidden_features}"
+            )