High-fidelity reproduction of central galaxy joint distributions with
  Neural Networks
The relationship between galaxies and haloes is central to the description of
galaxy formation, and a fundamental step towards extracting precise
cosmological information from galaxy maps. However, this connection involves
several complex processes that are interconnected. Machine Learning methods are
flexible tools that can learn complex correlations between a large number of
features, but are traditionally designed as deterministic estimators. In this
work, we use the IllustrisTNG300-1 simulation and apply neural networks in a
binning classification scheme to predict probability distributions of central
galaxy properties, namely stellar mass, colour, specific star formation rate,
and radius, using as input features the halo mass, concentration, spin, age,
and the overdensity on a scale of 3 $h^{-1}$ Mpc. The model captures the
intrinsic scatter in the relation between halo and galaxy properties, and can
thus be used to quantify the uncertainties related to the stochasticity of the
galaxy properties with respect to the halo properties. In particular, with our
proposed method, one can define and accurately reproduce the properties of the
different galaxy populations in great detail. We demonstrate the power of this
tool by directly comparing traditional single-point estimators and the
predicted joint probability distributions, and also by computing the power
spectrum of a large number of tracers defined on the basis of the predicted
colour-stellar mass diagram. We show that the neural networks reproduce
clustering statistics of the individual galaxy populations with excellent
precision and accuracy.
Natália V. N. Rodrigues, Natalí S. M. de Santi, Antonio D. Montero-Dorta, L. Raul Abramo
# High-fidelity reproduction of central galaxy joint distributions with Neural Networks

###### Abstract

The relationship between galaxies and haloes is central to the description of galaxy formation, and a fundamental step towards extracting precise cosmological information from galaxy maps. However, this connection involves several complex processes that are interconnected. Machine Learning methods are flexible tools that can learn complex correlations between a large number of features, but are traditionally designed as deterministic estimators. In this work, we use the IllustrisTNG300-1 simulation and apply neural networks in a binning classification scheme to predict probability distributions of central galaxy properties, namely stellar mass, colour, specific star formation rate, and radius, using as input features the halo mass, concentration, spin, age, and the overdensity on a scale of 3 \(h^{-1}\) Mpc. The model captures the intrinsic scatter in the relation between halo and galaxy properties, and can thus be used to quantify the uncertainties related to the stochasticity of the galaxy properties with respect to the halo properties. In particular, with our proposed method, one can define and accurately reproduce the properties of the different galaxy populations in great detail. We demonstrate the power of this tool by directly comparing traditional single-point estimators and the predicted joint probability distributions, and also by computing the power spectrum of a large number of tracers defined on the basis of the predicted colour-stellar mass diagram. We show that the neural networks reproduce clustering statistics of the individual galaxy populations with excellent precision and accuracy.

keywords: galaxies: statistics - cosmology: large-scale structure of Universe - methods: data analysis - methods: statistical

## 1 Introduction

Characterising the connection between the properties of galaxies and those of the underlying population of dark-matter (DM) haloes is one of the most crucial aspects to understand the large-scale structure (LSS) of the Universe. This link not only encapsulates fundamental information about the process of galaxy formation, but it is also a crucial step to optimise the extraction of cosmological constraints from galaxy maps.

The halo-galaxy connection is nowadays investigated using a variety of techniques (see, e.g., Wechsler and Tinker, 2018). On the one hand, empirical methods use DM-only simulations as the basis on top of which different analytical prescriptions are implemented in order to establish that connection. These techniques include sub-halo abundance matching (SHAM, e.g., Conroy et al., 2006; Behroozi et al., 2010; Trujillo-Gomez et al., 2011; Favole et al., 2016; Guo et al., 2016; Contreras et al., 2020, 2020; Hadzhiyska et al., 2021; Favole et al., 2022), halo occupation distributions (HODs, e.g., Berlind and Weinberg, 2002; Zehavi et al., 2005, 2018; Artale et al., 2018; Bose et al., 2019; Hadzhiyska et al., 2020; Xu et al., 2021) and empirical forward modelling (e.g., Becker, 2015; Moster et al., 2018; Behroozi et al., 2019). On the other hand, it is possible to model, with varying degrees of detail, the physical mechanisms that shape the process of galaxy formation. In this context, hydrodynamical simulations (e.g., Somerville and Dave, 2015; Naab and Ostriker, 2017; Pillepich et al., 2018, 2018; Springel et al., 2018; Villasescusa-Navarro et al., 2021, 2022) are perhaps the most ambitious efforts. These models employ known physics to simulate, at a sub-grid level, a variety of processes that are related to galaxy formation such as star formation, radiative metal cooling, and supernova, stellar, and black hole feedback - for reviews on this, see Somerville and Dave, 2015; Naab and Ostriker, 2017. This modelling can also be approached from a semi-analytic, less computationally demanding, perspective. These semi-analytic models (SAMs, e.g., White and Frenk, 1991; Guo et al., 2013) employ physically motivated recipes to mimic the galaxy formation processes.

