Inferring Cosmological Parameters with Evidential Physics-Informed Neural Networks

Abstract

We examine the use of a novel variant of Physics-Informed Neural Networks to predict cosmological parameters from recent supernovae and baryon acoustic oscillations (BAO) datasets. Our machine learning framework generates uncertainty estimates for target variables and the inferred unknown parameters of the underlying PDE descriptions. Built upon a hybrid of the principles of Evidential Deep Learning, Physics-Informed Neural Networks, Bayesian Neural Networks, and Gaussian Processes, our model enables learning the posterior distribution of the unknown PDE parameters through standard gradient-descent-based training. We apply our model to an up-to-date BAO dataset (Bousis et al. 2024) calibrated with the CMB-inferred sound horizon, and the Pantheon+ Sne Ia distances (Scolnic et al. 2018), examining the relative effectiveness and mutual consistency among the standard $\Lambda$CDM, wCDM and ${\Lambda }_s$CDM models. Unlike previous results arising from the standard approach of minimizing an appropriate ${\chi }^2$ function, the posterior distributions for parameters in various models trained purely on Pantheon+ data were found to be largely contained within the $2\sigma$ contours of their counterparts trained on BAO data. Our study illustrates how a data-driven machine learning approach can be suitably adapted for cosmological parameter inference.

1. Introduction

In this paper, we present a framework based on using a neural network to infer cosmological parameters from the Pantheon+ dataset of [1] and a recent collection of the BAO dataset presented in [2]. The neural network model we use is a surrogate model quantifying the luminosity distance L vs. redshift z relationship, and its weight parameters are obtained through maximizing the degree of adherence to the following one-dimensional ODE

$${\displaystyle \frac{dL}{dz}}-{\displaystyle \frac{L}{1+z}}-{\displaystyle \frac{c(1+z)}{H(z;\overrightarrow{\Omega })}}=0,$$

(1)

where $H(z;\overrightarrow{\Omega })$ is the Hubble function parametrized by $\overrightarrow{\Omega }$ and c is the speed of light. Equation (1) follows from the defining relation

$$L=c(1+z){\int }_0^{z}d\tilde{z}{\displaystyle \frac{1}{H(\tilde{z};\overrightarrow{\Omega })}},$$

(2)

where L is in units of Mpc, and its form is more convenient for us to infer unknown parameters $\overrightarrow{\Omega }$ of $H(z;\overrightarrow{\Omega })$ which depend on the underlying cosmological model assumed. In this work, we will consider two classes of deviations from the standard $\Lambda$CDM model described via the wCDM and ${\Lambda }_s$CDM models. The former refers to the standard $\Lambda$CDM model but with the equation of state parameter for dark energy w being not necessarily −1 (see ref. [3] for a recent study). In fitting it to the data, we take the free parameters of the wCDM model to be $\{H_0,{\Omega }_m,w\}$, where $H_0,{\Omega }_m$ are the Hubble constant and total matter density, respectively. For the ${\Lambda }_s$CDM model [4], here we take its free parameters to be $\{H_0,{\Omega }_m,z_t\}$, where $z_t$ is a transition redshift value from which the cosmological constant switches sign, representing a toy model of vacua transition from anti-de Sitter to de Sitter spacetime at some point in the early universe. These models are parametrically deformable to the standard $\Lambda$CDM model in the limits of $w\to -1$ for the wCDM model and $z_t\to \infty$ for the ${\Lambda }_s$CDM model.

In standard regression techniques invoking the principle of maximum likelihood estimation, cosmological parameters are inferred through minimizing a ${\chi }^2$ likelihood of the form (see [5])

$$-2log\left(\mathcal{L}\right)={\chi }^2=\Delta {\overrightarrow{D}}^T{C}^{-1}\Delta \overrightarrow{D},$$

(3)

where C is the covariance matrix expressing uncertainties and $\Delta D_k\equiv L_k-L_{model}\left(z_k\right)$ is the parameter residuals, with $L_k$ being an observed value and $L_{model}\left(z_k\right)$ being the corresponding theoretical estimate computed with Equation (2).

A fundamental difference between using (3) and a neural network-based approach is that the latter is structured around a surrogate model $\mathcal{M}\left(z\right)$ that represents the target variable as a function of the input variable, apart from an inference of the unknown parameters. Analogous to the pure numerical solution equipped with the best-fit parameters that minimizes (3), the final model $\mathcal{M}\left(z\right)$ is characterized by a set of network parameters that correspond to the minimum of a loss function that generalizes (3). For a multilayer-perceptron model trained using just a mean-squared error loss term, model training then translates to solving (3) through a gradient-descent-based approach with $L_{model}$ being the neural network. For more complicated frameworks such as that of Physics-Informed Neural Networks (PINNs) [6] where PDE constraints are simultaneously imposed, the loss function can be much more complex than (3). In this work, we examine the use of $\mathcal{M}\left(z\right)$ as an independent data-driven model to infer probability distributions for the parameters from data, in a manner consistent with Bayesian principles. To do so, we need a framework that ideally yields $\mathcal{M}\left(z\right)$ together with its predictive uncertainty. It should also yield the posterior distribution for each unknown parameter of the cosmological model upon completion of model training.

Recently in [7,8], Evidential Physics-Informed Neural Networks (or E-PINNs for short) was proposed as a framework for PDE-based scientific modeling that encapsulates uncertainty quantification robustly. It realizes a hybrid implementation of the algorithms of Evidential Deep Learning [9,10] and those of PINN. In [8], a principled approach was proposed for constructing priors for the unknown parameters and the learnable loss weight of the PDE residual term which is taken as the likelihood function for the unknown parameters. Gradient-descent-based training then translates to the maximum a posteriori learning of the distribution of the unknown parameters and weights of the surrogate model $\mathcal{M}\left(z\right)$. In this paper, we will use E-PINN as the machine learning framework for learning cosmological parameters from the Pantheon+ and BAO datasets.

