Score-based density estimators for SBI

[jsvetter]

What does this implement/fix? Explain your changes

This PR implements score-based methods for SBI (and related methods like flow matching). The first goal is to have a running version for score-based NPE. Later-on NLE, sequential methods and methods that can deal with multiple observations can be tackled.

This includes new base-classes for the required vectorfield-estimators and posteriors.

• adds a PosteriorScoreBasedPotential and ScorePosterior
• adds gradient method to all potentials, e.g., for running gradient ascent for MAP. (NPSE can directly use the score then, other potentials use autograd on the log_prob)

Left for a future PR (see #1226 ):

• iid_bridge for IID data as in Geffner et al.
• unification of flow matching and score matching neural nets and posteriors.

Does this close any currently open issues?

fixes #962
fixes #956

[manuelgloeckler]

Related to #963 and #966 .

[janfb]

@pfuhr do you agree to your contributions being re-licensed under Apache 2.0 in the future?

[pfuhr]

Yes, I agree that my contributions are re-licensed under Apache 2.0 in the future.

[manuelgloeckler]

Alright, I integrated everything with the current main and fixed some problematic side effects/bugs we introduced. The following stuff still has to be implemented/improved:

1. NSPE interface (MINOR)

• Check docstrings (as this is user-exposed, might be some copy pasta)
• build_posterior we do have two possibilities i.e. sample_with=”sde” will return a ScorePosterior and sample_with=”ode” will return a DirectPosterior (not implemented, may be very similar to flow matching implementation feat: flow matching methods #1049 , but in the end should really just build a ZUKO CNFs). Some stuff has been commented out and needs to be checked.
• Convergence checks (now using EMA with low tolerance...)

2. ScorePosterior (MEDIATE)

• sample should interface the score-based-samplers (dependent on 3)). But we need to properly shape conditions … . Not sure if need to use “accept_reject” for prior bounds, this anyway should in the future be done within the diffusion process.
• log_prob Not sure if needed, if one wants an log_prob one should build a direct posterior. But maybe still good for map.
• map Should use the score_net directly as gradient (no backprop on potentials) i.e. sample a few candidates, run gradient accent where the gradients is just the output of the score_net at t=T_min (Optional, create an issue to do this in the future).

3. score_based_potential (MINOR)

• This is more conceptually/abstraction. It just is not a potential but the gradient of a potential and behaves differently than the other potentials (might need another class?). Although this should, as the other potentials, mediate IID stuff. (These will also change in @gmoss13 feat: batched sampling for MCMC #1176. Any suggestions @janfb ?

4. Score based sampler (MAJOR)

• This needs to be rewritten to be a general sampling interface (we only need to implement one for now, but it must be extendable)
  • Should have a “predict/step” which steps backward in time (with euler/DDIM/…)
  • Should have an optional “correct” which stays at the same time (Langevin/Gibbs …)
• Needs to be vectorized i.e. must support batched x at best (but this also loosely depends on feat: batched sampling for MCMC #1176)

5. Score nets/estimator It looks quite good; I might make a few minor changes to improve the architecture. (MEDIATE)

• input z_score should be time-dependent (as we add noise to it so, std/mean will change)
• output_scores should be preconditioned by analytic approximation.
• Times should be preconditioned by how fast mean/std changes (if they do not change a lot, the score should also not change much).
• Add option to add covariate to loss to reduce variance.
• Improve MLP (maybe with layernorm and time insertion in each layer.

6. Tests: We need to extend and check tests. We currently have:

• linearGaussian_nspe_test.py Currently only tests sample
• score_estimator_tests.py Tests shaping conventions.
• score_tests.py: Does not yet do something but implements some analytic scores. We need to test the samplers which we can also do with the analytic scores!

If you have time/want to do some of the points @rdgao @jsvetter, let me know. We can also chat about what should be done in more detail.

[codecov]

Codecov Report

Attention: Patch coverage is 93.56110% with 49 lines in your changes missing coverage. Please review.

Project coverage is 78.20%. Comparing base (b3254ed) to head (bc00bd5).
Report is 13 commits behind head on main.

Files	Patch %	Lines
sbi/neural_nets/estimators/score_estimator.py	92.50%	12 Missing ⚠️
sbi/samplers/score/correctors.py	52.17%	11 Missing ⚠️
sbi/samplers/score/score.py	85.18%	8 Missing ⚠️
sbi/inference/npse/npse.py	96.47%	5 Missing ⚠️
sbi/neural_nets/score_nets.py	95.23%	5 Missing ⚠️
sbi/inference/posteriors/score_posterior.py	95.52%	3 Missing ⚠️
sbi/inference/potentials/score_based_potential.py	97.01%	2 Missing ⚠️
sbi/inference/base.py	83.33%	1 Missing ⚠️
sbi/inference/potentials/base_potential.py	50.00%	1 Missing ⚠️
sbi/samplers/score/predictors.py	97.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1015      +/-   ##
==========================================
- Coverage   85.44%   78.20%   -7.24%     
==========================================
  Files         101      110       +9     
  Lines        7941     8632     +691     
==========================================
- Hits         6785     6751      -34     
- Misses       1156     1881     +725     

Flag	Coverage Δ
unittests	78.20% <93.56%> (-7.24%)	⬇️

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

Files	Coverage Δ
sbi/analysis/tensorboard_output.py	86.41% <100.00%> (ø)
sbi/inference/__init__.py	100.00% <100.00%> (ø)
sbi/inference/fmpe/fmpe_base.py	94.28% <100.00%> (-0.06%)	⬇️
sbi/inference/npse/__init__.py	100.00% <100.00%> (ø)
sbi/inference/posteriors/base_posterior.py	85.88% <100.00%> (-0.17%)	⬇️
sbi/inference/posteriors/direct_posterior.py	98.79% <100.00%> (ø)
sbi/inference/posteriors/mcmc_posterior.py	85.93% <100.00%> (-0.77%)	⬇️
...inference/potentials/likelihood_based_potential.py	100.00% <100.00%> (ø)
.../inference/potentials/posterior_based_potential.py	95.34% <100.00%> (ø)
sbi/inference/snle/snle_base.py	93.87% <100.00%> (ø)
... and 34 more

... and 22 files with indirect coverage changes

[janfb]

partial review.

[gmoss13]

Left quite a few comments, but most can be addressed pretty easily. I think we're almost ready to merge this!

[michaeldeistler]

Got tired towards the end, but reviewed about 80%

[gmoss13]

Awesome! Pushed a couple of small doc fixes but otherwise happy to merge!

[janfb]

Looks all good now. Great effort everyone 👏