Bayesian Inference of Primordial Magnetic Field Parameters from CMB with Spherical Graph Neural Networks

TL;DR

This work introduces a Bayesian graph deep-learning framework that extends the spherical CNN DeepSphere with Bayesian neural network components to infer five cosmological parameters, including PMF-related ones, from simulated full-sky CMB maps. By training on maps that incorporate PMF effects through passive and vector modes, the model achieves near $R^{2}=0.9$ on key parameters and provides predictive distributions via DenseFlipout and a MultivariateNormalTriL output. Post-hoc uncertainty calibration using Variance Scaling and GPNormal yields well-calibrated predictive intervals, addressing over/underconfidence in the Bayesian estimates. The approach preserves the spherical geometry of the data, demonstrates robust generalization, and offers a scalable, uncertainty-aware tool for PMF-enabled cosmological inference in the era of precision CMB observations.

Abstract

Deep learning has emerged as a transformative methodology in modern cosmology, providing powerful tools to extract meaningful physical information from complex astronomical datasets. This paper implements a novel Bayesian graph deep learning framework for estimating key cosmological parameters in a primordial magnetic field (PMF) cosmology directly from simulated Cosmic Microwave Background (CMB) maps. Our methodology utilizes DeepSphere, a spherical convolutional neural network architecture specifically designed to respect the spherical geometry of CMB data through HEALPix pixelization. To advance beyond deterministic point estimates and enable robust uncertainty quantification, we integrate Bayesian Neural Networks (BNNs) into the framework, capturing aleatoric and epistemic uncertainties that reflect the model confidence in its predictions. The proposed approach demonstrates exceptional performance, achieving $R^{2}$ scores exceeding 0.89 for the magnetic parameter estimation. We further obtain well-calibrated uncertainty estimates through post-hoc training techniques including Variance Scaling and GPNormal. This integrated DeepSphere-BNNs framework not only delivers accurate parameter estimation from CMB maps with PMF contributions but also provides reliable uncertainty quantification, providing the necessary tools for robust cosmological inference in the era of precision cosmology.

Figures (6)

• Figure 1: Representation of the neural network architecture used for cosmological parameter inference. The input consists of full-sky spherical maps at resolution nside = 1024, which are first normalized using a normalization layer. The architecture includes three convolutional blocks with HealpyConv layers, each followed by batch normalization (BN), ReLU activation, and average HealpyPooling layers. The third block has GlobalPooling. The extracted features are passed to Bayesian dense layers, implemented using DenseFlipOut, followed by a MultivariateNormalTriL output layer that predicts the posterior distribution of the target parameters: cold dark matter density $\omega_{\mathrm{c}}$, baryon density $\omega_{\mathrm{b}}$, scalar amplitude $A_{s}$, magnetic field strength $B_{\mathrm{1Mpc}}$, and damping scale ratio $\beta = \log_{10}(\tau_{\nu}/\tau_{\mathrm{B}})$.
• Figure 2: Predicted versus ground truth true values for the five cosmological parameters using the full. Each subplot corresponds to one parameter, with the black dashed line representing the ideal prediction. Shaded bars denote one standard deviation.
• Figure 3: Uncertainty calibration curves (reliability diagrams) for each predicted parameter using the full dataset. The plots compare predicted Excepted quantile (x-axis) with the observed frequencies (y-axis). Perfect calibration corresponds to the dashed red diagonal. (a) Uncalibrated model outputs, (b) Post-hoc calibration using VarianceScaling. (c) Post-hoc calibration using GPNormal.
• Figure 4: Scatter plots of the minimum vs. maximum pixel values per image before and after outlier removal. The left panel shows the raw data, where extreme outliers reach values up to $1\times10^{30}$. Panels (a) correspond to the full dataset, and panels (b) to the prim-vp dataset. The left panels show the distribution of the minimum and maximum pixel values per image before cleaning. The right panel displays the cleaned dataset after applying the thresholds with value $1\times10^{6}$. Dashed lines indicate the minimum and maximum thresholds used for cleaning.
• Figure 5: Predicted versus ground truth true values for the five cosmological parameters using the prim-vp. Each subplot corresponds to one parameter, with the black dashed line representing the ideal prediction. Shaded bars denote one standard deviation.