Advancing Normalizing Flows for Correlated Non-Gaussian Fluctuations in Cosmological Simulations

We present a technique to improve the accuracy and training efficiency of normalizing flows for multiple images in the context of cosmology. Normalizing flows are powerful deep generative models that can learn complex probability distributions through invertible transformations applied to a simple distribution. They are well-suited for both image generation and density estimation, enabling precise likelihood evaluation and efficient sampling. However, due to the inherent constraint of invertibility, normalizing flows exhibit limited expressiveness when compared to other well-known models like GANs. Yet, past research has demonstrated that this limitation can be addressed by employing more expressive priors, such as resampled normal Gaussian or correlated priors, effectively enhancing the capabilities of normalizing flows. Our ongoing work focuses on developing a technique to enable normalizing flows to capture correlated non-Gaussian fluctuations in realistic mm-wave extragalactic foreground simulations, which consist of multiple components, based on N-body simulations. As part of this endeavor, we investigate and compare the efficiency of normalizing flows in generating two-component correlated non-Gaussian maps using various priors. The three priors considered are as follows: the normal distribution (the simplest setup), the correlated prior (which incorporates the auto-spectrum of each target maps but not their correlations), and the component-correlated prior (which incorporates both the auto- and the cross-spectra of the target maps). We find that, when dealing with $f_{nl}$ type local non-Gaussianity the application of correlated and component-correlated priors results in more accurate representations of the target distribution compared to other approaches. We will next apply this on realistic simulations of the CMB lensing convergence and extragalactic foreground fields, with the goal of field-level inference with mm-wave data.