Bayesian Posteriors with Stellar Population Synthesis on GPUs

TL;DR

This paper introduces GPU-accelerated techniques for stellar population synthesis modeling, significantly speeding up galaxy property inference from spectral data, enabling large-scale Bayesian analysis of galaxy populations.

Contribution

The paper presents novel GPU-based methods and approximations that accelerate SPS predictions and Bayesian inference, making large galaxy sample analysis feasible.

Findings

Speed-up factor of 50 for photometry calculations

Bayesian inference rate of 1000 galaxy posteriors per minute on a single GPU

Incorporation of burstiness modeling with negligible additional computation

Abstract

Models of Stellar Population Synthesis (SPS) provide a predictive framework for the spectral energy distribution (SED) of a galaxy. SPS predictions can be computationally intensive, creating a bottleneck for attempts to infer the physical properties of large populations of individual galaxies from their SEDs and photometry; these computational challenges are especially daunting for near-future cosmology surveys that will measure the photometry of billions of galaxies. In this paper, we explore a range of computational techniques aimed at accelerating SPS predictions of galaxy photometry using the JAX library to target GPUs. We study a particularly advantageous approximation to the calculation of galaxy photometry that speeds up the computation by a factor of 50 relative to the exact calculation. We introduce a novel technique for incorporating burstiness into models of galaxy star formation history that captures very short-timescale fluctuations with negligible increase in computation time. We study the performance of Hamiltonian Monte Carlo (HMC) algorithms in which individual chains are parallelized across independent GPU threads, finding that our pipeline can carry out Bayesian inference at a rate of approximately $1000$ galaxy posteriors per minute on a single GPU. Our results provide an update to standard benchmarks in the literature on the computational demands of SPS inference; our publicly available code enables previously-impractical Bayesian analyses of large galaxy samples, and includes several standalone modules that could be adopted to speedup existing SPS pipelines.