In this paper, we investigate the halo-galaxy connection from a machine learning (ML) perspective. The issue of the halo-galaxy connection has been addressed using ML by many works (e.g., Kamdar et al., 2016; Agarwal et al., 2018; Calderon and Berlind, 2019; Jo Kim, 2019; Man et al., 2019; Yip et al., 2019; Zhang et al., 2019; Jo and Kim, 2019; Kasmanoff et al., 2020; Delgado et al., 2021; McGibbon and Khochfar, 2021; Shao et al., 2021; Lovell et al., 2022; Stiskalek et al., 2022; de Andres et al., 2022; Jespersen et al., 2022; Chittenden and Tojeiro, 2023). In de Santi et al. (2022) we provide a ML suite combining some of the most powerful, well-known models in the literature
to predict central galaxy properties using host halo properties. All the applied methods, however, are designed to return a single value for each galaxy property, independently of the remaining properties. However, there are many complex interrelated processes involved in the formation and evolution of galaxies, and their properties cannot be precisely determined by halo properties alone. Therefore, a model that proposes to map the relation between galaxies and host haloes should encode not only the correlations between galaxy properties, but also the uncertainties due to the stochastic aspects of galaxy formation. In other words, any given halo could host a central galaxy with a variety of properties and, hence, a model should return joint probability distributions for the possible values of those galaxy properties, instead of a single one.

The ML suite from our precursor work (de Santi et al., 2022) provided encouraging results in terms of single-point estimation metrics, such as the Pearson correlation coefficient between true and predicted values, especially for stellar mass, which is highly correlated with halo mass. However, deterministic models that try to predict individual galaxy properties can be biased towards the most frequent values, and thus fail to recover the overall distributions of the galaxy properties. In that paper, this issue is treated as an imbalanced data problem, i.e., despite of the fact that different output values could be associated with some fixed set of halo properties, the machine tends to assign the most frequent values. To address this problem, we made use of a data augmentation technique to increase the weight of the less represented instances, which allowed us to better recover the under-represented populations, but still in a way that each halo is assigned a single, individual value for each central galaxy property (de Santi et al., 2022).

In the present work, we proceed by predicting probability distributions with neural networks (NNs) with a binning classification scheme, which we refer to as NN\({}_{\rm class}\), for the same central galaxy properties as de Santi et al. (2022), namely, stellar mass, \(g-i\) colour, specific star formation rate, and galaxy radius. This not only enables us to recover the overall distributions of the galaxy properties from the IllustrisTNG300-1 (hereafter, TNG300) sample, but also to capture the intrinsic scatter in the halo-galaxy mapping by providing, for each halo, the probability distributions associated with its central galaxy properties. We also train NN\({}_{\rm class}\) to predict the galaxy properties jointly, finding that the joint distributions recover correlations that are lost when predicting univariate distributions independently. ML probability-based descriptions have been used in related contexts, in particular with NNs, such as photometric redshift estimation (e.g., Lima et al., 2022), dynamical mass of galaxy clusters estimation (e.g., Ho et al., 2021; Ramanah et al., 2020) and recently in the halo-galaxy connection (e.g., Stiskalek et al., 2022).

In order to study how NN\({}_{\rm class}\) captures the intrinsic stochasticity in the halo-galaxy connection, we analyse the shape of the distributions of individual galaxies, which gives some insights on the contribution of secondary halo properties. Moreover, we analyse how this uncertainty affects clustering statistics, namely the power spectrum. Our technique enables us to define as many galaxy populations as wished, and to analyse to what extent those populations occupy the same types of haloes. We explore this flexibility by computing the power spectrum of a large number of galaxy populations (tracers), selected on the basis of the colour-stellar mass diagram.

The paper is organised as follows. The IllustrisTNG data and the chosen set of halo and galaxy properties are described in SS2. In SS3, we explain how we applied NNs to predict joint probability distributions. Section 4 analyses the quality of the results obtained with the NNs by comparing the predictions with the IllustrisTNG catalogue. In SS5, we present our results in terms of the power spectra of several galaxy populations. Finally, we outline our main conclusions in SS6, and discuss our plans for future improvements and applications.

## 2 Data

Our analysis is based on data from the IllustrisTNG magnetohydrodynamical cosmological simulation (Pillepich et al., 2018, 2018; Nelson et al., 2018; Marinacci et al., 2018; Naiman et al., 2018; Springel et al., 2018; Nelson et al., 2019). This simulation suite, which was generated using the AREPO moving-mesh code (Springel, 2010), is an improved version of the previous Illustris simulation (Vogelsberger et al., 2014, 2014, 2014). IllustrisTNG features a variety of updated sub-grid models accounting for star formation, radiative metal cooling, chemical enrichment from SNII, SNIa, and AGB stars, as well as feedback mechanisms (including stellar and super-massive black hole feedback). These models were calibrated to reproduce an array of observational constraints, such as the \(z=0\) galaxy stellar mass function and the cosmic SFR density, to name but a few (see the aforementioned references for more information). The IllustrisTNG simulation adopts the standard \(\Lambda\)CDM cosmology (Planck Collaboration et al., 2016), with parameters \(\Omega_{\rm m}=0.3089\), \(\Omega_{\rm b}=0.0486\), \(\Omega_{\Lambda}=0.6911\), \(H_{0}=100\,h\,{\rm km\,s^{-1}Mpc^{-1}}\) with \(h=0.6774\), \(\sigma_{\rm g}=0.8159\), and \(n_{S}=0.9667\).

