# Monitoring Shortcut Learning using Mutual Information

Mohammed Adnan  Yani Ioannou  Chuan-Yung Tsai  Angus Galloway  H.R. Tizhoosh
Graham W. Taylor

###### Abstract

The failure of deep neural networks to generalize to out-of-distribution data
is a well-known problem and raises concerns about the deployment of trained
networks in safety-critical domains such as healthcare, finance and autonomous
vehicles. We study a particular kind of distribution shift — _shortcuts_ or
_spurious correlations_ in the training data. Shortcut learning is often only
exposed when models are evaluated on real-world data that does not contain the
same spurious correlations, posing a serious dilemma for AI practitioners to
properly assess the effectiveness of a trained model for real-world
applications. In this work, we propose to use the mutual information (MI)
between the learned representation and the input as a metric to find where in
training the network latches onto shortcuts. Experiments demonstrate that MI
can be used as a domain-agnostic metric for monitoring shortcut learning.

shortcuts, spurious correlation, information bottleneck

##  1 Introduction

Our understanding of ‘how’ and ‘what’ neural networks learn is limited, which
raises concern about the deployment of neural networks in safety-critical
domains. Despite achieving state-of-the-art performance on benchmark datasets,
neural networks may fail to generalize in real-world settings or for out-of-
distribution data (Koh et al., 2021). For example, models trained for cancer
detection may not generalize on data from a new hospital (Castro et al., 2020;
Perone et al., 2018; AlBadawy et al., 2018) and self-driving cars may not
generalize to new lighting conditions or object poses (Alcorn et al., 2018;
Dai & Van Gool, 2018). One reason why models may fail in real-world settings
could be attributed to learning _shortcuts_ (Geirhos et al., 2020) from the
training data. A _shortcut_ is a type of distribution shift where spurious
correlations exist only in the training data, resulting in the learning of
non-intended or easy-to-learn discriminatory features which work well on the
training and test dataset but not on out-of-distribution real-world datasets
(Wiles et al., 2022; Geirhos et al., 2020). Shortcuts can arise due to dataset
biases or the model using ‘trivial’ or unintended features like high-frequency
noise patterns or the image background for the classification task. For
example, an Inception-V3 model trained to detect hip fractures was found to
use scanner information for learning discriminatory features (Badgeley et al.,
2019); deep learning systems trained to detect COVID-19 from chest radiographs
can rely on confounding factors (shortcuts) rather than medical pathology
(DeGrave et al., 2021). While our understanding of shortcuts and how they
arise is still developing, a helpful tool to practitioners deploying machine
learning models in safety-critical domains with a high cost of failure would
be to monitor shortcuts during the training phase. Although the phenomenon of
shortcut learning is widely known, there is no effective method available to
monitor shortcuts being learned. Interpretable machine-learning methods such
as feature attribution, Grad-CAM (Selvaraju et al., 2017), and LIME (Ribeiro
et al., 2016) have been used to understand a model’s dependency on spurious
correlations. However, it has been shown that such post-hoc explanations are
ineffective (Adebayo et al., 2021; Alqaraawi et al., 2020; Chu et al., 2020).

In this work, we show that shortcut learning can be understood using the
information bottleneck framework (Tishby et al., 1999; Tishby & Zaslavsky,
2015), by using the mutual information (MI) between the inputs and the learned
representation to monitor shortcut learning. We use the neural tangent kernel
(NTK) (Jacot et al., 2018) to study the training evolution of shortcut
learning. We design experiments using synthetic and complex real-world data to
demonstrate the relationship between (MI) and shortcut learning, and show that
MI can be used as a metric for domain-agnostic assessment of shortcuts. We
find that compression as measured by MI is associated with the tendency to
learn shortcuts.

##  2 Background

### Shortcut learning:

Wiles et al. (2022) defined shortcuts or spurious correlations as a type of
_distribution shift_ such that two or more attributes are correlated at
training time, but not for the test data, where they are independent. In a
more general sense, shortcuts are _easy-to-learn_ decision rules that can be
exploited in the absence of distribution shift (i.e. on standard benchmarks)
but fail to transfer to more challenging and diverse testing conditions, such
as real-world datasets (Geirhos et al., 2020).

### Information bottleneck method:

The information bottleneck (IB) can be viewed as a rate-distortion problem
cast entirely in terms of mutual information (MI)—denoted “$I(V;Z)=I(Z;V)$”
(Tishby et al., 1999). A  _distortion_ function measures how well a relevant
variable $V$ is predicted from another variable $Z$, where $Z$ is usually a
compressed representation of the input $X$. The  _rate_ refers to the
complexity of $Z$, which is less than or equal to $X$. IB is a general method
for data compression but has been advanced as an explanatory tool predictive
of learning and generalization of neural networks (NNs). It has been suggested
that NNs trained by SGD may learn compressed representations $Z$ of their
input, making them insensitive to data idiosyncrasies, yet maintain sufficient
_relevant_ information for predicting the output $Y$ (e.g. class labels)
(Tishby & Zaslavsky, 2015; Shwartz-Ziv & Tishby, 2017). This trade-off between
compression and preserving task relevant information is optimized by the
notion of “minimal sufficient statistics” (Cover & Thomas, 1991). The IB view
suggests that NNs trained by cross-entropy loss may implicitly minimize the
following Lagrangian:

$\min I(X;Z)-\beta I(Z;Y),$ (1)

enabling them to implement minimal sufficient statistics for different
$\beta$-constraints on the error.111To what extent the relationship between IB
and deep learning holds in general is the subject of ongoing debate (Saxe et
al., 2018; Jacobsen et al., 2018; Goldfeld et al., 2019).

### Neural tangent kernel:

Distribution-free estimation of MI for high-dimensional data is challenging
and often intractable. One workaround to this problem is using an infinite
ensemble of infinite-width neural networks (Shwartz-Ziv & Alemi, 2020) which
confer tractable bounds on MI. The neural tangent kernel (NTK) (Jacot et al.,
2018) is a kernel that describes the evolution of infinite-width neural
networks during their training by gradient descent, thus allowing the
systematic study of neural networks using tools from kernel methods. Infinite-
width neural networks behave as linear functions, and their training evolution
can be fully described by the NTK. Shwartz-Ziv & Alemi (2020) used the fact
that the output of an infinite ensemble of infinitely-wide neural networks
initialized with Gaussian weights and biases and trained by MSE loss is a
conditional Gaussian distribution. This allows tractable computation of (MI)
between the representation $Z$ and the targets $Y$: $I(Z;Y)$, and the MI
between $Z$ and the inputs $X$ during training: $I(X;Z)$.

Figure 1: Sample images from MNIST dataset. A small patch is added to all
images of even digits as a shortcut.

##  3 Shortcuts and information bottleneck

Our hypothesis is that MI as measured during the evolution of a network’s
parameters can be used to monitor shortcut learning as it occurs. Shortcuts
allow networks to learn a maximally compressed representation $Z$, i.e.,
$I(X;Z)$ is considerably reduced and thus can be used as a metric to monitor
exploitation of shortcuts.

To support our hypothesis, we design controlled datasets containing spurious
correlations and measure $I(X;Z)$ during the training evolution using the NTK.
Using insights from the concept of a previously introduced “mutual information
plane” (Shwartz-Ziv & Tishby, 2017), we can observe shortcut learning as it
happens and find the time where the network stops exploring the preferred
(generalizable) solution space and latches onto the spurious signal.

##  4 Experiments and observations

### Experimental Setup:

For each experiment, we plot $I(X;Z)$ and $I(Z;Y)$ w.r.t. time, generalization
error/loss on clean data (without spurious correlations), and the information
plane ($I(Z;Y)$ vs. $I(X;Z)$). Note that we calculate the upper bound of
$I(X;Z)$, and are only interested in the training dynamics rather than the
precise value of MI.

### MNIST with synthetic shortcut:

We train a model to classify MNIST images into odd and even digits. We add a
small white patch on one corner of _all even_ digits of the MNIST training
dataset as a spurious correlation (Figure 1). The network can use the patch to
accurately classify the images into odd and even. We compare the MI during the
training evolution on datasets with and without shortcuts.

(a) Information Plane

(b) Clean Test Loss

(c) $I(X;Z)$

(d) $I(Z;Y)$

Figure 2: Comparison of training evolution with and without synthetic
shortcuts on MNIST dataset. An upper bound on MI is plotted. (a) Information
plane plotting $I(X;Z)$ versus $I(Z;Y)$ during training. (b) Loss w.r.t. the
clean test data without shortcuts—the solid blue line corresponds to a model
trained with shortcuts, the broken orange line is without shortcuts. (c) Plot
of $I(X;Z)$ versus training time step. (d) Plot of $I(Y;Z)$ during the
training evolution. Animated GIF of the plot can be viewed here.

In 2(a) and 2(c), we observe that the mutual information $I(X;Z)$ increases
initially during training but then latches onto the shortcuts, after which the
mutual information decreases sharply. It can also be observed in 2(b) that
generalization error increases after the point at which MI starts to decrease,
indicating that the network explores the more optimal region of the solution
space in the initial training epochs before discovering shortcuts. This is
consistent with the findings of Schwartz-Ziv and Tishby (2017) about the
behaviour of SGD. In the initial phase, SGD explores the multidimensional
space of solutions. When it begins converging, it arrives at the diffusion
phase in which the network learns to compress (Shwartz-Ziv & Tishby, 2017). In
both settings, the model achieves high training accuracy, i.e., $I(Z;Y)$
(2(d)) but the difference in test set loss is striking (2(b)).

Figure 3: Saliency map of model trained on MNIST images with a small patch on
even images as a shortcut.

### Visualisation:

To visually verify that the network is learning the shortcut in the above
experiment, we generate a saliency map. We use the finite-difference
estimation to find the gradient of the class probability w.r.t. to the input
pixels. The network predominantly uses the shortcut in the image to predict
the class (Figure 3).

### Effect of partially correlated shortcuts:

In real-world data, shortcuts are often partially correlated with the output,
i.e., the model cannot classify with $100\%$ accuracy using only the
shortcuts. To understand the effect of partially correlated shortcuts on the
training dynamics, we construct different datasets with varying degrees of
shortcut efficacy. Instead of corrupting all the even images, we add a small
white patch on one corner only to a specific percentage (50%–100%) of even
images. We plot the MI for varying degrees of corruption in Figure 4 for 1000
training points sampled uniformly.

