Learning the Universe: Learning to Optimize Cosmic Initial Conditions with Non-Differentiable Structure Formation Models

TL;DR

This work tackles reconstructing cosmological initial conditions from non-linear, non-differentiable forward models by introducing LULO, a gradient-free optimizer that learns to update initial-condition fields from data discrepancies while keeping full physics simulations in the loop. The neural optimizer, trained on high-dimensional $128^3$ volumetric data, provides update directions without differentiating through the simulator and uses a line search to adjust step sizes, enabling efficient navigation of non-differentiable landscapes. Empirically, LULO recovers initial conditions from $M_{200\mathrm{c}}$ halo catalogs with cross-correlations $\geq80\%$ down to non-linear scales ($k\sim1\,h\mathrm{Mpc}^{-1}$), and accurately reproduces power spectra, bispectra, halo mass functions, and peculiar velocities. This gradient-free, physics-in-the-loop approach offers a scalable path to non-linear field-level inference, compatible with arbitrary forward models and future enhancements such as higher-resolution simulations and additional cosmological probes.

Abstract

Making the most of next-generation galaxy clustering surveys requires overcoming challenges in complex, non-linear modelling to access the significant amount of information at smaller cosmological scales. Field-level inference has provided a unique opportunity beyond summary statistics to use all of the information of the galaxy distribution. However, addressing current challenges often necessitates numerical modelling that incorporates non-differentiable components, hindering the use of efficient gradient-based inference methods. In this paper, we introduce Learning the Universe by Learning to Optimize (LULO), a gradient-free framework for reconstructing the 3D cosmic initial conditions. Our approach advances deep learning to train an optimization algorithm capable of fitting state-of-the-art non-differentiable simulators to data at the field level. Importantly, the neural optimizer solely acts as a search engine in an iterative scheme, always maintaining full physics simulations in the loop, ensuring scalability and reliability. We demonstrate the method by accurately reconstructing initial conditions from $M_{200\mathrm{c}}$ halos identified in a dark matter-only $N$-body simulation with a spherical overdensity algorithm. The derived dark matter and halo overdensity fields exhibit $\geq80\%$ cross-correlation with the ground truth into the non-linear regime $k \sim 1h$ Mpc$^{-1}$. Additional cosmological tests reveal accurate recovery of the power spectra, bispectra, halo mass function, and velocities. With this work, we demonstrate a promising path forward to non-linear field-level inference surpassing the requirement of a differentiable physics model.

Figures (15)

• Figure 1: High-level overview of LULO (Learning the Universe by Learning to Optimize) which aims to fit complex models to data by reconstructing the three-dimensional initial conditions. The process consists of two components: 1) applying a high-fidelity physics simulator $\mathscr{S}$, and 2) updating the initial conditions to minimize discrepancies $\Delta \boldsymbol{d}$ between the simulator output and the data. Importantly, any simulator model, including fully non-linear and non-differentiable ones, is supported. The neural optimizer, pre-trained via a supervised approach to learn how to map data discrepancies to updates in the initial conditions, proposes an update direction $\Delta \boldsymbol{x}$ across all initial condition voxels simultaneously. In the current implementation, the step size $\gamma_t$ is optimized via a line search algorithm that requires running the simulator (i.e., $\gamma_t=\gamma_t(\mathscr{S}, \Delta \mathbf{x}_t)$; see details in sections \ref{['sec:update_direction']}). The iterative process continues until the simulation output aligns with the data, as shown in the $2d$-slices after eight optimization steps. The slices show the evolved non-linear dark matter density field in a cubic box with side length $250h^{-1}$ Mpc and the corresponding dark matter halo field as produced by a non-differentiable spherical overdensity algorithm.
• Figure 2: Learning the relationship between changes in the output of a cosmological simulator and corresponding changes in the initial conditions involves two key steps: generating training data (left) and training a machine learning model with a specific architecture (right). We perturb the initial conditions using an operator $\mathcal{P}$ that preserves the statistical properties of the Gaussian initial conditions, namely zero mean and unit variance. The resulting set of difference fields $\{\Delta \boldsymbol{d}\}$, where $\Delta \boldsymbol{d} = \boldsymbol{d}_a - \boldsymbol{d}_b$, serves as input to the neural optimizer model during training. The neural optimizer is tasked with predicting the corresponding update, $\Delta \boldsymbol{x} = \boldsymbol{x}_a - \boldsymbol{x}_b$, in the initial conditions. The chosen architecture is a convolutional V-Net model, which effectively handles structures and their correlations across multiple scales. To prioritize learning accurate update directions we use a cosine similarity loss during training.
• Figure 3: The performance of the neural optimizer is evaluated through the cross-correlation between the predicted update direction $\Delta \boldsymbol{x}_{\mathrm{pred}}$ in the initial conditions and the ground truth $\Delta \boldsymbol{x}_{\mathrm{true}}$. The mean and standard deviation of the cross-correlation across all samples in the validation set are shown. A high alignment between the prediction and the truth down to scales of approximately $0.3h$ Mpc$^{-1}$ is obtained, with the correlation decreasing at smaller scales. This provides insight into which scales in the initial conditions are crucial for correcting data discrepancies at our particular resolution.
• Figure 4: The reconstructed initial conditions (bottom row; left), and the corresponding forward simulation output at $z=0$ in terms of dark matter overdensity (centre left), the Lagrangian velocity field (centre right), and mass-weighted halo count overdensity field (right) are compared with the ground truth (top row). The initial conditions have been smoothed at $10h^{-1}$ Mpc to reveal the visual alignment. A high correlation over these scales suffices for a highly accurate reconstruction in the data space, even at the smallest non-linear scales. In particular, the density and velocity fields, despite only being by-products and not used to constrain the initial conditions, are well reconstructed. Note that for the mass-weighted field, all haloes as found by the non-differentiable AHF algorithm have been included. Low-mass haloes, being the most numerous, are often not accurately reproduced because of the low number of particles per halo (see section \ref{['sec:halo_reconstruction']}). The majority of haloes are nonetheless observed to be accurately recovered.
• Figure 5: The quality of the reconstruction is quantified through the power spectra (top) and cross-correlation (bottom) of the reconstructed initial conditions (left), the dark matter overdensity (center), the mass-weighted halo count overdensity and halo count overdensity (right). In the residual panels, we display the transfer function, i.e. the square root of the ratio of the power spectra of the reconstructions with the ground truth. With initial conditions matching to within $10\%$ down to scales of $\sim 0.3h$ Mpc$^{-1}$, the final density field and the halo fields match down to highly non-linear scales of $\sim 2.1h$ Mpc$^{-1}$ and $\sim 2.7h$ Mpc$^{-1}$, respectively. The cross-correlations further demonstrate that high correlation in the initial conditions at large scales ($80\%$ at $\sim0.23h$ Mpc$^{-1}$) is sufficient for an $80\%$ correlation in the dark-matter overdensity and halo fields beyond scales of $1h$ Mpc$^{-1}$. This reflects the gravitational collapse of proto-structures into final structures as discussed in section \ref{['sec:reconstruction_powspec_etc']}. The linear $\Lambda$CDM power spectrum at $z=0$ is shown as a reference to highlight the non-linearities involved.

Theorems & Definitions (1)

• proof