The ML methodology that we developed in this work to reproduce the halo-galaxy connection is applied to galaxy clustering in terms of the power spectrum. For this reason, in order to minimise cosmic variance, we chose to analyse the largest box available in the database, TNG300, spanning a side length of \(205\;h^{-1}{\rm Mpc}\) with periodic boundary conditions. TNG300 contains \(2500^{3}\) DM particles of mass \(4.0\times 10^{7}\;h^{-1}{\rm M_{\odot}}\) and \(2500^{3}\) gas cells of mass \(7.6\times 10^{6}\;h^{-1}{\rm M_{\odot}}\). The adequacy of TNG300 in the context of clustering science has been extensively proven in a variety of analyses (see, e.g., Contreras et al., 2020; Gu et al., 2020; Hadzhiyska et al., 2020; Montero-Dorta et al., 2020; Shi et al., 2020; Hadzhiyska et al., 2021; Montero-Dorta et al., 2021, 2021; Favole et al., 2022; de Santi et al., 2022).

In this work, we employ both galaxy and DM halo information from TNG300. DM haloes in the entire IllustrisTNG suite are identified using a friends-of-friends (FOF) algorithm based on a linking length of 0.2 times the mean of the inter-particle separation (Davis et al., 1985). As in de Santi et al. (2022), the following halo properties are used as input features to train the NNs:

* _Virial mass_ (\(M_{\rm vir}[h^{-1}{\rm M_{\odot}}]\)), which is computed by adding up the mass of all gas cells and particles contained within the virial radius \(R_{\rm vir}\) (based on a collapse density threshold of \(\Lambda_{c}=200\)). In order to ensure that haloes are well resolved, we impose a mass cut \(\log_{10}(M_{\rm vir}[h^{-1}{\rm M_{\odot}}])\geq 10.5\), corresponding to at least 500 dark matter particles.
* _Virial concentration_ (\(c_{\rm vir}\)), defined in the standard way as the ratio between the virial radius and the scale radius, i.e., \(c_{\rm vir}=R_{\rm vir}/R_{\rm s}\). \(R_{\rm s}\) is obtained by fitting the DM density profiles of individual haloes with a NFW profile (Navarro et al., 1997).
* _Halo spin_ (\(\lambda_{\rm halo}\)), for which we follow the Bullock et al. (2001) definition: \(\lambda_{\rm halo}=|J|/\sqrt{2}M_{\rm vir}V_{\rm vir}R_{\rm vir}\). Here, \(J\) and \(V_{\rm vir}\) are the angular momentum of the halo and its circular velocity at \(R_{\rm vir}\), respectively.
* _Halo age_, parametrised as the half-mass formation redshift \(z_{1/2}\). This parameter corresponds to the redshift at which half of the present-day halo mass has been accreted into a single subhalo for the first time. The formation redshift is measured following
the progenitors of the main branch of the subhalo merger tree computed with sublink, which is initialised at \(z=6\).
* The _overdensity_ around haloes on a scale of \(3\ h^{-1}\mathrm{Mpc}\left(\delta_{3}\right)\), defined as the number density of subhaloes within a sphere of radius \(R=3h^{-1}\mathrm{Mpc}\), normalised by the total number density of subhaloes in the TNG300 box (e.g., Artale et al., 2018; Bose et al., 2019).

On the other hand, subhaloes (i.e., gravitationally bound substructures) are identified in IllustrsTNG using the subfind algorithm (Springel et al., 2001; Dolag et al., 2009). Subhaloes containing a non-zero stellar mass component are labelled as galaxies. Again, following de Santi et al. (2022) for consistency, TNG300 galaxies are characterised in this work using the following basic properties:

* The _stellar mass_ (\(M_{*}\) [\(h^{-1}\mathrm{M}_{\odot}\)]), which includes all stellar particles within the subhalo. In order to ensure that galaxies are well resolved, we impose a mass cut \(\log_{10}(M_{*}[h^{-1}\mathrm{M}_{\odot}])\geq 8.75\), corresponding to at least 50 gas cells.
* The _colour_\(g-i\), computed from the rest-frame magnitudes, which are obtained in IllustrsTNG by adding up the luminosities of all stellar particles in the subhalo (Buser, 1978). Note that the specific choice of colour is rather arbitrary. We have checked that using other combinations (i.e., \(g-r\)) provides similar results.
* The _specific star formation rate_ (sSFR [\(\mathrm{yr}^{-1}h\)]), which is the star formation rate (SFR) normalised by stellar mass. The SFR is computed by adding up the star formation rates of all gas cells in the subhalo. Note that around 14% of the galaxies at redshift \(z=0\) in TNG300 have SFR\(=0\). In order to avoid numerical issues, we have adopted the same approach as in de Santi et al. (2022), assigning to these objects' artificial values of SFR sampled from a Gaussian distribution \(\mathcal{N}(\mu=-13.5,\sigma=0.5)\).
* i.e., the comoving radius containing half of the stellar mass in the subhalo.

## 3 Methodology

