Predicting the Initial Conditions of the Universe using a Deterministic
  Neural Network
Finding the initial conditions that led to the current state of the universe
is challenging because it involves searching over an intractable input space of
initial conditions, along with modeling their evolution via tools such as
N-body simulations which are computationally expensive. Recently, deep learning
has emerged as a surrogate for N-body simulations by directly learning the
mapping between the linear input of an N-body simulation and the final
nonlinear output from the simulation, significantly accelerating the forward
modeling. However, this still does not reduce the search space for initial
conditions. In this work, we pioneer the use of a deterministic convolutional
neural network for learning the reverse mapping and show that it accurately
recovers the initial linear displacement field over a wide range of scales
($<1$-$2\%$ error up to nearly $k\simeq0.8$-$0.9 \text{ Mpc}^{-1}h$), despite
the one-to-many mapping of the inverse problem (due to the divergent backward
trajectories at smaller scales). Specifically, we train a V-Net architecture,
which outputs the linear displacement of an N-body simulation, given the
nonlinear displacement at redshift $z=0$ and the cosmological parameters. The
results of our method suggest that a simple deterministic neural network is
sufficient for accurately approximating the initial linear states, potentially
obviating the need for the more complex and computationally demanding backward
modeling methods that were recently proposed.
Vaibhav Jindal, Albert Liang, Aarti Singh, Shirley Ho, Drew Jamieson
# Predicting the Initial Conditions of the Universe using Deep Learning

###### Abstract

Finding the initial conditions that led to the current state of the universe is challenging because it involves searching over a vast input space of initial conditions, along with modeling their evolution via tools such as N-body simulations which are computationally expensive. Deep learning has emerged as an alternate modeling tool that can learn the mapping between the linear input of an N-body simulation and the final nonlinear displacements at redshift zero, which can significantly accelerate the forward modeling. However, this does not help reduce the search space for initial conditions. In this paper, we demonstrate for the first time that a deep learning model can be trained for the reverse mapping. We train a V-Net based convolutional neural network, which outputs the linear displacement of an N-body system, given the current time nonlinear displacement and the cosmological parameters of the system. We demonstrate that this neural network accurately recovers the initial linear displacement field over a wide range of scales (\(<1\)-\(2\%\) error up to nearly \(k=1\ \mathrm{Mpc}^{-1}\,h\)), despite the ill-defined nature of the inverse problem at smaller scales. Specifically, smaller scales are dominated by nonlinear effects which makes the backward dynamics much more susceptible to numerical and computational errors leading to highly divergent backward trajectories and a one-to-many backward mapping. The results of our method motivate that neural network based models can act as good approximators of the initial linear states and their predictions can serve as good starting points for sampling-based methods to infer the initial states of the universe.

Machine Learning Department, Carnegie Mellon University, Pittsburgh, PA, USA \({}^{2}\)Max Plank Institute for Astrophysics, Garching, Germany \({}^{3}\)Center for Computational Astrophysics, Flatiron Institute, New York, NY, USA; \({}^{4}\)Department of Physics & Center for Data Science, New York University, New York, NY, USA; \({}^{5}\)Department of Astrophysical Sciences, Princeton University, Princeton, USA; Correspondence to: Aarti Singh \(<\)arti@cs.cmu.edu\(>\), Shirley Ho \(<\)shirleyho@flatironinstitute.org\(>\).

## 1 Introduction

The evolution of our universe can be uniquely determined by its initial conditions and the laws of physics governing its dynamics. To understand this cosmic history, astrophysicists use a large number of surveys (Amendola et al., 2018; Eisenstein et al., 2011; Spergel et al., 2015) and simulations (Springel et al., 2001; Bagla, 2002; Villaescusa-Navarro et al., 2020). These simulations compute the gravitational evolution of a system of N-body particles given a set of nearly uniform and typically Gaussian initial conditions, representing the early universe. These forward simulations, however, are computationally expensive and require a large amount of time and computational resources.

In recent years, deep-learning techniques have been shown to be extremely helpful in accelerating the forward modeling process (He et al., 2019; Jamieson et al., 2022). These deep-learning models learn the mapping between pairs of inputs and outputs from numerical N-body simulations and act as fast and accurate approximators for these simulators. These deep-learning surrogates speed up the forward modeling process by several orders of magnitude. Since neural networks are theoretically proven to be universal function approximators (Hornik et al., 1989), the forward modeling of N-body simulations fits perfectly into the regime of neural networks.

