Wrapping Up My GSoC 2026 Journey with sbi

This is my final work product for Google Summer of Code 2026, where I spent the summer redesigning how sbi builds neural networks.

What the project is about?

sbi does Bayesian inference for simulators you cannot write a likelihood for. You have a simulator and a prior, you run simulations, and a neural network learns the posterior over the simulator's parameters from those runs. It implements a whole family of methods (NPE, NLE, NRE, FMPE, NPSE and the mixed variants), and every one of them trains some kind of neural density estimator.

My project was about how you pick that network and configure it. Over the summer I replaced the old string and factory-function interface with typed configuration objects, one class per model, across every estimator family in the library. The goal was that a setting the chosen model cannot use fails immediately when you write it, rather than being quietly dropped on the way to the network, and that existing code keeps working through a deprecation path instead of breaking.

Organisation: NumFOCUS | Sub-organisation: sbi-dev/sbi | Mentors: Jan Teusen (@janfb) & Nicholas Junge (@nicholasjng)

What was the problem?

The old API had three ways to specify a network, and they had grown separately:

1. A string: NPE(prior, density_estimator="nsf"). Easy, but no way to set hyperparameters.
2. A factory function: posterior_nn(model="nsf", hidden_features=64). It returned an opaque closure. You could not inspect it, print it or serialize it, and it accepted any keyword you gave it.
3. A custom module, for full control.

The factory functions were the problem. They spanned a whole family of models, so their signature was the union of every setting any model might want. A setting that the chosen model did not use was quietly discarded. There was no type information anywhere, so editors could not autocomplete and type checkers could not help.

The goal was a layered API where each layer is typed and inspectable, mistakes fail at configuration time instead of at training time, and existing code keeps working through a deprecation path rather than breaking.

How is the problem solved now?

Every estimator family now has typed config objects:

from sbi.inference import NPE, NRE
from sbi.neural_nets import ZukoNSFConfig, ResNetClassifierConfig

trainer = NPE(prior, density_estimator=ZukoNSFConfig(hidden_features=64, num_transforms=8))
trainer = NRE(prior, classifier=ResNetClassifierConfig(hidden_features=64, num_blocks=3))

They are frozen dataclasses, so they print cleanly, you can log them in an experiment record, and your editor autocompletes the fields. Strings still work and emit a FutureWarning that names the class to switch to. Custom modules still work.

The part I am most happy with is what happens when you are wrong. There are three distinct failures, all before training:

NSFConfig(hiden_features=64)          # TypeError  - misspelled name
NSFConfig(z_score_input="strucured")  # ValueError - misspelled value
MAFConfig(num_bins=8)                 # TypeError  - not a setting this model has

What each PR did

#1872 laid the groundwork. The shared base class, the supporting types, and a protocol rename to free up the name "Builder" for the new objects.

#1877 added the first real builder and #1882 wired it into the NPE trainers, with the deprecation path for strings.

The most useful thing that happened here was something we removed. The original design had build() take a BuildContext object carrying shape, device and z-scoring statistics. While implementing it, Jan noticed the parameter was never actually read: everything it held could be derived from the training batch that was already being passed in. We cut it. It was a premature abstraction, and recognising that early saved a lot of machinery that would have needed maintaining.

#1904 extended the builder to NLE, and this is where the design got interesting.

NPE models the parameters given the data. NLE models the data given the parameters. Same builder, opposite roles. A signature of build(batch_theta, batch_x) reads correctly for one and backwards for the other, which is exactly how you get a bug where the standardization lands on the wrong variable.

So build() became build(batch_input, batch_condition). input is whatever the model is modelling and condition is whatever it is conditioning on, and the trainer decides which is which. The config never has to know. That decision carried cleanly through every family that came after, including the vector-field one that landed months later.

#1912 added the mixed density estimators for MNLE and MNPE, where part of the data is discrete.

#1920 added the NRE classifiers and bundled in a set of fixes to the shared base class. The reasoning was timing: this was the first PR proving the base served a third family, so fixing it there meant NRE and everything after inherited the fixes instead of needing a retrofit.

What went in:

• Invalid Literal values now raise at construction. A misspelled field name already raised. A misspelled value on a correctly named field did not, and stayed silent until training.
• Model-incompatible settings raise, by inspecting the target build function's signature.
• frozen=True on the configs, so they are immutable.
• A custom __repr__ that prints only what you actually set.
• Uniform public API exposure, since one builder had been left unexported.