NNs are designed to learn how to map an instance, which is characterised by some set of input features \(X\), to a set of output features \(Y\), by weighting and combining the input features. These weights are fitted by minimising a loss function with some optimiser.

In this work, the input features are the halo properties and the outputs are the galaxy properties introduced in SS2. Starting with a sample where the target value \(Y\) is known for all instances (the TNG300 catalogue), we split it into training, validation and test sets. The training set is used to fit the model parameters (weights). The validation set is used to monitor overfitting, i.e., to ensure that the model is properly generalising to data outside of the training set, and to fit the model's hyperparameters1. The test set remains completely blind to the training and validating procedures, and can thus be used to infer the performance of the model when applied to entirely new instances. The training, validation and test sets contain, respectively, 48%, 12% and 40% of the initial sample of 174,527 objects from the TNG300 catalogue.

Footnote 1: In a NN, the model’s parameters are the weights to be learned automatically, while the hyperparameters are the number of layers, neurons, number of epochs, etc., which are often chosen manually.

Our goal is to predict central galaxy properties from a set of halo properties. In the context of ML, this would in principle fall in the category of a supervised regression problem. However, traditional regression models are designed to output single values, while any given halo could host many different central galaxies (since the set of halo properties that we use as inputs do not determine exactly the outcome of the galaxy formation process in terms of the precise values of the galaxy properties). This is reflected, as an example, in the well-known scatter in the stellar-to-halo mass relation (Wechsler, 1994; Stiskalek et al., 2022). Therefore, in order to incorporate this uncertainty, we need a model that returns not only a single best-estimate value for each galaxy property, but some proxy for the probability distribution for those properties.

In this paper, we have addressed this issue by converting the regression problem into a classification. The idea is to define \(K\) classes by splitting each galaxy property into \(K\) intervals, or bins. Just like in the usual classification tasks, the model will return a score associated with each class (bin). These scores add up to one, giving a probabilistic interpretation of the output. This approach has been widely used, as an example, in the context of photometric redshift estimation (Sadeh et al., 2016; Pasquet et al., 2019; Lima et al., 2022). We refer to our method, which is based on training NNs classifiers, as \(\mathrm{NN}_{\mathrm{class}}\).

As a starting point, we train four models to predict each galaxy property individually as univariate distributions, i.e., we have separate models to predict \(P(M_{*})\), \(P(g-i)\), \(P(\mathrm{sSFR})\), \(P(R_{1/2}^{(*)})\). As we discuss in 84, this approach is sufficient to recover the overall distribution \(P(Y)\) for a given sample. However, this does not guarantee, _a priori_, that the joint distributions are well reproduced. Therefore, we proceed to predict jointly pairs of properties, namely \(P(M_{*},g-i)\), \(P(M_{*},\mathrm{sSFR})\), \(P(g-i,\mathrm{sSFR})\) and \(P(R_{1/2}^{(*)},M_{*})\). Our strategy is similar to the univariate \(P(Y)\) case: we make a grid in the \(\{Y_{1},Y_{2}\}\) subspace in such a way that the output corresponds to pixels in this grid. Although in this paper we restrict ourselves to only two galaxy properties when predicting joint distributions, a similar approach could be used, in principle, to characterize galaxies and define populations using an arbitrary number of properties. This generalisation will be implemented in an upcoming paper.

Unless otherwise stated, for all the results shown here we set \(K=50\) classes for each one of the central galaxy properties, in equally spaced bins. For stellar mass, for example, this corresponds to bins of 0.085 dex. We must draw attention to the fact that this choice of binning is arbitrary. We have tried different numbers of bins, finding similar results in terms of the recovery of the distributions. Note that more refined versions of NNs that output distributions without binning the properties, and thus keeping it as a regression problem, already exist in the literature. In the context of photo-z estimation, Lima et al. (2022), for example, compares different types of NNs that return distributions, such as Mixture Density Networks (Bishop, 1994), Bayesian NNs, and also NNs following a similar strategy as in this work, with a binning classification scheme. Ho et al. (2021) estimate the probability distribution of the dynamical mass of galaxy clusters and also compare several types of NNs, including a classifier which is similar to our \(\mathrm{NN}_{\mathrm{class}}\). In the context of the halo-galaxy connection, Stiskalek et al. (2022) model the stellar-to-halo mass relation scatter with a Gaussian distribution and train an ensemble of NNs that predicts the mean and standard deviation. We found the binned classification to be a simpler approach that works as a proof of concept. A more careful exploration of alternative methods is left as future refinements.

Throughout the analysis, we compare our \(\mathrm{NN}_{\mathrm{class}}\) method with the deterministic models developed by de Santi et al. (2022), which we use as our baseline. In that work, several ML models are combined
to return a final, consensus output for the same galaxy properties described in SS2. The two consensus estimators are built from either the "Raw" models, which were trained with the original TNG300 sample, or the "SMOGN" models, which were trained using a data-augmented version of that data set. The SMOGN models were developed because of the difficulty for Raw models to recover the least frequent values of galaxy properties - i.e., to reproduce the tails of the distributions. The SMOGN data augmentation technique is a strategy to handle imbalanced data sets, whereby additional objects are artificially introduced in the training sample in order to force the machine to give more importance to less represented objects (Kunz, 2019).