The problem of inferring the initial state of the universe, or the input to an N-body simulation that generates a specific redshift zero (current time) nonlinear displacement field, is an inverse problem and poses a completely different challenge. Inverse problems are hard as they require a search over a large space of input configurations (the potential initial conditions of the universe), and typically involve one-to-many mapping if learned as a reverse mapping. Standard neural networks are one-to-one mappings, and not expected to work well in these problems. Sampling approaches such as Hamiltonian Monte Carlo (Radford, 2012) based on Bayesian priors are computationally very expensive. One could resort to more complex generative neural networks such as generative adversarial networks (GANs) (Goodfellow et al., 2014), normalizing flows (Rezende & Mohamed, 2015), diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020), etc. but they often either fail to converge, are unstable, or require excessive computational
time (cf. Che et al., 2016; Chen et al., 2023; Salimans & Ho, 2022).

Importantly, simulations in scientific fields such as cosmology are often deterministic, and therefore they are reversible in principle. The one-to-many backward problem arises primarily due to numerical and computational errors which get exacerbated only at small scales which are dominated by nonlinear effects. This can cause the backward trajectories to be highly divergent. This motivates the approach we demonstrate here: training a standard deterministic neural network to learn the reverse map and output the initial states of cosmological N-body simulations using the final state of displacements as input. We show that despite the one-to-many nature of the reverse mapping at small scales, a simple neural network can do an excellent job of predicting the initial states not only at large scales but even down to relatively small scales (\(k>0.1\ \mathrm{Mpc}^{-1}\,h\)) where the nonlinear dynamics of gravitational clustering become important. Our model continues to have \(<1-2\%\) error down to \(k=0.8\)-\(0.9\ \mathrm{Mpc}^{-1}\,h\). The inverse model we train is only slightly less accurate than the forward model based on the same architecture (Jamieson et al., 2022). Our results empirically motivate the use of neural networks as approximate inverse-mapping black boxes that could generate reliable initial states for a given output state, which could then be used to speed up the more fine-grained sampling-based inverse modeling methods.

## 2 Background

Consider an N-body system with particles distributed on a uniform grid with positions \(\mathbf{q}\). Let \(\boldsymbol{\Psi}_{ZA}(\mathbf{q})\) be their linear ZA approximation at redshift \(z=0\). Thus, the final positions of the particles when they evolve linearly according to the Zel'dovich approximation is

\[\mathbf{x}_{lin}(\mathbf{q})=\mathbf{q}+\boldsymbol{\Psi}_{ZA}(\mathbf{q}). \tag{1}\]

Let the final nonlinear displacement of the particle initially at grid site \(\mathbf{q}\) be \(\boldsymbol{\Psi}_{NL}(\mathbf{q})\). Thus, the final positions of the

Figure 1: Qualitative comparison of the \(x,y,\) and \(z\) displacements for a \(128\times 128\) slice of particles from a linear field sampled using the training parameters and the corresponding linear field predicted by our inverse model.
particles at redshift \(z=0\) under nonlinear evolution is

\[\mathbf{x}_{nonlin}(\mathbf{q})=\mathbf{q}+\mathbf{\Psi}_{NL}(\mathbf{q}). \tag{2}\]

In this work, we investigate the problem of predicting the linear displacement field at redshift \(z=0\), given the nonlinear displacement field and the cosmological parameters defining the N-body evolution.

## 3 Methodology

We train a CNN-based neural network to predict the initial state of an N-body system of particles evolving under gravity on an expanding cosmological background, given the final state. Our CNN takes the nonlinear displacement field \(\mathbf{\Psi}_{NL}(\mathbf{q})\) at redshift \(z=0\) (current time) and the value of \(\Omega_{m}\) as the input, and predicts the linear displacement field \(\mathbf{\Psi}_{ZA}(\mathbf{q})\), i.e., the Zel'dovich approximation (ZA) at redshift \(z=0\). We directly use the CNN architecture used in (Jamieson et al., 2022) and train it using 100 pairs of nonlinear and linear displacement fields. In terms of the training procedure, our method is almost identical to (Jamieson et al., 2022), with the only difference being that we reversed the inputs and outputs of the neural network. We now input the nonlinear displacement field to our CNN and ask it to predict the linear displacement field. This is exactly opposite to what was being done by (Jamieson et al., 2022). The source code of our implementation and experiments is available at github.com/vaibhavjindal/map2map/tree/inverse_mapping.

