feat(vfpe): opt-in composed standardization for scale-equivariant FMPE/NPSE (#1680)

sbi/inference/posteriors/vector_field_posterior.py

@@ -235,6 +235,18 @@ def sample(
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         is_iid = x.shape[0] > 1
+        # Composed standardization integrates the SDE/ODE in z-space; the iid and
+        # guidance score combinations evaluate the (theta-space) prior on the z-space
+        # state, which is not yet transformed. Guard explicitly instead of returning
+        # silently-wrong samples. (Single-observation sampling is fully supported.)
+        if self.vector_field_estimator._compose_standardization and (
+            is_iid or guidance_method is not None
+        ):
+            raise NotImplementedError(
+                "compose_standardization does not yet support iid (x with batch>1) or "
+                "guided sampling. Use a single observation, or disable "
+                "compose_standardization."
+            )
         self.potential_fn.set_x(
             x,
             x_is_iid=is_iid,
@@ -401,6 +413,13 @@ def _sample_via_diffusion(
                 "This may indicate numerical instability in the vector field."
             )
 
+        # Composed standardization (opt-in): the estimator samples in z-space.
+        # Unstandardize theta = shift + scale * z before returning so that
+        # rejection (within_support on the prior) and the final samples are in
+        # original theta space. No-op when the flag is off.
+        if self.vector_field_estimator._compose_standardization:
+            samples = self.vector_field_estimator.from_z(samples)
+
         return samples
 
     def sample_via_ode(
@@ -429,6 +448,13 @@ def sample_via_ode(
             torch.Size((num_samples,))
         )
 
+        # Composed standardization (opt-in): the ODE samples in z-space.
+        # Unstandardize theta = shift + scale * z before returning so that
+        # rejection (within_support on the prior) and the final samples are in
+        # original theta space. No-op when the flag is off.
+        if self.vector_field_estimator._compose_standardization:
+            samples = self.vector_field_estimator.from_z(samples)
+
         return samples
 
     def log_prob(
@@ -457,6 +483,15 @@ def log_prob(
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         is_iid = x.shape[0] > 1
+        # See sample(): the iid score combination evaluates the theta-space prior on the
+        # z-space state under composed standardization. Guard rather than return wrong
+        # values; single-observation log_prob is exact (affine Jacobian correction).
+        if self.vector_field_estimator._compose_standardization and is_iid:
+            raise NotImplementedError(
+                "compose_standardization does not yet support iid (x with batch>1) "
+                "log_prob. Use a single observation, or disable "
+                "compose_standardization."
+            )
         self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {}))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
@@ -516,6 +551,17 @@ def sample_batched(
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
         """
+        # Composed standardization integrates the SDE/ODE in z-space. Batched
+        # sampling (and the AutoGaussCorrectedScoreFn it feeds, which builds
+        # theta-space precisions) is not yet transformed for z-space and would
+        # return silently-wrong samples. Guard explicitly. (Single-observation
+        # sampling via sample() is fully supported.)
+        if self.vector_field_estimator._compose_standardization:
+            raise NotImplementedError(
+                "compose_standardization does not yet support sample_batched "
+                "(batched / multi-observation sampling). Use a single observation "
+                "via sample(), or disable compose_standardization."
+            )
         num_samples = torch.Size(sample_shape).numel()
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         condition_dim = len(self.vector_field_estimator.condition_shape)
@@ -647,6 +693,13 @@ def map(
         Returns:
             The MAP estimate.
         """
+        if self.vector_field_estimator._compose_standardization:
+            raise NotImplementedError(
+                "MAP is not yet supported with compose_standardization. "
+                "The potential gradient is computed in standardized z-space, so "
+                "gradient ascent in theta-space would be incorrect."
+            )
+
         if x is not None:
             raise ValueError(
                 "Passing `x` directly to `.map()` has been deprecated."

sbi/inference/potentials/vector_field_potential.py

@@ -118,6 +118,25 @@ def set_x(
                 `IIDScoreFunction`.
             ode_kwargs: Additional keyword arguments for the neural ODE.
         """
+        # Composed standardization runs the SDE/ODE in z-space; the iid and guided
+        # score combinations evaluate the (theta-space) prior on the z-space state,
+        # which is not transformed. Guard at this chokepoint so the broken combination
+        # cannot be reached even via direct potential use (not just via
+        # VectorFieldPosterior.sample/log_prob).
+        if bool(
+            getattr(self.vector_field_estimator, "_compose_standardization", False)
+        ):
+            if x_is_iid:
+                raise NotImplementedError(
+                    "compose_standardization does not yet support iid (x with "
+                    "batch>1). Use a single observation, or disable "
+                    "compose_standardization."
+                )
+            if guidance_method is not None:
+                raise NotImplementedError(
+                    "compose_standardization does not yet support guided sampling. "
+                    "Disable guidance, or disable compose_standardization."
+                )
         super().set_x(x_o, x_is_iid)
         self.iid_method = iid_method or self.iid_method
         self.iid_params = iid_params
@@ -155,8 +174,40 @@ def __call__(
         )
         self.vector_field_estimator.eval()
 
+        # Composed standardization (opt-in): the flow operates in z-space. Map the
+        # incoming theta -> z for the flow, and correct the log-prob with the
+        # affine Jacobian: log p_theta(theta) = log p_z(z) - sum(log scale).
+        # within_support below still uses the ORIGINAL theta. No-op when off.
+        compose = self.vector_field_estimator._compose_standardization
+        if compose:
+            flow_input = self.vector_field_estimator.to_z(theta_density_estimator)
+            log_abs_det = self.vector_field_estimator.log_abs_det()
+        else:
+            flow_input = theta_density_estimator
+
         with torch.set_grad_enabled(track_gradients):
             if self.x_is_iid:
+                if compose:
+                    # Backstop for a directly-forced `_x_is_iid`: compose + iid is
+                    # unreachable through the public API (guarded in `set_x`).
+                    # Reaching here means `_x_is_iid` was set directly, and continuing
+                    # would return a partially-incorrect log-prob -- the affine
+                    # Jacobian is applied but the prior is not yet expressed in
+                    # z-space. Raise rather than return quiet wrong numbers, matching
+                    # the set_x guards.
+                    raise NotImplementedError(
+                        "compose_standardization does not yet support iid (x with "
+                        "batch>1) log_prob. This path is guarded in set_x; reaching "
+                        "it means _x_is_iid was set directly, which would return a "
+                        "partially-incorrect log-prob (affine Jacobian applied, prior "
+                        "not yet expressed in z). Use a single observation, or disable "
+                        "compose_standardization."
+                    )
+                    # future iid support: lift this backstop AND the set_x guard
+                    # once the prior is transformed to z (see
+                    # followup/compose-iid-guidance-map). Retained z-space affine-
+                    # Jacobian correction to restore then:
+                    #   iid_posteriors_prob -= num_iid * log_abs_det
                 assert self.prior is not None, (
                     "Prior is required for evaluating log_prob with iid observations."
                 )
@@ -166,10 +217,7 @@ def __call__(
                 num_iid = self.x_o.shape[0]  # number of iid samples
                 iid_posteriors_prob = torch.sum(
                     torch.stack(
-                        [
-                            flow.log_prob(theta_density_estimator).squeeze(-1)
-                            for flow in self.flows
-                        ],
+                        [flow.log_prob(flow_input).squeeze(-1) for flow in self.flows],
                         dim=0,
                     ),
                     dim=0,
@@ -180,7 +228,9 @@ def __call__(
                     theta_density_estimator
                 ).squeeze(-1)
             else:
-                log_probs = self.flow.log_prob(theta_density_estimator).squeeze(-1)
+                log_probs = self.flow.log_prob(flow_input).squeeze(-1)
+                if compose:
+                    log_probs = log_probs - log_abs_det
             # Force probability to be zero outside prior support.
             in_prior_support = within_support(self.prior, theta)
 
@@ -199,6 +249,11 @@ def gradient(
     ) -> Tensor:
         r"""Returns the potential function gradient for score-based methods.
 
+        Coordinate convention: under compose_standardization the SDE sampler calls
+        gradient() on a z-space state, while MAP (which would need theta-space
+        gradients) is guarded; so gradient() always operates in the estimator's
+        current coordinate.
+
         Args:
             theta: The parameters at which to evaluate the potential gradient.
             time: The diffusion time. If None, then `t_min` of the

sbi/neural_nets/estimators/base.py

@@ -386,6 +386,65 @@ def __init__(
             embedding_net if embedding_net is not None else nn.Identity()
         )
 
+        # ---- Composed standardization (opt-in; default OFF) ----
+        # When enabled, the estimator stays PURE z-space and an invertible per-dim
+        # affine transform theta = shift + scale * z is applied ONLY at the
+        # boundaries (loss input, sample output, log_prob input). Buffers default
+        # to the identity transform so the flag-OFF path is byte-identical.
+        self.register_buffer(
+            "_theta_shift", torch.zeros(1, *self.input_shape, dtype=torch.float32)
+        )
+        self.register_buffer(
+            "_theta_scale", torch.ones(1, *self.input_shape, dtype=torch.float32)
+        )
+        # Persisted as a buffer (not a plain attribute) so that compose mode survives
+        # save/load: a checkpoint trained with composition reloads as compose-enabled.
+        self.register_buffer(
+            "_compose_standardization", torch.tensor(False), persistent=True
+        )
+        # Plain Python bool mirror — NOT a buffer, NOT persisted. Allows hot-path
+        # guards to short-circuit before touching the tensor buffer (a tensor
+        # truth-test forces a host/device sync on CUDA). Kept in sync by
+        # _wire_compose (build path) and _load_from_state_dict (checkpoint path).
+        self._compose_enabled: bool = False
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        r"""Gracefully load checkpoints from before composed standardization existed.
+
+        The three compose buffers (``_theta_shift``/``_theta_scale``/
+        ``_compose_standardization``) are handled atomically:
+
+        * NONE present (legacy pre-compose checkpoint): inject the current
+          identity / disabled defaults so the estimator loads as composition-OFF
+          (its original behavior).
+        * ALL present (compose-aware checkpoint): load normally.
+        * PARTIAL (some present, some missing): raise rather than silently inject
+          identity defaults — a checkpoint with ``_compose_standardization=True``
+          but a missing ``_theta_shift``/``_theta_scale`` would otherwise load a
+          silently-wrong identity affine.
+        """
+        compose_names = ("_theta_shift", "_theta_scale", "_compose_standardization")
+        present = [n for n in compose_names if prefix + n in state_dict]
+        if present and len(present) != len(compose_names):
+            missing = [n for n in compose_names if prefix + n not in state_dict]
+            raise RuntimeError(
+                "Partial composed-standardization checkpoint: present "
+                f"{[prefix + n for n in present]} but missing "
+                f"{[prefix + n for n in missing]}. Refusing to inject identity "
+                "defaults for the missing buffers (that would silently produce a "
+                "wrong affine). Load a checkpoint with either all three compose "
+                "buffers or none of them."
+            )
+        if not present:
+            # Legacy checkpoint: inject identity / disabled defaults (compose-OFF).
+            for name in compose_names:
+                if hasattr(self, name):
+                    state_dict[prefix + name] = getattr(self, name).clone()
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        # Sync the plain-Python mirror AFTER the buffer is loaded so the hot-path
+        # short-circuit in _check_compose_internal_stats_unit stays correct.
+        self._compose_enabled = bool(self._compose_standardization)
+
     @property
     def embedding_net(self) -> nn.Module:
         r"""Return the embedding network if it exists."""
@@ -420,6 +479,59 @@ def std_base(self) -> Tensor:
         (the initial noise at time t=T)."""
         return self._std_base
 
+    # -------------------------- COMPOSED STANDARDIZATION HELPERS ----------------
+
+    def to_z(self, theta: Tensor) -> Tensor:
+        r"""Map original-space parameters to standardized z-space.
+
+        ``z = (theta - shift) / scale``. Identity when composed standardization
+        is disabled (shift=0, scale=1).
+        """
+        return (theta - self._theta_shift) / self._theta_scale
+
+    def from_z(self, z: Tensor) -> Tensor:
+        r"""Map standardized z-space parameters back to original space.
+
+        ``theta = shift + scale * z``. Identity when composed standardization
+        is disabled (shift=0, scale=1).
+        """
+        return self._theta_shift + self._theta_scale * z
+
+    def log_abs_det(self) -> Tensor:
+        r"""Log absolute determinant of the affine z->theta Jacobian.
+
+        ``sum(log scale)`` over the input dimensions. Used to correct the
+        flow log-prob (which is computed in z-space) back to theta-space:
+        ``log p_theta(theta) = log p_z(z) - sum(log scale)``.
+        """
+        return torch.log(self._theta_scale).sum()
+
+    def _check_compose_internal_stats_unit(self) -> None:
+        r"""Enforce the composed-standardization invariant: unit internal stats.
+
+        When composition is enabled the estimator must stay PURE z-space, which
+        requires the internal input-norm stats ``mean_0``/``std_0`` to be exactly
+        ``0``/``1`` (the network sees unit-z input; the boundary affine carries the
+        original-space scale). The build path forces this, but a manually assembled
+        estimator or a tampered checkpoint could re-enable composition with
+        non-unit stats, silently producing an inconsistent model. Reject it at
+        use-time. Short-circuits cheaply when composition is off (the common path).
+        """
+        # Short-circuit on the plain Python bool mirror (no tensor truth-test,
+        # no host/device sync). The compose-OFF path (the common case, called on
+        # every ODE/SDE step) pays zero tensor access.
+        if not self._compose_enabled:
+            return
+        if not ((self.mean_0 == 0).all() and (self.std_0 == 1).all()):
+            raise ValueError(
+                "compose_standardization is enabled but the internal input-norm "
+                "stats mean_0/std_0 are not unit (0/1). Under composition the "
+                "estimator must stay pure z-space (the boundary affine carries the "
+                "original-space scale). This estimator is internally inconsistent "
+                "— rebuild it via build_vector_field_estimator(..., "
+                "compose_standardization=True), which forces unit stats."
+            )
+
     # -------------------------- ODE METHODS --------------------------
 
     @abstractmethod

sbi/neural_nets/estimators/flowmatching_estimator.py

@@ -109,6 +109,33 @@ def __init__(
         self.register_buffer("mean_0", mean_0_tensor)
         self.register_buffer("std_0", std_0_tensor)
 
+    def _check_compose_baseline_compatible(self) -> None:
+        """Guard the unsupported gaussian_baseline + composed-standardization pair.
+
+        ``gaussian_baseline`` derives the velocity from the (original-space) data
+        statistics ``mean_0``/``std_0``; composed standardization instead maps
+        ``theta -> z`` so the network sees standardized inputs. The two are
+        mutually exclusive (see also the build-time guard in ``_wire_compose``).
+
+        The public build path refuses the pair up front; this runtime check is the
+        backstop for estimators assembled manually, or loaded from a checkpoint
+        that predates the build-time guard, where ``_compose_standardization`` is
+        re-activated alongside ``gaussian_baseline``.
+        """
+        # ``gaussian_baseline`` is a plain Python bool, so checking it first lets
+        # the common path (and every ODE step) short-circuit before touching the
+        # ``_compose_standardization`` tensor buffer (whose truth test would force a
+        # device sync on CUDA).
+        if self.gaussian_baseline and bool(self._compose_standardization):
+            raise ValueError(
+                "gaussian_baseline and composed standardization cannot be used "
+                "together: the analytical baseline velocity is derived from "
+                "mean_0/std_0 in the original data space, while composed "
+                "standardization feeds the network standardized z (theta -> z). "
+                "Enable only one (set gaussian_baseline=False to use composed "
+                "standardization)."
+            )
+
     def _get_time_dependent_stats(self, time: Tensor) -> Tuple[Tensor, Tensor]:
         """Compute time-dependent mean and std for z-scoring.
 
@@ -209,6 +236,9 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
         Returns:
             The estimated vector field in original space.
         """
+        self._check_compose_baseline_compatible()
+        self._check_compose_internal_stats_unit()
+
         # Continue with standard processing (broadcast shapes etc.)
         batch_shape_input = input.shape[: -len(self.input_shape)]
         batch_shape_cond = condition.shape[: -len(self.condition_shape)]
@@ -283,6 +313,15 @@ def loss(
         Returns:
             Loss value.
         """
+        self._check_compose_baseline_compatible()
+        self._check_compose_internal_stats_unit()
+
+        # Composed standardization (opt-in): the estimator is trained PURELY in
+        # z-space. Standardize theta -> z at the very top; everything below is
+        # unchanged. When the flag is off, this is a no-op (shift=0, scale=1).
+        if self._compose_standardization:
+            input = self.to_z(input)
+
         # Randomly sample time steps
         if times is None:
             times = torch.rand(input.shape[:-1], device=input.device, dtype=input.dtype)