We observe that as the effectiveness of the shortcut increases, $I(X;Z)$
converges to a lower value indicating the ability to perform more compression.
We also note an interesting behaviour during the training evolution: when the
shortcut is partially correlated, MI does not decrease significantly as
compared to the 100% effective shortcut. We speculate that, while in these
cases (e.g. 80–90% shortcut efficacy) the model is able to recover some
generalization ability, its ability to discover high generalization solutions
is irrevocably deteriorated once it discovers the minima corresponding to the
shortcut solution.

(a) Information Plane

(b) Clean Test Loss

(c) $I(X;Z)$

(d) $I(Z;Y)$

Figure 4: Effect of shortcut effectiveness on the MI trajectory. Shortcut is
added to different percentages of even images in each experiment. With
increase in shortcut effectiveness, MI converges to a lower value. Animated
GIF of the plot can be viewed here.

### CelebA with natural shortcuts:

We also test our hypothesis on a dataset containing natural shortcuts. We
curate images from the CelebA dataset such that all images tagged as male have
the “black hair color” attribute, while images tagged as female have the
“blonde hair color” attribute. We train the network to classify facial images
into the male and female categories222The CelebA dataset only provides binary
labels and we do not know how the gender attribute was assigned. Therefore it
should be considered as nothing more than an arbitrary class in this
experiment., while the network may use hair color as a shortcut for accurately
classifying the images. Since the dataset is not controlled, some images from
both classes have a black color in different parts of the image (background,
clothes, etc.), reducing the effectiveness of the hair color attribute. We
plot the MI trajectory with and without shortcuts for 100 sample training
points sampled uniformly in $\log$ scale (Figure 5). We observe MI profile
similar to Figure 4. On data without the shortcut, MI increases consistently,
while in the presence of shortcuts, MI converges to a lower value, therefore
validating our hypothesis on real-world data with natural shortcuts.

(a) Information Plane

(b) Clean Test Loss

(c) $I(X;Z)$

(d) $I(Z;Y)$

Figure 5: Mutual information profile for CelebA dataset with hair attribute as
shortcut. Plot of $I(X;Z)$ in (a) and (c) show that shortcuts results in
reduced $I(X;Z)$. Animated GIF of the plot can be viewed here.

### Effect of shortcut on the loss landscape:

We visualize the loss landscape of neural networks to understand the effect of
shortcuts on the optimization trajectory. We plot loss along a linear path
connecting the initial parameter $\theta_{\text{o}}$ and converged parameter
$\theta^{*}$ in the weight space (Goodfellow et al., 2014) and polar
coordinates ($r_{t},\phi_{t}$) plot measuring the deviation from the linear
line between $\theta_{i}$ and $\theta^{*}$ (Figure 6). We parameterize the
line with $\alpha$ such that $\theta=(1-\alpha)\theta_{i}+\alpha\theta^{*}$.
Polar coordinates can be calculated using
$r_{t}=\frac{||\triangle\theta_{t}||}{||\triangle\theta_{\text{o}}||}$ and
$\phi_{t}=\arccos{\frac{\triangle\theta_{t}\times\triangle\theta_{\text{o}}}{||\triangle\theta_{t}||\times||\triangle\theta_{\text{o}}||}}$,
where $\triangle\theta_{t}=\theta_{t}-\theta^{*}$. We observe that the loss
landscape around $\theta^{*}$ in the case of shortcuts is surprisingly flat as
compared to the valley-like shape for a model trained on data not containing
shortcuts using the MNIST dataset. The polar plot shows that the optimizer
deviates less from the linear trajectory when trained with shortcuts.

(a) Radial plot

(b) Loss landscape visualization

Figure 6: Visualization of loss landscape on MNIST dataset. (a). polar
coordinates ($r_{t},\phi_{t})$ measuring the deviation from the linear path
between initialisation and converged parameters in the weight space during the
optimization. (b). 1-D visualization of loss landscape.

##  5 Conclusion and future work

In this work, we sought to understand why networks tend to to learn shortcuts
through the lens of the information bottleneck method. We showed that mutual
information can be used to monitor training dynamics w.r.t. shortcut learning
without using any domain knowledge; this is an advantage compared to methods
adapted by the interpretable ML literature, where domain knowledge is
required. However, we used the NTK to estimate mutual information, limiting
our approach to infinite-width neural networks. In the future, we will address
this limitation and develop methods to avoid shortcut learning using tractable
MI estimators that place less severe constraints on the model architecture
(Selby et al., 2022; Gabrié et al., 2018). We are also interested in exploring
the relationship between the curvature of the solution minima and shortcut
learning.