For our experiments, we train our model using simulations of \(128^{3}\) particles in a square box with a side-length of \(250\ \mathrm{Mpc}/h\). The particles are distributed uniformly across this grid with a mean separation of \(1.95\ \mathrm{Mpc}/h\) between two adjacent particles of the grid. This mean separation corresponds to a Nyquist wavenumber of \(k=1.608\ \mathrm{Mpc}^{-1}\ h\), the theoretical limit beyond which we can't trust the predictions of any model for this setting. For the training data, we randomly generated 100 linear fields for a fixed set of cosmological parameters (\(\Omega_{m}=0.300,\Omega_{b}=0.050,h=0.700,n_{s}=0.965,\sigma_{8}=0.799\)). We then use the emulator by (Jamieson et al., 2022) to generate the nonlinear displacement fields for these linear fields. We use this small set of pairs of nonlinear and linear fields to train our model.

The neural network architecture involves four convolutional layers, which determine its field-of-view to be \(96^{3}\). That is, a focal particle's initial displacement is predicted based on its environment out to its \(48^{\mathrm{th}}\) neighbour on the initial particle grid, corresponding to a distance of \(93.75\ \mathrm{Mpc}/h\). This finite field-of-view has a benefit: the evolution of the system on large scales is accurately described by linear theory, so the V-net model preserves this linear evolution on scales larger than the field-of-view while allowing for nonlinear evolution at smaller (medium range of) scales.

On very small scales, the particles clustered tightly into dark matter halos where their orbits become complicated and difficult to predict. This small-scale clustering imposes a resolution limit on the inverse mapping, making it a one-to-many mapping as numerical errors and floating point precision make the neural network model unable to fully encode the detailed textures of this virialized motion. This limitation is also present in the forward model, which tends to make halos more diffuse on small scales than in the simulations, blurring the details of these sharp structures. For the inverse mapping, the initial linear displacements are blurred for particles that end up inside of halos, limiting the accuracy of the inverse model predominantly on these very small scales.

To evaluate the in-distribution performance of our model, i.e., the performance on fields corresponding to the training set cosmological parameters, we analyze the results of our model on unseen pairs of nonlinear and linear fields with the same set of cosmological parameters as the training data (Sections 4.1, 4.2 and 4.3). To evaluate how well our model performs on unseen cosmological parameters, we present the out-of-distribution (OOD) evaluation on the Quijote simulation suite (Villaescusa-Navarro et al., 2020) in Section 4.4.

## 4 Results and Analysis

In the following section, we showcase our method's outputs both qualitatively and quantitatively with a variety of method and summary statistics.

### Qualitative Analysis

Figure 1 shows the \(x,y,\mathrm{and}\ z\) direction displacements for a \(128\times 128\) slice of the linear displacement field and the corresponding linear field predicted by our inverse model. Qualitatively, the predictions of our model match very well with the original linear field that we wanted to predict.

To further evaluate the quality of the prediction, we used the forward-direction emulator again to generate the nonlinear field when fed the predicted linear field. We qualitatively compare this regenerated nonlinear field with the original nonlinear field in Figure 2. We find these displacement fields match very nicely qualitatively. The residuals are also not significant (especially one should take into account that the color bar at the residual plot is approximately 10 times smaller).

### One-Point Statistics

The initial conditions of the simulation are set up in Fourier space, with each mode of the displacement field drawn from a random Gaussian distribution with a variance determined by the linear power spectrum \(P(k)\), which is determined by the cosmological parameters. This construction yields
a coordinate-space displacement field that is also Gaussian, so its statistics are uniquely determined by the two-point correlation function, which is simply the Fourier transform of the power spectrum. To demonstrate that we have accurately recovered the initial conditions of the simulation we must show that the output of our model has both the correct power spectrum and that its statistical distribution is Gaussian.

Figure 3: Two-point correlation comparison between the original linear field (target) and the linear field predicted by our model (prediction). For an exact prediction, the variation of power with wavenumber should be exactly similar for both the fields, and the values of transfer function fractional error and stochasticity should be exactly zero.

Figure 2: Qualitative comparison of the \(x,y,\) and \(z\) displacements for a \(128\times 128\) slice of particles from the original nonlinear field and the nonlinear field generated from the linear field predicted by our inverse model. The nonlinear field is generated from the forward modeling emulator (Jamieson et al., 2022). The slice shown here is the same as the one shown in Figure 1 for a one-to-one comparison.
We plot the histogram of probability density of the displacements of the particles for the original linear field and the one predicted by our model (Figure 5). Each particle has three displacements corresponding to the \(x,y,\) and \(z\) directions. Since the linear fields are sampled from a Gaussian distribution determined by the cosmological parameters, we expect the distribution of displacements to match \(\mathcal{N}(0,5.275)\) for the simulation setting with \(128^{3}\) particles in a box size of 250 Mpc/\(h\). Specifically, the variance of displacements along the \(i^{\rm th}\) Cartesian direction is given by:

\[\sigma_{i}^{2}=\int\frac{\mathrm{d}k^{3}}{(2\pi)^{3}}\frac{(k^{i})^{2}}{k^{4}} P(k), \tag{3}\]

where the integral is over all wave vectors in Fourier space, \(k\equiv|\vec{k}|\) is the magnitude of the wave vector and \(k^{i}\) is its \(i^{\rm th}\) component. We numerically evaluate this integral on a Fourier space grid with the same dimensions as the initial simulation grid.

From Figure 5 we see that the statistics of our model output agree well with the expected Gaussian distribution. The bin counts also match the particular realization from this distribution of the target data out to displacements of \(\sim 10~{}{\rm Mpc}/h\), or 2\(\sigma_{i}\). Extremely large displacements indicate a particle either flowing outwards in a rare, extremely underdense environment or inwards towards a rare, high-density peak. Based on the residuals, we see that the tails of the distribution predicted by the model are somewhat smaller than the target data, indicating the model misplaces these particles. This is unsurprising due both to the rarity of these trajectories and to the fact that these particles are the most affected by extreme non-linearities in their environments, which exacerbates the one-to-many problem.

### Two-Point Correlation Comparison

The displacement power spectrum for a displacement field \(\mathbf{\Psi}\) for wavenumber \(k\) is defined as

\[P(k)=\sum_{i\in\{x,y,z\}}\langle\mathbf{\Psi}_{i}(k)\mathbf{\Psi}_{i}(k)\rangle. \tag{4}\]

Using this definition of power spectrum in the Fourier space, we can now define the transfer function as

\[T(k)=\sqrt{P(k)}, \tag{5}\]

and the correlation coefficient as

\[r(k)=\frac{P_{pred\times true}(k)}{\sqrt{P_{pred}(k)P_{true}(k)}}, \tag{6}\]

Figure 4: Two-point correlation comparison between the original nonlinear field (target) and the nonlinear field generated from the predicted linear field given by the inverse model. For an exact prediction, the variation of power with wavenumber should be exactly similar for both the fields, and the values of transfer function fractional error and stochasticity should be exactly zero.

Figure 5: Distribution of the displacements of particles for a given linear field and the corresponding linear field predicted by our model. The distribution has been calculated by considering the \(x,y,\) and \(z\) displacements of \(128^{3}\) particles. Relative error denotes the relative errors in probability density for different bins.
where \(P_{pred}(k)\) is the displacement power spectrum predicted by our neural network, \(P_{true}(k)\) is the ground truth power spectrum, and \(P_{pred\times true}(k)\) is the cross power spectrum between the predicted and the ground truth fields. Using these two quantities, we define the transfer function fractional error,

\[\frac{\Delta T(k)}{T(k)}=\sqrt{\frac{P_{pred}(k)}{P_{true}(k)}}-1, \tag{7}\]

to measure the discrepancy between amplitudes of the predicted and the true fields. We also define stochasticity,

\[1-r^{2}(k)=1-\frac{P_{pred\times true}^{2}(k)}{P_{pred}(k)P_{true}(k)}, \tag{8}\]

to capture the excess fraction of correlation in the prediction of our model that cannot be accounted for in the target data. For an ideal match between the target and the predicted field, the values of both these quantities should be exactly zero. Figure 3 shows the performance of our model in terms of these quantities for an original linear displacement field and its corresponding predicted linear field. We see that from large scales down to scales where \(k\simeq 0.5\ \mathrm{Mpc}^{-1}\ h\) the model achieves percent-level accuracy for both the transfer functions and stochasticities. Note that non-linearities become important at scales where \(k>0.1\ \mathrm{Mpc}^{-1}\ h\), so the model is able to accurately learn the inverse mapping even in the moderately nonlinear regime for in-distribution fields.

To further test the reliability of our model, we generated the nonlinear displacement field from our predicted linear field by using the forward direction emulator (Jamieson et al., 2022). Since the mapping from the nonlinear displacement field to the linear displacement field is one-to-many, it becomes critical to see that the nonlinear field which gets generated from our predicted linear field matches the actual nonlinear field. Figure 4 shows the power spectra comparison between these two nonlinear displacement fields. We see that there is an excellent match between the powers of the two fields. The transfer function fractional errors and the stochasticity plots show a good match for large to medium scales (\(k<0.5\ \mathrm{Mpc}^{-1}\ h\)).

