Machine Learning for Astroparticle Physics:
A Crash-course in SBI

Lecture 3b - Inference heads and parameter models

Christoph Weniger, University of Amsterdam (GRAPPA)

Today's Roadmap

The SBI pipeline — the simulator

Simulator — encodes the physics.
Maps parameters \(\boldsymbol\theta\) to data realisations \(\mathbf{x} \sim p(\mathbf{x}\mid\boldsymbol\theta)\).

Examples:

• Strong-lensing image generators (source + lens model + PSF + noise).
• Gravitational-wave waveform models (LALSuite, IMRPhenom).
• Cosmological N-body / hydro (Quijote, CAMELS).
• Particle-physics event generators (Pythia, MadGraph).

⚠ Failure mode: wrong model → wrong inference. Misspecification of the simulator propagates into the posterior, often invisibly.

The SBI pipeline — the inference loss

Inference loss \(\mathcal{L}\) — defines the statistical question.
The choice of loss is what turns the network into a particular kind of estimator.

The loss decides what we're solving:

• NPE → NLL of \(\boldsymbol\theta^\star\) under \(q_\phi(\boldsymbol\theta\mid\mathbf{x})\).
• NRE → binary cross-entropy on a joint-vs-marginal classifier.
• NLE → NLL of \(\mathbf{x}^\star\) under \(q_\phi(\mathbf{x}\mid\boldsymbol\theta)\).
• Flow matching → velocity-regression on a probability path.
• Diffusion → denoising score matching.

Same simulator, same network, different loss ⇒ different statistical task.

The SBI pipeline — the embedding network

Embedding network \(T_\phi(\mathbf{x})\) — data compression.
Maps high-dimensional observations to a low-dimensional summary that the inference head can chew on.

Why we need it:

• Strong lensing: \(\mathbf{x}\) is a \(256\times 256\) image, \(\sim\!65{,}000\) pixels.
• Gravitational waves: \(\sim\!16{,}000\) time-series samples per detector.
• Cosmology: density fields of \(\sim\!10^9\) voxels.

A raw MLP would have too many parameters and no inductive bias. A CNN (for images), Transformer (for sequences), or GNN (for catalogues) does the compression. Trained jointly with the inference head, so \(T_\phi\) learns features informative for the posterior, not for any predefined label. Architecture choices are the topic of Lec 4.

NPE, NRE, NLE, one table to rule them all

What does an NPE network actually learn?

In Lec 2 we trained \(q_\phi(\boldsymbol\theta\mid\mathbf{x})\) on a fixed dataset of pairs \((\boldsymbol\theta_n, \mathbf{x}_n)\) by minimising the NLL. Today the pairs come from a simulator, \(\boldsymbol\theta_n \sim p(\boldsymbol\theta)\), \(\mathbf{x}_n \sim p(\mathbf{x}\mid\boldsymbol\theta_n)\). Same loss, same network. This is neural posterior estimation (NPE).

The NLL loss is just maximum likelihood: maximise \(\prod_i q_\phi(\boldsymbol\theta^{(i)}\mid\mathbf{x}^{(i)})\), the joint probability of all observed parameters given their data, which is the same as minimising \(-\!\sum_i \log q_\phi(\boldsymbol\theta^{(i)}\mid\mathbf{x}^{(i)})\).

Claim: at the minimum of \(\mathcal{L}\), the learned density \(q_\phi(\boldsymbol\theta\mid\mathbf{x})\) is the Bayesian posterior.

Papamakarios & Murray 2016 (mixture-density); Lueckmann et al. 2017; Greenberg et al. 2019; review: Cranmer, Brehmer, Louppe (PNAS 2020).

Training data is a distribution

The simulator pairs \(\mathcal{D} = \{(\boldsymbol\theta^{(i)}, \mathbf{x}^{(i)})\}_{i=1}^{N}\) define an empirical joint distribution:

\[ \hat p(\boldsymbol\theta, \mathbf{x}) \;=\; \frac{1}{N}\sum_{i=1}^{N}\, \delta(\boldsymbol\theta-\boldsymbol\theta^{(i)})\, \delta(\mathbf{x}-\mathbf{x}^{(i)}) \]

Infinite-data limit — minimising NLL gives the posterior

Take two idealisations: (i) \(N\to\infty\), so \(\hat p \to p(\boldsymbol\theta, \mathbf{x}) = p(\boldsymbol\theta)\,p(\mathbf{x}\mid\boldsymbol\theta)\); (ii) the network is infinitely flexible, so \(q(\boldsymbol\theta\mid\mathbf{x})\) is any function with \(\int q(\boldsymbol\theta\mid\mathbf{x})\,d\boldsymbol\theta = 1\) for every \(\mathbf{x}\).

\[ \mathcal{L}[q] \;\longrightarrow\; -\!\iint p(\boldsymbol\theta, \mathbf{x})\, \log q(\boldsymbol\theta\mid\mathbf{x})\, d\boldsymbol\theta\, d\mathbf{x} \]

The learned \(q_\phi(\boldsymbol\theta\mid\mathbf{x})\) is the Bayesian posterior implied by the simulator. A finite network just interpolates this idealised answer.

Infinite-data limit — two proofs

Starting from \(\mathcal{L}[q] = -\!\iint p(\boldsymbol\theta, \mathbf{x})\, \log q(\boldsymbol\theta\mid\mathbf{x})\, d\boldsymbol\theta\, d\mathbf{x}\), we want the minimiser \(q^\star\).

Same answer, two lenses: (a) is mechanical (stationary point of a constrained functional); (b) is geometric (NLL = KL up to a \(q\)-independent constant, so minimising NLL = projecting \(p(\cdot\mid\mathbf{x})\) onto the model family under KL).

Heteroscedastic noise — variance as a second output

Our first model baked in a single \(\sigma_\theta\). Make \(\sigma\) a function of \(x\); same network, two heads: \(\mu_\phi(x)\) and \(\log\sigma_\phi(x)\). The NLL trades off mean accuracy against calibrated width.

Toy: noise is small on the left, large on the right. The grey band is the learned \(\sigma(x)\), it tracks the true scale instead of averaging over \(x\).

The Gaussian-head network — the simplest inference head

// Train q_phi(theta|x) on simulator pairs (theta_n, x_n)
for step = 1, 2, ... :
    sample mini-batch B of pairs
    mu, log_sigma = f_phi(x)
    L = mean( 0.5*((theta-mu)/sigma)^2 + log_sigma )
    phi <- Adam(phi, grad_phi L)

Live demo — Gaussian-head NN on a linear forward model

Forward model: \(t_i = a + b\,x_i + \varepsilon\), 20-point grid, \(\varepsilon\sim\mathcal{N}(0,0.25^2)\). Priors \(a\sim\mathcal U[-1,1]\), \(b\sim\mathcal U[-2,2]\). Network: MLP \([20{\to}32{\to}16{\to}5]\), output \((\mu_a,\mu_b,\log\sigma_a,\log\sigma_b,\rho)\).

Left: a single observation \(\mathbf{x}\). Right: the learned posterior \(q_\phi(a,b\mid\mathbf{x})\) (orange ellipse) vs the true \((a^\star,b^\star)\) (green star). The same trained network handles any new \(\mathbf{x}\), amortised inference.