Photometric Redshifts for Cosmology: Improving Accuracy and Uncertainty Estimates Using Bayesian Neural Networks

As large astronomical surveys come online in the next few years both from the ground with LSST and the Vera Rubin Telescope and in space with Euclid, many researchers are turning to machine learning to handle the exponentially increasing influx of data. However, common machine learning methods often only provide point estimates and do not generally provide accurate confidence intervals for specific predictions [12, 13, 3]. Accurate uncertainty estimates, whether from machine learning or elsewhere, are critical for measurements from these surveys because these measurements and uncertainties are not the end goals; rather, they are subsequent inputs into inference to constrain models of our Universe.

One crucial goal of these surveys is to determine expansion history of the universe and with it the parameters describing dark energy. This determination relies ultimately on accurately and precisely measuring the redshifts of hundreds of millions of galaxies. Spectroscopically measuring galaxy redshifts, where the light is split into hundreds of small bins of wavelength, is time consuming and practically impossible for the necessary sample size to constrain cosmological parameters. Instead of measuring detailed galaxy spectra to determine redshift, one can take images of galaxies in a few large bins of wavelength (photometry). While galaxy photometry contains information about redshift, the observed variation between the intrinsic properties of galaxies makes it difficult to model a-priori. Astronomers have adopted data driven approaches using machine learning methods for redshift estimates using photometry [24, 1, 4, 27, 14, 29].

In this work, we investigate photometric redshift (photo-z) estimation using Bayesian neural networks (BNN), a type of probabilistic neural network (NN) [15]. Probabilistic neural networks, conceptualized in the 1990s [26], have previously been limited in their ability to process the size of data required for performing photo-z estimation for large-scale surveys, because of the complexity of their computation. However, recent breakthroughs in conceptual understanding and computational capabilities (e.g. [10, 9]) now make probabilistic deep learning possible for cosmology. Probabilistic deep learning such as BNN has many advantages compared to traditional neural networks, including better uncertainty representations, better point predictions, and offers better interpretability of neural networks because they can be viewed through the lens of probability theory. In this way we can draw upon decades of development in Bayesian inference analyses. We compare the performance of a BNN to a fully-connected non-probabilistic NN and evaluate the accuracy of the confidence intervals for the probabilistic predictions. To our knowledge this is the first application of BNN for photo-z estimation. We compare the NN to BNN in order to assess the effect of incorporating Bayesian statistics to photo-z estimation; we intentionally compose both models to be as similar as possible.

We compare the performance of the NN and BNN on the data discussed in §2.1 in Table 2. The percentage of point-source outlier predictions as defined by Eq. 1 differ for the NN and BB. For the NN, only 6.5% of points lie outside of the lines as outliers for the NN and 17.3% for the BNN. With the modified outlier metric (Eq. 2) discussed in §2.3 we obtain $O_{b}=6.3\%$. Fig. 2 contains results from an example determination with a BNN and non-Bayesian NN, where results are divided into bins of size $z=0.1$ and averaged. We note that both models generally perform worse at higher redshifts, which is due in large part to the reduced signal to noise for distant dim sources and also the disproportionate number of high redshift sources (z > 2.5) compared to low redshift sources. The latter point is an unavoidable attribute of similar datasets. We note that low redshifts are generally over-covered, indicating that their photo-z uncertainties are over-estimated, while the photo-z uncertainties of high redshift galaxies are generally under-estimated.

While our goal in this work is to compare the two types of NNs, we provide a comparison to LSST requirements [6] (Table 2) for reference. The NN meets the LSST goal for outlier rate and bias. We believe both models could be further optimized for these requirements.

Compared to the non-probabilistic NN, the BNN has the advantage of producing uncertainty constraints on every prediction, which are necessary for using photo-z estimation as a probe of cosmological parameters. Based on the results in this work, the BNN has the disadvantage of generally producing worse point estimates compared to the NN, however, pairing photo-z predictions with uncertainties provides a more robust look into the quality of photo-z predictions. We note that the BNN in this work was designed with the intention to closely resemble the non-probabilistic NN, and therefore our findings may not generalize to other BNN models. Optimal photo-z performance from the BNN may require significant model adjustments, which is an on-going study. Fig. 2 (right) visualizes the $O$ and $O_{b}$ rates per redshift bin; as expected, the number of outliers decreased when considering the photo-z uncertainty. The uncertainties produced by the BNN are larger than expected for $0<z<2.5$ and are underestimated in $2.5<z<4$. This is not necessarily a flaw inherent to a BNN, but it is worth investigating how to develop better uncertainty estimates with broad-band photometry. It is possible that one source of the over-estimation of photo-z uncertainties results from a disparity between the complexity present in the band magnitudes compared to the BNN model; we use five photometric band fluxes paired with a single spectroscopic redshift per galaxy for training, while the model parameters optimized during the training process of a NN can easily reach into the thousands. In a future work we will apply galaxy photometric images to a Bayesian convolutional neural network, which we believe will enhance the information present in the band fluxes.