Emulation of SPHEREx Galaxy Power Spectra I: Neural Network Details and Optimization

Abstract

We present neural networks to generate redshift-space galaxy power spectrum multipoles for multiple tracer and redshift bins simultaneously given a set of input cosmology and galaxy bias parameters. This emulator utilizes a combination of fully-connected layers and transformer architecture to accurately predict galaxy power spectrum multipoles $900$ times faster than the SPHEREx pipeline. We quantify network performance using both $Δχ^2$, and likelihood contours for simulated SPHEREx analyses, using two correlated tracer bins and two independent redshift bins. After optimizing network architecture, the loss function, and training set sampling strategy, we achieve $\operatorname{Med}\left( Δχ^2\right) = 0.069$ when comparing to our testing set. At the contour-level our emulator agrees with EFT predictions over a realistic parameter range, with an average 1D best-fit shift of $0.078σ$ and $0.82 \%$ change in 1D error bars. These results demonstrate the feasibility of using neural-network emulators to accelerate SPHEREx redshift-space power-spectrum analyses.

Figures (7)

• Figure 1: (Left): Schematic of our full galaxy power spectrum emulator. Cosmology and galaxy bias parameters are normalized and passed to six independent neural networks, each responsible for generating the auto or cross power spectrum for a particular tracer and redshift bin. We then concatenate the output from these networks to output the galaxy power spectrum for two tracer and two redshift bins. (Right): Architecture of an individual network in the full emulator. Relevant parameters first pass through a series of four MLP blocks, followed by an embedding layer. After that, we proceed through one transformer encoder block using the same architecture as Ref. LSST-emulator-2, before finally outputting normalized multipoles for a specific redshift and tracer bin.
• Figure 2: (Top): Histogram of $\Delta \chi_{tot}^2$ values for emulators trained on training sets sampled with a Latin hypercube (gray) and a uniform hypersphere (purple). We see two orders of magnitude difference between the two sampling strategies. (Bottom): Median $\Delta \chi_{tot}^2$ with respect to the total number of training samples for each sampling strategy. Circled points correspond to the histograms in the top panel. In both cases, performance improves roughly following a power law.
• Figure 3: (Left): Histograms of $\Delta \chi_{(z, tt)}^2$ for each individual network in the final optimized emulator. Each panel denotes a specific redshift bin, and we denote auto/cross power spectra with solid and dashed lines respectively. (Right): Histogram of $\Delta \chi_{tot}^2$ calculated from the full optimized emulator output, which corresponds to the purple hypersphere-sampled histogram in Fig.\ref{['fig:tset_size']}.
• Figure 4: Heatmap showing the median $\Delta \chi_{tot}^2$ of our emulator throughout parameter space, Each pixel represents the median value from samples in the test set whose parameters fall in that specific range, represented via Eq. \ref{['eq:heatmap_calc']}. Other than some regions along the edges, the emulator performs similarly across the entirety of parameter space.
• Figure 5: Parameter contours from simulated likelihood analyses with one tracer and one redshift bin, using pure EFT calculations (gray) and our power spectrum emulator (red). Dashed lines indicate the fiducial cosmology at which our simulated data vector is generated. Both contours lie almost directly on top of each-other, indicating our emulator recovers the expected posterior distribution.