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

Abstract

Deep learning has emerged as a transformative methodology in modern cosmology, providing powerful tools to extract meaningful physical information from complex astronomical data. This paper implements a novel Bayesian graph deep learning framework for estimating key cosmological parameters in a primordial magnetic field (PMF) cosmology 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 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 delivers accurate parameter estimation from CMB maps with PMF contributions while providing reliable uncertainty quantification, enabling robust cosmological inference in the era of precision cosmology.

1. Introduction

Cosmic Microwave Background (CMB) radiation is a cornerstone of modern cosmology as it provides a snapshot of the physics of the early Universe [1]. The CMB is sensitive to physical densities because the evolution of acoustic oscillations in the primordial plasma depends directly on them, rather than merely on their fractional contribution to the total energy density [2]. Precise estimation of key cosmological parameters, such as the physical baryon density ${\omega }_b$, cold dark matter density ${\omega }_c$, reionization optical depth $\tau$, scalar spectral index $n_s$, the amplitude of the primordial power spectrum $A_s$, and Hubble parameter $H_0$, is essential for understanding the composition and evolution of the Universe. In particular, the Hubble parameter, which describes the expansion rate of the Universe, has garnered significant attention due to the so-called Hubble tension [3,4]. This term refers to a persistent and statistically significant discrepancy between the values derived from independent cosmological measurements of the early Universe and those obtained from local data [5]. Primordial magnetic fields (PMFs) [6,7] offer a potential solution to the Hubble tension that is not inherently fine-tuned for that purpose. Previous studies, such as [8,9,10], have shown that PMFs can induce baryon inhomogeneities, which accelerate the recombination process and consequently increase the CMB-inferred value of $H_0$, bringing it into closer agreement with local measurements. The PMF scenario is particularly compelling because it is grounded in well-established theoretical frameworks that predate the Hubble tension and have long been investigated as a possible origin of the large-scale magnetic fields observed in galaxies and clusters; see, e.g., [11,12,13].

Beyond their potential role in addressing cosmological tensions, PMFs serve as valuable probes of early-Universe physics. They are motivated by high-energy processes, such as inflationary magnetogenesis and cosmological phase transitions [14,15,16,17]. As such, constraints on PMFs derived from CMB observations offer a promising avenue for investigating the origin and evolution of cosmic magnetism, while also providing natural predictions in certain extensions of the Standard Model [18]. Furthermore, PMFs imprint distinct signatures on the CMB by sourcing scalar, vector, and tensor perturbations, thereby modifying both temperature and polarization anisotropies [19]. PMFs can generate vector modes and B-mode polarization patterns that are absent in the standard $\Lambda$CDM model, as well as induce Faraday rotation of the CMB polarization plane [20], producing a frequency-dependent signal that upcoming experiments, such as the Simons Observatory, LiteBIRD, and CMB-S4, aim to detect [21].

On the other hand, the choice of inference methodology and statistical framework is a critical determinant of the precision and reliability of constraints on cosmological parameters or PMF models [22]. Robust bounds depend on the accurate modeling of the likelihood, thorough treatment of systematics, well-motivated priors, and a careful handling of parameter degeneracies. In addition, traditional techniques applied to the power spectrum are computationally expensive and may not fully exploit the non-Gaussian information present in the CMB maps. This creates a compelling opportunity for deep learning techniques to provide faster and potentially more accurate estimates by learning directly from high-resolution maps [23].

A primary challenge in applying standard deep learning architectures to CMB analysis is the data’s inherent spherical geometry. Conventional convolutional neural networks, designed for flat grids, introduce distortions and are ill-suited for processing data on the sphere. From that perspective, this work is intended to address this issue using the spherical convolutional neural network DeepSphere [24,25,26]. DeepSphere uses a pixelation scheme to represent spherical data, called HEALPix [27], which divides the sphere into pixels of equal area in a hierarchical way, allowing the neural network to efficiently perform convolutions and analysis of the celestial sphere [28].

Beyond geometric challenges, robust inference requires not only accurate predictions but also reliable uncertainty quantification. Standard deep learning models provide deterministic point estimates, failing to capture the epistemic uncertainty associated with limited data or model limitations. To this end, we integrate Bayesian Neural Networks (BNNs) [29] into our architecture. BNNs learn a probability distribution over the model weights, naturally quantifying predictive uncertainty.

This work is organized as follows: Section 2 provides background on graph neural networks, DeepSphere, and Bayesian Neural Networks. Section 3 details our methods, including the simulation pipeline, data validation, network architecture, and training procedure. Section 4 presents the results and analysis of the model’s performance and uncertainty calibration. Finally, Section 5 and Section 6 conclude the paper and discuss future research directions.

2. Neural Networks

Deep learning has become an essential tool in modern cosmology, transcending the limitations of traditional statistical methods to address the challenges posed by the massive, complex, and high-dimensional nature of CMB [30,31,32]. This research leverages convolutional operations extended to graph neural networks to predict cosmological parameters; specifically, we use the DeepSphere framework [24]. DeepSphere offers an efficient implementation of spherical convolutions by representing the celestial sphere as a graph and performing spectral or spatial graph convolutions, thus preserving the global geometry of the CMB data. Furthermore, we use the Bayesian Neural Networks approach, which learn the mean and variance to define a probability distribution over the weights. This approach explicitly quantifies epistemic uncertainty, reflecting the model’s confidence in its predictions arising from limited data or inherent model limitations [29,33].