## References

  * Adebayo et al. (2021) Adebayo, J., Muelly, M., Abelson, H., and Kim, B.  Post hoc explanations may be ineffective for detecting unknown spurious correlation.  In _International Conference on Learning Representations_ , 2021.
  * AlBadawy et al. (2018) AlBadawy, E. A., Saha, A., and Mazurowski, M. A.  Deep learning for segmentation of brain tumors: Impact of cross-institutional training and testing.  _Medical physics_ , 45(3):1150–1158, 2018.
  * Alcorn et al. (2018) Alcorn, M. A., Li, Q., Gong, Z., Wang, C., Mai, L., Ku, W.-S., and Nguyen, A.  Strike (with) a pose: Neural networks are easily fooled by strange poses of familiar objects, 2018.
  * Alqaraawi et al. (2020) Alqaraawi, A., Schuessler, M., Weiß, P., Costanza, E., and Berthouze, N.  Evaluating saliency map explanations for convolutional neural networks: a user study.  In _Proceedings of the 25th International Conference on Intelligent User Interfaces_ , pp. 275–285, 2020.
  * Badgeley et al. (2019) Badgeley, M. A., Zech, J. R., Oakden-Rayner, L., Glicksberg, B. S., Liu, M., Gale, W., McConnell, M. V., Percha, B. L., Snyder, T. M., and Dudley, J. T.  Deep learning predicts hip fracture using confounding patient and healthcare variables.  _NPJ Digital Medicine_ , 2, 2019.
  * Castro et al. (2020) Castro, D. C., Walker, I., and Glocker, B.  Causality matters in medical imaging.  _Nature Communications_ , 11, 2020.
  * Chu et al. (2020) Chu, E., Roy, D., and Andreas, J.  Are visual explanations useful? a case study in model-in-the-loop prediction.  _arXiv preprint arXiv:2007.12248_ , 2020.
  * Cover & Thomas (1991) Cover, T. M. and Thomas, J. A.  _Elements of Information Theory_.  John Wiley & Sons, Inc., New York, 1991.
  * Dai & Van Gool (2018) Dai, D. and Van Gool, L.  Dark model adaptation: Semantic image segmentation from daytime to nighttime, 2018.
  * DeGrave et al. (2021) DeGrave, A. J., Janizek, J. D., and Lee, S.-I.  Ai for radiographic covid-19 detection selects shortcuts over signal.  _Nature Machine Intelligence_ , 3(7):610–619, 2021.
  * Gabrié et al. (2018) Gabrié, M., Manoel, A., Luneau, C., Macris, N., Krzakala, F., Zdeborová, L., et al.  Entropy and mutual information in models of deep neural networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Geirhos et al. (2020) Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A.  Shortcut learning in deep neural networks.  April 2020.
  * Goldfeld et al. (2019) Goldfeld, Z., Van Den Berg, E., Greenewald, K., Melnyk, I., Nguyen, N., Kingsbury, B., and Polyanskiy, Y.  Estimating Information Flow in Deep Neural Networks.  In _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2299–2308. PMLR, 2019.
  * Goodfellow et al. (2014) Goodfellow, I. J., Vinyals, O., and Saxe, A. M.  Qualitatively characterizing neural network optimization problems.  _arXiv preprint arXiv:1412.6544_ , 2014.
  * Jacobsen et al. (2018) Jacobsen, J.-H., Smeulders, A. W. M., and Oyallon, E.  I-RevNet: Deep Invertible Networks.  In _International Conference on Learning Representations_ , 2018.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  June 2018.
  * Koh et al. (2021) Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., et al.  Wilds: A benchmark of in-the-wild distribution shifts.  In _International Conference on Machine Learning_ , pp. 5637–5664. PMLR, 2021.
  * Perone et al. (2018) Perone, C. S., Ballester, P., Barros, R. C., and Cohen-Adad, J.  Unsupervised domain adaptation for medical imaging segmentation with self-ensembling, 2018.
  * Ribeiro et al. (2016) Ribeiro, M. T., Singh, S., and Guestrin, C.  ”why should i trust you?”: Explaining the predictions of any classifier, 2016.
  * Saxe et al. (2018) Saxe, A. M., Bansal, Y., Dapello, J., Advani, M., Kolchinsky, A., Tracey, B. D., and Cox, D. D.  On the Information Bottleneck Theory of Deep Learning.  In _International Conference on Learning Representations_ , 2018.
  * Selby et al. (2022) Selby, K. A., Rashid, A., Kobyzev, I., Rezagholizadeh, M., and Poupart, P.  Learning functions on multiple sets using multi-set transformers.  In _The 38th Conference on Uncertainty in Artificial Intelligence_ , 2022.  URL https://openreview.net/forum?id=HzMEEOUs5x5.
  * Selvaraju et al. (2017) Selvaraju, R. R., Cogswell, M., Das, A., Vedantam, R., Parikh, D., and Batra, D.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 618–626, 2017.
  * Shwartz-Ziv & Alemi (2020) Shwartz-Ziv, R. and Alemi, A. A.  Information in infinite ensembles of infinitely-wide neural networks.  In Zhang, C., Ruiz, F., Bui, T., Dieng, A. B., and Liang, D. (eds.), _Proceedings of The 2nd Symposium on Advances in Approximate Bayesian Inference_ , volume 118 of _Proceedings of Machine Learning Research_ , pp. 1–17. PMLR, 08 Dec 2020.  URL https://proceedings.mlr.press/v118/shwartz-ziv20a.html.
  * Shwartz-Ziv & Tishby (2017) Shwartz-Ziv, R. and Tishby, N.  Opening the black box of deep neural networks via information. 2017.  _arXiv preprint arXiv:1703.00810_ , 2017.
  * Tishby & Zaslavsky (2015) Tishby, N. and Zaslavsky, N.  Deep learning and the information bottleneck principle.  In _2015 ieee information theory workshop (itw)_ , pp. 1–5. IEEE, 2015.
  * Tishby et al. (1999) Tishby, N., Pereira, F. C., and Bialek, W.  The information bottleneck method.  In _Allerton Conference on Communication, Control and Computing_ , 1999.
  * Wiles et al. (2022) Wiles, O., Gowal, S., Stimberg, F., Rebuffi, S.-A., Ktena, I., Dvijotham, K. D., and Cemgil, A. T.  A fine-grained analysis on distribution shift.  In _International Conference on Learning Representations_ , 2022.