### Out-Of-Distribution Evaluation

Given that our model has been trained on simulations generated by a neural network emulator (Jamieson et al., 2022) with fixed cosmological parameters (\(\Omega_{m}=0.300,\Omega_{b}=0.050,h=0.700,n_{s}=0.965,\sigma_{8}=0.799\)), there is a legitimate concern about its ability to generalize to actual N-body simulations with different sets of cosmological parameters. In order to assess this out-of-distribution (OOD) performance, we have conducted a series of experiments on simulations from the Quijote suite (Villaescusa-Navarro et al., 2020).

The Quijote suite comprises 2000 linear and nonlinear fields, each consisting of \(512^{3}\) particles that are uniformly distributed in a cube with a side length of \(1000\ \mathrm{Mpc}/h\). To quantify the dissimilarity between these simulations and our training data, we rank them based on the sum of their relative percentage differences in \(\Omega_{m}\) and \(\sigma_{8}\) from our training values, since these two parameters have the most significant impact on the simulated distributions. This dissimilarity metric is then used to evaluate the model's performance on different percentiles, with the 100th percentile simulation representing the one on which we expect our model to perform the worst. Notably, our model's performance is unaffected by changes in the box size, as long as the mean particle density remains constant.

In order to evaluate the model's performance on out-of-distribution (OOD) Quijote simulation data, we qualitatively compare the \(x,y,\) and \(z\) displacements of a \(512\times 512\) slice from the linear and nonlinear fields, similar to Section 4.1. The slices for the 0th percentile (\(\Omega_{m}=0.296,\Omega_{b}=0.067,h=0.523,n_{s}=1.091,\sigma_{8}=0.806\)) and the 20th percentile (\(\Omega_{m}=0.311,\Omega_{b}=0.067,h=0.673,n_{s}=0.993,\sigma_{8}=0.977\)) simulations are presented in Figures 8 and 11 respectively. Additionally, the one-point statistics for these simulations are provided in Figures 6 and 9. It is worth mentioning that the Quijote simulations consist of \(512^{3}\) particles in a box with a side length of \(1000\ \mathrm{Mpc}/h\). Hence, based on Equation 3, the theoretical distributions for these simulations are \(\mathcal{N}(0,6.193)\) and \(\mathcal{N}(0,7.159)\) respectively. Furthermore, we present the two-point correlation comparisons for these simulations in Figures 7 and 10.

For completeness, we present the results of our power spectra analysis of the actual linear field and our model's predictions for a range of percentiles of the dissimilarity metric in Figure 12. We are able to match the in-distribution performance for the 0th percentile simulation and we are able to do reasonably well for the 20th percentile simulation for medium scales. The performance gradually gets worse as we move towards the higher percentiles of the dissimilarity metric. Additionally, in Figure 13, we compare the actual nonlinear fields with the nonlinear fields generated by passing the output of our model through the forward direction emulator (Jamieson et al., 2022). Once again, we see a good match for lower percentiles of the dissimilarity metric which progressively gets worse with increasing percentiles.
we show that neural networks are able to do recover a linear field which matches the original linear field for a wide range of scales, including scales affected by the nonlinear physics of gravitational clustering. The final simulation state is also accurately recovered after first being inverse mapped by our model, and then being forward mapped back to the final state. We empirically demonstrate that the model generalizes reasonably well to OOD cosmological parameters. This indicates that despite the ill-defined nature of the inverse-mapping problem, neural networks can still be successfully trained to accurately predict the linear fields at large and medium scales.

## 6 Future Work

In this work, we showed that neural networks are successful in predicting the initial states of the N-body simulation at a wide range of scales, going down to small scales where nonlinear gravitational effects kick in. The inherent one-to-many mapping between the final states and the initial states at smaller scales makes it impossible for N-body simulators to make such predictions. However, inverse neural network models, such as the one presented in this work, can be trained to predict approximate initial states of N-body simulations which are correct for a wide range of scales. These approximate initial states could then be used as initializations by further downstream sampling-based methods (e.g. Hamiltonian Monte Carlo, Active Learning, etc.) to refine their search for better initial states which are correct at even smaller scales.

Other directions for future work involve training on larger set of cosmological parameters, improving OOD performance, and uncertainty quantification for the initial states.

## Acknowledgements

This work is supported by Simons Foundation grant on Learning the Universe.
