Mitigating Simulator Dependence in AI Parameter Inference for the Epoch of Reionization: The Importance of Simulation Diversity

TL;DR

The paper tackles simulator bias in AI-based EoR parameter inference from the 21cm signal. It tests whether training on data from four simulators (three 21cmFAST variants and zreion) improves generalization to unseen simulators by predicting $\Delta z = z_{25}-z_{75}$ from $30\times256\times256$ lightcones using a CNN. Results show substantial gains in out-of-distribution accuracy as simulators are added to the training set, with average $\mathrm{MSE}$ dropping from ~0.32 for single-simulator models to ~0.11 for models trained on three simulators, demonstrating the value of simulation diversity to mitigate simulator-specific biases. The findings suggest building large, diverse training datasets spanning multiple reionization models to improve robustness when applying AI-based inferences to real 21cm observations.

Abstract

The 21cm signal of neutral hydrogen contains a wealth of information about the poorly constrained era of cosmological history, the Epoch of Reionization (EoR). Recently, AI models trained on EoR simulations have gained significant attention as a powerful and flexible option for inferring parameters from 21cm observations. However, previous works show that AI models trained on data from one simulator fail to generalize to data from another, raising doubts about AI models' ability to accurately infer parameters from observation. We develop a new strategy for training AI models on cosmological simulations based on the principle that increasing the diversity of the training dataset improves model robustness by averaging out spurious and contradictory information. We train AI models on data from different combinations of four simulators, then compare the models' performance when predicting on data from held-out simulators acting as proxies for the real universe. We find that models trained on data from multiple simulators perform better on data from a held-out simulator than models trained on data from a single simulator, indicating that increasing the diversity of the training dataset improves a model's ability to generalize. This result suggests that future EoR parameter inference methods can mitigate simulator-specific bias by incorporating multiple simulation approaches into their analyses.

Figures (9)

• Figure 1: A comparison of the different algorithms used in this work to simulate the evolution of the EoR brightness temperature field. For this comparison, each EoR history uses the same underlying density field, though in the actual datasets no two instances share the same initial density field. For this figure, input parameters were adjusted in order to produce instances with similar reionization histories while still showcasing the qualitative differences between simulators. For example, note how Dataset CV results in smaller more jagged bubbles while Dataset FS produces comparatively larger and rounder bubbles, reflecting their respective bubble-flagging methods.
• Figure 2: Distribution of $\Delta z$ across all four datasets. It is important that these distributions overlap, as AI models traditionally struggle to predict on label ranges outside that of their training data (Sooknunan2024).
• Figure 3: A simple diagram of our CNN. Our model consists of three 2D convolutional downsampling layers, a global pooling layer, and three linear layers, with a resulting scalar output. Our model takes data cubes of size $30\times256\times256$ as inputs to the network. Each cube is made up of a set of 30 brightness temperature slices sampled evenly along the redshift axis, treated as color channels such that the convolutional kernel convolves over the xy plane.
• Figure 4: Predicted vs. true $\Delta z$ of single-dataset models across all four datasets. Note the clear difference between in-distribution and out-of-distribution performance. For the models trained on Datasets CV, FS, and ZR, the model will not predict values of $\Delta z$ below the minimum value seen during training, resulting in a horizontal line of predictions at the minimum. This is likely due to the sparsity of the trained model's latent space, an issue that could be rectified by tweaking the model architecture or further hyperparameter optimization, though doing so is unlikely to result in a significant boost in accuracy.
• Figure 5: Predicted vs. true $\Delta z$ for models trained on two datasets.