#  Physical Benchmarking for AI-Generated
Cosmic Web

Xiaofeng Dong
Department of Physics,
University of Chicago,
Chicago, IL 60637, USA
High Energy Physics Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>
Nesar Ramachandra
Computational Science Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>
Salman Habib
Computational Science Division,
High Energy Physics Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>
Katrin Heitmann
High Energy Physics Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>
Michael Buehlmann
High Energy Physics Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>
Sandeep Madireddy
Mathematics and Computer Science Division,
Argonne National Laboratory,
Lemont, IL 60439, USA
<EMAIL_ADDRESS>

###### Abstract

The potential of deep learning based image-to-image translations has recently
drawn a lot of attention; one intriguing possibility is that of generating
cosmological predictions with a drastic reduction in computational cost. Such
an effort requires optimization of neural networks with loss functions beyond
low-order statistics like pixel-wise mean square error, and validation of
results beyond simple visual comparisons and summary statistics. In order to
study learning-based cosmological mappings, we choose a tractable analytical
prescription – the Zel’dovich approximation – modeled using U-Net, a
convolutional image translation framework. A comprehensive list of metrics is
proposed, including higher-order correlation functions, conservation laws,
topological indicators, dynamical robustness, and statistical independence of
density fields. We find that the U-Net approach does well with some metrics
but has difficulties with others. In addition to validating AI approaches
using rigorous physical benchmarks, this study motivates advancements in
domain-specific optimization schemes for scientific machine learning.

##  1 Introduction

In the current era of ‘precision cosmology’, cosmological N-body simulations
play a key role in the study of the large-scale structure of the universe, and
are essential to interpreting the unprecedented amount of observational data
available from sky surveys [1]. In exploring cosmological parameter space, it
is computationally challenging to evolve billions and trillions of particles
at high enough resolution, even for a single simulation. Since large numbers
of simulations are needed for tasks such as covariance estimation [2], the
computing requirements can easily become prohibitive.

Deep learning (DL) has exhibited great potential in learning highly complex
functions and mappings, thereby attracting attention in various scientific
domains, including cosmology [3]. In particular, if DL methods can
successfully capture the complexity of cosmic evolution, they can directly
address the computational bottleneck, and even eliminate it in some cases.
However, issues such as interpretability, as well as the ability of DL models
to make sufficiently accurate and physically meaningful predictions, are yet
to be fully explored. For example, can DL models, without any domain knowledge
of cosmology, capture all the intricacies of nonlinear gravitational
clustering? And in what way can AI-based emulators complement standard
numerical approaches?

In this ongoing work, we model cosmic evolution using the widely adopted U-Net
approach [4], applying it to a theoretically well-understood prescription, the
Zel’dolvich approximation (ZA) [5]. Since the ZA is a simple dynamical mapping
scheme, it has the benefit of providing a clear physical picture while
enabling a comprehensive study of the neural network’s physical
interpretability.

##  2 Approximation for structure formation of the Universe

#### Dynamical approximation :

The Zel’dolvich approximation is the first-order solution to the Lagrangian
perturbation theory of structure formation. It describes the motion of mass
elements: an initially uniform distribution given by Lagrangian coordinates
$\vec{q}$ is displaced by

$\vec{x}(t)=\vec{q}+b(t)\vec{S}(\vec{q}),$ (1)

where $\vec{x}(r)$ are the comoving coordinates, $b(t)$ is the linear growth
rate of fluctuations and $\vec{S}(\vec{q})=\vec{\nabla}\Phi(\vec{q})$ is the
gradient of the initial gravitational potential $\Phi(\vec{q})$. The potential
$\Phi(\vec{q})$ is determined by the primordial density fluctuations
$\delta(x)$ via the Poisson equation, $\nabla^{2}\Phi=4\pi G\delta$, where $G$
is the Newtonian constant of gravitation. Initial density perturbations are
realizations of a Gaussian random field, which is fully specified by a (given)
power spectrum. Uniformly placed particles representing mass elements are then
moved according to the ZA.

#### Generation of the Dataset:

For an initial power spectrum $P(k)$, the displacement field
$\textit{S}(\textit{q})$ is generated using FFT-based techniques, and the
particles are transported from a regular grid via ZA displacements. We
generate 2000 pairs of ZA-evolved particles at two different time steps of the
evolution, with respective scale factors $a_{0}=0.0464$ (an “early” snapshot
with redshift $z_{0}=20.5$) and $a_{1}=0.215$ (a “late” snapshot with redshift
$z_{1}=3.6$). The evolution of the so-called “cosmic web” in between these two
moments is given by the ZA (from Equation 1), solely determined by the initial
potential gradient. The displacement fields at these two times form the basis
of the training/validation datasets. Data is split such that 1000 pairs are
used for training, 500 for validation and 500 for testing. Each dataset is
generated with a different random seed to ensure their statistical
independence. The dataset consists of $64^{3}$ N-body particles in a volume of
64 $h^{-1}$Mpc and their displacements in three spatial dimensions.

