Design and optimization of neural networks for multifidelity cosmological emulation

TL;DR

The paper tackles the computational bottleneck of cosmological emulation on nonlinear scales by replacing Gaussian-process surrogates with a neural-network–based multifidelity framework, T2N-MusE. It introduces a 2-step multifidelity architecture, a 2-stage hyperparameter optimization, a 2-phase training strategy, and per-redshift PCA to handle high-dimensional outputs, demonstrating substantial accuracy gains on the Goku simulation suite. The results show mean LOOCV error reductions by over a factor of 5 and worst-case reductions by ~8× compared with prior GP-based approaches, culminating in the production of GokuNEmu, a highly capable matter power spectrum emulator. The approach provides scalable, efficient training for large, high-dimensional cosmological parameter spaces and can be extended to other summary statistics with publicly available code.

Abstract

Accurate and efficient simulation-based emulators are essential for interpreting cosmological survey data down to nonlinear scales. Multifidelity emulation techniques reduce simulation costs by combining high- and low-fidelity data, but traditional regression methods such as Gaussian processes struggle with scalability in sample size and dimensionality. In this work, we present T2N-MusE, a neural network framework characterized by (i) a novel 2-step multifidelity architecture, (ii) a 2-stage Bayesian hyperparameter optimization, (iii) a 2-phase $k$-fold training strategy, and (iv) a per-$z$ principal component analysis strategy. We apply T2N-MusE to selected data from the Goku simulation suite, covering a 10-dimensional cosmological parameter space, and build emulators for the matter power spectrum over a range of redshifts with different configurations. We find the emulators outperform our earlier Gaussian process models significantly and demonstrate that each of these techniques is efficient in training neural networks or/and effective in improving generalization accuracy. We observe a reduction in the mean error by more than a factor of five and in the worst-case error by approximately a factor of eight in leave-one-out cross-validation, relative to previous work. This framework has been used to build the most powerful emulator for the matter power spectrum, GokuNEmu, and will also be used to construct emulators for other statistics in future.

Figures (13)

• Figure 1: Examples of the original and modified 2-step MF NN architectures. Both architectures have the same $NN_\mathrm{L}$ (step 1: the LF NN) but different $NN_\mathrm{LH}$ (step 2: the NN used to correct the LF output). The original $NN_\mathrm{LH}$ (a) approximates the correlation between the LF and HF functions, with $(\mathbf{x}, \mathbf{y}^\mathrm{L})$ as input and $\mathbf{y}^\mathrm{H}$ as output. The modified $NN_\mathrm{LH}$ (b) learns the mapping from $\mathbf{x}$ to the ratio of $\mathbf{y}^\mathrm{H}$ to $\mathbf{y}^\mathrm{L}$, $\mathbf{r} = \mathcal{G} (\mathbf{x})$, and the final HF output is the element-wise product of the LF output with the correction ratio $\mathbf{r}$.
• Figure 2: Overview of the workflow of training a highly optimized NN. The workflow consists of three main steps: data compression, hyperparameter optimization, and training the final model. When optimizing the hyperparameters, a large space is explored in the initial search stage, and then a smaller space is searched in the second stage (fine-tuning). Each evaluation of the hyperparameters involves both training and validation of an NN.
• Figure 3: Illustration of the two-phase $k$-fold training and cross-validation strategy for $NN_\mathrm{L}$, assuming a total of 9 samples, of which 3 (orange circles) are supposed to be tested against (i.e., the HF cosmologies). In phase 1, the model is trained on the remaining 6 samples (blue circles) using 3 separate runs with different random seeds, and validated on the 3 held-out test samples. In phase 2, we perform regular $k$-fold training and validation, with the initial model (weights and biases) set to the best model found in phase 1.
• Figure 4: LOO errors of the emulators built with approaches Base, Mid, and Optimal. Redshifts are color coded. The solid lines are the error averaged over cosmologies, and the corresponding shaded regions indicate the range of individual cosmologies. The gray-shaded area marks the region where the error is less than 1%. Each model is titled with the name of the approach and its overall validation error.
• Figure 5: Summary comparison of the LOO errors of the emulators built with different approaches. Black crosses indicate the mean validation error of each approach, while gray crosses show the worst-case errors (the maximal error over all test points).