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

[test1card]

What does this PR do?

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

This adds an opt-in per-dimension composed standardization
(compose_standardization=True). The estimator is trained and sampled in a
standardized z-space, with an invertible affine map theta = shift + scale * z
composed at the boundaries:

• loss input standardized;
• sample output unstandardized back to theta;
• log_prob corrected by the affine Jacobian (-sum log scale);
• prior-support rejection performed in original theta-space.

Per-dimension scaling handles heterogeneous parameter scales (mixed O(1) and
O(1e-6) dimensions). The flag defaults to off; with it off, behavior is
byte-identical to current sbi, and legacy checkpoints load via identity buffers.

MAP raises NotImplementedError in compose mode: the score potential is computed
in z-space, so gradient ascent in theta-space would require a chain-rule
rescaling of the gradient. This is left as a documented follow-up; single-
observation sampling and log_prob are fully supported.

Changes

• neural_nets/estimators/base.py: standardization buffers and helpers, legacy-checkpoint shim
• neural_nets/estimators/flowmatching_estimator.py, score_estimator.py: standardize loss input
• inference/posteriors/vector_field_posterior.py: unstandardize samples; guard MAP/iid/guided under compose
• inference/potentials/vector_field_potential.py: affine Jacobian in log_prob
• neural_nets/net_builders/vector_field_nets.py: opt-in builder kwarg
• tests/test_scale_equivariance.py: homogeneous + heterogeneous equivariance, opt-in-off regression

Validation

Run on an editable install:

• FMPE and NPSE recover equivariance at scale s = 1e-5 (both SDE and ODE paths),
  covered by tests/test_scale_equivariance.py (homogeneous, heterogeneous).
• Full non-slow test suite: 6062 passed, 0 failed.

Does this close any issues?

Relates to #1680. That issue was closed by #1681, but #1681 only touched the
z_score_x off-branch and did not address the default-path scale-equivariance of
FMPE/NPSE; this PR completes that case.

Checklist

• Tests added (scale-equivariance, heterogeneous, opt-in-off regression)
• Slow tests marked
• Code commented where non-obvious
• Default behavior unchanged (flag off = byte-identical)

[janfb]

Thanks a lot for this, really nice first contribution @test1card!
I just had a first look and I think you caught a subtle problem. I appreciate that you backed it with equivariance tests and a clear opt-in, default-off design. I could reproduce both the failure and the recovery locally, so the core idea clearly works.

One pointer before a deeper review: please take a look at #1752 which is more recent than #1681 and changed how FlowMatchingEstimator handles z-scoring. On the current default path, FMPE already z-scores θ by per-dim data statistics (via the mean_0/std_0 buffers + time-dependent normalization, plus an optional Gaussian baseline). That changes the framing a bit and it means your new _theta_shift/_theta_scale buffers likely overlap with statistics the estimator already stores. Here, I think it would be nice to explore whether we can combine the two rather than maintain a second, parallel standardization mechanism.

A few other things I'll want to look at more closely, at a high level:

• the interaction with Re-enable time-dependent z-scoring for Flow Matching #1752's gaussian_baseline mode (currently untested together with the flag),
• the scope: this is single-observation only right now (MAP / iid / guided sampling raise NotImplementedError), which is a fine starting point but worth confirming as intended,
• a couple of edge cases in the sampling/loading paths.

I won't be able to do the detailed line-by-line review until July, sorry for that. In the meantime, looking at #1752 and thinking about how the new standardization could reuse the existing mean_0/std_0 machinery would be the most useful direction. Thanks again for tackling this!

[codecov]

Codecov Report

❌ Patch coverage is 99.03846% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 87.67%. Comparing base (9cf0d9e) to head (2f717aa).
⚠️ Report is 22 commits behind head on main.

