Explaining dark matter halo density profiles with neural networks
We use explainable neural networks to connect the evolutionary history of
dark matter halos with their density profiles. The network captures independent
factors of variation in the density profiles within a low-dimensional
representation, which we physically interpret using mutual information. Without
any prior knowledge of the halos' evolution, the network recovers the known
relation between the early time assembly and the inner profile, and discovers
that the profile beyond the virial radius is described by a single parameter
capturing the most recent mass accretion rate. The results illustrate the
potential for machine-assisted scientific discovery in complicated
astrophysical datasets.
Luisa Lucie-Smith, Hiranya V. Peiris, Andrew Pontzen
# Explaining dark matter halo density profiles with neural networks

###### Abstract

We use explainable neural networks to connect the evolutionary history of dark matter halos with their density profiles. The network captures independent factors of variation in the density profiles within a low-dimensional representation, which we physically interpret using mutual information. Without any prior knowledge of the halos' evolution, the network recovers the known relation between the early time assembly and the inner profile, and discovers that the profile beyond the virial radius is described by a single parameter capturing the most recent mass accretion rate. The results illustrate the potential for machine-assisted scientific discovery in complicated astrophysical datasets.

**Introduction** - Machine-assisted scientific discovery has recently become a key frontier in artificial intelligence (AI) research [1]. It is useful to build in characteristics of good physical models into such frameworks for artificial scientific discovery. Such desiderata include the compression of information within a dataset into a set of minimal ingredients which: can accurately predict new data; can be usefully contextualized within the broader scientific domain; and generalize beyond the specific setting in which the model was originally fitted to explain other aspects of the physical system. The required compression can be achieved through'representation learning'; recent works have shown that neural networks can learn low-dimensional representations which match the already-known parameters describing simple physical systems [2]. However, extracting new knowledge from deep learning models requires the development of tools that can explain latent representations in terms of the physics they represent, without knowing the relevant physical parameters _a priori_.

In this work, we demonstrate the use of an _explainable AI_ framework to address an open problem in cosmological structure formation. In the modern picture of structure formation, galaxies form at the center of extended, overdense 'halos' of dark matter, which originate from small fluctuations in the density of matter in the early Universe and undergo highly non-linear dynamical processes throughout their evolution [3; 4; 5; 6; 7]. The dynamical history of a halo then determines the way in which matter is distributed within a halo today. The matter distribution within a halo (known as its 'density profile') strongly affects cosmological analyses that connect the observed galaxy distribution with theoretical predictions of the underlying matter distribution [8], as well as direct and indirect searches for dark matter [9]. Current models of halo density profiles rely on empirical fitting functions calibrated to simulations, such as the Navarro-Frenk-White (NFW) profile [10] and the Einasto profile [11]. Halo density profiles are known to show universal shapes across 20 orders of magnitude in mass [12; 13]. However, no general consensus on an explanation from first principles for the origin of halo density profiles exists.

In Ref. [14], we used an interpretable deep learning framework to create a model of the independent degrees of freedom in the spherically-averaged density profiles of dark matter halos. The model, which we denoted an _interpretable variational encoder_ (IVE), generated a compact, low-dimensional latent representation of the 3D density field in a region containing the halo, which contains all the information used by the neural network to predict the profile. The representation is interpretable because it is disentangled, i.e. we require that each latent component captures different, independent factors of variation in the profiles. We found that three components are required (and are sufficient) for modeling the profiles out to the halo outskirts: one component describes the overall normalization of the profile, the second the shape of the profile within the virial radius, and the third the shape of the profile beyond the virial radius that is affected by infalling material.

In this _Letter_, we turn to the physical interpretation of the learnt IVE latent representation in terms of the evolution history of halos that led to the final halo density profiles. Although the network was trained only with information about the present-day density field, we explore whether the latent parameters carry memory of the evolution history of the halos. We measure the information encoded within each latent about the halos' evolution history using the information-theoretic measure of _mutual information_ (MI). By this metric, the IVE representation and the NFW parametrization similarly highlight a dependence of the profile on physical accretion history. However the IVE additionally allows us to measure the connection between a halo's recent evolution history and the density in its far outskirts, something that the NFW profile does not capture.