We examine the differences in the inferred cosmological parameters when E-PINN is trained on these datasets separately and examine how E-PINN differentiates among the alternative cosmological models with respect to each dataset and their synthesis. In the aspect of the model training algorithm, while still leveraging the basic framework of E-PINN as proposed in [8], we also incorporate Gaussian Process Regression [11] within the training algorithm in a few ways to refine the parameter inference process. Gaussian Process is used to guide the construction of the prior distributions for $\overrightarrow{\Omega }$. The predictive variance values provided by Gaussian Process Regression are employed as proxy targets to supervise the learning of epistemic uncertainty in our model. Although our primary motivation for incorporating Gaussian Processes into the framework stems from the relatively small size of the BAO dataset [2], the methods we propose are readily transferable to other scientific modeling problems and extend the versatility of the E-PINN toolkit of [8].

Previous to our work, there have been studies [12,13,14,15,16,17,18] related to the use of neural network-based models for analyzing cosmological data and the inference of parameters. In [12], PINN was applied to the Union 2.1 dataset, and an uncertainty framework was proposed where the perceptron model’s outputs were taken to be the (mean) luminosity distance and its associated uncertainty, with the loss function being the log-likelihood of the Gaussian with the outputs as its moments. For us, following the framework of Evidential Deep Learning, we assert a prior distribution (normal-inverse-gamma) for the mean and variances, integrating them out to obtain a t-distribution as the marginal likelihood. Our model’s outputs then correspond to the learnable parameters of this higher-order distribution. In contrast to our work, in [12], there was no methodology proposed to infer unknown cosmological parameters from the data. In [13], the authors essentially used PINN (see Equation (33) of [13]) and were focused primarily on whether PINN can be used to reproduce numerical solutions of PDE (in the context of cosmological models). Ref. [13] performed parameter inference but it was performed using the standard regression method of minimizing the ${\chi }^2$. In [14,19], no PINN-related formalism was invoked but a perceptron model trained on simulated data generated based on some chosen fiducial values of the cosmological parameters and synthetic noise added to the redshift. The statistical inference was performed using the standard ${\chi }^2$ method as in [13], rather than through a learned posterior distribution supported by a data-informed prior in our framework.

Our paper is organized as follows. We begin by presenting the theoretical formulation of the E-PINN model in Section 2, including how we invoked Gaussian Processes to enhance the original framework of [8]. This is followed by a discussion on methodology such as model training implementation details, metrics, etc., in Section 3. Our main results on the cosmological parameters are collected in Section 4. We end with a summary and some comments on the relevance of our work to the Hubble tension problem [20] in Section 5. Appendix A contains some technical details related to the loss function, including a detailed derivation for the hyperparameters of the prior for the PDE residual loss weight, while Appendix B gathers various corner plots for the posterior distributions predicted by our models. Our study illustrates how a data-driven machine learning approach can be suitably adapted for cosmological parameter inference.

2. Model Formulation

In this section, we introduce the main ideas and practical implementation of E-PINN, and explain how we extend the original algorithm of [8] by incorporating Gaussian Processes to construct the parameters’ prior and supervise learning of the epistemic uncertainty. We refer the reader to [8] for a more technical exposition of E-PINN.

2.1. Basic Structure of the Neural Network and the Loss Function

For our purpose (and for the general context of regression), we take the base neural network of E-PINN to be a multilayer perceptron $\mathcal{M}(z,\overrightarrow{w})$, where z denotes its input, and $\overrightarrow{w}$ its weight parameters. In our work here, z is the redshift value associated with the measured luminosity distance of some astrophysical object/phenomenon. Our model $\mathcal{M}(z,\overrightarrow{w})$ has four output neurons $\{\alpha ,\beta ,\nu ,\gamma \}$, with one of them ($\gamma$) pertaining to the mean of the target (luminosity distance) and the other three related ($\alpha ,\beta ,\nu$) to the predictive uncertainty as follows:

$${\sigma }_p^2={\displaystyle \frac{\beta }{(\alpha -1)}}\left({\displaystyle \frac{1}{\nu }}+1\right).$$

(4)

The uncertainty ${\sigma }_p^2$ is the variance of a probability distribution that describes the statistical fluctuation of the observed target about its mean. This probability distribution is a t-distribution defined by the following density function:

$$P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))={\displaystyle \frac{\Gamma \left(\alpha +\frac{1}{2}\right)}{\Gamma \left(\alpha \right)\sqrt{2\pi \beta (1+\nu )/\nu }}}{\left(1+{\displaystyle \frac{{(L_{obs}-\gamma )}^2}{2\beta (1+\nu )/\nu }}\right)}^{-(\alpha +{\displaystyle \frac{1}{2}})},$$

(5)

where $L_{obs}$ is the observed luminosity distance, and $\gamma$ is the model prediction (to be interpreted as a mean value), with $\alpha ,\beta ,\nu$ being its other parameters which set the scale for its variance (Equation (4)) and other moments. (In Appendix A, we provide more details on how this is related to Gaussian distributions.) In Equation (5), $P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))$ denotes the probability of observing the data $\mathcal{D}$ given the model $\mathcal{M}\left(\overrightarrow{w}\right)$. Here, and henceforth, we suppress the input z for simplicity, and all variables on the RHS of Equation (5) are to be understood as functions of z which is the input variable of the model, with $\alpha ,\beta ,\nu ,\gamma$ being the model’s outputs. For them to be interpreted as what we have described thus far, we need a suitable loss function that is derived based on these principles. Here, we take the negative log-likelihood of Equation (5) to be part of our complete loss function:

$${\mathcal{L}}_{data}=-log\left[P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))\right].$$

(6)

If we do not allude to any PDE constraints, i.e., asserting that Equation (1) describes the data, then ${\mathcal{L}}_{data}$ would constitute our complete loss function, and model training typically means using a gradient-descent-based method to find weights $\overrightarrow{w}$ such that they minimize ${\mathcal{L}}_{data}$. We call the set of points $(z_{obs},y_{obs})$ the training dataset, with model training translating to iterations in the gradient-descent-based algorithm.

Now, let us bring in the putative PDE description for the observed data. For this work, this would be Equation (1) with some ansatz for $H(z;\overrightarrow{\Omega })$ that is in turn derived from cosmological approximations in General Relativity, etc. With the parameters $\overrightarrow{\Omega }$ unknown, we add the following loss term:

$${\mathcal{L}}_{pde}=-log\left[P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\right],P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\sim exp\left[-{\displaystyle \frac{1}{2{\sigma }_R^2}}\sum _{k=1}^{N_D}{\mathcal{R}}_k^2\left(\mathsf{\partial }\gamma ,\gamma ,z_k,\overrightarrow{\Omega }\right)\right],$$