The specifications of \(\text{NN}_{\text{class}}\) are described as follows. We use the categorical cross-entropy loss function and the adam optimiser to train the networks. The architecture may change depending on the galaxy properties to be predicted. In general, our developed networks have a single intermediate layer, with a number of neurons that typically depends on whether the output is an univariate or a joint distribution. We use the L2 regularisation, which applies a penalty proportional to the square of the model's weights. The number of epochs (iterations) is constrained with an early-stopping criteria based on the validation set loss. In the intermediate layers we used the ReLU function as activation, while in the output layer we use the Softmax function, which is similar to the Sigmoid function, but it normalises the output in such a way that the scores of the \(K\) classes add up to one. In this way, the \(\text{NN}_{\text{class}}\) output works as a proxy for a probability in bins of galaxy properties.

## 4 Results

Fig. 1 shows the distributions of the galaxies in the test set. The first column is the truth table, the TNG300 catalogue. The second column is the \(\text{NN}_{\text{class}}\) prediction of univariate distributions, i.e., galaxy properties predicted independently. With the univariate distributions we can compute the joint distributions as \(P(Y_{1})\cdot P(Y_{2})\), which are shown in the heatmap diagrams. The third column is the \(\text{NN}_{\text{class}}\) prediction for the joint distributions \(P(Y_{1},Y_{2})\), which can be integrated to recover the univariate distributions \(P(Y)\) shown in the marginal plots from the third column, i.e.:

\[P(Y_{i})=\int P(Y_{i},Y_{j})dY_{j}. \tag{1}\]

The univariate distributions predicted by \(\text{NN}_{\text{class}}\), shown in black solid lines in the second-column plots of Fig. 1, are in excellent agreement with the true distributions from TNG300, shown in gray shaded regions. They also reproduce fairly well the joint distributions \(P(Y_{1})\cdot P(Y_{2})\) for most cases. The \(P(g-i)\cdot P(\text{sSFR})\) joint distribution, however, fails to reproduce the shape of the distribution for redder colours and lower sSFRs. According to this prediction, red galaxies could have virtually any value of sSFR, while what we actually observe in TNG300 is that as galaxies move from the blue to the red the peak, their sSFRs decrease. This important feature is recovered when \(\text{NN}_{\text{class}}\) is trained to predict \(P(g-i,\text{sSFR})\) jointly (third column in Fig. 1).

The above result indicates that our input halo properties alone are unable to predict accurately the correlations between colour and sSFR. The model would need additional features in order to capture this relation. It is interesting, however, that we can overcome this limitation by predicting the joint distribution directly using only the presented halo properties. This exercise indicates that, in order to robustly assign galaxies to haloes, with all the properties consistently correlated, the properties should be predicted together. Note that, in principle, one could define galaxy populations based on as many parameters as wished. Therefore, in the most general case, we would have an \(N\)-dimensional distribution associated to each host halo.

As a complementary analysis, Fig. 2 shows two additional well-known relations in the context of the halo-galaxy connection: the stellar-to-halo mass relation, and the galaxy size-halo mass relation obtained with TNG300 and with \(P(M_{*})\) and \(P(R_{1/2}^{(*)})\) predicted by \(\text{NN}_{\text{class}}\).

Figures 1 and 2 allow for a visual inspection of the results. In order to quantify the similarity between the distributions, we have performed the Kolmogorov-Smirnov (KS) test, which measures the maximum distance between cumulative distributions (for more details, see Ivezic et al., 2014):

\[\Delta=max(|F_{1}-F_{2}|). \tag{2}\]

The results are shown in Table 1. For comparison, we also show the values obtained with our baseline models, Raw and SMOGN, from de Santi et al. (2022). Once again, we see that for most cases the independent prediction of univariate distributions reproduce fairly well the joint distributions, except for colour and sSFR. In all cases, \(\text{NN}_{\text{class}}\) provides significantly lower values as compared to Raw and SMOGN.

So far, we have focused on the combined distributions for the entire test sample. We now turn our attention to individual objects and the probability distributions that our ML machinery predicts for them. In particular, Fig. 3 displays, in a similar format to that of Fig. 1, some examples of the joint probability distribution \(P(M_{*},g-i)\) for three illustrative cases: a red object, a blue object, and an object lying at the so-called green valley region (from left to right). In each panel, the host halo mass is specified on the top, whereas the true TNG300 values of stellar mass and colour are shown as the dashed lines. As a reference, we also include in the marginal plots the distributions of the objects in the test set within a bin of \(\pm 0.1\) in halo mass around the values indicated on the top of the plots.

The first thing to notice from Fig. 3 is that the distributions are significantly narrower along the x-axis, as compared to the y-axis. This is of course expected, since stellar mass is the galaxy property that displays a tighter relation with the halo properties (particularly with halo mass), and therefore is the easiest to predict. It is also noteworthy that not all distributions can be well approximated by a Gaussian distribution. Some distributions are significantly skewed or, depending on halo mass, even bimodal, reflecting the well-known colour/sSFR bimodality of the galaxy population (e.g., Baldry et al., 2004).

