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

FMPE and NPSE amortized posteriors are not equivariant under an exact reparameterization
theta -> c*theta when parameters are far from O(1) scale: the flow/SDE integrate against a
unit base, so calibration collapses despite default z_score_theta. PR #1681 only touched the
z_score_x off-branch (no-op on the default path).

This adds an OPT-IN per-dim composed standardization (compose_standardization=True): the
estimator is trained and sampled in standardized z-space with an invertible affine transform
theta = shift + scale * z composed at the boundaries (loss input standardized; sample output
unstandardized; log_prob corrected by the affine Jacobian; prior-support rejection in original
theta). Per-dim handles HETEROGENEOUS scales (mixed O(1) and O(1e-6) dims). Default OFF =>
behavior byte-identical to current sbi.

MAP raises NotImplementedError in compose mode (the score potential is z-space; documented
limitation). Files: estimators/base.py (buffers+helpers), flowmatching_estimator.py +
score_estimator.py (loss standardize), vector_field_posterior.py (unstandardize samples + MAP
guard), vector_field_potential.py (log_prob Jacobian), net_builders/vector_field_nets.py
(opt-in builder kwarg). Adds tests/test_scale_equivariance.py (homogeneous + heterogeneous +
opt-in-off regression).

Validated on editable clone: FMPE/NPSE recover equivariance at s=1e-5 (SDE+ODE); canonical
BoxUniform two_moons+SLCP recover (miscal ~0.02 vs default ~0.48); full non-slow pytest 6062
passed / 0 failed.

Author: test1card

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))
@@ -647,6 +682,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,6 +174,17 @@ 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:
                 assert self.prior is not None, (
@@ -166,21 +196,29 @@ 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,
                 )
+                if compose:
+                    # NOTE: currently UNREACHABLE -- composed standardization
+                    # guards iid in `set_x` (the iid score combination would
+                    # evaluate the theta-space prior on a z-space state). Retained
+                    # for future iid support: each of the num_iid z-space flow
+                    # log-probs is corrected by the affine Jacobian. (The
+                    # prior-coordinate transform still needs doing before the guard
+                    # is lifted.)
+                    iid_posteriors_prob = iid_posteriors_prob - num_iid * log_abs_det
                 # Apply the adjustment for iid observations i.e. we have to subtract
                 # (num_iid-1) times the log prior.
                 log_probs = iid_posteriors_prob - (num_iid - 1) * self.prior.log_prob(
                     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)
 

sbi/neural_nets/estimators/base.py

@@ -386,6 +386,38 @@ 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
+        )
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        r"""Gracefully load checkpoints from before composed standardization existed.
+
+        Old checkpoints lack the ``_theta_shift``/``_theta_scale``/
+        ``_compose_standardization`` buffers. Rather than failing a strict load, inject
+        the current (identity / disabled) defaults for any missing compose buffer so the
+        legacy estimator loads as composition-OFF (its original behavior). Newer
+        checkpoints carry the buffers and load normally.
+        """
+        for name in ("_theta_shift", "_theta_scale", "_compose_standardization"):
+            key = prefix + name
+            if key not in state_dict and hasattr(self, name):
+                state_dict[key] = getattr(self, name).clone()
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     @property
     def embedding_net(self) -> nn.Module:
         r"""Return the embedding network if it exists."""
@@ -420,6 +452,33 @@ 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()
+
     # -------------------------- ODE METHODS --------------------------
 
     @abstractmethod

sbi/neural_nets/estimators/flowmatching_estimator.py