(7)

where k is a subscript denoting each data point, $N_D$ is the number of observations, and suppressing the data point index, we define

$$\mathcal{R}\left(\mathsf{\partial }\gamma ,\gamma ,z,\overrightarrow{\Omega }\right)\equiv {\displaystyle \frac{d\gamma }{dz}}-{\displaystyle \frac{\gamma }{1+z}}-{\displaystyle \frac{c(1+z)}{H(z;\overrightarrow{\Omega })}},$$

which refers to the RHS of Equation (1), with $\gamma$ representing the mean luminosity distance. The parameter ${\sigma }_R^2$ is a learnable parameter that evolves during model training, and from Equation (7), we can see that it acts as a weighting factor quantifying the magnitude of this loss term relative to other loss terms in the complete loss function. Equation (7) guides the model towards adhering to the ODE (1) when added to the complete loss function to be minimized. We also used the symbol $P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })$ to denote the interpretation of this term as the likelihood of obtaining $\mathcal{M}\left(\overrightarrow{w}\right)$ as a surrogate model assuming the parameters $\overrightarrow{\Omega }$.

In standard PINN, ${\sigma }_R^2$ is a free parameter and, to our knowledge, there is no principled approach towards determining its choice. For our E-PINN model, we lift ${\sigma }_R^2$ to be a learnable parameter that evolves as the model shifts towards a minimum in the loss landscape. We regularize this procedure by introducing another loss term that represents the negative logarithm of the prior density function for ${\sigma }_R^2$ as follows:

$${\mathcal{L}}_{\pi \left({\sigma }_R^2\right)}=-log\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r),\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r)\equiv {\displaystyle \frac{{\beta }_r^{{\alpha }_r}}{\Gamma \left({\alpha }_r\right)}}{\sigma }_R^{-2({\alpha }_r+1)}{e}^{-{\displaystyle \frac{{\beta }_r}{{\sigma }_R^2}}},$$

(8)

In Appendix A.3, we furnish more detailed explanations of how (8) is derived, in particular showing that one can set ${\alpha }_r,{\beta }_r$ of (8) such that these values align consistently with other aspects of our framework.

Thus far, we have motivated the inclusion of three different loss terms ${\mathcal{L}}_{data},{\mathcal{L}}_{pde},$ ${\mathcal{L}}_{\pi \left({\sigma }_R^2\right)}$, with each of them being interpretable as the negative logarithm of some density/likelihood function. If we follow the principles of Bayesian statistics, we can further consider incorporating a prior density function $\pi \left(\overrightarrow{\Omega }\right)$ for the unknown parameters. In Appendix A.2, we show that such a prior can be derived from considering the family of numerical solutions to the ODE and the model that is trained purely on data based on ${\mathcal{L}}_{data}$. Corresponding to $\pi \left(\overrightarrow{\Omega }\right)$, we thus introduce another loss term ${\mathcal{L}}_{\pi (\Omega )}=-log\left(\pi (\Omega )\right)$. Taking into account all the loss terms considered so far, we have the loss function

$$\begin{array}{ccc}\hfill \mathcal{L}& =& {\mathcal{L}}_{data}+{\mathcal{L}}_{pde}+{\mathcal{L}}_{\pi \left({\sigma }_R^2\right)}+{\mathcal{L}}_{\pi (\Omega )}\hfill \\ & =& -log\left[P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r)\pi \left(\overrightarrow{\Omega }\right)\right],\hfill \end{array}$$

(9)

with the product $P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })$ being the joint likelihood function for $\overrightarrow{\Omega }$. Upon completion of model training, we can place confidence intervals on the model’s predictions using ${\sigma }_p^2$ of Equation (4). Since we infer the posterior distribution of $\overrightarrow{\Omega }$ at the end of model training, our framework thus appears as a maximum a posteriori estimation of $\overrightarrow{\Omega }$, or more precisely1 a MLE estimation that is regularized by a prior $\pi \left(\overrightarrow{\Omega }\right)$.

From Equation (5), we can see that while the empirical data $L_{obs}$ correlates with the mean output $\gamma$ (through the factor ${(L_{obs}-\gamma )}^2$), there is no other external information supervising the three other uncertainty-related outputs $\alpha ,\beta ,\nu$. In our context, the cosmological datasets are already equipped with uncertainty estimates which we can potentially use to refine model training. In the following Section 2.2, we discuss this issue in detail and introduce Gaussian Processes as a complementary tool to supervise the learning of the model’s uncertainty. For model training, this further introduces a couple more loss terms to be added to (9) for our complete loss function.

2.2. Using Gaussian Processes to Supervise Uncertainties

Although model training can proceed without additional information on data uncertainties, the datasets selected for our work here are already equipped with measurement uncertainties—for the BAO data, these were computed in [2] from raw uncertainties of each sample as collected in Table 1 of [2], whereas for the Pantheon data, we used the diagonal elements of the covariance matrix presented in [21]. They correspond to what is known in the machine learning terminology as ‘aleatoric uncertainties’ which generally refer to uncertainties that are measurement or observation-related such as noise, etc. We add a simple mean-squared-loss term in the form

$${\mathcal{L}}_{alea}=\mathbb{E}\left({\displaystyle \frac{\beta }{\alpha -1}}-{\sigma }_a^2\right),$$

(10)

where ${\sigma }_a^2$ denotes the measurement’s statistical variance for each point, and $\mathbb{E}$ denotes taking the average over all the training samples. In Appendix A, we explain more carefully why the factor of ${\displaystyle \frac{\beta }{\alpha -1}}$ arises in (10).

Now there is another notion of uncertainty known as the ‘epistemic uncertainty’, which pertains to the model itself instead of the data measurement process. In contrast to aleatoric uncertainty, this is a quantity that characterizes the degree of data sufficiency and model complexity. It can be supervised in model training if there is some independent knowledge of the model variance. Here, we use a Gaussian Process Regression model to furnish information on the epistemic uncertainty distribution. A Gaussian Process (GP) is essentially a distribution over functions. Ref. [11] Denoting the GP by $f\left(z\right)$, schematically

$$f\left(z\right)\sim \mathcal{G}\mathcal{P}\left(m\left(z\right),k(z,z^{\prime })\right),$$

