score-based diffusion models as density estimators

[janfb]

Is your feature request related to a problem? Please describe.

Score estimation has been introduced for SBI, see

• Compositional score estimation, good for iid data: https://proceedings.mlr.press/v202/geffner23a.html
• sequential score estimation, for NLE and NPE: https://openreview.net/references/pdf?id=anxDgYxrx

it would be great to integrate these approaches into sbi.

Describe the solution you'd like

• Add score-based diffusion models as density estimators here: https://github.com/sbi-dev/sbi/tree/main/sbi/neural_nets/density_estimators
• Add neural posterior score estimation as inference method, e.g., here: https://github.com/sbi-dev/sbi/tree/main/sbi/inference/snpe
• add neural likelihood score estimation as inference method here: https://github.com/sbi-dev/sbi/tree/main/sbi/inference/snle
• add the option of compositional score estimation for iid data.

(feel free to move the last three points to separate issues, or create individual PRs for each).

[jsvetter]

I would also be interested in doing this!

[pfuhr]

me too!

[pfuhr]

I am not part of the org yet, though would like to be assigned.

[rdgao]

rough work plan for the "basic" components:

• in neural_nets, something like a VectorFieldEstimator super class that is the generic class for specific subclasses, including ScoreEstimator and FlowMatchEstimator, corresponding to score matching (i.e., denoising diffusion model) and flow matching, respectively, with their specific losses
• a VectorFieldSampler class in samplers that receives the VectorFieldEstimator and samples by, e.g., solving the ODE/SDE
• a VectorFieldPosterior object in inference/posteriors/ that receives the VectorFieldEstimator and sets up a VectorFieldSampler to create the posterior distribution, similar to DirectPosterior

does this sound about right? (also @janfb, @michaeldeistler, etc.)

[jsvetter]

torchdiffeq seems like a reasonable choice for ODEs in pytorch. It's also used in dingo! Alternatives are torchDyn and torchode.

[rdgao]

I don't know how to link the branch to the issue here but it's in score_estimator