**Background** - We begin by briefly reviewing the current understanding of the physics of halo density profiles. The NFW profile is the most widely used fitting function for the halo density profile. It is a two-parameter functional form given by

\[\rho(r)=\frac{\rho_{s}}{r/r_{s}\left(1+r/r_{s}\right)^{2}}\,, \tag{1}\]

where \(r_{s}\) and \(\rho_{s}\) are the scale radius and characteristic density, respectively. The scale radius is often re-written in terms of a concentration parameter \(c\equiv r_{200m}/r_{s}\), so that the NFW profile can be parametrized by a virial radius \(r_{200m}\) (or virial mass) and concentration \(c\). The virial radius \(r_{200m}\) is typically
adopted as a proxy for the halo boundary, and defined as the radius which contains a mean density that is 200 times the mean density of the Universe. High-resolution simulations have revealed this functional form to be 'universal': it provides a good fit to stacked profiles of halos for a large range of halo masses [13; 15], for several different cosmological models [16; 17; 18; 19; 20], and even in the absence of hierarchical growth [21; 22; 23]. This suggests that universal density profiles are a generic feature that arises from collisionless gravitational collapse.

Despite the lack of a first-principles explanation for the self-similarity of halo density profiles, some insights have been gained from studying the correlation between the NFW concentration and summary statistics of the halo evolution process. Mass, concentration and halo formation time all correlate: on average, low-mass halos assemble earlier and have higher characteristic densities (or concentration), reflecting the larger background density at earlier times [24; 25; 26; 19; 10; 27]. This description can explain the qualitative trend of the mean concentration as a function of halo mass, but not the large residual scatter in concentration seen in simulations [25; 28]. It is also limited to the simplest summary statistic of the halo evolution history, i.e. the halo formation time. Further work has suggested that the self-similarity of halos may be related to the self-similarity of the halo mass assembly history [15], although this has only been validated on stacked profiles of well-behaved,'relaxed' halos.

The situation worsens when modeling profiles beyond the virial radius: the halo outskirts strongly deviate from the NFW form due to the presence of the splashback radius, where particles reach the apocenter of their first orbit. Recent work has focused on modeling the location of the splashback radius, finding that it is sensitive to the late-time mass accretion rate [29; 30; 31; 32; 33; 34]. Modeling the full shape of the outer profile remains a difficult task due to its intrinsically non-equilibrium nature, leading to a reliance on multi-parameter fitting functions with little physical explainability [35; 36].

**Deep learning model** - The IVE architecture used in this work has two main components: the encoder, mapping the 3D density field to a low-dimensional latent representation, and the decoder, mapping the latent representation and the query radius \(\log(r)\) to the output profile \(\log[\rho(r)]\). By design, all the information used by the model to predict the density profiles is captured within the latent representation. An illustration of the model is shown in the top-half of Fig. 1. The encoder is a 3D convolutional neural network (CNN) with parameters \(\phi\) that maps the inputs \(\mathbf{x}\) to a multivariate distribution in the latent space \(p_{\phi}(\mathbf{z}|\mathbf{x})\). We choose the latent representation to be a set of independent Gaussians, \(p_{\phi}(\mathbf{z}|\mathbf{x})=\prod_{i=1}^{L}\mathcal{N}(\mu_{i}(\mathbf{x}),\sigma_{i} (\mathbf{x}))\), where \(L\) is the dimensionality of the latent space; under this assumption, the encoder maps the inputs \(\mathbf{x}\) to the vectors \(\mu=\mu_{i},..,\mu_{L}\) and \(\sigma=\sigma_{i},..,\sigma_{L}\). The decoder of the IVE consists of another neural network model with parameters \(\theta\) that maps a sampled latent vector \(z\sim p_{\phi}(\mathbf{z}|\mathbf{x})\) and a value of the query \(\log(r)\) to a single predicted estimate for \(\log[\rho_{\rm pred}(r)]\).

A crucial aspect of the IVE that makes the latent space interpretable is that it is _disentangled_: independent factors of variation in the density profiles are captured by different, independent latents. This is achieved through the design of a loss function that minimizes the mean squared error between predicted and ground-truth profiles, while simultaneously maximizing the degree of independence between the latent variables by encouraging those to be as close as possible to independent Gaussians of mean 0 and variance 1 [37]. More details on the encoder and decoder architectures and the loss function are presented in Ref. [14].