where $m\left(z\right)$ is the mean function and $k(z,z^{\prime })$ is the covariance kernel function. Here we took $k(z,z^{\prime })=\exp \left(-{(z-z^{\prime })}^2/2l^2\right)$, a RBF function with a characteristic length scale l that we determine by maximizing the log marginal likelihood using the implementation in scikit-learn [22]. This length scale can be understood as a factor that controls the rate at which the correlation (as modeled by the kernel) decays with increasing separation. A relatively larger choice of l implies a kernel that is smoother and more slowly varying. Conditioned on an observed set of data ${\{z_i,y_i\}}_{i=1}^{N_D}$, the posterior distribution for f evaluated at some arbitrary redshift $\tilde{z}$ is a Gaussian distribution $\mathcal{N}(\tilde{\mu },{\tilde{\sigma }}_e^2)$ with moments

$$\tilde{\mu }=k(\tilde{z},z){[k(z,z)+{\sigma }_a^2\mathbb{I}]}^{-1}y,{\tilde{\sigma }}_e^2=k(\tilde{z},\tilde{z})-k(\tilde{z},z){[k(z,z)+{\sigma }_a^2\mathbb{I}]}^{-1}k(z,\tilde{z}),$$

(11)

where ${\sigma }_a^2$ is the aleatoric uncertainty, and ${\sigma }_e^2$ is used to supervise the learning of the epistemic uncertainty. To see why this is a natural choice, we recall that our framework assumes an auxiliary Gaussian target (the luminosity distance in our context) with normal-inverse-gamma distribution being the prior for its mean and variance, and the epistemic uncertainty being the expectation value of the auxiliary Gaussian’s variance. Thus, the GP variance ${\tilde{\sigma }}_e^2$ in (11) is a natural candidate for supervising the learning of the epistemic uncertainty. Like aleatoric uncertainty in (10), we introduce an additional mean-squared loss term of the form

$${\mathcal{L}}_{epi}=\mathbb{E}\left({\displaystyle \frac{\beta }{\nu (\alpha -1)}}-{\sigma }_e^2\right),$$

(12)

where ${\sigma }_e^2$ is the GP variance at each training datapoint, and we are averaging over the training dataset. In Appendix A, we explain more carefully how the factor of ${\displaystyle \frac{\beta }{\nu (\alpha -1)}}$ arises in (12). Adding both uncertainty-related losses to Equation (9), our complete loss function is thus

$$\begin{array}{ccc}\hfill \mathcal{L}& =& {\mathcal{L}}_{data}+{\mathcal{L}}_{pde}+{\mathcal{L}}_{\pi \left({\sigma }_R^2\right)}+{\mathcal{L}}_{\pi (\Omega )}+{\lambda }_a{\mathcal{L}}_{alea}+{\lambda }_e{\mathcal{L}}_{epi}\hfill \\ & =& -log\left[P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r)\pi \left(\overrightarrow{\Omega }\right)\right]+{\lambda }_a{\mathcal{L}}_{alea}+{\lambda }_e{\mathcal{L}}_{epi}.\hfill \end{array}$$

(13)

The addition of the loss terms (10) and (12) guides the learning of the uncertainty-related model outputs $\alpha ,\beta ,\nu$ to complement how the observed data supervises the learning of the mean target variable $\gamma$. We weighted each loss term with tunable coefficients ${\lambda }_e,{\lambda }_a$ that can be adjusted as hyperparameters to yield a good error calibration at the end of model training.

2.3. A Summary List of Model Implementation

For clarity, in the following, we provide a brief overview of the model implementation process. Our framework is structured around a two-phase training algorithm where in the first phase, the neural network is trained purely on the empirical dataset without alluding to the ODE in Equation (1). Apart from being used to derive one of the loss term (${\mathcal{L}}_{\pi \left({\sigma }_R^2\right)}$), this initial model also furnishes the general initial conditions (e.g., initial values for some of the learnable parameters) for model training in the second phase which uses the full loss function of (13).

(I)

The loss function in this training phase consists of three loss terms: the data loss term (6), and the aleatoric (10) and epistemic (12) loss terms:

$${\mathcal{L}}_{1\mathrm{st}\mathrm{phase}}=-log\left[P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))\right]+{\lambda }_a{\mathcal{L}}_{alea}+{\lambda }_e{\mathcal{L}}_{epi}.$$

(14)

(II)

Independently, a GP regression model is fitted to data so as to gain epistemic uncertainty information for supervising ${\mathcal{L}}_{epi}$, and for constructing $\pi \left(\overrightarrow{\Omega }\right)$ (see Appendix A.2 for full details)

(III)

Upon convergence of the purely data-fitted model, we can now determine the prior $\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r)$ following the steps described in Appendix A.3. (This essentially involves solving for ${\alpha }_r,{\beta }_r$ using (A14) and (A16).)

(IV)

We then proceed with the second phase of model training. This phase of training refines the purely data-fitted model such that it conforms to the presumed PDE description. Apart from the model’s weights, $\{\overrightarrow{\Omega },{\sigma }_R^2\}$ are also learnable parameters. In this final phase, the model is trained using the full loss function

$$\mathcal{L}=-log\left[P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\pi ({\sigma }_R^2;{\alpha }_r,{\beta }_r)\pi \left(\overrightarrow{\Omega }\right)\right]+{\lambda }_a{\mathcal{L}}_{alea}+{\lambda }_e{\mathcal{L}}_{epi},$$

(15)

with each of the six individual loss terms defined in Equations (6)–(8), (10), (12) and (A9).

(V)

Upon completion of training, the model predictions are expressed by the target variable $\gamma$ while confidence bands can be constructed from $\alpha ,\beta ,\nu$. We also infer the PDE parameters $\overrightarrow{\Omega }$ with its uncertainty as defined by the median and credible intervals of the posterior distribution.

In [8], this framework was validated and compared against two other more popularly known models (Bayesian Physics-Informed Neural Networks and Deep Ensemble) for its uncertainty quantification and accuracy in recovering unknown PDE parameters. The controlled case studies used in [8] were nonlinear second-order differential equation in 1D (Poisson equation with Gaussian source) and 2D (Fisher–KPP equation). As shown in [8], the most significant superiority of E-PINN over these two standard uncertainty quantification frameworks is that it yields uncertainty estimates that are calibrated much more consistently (i.e., the confidence intervals deduced are much more consistent with actual model errors). In our context of inferring cosmological parameters, E-PINN thus appears to be well-suited since having a robust uncertainty quantification framework that can generate reliable posterior distributions is particularly critical for our purpose.