###  2.1 Snapshot-to-snapshot translation using convolutional neural networks

Figure 1: Comparison of projected densities between the predicted cosmic web
and ground truth. Densities are derived from the displacement field with a
Cloud-In-Cell (CIC) method, and summed over one axis. The ‘input’ panel (first
panel) presents the density field projection of the early ($z=z_{0}$)
snapshot. The colorbar in each panel shows the magnitude of the dark matter
density $\rho(\textbf{x}$).

A number of DL-based unsupervised generative models and supervised
interpolation models have been applied in cosmological data creation [6, 7,
8]. Except for a few applications where the loss functions are tailored to the
physical problem, most of the applications are domain-agnostic. That is, all
the information about the underlying physics is entirely learned from the
training data. These data-driven models have shown good accuracy with
validation datasets, albeit with respect to metrics that resemble the loss
function.

#### Architecture:

We choose an image-to-image translation framework [9] originally used for
biomedical images and recently applied in a cosmological setting [10, 4, 11,
12, 13]. We adopt the 3D U-Net architecture used in [4] with 15 convolution
and deconvolution layers. The architecture has a contracting path and then an
expanding path. A basic block consists of a convolution layer, batch
normalization layer and a Rectified Linear Unit (ReLU) layer. Mean squared
error between the Lagrangian co-ordinates is used as the loss function, with
L2 regularization. We refer the reader to Refs. [9, 4] for further details of
the network architecture.

#### Model Training:

During the training of the 3D U-Net, a pair of initial-final displacement
fields is selected randomly from the dataset as the input, and the neural
network is trained to learn the mapping between the $z_{0}$ and $z_{1}$
displacement fields. After the training is completed, the model is run on a
set of test data, independent from training or validation, in order to compare
model predictions with the ground truth (the ZA result).

##  3 Validation Metrics

Figure 2: Different physical metrics for assessment of the NN predictions for
the ZA-evolved cosmic web. Top left: Density PDF curves for the NN prediction
($z_{1}$), true ($z_{1}$) and input ($z_{0}$) density fields. Top right and
bottom left: Matter power spectrum and bispectrum comparison between the
predicted density field and ground truth. Bottom right: Percolation transition
curve for the truth, prediction and input fields. Shaded area indicates
standard deviations over ensemble of 30 realizations.

Once the loss has converged (after 10 training epochs), we first visualize the
reconstructed cosmic web. Figure 1 shows the U-Net generated density
(projected along the $x$-axis) of the cosmic density field, demonstrating a
reasonably good agreement with the ZA result. We next implement a number of
more quantitative tests. In particular, we look at measures that quantify the
matter distribution in the universe and topological connectivity of large
scale structures. Details of these metrics are given below and summarized in
Fig. 2.

#### Pixel-wise Comparison:

The most straightforward comparison is the relative error in predicted
densities or displacements. For a density field $\rho(\textbf{x})$, we define
a relative error
$|\rho(\textbf{x})_{pred}-\rho(\textbf{x})_{truth}|/\rho(\textbf{x})_{truth}$,
where $\rho(\textbf{x})_{pred}$ is the U-Net prediction, and
$\rho(\textbf{x})_{truth}$ is the ZA result. The last panel of Fig. 1 shows
the relative difference for one of the realizations. In addition, the one-
point probability distribution function (PDF) is used to assess the prediction
field, which is also an important measure for nonlinear clustering in cosmic
structure. The density PDF measurements of prediction and ground truth are
shown in the first panel of Fig. 2. While relatively unbiased, we note that
the density PDF error rises with density, which is concerning if the
simulation is expected to capture correct clustering physics.

#### Matter Power Spectrum:

The power spectrum is the Fourier transform of the 2-point correlation
function in real space, and is widely used as a statistical measure of the
density field. The matter overdensity
$\delta(\textbf{x})=({\rho(\textbf{x})-\overline{\rho}})/{\overline{\rho}}$,
where $\overline{\rho}$ is the mean density. Writing its Fourier transform
dual as $\delta(\textbf{k})$, the power spectrum $P(k)$ is defined by
$\langle\delta(\textbf{k})\delta(\textbf{k}^{\prime})\rangle=(2\pi)^{3}P(k)\delta^{3}(\textbf{k}-\textbf{k}^{\prime})$.
Measurements of the predicted power spectra and ground truth $P(k)$ are shown
in the second panel of Fig. 2, where we observe differences of less than 5%,
although there is a $k$-dependent bias in the error behavior.

#### Three-Point Correlation – Bispectrum:

Apart from two-point correlation functions, higher-order statistics can also
be used to characterise the density field, and are especially useful for
studying the late stages of structure formation where the evolution is highly
nonlinear and non-Gaussian. Adapting similar symbol conventions as above, the
bispectrum $B(k_{1},k_{2},k_{3})$ [14] is defined as:
$\langle\delta(\textbf{k}_{1})\delta(\textbf{k}_{2})\delta(\textbf{k}_{3})\rangle=(2\pi)^{3}B(k_{1},k_{2},k_{3})\delta^{3}(\textbf{k}_{1}+\textbf{k}_{2}+\textbf{k}_{3})$.
Equilateral triangles in $k$ space are selected for our bispectra analysis
shown in the third panel of Fig. 2. Deviation of the bispectra of the U-Net
generated fields with respect to ZA bispectra is considerably higher compared
to the lower order metrics of the one-point PDF and matter power spectra.

#### Topological metrics – Percolation Analysis:

Topological statistics can be used to directly address the issue of
connectivity of spatial structures, e.g., clusters of galaxies [15]. For a
given density threshold, the volume fraction of the excursion set $f_{ES}$
(i.e., the regions over the threshold) and the volume fraction of the largest
region $f_{1}$ are computed. The behavior of $f_{1}/f_{ES}$ as a function of
$f_{ES}$ (last panel of Fig. 2) displays a rapid transition from zero to 1;
this signifies a topological shift from isolated clusters to a fully connected
field. For the input density field ($z=z_{0})$, the percolation transition
happens at a different density threshold than $z=z_{1}$ case, indicating a
change in the topological behaviour. For this metric the NN results are very
good and in excellent agreement with those from the ZA.

(a)

(b)

Figure 3: Left panel: illustration of the cross-power test, showing the
density fields and cross-correlation power spectra computed in between them.
Right panel: Mean and standard deviation of the three cross-power spectra
(normalized by the auto-power spectra).

#### Cross-power Test:

The accurate cosmic evolution follows fixed dynamical rules, independent of
the initial conditions. However, the neural network emulating the evolution
might show discernible bias inherited from the training dataset. Specifically,
we need to investigate whether the NN prediction induces otherwise non-
existent correlations among the outputs, in contrast to the ZA-evolution of
fields, where independent initial conditions should result in independent
evolved fields. An example test of a cross-power spectral measurement is shown
in Fig. 3(b), where we generate two initial conditions, independent of the
training set, measure their cross-power spectra, then evolve them separately
by ZA and NN, and then measure the output cross-power between them again (Fig.
3(a) illustrates this validation route). From Fig. 3(b), the cross-power shows
that, at least for 2-point statistics, there are no hidden correlations among
the NN-predicted density fields.

##  4 Discussion

While DL-based translation networks and generative models have demonstrated
significant potential, critical tests have to be made before they are applied
in high precision scientific fields like cosmology (where we are interested in
error metrics aiming at accuracies of better than 1% [16]). In this work, we
have demonstrated a subset of the validation schemes to motivate the
scientific machine learning community towards using comprehensive sets of
benchmarking tools. Starting from a tractable approximation to the complex
dynamics of cosmic structure formation, we demonstrate the use of metrics that
were not part of the optimization of the neural network. For instance,
topological transitions and bispectra are higher-order emergent features,
agreement with which is not trivial in NN-based structure formation models.
These checks not only serve as robust validation tools, but also provide a
clear direction for designing future cosmic structure emulation schemes.

We note that although the DL results appear to be promising and close to the
true physical results, in order for neural networks to compete with
simulations, eventually all errors must be less than 1% for clustering
metrics, and the spatial dynamic range of the 3D NN results must be
significantly extended, from a part in a hundred to parts per million. An
interesting and long road lies ahead.

##  5 Acknowledgements

The authors would like to thank Ben Gutierrez for useful discussions. Work at
Argonne National Laboratory was supported by the U.S. Department of Energy,
Office of High Energy Physics. Argonne, a U.S. Department of Energy Office of
Science Laboratory, is operated by UChicago Argonne LLC under contract no. DE-
AC02- 06CH11357. This material is also based upon work supported by the U.S.
Department of Energy, Office of Science, Office of Advanced Scientific
Computing Research and Office of High Energy Physics, Scientific Discovery
through Advanced Computing (SciDAC) program. The training is carried out on
Swing, a GPU system at the Laboratory Computing Resource Center (LCRC) of
Argonne National Laboratory.