**Methods** - We trained two IVE models for different tasks: one (\(\text{IVE}_{\rm virial}\)) was trained to model the density profile up to the halo virial radius \(r_{200\rm m}\), and the second (\(\text{IVE}_{\rm infall}\)) was trained to model profiles beyond the halo boundary out to \(2\,r_{200\rm m}\). The former is used for direct comparison with the NFW profile, which is also designed to model the profile out to the virial radius, and the latter is used to investigate the less studied halo outer profile. The innermost radius of the profiles we consider is \(r_{\rm min}=3\,\epsilon\), where \(\epsilon\) is the gravitational softening of the simulation; this choice ensures that we can robustly trust the inner profile. The inputs are given by the 3D density field within a \(N=131^{3}\) sub-box of size \(L_{\rm sub-box}=0.4\,\text{Mpc}\,h^{-1}\) for the \(\text{IVE}_{\rm virial}\) model, and of size \(L_{\rm sub-box}=0.6\,\text{Mpc}\,h^{-1}\) for the \(\text{IVE}_{\rm infall}\) one. For the latter, we further restricted our analysis to halos with \(r_{200\rm m}{\leq 150\,\text{kpc}\,h^{-1}}\). Further discussion on the training data of the \(\text{IVE}_{\rm virial}\) and \(\text{IVE}_{\rm infall}\) models is presented in Ref. [14]. To compare the \(\text{IVE}_{\rm virial}\) results with NFW, we fitted the NFW formula in Eq. (1) to each halo's density profile using least-squares minimization, and recovered the best-fitting parameters \(r_{s}\) and \(\rho_{s}\). The concentration was then derived using \(c=\)\(r_{200\rm m}/r_{s}\). A description of the simulations used for training and testing the IVE models can be found in the Supplemental Material.

The first step of the analysis was to verify that the IVE models learn to predict the density profiles at \(\sim\!5\%\) accuracy, comparable to the accuracy of NFW fits (see Supplemental

Figure 1: A neural network is trained to discover the underlying degrees of freedom in halo density profiles in the form of a latent representation, when presented with the full 3D density structure of a halo. We physically interpret the discovered representation by measuring the MI between the latent parameters and the assembly history of the halos.
Material). Crucially, only the \(z=0\) snapshots were used for training the IVE models to construct the latent representations mapping the 3D density field to dark matter halo profiles - i.e., the model had no access to the merger histories of the halos during training. The resulting disentangled latent space directly corresponds to the underlying degrees of freedom in the halo density profiles. Following recent works [38, 39, 14], we then used the MI to (i) quantify the information captured by the latent space about the halo density profiles and (ii) connect the IVE latents to the halo's evolution history, showing how the latter determines the present-day density profile.

The MI was estimated using GMM-MI [39], which performs density estimation using Gaussian mixtures and provides MI uncertainties through bootstrap. We first measured the MI between each latent and the density profile \(\rho(r)\); this allows us to directly link each latent to a degree of freedom in the profile that affects its shape over a certain radial range. We then measured the MI between each latent and the mass assembly history of each halo. This in turn allowed us to connect each degree of freedom describing the density profile of the halo directly to characteristics of the halos' evolution which determines that component.

**Results** - Figure 2 quantifies the information contained within the latents of the \(\text{IVE}_{\text{infall}}\) (top panel) and the \(\text{IVE}_{\text{virial}}\) (bottom panel) models about the ground-truth density profiles1. We show the MI between each latent parameter and the ground-truth profiles, which we denote as \(\text{MI}_{\rho_{\text{true}}(r)}\). The three latents discovered by the \(\text{IVE}_{\text{infall}}\) describe (i) the normalization of the profile, which dominates the variation in the profiles out \(\sim r_{200\text{m}}/2\), (ii) the shape of the inner profile, which becomes informative on radial scales approaching \(r_{200\text{m}}\), and (iii) the shape of the outer profile beyond \(r_{200\text{m}}\). The first two are analogous to the two NFW parameters, mass and concentration, respectively. A closer comparison between the inner shape latent of the \(\text{IVE}_{\text{virial}}\) model and concentration (bottom panel of Fig. 2) shows that both parameters carry information about the density in the core and on radial scales close to \(r_{200\text{m}}\). This bimodality is due to a compensation effect between the density in the inner region and that close to the virial boundary: at fixed normalization, halos with denser cores become less dense in the outskirts, and vice versa. The \(\text{MI}_{\rho_{\text{true}}(r)}\) of the inner latent is shifted towards larger radii compared to that of concentration, suggesting that the former is sensitive to variations in the shape of the profile on larger radial scales than the latter; this distinction will become relevant when physically interpreting the latent and comparing it to concentration.