The red galaxy on the left-hand panel shows very little scatter in colour. This is typically the case for red galaxies hosted by haloes with \(\log_{10}(M_{\text{vir}}[h^{-1}M_{\odot}])\gtrsim 12.5\). By visually inspecting Fig. 1 and Fig. 2, we can get a sense as to why this happens: massive haloes are typically populated by massive galaxies, since the scatter in the stellar-to-halo mass relation is small. Massive galaxies are almost exclusively very red, which explains why the machine predicts a very narrow distribution of colours from the set of halo properties employed. The situation is very different for the blue galaxy featured in the middle panel. In this case, the predicted colour distribution is much broader than that for the red galaxy. Here, the host halo mass is much smaller, which implies a larger scatter in the stellar-to-halo mass relation. On top of that, blue galaxies intrinsically display a wide range of colours. All this uncertainty is captured by the machine in terms of a wider colour distribution.

Finally, the green-valley galaxy on the right-hand panel of Fig. 3 represents the most extreme case of the three, where the colour
Figure 1: Distributions of galaxy properties. From top to bottom: colour \(v\). stellar mass, sSFR \(v\). stellar mass, sSFR \(v\). colour, and radius \(v\). stellar mass. The first column shows the true distributions from TNG300. The second column shows the distributions computed from the univariate distributions as predicted by \(\rm NN_{class}\) – i.e., predicted independently from each other. The third column shows the joint distributions as predicted by \(\rm NN_{class}\). The grey shaded regions in the marginal plots correspond to the TNG300 distributions, while the black solid lines correspond to the \(\rm NN_{class}\) predictions. The univariate distributions shown in the third column plots were computed by marginalising the joint distributions.
degeneracy produces a bimodal distribution. These objects are caught between two intrinsically different populations, i.e., the blue cloud and the red sequence. The analysis of individual distributions reveals that these objects are the ones that display a weaker relation with the properties of their host haloes (at least the ones analysed in this work). As discussed in de Santi et al. (2022), these objects exemplify the most clear case where halo properties alone seem insufficient to predict the colour/sSFR, thus emphasising the advantages of our probability-based methodology.

This probability distribution description on an individual-object basis allows us to explore the dependence of galaxy properties on secondary halo properties at fixed halo mass (a dependence that is closely related to the so-called galaxy assembly bias effect, see, e.g., Wechsler and Tinker, 2018; Sato-Polito et al., 2019; Montero-Dorta et al., 2020, 2021). In particular, we have analysed the dependence of \(P(M_{*},g-i)\) on halo age at fixed halo mass for green-valley objects. To this end, we selected objects in the test sample with predicted colour within the range \(0.80<g-i\leq 1.05\) and halo masses of \(11.8<\log_{10}(M_{\rm vir}[h^{-1}{\rm M}_{\odot}])<12.2\) (we have checked that choosing a narrower halo mass range would not alter our results significantly). This subset was subsequently split by halo age (taking the 15% and 85% quantiles). For younger haloes, a stack of all distributions still reveals some bimodality in colour, albeit with a stronger preference for the blue peak. The predicted probability distribution for green-valley galaxies in older haloes is, conversely, much more skewed towards redder colours. The tail of the distribution for these objects still covers the green valley, which means that in some realisations these host haloes will be populated by a green-valley central galaxy (although the probability for this to happen is low). These results are reassuring in terms of the robustness of our methodology, demonstrating that our probability description is capable of capturing secondary halo dependencies.

## 5 Power Spectrum

With the help of the method presented in this work we have greater flexibility to define different tracers based on galaxy properties. In this section, we explore the performance of \(\rm NN_{class}\) in terms of the accuracy with which we can reproduce the power spectra of those tracers. We compute spectra for tracers in the test set, using the python package nbodykit(Hand et al., 2018). For the truth TNG300 catalogue we use the positions of the central galaxies, but for the predictions we use the positions of the host haloes. Once again, we compare \(\rm NN_{class}\) with the baseline models from de Santi et al. (2022). As a complementary analysis, in Appendix B we compare the power spectra of tracers defined according to the same criteria of that previous work, which are based on individual galaxy properties.

Since TNG300 is a single box, the uncertainties of the spectrum on each bandpower \(k_{i}\), for each tracer \(\alpha\), are computed according to the theoretical (Gaussian) covariance, i.e.:

\[\frac{\sigma_{\alpha,i}^{2}}{P_{\alpha,i}^{2}}=\frac{2}{V\bar{V}}\left(\frac{ 1+\bar{n}_{\alpha}P_{\alpha,i}}{\bar{n}_{\alpha}P_{\alpha,i}}\right)^{2}, \tag{3}\]

with \(\bar{V}=4\pi k_{i}^{2}\Delta k/(2\pi)^{3}\), and the residuals are defined as

\[\frac{\left(p_{\alpha,i}^{\rm pred}-p_{\alpha,i}^{\rm TNG300}\right)^{2}}{ \sigma_{\alpha,i}^{2}}\,. \tag{4}\]

Our choice of tracers is driven by the fact that the target selection in galaxy surveys often rely on the analysis of colour-magnitude diagrams (see e.g. Eisenstein et al., 2001, 2011; Zhou et al., 2020). One of the most common ways to define galaxy populations is in terms of the red sequence and the blue cloud, which can also be clearly distinguished in the colour-stellar mass diagram, as shown in Fig.1. They are two distinct populations with different biases, hence their interest for studies of large scale structure.

