LINX: A Fast, Differentiable, and Extensible Big Bang Nucleosynthesis Package

The Big Picture

In the first few minutes after the Big Bang, the universe was a cosmic pressure cooker. Temperatures in the billions of degrees. Protons and neutrons smashing together to forge the lightest elements: hydrogen, helium, lithium. This process, Big Bang Nucleosynthesis (BBN), left a chemical fingerprint that scientists can still read today by measuring elemental abundances in ancient, pristine gas clouds billions of light-years away. The agreement between theory and observation is one of the great triumphs of modern cosmology.

But computing those predictions accurately, and fast enough to be useful, has been a persistent bottleneck. Cosmologists need to run BBN codes millions of times to fit models to data, testing vast combinations of physical parameters and hunting for new physics hiding in a slight mismatch between predicted and observed helium abundance.

Existing tools work, but weren’t built for precision cosmology at today’s scale. They’re slow. They lack differentiability, the ability to compute how outputs respond to small changes in inputs, which modern statistical methods require. And they don’t integrate easily into the analysis workflows now standard elsewhere in cosmology.

A team of researchers at NYU, Princeton, MIT, and IAIFI has built a new BBN code from scratch, one that speaks the language of modern machine learning. LINX (Light Isotope Nucleosynthesis with JAX) is fast, fully differentiable, and slots directly into gradient-powered inference pipelines.

LINX brings Big Bang Nucleosynthesis into the era of differentiable programming. It makes gradient-based Bayesian inference practical for the first time and puts joint CMB+BBN analyses within reach of a laptop.

How It Works

The core of LINX is JAX, Google’s numerical computing library. JAX compiles Python to fast machine instructions and can automatically compute derivatives of any function. This isn’t just a software engineering detail; it changes what cosmologists can actually do with a BBN code.

Traditional BBN codes solve a coupled system of nuclear reaction network equations: differential equations tracking how abundances of hydrogen, deuterium, helium-3, helium-4, lithium-7, and other light isotopes evolve as the universe cools. LINX implements the same physics, but keeps every step differentiable end-to-end:

1. Thermodynamic background: LINX computes how temperature, energy density, and expansion rate evolve in the early universe, including contributions from photons, electrons, positrons, and neutrinos, with precision inputs from the code nudecBSM.
2. Weak reaction rates: The neutron-to-proton ratio, a key input to element formation, is set by weak interactions (responsible for certain types of radioactive decay) in the first seconds. LINX computes these rates using a fast numerical scheme that preserves accuracy while cutting computation time.
3. Nuclear network integration: Using JAX’s equation-solving tools, LINX integrates the full reaction network forward in time, tracking isotope abundances from ~10 MeV down to ~10 keV (energy units corresponding to the universe’s temperature at those epochs) until abundances freeze out and reach their final values.
4. Likelihood evaluation: Predicted abundances are compared directly to observations and combined with cosmic microwave background data for joint statistical analysis.

Differentiability means users can compute gradients of final elemental abundances with respect to any input, analytically rather than numerically. Those inputs can be anything: baryon density, neutrino count, nuclear reaction rates, or exotic new-physics parameters.

This is what makes Hamiltonian Monte Carlo (HMC) and variational inference viable for BBN. These gradient-based methods navigate parameter space by following mathematical slopes rather than wandering randomly. In high-dimensional problems, they can be orders of magnitude more efficient than the random-walk samplers BBN codes have historically been stuck with.

In validation tests, LINX matches the established code PRIMAT to better than 0.1% for helium-4 abundance and within a few percent for deuterium and lithium-7, well within current observational uncertainties. A single evaluation takes milliseconds on a standard laptop. After JIT compilation, full parameter scans run orders of magnitude faster than comparable legacy codes.

Why It Matters

Cosmology is entering an era of extraordinary precision. Experiments like the Simons Observatory and CMB-S4 will deliver exquisite measurements of the cosmic microwave background, and those measurements demand equally precise theoretical predictions. The angular pattern of temperature fluctuations in the CMB depends on the primordial helium abundance, which a BBN code must compute before a CMB simulation code can even begin. BBN sits at a chokepoint in the analysis pipelines for these next-generation experiments.

LINX removes that chokepoint.

By combining LINX with differentiable CMB emulators like CosmoPower, the team built a fully differentiable joint CMB+BBN pipeline that constrains the baryon density and effective number of neutrino species simultaneously, all on personal hardware. No supercomputer required. That makes it practical to search for light dark matter species, non-standard neutrino interactions, and exotic particle decays during BBN. Analyses that previously required weeks of cluster time become afternoon projects.

LINX is built to grow, too. Because it’s written in high-level Python with JAX, adding new physics (say, a particle species that injects entropy during nucleosynthesis) means modifying a few Python functions, not rewriting Fortran from scratch.

With LINX, the cosmic chemical record of the Big Bang becomes accessible to the full toolkit of differentiable programming, applied to one of cosmology’s oldest and most powerful observational probes.