Footnote 1: We verify the conclusions of our previous work in Ref. [14] at higher precision using the new GMM-MI estimator [39].

We now move on to a physical interpretation of the latents in relation to characteristics of the halos' evolution histories. Recall that the network did not have access to this information during training. The interpretation of the normalization latent is straightforward: it captures the \(z=0\) mass of the halo, \(M_{200\text{m}}\). Their MI is \(\sim 2.07\pm 0.01\) nats, implying a strong correlation between the two. This also matches expectations from the literature [10, 11], as halo mass also controls the normalization in the NFW and Einasto fitting functions. To physically interpret the inner and outer shape latents, we measure

Figure 2: The MI between the latent parameters and the ground-truth halo profiles \(\rho_{\text{true}}(r)\) for the \(\text{IVE}_{\text{infall}}\) (_top_) and the \(\text{IVE}_{\text{virial}}\) (_bottom_) models. In the \(\text{IVE}_{\text{infall}}\) case, we also show MI with the NFW concentration. (For clarity we do not show the normalization latent for \(\text{IVE}_{\text{infall}}\), since it behaves identically to the \(\text{IVE}_{\text{virial}}\) normalization latent.)

Figure 3: The MI between the latent parameters and the mass accretion histories (denoted \(\text{MI}_{M(z)}\); top row), and that between the latent parameters and the mass accretion rate (denoted \(\text{MI}_{dM(z)/dz}\); bottom row). The inner shape latent and the NFW concentration carry memory of the early-time mass assembly history, as well as the later-time mass accretion rate. The outer shape latent carries information about the halos’ most recent mass accretion rate over the past dynamical time (indicated by the arrow).
their MI with two quantities that describe the assembly history of the halos over cosmic time. The first is the mass accretion history, \(M_{\rm 200m}(z)/M_{\rm 200m}(z=0)\), which describes the evolution of the halo mass as a function of time \(M_{\rm 200m}(z)\) normalized to the present-day halo mass \(M_{\rm 200m}(z=0)\). The second is the mass accretion rate \(\Gamma(t)\equiv\Delta\ln M_{\rm 200m}(a)/\Delta\ln a\)[33], which describes the rate of change in halo mass with respect to the scale factor \(a(t)\). The value of the accretion rate depends on the time interval used to compute the change in mass and scale factor; we compute \(\Gamma(t)\) by taking the finite difference of the halo masses at each consecutive timestep in the simulation.

Figure 3 shows the MI between the latents and mass accretion history (MI\({}_{M(z)}\); top row) and that between the latents and the mass accretion rate (MI\({}_{\rm dM(z)/dz}\); bottom row). We first focus on the inner shape latent, which we compare to the NFW concentration. The MI\({}_{M(z)}\) of the inner shape latent increases with time during the early formation period, peaks at \(z\sim 1\), and declines rapidly towards \(z=0\); recall that this is the MI with the mass assembly history normalized to the present-day halo mass. This result reveals that the inner shape latent is sensitive to the early assembly history of halos. The MI\({}_{dM(z)/dz}\) of the same latent reveals that the latter is also sensitive to the later time mass accretion rate. This dual dependence explains the bimodal shape of the MI between the inner latent and the profile (Fig. 2, bottom panel): the early assembly phase determines the shape of the profile in the innermost region of the halo, while the later time mass accretion rate determines the shape of the profile close to the virial radius. We further validate this interpretation in the Supplemental Material.

