Normalizing Flows for Random Fields in Cosmology

Adam Rouhiainen
Moritz Münchmeyer
Department of Physics
University of Wisconsin-Madison
Madison, WI 53706
rouhiainen@wisc.edu

Utkarsh Giri
Perimeter Institute for Theoretical Physics
Waterloo, ON N2L 2Y5
utkarshgiri18@gmail.com

Abstract
We study the use of normalizing flows to represent the field-level probability
density distribution of random fields in cosmology such as the matter and radiation
distribution. We evaluate the performance of the real NVP flow for sampling ofGaussian and near-Gaussian random fields, and N-body simulations, and check thequality of samples with different statistics such as power spectrum and bispectrum
estimators. We explore aspects of these flows that are specific to cosmology, suchas flowing from a physical prior distribution and evaluating the density estimation
results in the analytically tractable correlated Gaussian case.

1

Introduction and brief review of normalizing flows

Normalizing flows are a major recent development in probabilistic machine learning [1]. A flurry ofrecent work has shown their usefulness for a wide range of applications, such as image generation[2, 3], variational inference [4], and likelihood-free inference [5, 6]. In cosmology, normalizing
flows have recently been used to represent the posterior distribution of summary statistics such as thepower spectrum or cosmological parameters [7, 8]. In the present paper we use normalizing flows torepresent the probability density of fields such as the matter distribution directly at field level.
A normalizing flow is a natural way to construct flexible probability distributions by transforming
a simple base distribution (often Gaussian) into a complicated target distribution. This is done by
applying a series of learned diffeomorphisms to the base distribution. Given a base distribution pu (u)−1of a random variable u, the target distribution px (x) is given by px (x) = pu (u) |det JT (u)| ,where T is the transformation (the “flow"), x = T (u), and JT is its Jacobian. We can construct atransformation T = TK ◦ · · · ◦ T1 by composing simple transformations Tk , depending on learned
parameters and are parametrized using neural networks. In this way very expressive densities canbe constructed. Assuming z0 = u and zK = x, the transformation at each step k is zk = Tk (zk−1 )PKand the Jacobian determinant is log |JT (z)| = k=1 log |JTk (zk−1 )|. Once the flow is learned, twobasic statistical operations are performed efficiently: density evaluation and sampling.
While not as expressive as GANs or VAEs, it is a priori plausible that normalizing flows could beparticularly strong at representing PDFs of fields in cosmology. Many fields in cosmology, suchas the matter distribution, start as Gaussian fields in the far past. They then become progressively
more non-Gaussian with time due to non-linear interactions. In the same way, a Gaussian field basedistribution of a normalizing flow becomes progressively more non-Gaussian by the applicationsof more transformations Tk . Such a normalizing flow is therefore a natural candidate to represent
somewhat (non-)Gaussian PDFs of matter fields at late times.
In this work, we utilize the real-valued non-volume preserving (real NVP) flow [3]. The real NVP
flow is expressive and is fast both for sampling and inference, and has also recently received a lot ofattention for its use to represent PDFs in lattice field theory (see e.g. [9]). The details of our network
Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

resemble [9], which we have modified for this work under the CC BY license. We use 3 convolutional
layers with kernel size 3 and leaky ReLU activation functions. The number of channels is in this order:
1 (the scalar PDF values), 16 (arbitrary number of feature maps), 16 (arbitrary number of featuremaps), 2 (the output variables s and t). As in [9], we stack K = 16 affine coupling layers, each with
their own CNN to parametrise the affine transformation s and t. We use stride 1 convolutions andno pooling. To implement periodic boundary conditions we use circular padding. To train the flow,
we minimize the forward Kullback–Leibler (KL) divergence [1], the relative entropy from the target
distribution p∗x (x) to the base distribution px (x). As a function of learned parameters ϕ,L(ϕ) ≈ −N1 Xlog pu T −1 (xn ; ϕ) + log |det JT −1 (xn ; ϕ)| + const.,
N n=1

(1)

where T is the flow transformation. We use the Adam optimizer to minimize the loss with respect
to the parameters ϕ, with a learning rate of 0.001. We also normalize the training samples to unit
variance and use a batch size of 128. Our results are reproducible in Jupyter Notebooks available atgithub.com/SubmissionForPapers/NormalizingFlowsCosmology.

2Application to Gaussian fields with density evaluation

We train our real NVP flow on samples from a correlated Gaussian field, with a CMB temperature
power spectrum. Here we flow from an uncorrelated Gaussian prior distribution. We use patches of642 pixels covering a sky angle side length of 4 degrees. In this first example, we use infinite on
the fly created training samples. Training time was about 40 hours on an RTX 3090 and used about
3.8 GB of GPU memory. Samples from the prior, model, and target (training) distribution are shownin Fig. 1 (left). By eye, the model samples look like the training data.
Gaussian task
Prior samples Prior samples

Local non-Gaussian task
Prior samples Prior samples

Model samples Model samplesModel samples Model samples

Target samples Target samples

Target samples Target samples

Figure 1: Prior samples (top), flow samples (middle), and training samples (bottom). Left: Gaussian
local
field samples. Right: Local non-Gaussian field samples with f˜N
L = 0.2. The non-Gaussian training
data is generated by drawing Gaussian maps from a CMB power spectrum on a 4 degree sky patch
represented on 642 pixels and making them non-Gaussian with Eq. 2. The network makes thecorrelated Gaussian prior samples more non-Gaussian, with extrema becoming more pronounced.
For some of the data analysis applications discussed in the introduction, we need to use the flowin reverse direction for density evaluation. There are two different tasks we can consider here:
in distribution density evaluation (IID: independent, identical distribution) and out of distribution(OOD) density evaluation; we focus here on IID density evaluation. We would like to verify that
IID samples, when run backwards through the flow (with uncorrelated Gaussian noise prior), areassigned a probability that corresponds to their true probability. Here the fact that we start with atractable Gaussian distribution allows us to compare the flow probability with the exact analytic
probability. We sample 10,000 new IID samples x from the same distribution as the training data2

Table 1: Local non-Gaussianity measurements
local
f˜N
L training

local
f˜N
L flowequif˜N
L training

equi
f˜N
L flow

0.2 ± 0.074
0.05 ± 0.024

0.187 ± 0.085
0.047 ± 0.029

0.150 ± 0.044
0.037 ± 0.017

0.146 ± 0.040
0.035 ± 0.017

and reverse flow them to obtain log P flow (x). For the configuration above, we find that the crosscorrelation coefficient between log P flow (x) and log P true (x) is r ≃ 0.993. We plot these quantities
for 200 random example maps in Fig. 2. When we limit training data to 1,000 Gaussian maps, thecross correlation coefficient dropped to about r ≃ 0.97, while with 10,000 training maps we foundr ≃ 0.98. The flow thus learns density evaluation on IID samples rather well.
150flow
truthlog P(x)10050050100150

0255075100map nr.

125150175

200

Figure 2: Density evaluation log P (x) result from the flow compared to the true log probability ofthe sample (mean shifted to zero), for 200 maps with 642 pixel resolution. The cross-correlation
coefficient is 0.993, estimated from 10,000 maps.

3

Local non-Gaussianity and correlated prior

We consider the simplest non-Gaussian random field in cosmology, the local non-Gaussianity. Local
non-Gaussianity [10] is generated by transforming a correlated Gaussian field ϕG (x) as
localϕN G (x) = ϕG (x) + f˜N
ϕG (x)2 − ϕG (x)2(2)LTo draw samples from this distribution, we square the Gaussian field and add it to the original fieldwith some amplitude. This form of non-Gaussianity is generated in cosmology for example by multilocal
field inflation [11]. We normalize ϕG to variance 1 so that f˜N
L = 1 indicates O(1) non-Gaussianity.
This time we flow from a correlated Gaussian field with the correct power spectrum. This makes thetraining much more efficient as the flow has a better starting point. We show samples from the prior,model, and training distribution in Fig. 1 (right). The samples are indistinguishable by eye from thetraining data. We show the power spectrum of the samples, prior, and training data in Fig. 3 (left).
The flow thus learns to induce the right non-Gaussianity while keeping the power spectrum intact.
The correlated prior also improves the training convergence. We measure the non-Gaussianity inthe samples by estimating the amplitude of two non-Gaussian templates, the local and equilateral
non-Gaussianity. In the present case the training data has local non-Gaussianity by definition, however
the equilateral template has some overlap with the local template, and it is useful to measure both in
the model samples whose non-Gaussianity can be non-local. We record mean estimated values andper sample variance of the estimated non-Gaussianity in Table 1, averaging over 10,000 training andflow samples. The flow is accurate to about 5% in the mean and 10% in the per sample variance.

4

Non-Gaussian fields from N-body simulations

Here we tackle the important use case of representing field PDFs from N-body simuations. We use
100 high-resolution simulations from the Quijote [12] suite of N-body simulation data (providedunder the MIT License) to generate training patches of the 2D matter field. The simulations were run
in a box size of 1 h−1 Gpc and use 10243 dark matter particles. We use snapshots generated at z = 2
to estimate the matter density by painting the matter particles on a 3D mesh of size 10243 using theCloud-in-Cell [13] mass assignment scheme implemented in nbodykit [14]. We are able to resolve3

10 12C

2C

2 × 100

C trueC flow
C prior100

C trueC flow
C prior

10010 2
6 × 10 10500 1000 1500 2000 2500 3000 3500

05000

10000

15000

20000

Figure 3: Left: Comparison of the prior, model, and training set power spectrum averaged over10,000 maps, for the fNL = 0.2 example. With the correlated prior, the BAO wiggle is resolved
perfectly. Right: Projected matter power spectrum of the N-body simulation samples (red), trained
flow model (black) and correlated input prior (green), averaging over 10,000 maps.
Table 2: N-body non-Gaussianity measurements
localf˜N
L traininglocalf˜NL flow

equi
f˜N
L trainingequif˜NL flow

0.85 ± 0.041

0.78 ± 0.29

0.61 ± 0.24

0.62 ± 0.24

nonlinear scales of size k ≳ 1 h−1 Mpc. We divide the simulation volume into boxes of size 1283
and project the density to get a realization of a 2D field. We generated 51,200 independent patches.
Limited by GPU memory, larger maps would require a method to patch these 1283 maps together.
We train the same real NVP flow as the previous sections. Training time was 50 hours on an
RTX 3090, using 11 GB of memory. We found faster convergence and better sample quality from acorrelated Gaussian prior with the right power spectrum compared to flowing from Gaussian noise. InFig. 3 (right) we find a near perfect power spectrum of the flow. Non-Gaussianity measurements arerecorded in Table 2, sampling over 10,000 training and model maps. The flow is accurate to 5 to 10%.

Model
samplesN-bodymodel
samplesTargettarget
samplesN-bodysampleFigure 4: N-body simulation samples trained on real NVP flow. Each sample is derived from a 3D
volume of (125 h−1 Mpc)3 by projecting the matter field such that the 2D map has a resolution of1282 pixels. Left: Model samples from the trained flow. Right: Sample from the training data.

5Conclusion

The goal of this study is to evaluate how good normalizing flows are at representing PDFs of random
fields in cosmology, both for sampling and density evaluation. We find that we can train flows whose
samples have power spectra and non-Gaussianity accurate to a few percent. We tested the inverse
operation of density evaluation in the analytically tractable Gaussian case and found percent levelaccuracy in log P . Using the non-Gaussianity estimation formalism of cosmology [15], we compared
the non-Gaussianity measurement of the flow distribution to the target distribution. Measuring threepoint correlation functions in statistical machine learning may be generally useful to quantify theperformance of normalizing flows in representing PDFs. In upcoming work we will apply these flows
to inference problems in cosmology. For example, analysing a small but deep sky patch with HMC
such as for CMB lensing in [16] would benefit from a non-Gaussian prior distribution of the lensing
4

potential represented with a normalizing flow. Normalizing flows are also suitable for variational
inference of the initial conditions of the matter distribution of the universe. For some applications
it would be useful to make the flow PDF conditionally dependent on astrophysical or cosmologicalparameters. Memory constraints will require a patching scheme to represent larger maps, or 3D maps.

References[1] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji
Lakshminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXiv e-prints,
arXiv:1912.02762, December 2019.
[2] Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative Flow with Invertible 1x1 Convolutions. arXiv e-prints, arXiv:1807.03039, July 2018.
[3] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP.
arXiv e-prints, arXiv:1605.08803, May 2016.
[4] Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows.
arXiv e-prints, arXiv:1505.05770, June 2016.
[5] George Papamakarios, Theo Pavlakou, and Iain Murray. Masked Autoregressive Flow forDensity Estimation. arXiv e-prints, arXiv:1705.07057, May 2017.
[6] Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference.
arXiv e-prints, arXiv:1911.01429, November 2019.
[7] Justin Alsing, Tom Charnock, Stephen Feeney, and Benjamin Wandelt. Fast likelihood-freecosmology with neural density estimators and active learning. Mon. Not. Roy. Astron. Soc.,
488(3):4440–4458, February 2019.
[8] Ana Diaz Rivero and Cora Dvorkin. Flow-Based Likelihoods for Non-Gaussian Inference.
Phys. Rev. D, 102(10):103507, July 2020.
[9] Michael S. Albergo, Denis Boyda, Daniel C. Hackett, Gurtej Kanwar, Kyle Cranmer, Sébastien
Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan. Introduction to Normalizing
Flows for Lattice Field Theory. arXiv e-prints, arXiv:2101.08176, January 2021.
[10] Eiichiro Komatsu and David N. Spergel. Acoustic signatures in the primary microwave background bispectrum. Phys. Rev. D, 63:063002, February 2001.
[11] Christian T. Byrnes and Ki-Young Choi. Review of local non-Gaussianity from multi-field
inflation. Adv. Astron., 2010:724525, February 2010.
[12] Francisco Villaescusa-Navarro, ChangHoon Hahn, Elena Massara, Arka Banerjee, Ana Maria
Delgado, Doogesh Kodi Ramanah, Tom Charnock, Elena Giusarma, Yin Li, Erwan Allys, andet al. The Quijote simulations. The Astrophysical Journal Supplement Series, 250(1):2, August2020.
[13] R. W. Hockney and J. W. Eastwood. Computer Simulation Using Particles. 1981.
[14] Nick Hand, Yu Feng, Florian Beutler, Yin Li, Chirag Modi, Uros Seljak, and Zachary Slepian.
nbodykit: an open-source, massively parallel toolkit for large-scale structure. Astron. J.,
156(4):160, September 2018.
[15] J. R. Fergusson, D. M. Regan, and E. P. S. Shellard. Rapid Separable Analysis of Higher Order
Correlators in Large Scale Structure. Phys. Rev. D, 86:063511, September 2012.
[16] M. Millea et al. Optimal CMB Lensing Reconstruction and Parameter Estimation with SPTpol
Data. arXiv e-prints, arXiv:2012.01709, December 2020.

5

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes] See Section 1, Section 4, andSection 5.
(c) Did you discuss any potential negative societal impacts of your work? [No] We modify
an existing generative model for cosmological applications.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] See Section 1.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)? [Yes] See Section 1.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] See Table 1 and Table 2.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes] See Section 2 and Section 4.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See Section 1 for thereal NVP code and Section 4 for the Quijote data.
(b) Did you mention the license of the assets? [Yes] See Section 1 for the real NVP modeland Section 4 for the Quijote data.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]See Section 1.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [No] The Quijote data was obtained from physical simulations.The
licenses for the real NVP code and Quijote data give us permission, and we used themfor their intended purposes.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [No] We don’t believe it does; we use cosmological
data from simulations.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]

6