3. Methodology

3.1. On the Datasets and Some Limitations

The BAO dataset collected in Table 1 of [2] consists of 32 measurements. It is a list of transverse BAO measurements of the comoving angular diameter distance $D_M/r_d$, where $r_d$ is the sound horizon scale at the end of the baryonic drag epoch. These samples include recent data such as those made by DESI [23,24], the Sloan Digital Sky Survey (SDSS) [25,26,27,28,29,30,31,32,33,34,35,36,37,38] and the Dark Energy Survey (DES) [39,40]. As described in [2], the samples involved anisotropic BAO analyses which incorporate the full 3D galaxies’ distributions, often based on some fiducial cosmological model to convert observed angles and redshifts into physical distances. In [2], the sound horizon $r_d$ was taken to be 147.18 Mpc following Planck18 report [41], and here we adopted the same value for $r_d$ in when translating values of $D_M/r_d$ in Table 1 of [2] to $D_L$. This is a limitation of our work which, in principle, can be overcome by deriving expressions for $r_d$ for each cosmological models (equipped with unknown, learnable parameters) and then replacing numerical luminosity distance targets with $r_d\left(\overrightarrow{\Omega }\right)\times (1+z)N_{data}$ where $N_{data}$ is the numerical $D_M/r_d$ value in Table 1 of [2]. In practice, this would complicate the gradient-descent-based model training because $r_d\left(\overrightarrow{\Omega }\right)$ can only be expressed through a numerical integral and not an explicit function of $\overrightarrow{\Omega }$. An ideal approach would be to adopt a model-independent value for $r_d$ if possible. Interestingly, we note that lowering the sound horizon to 140 Mpc recently proposed by Liu et al. in [42]) to be a model-independent result would naively yield the BAO dataset to be visually compatible with that of Pantheon+ data on the $(D_L,z)$ plane. Another model-independent determination of $r_d$ was performed in [43], which employed another set of combined SNe Ia and BAO data to yield a $r_d$ value of $101.0\pm 2.3$ in units of $h_0^{-1}$ Mpc. Since we have chosen to use the collection of BAO data of [2] for definiteness to enable comparison of the posterior distributions and best-fit values to the standard analysis performed in [2], we used the same Planck18 $r_d=147.18$ Mpc value as was used in [2], leaving a more extensive investigation of $r_d$ dependence for future work.

Another caveat that we would like to emphasize is that the data points presented in Table 1 of [2] are however characterized by correlations2, since several entries are obtained from surveys that partially overlap in sky coverage and source populations (e.g., galaxies and quasars). For example, the pair of points labeled $N=1,2$ in [2] with $(z,D_M/r_d)$ values being (0.32, 8.54), (0.32, 8.76) both originated from SDSS-III data. A limitation of our model is that it does not enable any external knowledge of correlations among points to be used directly in the learning of weights. The fundamental reason is that the t-distribution (A1) underpinning the data loss term is derived from marginalizing over the means and variances of products of univariate Gaussian distributions (rather than multivariate ones). On the other hand, we note that there are also pairs of points in Table 1 of [2] sharing the same redshifts but originating from different surveys, e.g., $N=5\left(\mathrm{BOSS}\right),6\left(\mathrm{DESI}\right)$, $N=18\left(\mathrm{DES}\right),19\left(\mathrm{eBOSS}\right)$.

The Pantheon+ dataset [1] provides Type Ia supernovae (SNe Ia) luminosity distance and distance moduli measurements for redshifts in the range $z\in [0.001,2.3]$, calibrated by the second rung of the distance ladder using Cepheids with the absolute magnitude being $M_B=-19.25\pm 0.01$. The samples consists of 1701 light curves of 1550 spectroscopically confirmed SNe Ia. The data together with the uncertainties can be found at their GitHub website: https://github.com/PantheonPlusSH0ES, accessed on 2 December 2025. A limitation of our usage of this dataset is that we only used the diagonal elements of the covariance matrix for supervising the aleatoric uncertainty. Similar to the case for the BAO dataset, any external knowledge of correlations among data points cannot be used in our formalism. There are no output variables of our neural network model that can interpreted as correlations (or their functions) between different data points. By construction, the four outputs of our model are related to the mean and uncertainty of the luminosity distance at each individual data point. Any generalization of our model to enable the full covariance matrix to be used would have to admit a much larger output layer with targets that are related to the correlations. Each cosmological model is associated with a different Hubble function. In our work here, we set the curvature density term to be zero for simplicity, leaving generalizations that treat it as a learnable parameter for future work. The wCDM and ${\Lambda }_s$CDM models are defined as follows:

$$\begin{array}{ccc}\hfill {\displaystyle \frac{H_{w\mathrm{CDM}}\left(z\right)}{H_0}}& =& {\left({\Omega }_m{(1+z)}^3+(1-{\Omega }_m){(1+z)}^{3(1+w)}\right)}^{1/2},\hfill \end{array}$$

(16)

$$\begin{array}{ccc}\hfill {\displaystyle \frac{H_{{\Lambda }_s\mathrm{CDM}}\left(z\right)}{H_0}}& =& {\left({\Omega }_m{(1+z)}^3+(1-{\Omega }_m)\mathrm{sgn}(z_t-z)\right)}^{1/2},\hfill \end{array}$$

(17)

where we assume the Planck-measured radiation density parameter ${\Omega }_r\sim 9.26\times {10}^{-5}\sim 0$ for simplicity. We take the free parameters of the wCDM model to be $\{H_0,{\Omega }_m,w\}$ where w is the dark energy equation of state parameter. For ${\Lambda }_s$CDM model, its free parameters are taken to be $\{H_0,{\Omega }_m,z_t\}$ where $z_t$ is a transition redshift value from which the cosmological constant switches sign representing a toy model of vacua transition from anti-de Sitter to de Sitter spacetime at some point in the early universe. For computational convenience, here we use a hyperbolic tangent function as a smooth representation of the signum function. These models are parametrically deformable to the standard $\Lambda$CDM model in the limits of $w\to -1$ for the wCDM model and $z_t\to \infty$ for the ${\Lambda }_s$CDM model. In Figure 1, we plot the Pantheon and BAO datasets equipped with Gaussian Process regression curves and uncertainty bands.