In a similar fashion, we defined seven tracers (\(\alpha=1,\ldots,7\)) based on the colour-stellar mass diagram, \(P(M_{*},g-i)\). We split red galaxies (\(g-i>1.05\)) into lower (\(\alpha=1\)) and higher (\(\alpha=2\)) stellar masses. Conversely, "green-valley" galaxies (defined as \(0.80<g-i\leq 1.05\)) are split into three mass bins, leading to populations \(\alpha=3,4,5\). Finally, blue galaxies (\(g-i\leq 0.8\)) are separated into lower (\(\alpha=6\)) and higher (\(\alpha=7\)) stellar mass bins. This selection is outlined in Table 2, and it is represented in the lower right corner of Fig. 4.

An interesting feature of the probabilistic approach is that each

\begin{table}
\begin{tabular}{l c c c c c c c} \hline
**1D KS** & \(P(Y)\) & Raw & SMOGN & **2D KS** & \(P(Y_{1})\cdot P(Y_{2})\) & \(P(Y_{1},Y_{2})\) & Raw & SMOGN \\ \hline \(P(M_{*})\) & 0.002 & 0.064 & 0.064 & \(P(M_{*},g-i)\) & 0.010 & 0.005 & 0.183 & 0.163 \\ \(P(g-i)\) & 0.004 & 0.181 & 0.116 & \(P(M_{*},\rm sSFR)\) & 0.012 & 0.009 & 0.253 & 0.209 \\ \(P(\rm sSFR)\) & 0.004 & 0.213 & 0.168 & \(P(g-i,\rm sSFR)\) & 0.110 & 0.009 & 0.266 & 0.176 \\ \(P(R_{1/2}^{(\alpha)})\) & 0.009 & 0.217 & 0.110 & \(P(M_{*},R_{1/2}^{(\alpha)})\) & 0.015 & 0.007 & 0.217 & 0.150 \\  & & & & & \(P(M_{\rm vir},M_{*})\) & 0.008 & – & 0.064 & 0.064 \\  & & & & & \(P(M_{\rm vir},R_{1/2}^{(\alpha)})\) & 0.012 & – & 0.217 & 0.110 \\ \hline \end{tabular}
\end{table}
Table 1: KS test values for univariate (1D) and joint (2D) distributions computed with the NNs and the baseline models.

Figure 2: Stellar-to-halo mass relation (top) and galaxy size–halo mass relation (bottom) from the TNG300 catalogue (left) and from \(\rm NN_{class}\) predictions (right).
galaxy is generated through a realisation of a probability distribution spreading over many bins. As a consequence, we can build many catalogues of central galaxy properties by drawing values \(y_{1},y_{2}\) from \(P(Y_{1},Y_{2})\). We have performed \(r=42\) realisations of \(P(M_{*},g-i)\), leading to as many values of \(M_{*}\) and \(g-i\) for each halo. We then compute the spectrum of each of these samples, and from that the mean and variance of the spectra. For the mean spectrum \(\hat{P}_{\alpha,i}\), we compute the uncertainties according to Eq. (3).

Fig. 4 shows the power spectra and residuals of the seven tracers defined in terms of \(P(M_{*},g-i)\) - see Table 2. Tracers \(\alpha=3,4\) are relatively rare, hence their corresponding regions in colour-stellar mass space are poorly populated by single-point estimators. Therefore, a model that predicts galaxies in these regimes improves the quality of the fit considerably - i.e., it reduces the \(\chi^{2}\). We had already seen an improvement with the SMOGN models, which better recover this region as compared to the Raw models, but with NN\({}_{\rm class}\) this improvement is even more pronounced. There are only a few \(\alpha=5\) galaxies in TNG300, which makes this population very sparse. In particular, it has the largest variance over realisations. Conversely, all models are equally good at reproducing the power spectra of tracer populations closer to the peaks of the probability distributions: for \(\alpha=1,2,6,7\), the \(\chi^{2}\) is comparable between all models.

As discussed above, we are able to draw multiple samples from the probabilities predicted by NN\({}_{\rm class}\). Each realisation leads to slightly different power spectra, as can be seen in Fig. 4. By computing the variance of the multiple \(P(k)\) we can assess the uncertainties due to the intrinsic stochasticity in the halo-galaxy connection. Fig. 5 compares the relative errors \(\sigma^{2}/P_{\rm TNG300}^{2}(k)\) computed using \(\sigma_{\rm CV}^{2}\), from Eq. 3 (which encodes the uncertainty due to cosmic variance, CV), with \(\sigma_{\rm NN_{class}}^{2}\), which encodes the statistical uncertainties in the halo-galaxy connection estimated with NN\({}_{\rm class}\). As we already saw in Fig. 4, the cosmic variance error bars are typically larger than the scatter in the power spectra due to the multiple realisations of the NN\({}_{\rm class}\) probabilities. The contribution of \(\sigma_{\rm NN_{class}}^{2}\) seems more relevant for the tracer population 5, which is very sparse. However, for all tracers \(\sigma_{\rm CV}^{2}\) decreases for smaller scales (due to the Fourier bin volume), while \(\sigma_{\rm NN_{class}}\) remains approximately constant. Therefore, the relative contribution of \(\sigma_{\rm NN_{class}}\) for the total error budget of the power spectra appears to become more important at smaller scales.

