Re-enable time-dependent z-scoring for Flow Matching

[satwiksps]

Description

This PR re-introduces z-scoring for Flow Matching estimators using a time-dependent normalization approach and adds a Gaussian Baseline for improved training stability.

As discussed in #1623, standard z-scoring is problematic because the network input evolves from data to noise. This implementation provides two distinct normalization modes to handle this evolution while maintaining training stability.

Corrected Normalization Statistics:
Since we define $t=0$ as Data and $t=1$ as Noise, the statistics are handled based on the chosen mode:

1. Gaussian Baseline (gaussian_baseline=True): Normalizes inputs to $N(0, 1)$ across the entire path. The drift signal is handled by the hard-coded affine baseline.

   $$\mu_t = (1 - t) \cdot \mu_{data}$$
   $$\sigma_t = \sqrt{(1 - t)^2 \sigma_{data}^2 + t^2}$$

2. Variance Only (gaussian_baseline=False): Normalizes variance while preserving the raw data location at $t=0$. This ensures the network can still learn the drift signal when no baseline is used.

   $$\mu_t = t \cdot \mu_{data}$$
   $$\sigma_t = \sqrt{t^2 \sigma_{data}^2 + (1 - t)^2}$$

Gaussian Baseline:
We implemented an affine vector field baseline (enabled by default). The network now learns the residual vector field with respect to the optimal Gaussian probability path, significantly improving convergence on shifted datasets.

Related Issues/PRs

• Closes Add back z-scoreing for flow matching #1623

Changes

1. sbi/neural_nets/net_builders/vector_field_nets.py: Updated build_vector_field_estimator to calculate training data statistics, accept the gaussian_baseline flag, and pass them to the estimator.
2. sbi/neural_nets/estimators/flowmatching_estimator.py:
   • Buffer Management: Registered mean_1 and std_1 as buffers and expanded them to match input_shape to ensure compatibility with multi-dimensional data in CI.
   • Split Logic: Implemented the split logic in forward() to support both Gaussian Baseline (residual learning) and Variance-only (signal preserving) modes.
   • Numerical Stability: Added a small epsilon (1e-5) to variance calculations to prevent division-by-zero errors.
3. tests/linearGaussian_vector_field_test.py:
   • Added test_fmpe_time_dependent_z_scoring_integration: Verifies statistics population, buffer registration, and forward pass shapes.
   • Added test_fmpe_shifted_data_gaussian_baseline: Verifies that the Gaussian Baseline outperforms variance-only scaling on shifted data ($U[95, 105]$) with robust simulation counts ($N=2000$).

Verification

Verification

• Shifted Data Benchmark: Confirmed that gaussian_baseline=True achieves lower validation loss and faster convergence than variance-only scaling on a shifted 1D prior ($U[95, 105]$).
• Integration Tests: All new tests pass, confirming correct buffer registration and stability with z_score_x='independent'.
• Benchmarks: I ran the sbi benchmarks locally (pytest --bm --bm-mode fmpe) to check for stability and performance. All 12 tests passed successfully.

[codecov]

Codecov Report

❌ Patch coverage is 75.00000% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.82%. Comparing base (42d89f3) to head (8e37f05).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...i/neural_nets/estimators/flowmatching_estimator.py	74.07%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1752      +/-   ##
==========================================
- Coverage   87.88%   87.82%   -0.06%     
==========================================
  Files         140      140              
  Lines       12726    12777      +51     
==========================================
+ Hits        11184    11222      +38     
- Misses       1542     1555      +13     

Flag	Coverage Δ
fast	82.71% <75.00%> (?)

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

Files with missing lines	Coverage Δ
sbi/neural_nets/factory.py	95.06% <ø> (+1.08%)	⬆️
sbi/neural_nets/net_builders/vector_field_nets.py	94.00% <100.00%> (+0.03%)	⬆️
...i/neural_nets/estimators/flowmatching_estimator.py	85.58% <74.07%> (-11.09%)	⬇️