3.2. Model Training Setup and Implementation Details

In the following, we furnish some details of the model training, organizing them in terms of the dataset that was used. The base surrogate model $\mathcal{M}\left(z\right)$ was implemented as a fully connected perceptron with two hidden layers of 32 neurons each and an output layer of 4 neurons corresponding to the $\{\alpha ,\beta ,\nu ,\gamma \}$ variables of (A1). For each training dataset, the same initial model ${\mathcal{M}}_0$ was used for training the different cosmological model-based neural networks. The finite parameter domains were chosen to be ${\Omega }_m\in (0.10,0.55),$ $h_0\in (0.50,0.90),w\in (-2.0,-0.01),z_t\in (1.5,3.5)$. We implemented Gaussian Process (GP) regression with a radial-basis function kernel via the scikit-learn library [22], with the optimized kernel’s characteristic length-scales being 2.85 for the Pantheon+ data, 3.46 for the BAO data, and 2.89 for the combined dataset.

For models trained on the Pantheon+ and the combined Pantheon+BAO datasets, in the initial training phase, we used a learning rate of $5\times {10}^{-6}$ for the first $5\times {10}^4$ epochs and ${10}^{-6}$ for the subsequent ones with the total number of epochs being ${10}^6$. The data uncertainty hyperparameters were taken to be ${\lambda }_e={\lambda }_a=1$. For the second phase, the learning rate was $5\times {10}^{-6}$ for the first $1.2\times {10}^6$ epochs followed by $1\times {10}^{-6}$ for another $1\times {10}^6$ epochs. On the other hand, for the smaller BAO dataset, convergence was attained for various cosmological models in $3\times {10}^5$ epochs with a learning rate of $2\times {10}^{-5}$ for the first $2\times {10}^5$ epochs and followed by $2\times {10}^{-6}$ for the remaining ones. The data uncertainty hyperparameters were taken to be ${\lambda }_e={\lambda }_a={10}^8$. Final relative tolerance was of the order $\sim {10}^{-7}$ for Pantheon+ data-based models and higher at $\sim {10}^{-5}$ for BAO data-based ones.

Table 1. Prior density function for each parameter was taken to be univariate Gaussians, of which the means and standard deviations are tabulated here for all three models trained on each dataset. The means and variances are the modes and variances of $f\left(\overrightarrow{\Omega }\right)$ so that the Gaussian priors are representative of the highest density regions of $f\left(\overrightarrow{\Omega }\right)$.

	$\mathbf{\Lambda }$CDM	${\mathbf{\Lambda }}_{\mathit{s}}$CDM	wCDM
Pantheon+ data	${\Omega }_m=0.357\pm 0.164$, $h_0=0.729\pm 0.119$	${\Omega }_m=0.357\pm 0.164$, $h_0=0.729\pm 0.119$, $z_t=2.520\pm 0.755$	${\Omega }_m=0.376\pm 0.164$, $h_0=0.769\pm 0.129$, $w=-1.553\pm 0.725$
BAO data	${\Omega }_m=0.357\pm 0.160$, $h_0=0.671\pm 0.127$	${\Omega }_m=0.366\pm 0.159$, $h_0=0.663\pm 0.128$, $z_t=2.643\pm 0.752$	${\Omega }_m=0.339\pm 0.161$, $h_0=0.720\pm 0.137$, $w=-1.472\pm 0.714$
Combined Pantheon+BAO	${\Omega }_m=0.238\pm 0.163$, $h_0=0.737\pm 0.121$	${\Omega }_m=0.247\pm 0.163$, $h_0=0.729\pm 0.121$, $z_t=2.684\pm 0.754$	${\Omega }_m=0.256\pm 0.163$, $h_0=0.794\pm 0.132$, $w=-1.553\pm 0.717$

shows the parameters’ prior densities for each model. These parameters were determined from the empirical distribution that measures the likelihood of each parametrized family of numerical solutions of the PDE using its deviations from the corresponding purely data-fitted model.

3.3. On Empirical Coverage Probability and Log Model Evidence

Upon completion of model training, we assess the uncertainty quantification through computing the empirical coverage probability (ECP). The ECP at level $1-\alpha$ is the proportion of observed target values that fall within the corresponding t-distribution-based confidence band of (A1). To assess the degree of calibration, one can compare the ECP values to their nominal target level ($1-\alpha$) (nominal coverage probabilities). On the ECP vs. NCP plane, a robust uncertainty quantification would yield a curve that is close to the straight line joining the origin to (1, 1). A representative index would be the mean of the absolute discrepancy between the ECP and NCP. For each model, we compute this mean calibration error (MCE) (see also [44]) and examined plots of ECP vs. NCP, finding that all MCE are very small ≲0.05, with models trained on Pantheon data being better calibrated with an MCE that is 0.1 smaller than those trained on BAO data. Most crucially, none of the 9 models had a ECP curve that is dominantly above or below the ideal line which would have indicated a systematic bias (see Figure 2 for a couple of illustrative ECP plots).

The loss function of our model is, up to a normalization factor, the posterior distribution. The completion of model training yields quantities are directly related the log model evidence that can be further used to discriminate between models. Integrating out the parameters $\overrightarrow{\Omega }$, the model likelihood M and its logarithm are

$$\begin{array}{ccc}\hfill M& =& \int d\Omega P\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\pi \left(\overrightarrow{\Omega }\right),\hfill \\ \hfill logM& =& logP\left(\mathcal{D}\right|\mathcal{M}\left(\overrightarrow{w}\right))+log\left(\int d\Omega P\left(\mathcal{M}\left(\overrightarrow{w}\right)\right|\overrightarrow{\Omega })\pi \left(\overrightarrow{\Omega }\right)\right),\hfill \end{array}$$

(18)

where $\overrightarrow{w}$ are the final model weights and biases. The log model evidence (log M) is often used as a measure to quantify and compare support between competing statistical models from data [5,45]. In

Table 2. Table of inferred parameters (posterior medians with 0.68 C.I.) and logarithm of model evidence (log M). For each dataset (combined, Pantheon+, BAO), the cosmological model with the highest log M is shaded in gray.

