Multiscale Flow for Robust and Optimal Cosmological Analysis

TL;DR

This work tackles the challenge of extracting maximal, robust cosmological information from nonlinear 2D weak-lensing fields by learning the full field-level likelihood $p(x|y)$ with a Haar wavelet–based multiscale normalizing flow. The method factorizes the likelihood across scales via multiresolution analysis, modeling each scale’s term with dedicated conditional NF blocks and combining them to recover the full density. A two-stage training procedure—generative likelihood optimization followed by discriminative calibration using $p(y|x)$ samples—yields calibrated posteriors and enables robust detection of distribution shifts, such as baryonic effects, while preserving strong cosmological constraints. The results show substantial improvements in figure-of-merit over traditional summaries (power spectrum, peak counts, scattering transform, CNN), effective baryon marginalization, and the ability to generate realistic mock weak-lensing maps, with potential applicability to other intensity maps and 3D fields.

Abstract

We propose Multiscale Flow, a generative Normalizing Flow that creates samples and models the field-level likelihood of two-dimensional cosmological data such as weak lensing. Multiscale Flow uses hierarchical decomposition of cosmological fields via a wavelet basis, and then models different wavelet components separately as Normalizing Flows. The log-likelihood of the original cosmological field can be recovered by summing over the log-likelihood of each wavelet term. This decomposition allows us to separate the information from different scales and identify distribution shifts in the data such as unknown scale-dependent systematics. The resulting likelihood analysis can not only identify these types of systematics, but can also be made optimal, in the sense that the Multiscale Flow can learn the full likelihood at the field without any dimensionality reduction. We apply Multiscale Flow to weak lensing mock datasets for cosmological inference, and show that it significantly outperforms traditional summary statistics such as power spectrum and peak counts, as well as novel Machine Learning based summary statistics such as scattering transform and convolutional neural networks. We further show that Multiscale Flow is able to identify distribution shifts not in the training data such as baryonic effects. Finally, we demonstrate that Multiscale Flow can be used to generate realistic samples of weak lensing data.

Figures (7)

• Figure 1: Illustration of Multiscale Flow model. The input map $x_{2^n}$ with resolution $2^n$ is iteratively processed with a set of low pass filters ($\phi$), high pass filters ($\psi_1$, $\psi_2$, $\psi_3$) and downsampling ($\downarrow 2$), resulting in a series of detailed maps $x_{2^{n-1}}^d, x_{2^{n-2}}^d, \cdots, x_{2^k}^d$ and an approximation map $x_{2^k}$. These maps are then transformed by several NF blocks to Gaussian latent maps $z_{2^{n-1}}^d, z_{2^{n-2}}^d, \cdots, z_{2^k}^d, z_{2^k}$, where each NF block is composed of an actnorm layer, an invertible $1\times 1$ convolution, and an affine coupling layer (Equation \ref{['eq:affine1']}, \ref{['eq:affine2']}), as shown on the top left of this figure. The NF transformation is conditioned on the conditional variable $y$ and approximation maps, which are represented by dashed arrows in the illustration. The log-likelihood of the input map $x_{2^n}$ can be calculated with Equation \ref{['eq:decomposition']}.
• Figure 2: Percentage of test data that fall outside $95\%$ confidence region for different $\lambda$ values. A perfectly calibrated posterior has $5\%$ outliers. The shaded region shows the uncertainty due to finite number of test data. This measurement is made on weak lensing maps with $64^2$ resolution and $n_g=30 \mathrm{arcmin}^{-2}$ galaxy density.
• Figure 3: Multiscale Flow posterior comparison of different scales on 20 test data with galaxy number density $n_g=30 \mathrm{arcmin}^{-2}$.
• Figure 4: Comparison of posterior distributions between different scales of Multiscale Flow and power spectrum on a $3.5 \times 3.5 \mathrm{deg}^2$ convergence map with $n_g=20 \mathrm{arcmin}^{-2}$.
• Figure 5: Top panel: scale-dependent posterior analysis of a baryon-corrected convergence map using Multiscale Flow trained on dark-matter-only maps (left), and Multiscale Flow trained on BCM maps (right). Bottom panel: ROC curve of identifying distribution shift with $\log p$ (left) and $\Delta \log p$ (right). The "small scales" in the lower left panel represent combining the three small scale terms. In these experiments, we consider $30 \mathrm{arcmin}^{-2}$ galaxy shape noise.