[satwiksps]

It seems tests/torchutils_test.py::TorchUtilsTest::test_searchsorted is consistently failing in the CI with an execnet.gateway_base.DumpError.

Since this failure is in torchutils_test.py (which I haven't touched) and appears to be a serialization issue with pytest-xdist masking a local assertion error, I believe it is unrelated to my changes in flowmatching_estimator.py ?

The actual Flow Matching benchmarks and integration tests for this PR passed successfully though

[janfb]

It seems tests/torchutils_test.py::TorchUtilsTest::test_searchsorted is consistently failing in the CI with an execnet.gateway_base.DumpError.

Since this failure is in torchutils_test.py (which I haven't touched) and appears to be a serialization issue with pytest-xdist masking a local assertion error, I believe it is unrelated to my changes in flowmatching_estimator.py ?

The actual Flow Matching benchmarks and integration tests for this PR passed successfully though

It seems tests/torchutils_test.py::TorchUtilsTest::test_searchsorted is consistently failing in the CI with an execnet.gateway_base.DumpError.

Since this failure is in torchutils_test.py (which I haven't touched) and appears to be a serialization issue with pytest-xdist masking a local assertion error, I believe it is unrelated to my changes in flowmatching_estimator.py ?

The actual Flow Matching benchmarks and integration tests for this PR passed successfully though

Yes, this is unrelated and popped up here by chance or because of an unrelated change in a downstream package. I pushed a fix to this branch ✅

[janfb]

Thanks for working on this @satwiksps !

Overall, this looks exactly right. However, after reviewing the code and tracing through the flow matching implementation, I believe the z-scoring formula is inverted relative to the interpolation convention (quite confusing!)

The interpolation in the loss function is:

theta_t = (1 - t) * theta_data + t * theta_noise

So the expected input mean at each time is:

E[θ_t] = (1-t) * mean_data + t * 0 = (1-t) * mean_data

Current PR formula:

• mu_t = t * mean_1
• var_t = (t * std_1)² + (1 - t)²

This gives mu_t = 0 at t=0 and mu_t = mean_data at t=1 — exactly backwards.

Correct formula should be:

• mu_t = (1 - t) * mean_1
• var_t = ((1 - t) * std_1)² + t²

The formula only matches at t=0.5 and is maximally wrong at the boundaries.

Note on zuko's sampling: I had to dig a bit but in zuko, NormalizingFlow.sample() uses transform.inv() which integrates backward (t1→t0), so training and sampling conventions do align — the issue is purely the z-scoring formula.

To verify this, I suggest the following test: The standard linear Gaussian test, but with uniform prior between 95 and 100, and with data x_o centered at 100 (far from N(0,1)). With the inverted formula, C2ST should degrades significantly compared to no z-scoring and it should be fixed (c2st close 0.5) with the correct formula.

Can you confirm this (maybe I got confused with the integration directions after all)?

[manuelgloeckler]

Hey @satwiksps !

Thanks for the contribution! I checked with main and as of now it does I guess on average perform very similar if not a bit worse than before (although, I think thats mostly fine i.e. these tasks).

I wonder if it would make sense to improve the "preconditioning" a bit more (see comments).

[janfb]

Hey @satwiksps !

Thanks for the contribution! I checked with main and as of now it does I guess on average perform very similar if not a bit worse than before (although, I think thats mostly fine i.e. these tasks).

I wonder if it would make sense to improve the "preconditioning" a bit more (see comments).

Thanks for adding the comparison to main. What could happen here is that the benchmarking tasks are not discriminative w.r.t. to z-scoring, no? I.e., we need a task that benefits from z-scoring.

[janfb]

Alright, I looked at it again and I realized that my proposal was actually incorrect. The formulas I proposed would result in total normalization, i.e., "independent" z-scoring, where all time steps have equal zero mean after z-scoring and we lose valuable time-depenedent information - sorry @satwiksps , your formulas where actually correct!

What Manuel proposed is great, we z-score with respect to the Gaussian baseline, e.g., what one would expect when the posterior is actually Gaussian. Then the flow matching network only has to learn the residual from this ideal baseline (please correct me @manuelgloeckler if this intuition is inaccurate).

I tested this locally with the following setup:

• Prior: BoxUniform([95, 105]), x_o=100
• Simulator: x = theta + 0.5 * noise
• Reference posterior: N(x_o, 0.5²I)
• 3000 simulations.

Results:

Formula	C2ST	Description
Gaussian	0.631	Gaussian baseline + residual learning
var_only	0.772	Variance scaling only
pr	0.774	PR's time-dependent z-scoring
static	0.796	Static mean subtraction
none	0.865	No z-scoring
independent z-scoring	0.922	"Correct" mean formula

Thus, @satwiksps I suggest you implement both options, your proposal and Manuel's proposal and add the test as a new z-scoring test and confirm the results.
@manuelgloeckler I think it would be good to have both options as the gaussian baseline assumption can be suboptimal when the posterior is multi-modal or skewed?

[manuelgloeckler]

@janfb The preconditioning is with respect to the "prior" not the posterior (as this would require regression from x). I don't think that it will "hurt" in almost all cases i.e. FM nets are initialized to output zero hence effectively will let the initialized network sample from a mass covering Gaussian approximation of the prior (and everything else needs to be learned).

Nonetheless having an option to disable it is always good.

Agree that the benchmark tests are not really sensitive to the z-scoreing, but as we usually enable z-scoreing by default it shouldn't hurt performance even if its not necessary. But as said the deviation is small enough to be fine (and might improve with the additional baseline).

[janfb]

Thanks for the update @satwiksps ! looks good, I just have one crucial question on the standard z-scoring formulas again, please check 🙏

[manuelgloeckler]

Thanks for you contribution.

I think the formula is still a bit off (but it also was never very clearly defined by us anyway (: ).

I do have a minor suggestion on t he mean,var buffers as well as the gaussian baseline test, which should be addressed (see comments). Once this done, we can merge it :)

Kind regards,
Manuel

[janfb]

Hi @manuelgloeckler and @satwiksps

I tested this locally using a linear Gaussian test with shifted prior to U(95, 105). We have the option to just z-score the time vector or to additionally use the Gaussian baseline assumption. It turns out that the formula that works best is the one that normalized the time vector to 0 across time (called zscore_true_marginal below).

I also tested the Gaussian baseline option, once with the formulas dirived from the Flow Matching velocity objective (baseline_velocity), and once the one that Manuel proposed which come from the score matching Ansatz I believe (baseline_position). Here we see that the former works a bit better.

• zscore_true_marginal: C2ST = 0.514 +/- 0.011
• baseline_velocity: C2ST = 0.529 +/- 0.014
• baseline_position: C2ST = 0.570 +/- 0.036
• zscore_initial_pr: C2ST = 0.640 +/- 0.011

I added all options as options in the internal code for and a smoke test comparing all these options. But just for reference. In the next commit I will clean things up.

So, I suggest we go with the zscore_true_marginal by default and offer the velocity Gaussian baseline as additional option with default False

[manuelgloeckler]

Thanks!

The implementation looks good. I am just a bit confused/concerned about the assumed integration direction. But I can be wrong there, can you point me to the part were this "switch" happens?

[janfb]

@manuelgloeckler thanks for the review and checking again the integration direction. It should now be fixed.

I will run slow tests again to make sure all is clean, and then we can merge / do you approve?

[janfb]

Slow tests are passing and formulas have been clarified. Thanks again @satwiksps for the initial hard work on this one! 👏 🚀

I will leave this final review and approval to @manuelgloeckler !

[manuelgloeckler]

Great, thanks all!