Model	Dataset	${\mathit{h}}_{\mathbf{0}}$	${\mathbf{\Omega }}_{\mathit{m}}$	$\mathit{w}\mathbf{\left(}\mathit{w}\mathbf{CDM}\mathbf{\right)}$, ${\mathit{z}}_{\mathit{t}}\mathbf{\left(}{\mathbf{\Lambda }}_{\mathit{s}}\mathbf{CDM}\mathbf{\right)}$	$\mathbf{log}\mathit{M}$
	Pantheon+	$0.{729}_{-0.024}^{+0.033}$	$0.{357}_{-0.092}^{+0.101}$		3687
$\Lambda$CDM	BAO	$0.{680}_{-0.082}^{+0.090}$	$0.{357}_{-0.110}^{+0.110}$		59.8
	Combined	$0.{745}_{-0.041}^{+0.033}$	$0.{320}_{-0.092}^{+0.092}$		3717
	Pantheon+	$0.{745}_{-0.024}^{+0.033}$	$0.{385}_{-0.083}^{+0.073}$	$-1.{431}_{-0.366}^{+0.406}$	3685
wCDM	BAO	$0.{712}_{-0.098}^{+0.090}$	$0.{348}_{-0.110}^{+0.110}$	$-1.{310}_{-0.447}^{+0.528}$	60.2
	Combined	$0.{769}_{-0.057}^{+0.057}$	$0.{293}_{-0.092}^{+0.110}$	$-1.{350}_{-0.406}^{+0.447}$	3712
	Pantheon+	$0.{729}_{-0.024}^{+0.033}$	$0.{357}_{-0.083}^{+0.101}$	$2.{520}_{-0.571}^{+0.571}$	3688
${\Lambda }_s$CDM	BAO	$0.{671}_{-0.073}^{+0.090}$	$0.{366}_{-0.110}^{+0.110}$	$2.{602}_{-0.612}^{+0.531}$	61.0
	Combined	$0.{737}_{-0.033}^{+0.041}$	$0.{274}_{-0.083}^{+0.101}$	$2.{602}_{-0.571}^{+0.571}$	3650

, we display the log model evidence for each model as a comparison index among models trained on the same dataset. In Figure 3, we show the evolution of log M for a couple of models together with the associated loss function. All models were checked to display convergence with a relative tolerance < ${10}^{-4}$ in both the loss and log M term.

4. Results

We collect the inferred parameters together with their uncertainties in Table 2 below. Generally, we found that for each class of cosmological models, the neural networks trained separately on Pantheon and BAO data exhibited a systematic difference evident in the residuals in their luminosity–redshift curves (see Figure 4) and the inferred posterior distributions of the parameters (see Figure 5). When trained on the combined dataset, all models yielded higher $h_0$ and lower ${\Omega }_m$ compared to when being trained only on the Pantheon+ dataset, with ${\Lambda }_s$CDM being associated with a clearly lower log model evidence compared to $\Lambda$CDM and wCDM models.

4.1. On Tensions Between Models Trained Separately on Pantheon+ and BAO Data

We examined the difference in the joint marginal distributions of $h_0$ and ${\Omega }_m$ for the three models, each trained separately on the Pantheon+ and BAO data. Figure 5 shows the 68% and 95% contours for each model. The $\Lambda$CDM and ${\Lambda }_s$CDM models yielded similar distributions with the Jensen–Shannon divergence [46] between the Pantheon and BAO data-based distributions being 2.495 and 2.592, respectively, while that of wCDM model was characterized by the lowest Jensen–Shannon divergence of 2.342. The posterior distributions for the cosmological parameters in various models trained purely on Pantheon+ data were found to be largely contained within the $2\sigma$ contours of their counterparts trained on BAO data (Figure 5). This is in stark contrast to Figure 4 of [2], where the posterior distributions did not overlap at $3\sigma$.

For each cosmological model, we consider the differences in the predicted luminosity–distance curves resulting from the model being trained purely on either BAO or Pantheon+ datasets. The normalized residuals between the predictions of the BAO-trained and Pantheon-trained models are shown in Figure 4. We found that these residuals exhibited strong deviations from the $\mathcal{N}(0,1)$ distribution ($p\approx 0$) associated with statistical noise. This indicates the presence of dataset-dependent systematic effects, whereby each dataset favors a different best-fit model.

4.2. On Models Trained on the Combined Pantheon+ and BAO Data

When trained on the combined dataset, all three models yielded similar prediction curves as depicted in Figure 6 below. Figure 7 shows the marginal posteriors for $h_0,{\Omega }_m$. The $\Lambda$CDM and wCDM models showed the highest log Bayes factor, and all three models yielded posterior medians of $h_0,{\Omega }_m$ that agree within one standard deviation. The posterior medians for $h_0$ for all models were all larger than $0.73$, with a standard deviation falling within $(0.03,0.06)$. Each model yielded lower values of ${\Omega }_m$ and higher values of $h_0$ than when trained on the individual datasets separately.

5. Discussion

In this work, we have applied E-PINN—a novel variant of Physics-Informed Neural Networks—to predict cosmological parameters from recent supernovae [21] and baryon acoustic oscillations (BAO) datasets [2]. Built upon a hybrid of the principles of Evidential Deep Learning, Physics-Informed Neural Networks, and Bayesian Neural Networks, our model enables learning the posterior distribution of the unknown PDE parameters through standard gradient-descent-based training. We also introduced a novel refinement of the original E-PINN framework [7,8] that integrates Gaussian Processes into its algorithm, enabling supervised learning of epistemic uncertainty and the construction of prior functions for the model parameters.

Our neural network-based approach introduces a higher degree of model independence relative to standard regression-based statistical analysis since the perceptron model does not descend from any solutions of some presumed cosmological model while being a fundamental part of the learnable likelihood function. Our approach fundamentally differs from the usual statistical analysis in a few ways: (i) instead of some uniform prior, we use a data-informed prior, constructed to represent an empirical distribution derived from the deviations between the observed data trend and the numerical solution of the presumed PDE; (ii) the loss function that is minimized is generalized from the negative log-likelihood of a Gaussian to a combination of terms (Equation (15)) that incorporates both PDE constraints and data loss terms; (iii) Gaussian Process Regression is invoked to supervise learning of epistemic uncertainty; (iv) the surrogate perceptron model extends the standard approach of only using families of PDE solutions for best-fit estimation, enabling the identification of regions where data trends deviate from the presumed PDE descriptions.