@@ -283,6 +283,12 @@ def loss(
         Returns:
             Loss value.
         """
+        # 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)

sbi/neural_nets/estimators/score_estimator.py

@@ -248,6 +248,12 @@ def loss(
             MSE between target score and network output, scaled by the weight function.
 
         """
+        # 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)
+
         # Sample times from the Markov chain, use batch dimension
         if times is None:
             times = self.train_schedule(input.shape[0])

sbi/neural_nets/net_builders/vector_field_nets.py

@@ -41,6 +41,11 @@ class _VectorFieldBaseConfig(_EstimatorConfigBase):
     sinusoidal_max_freq: Optional[float] = None
     fourier_scale: Optional[float] = None
 
+    # Composed standardization (opt-in scale-equivariance for FMPE / NPSE).
+    # When True, the estimator is trained/sampled in standardized z-space with an
+    # invertible per-dim affine transform applied only at the boundaries.
+    compose_standardization: Optional[bool] = None
+
     # MLP-specific
     layer_norm: Optional[bool] = None
     skip_connections: Optional[bool] = None
@@ -126,6 +131,7 @@ def build_vector_field_estimator(
         VectorFieldNet,
     ] = "mlp",
     gaussian_baseline: bool = False,
+    compose_standardization: bool = False,
     **kwargs,
 ) -> Union[FlowMatchingEstimator, ConditionalScoreEstimator]:
     """Builds a vector field estimator (flow matching or score matching) with the given
@@ -151,6 +157,12 @@ def build_vector_field_estimator(
         gaussian_baseline: If True, use analytical Gaussian baseline velocity
             derived from Bayes' rule. The network then only learns the residual.
             Only used when estimator_type="flow". Defaults to False.
+        compose_standardization: If True (opt-in), compose an invertible per-dim
+            affine standardization with the flow so the estimator becomes
+            scale-equivariant. The estimator is trained and sampled in
+            standardized z-space (per-dim mean/std computed from training theta),
+            and theta = shift + scale * z is applied only at the boundaries.
+            Defaults to False (behavior identical to current sbi).
         **kwargs: Additional arguments forwarded to the estimator and network
             constructors.  Valid keys are defined by ``ScoreEstimatorConfig``
             and ``FlowEstimatorConfig``; validation happens in the upstream
@@ -225,22 +237,51 @@ def build_vector_field_estimator(
     else:
         mean_0, std_0 = 0, 1
 
+    # Composed standardization (opt-in): compute the per-dim affine transform
+    # theta = shift + scale * z from the training theta. The estimator stays PURE
+    # z-space, so we force the internal input-norm stats (mean_0/std_0) to 0/1 to
+    # be consistent with unit-z input. The base distribution is left in z-space
+    # (FMPE 0/1, VE sigma_max); we do NOT scale the base by `scale`.
+    compose_shift = None
+    compose_scale = None
+    if compose_standardization:
+        # Per-dim standardization (independent), matching the validated reference.
+        compose_shift, compose_scale = z_standardization(batch_x, structured_dims=False)
+        compose_scale = compose_scale.clamp_min(1e-20)
+        # Force the internal input-norm stats to unit (the estimator now sees z).
+        # Use PER-DIM tensors (not scalars) so the VE base buffers (mean_t/std_t,
+        # computed via broadcast_to) stay contiguous and load_state_dict-safe.
+        mean_0 = torch.zeros_like(compose_shift)
+        std_0 = torch.ones_like(compose_scale)
+
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
     embedding_net_y = (
         nn.Sequential(standardizing_net(batch_y, structured_y), embedding_net)
         if z_score_y_bool
         else embedding_net
     )
 
+    def _wire_compose(estimator):
+        """Wire the per-dim affine standardization onto the estimator (opt-in)."""
+        if compose_shift is not None and compose_scale is not None:
+            shift = compose_shift.reshape(1, *estimator.input_shape).float()
+            scale = compose_scale.reshape(1, *estimator.input_shape).float()
+            estimator._theta_shift.copy_(shift)
+            estimator._theta_scale.copy_(scale)
+            estimator._compose_standardization.fill_(True)
+        return estimator
+
     if estimator_type == "flow":
-        return FlowMatchingEstimator(
-            net=vectorfield_net,
-            input_shape=batch_x[0].shape,
-            condition_shape=batch_y[0].shape,
-            embedding_net=embedding_net_y,
-            mean_0=mean_0,
-            std_0=std_0,
-            gaussian_baseline=gaussian_baseline,
+        return _wire_compose(
+            FlowMatchingEstimator(
+                net=vectorfield_net,
+                input_shape=batch_x[0].shape,
+                condition_shape=batch_y[0].shape,
+                embedding_net=embedding_net_y,
+                mean_0=mean_0,
+                std_0=std_0,
+                gaussian_baseline=gaussian_baseline,
+            )
         )
     elif estimator_type == "score":
         # Choose the appropriate score estimator based on SDE type
@@ -272,14 +313,16 @@ def build_vector_field_estimator(
             vp_keys = ["beta_min", "beta_max"]
             estimator_kwargs = {k: kwargs[k] for k in vp_keys if k in kwargs}
 
-        return estimator_cls(
-            net=vectorfield_net,
-            input_shape=batch_x[0].shape,
-            condition_shape=batch_y[0].shape,
-            embedding_net=embedding_net_y,
-            mean_0=mean_0,
-            std_0=std_0,
-            **estimator_kwargs,
+        return _wire_compose(
+            estimator_cls(
+                net=vectorfield_net,
+                input_shape=batch_x[0].shape,
+                condition_shape=batch_y[0].shape,
+                embedding_net=embedding_net_y,
+                mean_0=mean_0,
+                std_0=std_0,
+                **estimator_kwargs,
+            )
         )
     else:
         raise ValueError(f"Unknown estimator type: {estimator_type}")