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

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$ – $Ω_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.

💡 Research Summary

The paper addresses a fundamental bottleneck in modern cosmological inference: the computational cost of repeatedly evaluating theoretical observables (e.g., the cosmic‑shear angular power spectrum, APS) for millions of parameter sets during Bayesian sampling. Traditional emulators such as CosmoPower, CONNECT, and the Boruah et al. models accelerate this step but are typically specialized to a single survey configuration or redshift distribution. When a new galaxy sample with a different redshift distribution (hence a different lensing kernel) is introduced, a new emulator must be trained from scratch, which defeats the purpose of rapid inference.

To overcome this limitation, the authors explore Model‑Agnostic Meta‑Learning (MAML), a meta‑learning algorithm that learns an initialization of neural‑network weights which can be fine‑tuned to a new task with only a few gradient steps and a small amount of data. In the cosmological context, each “task” corresponds to a specific redshift distribution N(z) for the source galaxies. The authors generate a suite of synthetic tasks by sampling a broad range of N(z) shapes and, for each, compute the corresponding APS using a Boltzmann code (CLASS). They then train a feed‑forward network using first‑order MAML combined with the Adam optimizer, sharing Adam’s moment estimates across tasks to reduce memory overhead.

Three emulators are compared:

1. MAML‑trained meta‑emulator – pre‑trained on many N(z) tasks, then fine‑tuned on a new distribution with only O(100) samples.
2. Single‑task pre‑trained emulator – trained on a single redshift distribution and fine‑tuned on the new one.
3. Untrained (scratch) emulator – no pre‑training; it is trained from the same O(100) samples.

Performance is evaluated on two fronts. First, the mean‑squared error (MSE) on a held‑out test set of APS values shows that the MAML emulator achieves an MSE of ~0.001, substantially lower than the single‑task (≈0.004) and scratch (≈0.012) models. Second, the authors embed each emulator in a full MCMC pipeline to infer the cosmological parameters S₈ and Ωₘ. They quantify the similarity of the resulting posterior to the “ground‑truth” posterior obtained with direct Boltzmann‑code calculations using the Battacharyya distance (BD). The MAML emulator yields BD = 0.008, whereas the single‑task and scratch models give BD = 0.038 and 0.243 respectively. Thus, the meta‑learned emulator reproduces the true posterior almost indistinguishably, while requiring far fewer fine‑tuning samples.

Computationally, the meta‑training phase consumes roughly 12 hours on four GPUs, after which fine‑tuning to a new N(z) takes less than a minute. By contrast, achieving comparable accuracy with the scratch emulator would require thousands of training points and many hours of GPU time. The authors therefore argue that the upfront cost of meta‑training is amortized over many future survey configurations, making the approach highly attractive for upcoming large‑scale weak‑lensing surveys (LSST, Euclid, Roman).

The paper also discusses limitations. The current study only varies the redshift distribution; other systematic effects such as baryonic feedback, intrinsic alignments, and photometric‑redshift errors are not included. The task distribution used for meta‑training may not fully capture the diversity of real survey conditions, and the scalability of MAML to higher‑dimensional parameter spaces (e.g., modified gravity, non‑linear clustering) remains to be demonstrated. Future work is suggested in three directions: (i) expanding the task set to include a broader suite of systematics and cosmological models, (ii) integrating the meta‑emulator with iterative training schemes (e.g., CONNECT) to further accelerate adaptation, and (iii) quantifying uncertainties in the meta‑parameters themselves to propagate them into final cosmological constraints.

In summary, this work provides the first demonstration that meta‑learning, specifically MAML, can be successfully applied to cosmological emulation. It shows that a single meta‑trained network can be rapidly adapted to new lensing kernels with minimal data, delivering both high predictive accuracy and faithful parameter inference. This methodology promises to democratize fast cosmological analysis, reducing dependence on massive HPC resources and enabling more flexible, survey‑agnostic pipelines for the next generation of cosmological experiments.