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 is necessary for far more powerful statistical methods. And they resist integration into the analysis workflows now standard elsewhere in cosmology.

A team led by 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, enabling gradient-based Bayesian inference for the first time and making joint CMB+BBN analyses tractable on a laptop.

How It Works

The core innovation in LINX is its foundation in 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 mathematical technique that dramatically speeds up the calculation without sacrificing accuracy.
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 that correspond 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, precise measures of how final elemental abundances respond to any change in the inputs, 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 possible. These gradient-based statistical methods navigate parameter space by following mathematical slopes rather than wandering randomly, and they are exponentially 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 that 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. This makes richer searches for new physics practical: light dark matter species, non-standard neutrino interactions, exotic particle decays during BBN. Analyses that previously required weeks of cluster time become afternoon projects.

LINX is also built to grow. 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.

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