Meta-learning lets a cosmology emulator adapt to new galaxy depth maps quickly

The work

Technical · from the record

The authors trained a neural network to stand in for the direct theoretical computation of the cosmic shear angular power spectrum, using the Model-Agnostic Meta-Learning (MAML) algorithm so that one set of network parameters could be fine-tuned to new galaxy redshift distributions without taking any description of the distribution as input. The meta-trained emulator was fine-tuned on a novel redshift distribution with 100 samples over 64 epochs and compared with an emulator pre-trained on a single redshift distribution and one trained from scratch. In an MCMC analysis, the MAML emulator's posterior had a Bhattacharyya distance of 0.008 from the posterior computed from theory in the S8-Omega_m plane, compared with 0.038 for the single-task pre-trained emulator and 0.243 for the emulator with no pre-training. An emulator trained from scratch matched or exceeded the fine-tuned MAML emulator once more than about 8,000 training samples were provided for the in-distribution test task, and after 4,000 samples for an out-of-distribution multi-modal test distribution.

How AI was used

Cosmic shear angular power spectrum data vectors were generated from theory with the Core Cosmology Library for randomly drawn Smail-type and Gaussian source redshift distributions, each split into five tomographic bins with added noise, and for Latin hypercube samples of five cosmological parameters plus five tomographic bin mean-redshift shifts. Each unique redshift distribution constituted a meta-learning task and each data vector within it a shot. A network combining a projection layer reshaped into 2D, dilated convolutional layers and a final fully connected layer was built in PyTorch with dropout of 0.2; data vectors were log-transformed and both inputs and outputs standardised. The network was meta-trained with first-order MAML using mean squared error loss in inner and outer loops, an initial inner learning rate of 0.001 and outer learning rate of 0.01, with the Adam state shared between loops and 60% of each task's samples used as the support set. A grid search over 64 combinations of task count, shot count and task batch size led to a configuration of 20 tasks, 500 samples per task and 5 tasks per batch. Two comparison emulators were trained on the same architecture: one pre-trained on a single Gaussian redshift distribution with 10,000 samples, and one trained from scratch on the test task. Emulators were fine-tuned to a novel task for 64 epochs with varying numbers of samples, then run over unseen data vectors and inside an emcee likelihood to produce posterior samples.