2.1. Graph Neural Networks and DeepSphere

Graph neural networks (GNNs) are a class of supervised deep learning models [34] which can learn functions over graphs and are a leading approach for building predictive models on graph-structured data $\mathcal{G}(\mathcal{V},\mathcal{E})$ [35], where nodes $\mathcal{V}$ are objects interconnected by edges $\mathcal{E}$, which are relationships [36]. GNNs represent a formidable progression in deep learning, designed to process data structures in non-Euclidean domains where Convolutional Neural Networks (CNNs) have limitations [37,38]. The success of GNNs in handling arbitrary relational data has led to the broader field of geometric graphs, which aims to generalize deep learning operations to diverse non-Euclidean manifolds [39]. In some domains, the nodes in the graph are embedded in a 3D space; therefore, the coordinates could be used as features of the nodes [35]. In cosmology, objects like dark matter halos and galaxy clusters are represented as nodes connected by edges that trace cosmic filaments and spatial proximity [34]. When working with geometrical graphs, it is important to consider the symmetry of the task with respect to translations and rotations of the frame of reference [35], and to generalize data defined on a curved surface, such as CMB maps with PMFs and weak lensing. The challenge is to develop a method that respects spherical symmetry [30]. This challenge drives the development of Spherical Convolutional Neural Networks (SCNNs) such as DeepSphere (https://github.com/deepsphere (accessed on 19 January 2026)) [24,25,26], which is computationally efficient and well adapted to the spherical geometry of the CMB. This adaptation is essential because the CMB is naturally represented on the two-sphere, ${\mathbb{S}}^2$, and is typically pixelized using the HEALPix scheme [27]. The HEALPix pixels are treated as nodes of a graph, and local neighborhoods are exploited through spherical convolutions defined directly on the pixelization. With many standard CNN operations, such as convolutions and pooling, using filters restricted to being radial, DeepSphere can be made invariant or equivariant to rotation. This ensures that the extracted features are independent of the map’s orientation, a necessary constraint for robust inference. DeepSphere has been shown to be computationally more efficient than harmonic-based methods and yields superior performance in tasks such as cosmological model discrimination, particularly in scenarios involving high noise levels or partial sky coverage, confirming its utility for next-generation surveys [24].

2.2. Bayesian Neural Networks

Standard deep learning models are characterized by their deterministic nature, learning single-point estimates for all model parameters. While it is effective for prediction, those models fail to quantify the epistemic uncertainty inherent in the data and prediction process, a critical requirement for robust scientific inference [40]. Bayesian Neural Networks offer a principled approach to quantifying uncertainty in deep learning models. They treat the model parameters as random variables governed by probability distributions, thereby capturing epistemic uncertainty [41]. Inference in BNNs involves computing the posterior distribution, $p\left(\mathrm{w}\right|\mathcal{D})$, over the weights, $\mathtt{w}$, given the training data $\mathcal{D}$. This is commonly achieved via scalable methods such as variance inference or Monte Carlo Dropout [29,33]. Variance inference is the preferred and scalable method for training BNNs, which finds a distribution $q\left(\mathrm{w}\right|\theta )$ that minimizes the difference to the true posterior $p\left(\mathrm{w}\right|\mathcal{D})$ [41] as well as a prior distribution on the weights $p\left(\mathrm{w}\right)$. This step is reached by minimizing the KullBack–Leibler (KL) divergence between the two distributions, and is equivalent to maximizing the evidence lower bound (ELBO) [29], which includes a term associated with the negative log-likelihood (NLL), as [31]

$$\mathrm{ELBO}\sim \mathrm{KL}-\mathrm{NLL}=\mathrm{KL}\left(\mathrm{q}\left(\mathrm{w}\right|\theta )\right|\left|\mathrm{p}\left(\mathrm{w}\right)\right)-({\displaystyle \frac{1}{2}}log|\Sigma |+{\displaystyle \frac{1}{2}}{(\mathrm{y}-\mu )}^T{\Sigma }^{-1}(\mathrm{y}-\mu )),$$

(1)

where $\Sigma$ is the predicted covariance matrix, $\mathbf{y}$ is the observed data, and $\mu$ is the predicted mean vector of the model. The application of BNNs to the estimation of cosmological parameters has been investigated in recent works [31,32]. The combination of GNNs and BNNs results in a powerful framework, which not only achieves high predictive accuracy but crucially reliable estimates of epistemic uncertainty.

3. Methods

To analyze the simulated CMB with PMF maps and extract cosmological information, we employed DeepSphere, a graph-based convolutional neural network architecture specifically designed for spherical data [24]. Specifically, we make use of the HealpyPseudoConv layers, which implement convolution operations on the sphere by aggregating information from neighboring pixels according to the hierarchical structure of HEALPix. In addition, we employ HealpyPooling layers, which perform a coarsening of the sphere by reducing the resolution defined by nside. This combination of spherical convolution and pooling layers allows the network to efficiently capture both large-scale and small-scale anisotropy patterns of the CMB, while remaining computationally feasible for high-resolution data. Compared to planar projections, the spherical approach has the advantage of avoiding distortions and preserving the isotropy of the signal, which is crucial for cosmological inference tasks [27]. In cosmological applications, DeepSphere has been applied to tasks such as weak lensing map classification, cosmological parameter estimation, and non-Gaussianity detection [24]. In our study, we extend its use to the problem of predicting five cosmological parameters—the cold matter density ${\omega }_{\mathrm{c}}$, the baryon density ${\omega }_{\mathrm{b}}$, the amplitude $A_s$, the amplitude of the PMF $B_{1\mathrm{Mpc}}$, and the magnetic damping scale relation $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$, including PMF contributions directly from simulated CMB maps.

3.1. Simulation of CMB Maps with PMF

We generated CMB temperature and polarization anisotropy maps using a modified pipeline based on CAMB (https://camb.info/ (accessed on 19 January 2026)) [42,43], specifically, the patch for CAMB that computes CMB sourced by PMFs [22,43], and Healpy (https://healpix.sourceforge.io/ (accessed on 19 January 2026)) [27]. The latter was used to pixelize and manipulate the spherical CMB maps and to generate random CMB realizations in Healpix format from the input power spectrum. Following the implementation in [22], CAMB is modified to incorporate the effects of PMFs by treating them as a stochastic, non-helical Gaussian field. The implementation follows a covariant formalism where the magnetic energy–momentum tensor acts as an active source in the linear Einstein–Boltzmann equations. Specifically, the code accounts for scalar, vector, and tensor contributions, scaling the comoving magnetic field as $B\propto a^{-2}$ under the ideal MHD approximation. To ensure physical consistency, the patch recalculates the correlation functions for a power-law spectrum $P_B\left(k\right)\propto k^{n_B}$ including an ultraviolet cutoff at the damping scale $k_D$ and explicitly adds both passive and compensated magnetic modes (generated before and after neutrino decoupling). Furthermore, the code employs asymptotic approximations for the convolution integrals to maintain numerical stability across different spectral index regimes. As a result, this CAMB’s variant generates the full set of angular power spectra (TT, EE, BB, and TE), where the PMF-induced vector and tensor modes provide distinct signatures, particularly in the BB-mode signal at high multipoles. These spectra are then passed to Healpy, which synthesizes them into simulated full-sky CMB maps in the form of Stokes I, Q, and U parameters. This framework incorporates standard $\Lambda \mathrm{CDM}$ physics and assumes a flat Universe. It is restricted to three massless neutrinos and passive, compensated magnetic modes that are mutually uncorrelated [22]. In addition, we assume that primordial magnetic fields were generated within the simplest inflationary magnetogenesis scenario. Under this model, the magnetic spectral index naturally approaches the nearly scale-invariant value $n_b=-2.9$ [44]. The total theoretical angular power spectrum is given by

$${\mathrm{C}}_l^{\mathrm{theor}}={\mathrm{C}}_l^{\mathrm{prim}}+{\mathrm{C}}_l^{\mathrm{pass}}+{\mathrm{C}}_l^{\mathrm{comp}},$$

(2)

comprising three main components: the primordial (non-magnetic) spectrum from standard cosmological perturbations ${\mathrm{C}}_l^{\mathrm{prim}}$, the passive magnetic contribution ${\mathrm{C}}_l^{\mathrm{pass}}$, and the compensated magnetic contribution ${\mathrm{C}}_l^{\mathrm{comp}}$. Scalar perturbations are presented in both the compensated ${\mathrm{C}}_l^{\mathrm{comp}}$ and passive ${\mathrm{C}}_l^{\mathrm{pass}}$ modes; vector perturbations are included in the compensated contribution ${\mathrm{C}}_l^{\mathrm{comp}}$ and tensor perturbations are included in the passive contribution ${\mathrm{C}}_l^{\mathrm{pass}}$ (more details in Appendix A.3). In this paper, we adopt the following notation: full to refer to the dataset including scalar, vector, and passive magnetic contributions, and prim-vp to refer to the dataset excluding the scalar compensated mode. Two independent datasets were constructed, each containing 1000 power spectra with 5 realizations, so that each dataset, full and prim-vp, contains 5000 maps. The cosmological parameters were set to Planck best-fit cosmology [45], with  $h=0.67556$, which is the dimensionless reduced Hubble constant, defined as $h=H_0/100\left({\mathrm{kms}}^{-1}{\mathrm{Mpc}}^{-1}\right)$, the scalar spectral index $n_s$ as $0.9665$, and the magnetic spectral index $n_{\mathrm{B}}$ as $-2.9$.

Table 1. Ranges of cosmological and PMF parameters used for dataset generation and the measured values according to Planck best-fit cosmology results [45]. Here, ${\omega }_{\mathrm{c}}={\Omega }_{\mathrm{c}}{h}^2$ and ${\omega }_{\mathrm{b}}={\Omega }_{\mathrm{b}}{h}^2$ represent the cold dark matter and baryon density parameters, respectively; both are dimensionless. $A_s$ is the amplitude of the primordial scalar perturbations given in ${\mathrm{Mpc}}^{-3}$. $B_{1\mathrm{Mpc}}$ denotes the magnetic field strength smoothed over 1 Mpc scale, expressed in nanogauss (nG). $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$ is the logarithmic ratio of the magnetic damping scale, representing the ratio of the time ${\tau }_{\nu }$ of neutrino decoupling to the time ${\tau }_{\mathrm{B}}$ at which the magnetic field was generated in the early Universe, a dimensionless parameter.

Parameter	Range	Planck Ref.
${\omega }_{\mathrm{c}}$	[0.05, 0.50]	$0.112\pm 0.0012$
${\omega }_{\mathrm{b}}$	[0.005, 0.05]	$0.0226\pm 0.00015$
$A_s$	[1.015, 4.015] $\times {10}^{-9}$	$2.105\pm 0.03\times {10}^{-9}$
$B_{1\mathrm{Mpc}}$	[5, 15]	<5.8 at 95% CL [46]
$\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$	[5, 20]	-

summarizes the ranges of cosmological and primordial magnetic field (PMF) parameters used in the generation of the datasets used in this study. These parameters span values chosen to cover the relevant physical scenarios considered in our analysis.

Table 1 summarizes the cosmological and PMF parameters adopted for dataset generation, along with the corresponding reference values from the Planck best-fit $\Lambda$CDM cosmology [45]. We selected prior ranges significantly wider than the current tight constraints from the Planck mission. This choice allows us to test the robustness and generalization of our neural network: by training over a large parameter volume, the model learns the underlying physical relationships between cosmological parameters and the CMB maps rather than overfitting to a narrow region near the current best-fit values. Such an approach is common in early-stage machine learning applications for cosmology. The PMF-related parameters, $B_{1\mathrm{Mpc}}$ and $\beta$, extend beyond the standard $\Lambda$CDM framework. However, according to Plank results: constrains on primordial magnetic field (2015) [46] and Zucca et al. [22], representative values of $B_{1\mathrm{Mpc}}=5\mathrm{nG}$ and $\beta =17$ were used in their analysis of PMF-induced anisotropies, which lie close to the upper and lower limits of the ranges explored in this study.

3.2. Dataset and Validation

The CMB is modeled as a Gaussian random field, meaning that it is not a unique deterministic sky but a statistical ensemble of possible realizations, all of them characterized by the same angular power spectra $C_l$. While most values in these realizations cluster near the mean, a finite number of modes in the sky, especially at large angular scales where cosmic variance is high inevitably, leads to extreme fluctuations. These statistically plausible outliers are a natural consequence of cosmic variance [45]. The broad parameter ranges chosen for dataset generation (see Table 1) were adopted following the analysis in [22], where they investigated constraints on PMF from Planck mission and South Pole Telescope data. Their choice of ranges is motivated by both observational and theoretical considerations. The damping scale ratio ($\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$), for instance, can capture a variety of physically plausible PMF generation scenarios, including inflationary and phase-transition models. Standard cosmological parameters such as ${\omega }_{\mathrm{c}}$, ${\omega }_{\mathrm{b}}$, and $A_s$ were also varied within extended ranges compared to Planck priors. This approach is essential for exploring degeneracies between PMF effects and baseline cosmology. Furthermore, we implemented a data-cleaning procedure to remove maps with pixel values outside physically reasonable limits. The minimum and maximum pixel values of each image were examined, and a per-pixel threshold criterion was applied to flag anomalous maps. Maps exceeding these thresholds were classified as outliers and excluded from the dataset. Similar pre-processing and anomaly-rejection strategies have been employed in recent works applying deep learning to CMB data, ensuring robust inference of cosmological parameters [47]. We found extreme outliers in some maps with magnitudes up to ${10}^{30}$, which are effectively removed after thresholding, resulting in a clean and coherent data range below ${10}^6$ (more details in Appendix A.2). From the initial 5000 CMB maps generated with CAMB, full reaches 67 anomalous maps discarded, leaving 4933 valid maps, and prim-vp reaches 83 anomalous maps discarded, leaving 4917 valid maps. The cleaned dataset was randomly partitioned into training, validation, and test subsets using a two-step procedure. First, $70\%$ of the data was allocated to the training set, while the remaining $30\%$ was reserved for further splitting. This reserved portion was then evenly divided into validation and test sets, each comprising $15\%$ of the original dataset. In addition, we generated an independent dataset consisting of 400 simulated CMB maps with one realization, which serves as an external validation set for the model. This dataset was built under the same physical assumptions as the training data but with a distinct random seed. The main purpose of this external dataset is to evaluate the generalization capacity when exposed to unseen data distributions [48]. With this, the test dataset measures the model’s generalization during training, the validation dataset monitors overfitting and manages hyperparameter tuning, and the external dataset assesses generalization to unseen data.

3.3. Architecture

The network architecture was inspired by the DeepSphere global regression design [24] and implemented using TensorFlow with HealpyGCNN layers. It was designed to combine spherical convolutional feature extraction with Bayesian dense layers for probabilistic regression. The overall architecture consists of three main stages: normalization layer, hierarchical spherical convolution and pooling layers, and Bayesian dense layers for regression, as depicted in Figure 1.

The input consists of full-sky spherical maps at resolution nside = 1024; then, we create samples by dividing the complete spheres in patches based on HEALPix sampling of order = 2, which corresponds to a second-level subdivision of the HEALPix grid. This procedure effectively group pixels into larger, non-overlapping patches that preserve the topological structure of the sphere. We applied an input normalization layer (N) based on the mean and variance of the training dataset by iterating over all input maps. These statistics were stored and later used to initialize the Normalization layer (https://keras.io/api/layers/preprocessing_layers/numerical/normalization/ (accessed on 19 January 2026)) in Keras which standardizes the inputs. We employed a sequence of HealpyPseudoConv and HealpyPool layers. These layers exploit the hierarchical structure of the HEALPix pixelization. Each block contains a HealpyPseudoConv layer (GC) with $p=1$, meaning that the convolution considers only the immediate neighbors on the sphere, and an increasing number of feature maps ($F_{\mathrm{out}}=64,128,256$), a batch normalization layer (BN), and a ReLU activation. Between blocks, we introduced spherical pooling layers (HealpyPooling, P) with average reduction of the resolution, nside. As the network gets deeper, the number of feature maps ($F_{\mathrm{out}}$) increases from 64 to 128 to 256, while the spatial resolution decreases due to pooling. This is a standard and powerful design pattern, used in architectures such as AlexNet [49] or ResNet [50]. The outputs of the last spherical convolutional block are added using a global average pooling layer (AV). This reduces the feature maps to a compact vector representation that can be passed to fully connected layers. To incorporate predictive uncertainty, we employed three DenseFlipout layers [51] from TensorFlow Probability with 256, 128, and 64 units, each using a ReLU activation. The Bayesian treatment is achieved through variational inference where weights are modeled as distributions with a posterior approximated by the Flipout estimator. The final regression layer is a Bayesian multivariate distribution. Specifically, we used DenseFlipout (https://www.tensorflow.org/probability/api_docs/python/tfp/layers/DenseFlipout (accessed on 19 January 2026)) to output the parameters of a MultivariateNormalTriL (https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/MultivariateNormalTriL (accessed on 19 January 2026)) distribution [52] of dimension five (corresponding to ${\omega }_{\mathrm{c}}$, ${\omega }_{\mathrm{b}}$, $A_s$, $B_{1Mpc}$, and $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$).

3.4. Training

The regression network was trained to predict the five cosmological parameters from simulated CMB maps using a probabilistic approach. The input shape was defined as $(32,N_{\mathrm{pix}},3)$, where 32 corresponds to the mini-batch size (i.e., the number of training examples processed simultaneously); $N_{\mathrm{pix}}$ denotes the number of padded pixels per map and the factor of three corresponds to the channels of the simulation. The model was compiled with the Adam optimizer, using a learning rate of $1\times {10}^{-4}$, and the negative log-likelihood ($\mathtt{NLL}$) along with KL as the loss function. In addition, mean squared error (MSE), mean absolute error (MAE), and the coefficient of determination ($R^2$ score) were monitored as complementary metrics. The $R^2$ score was computed using the r2score scikit-learn (https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html (accessed on 19 January 2026)) implementation, which measures the proportion of variance in the true cosmological parameters that can be explained by the model predictions [53]. It is defined as

$$R^2=1-{\displaystyle \frac{{\sum }_i{(y_i-{\widehat{y}}_i)}^2}{{\sum }_i{(y_i-\overline{y})}^2}},$$

(3)

where $y_i$ are the true parameter values, ${\widehat{y}}_i$ are the corresponding predictions, and $\overline{y}$ is the mean of the true values. A value of $R^2=1$ indicates perfect predictions, while values close to zero imply that the model performs similarly to using the mean of the data.

Regarding callbacks, three training callbacks were included: (1) EarlyStopping with a patience of 25 and automatic restoration of the best weights, (2) ModelCheckpoint for saving the optimal weights during training, and (3) ReduceLROnPlateau, which decreased the learning rate by a factor of 0.6 after five consecutive non-improving validation steps, with a floor of ${10}^{-6}$. These strategies have been widely adopted in deep learning to prevent overfitting and to stabilize optimization in non-convex loss landscapes [54]. Training was performed for 300 epochs with a batch size of 32. The total number of trainable parameters is approximately $3.8\times {10}^5$. The architecture combined Healpy-based convolutional layers for feature extraction with Bayesian dense layers for uncertainty-aware regression, resulting in an end-to-end model capable of mapping spherical CMB realizations to cosmological parameters.

3.5. Uncertainty Calibration

Uncertainties in deep neural networks tend to be under- or over-confident in their predictions. One way to diagnose the quality of uncertainty is through reliability plots [55]. Calibrating neural networks during training is not efficient in most cases because it affects the model’s performance. However, post hoc methods applied after training preserve the precision achieved during the training period [31]. We leverage post hoc calibration to improve the reliability of the uncertainty estimates. In the context of Bayesian regression, well-calibrated uncertainty estimates are essential for ensuring that predictive intervals reflect true empirical coverage. We employed the NetCal calibration framework (https://github.com/EFS-OpenSource/calibration-framework (accessed on 19 January 2026)), which has been extended to support probabilistic regression tasks. These techniques aim to adjust the predictive distribution outcomes such that the nominal coverage of predictive intervals aligns with their observed coverage. In fact, we implemented two post hoc calibration methods introduced in [48], the Variance Scaling and GPNormal. The VarianceScaling method applies a correction factor, $\alpha$, to rescale the predictive variance ${\sigma }_{\mathrm{cal}}^2=\alpha ·{\sigma }_{\mathrm{pred}}^2$ and re-calibrate the uncertainties. GPNormal, on the other hand, uses a Gaussian process to map the predictive mean and variance to calibrated uncertainty estimates, allowing for more flexible non-linear transformation of the uncertainty space [48].

4. Results

In this section we report the results obtained using two simulated datasets that differ in their scalar components: full and prim-vp. Both datasets were used to train and evaluate the same neural network model described previously, with the goal of predicting the five cosmological parameters listed in Table 1. To assess the model’s performance, we evaluated it on the respective validation and test sets and computed two key metrics—mean squared error (MSE) and mean absolute error (MAE)—along with the loss function described by Equation (1). Additionally, an external dataset was employed as a held-out test set. This dataset was never used during training or hyperparameter tuning and serves solely to evaluate the model’s generalization performance.

Table 2. Comparison of performance (Loss, MSE, and MAE) evaluated on the validation, test, and external sets of both datasets. A lower value indicates better performance.

Set	FULL			PRIM-VP
	Loss	MSE	MAE	Loss	MSE	MAE
Validation	5250.468	0.055	0.154	5254.806	0.066	0.174
Test	5250.208	0.050	0.147	5254.857	0.068	0.176
External	5250.396	0.052	0.151	5255.045	0.068	0.176

summarizes the evaluation results for both datasets.

We observe that the model reaches comparable training stability across both datasets. However, the full outperformed the prim-vp dataset across all three metrics. Since MSE penalizes larger errors more heavily, its lower value shows that the model was more effective at suppressing outlier predictions when trained with the full scalar input. These results demonstrate that the scalar input in the full dataset enhances the model’s capacity to learn the underlying mapping between CMB features and cosmological parameters, including the effects of primordial magnetic fields.

Table 3. $R^2$ scores on the test set for each predicted cosmological parameter. The models correspond to full (and passive components) and prim-vp.

Parameter	FULL	prim-vp
${\omega }_c$	77.3%	67.1%
${\omega }_b$	98.5%	97.3%
$A_s$	62.6%	56.2%
$B_{1\mathrm{Mpc}}$	99.2%	98.2%
$\beta ={log}_{10}({\tau }_{\nu }/{\tau }_{\mathrm{B}})$	89.5%	85.6%

reports the $R^2$ for each parameter in both experiments. Beyond the quantitative metrics, these results also highlight how the inclusion of the scalar contribution affects the network’s ability to disentangle the physical dependencies among parameters. The improvement observed in the full configuration, particularly for ${\omega }_{\mathrm{c}}$ and $\beta$, shows that scalar information provides complementary features that help the model to predict with better results. Baryon density and the magnetic field are the best-predicted parameters in both datasets with high precision, while the magnetic damping scale obtained the lowest value, especially in the prim-vp case, which lacks the scalar component. The strong $R^2$ achieved for $B_{1\mathrm{Mpc}}$ in both datasets indicates that the network effectively captures the imprint of the magnetic field amplitude on the CMB spectra, despite the complex interplay between PMF-induced and standard cosmological signals. This result is encouraging, as it shows that machine learning methods can detect magnetic signatures even when trained on broad parameter priors.

Figure 2 shows the predicted values versus the ground truth for each of the five parameters in the full dataset. The red dashed line represents the perfect match. Most predictions align well along this line, especially for $B_{1\mathrm{Mpc}}$ and ${\omega }_b$, confirming the strong performance seen in the $R^2$ scores. It is also worth noting that the relatively moderate performance for $A_s$ and ${\omega }_c$ reflects the inherent degeneracies between these parameters and others in the CMB, a well-known limitation in cosmological parameter estimation. Beyond the $R^2$ scores, the regression performance can also be assessed through the complementary metrics reported in Table 2, the mean squared error (MSE) and mean absolute error (MAE). The low values achieved by both metrics indicate that the model not only captures the global variance of the data (as reflected in the high $R^2$) but also maintains a small average deviation between predictions and ground truth values. This consistency across metrics and the consistency between the validation, test, and external sets supports the model’s ability to generalize over unseen data and reinforces the reliability of its predictions. Together, these results confirm that the proposed architecture achieves both precision and robustness in learning the nonlinear mapping between CMB features and cosmological parameters.

The error bars represent the predicted standard deviation of the output distribution. Parameters with lower $R^2$ tend to have larger uncertainty, indicating the model is also less confident when predictions are less accurate. Nevertheless, the resulting uncertainty estimates may still be uncalibrated due to the model misspecification or approximation errors [56]. In Figure 3 we display the reliability diagrams when we use VarianceScaling and GPNormal calibrations for the five output dimensions. Each plot compares the expected quantile against the observed frequency.

Regarding the uncalibrated curves (see Figure 3a), the parameters with the highest $R^2$ scores—${\omega }_b$ and $B_{1\mathrm{Mpc}}$—show clear overconfidence from 0.2 to 0.8 in the expected quantile. In $B_{1\mathrm{Mpc}}$ the overconfidence is more remarkable in quantile 0.6. The predicted confidence intervals are too narrow and capture fewer true values than expected; even when point predictions are highly accurate, the associated uncertainty estimates can still be miscalibrated. On the other hand, the parameters $A_s$ and $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_B)$, which have lower $R^2$ scores, show well-calibrated, even slightly under-confident behavior of $A_s$ in quantile 0.5, resulting in curves that lie above the diagonal. After applying post hoc calibration (see Figure 3b,c), both VarianceScaling and GPNormal improve the alignment between predicted confidence levels and empirical frequencies, thereby reporting reliable uncertainties for our results. After applying the calibration procedure, it is important to fine-tune the model incorporating the recalibration adjustments to ensure that these corrections are consistent with the learned representations. Otherwise, post hoc corrections may lead to a mismatch between the adjusted uncertainty estimates and the actual error distribution. This principle has been discussed in the recent literature, where integrating calibration directly into the training loop has been shown to improve both reliability and generalization [57].

Importantly, these calibration effects have direct implications for the physical interpretation of the inferred PMF constraints. In the uncalibrated case, the overconfident uncertainty estimates on $B_{1\mathrm{Mpc}}$ would lead to artificially tight confidence intervals, potentially resulting in spurious detections or overly aggressive exclusions of PMF scenarios. After calibration, the posterior intervals broaden in regimes where the PMF signal is weak or degenerate with cosmic variance and instrumental noise, yielding conservative and statistically meaningful upper limits on $B_{1\mathrm{Mpc}}$. Conversely, when the PMF imprint on the CMB is informative, the calibrated uncertainties contract appropriately, allowing robust exclusion of specific PMF models and enabling principled model comparison. These features are particularly relevant for forthcoming CMB experiments such as the CORE [58] and LiteBIRD [59], where subtle PMF signatures coexist with complex systematics and potential simulation-to-data mismatches. In this context, calibrated uncertainties provide a safeguard against overinterpretation, ensuring that model selection and hypothesis testing are driven by genuine physical information rather than miscalibrated confidence estimates. As a result, the calibrated framework strengthens the robustness and credibility of PMF constraints derived from next-generation CMB observations.

5. Discussion

We discuss the main experimental and architectural challenges found during the implementation of the Bayesian graph neural network. The computational process was executed using two NVIDIA TESLA V100-SXM2 GPUs running on the Google Cloud Platform. The initial architecture suggested by the main DeepSphere paper [24]—i.e., $NN=AV\circ GC\circ \sigma \circ GC\circ P\circ \sigma \circ BN\circ GC$—proved effective for predicting a single cosmological parameter. Here, $GC$ denotes a graph convolutional layer implemented on the HEALPix sphere, $\sigma$ is the activation function (ReLU), $BN$ corresponds to batch normalization, P denotes pooling, and $AV$ represents a global average layer used to aggregate the node features into a single vector. The normalization layer showed instability and tended toward infinite values because of large outputs in the input data. After the data cleaning phase to remove these outliers, the normalization layer stabilized and assumed reliable values. Encouraged by the promising performance of the model when predicting a single cosmological parameter, we generated a new dataset to explore the joint prediction of two parameters by mixing various PMF components. Nevertheless, the original architecture was not scalable and failed to converge when trained to predict multiple parameters simultaneously. To address this issue, the architecture was modified by adding a similar block to the original structure, and the number of convolutional filters was increased. Regarding the pooling operation, experimenting with both Average and Maximum strategies revealed a performance dependency on the output size. The Max pooling yielded better results for predicting a single parameter, whereas AVG pooling provided superior performance for the simultaneous prediction of multiple cosmological parameters. Subsequent success with two parameters lead to the creation of datasets for three parameters and then four parameters. In this scenario, three dense layers were applied instead of two for better performance. We identified and selected the two best performing datasets, full and prim-vp, chosen based on preliminary results that showed the highest $R^2$ scores across the four parameters (see

Table A1. Coefficient of determination ($R^2$) for the prediction of four cosmological parameters across all simulated datasets. Each column corresponds to one target parameter: the cold dark matter density (${\omega }_c$), the scalar amplitude ($A_s$), the magnetic field strength smoothed at $1\mathrm{Mpc}$ ($B_{1\mathrm{Mpc}}$), and magnetic damping scale $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_B)$. The first column reports the mean $\langle R^2\rangle$ over all parameters.

Datasets	$\langle {\mathit{R}}^2\rangle$	${\mathit{\omega }}_{\mathit{c}}$	${\mathit{A}}_{\mathit{s}}$	${\mathit{B}}_{1\mathbf{Mpc}}$	$\mathit{\beta }={log}_{10}({\mathit{\tau }}_{\mathit{\nu }}/{\mathit{\tau }}_{\mathit{B}})$
full	0.894	0.911	0.758	0.993	0.913
prim-vp	0.880	0.891	0.842	0.923	0.863
svp	0.860	0.876	0.660	0.992	0.912
vp	0.800	0.959	0.258	0.996	0.985
prim-sv	0.783	0.920	0.956	0.977	0.281
prim-v	0.775	0.951	0.980	0.990	0.179
sv	0.541	0.955	0.069	0.992	0.148
prim-s	0.229	0.351	0.888	−0.119	−0.206
sp	0.209	−0.040	−0.102	0.725	0.251
prim-sp	0.054	−0.011	0.043	0.181	0.005

of Appendix A) which were then used to generate the final datasets for the prediction of the cosmological parameters. The training was conducted using a batch size of 32 and a learning rate of $1\times {10}^{-4}$. The training time was on average 8 s per step for 108 steps per epoch, resulting in approximately 860 s per epoch. In contrast to the deterministic model, the probabilistic model exhibited superior training stability and did not require stabilization techniques such as ReduceLROnPlateau or EarlyStopping to avoid divergence. The final model has 382,868 parameters. We have published a subset of the dataset on Zenodo (https://doi.org/10.5281/zenodo.17400503 (accessed on 20 October 2025)) for public access and further use, and the scripts used in this analysis are publicly available in the following repository BSCNN (https://github.com/JavierOrjuela/Bayesian-SphericalCNN (accessed on 19 January 2026)).

6. Conclusions

In this work, we have introduced a novel Bayesian graph deep learning framework for estimating key cosmological parameters from simulated full-sky CMB maps, explicitly incorporating contributions from PMFs through their passive and vector modes. By extending the DeepSphere spherical convolutional neural network with Bayesian Neural Network layers, our approach simultaneously leverages the geometry of spherical data and provides rigorous uncertainty quantification. The main conclusions of our study can be summarized as follows:

(i)

High predictive accuracy through integration of physical characteristics. The proposed B-GNN model achieved $R^2$ scores close to $90\%$ for most of the five cosmological parameters considered: ${\omega }_b$, ${\omega }_c$, $A_s$, $B_{1\mathrm{Mpc}}$, and $\beta ={log}_{10}({\tau }_{\nu }/{\tau }_B)$. Moreover, we reported a mean absolute error (MAE) of $0.151$ and a mean squared error (MSE) of $0.052$ on an independent test set, confirming its strong generalization. Furthermore, the uncertainty estimates provided by the BNN were excellently calibrated, as shown in Figure 3. We found that the inclusion of PMF-induced passive and vector modes was crucial for improving the model’s representational power. This is clearly reflected in the superior performance of the full dataset compared to prim-vp, demonstrating that embedding relevant physical information in the input enhances the accuracy and stability of the predictions.

(ii)

Reliable and calibrated uncertainty quantification. A central contribution of this work is the incorporation of predictive uncertainty estimates. Although the raw uncertainty outputs of the probabilistic model exhibited over-confidence, applying post hoc calibration methods such as VarianceScaling and GPNormal effectively corrected this behavior. The resulting calibrated confidence intervals accurately matched the empirical error distribution, thereby increasing the interpretability and reliability of the model’s predictions. This is particularly important in cosmology, where well-calibrated uncertainties are essential for meaningful parameter inference.

(iii)

Strong generalization and architectural robustness. The model showed stable performance across independent validation datasets, confirming its ability to generalize beyond the training distribution. This robustness is the result of a carefully designed training strategy, including regularization, optimal learning rate scheduling, and dropout. Importantly, unlike the deterministic baseline, the Bayesian model maintained stability without requiring additional early stopping or learning-rate reduction techniques, further supporting its suitability for real data applications.

(iv)

Advancement of cosmological parameter inference with Bayesian deep learning. Compared to traditional parameter estimation techniques, this framework integrates a spherical CNN architecture with Bayesian inference principles, enabling both accurate point predictions and principled uncertainty estimation within a single model. Once trained, the network provides near-instantaneous parameter estimates, bypassing the computational cost of traditional MCMC methods.

(v)

The approach presented in this paper is best viewed as complementary to traditional methods. In fact, hybrid schemes using Bayesian Neural Networks to accelerate MCMC sampling preserve Bayesian rigor while dramatically improving efficiency [23]. However, recent proposals might be advantageous. For example, some neural-network-based methods enable near-optimal map-level inference while naturally incorporating masking, inhomogeneous noise, and non-Gaussianities. Examples include neural implementations of Wiener filtering that match conjugate-gradient accuracy at a fraction of the cost and improve upon pseudo-correlators estimators for complex sky masks [60], as well as spherical CNNs that constrain primordial non-Gaussianity directly from maps, avoiding bispectrum-based information loss [61]. Simulation-based inference further bypasses explicit likelihood assumptions, allowing systematics and nonlinearities to be propagated consistently into posterior constraints, yielding tighter bounds than power-spectrum-only analyses in weak lensing [62] and 21 cm cosmology [63]. The benefits are especially pronounced for non-Gaussian, high-dimensional signals such as 21 cm tomography, where map-level neural inference exploits information beyond two-point statistics [64,65]. Bayesian Neural Networks and normalizing flows further enable calibrated uncertainties and fast likelihood-free inference at a fraction of the computational cost of conventional pipelines [66].

(vi)

Outlook and implications for future applications. Our results indicate that a well-calibrated Bayesian Neural Network can robustly infer cosmological parameters from complex, high-dimensional CMB data that include subtle physical effects such as primordial magnetism. This framework offers a flexible and powerful alternative to conventional inference pipelines, providing not only point estimates but also interpretable uncertainty information. These characteristics make it a promising tool for future analyses for the next generation of CMB experiments, such as the PICO [67], LiteBIRD [59], and CORE [58]. These missions will deliver ultra-deep, high-resolution polarization maps where the PMF’s B-mode signal will be a critical and computationally challenging component to isolate. Our work establishes a foundation for using advanced machine learning to harness this data, offering a powerful and complementary tool to traditional methods for probing fundamental physics in the early Universe.

In summary, this study demonstrates that probabilistic deep learning methods, when carefully designed and systematically calibrated, can significantly enhance the precision and reliability of cosmological parameter estimation from CMB data. The integration of physical information, spherical geometry, and Bayesian inference principles offers a new and powerful pathway to tackle open questions in cosmology and to explore beyond-standard physics.