Even though we see no evidence of a bias associated with this additional source of statistical uncertainties, the stochastic nature of the relationship between galaxies and their haloes may present further challenges for multi-tracer analyses of LSS (Seljak, 2009; McDonald & Seljak, 2009). The advantages of the multi-tracer technique are reliant upon the partial cancellation of cosmic variance that results from clustering measurements from different galaxy types that are assumed to reflect the same underlying dark matter density field - in that respect see also Abramo & Leonard (2013); Abramo et al. (2016). The "stochastic bias" associated with the nature of the galaxy-halo connection can dilute some of the expected cosmic variance cancellation. However, that stochastic component seems to affect mostly the power spectra on small scales, where non-linear effects already limit our ability to employ the multi-tracer technique effectively - see, e.g., Montero-Dorta et al. (2020a).

## 6 Discussion and Conclusions

Although there is an obvious relation between the baryonic and DM components of haloes, there is also mounting evidence that the properties of haloes alone are insufficient to reproduce the properties of galaxies, since the latter are shaped by a variety of galaxy-formation processes. On the other hand, ML regression models are traditionally designed to reproduce single-value statistics, and thus are ill-equipped to encode the intrinsic scatter in the halo-galaxy connection. Building on the recent work of de Santi et al. (2022), here we use the TNG300 hydrodynamical simulation in combination with NNs to map the connection between the properties of central galaxies and the properties of their hosting haloes. As in the aforementioned work, NNs are trained to reproduce the stellar mass, \(g-i\) colour, sSFR and radius of TNG300 galaxies based on a set of halo/environmental properties that include virial mass, concentration, formation redshift, spin, and overdensity (computed over scales of 3 \(h^{-1}\)Mpc). In order

\begin{table}
\begin{tabular}{l c c c} \hline Tracer & \(\log\,(M_{*}[h^{-1}{\rm M}_{\odot}])\) & \(g-i\) & \# objects \\ \hline \(\alpha=1\) & \((9.5,\,10.5)\) & \((1.05,\,)\) & \(4,073\) \\ \(\alpha=2\) & \((10.5,\,)\) & \((1.05,\,)\) & \(5,207\) \\ \(\alpha=3\) & \((9.5,\,)\) & \((0.80,\,1.05]\) & \(4,786\) \\ \(\alpha=4\) & \((9.5,\,10.5]\) & \((0.80,\,1.05]\) & \(5,950\) \\ \(\alpha=5\) & \((10.5,\,)\) & \((0.80,\,1.05]\) & \(1,267\) \\ \(\alpha=6\) & \((9.5,\,5]\) & \((,0.80]\) & \(29,695\) \\ \(\alpha=7\) & \((9.5,\,10.5]\) & \((,0.80]\) & \(18,432\) \\ \hline \end{tabular}
\end{table}
Table 2: Criteria for splitting central galaxies by stellar mass and colour, in order to define the tracers used in the power spectrum analysis.

Figure 3: \(P(M_{*},g-i)\) for individual objects predicted by NN\({}_{\rm class}\). The dashed green lines show the true values for stellar mass and colour from TNG300. The shaded regions in the marginal plots are the distributions of objects with similar halo mass as indicated on the top of the corresponding panel.
to alleviate the deficiencies of ML deterministic regression models, we have tested a different approach for the first time in the context of the halo-galaxy connection. The NNs are now trained to predict probability distributions instead of single-value statistics by means of a binning classification scheme. In essence, the distributions of galaxy properties are split into \(K\) narrow bins so that the NNs can associate a score to each of the \(K\) classes. This is performed in such a way that the output can be used as a proxy for the probability distributions of the central galaxy properties.

We have shown that this approach is in fact capable of producing bivariate distributions of galaxy properties, i.e., \(P(Y_{1},Y_{2})\), in outstanding agreement with those from TNG300 (here, \(\{Y_{1},Y_{2}\}\) is any pair of galaxy properties). These joint distributions can be compared with the product of the two 1D (disjoint) distributions, \(P(Y_{1})\)

Figure 4: Power spectra and residuals for seven tracers selected on the basis of the colour-stellar mass diagram (bottom right panel). The green solid lines correspond to TNG300, while the light purple solid lines correspond to spectra from \(r=42\) samples drawn from the probabilities predicted by NN\({}_{\rm class}\). The dark purple, thick dashed lines correspond to the mean of those realisations. The baseline models are shown in orange: darker dotted lines correspond to the Raw model and lighter dotted-dashed lines correspond to the SMOGN model.

Figure 5: Relative error for seven tracers selected based on the colour - stellar mass diagram. The variances are normalised by the TNG300 spectrum \(P_{T}(k)\) of each tracer \(\alpha\). Orange dotted lines correspond to the relative error computed with Eq.(3), purple dashed lines correspond to the relative error computed with NN\({}_{\rm class}\) and green solid lines correspond to the total relative error.