Differentiable simulation & gradient-based inference

If the simulator is differentiable, you can run it backwards: compute how its outputs change with the boson mass, and use that gradient to infer the mass from data. This article covers the differentiable-simulation idea and the inference it enables — a genuine novelty of the campaign.

What differentiable simulation means

A differentiable simulator is written so that automatic differentiation can propagate gradients through the entire computation — every timestep, every FFT. The output is not just a prediction but also its derivative with respect to any input parameter, exactly and cheaply.

Why it is powerful for inference

The inverse problem in cosmology is: given observed halo cores or a power spectrum, what boson mass produced them? With a differentiable model you can compute $\partial(\text{observable})/\partial m$ and descend the mismatch to the best-fit mass by gradient descent (figure) — far more efficient than a blind parameter scan, and it yields the sensitivity (Fisher information) for free.

Worked example — recovering the mass

Our differentiable JAXiON pipeline verifies the gradient against finite differences (AD/FD $=0.999$–$1.002$), then uses simulation-based inference (SBI) to recover the boson mass to $\sim$5% (e.g. $m_{22}=1.15^{+0.06}_{-0.05}$ for a true $1.10$). The credible intervals are calibrated (coverage 50/69/91% vs nominal 50/68/90). This is the first differentiable, adaptive Schrödinger–Poisson inference of the boson mass.

Differentiable simulation: the gradient $\partial(\text{loss})/\partial m$ points toward the boson mass that best matches the data, so gradient descent recovers $m$ efficiently — with calibrated uncertainties.

Simulation-based inference

Because the map from mass to observable can be non-monotone near the formation threshold, we pair gradients with simulation-based inference (SBI): train a neural density estimator on simulated (mass, observable) pairs to get a full posterior, with the gradients sharpening it. SBI also handles the intrinsic scatter (e.g. the core–halo dispersion) that a naive fit would ignore.

The payoff and its limits

Differentiable inference turns each expensive simulation into maximal information about the boson mass. Its limit is the same resolution wall (Topic 9.5): the gradients are only as good as the forward model, so the method shines for well-resolved observables (cores) and is bounded where the simulation cannot converge.

The recovered posterior on the boson mass ($\sim5\%$) from differentiable simulation-based inference.

In our research

This is a headline novelty of the campaign: the first differentiable, adaptive Schrödinger–Poisson solver plus differentiable SBI of the boson mass (recovered to ~5%, calibrated). It builds directly on JAXiON's differentiable spectral solver (Topic 9.3) and the imaginary-time soliton solver (Topic 5.5), and feeds the mass-tension analysis of Topic 10.

Key references

Cranmer, Brehmer & Louppe (2020), The frontier of simulation-based inference, PNAS 117, 30055 (arXiv:1911.01429).
Bradbury et al. (2018), JAX.
Hui, Ostriker, Tremaine & Witten (2017), Phys. Rev. D 95, 043541.

PreviousThe de Broglie wall and the fluid shortcut NextDwarf-galaxy cores: a light-boson signal