It also carried a rename I am glad we did. The configs used z_score_x and z_score_y, inherited from the low-level functions. Those names are ambiguous across families: for NPE, z_score_x standardizes the parameters, not x. They became z_score_input and z_score_condition, matching the role-neutral signature. The builder API had not shipped in a release yet, so the rename was free. A few weeks later it would have been a breaking change.

#1921 covered FMPE and NPSE. This was also the point where the working style changed. Instead of implementing a design handed to me, I started writing a short design proposal first and getting it reviewed before writing code. It felt slower for about a day and then obviously correct: a design mistake caught in a document costs an hour, and the same mistake caught in code review costs a week.

Partway through, we changed direction, and I think it made the whole API better. Until then, each family had one flat class holding every field any model in that family might want. To stop you setting a field your model does not use, I had built machinery: applicability lists, signature inspection, and drift tests to keep those lists honest as models changed.

The new idea deletes all of it by making the class itself the choice. One small config per model, carrying only the fields that model accepts. Then there is nothing to validate, because a field that does not apply simply does not exist, and Python's own argument checking raises for free.

#1975 piloted that on the marginal trainer with nine classes, one per Zuko flow. #1986 then converted the density, classifier and mixed families, replacing three flat builders with 15 density configs, 3 classifier configs and one mixed config.

Where things stand?

Merged and in sbi:

Open for review:

• #1997 finishes the mixed family. The internal mixed build function was still taking two APIs at once, the old flat arguments and a config, so it now takes only a config and the deprecated strings build one in the factory. It also fixes a compatibility bug I would never have caught by counting parameters: the flat API passed a wider spline tail_bound to every model that reads it, and keying that off the model name instead of the field silently narrowed the tail for one model while leaving the parameter count identical.

Draft, queued behind each other:

• #1996 converts the vector-field family to the per-model design. This one is the most interesting of the set. The builder picks along what looks like three axes at once (flow or score matching, which SDE, which network architecture), but the axes are not independent: the estimator type and the SDE together select exactly one class. So it becomes two composed objects, FlowMatchingConfig(net=MLPConfig(...)), which is seven classes instead of the sixteen you get from multiplying the axes out.
• #1987 is the documentation PR: a new how-to for the config API, a rework of the "abstraction levels" guide, the config classes added to the API reference, and the deprecated strings cleared out of the tutorials. It stays a draft until #1996 merges, because it documents the API that PR lands, then it gets rebased and reviewed.

The stacking matters here. Each of these sits on the one before it, so they merge in order and each gets rebased onto main once its predecessor lands. That is a habit I picked up the hard way, and it is why the diffs stay readable.

What is left to do?

Test consolidation. Each PR added its own test file, which was right during development and is now redundant. We agreed to consolidate them into fewer parametrized files once everything merges. The distinction that matters for CI cost is that tests which train are the expensive ones, while construction and default checks are cheap, so the consolidation should keep the cheap ones broad and the expensive ones few.

Retiring the legacy factory functions. posterior_nn, likelihood_nn and classifier_nn still exist and still work. They now build the new configs internally, so they are already thin. Once the deprecation window closes they can go, along with the legacy *Config validators behind them.

The stretch goal: data-source-agnostic builds. This is the one I did not get to, and it is the most interesting remaining piece. Right now z-scoring statistics are computed inside the build functions from a single in-memory batch. For datasets too large to hold in one tensor, those statistics would have to come from a pre-pass over a dataloader, or be accumulated online. The plan is to move statistics and shapes out of the build functions and into a context object that the trainer fills in, which is the BuildContext we deliberately deferred back in phase 1. The private types for it are already in the codebase as a reserved seam, so it is a local addition rather than a re-architecture. It is correctness-sensitive, since wrong z-scoring silently degrades inference rather than erroring, and it touches roughly 24 build functions, which is why it stayed a stretch goal.

Acknowledgement

I am really grateful to my mentor Jan Teusen for giving me this opportunity and for reviews that consistently found the thing I had not thought about, and for pushing me to write designs before code in the second half. To Nicholas Junge for the backup mentorship. And to the sbi maintainers for treating a student's PRs with the same seriousness as anyone else's. This was a great learning experience for me.

The coding period is over but the work is not. I am continuing with sbi through November to land the vector-field conversion and the documentation, and I intend to stick around after that as a permanent contributor.

I have also written some blogs every two weeks throughout the program describing my learnings in detail.

Signing off,
Satwik Sai Prakash Sahoo
GitHub | LinkedIn | Bluesky | X