With regards to the Hubble tension problem [20], the essential finding of our work is that the posterior distributions for cosmological parameters in various models trained purely on Pantheon+ data were found to be largely contained within the $2\sigma$ contours of their counterparts trained on BAO data (Figure 5). As tabulated in Table 2, the $h_0$ values were within about $2\sigma$ of one another as defined through the marginal distributions in $h_0,{\Omega }_m$, in contrast to those in [2] exhibiting more than $4\sigma$ tension as inferred from the standard approach of minimizing an appropriate ${\chi }^2$ function. The normalized residuals (Figure 4) indicated the presence of dataset-dependent systematic effects, where each dataset favors a different set of cosmological parameters—a trend supportive of some degree of Hubble tension that is consistent with [2]. In [2], the best-fit values for the $\Lambda$CDM model associated with the BAO dataset were $(h_0,\Omega )=(0.67,0.34)$, while the Pantheon dataset yielded $(h_0,\Omega )=(0.73,0.33)$. Their posterior distributions showed large deviations as depicted in Figure 4 of [2] where one can see that their probability contours at 3$\sigma$ do not even overlap. While our framework yielded parameter estimates similar to theirs—$(h_0,\Omega )=(0.68,0.36)$ based on BAO data and $(h_0,\Omega )=(0.73,0.36)$ based on Pantheon data—these posterior medians were inferred with larger uncertainties, with posterior distributions that showed a much larger degree of overlap in Figure 5, compared to Figure 4 of [2]. One contributing factor can be traced to the nature of the prior density functions that we derived for the parameters (see Table 1). Their uncertainties were slightly larger yet similar in order of magnitude to those characterizing the joint marginal distributions in Figure 5. In contrast, typically flat priors are used in the standard analysis, with the eventual parameter uncertainties only dependent on the minimization of the likelihood function in Equation (3). Apart from the prior, our framework also assumes a different likelihood function. Instead of Equation (3) which is expressed in terms of the solutions to the Friedmann equations, we have Equation (7). These differences led to our joint marginals as shown in Figure 5 to deviate from those in the standard analysis of [2]. Notably, our models’ inferences of $h_0,{\Omega }_m$ were equipped with much larger uncertainties (at least twice for $h_0$ and thrice for ${\Omega }_m$) compared to those arising from the standard analysis in [2].

Overall, our simulation results showed that a more data-informed approach can seemingly reduce statistical tensions between models trained separately on Pantheon and BAO data. This is fundamentally due to different sets of assumptions underlying the statistics of model parameters, most crucially the choice of prior and form of the likelihood functions, and how they relate to cosmological ones. A caveat is that the sound horizon $r_d$ was taken to be fixed at 147.18 Mpc following the Planck18 results [41] when translating $D_M/r_d$ in the BAO dataset of [2] to luminosity distances. While this enables us to compare our results fairly with those of [2], it should also be interpreted as a limitation since the tension is sensitive to $r_d$.

All initial models in the absence of PDE constraints arising from the presumed cosmological models appeared to suggest that the luminosity–redshift curve should flatten out towards higher redshift gradually, in contrast to the numerical solutions for all three cosmological models considered here. It would be interesting to observe if future empirical data from supernovae light curves at high redshift support this trend. With regards to model selection, we note that the log model evidence appeared to disfavor ${\Lambda }_s$CDM when the combined Pantheon and BAO data were taken into account but otherwise showed no other notable model preferences. The wCDM model yielded the highest $h_0$ values relative to the two other models irrespective of the dataset used.

An immediate future direction worth pursuing as a follow-up to our work here would be to use a model-independent sound horizon $r_d$ for training models on BAO data, or to lift it to be a learnable parameter. A recent analysis [42] inferred a value for $r_d\sim 140$ Mpc by leveraging time-delay measurements of gravitationally lensed quasars from H0LiCOW collaboration [47] in a model-independent approach. Such a value would naively reduce deviations between the models trained separately on Pantheon and BAO datasets as shown in Figure 4. An earlier model-independent analysis in [43] estimated $r_d=101.0\pm 2.3h_0^{-1}$ Mpc, taking into account a combined SNe Ia+BAO dataset. If we used our inferred $h_0$ value based on the model trained on the combined Pantheon+BAO data, this translates to $r_d\sim 135.6_{-3.5}^{+3.1}$ Mpc, which is notably close to the value determined in [42]. The degree of sensitivity to $r_d$ suggests that any future analysis should ideally treat $r_d,h_0$ and other cosmological parameters to be simultaneously learnable ones right from the outset, although this would increase computational cost and complexity. As noted in [42], future cosmological probes may bring in greater diversity of data sources, such as gravitational wave standard sirens [48], with which we can infer the sound horizon and other cosmological parameters. Our data-driven neural network methodology is poised to leverage such increasingly diverse observations to infer parameters with robust data-informed priors. More generally, in the aspect of machine learning techniques, we expect our proposed method of synergizing Gaussian Processes with E-PINN to be transferable to other scientific modeling problems, and to be particularly useful for contexts where data is relatively scarce and the learning of epistemic uncertainty then becomes crucial to the model training process.

Another follow-up work would be to explore generalizations of evidential deep learning that can facilitate the learning of known correlations among data points. Naively, one can consider the multivariate t-distribution obtained by marginalizing out means and covariances of a multivariate Gaussian with the Normal-Inverse-Wishart prior, but this would imply that the input’s dimensionality is fixed to the specific value of the training dataset size, and the number of target variables would be increased by over an order of ${10}^6$. Another line of approach would be to consider adjusting the relative weightage of data points that are correlated, for example taking some form of weighted mean of a pair of correlated points instead of treating them as separate entities within the sum over all training points. One would need to examine how to take such a weighted mean that is robustly consistent with the off-diagonal elements of the covariance matrix. In the absence of a framework that enables external knowledge of correlation to be used in the supervised learning, a cleaner approach would be to simply use datasets that are not characterized by significant correlations.3 For example, in the context of our work, it would thus be interesting to see how our various results change when we use BAO data that is sourced from SDSS [49], DESI [50], etc., separately instead of Table 1 of [2].