Neural networks stand in for slow cosmology code in parameter fitting

The work

Technical · from the record

The authors built a framework for specifying, training, packaging and distributing machine-learning emulators of Einstein-Boltzmann codes within the CosmoPower package, together with wrappers that load the packaged emulators into the Cobaya and CosmoSIS inference codes. Using it they trained dense neural network emulators of CAMB v1.5.0 outputs - CMB TT, TE, EE, BB and lensing-potential angular power spectra, linear and non-linear matter power spectra, background evolution quantities and derived parameters - for LCDM and four extension models, using 100000 training samples for LCDM and 120000 for the extended models. Compared with direct CAMB calculations the emulators stayed within 10% of a cosmic-variance-limited experimental uncertainty, with the exception of the w0wa emulator, where some small-scale CMB outliers reached about 80% of that uncertainty. In MCMC runs on simulated cosmic-variance-limited CMB data the emulator chains reproduced the CAMB posteriors and the input cosmology to within 0.1 sigma, with a single power spectrum evaluation taking about 0.1 s against about 20 s for CAMB.

How AI was used

Cosmological parameter sets were drawn by Latin hypercube sampling over the ranges in Table 2 and passed to CAMB v1.5.0 with Stage-IV accuracy settings to produce training spectra, background quantities and derived parameters, stored in HDF5 files; a small fraction of unphysical samples was discarded. Inputs and outputs were normalised by their means and standard deviations, most quantities were emulated in the logarithm, and the TE and lensing-potential spectra were first compressed with a principal component analysis (512 and 64 retained components respectively), with scree plots used to choose the number of components. Dense neural networks with four hidden layers of 512 neurons were then fitted to map parameters to spectra, trained with the Adam optimiser over successive learning iterations whose learning rates, batch sizes, validation split, gradient accumulation steps, patience values and maximum epochs are set in the packaging prescription. Separate emulators were trained for LCDM and for the +Neff, +sum m_nu, +Neff+sum m_nu and +w0wa extensions, and for the non-linear matter power spectrum and non-linear boost with HMCode baryonic feedback parameters as additional inputs. Accuracy was evaluated by passing held-back samples through each emulator and plotting errors against the CAMB outputs, both as fractional differences and relative to cosmic-variance-limited or Simons Observatory noise curves. The packaged emulators were then loaded through the Cobaya and CosmoSIS wrappers to run Monte Carlo posterior sampling on a simulated cosmic-variance-limited CMB data vector and on a Stage-IV-like 3x2pt large-scale-structure data set, with an optional fall-through to the native Einstein-Boltzmann code for parameters outside the trained range.