Design and optimization of neural networks for multifidelity cosmological emulation

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.

💡 Research Summary

The paper introduces T2N‑MusE, a neural‑network‑based framework for multifidelity (MF) cosmological emulation that dramatically improves upon previous Gaussian‑process (GP) approaches. The authors address the challenge of building accurate, fast surrogates for the matter power spectrum across a ten‑dimensional cosmological parameter space, using the Goku‑W simulation suite which provides 564 low‑fidelity (LF) and 21 high‑fidelity (HF) N‑body runs. The key innovations are fourfold: (i) a modified two‑step MF architecture that learns the ratio of HF to LF spectra rather than a direct mapping, thereby reducing the input dimensionality of the second network from d_in + d_out to d_in; (ii) a two‑stage Bayesian hyper‑parameter optimization that first performs a coarse global search and then fine‑tunes around the best candidates, using Gaussian‑process surrogate models and k‑fold cross‑validation to evaluate each configuration; (iii) a two‑phase k‑fold training strategy where the data are split into five folds, each fold is trained independently, and the final prediction is obtained by averaging the five models, which mitigates over‑fitting especially when HF data are scarce; and (iv) a per‑redshift principal component analysis (PCA) compression that applies PCA separately at each redshift, preserving nonlinear redshift evolution better than a global PCA.

The workflow proceeds as follows: the 384‑dimensional output (six redshifts × 64 k‑modes) is first compressed using either global or per‑z PCA; the LF network (NN_L) is trained on the compressed LF data; the ratio network (NN_LH) is trained on the HF/LF ratio r = y_H / y_L as a function of the cosmological parameters only; after hyper‑parameter optimization, the two networks are combined to predict HF spectra via ŷ_H = r̂ · ŷ_L. The authors evaluate each component’s contribution through ablation studies, showing that the ratio‑based architecture alone reduces mean absolute error (MAE) from ~2.8 % to ~1.1 %; adding Bayesian optimization brings MAE to ~0.6 %; the two‑phase k‑fold reduces the worst‑case error from ~1.5 % to ~1.2 %; and per‑z PCA yields a modest additional MAE improvement.

In leave‑one‑out cross‑validation, the final emulator (named GokuNEmu) achieves a mean relative error of <0.4 % and a worst‑case error of <1.2 %, compared with the previous GP‑MF model’s 2.2 % mean and 9.6 % worst errors. Training on an NVIDIA A100 GPU takes roughly three hours, including hyper‑parameter searches, while inference per sample is ~0.2 ms with <1.2 GB memory—orders of magnitude faster and more memory‑efficient than GP methods, which scale cubically with sample size.

The authors discuss the broader implications: the ratio‑based two‑step design efficiently transfers information from LF to HF regimes without inflating the network’s input size; the two‑stage Bayesian search systematically explores the hyper‑parameter space, avoiding manual tuning; the two‑phase k‑fold scheme provides robust uncertainty estimates and stabilizes training with limited HF data; and per‑z PCA respects the nonlinear redshift dependence of the power spectrum, improving compression fidelity.

Finally, the paper concludes that T2N‑MusE sets a new benchmark for MF cosmological emulation and can be extended to other statistics such as the bispectrum, Lyman‑α flux power, or baryonic feedback‑affected observables. Ongoing work includes applying the framework to higher‑order correlators and to simulations with additional physics, demonstrating the method’s scalability and versatility for next‑generation cosmological analyses.