Files with missing lines	Patch %	Lines
sbi/neural_nets/net_builders/vector_field_nets.py	97.29%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1884      +/-   ##
==========================================
- Coverage   87.89%   87.67%   -0.22%     
==========================================
  Files         143      143              
  Lines       13353    14488    +1135     
==========================================
+ Hits        11736    12703     +967     
- Misses       1617     1785     +168     

Flag	Coverage Δ
fast	81.63% <99.03%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/inference/posteriors/vector_field_posterior.py	80.50% <100.00%> (+2.27%)	⬆️
sbi/inference/potentials/vector_field_potential.py	91.91% <100.00%> (+0.92%)	⬆️
sbi/neural_nets/estimators/base.py	80.34% <100.00%> (+3.63%)	⬆️
...i/neural_nets/estimators/flowmatching_estimator.py	98.33% <100.00%> (+0.13%)	⬆️
sbi/neural_nets/estimators/score_estimator.py	92.01% <100.00%> (+0.12%)	⬆️
sbi/neural_nets/factory.py	95.06% <ø> (ø)
sbi/neural_nets/net_builders/vector_field_nets.py	93.97% <97.29%> (-0.04%)	⬇️

... and 10 files with indirect coverage changes

[test1card]

Thanks! I pushed an updated version of the branch, so the PR now shows the hardened commit rather than the earlier one. The #1752 pointer really helped, it turns out to be the same question as whether the two standardizations are redundant.

As far as i understood it correctly, #1752's z-scoring works on the network internally as it z-scores the input and rescales the velocity output. But forward() still returns velocity in original θ-coordinates, and the ODE/SDE integrates there from a base whose scale is set by the noise process, but not the data (unit for the FMPE flow, roughly σ_max for the score SDE),the network sees well-scaled inputs, but the base and the transport path are still stuck at that data-independent scale while the posterior sits down at scale s. Nothing pulls the measure or the path down to where the data is. That's the whole reason the collapse is still there even with #1752 in the base, it's conditioning the network one layer away from where the problem actually lives. Also, the stats come from one computation (_compute_theta_standardization). Under compose the estimator lives in z the whole time, mean_0/std_0 are pinned to 0/1 (build-time invariant, runtime backstop), and the boundary affine (_theta_shift/_theta_scale) is the thing holding the scale. The data is already unit in z, so #1752's data part collapses to unit and whatever time-dependent factor is left doesn't care about scale. One active θ-standardization, not two. Those new buffers are the boundary transform, not a second copy of mean_0/std_0 sitting next to the first. gaussian_baseline and compose can't both be on now (guarded at build, backstopped at runtime), because the baseline reads mean_0/std_0 as genuine data stats and compose has already turned them into unit-z. Single-obs is deliberate for this PR. iid, guided, and MAP raise NotImplementedError instead of handing back something wrong and quiet.

The update also adds the missing coverage tests and a runtime backstop on the iid log_prob path, plus the edge-case guards on the sampling/loading paths. The bigger scope, opening up the affine iid/guidance/MAP paths by pushing the prior into z-space, is on a separate follow-up branch, not yet finished. I am new to this as you've noticed, so should that follow-up go as its own PR after this one lands, or fold it in here?

And i wanted to say sorry because I leaned on LLM assistance a fair amount here, which the contributing guide asks contributors to flag. I went through and tested all of it myself, so if something's wrong it's on me.

No rush on the line by line!

[test1card]

Pushed a small follow-up, the Apache license headers on the four new test files.

I also noticed that the standardization path that with compose_standardization=True, a user-passed z_score_x='structured' is silently overridden as the composed affine always uses per-dimension statistics, and the internal input-norm is pinned to unit. After a rebase on current main, 'transform_to_unconstrained' will raise early via #1886, but 'structured' still falls through quietly. I have a small guard and tests ready that raises a ValueError in that case, but it changes API behavior,so I didn't want to push it into the diff mid-review.

Could you please recommend what to do with it? I know it has to be fixed, but the scope creep worries me.