Meta-learning for cosmological emulation: Rapid adaptation to new lensing kernels

Abstract

Theoretical computation of cosmological observables is an intensive process, restricting the speed at which cosmological data can be analysed and cosmological models constrained, and therefore limiting research access to those with high performance computing infrastructure. Whilst the use of machine learning to emulate these computations has been studied, most existing emulators are specialised and not suitable for emulating a wide range of observables with changing physical models. Here, we investigate the Model-Agnostic Meta-Learning algorithm (MAML) for training a cosmological emulator. MAML attempts to train a set of network parameters for rapid fine-tuning to new tasks within some distribution of tasks. Specifically, we consider a simple case where the galaxy sample changes, resulting in a different redshift distribution and lensing kernel. Using MAML, we train a cosmic-shear angular power spectrum emulator for rapid adaptation to new redshift distributions with only $O(100)$ fine-tuning samples, whilst not requiring any parametrisation of the redshift distributions. We compare the performance of the MAML emulator to two standard emulators, one pre-trained on a single redshift distribution and the other with no pre-training, both in terms of accuracy on test data, and the constraints produced when using the emulators for cosmological inference. We observe that within an MCMC analysis, the MAML emulator is able to better reproduce the fully-theoretical posterior, achieving a Battacharrya distance from the fully-theoretical posterior in the $S_8$ -- $\Omega_m$ plane of 0.008, compared to 0.038 from the single-task pre-trained emulator and 0.243 for the emulator with no pre-training.