The NFW concentration shows a similar picture to the inner shape latent. However, its MI\({}_{M(z)}\) peaks at earlier times (\(z\sim 0.55\)) compared to the inner shape latent. This implies that the inner shape latent carries information about the build up of mass onto the halo over a longer period of time than concentration, which therefore affects the inner halo structure (and the profile) out to larger scales. The sensitivity of the inner latent to later times/larger scales in the profile explains why the inner latent MI\({}_{\rho_{\rm trw}(r)}\) is shifted towards larger radial scales than that of concentration (Fig. 2). Moreover, we find that the absolute magnitude of the concentration MI\({}_{M(z)}\) is higher than that of the inner shape latent; this is because the closer to the halo core the stronger the correlation with the early assembly history due to halos accreting mass 'inside-out'. As a result, concentration, which is sensitive to the profile on smaller \(r\) than the latent, has a higher MI with the early assembly history than the latent. Finally, the NFW concentration is related to the later time mass accretion rate in a similar way to the inner latent.

Figure 3 also quantifies the information contained within the outer shape latent about the mass evolution history of halos. We find that the outer shape latent is primarily determined by the late-time mass accretion rate. In particular, the outer shape latent is sensitive to the accretion rate over the past \(\sim 5\) Gyr. This timescale corresponds to the halo dynamical time, \(t_{\rm dyn}\equiv 2\times r_{\rm 200m}/v_{\rm 200m}\), defined as the time it takes for material to cross the halo at a typical virial velocity \(v_{\rm 200m}=\sqrt{{\rm G}M_{\rm 200m}/r_{\rm 200m}}\). This result suggests that the outer profile is detemined by the accretion of matter infalling onto the halo, that is dynamically out of equilibrium and has not yet virialized within the halo.

**Discussion** - Our results show that the IVE framework has extracted a direct connection between the assembly history of cold dark matter halos and their density profiles, without having access to explicit information about the time evolution of the halos during training. This has deep implications for understanding the origin of universality in dark matter halos. The results suggest that the universality in the profiles, captured by three degrees of freedom alone, is a direct consequence of a universality in the halo assembly histories themselves, since the latents contain comparable amounts of information about both quantities.

Previous work [15] found a resemblance between the shape of the average mass accretion history, expressed in terms of the critical density of the Universe, and the average enclosed mass profile, expressed in terms of its enclosed density, for a selected set of 'well-behaved' halos of similar mass. In the halo outskirts, the profile has been linked to the dynamical accretion history of the halos primarily through the relation between the splashback radius and the mass accretion rate [30; 35]; existing models make use of multi-parameter fitting functions to capture the dynamical impact on the outer profile [36].

By contrast, within the IVE framework, the connection between the density profiles and the _entire_ mass accretion history or mass accretion rate is clearly elucidated through MI. This result was obtained directly using all halos in the simulations, without requiring a curated sample of well-behaved halos. The IVE rediscovers the known correlation between the inner profile and halo formation time [24; 25]; it then additionally demonstrates that the complexity of the dynamical, infalling material is encoded in only a single degree of freedom that captures the recent mass accretion rate. In future work, we will use the connection between assembly history, latents, and density profile captured by the IVE framework to build a model that can determine mass accretion histories from density profiles.

More broadly, our results represent progress towards enabling _new_ machine-assisted scientific discoveries, going beyond artificial rediscovery of known physical laws and concepts [40; 2; 41]. The key requirements for artificial scientific discovery are _interpretability_ and _explainability_; the IVE achieves these by generating a low-dimensional representation that disentangles the independent factors of variation in the output (interpretability) and can be explained in terms of the physics it represents through MI (explainability). The approach shows promise for gaining insight into other emergent properties of the cosmic large-scale structure which currently lack of physical explanations.

**Acknowledgments** - LLS thanks Benedikt Diemer and Simon White for useful discussions. This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement nos. 818085 GMGalaxies and
101018897 CosmicExplorer). This work has been enabled by support from the research project grant 'Understanding the Dynamic Universe' funded by the Knut and Alice Wallenberg Foundation under Dnr KAW 2018.0067. The work of HVP was additionally supported by the Goran Gustafsson Foundation for Research in Natural Sciences and Medicine. AP was additionally supported by the Royal Society. HVP and LLS acknowledge the hospitality of the Aspen Center for Physics, which is supported by National Science Foundation grant PHY-1607611. The participation of HVP and LLS at the Aspen Center for Physics was supported by the Simons Foundation. This work was partially enabled by the UCL Cosmoparticle Initiative.