## References

  * [1] Salman Habib, Adrian Pope, Hal Finkel, Nicholas Frontiere, Katrin Heitmann, David Daniel, Patricia Fasel, Vitali Morozov, George Zagaris, Tom Peterka, and et al.  Hacc: Simulating sky surveys on state-of-the-art supercomputing architectures.  New Astronomy, 42:49–65, Jan 2016.
  * [2] The LSST Dark Energy Science Collaboration.  The LSST Dark Energy Science Collaboration (DESC) Science Requirements Document.  arXiv e-prints, page arXiv:1809.01669, September 2018.
  * [3] Michelle Ntampaka, Camille Avestruz, Steven Boada, Joao Caldeira, Jessi Cisewski-Kehe, Rosanne Di Stefano, Cora Dvorkin, August E. Evrard, Arya Farahi, Doug Finkbeiner, Shy Genel, Alyssa Goodman, Andy Goulding, Shirley Ho, Arthur Kosowsky, Paul La Plante, Francois Lanusse, Michelle Lochner, Rachel Mandelbaum, Daisuke Nagai, Jeffrey A. Newman, Brian Nord, J. E. G. Peek, Austin Peel, Barnabas Poczos, Markus Michael Rau, Aneta Siemiginowska, Dougal J. Sutherland, Hy Trac, and Benjamin Wandelt.  The Role of Machine Learning in the Next Decade of Cosmology.  BAAS, 51(3):14, May 2019.
  * [4] Siyu He, Yin Li, Yu Feng, Shirley Ho, Siamak Ravanbakhsh, Wei Chen, and Barnabás Póczos.  Learning to predict the cosmological structure formation.  Proceedings of the National Academy of Sciences, 116(28):13825–13832, 2019.
  * [5] Y. B. Zel’Dovich.  Gravitational instability: an approximate theory for large density perturbations.  A&A, 500:13–18, March 1970.
  * [6] Siamak Ravanbakhsh, Francois Lanusse, Rachel Mandelbaum, Jeff Schneider, and Barnabas Poczos.  Enabling Dark Energy Science with Deep Generative Models of Galaxy Images.  arXiv e-prints, page arXiv:1609.05796, Sep 2016.
  * [7] Warren R Morningstar, Yashar D Hezaveh, Laurence Perreault Levasseur, Roger D Blandford, Philip J Marshall, Patrick Putzky, and Risa H Wechsler.  Analyzing interferometric observations of strong gravitational lenses with recurrent and convolutional neural networks.  arXiv preprint arXiv:1808.00011, 2018.
  * [8] Mustafa Mustafa, Deborah Bard, Wahid Bhimji, Zarija Lukić, Rami Al-Rfou, and Jan M. Kratochvil.  Cosmogan: creating high-fidelity weak lensing convergence maps using generative adversarial networks.  Computational Astrophysics and Cosmology, 6(1), May 2019.
  * [9] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-Net: Convolutional Networks for Biomedical Image Segmentation.  arXiv e-prints, page arXiv:1505.04597, May 2015.
  * [10] Elena Giusarma, Mauricio Reyes Hurtado, Francisco Villaescusa-Navarro, Siyu He, Shirley Ho, and ChangHoon Hahn.  Learning neutrino effects in Cosmology with Convolutional Neural Networks.  arXiv e-prints, page arXiv:1910.04255, October 2019.
  * [11] M. A. Aragon-Calvo.  Classifying the large-scale structure of the universe with deep neural networks.  MNRAS, 484(4):5771–5784, April 2019.
  * [12] Renan Alves de Oliveira, Yin Li, Francisco Villaescusa-Navarro, Shirley Ho, and David N. Spergel.  Fast and Accurate Non-Linear Predictions of Universes with Deep Learning.  arXiv e-prints, page arXiv:2012.00240, November 2020.
  * [13] Ziyong Wu, Zhenyu Zhang, Shuyang Pan, Haitao Miao, Xiaolin Luo, Xin Wang, Cristiano G. Sabiu, Jaime Forero-Romero, Yang Wang, and Xiao-Dong Li.  Cosmic Velocity Field Reconstruction Using AI.  ApJ, 913(1):2, May 2021.
  * [14] Johnathan Hung, James R. Fergusson, and E. P. S. Shellard.  Advancing the matter bispectrum estimation of large-scale structure: a comparison of dark matter codes, 2019.
  * [15] Sergei Shandarin, Salman Habib, and Katrin Heitmann.  Origin of the cosmic network in $\Lambda\mathrm{CDM}$: Nature vs nurture.  Phys. Rev. D, 81:103006, May 2010.
  * [16] Ofer Lahav and Andrew R Liddle.  The Cosmological Parameters (2019).  arXiv e-prints, page arXiv:1912.03687, December 2019.
  * [17] Andrea Petri, Zoltán Haiman, Lam Hui, Morgan May, and Jan M. Kratochvil.  Cosmology with minkowski functionals and moments of the weak lensing convergence field.  Phys. Rev. D, 88:123002, Dec 2013.

##  Appendix A Additional metrics

The metrics studied in this analysis are by no means exhaustive. We list a few
other metrics that will be incorporated into our validation scheme in the near
future.

#### Minkowski Functionals:

Minkowski Functionals (MFs) are topological descriptors of random fields that
are useful in describing their statistical properties [17]. Given a random
field $\rho(\vec{x})$, and its variance $<\rho^{2}>=\sigma_{0}^{2}$, we can
consider its excursion sets $\Sigma(\nu)={x>\nu\sigma_{0}}$ which consists of
all points that exceed the threshold $\nu\sigma_{0}$. Three MFs measure the
area, length of boundary and genus that is characteristic of excursion sets:

$V_{0}(\nu)=\frac{1}{A}\int_{\Sigma(\nu)}da$
$V_{1}(\nu)=\frac{1}{4A}\int_{\Sigma(\nu)}dl$ $V_{2}(\nu)=\frac{1}{2\pi
A}\int_{\Sigma(\nu)}\mathcal{K}dl$

where A is the area of the map, $da$ and $dl$ the area and boundary length
elements, and $\mathcal{K}$ is the curvature of the boundary.

#### Dynamical Robustness:

Apart from the density field, the matter velocity field is of equal
significance regarding the correct description of the evolution of the cosmic
web, and velocity-density cross correlations are crucial for this measurement.
For ZA evolution, explicit reversibility can be also checked to guarantee the
dynamical robustness of reconstruction.