Reconstructing the Universe with Variational self-Boosted Sampling

Abstract

Forward modeling approaches in cosmology have made it possible to reconstruct the initial conditions at the beginning of the Universe from the observed survey data. However the high dimensionality of the parameter space still poses a challenge to explore the full posterior, with traditional algorithms such as Hamiltonian Monte Carlo (HMC) being computationally inefficient due to generating correlated samples and the performance of variational inference being highly dependent on the choice of divergence (loss) function. Here we develop a hybrid scheme, called variational self-boosted sampling (VBS) to mitigate the drawbacks of both these algorithms by learning a variational approximation for the proposal distribution of Monte Carlo sampling and combine it with HMC. The variational distribution is parameterized as a normalizing flow and learnt with samples generated on the fly, while proposals drawn from it reduce auto-correlation length in MCMC chains. Our normalizing flow uses Fourier space convolutions and element-wise operations to scale to high dimensions. We show that after a short initial warm-up and training phase, VBS generates better quality of samples than simple VI approaches and reduces the correlation length in the sampling phase by a factor of 10-50 over using only HMC to explore the posterior of initial conditions in 64${}^3$ and 128${}^3$ dimensional problems, with larger gains for high signal-to-noise data observations.

arXiv commentA shorter version of this paper is accepted for spotlight presentation in Machine Learning for Astrophysics Workshop at ICML, 2022

References in corpus (16)

• Variational Inference: A Review for Statisticians
• Cosmology and Fundamental Physics with the Euclid Satellite
• Solving Large Scale Structure in Ten Easy Steps with COLA
• FastPM: a new scheme for fast simulations of dark matter and halos
• ELUCID - Exploring the Local Universe with reConstructed Initial Density field I: Hamiltonian Markov Chain Monte Carlo Method with Particle Mesh Dynamics
• Physical Bayesian modelling of the non-linear matter distribution: new insights into the Nearby Universe
• Towards optimal extraction of cosmological information from nonlinear data
• Cosmological Reconstruction From Galaxy Light: Neural Network Based Light-Matter Connection
• Reconstructing large-scale structure with neutral hydrogen surveys
• Translation and Rotation Equivariant Normalizing Flow (TRENF) for Optimal Cosmological Analysis
• FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver
• Geometric variational inference
• MUSE: Marginal Unbiased Score Expansion and Application to CMB Lensing
• Efficient Optimal Reconstruction of Linear Fields and Band-powers from Cosmological Data
• Primordial power spectrum and cosmology from black-box galaxy surveys
• CosmicRIM : Reconstructing Early Universe by Combining Differentiable Simulations with Recurrent Inference Machines

Cited by in corpus (7)

• Cosmology inference at the field level from biased tracers in redshift-space
• Bayesian Inference of Initial Conditions from Non-Linear Cosmic Structures using Field-Level Emulators
• Learning the Universe: Learning to Optimize Cosmic Initial Conditions with Non-Differentiable Structure Formation Models
• Field-level Reconstruction from Foreground-Contaminated 21-cm Maps
• Lagrangian space remapping and the angular momentum reconstruction from cosmic structures
• Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation
• On the Connection between Field-Level Inference and $n$-point Correlation Functions