The contributions from the authors are listed below: **L.L.-S.**: conceptualization; formal analysis; investigation; methodology; software; validation; visualization; writing - original draft, review & editing. **H.V.P.**: conceptualization; methodology; interpretation & validation; writing - review & editing. **A.P.**: methodology; writing - review & editing.

## Supplemental Material

### Simulations

We begin with a description of the simulations used for training the IVE models and subsequent interpretation of the learned representation. We ran four dark-matter-only \(N\)-body cosmological simulations using the publicly-available code GADGET-4 [42] assuming a _Planck_\(\Lambda\)CDM cosmological model [43]. We evolved \(N=512^{3}\) dark matter particles in a \((50\,\mathrm{Mpc}\,h^{-1})^{3}\) box from \(z=99\) to \(z=0\). The four simulations are based on different realizations of the initial Gaussian random field, generated using genetIC[44]. Three simulations were used for training the machine learning model and one was set aside for validation and testing.

Dark matter halos were identified at \(z=0\) using the SUB-FIND halo finder [42; 45], as done in Ref. [14]. We restricted our analysis to halos within the mass range \(\log_{10}(M/M_{\odot})\in[11,13]\), in order to fully resolve the inner profile of the lowest-mass halos and not be affected by small-number statistics at the high-mass end. To track the evolution history of the dark matter halos, we saved \(91\) snapshots of the simulations between \(z\sim 7\) and \(z=0\). We used the pynbody and tangos software packages [46] to construct the merger trees of every dark matter halo. tangos matches a halo with its successor in time based on the fraction of common particles between the two objects; the procedure is repeated for every timestep in the simulation, thus yielding halo merger trees from \(z\sim 7\) to \(z=0\). The merger trees were then used to track the mass of each halo's main progenitor over time.

### Predictions for the halo density profiles

Figure 4 shows the accuracy of the predictions of the IVE\({}_{\mathrm{virial}}\) and IVE\({}_{\mathrm{infall}}\) models. We show the mean and 90% confidence interval of the residuals \(\log_{10}[\rho_{\mathrm{predicted}}/\rho_{\mathrm{true}}]\), in every radial bin of the profile used for testing. Since every radial bin contains a different value of \(r\) for different halos, we plot the residuals as a function of \(r_{\mathrm{eff}}\) defined as the median of the distribution of radius values within each bin. The grey band shows the residuals of the NFW fits for comparison. Note that the IVE results include uncertainties in the latent distributions, whereas the NFW fits do not include uncertainties as they were obtained through least-squares optimization. The performance of the IVE models is consistent with that of the NFW profile, meaning that our model contains sufficient predictive accuracy to yield meaningful latent representations.

### Further physical interpretation of the inner shape latent

We present a further investigation on the physical interpretation of the inner shape latent. As shown in Fig. 2, the MI between the inner shape latent and the ground-truth density profile has a bimodal shape. The MI first peaks at \(r_{1}\sim 0.1\,r_{200\mathrm{m}}\) and then again at \(r_{2}\sim 0.9\,r_{200\mathrm{m}}\), meaning that the latent contains information about the shape of the profile in the inner region of the halo and close to the halo virial radius. When physically interpreting the latent, we found that the latent carries information about the early formation history of the halo and the later time mass accretion rate (Fig. 3). We now verify if the early formation history is responsible for the shape of the profile in the inner region, while the later time mass accretion rate is what determines that closer to the virial radius. To do so, we compute the MI between the ground-truth density at the first MI peak, \(\rho(r_{1})\), and both the halo mass assembly history and mass accretion rate. We then repeat the calculation for the ground-truth density at the second MI peak, \(\rho(r_{2})\). Fig. 5 shows the MI between the two ground-truth densities with the mass assembly history, \(M(z)\), in the top panel2

Figure 4: Mean and 90% confidence interval of the residuals \(\log[\rho_{\mathrm{predicted}}/\rho_{\mathrm{true}}]\) of the IVE\({}_{\mathrm{virial}}\) and IVE\({}_{\mathrm{infall}}\) models, as a function of \(r_{\mathrm{eff}}\) defined as the median radius in each bin. The grey band shows the NFW residuals.
because \(\rho(r)\) also depends on the overall normalization and therefore must be compared to \(M_{200\mathrm{m}}(z)\).
