Enhancing Photometric Redshift Estimation for LSST with a Hybrid LSTM-Mixture Density Network

Abstract

Accurate photometric redshift (photo-$z$) estimation and robust uncertainty quantification are essential for the LSST to achieve its precision cosmology goals. Traditional machine learning algorithms are largely restricted to point estimates, struggling to characterize the multimodal nature of redshift PDFs and the degeneracies within the color-redshift space. To address this, we present and validate the LSTM-MDNz architecture, which integrates sequential feature extraction with flexible probability density modeling to enhance both prediction accuracy and uncertainty calibration across a broad redshift range, thereby meeting the stringent data quality requirements necessitated by next-generation cosmological analysis. The LSTM-MDNz framework treats multi-band photometry as wavelength-ordered sequences, utilizing LSTM networks to capture non-linear evolutionary correlations across the SED. A Mixture Density Network (MDN) is then employed to explicitly model posterior PDFs via Gaussian mixture models (GMMs). Performance is evaluated on the HSC GalaxiesML dataset (which serves as a small-scale proxy for next-generation surveys like LSST) and benchmarked against the BNN architecture established by Jones et al. (2024). The proposed model consistently outperforms the BNN baseline, achieving a $\sim 10\%$ improvement in point-estimation accuracy (specifically across RMSE, MAE, scatter, and $\sigma_{\text{NMAD}}$) and a $\sim 20\%$ reduction in the rates of both general and catastrophic outliers. A uniform probability integral transform (PIT) distribution confirms well-calibrated probabilistic outputs. Furthermore, the PDF-based confidence metric $z_{\text{conf}}$ enables high-purity catalog construction: excluding just approximately $4\%$ of extremely low-confidence ($z_{\text{conf}} < 0.05$) samples reduces the overall outlier rate by $\sim 48\%$.