Improving Convolutional Neural Networks for Cosmological Fields with Random Permutation

The Large Scale Structure (LSS) of the universe contains crucial information about its late-time growth history and fundamental physics. Over the past few decades, a diverse array of probes has been employed to study these structures, each contributing significantly to our understanding of the universe. Prominent among these are the Stage-III LSS surveys, which include the Dark Energy Survey (DES) Dark Energy Survey Collaboration et al. (2016); Troxel et al. (2018); Amon et al. (2022), the Kilo-Degree Survey (KiDS) Kuijken et al. (2015); Heymans et al. (2021); Asgari et al. (2021), the Hyper Suprime-Cam Subaru Strategic Program (HSC) Dalal et al. (2023); Li et al. (2023); Sugiyama et al. (2023), and the Baryon Oscillation Spectroscopic Survey (Boss and eBOSS) Smee et al. (2013); Zhao et al. (2022); Alam et al. (2021).

The upcoming Stage-IV LSS surveys, including Rubin Observatory’s LSST Ivezić et al. (2019), the Roman Space Telescope Dore et al. (2019), and the recently launched Euclid mission Laureijs et al. (2011), in combination with the Dark Energy Spectroscopic Instrument Levi et al. (2019), are or will be gathering data soon. All these efforts will shed light on two of the most mysterious problems in physics: dark matter and dark energy.

One of the most important analytical tools used in these studies is the two-point correlation function. This function links observational measurements with theoretical predictions derived from perturbation theory. However, the two-point correlation function only captures the Gaussian information.

A variety of non-Gaussian statistics have thus been proposed to study cosmological fields. These include the three-point function and bispectrum Takada and Jain (2003); Halder et al. (2023); Takada and Jain (2004), peaks and voids Jain and Van Waerbeke (2000); Marian et al. (2009); Martinet et al. (2018); Harnois-Déraps et al. (2021); Zürcher et al. (2022a); Marques et al. (2023), Minkowski functionals Munshi et al. (2012); Kratochvil et al. (2012); Grewal et al. (2022), Betti numbers and persistent homology Feldbrugge et al. (2019); Parroni et al. (2021); Heydenreich et al. (2021, 2022), k-NN and CDF Banerjee and Abel (2023); Anbajagane et al. (2023), and wavelet transform based statistics Cheng et al. (2020); Cheng and Ménard (2021); Valogiannis and Dvorkin (2022); Hothi et al. (2023); Pedersen et al. (2022); Gatti et al. (2023a). See Ref. Euclid Collaboration et al. (2023) for a forecast with the settings of the Euclid Mission.

Most recently, machine learning has been used to extract cosmological information at the field level, particularly with Convolutional Neural Networks (CNN) Fluri et al. (2018); Ribli et al. (2019a); Gupta et al. (2018); Fluri et al. (2019); Ribli et al. (2019b); Matilla et al. (2020); Jeffrey et al. (2021a); Lu et al. (2022, 2023); Hortúa et al. (2020); Akhmetzhanova et al. (2024), Graph Neural Networks (GNN) Perraudin et al. (2019); Fluri et al. (2022); Makinen et al. (2022); Mishra-Sharma (2022), and generative models Dai and Seljak (2022); Ullmo et al. (2021); Yiu et al. (2022); Villaescusa-Navarro et al. (2021a); Dai and Seljak (2023).

Neural networks generally surpass traditional Gaussian statistics in constraining cosmologies, such as the power spectrum. However, the extent of the improvement can vary significantly depending on the realism of the simulations and the scales considered. When compared to other non-Gaussian statistics, the situation is more complicated. For instance, several studies Fluri et al. (2018); Gupta et al. (2018) have suggested that CNNs outperform peak counts methods, whereas others have pointed out that their enhancement over Minkowski functionals is only moderate Matilla et al. (2020). Moreover, there are scenarios where CNNs appear not to perform as effectively as other methods, such as the scattering wavelet transform Cheng et al. (2020). This observation raises the question that while neural networks offer substantial advancements in many areas, their efficacy can vary for cosmological fields. In other words, even though the CNN often contains thousands of free parameters to approximate the underlying function, it might not be suitable to directly apply the commonly used classification-oriented CNN in cosmology. Such insights are crucial for guiding future research and methodology choices in the field of cosmology, particularly in the integration and application of machine learning techniques.

In this paper, we explore modifications to the architecture of CNNs to enhance their performance in cosmological tasks. Cosmological maps such as weak lensing mass maps differ from images used to develop CNNs in that they are stochastic and low signal-to-noise. Moreover cosmological inference we use them for regression, not classification as is typically done in industry. So we explore a variety of regularization and data augmentation schemes in applying CNNs on lensing maps. Our preliminary findings suggest that employing a combination of maximum and average pooling improves the network’s ability to extract relevant information. We also introduce a novel regularization scheme rooted in our understanding of cosmological fields: the incorporation of random permutation layers designed to effectively augment the training data by sacrificing the information of large-scale correlation. We demonstrate the effectiveness of this straightforward technique for statistical noise levels expected in current and upcoming weak lensing surveys. We also include a number of sources of systematic uncertainty.

Additionally, we conduct comparative analyses with conventional regularization methods and alternative strategies that disrupt large-scale correlation. Our initial results indicate that the permutation operation directly applied to the latent space of the neural network consistently boosts accuracy in both the simplified and realistic simulation cases considered in this work. This suggests that by selectively filtering out large-scale information in the deeper layers, the neural network exhibits enhanced performance in extracting information from more critical scales. We anticipate that this simple regularization technique not only holds promise for improving weak lensing fields but may also for other stochastic fields, such as galaxy maps or 21-cm maps. We do not assert that our architecture represents the optimal model as we have not carried out a systematic study for all possible use cases.

The structure of this paper is as follows: In Sec. II we describe the simulations we used for this work, the preprocessing of the data, and the simulation-based inference setup. In Sec. III we discuss how the CNN for cosmological fields should be different from the standard design for image classification. In Sec. IV we present the random permutation layer as a regularization scheme that helps model generalization. We summarize the findings and outlook for future work in Sec. V and we perform extensive tests comparing to other models in the Appendices.

Unlike the prevalent use of CNNs in image recognition tasks, their application in cosmology typically addresses regression problems rather than image classification. In this context, the precision of the output is most important. Moreover, the input data in cosmology, in particular weak lensing fields, significantly differs from conventional images, such as those of cats and boats and others that are part of ImageNetDeng et al. (2009), the long-standing industry benchmark for image classification algorithms. As depicted in Fig. 1, the weak lensing field is stochastic, low signal-to-noise, and has correlations on all scales. We seek a CNN approach that is simple and generalizable, rather than adopting complex architectures like ResNet He et al. (2015) or VGGNet Simonyan and Zisserman (2014) that have been optimized for classification tasks on images in other domains. Prior research Fluri et al. (2018); Ribli et al. (2019a); Gupta et al. (2018); Jeffrey et al. (2020); Lu et al. (2022); Matilla et al. (2020); Lu et al. (2023) employing CNNs for analyzing weak lensing fields has yielded promising results using relatively simple architectures.

Our 6-layer CNN model, as shown in Fig. 2, incorporates a notable departure from previous designs Fluri et al. (2018); Ribli et al. (2019a); Gupta et al. (2018); Fluri et al. (2019); Ribli et al. (2019b); Matilla et al. (2020); Jeffrey et al. (2021a); Lu et al. (2022, 2023); Hortúa et al. (2020), where either the Maximum pooling or the Average pooling is used. Instead, we have combined the use of Maximum Pooling and Average Pooling layers. One motivation for us to use Maximum Pooling in the initial layers is driven by the observation that positive extreme values in our pixels correspond to galaxy clusters which are known to carry valuable cosmological information. Indeed, a study of CNNs using saliency maps for lensing mass maps  Matilla et al. (2020) found that clusters are more informative than voids in the presence of realistic noise. Conversely, deeper layers, correlating distant sky regions of the original map, may require equitable consideration. Here, our choice to switch to average pooling facilitates better generalization in the subsequent Multi-Layer Perceptron (MLP) layers. We test these choices as discussed below and in Appendix. A.

We tested both simulation cases 1 and 2 and also increased the effective galaxy density (lower noise) to very approximately represent future LSST and Roman weak lensing surveys (we do not attempt to simulate the redshift coverage of those surveys). However, the best pooling option does not appear to correlate with noise level. In fact, the average pooling is similar to or better than the mix-pooling. When training on $\Omega_{\mathrm{m}}$ though, the network fails to converge. We also tested extreme pooling with the largest absolute value (which then includes underdense extrema) and found no difference in the final results. For this reason, we use the mixed-pooling as our baseline choice. One caveat is that the extremely valued pixels could also be the ones most affected by noise or limitations in the simulations or unmodelled effects; this becomes more of a concern for smaller pixel values than ours. Adopting pooling layers in a Bayesian way could potentially overcome this problem Njieutcheu Tassi (2020).

It’s important to note that this architecture is not optimal for every scenario, such as when dealing with patches of different sizes or different choices of systematic uncertainties. Furthermore, we only optimize the architecture for the inference of $S_{8}$ while a better design for $\Omega_{\mathrm{m}}$ is likely different and merits follow-up work.

Nonetheless, our findings below show that simple, physics-inspired modifications to a CNN can significantly improve accuracy. This encourages the pursuit of tailored models for specific problems in cosmology, rather than relying on existing architectures primarily designed for image classification.

Adding randomness is a common way in machine learning to improve the generalization accuracy. In this section, we introduce such a technique: random permutation of the pixels in the deeper layers of the network that correspond to large scales.

Figure 2 illustrates our design and the positions we replace the normal design with the random permutation layers. The random permutation layer shuffles (randomizes) the position of each pixel of the feature map. In this work, the shuffle is different for each batch, each channel, and changes every epoch. Note that the four tomographic bins have the same random permutation (shuffle) applied ²²2This is due to the fact that we use the Conv2d function implemented in Pytorch https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html.

As an example, if we add the permutation layer before convolution layer 6 (Option 1), the input is the feature map of size 8 by 8 for 128 channels. The random permutation operation shuffles these 8x8 pixels randomly for each 128 channel differently. Each pixel in this feature map approximately corresponds to $512/8\times 6.8\approx 435.2\ \mathrm{arcmin}$. At this scale, the two-point correlation has a large sample variance. We find that by removing the information at large scales in the deeper layers, the neural network optimizes for the scales that are more relevant for cosmology. Moreover, it appears to extract some of the large-scale information in the earlier layers, though as discussed below we do not yet understand exactly how it does that. The shuffling operation can be inserted at different positions in the network, depending on the scale of which noise dominates over the signal. Note that this way of calculating scales is only approximate.

We test the positions of the shuffling operation for simulation cases 1 and 2. As shown in Fig. 3, the accuracy depends on where we perform the shuffling. For simulation case 1 it is the second last layer and for simulation case 2 it is the third last layer. The simplification effectively makes them 5-layer and 4-layer models respectively. In general, we expect the results to depend on the underlying variation of cosmologies, the noise level, tomographic binning, patch size, training set size, and other details. In particular, it will be interesting to see how increasing sample variance at large scales impacts the choice of the network Doux et al. (2022). To ensure a fair comparison, we also varied the depth of the non-shuffled CNN and found the 6-layer works the best (see Appendix. B).

The errors on the test set for both simulation cases are shown in Fig 4 and Fig 5 respectively. The posteriors for the two parameters are shown in Fig. 6. We find that the CNN with shuffling the standard deviation of $S_{8}$ is 31% smaller than the non-shuffle case.

Note that the CNN for the two cases is optimized separately, and indeed it is also different for $S_{8}$ and $\Omega_{\mathrm{m}}$. In particular, it is the 6-layer model for non-shuffle for $S_{8}$, the 4-layer model for non-shuffle for $\Omega_{\mathrm{m}}$, and both 4-layer model for shuffle-CNN (option 3). To make the comparison for the same CNN model with the same weight size, we use shuffle option 1 for $S_{8}$ as well and get a 17% improvement on the standard deviation of $S_{8}$. However, it is important to note that this change of layer depths only helps the shuffle-CNN, while the non-shuffle CNN suffers from insufficient depth in our case.

We applied the TARP coverage test with 58 simulations (as there are only 58 independent cosmologies for the $S_{8}-\Omega_{\mathrm{m}}$ plane), and found the posterior is well calibrated as shown in Fig. 7.

A pseudo-code in Pytorch is displayed in Algorithm. 1. Only the lines associated with the random permutation operation are shown. Note that the self_training variable might not be effective depending on the version of Pytorch. In that case, a customized variable should be used. This is also optional as in some cases random shuffling in both training and validation could boost the overall accuracy. Similar to this, we can also easily define the model to only shuffle the pixels after every N epoch.

CNNs are the most extensively developed deep learning approach for image analysis. In recent years they have also been applied for inferring cosmological parameters from weak lensing and other cosmological maps. We show that for the characteristics of cosmological fields like lensing mass maps, variations of the standard CNN architecture can lead to improved performance. We have presented a set of regularization and data augmentation methods and quantified the performance improvement for simulated lensing surveys. The simulated surveys mimic the Dark Energy Survey parameters in both statistical and systematic uncertainties, but we also consider the lower shape noise levels corresponding to Stage-IV surveys. The explorations in this paper are not definitive as the detailed implementation should be adapted to the specific applications.

We use random permutations (shuffling operation) in the deep layers of CNN as a novel regularization and data augmentation technique. We also use a mix of maximum and average pooling for varying layer depths. These simple modifications can enhance the performance compared to traditional CNN designs that have been adapted from image classification. Figure. 6 demonstrates the improvement in parameter inference. We also show in Fig. 8 that it can be expected to have better generalization accuracy. While the design requires optimization for specific problems, incorporating a random permutation layer emerges as a promising approach for cosmological fields. An open question that we have addressed only partially is whether CNNs like the ones we have tested lose some large-scale information. It is possible that the early layers of the CNN, prior to shuffling, have extracted the available signal. We find some evidence for this when considering the power spectrum – adding the power spectrum, including the large scales (small angular wavenumbers), to the MLP does not improve the performance further. We leave a detailed investigation for future work.

With the upcoming Stage-IV surveys (Euclid, Rubin-LSST, and Roman), we expect a substantial improvement in the quality and size of lensing mass maps. Deep learning approaches along with higher-order statistics and field-level inference can extract the vast amount of information from these maps that goes beyond standard 2-point statistics. While CNNs may not capture all the information contained in cosmological maps, they offer some clear advantages. As a model-specific data compressor, CNNs are flexible for a lot of different settings. For example, if the data has multiple tomographic bins as in current weak lensing surveys, CNNs can take them as different channels and capture the cross-channel information. If one prefers to analyze data with a shear field instead of the convergence field, one can also treat the two components as different channels and avoid the complexity of reconstructing the convergence map. Hence exploring improved performance with CNNs is likely to be of value in cosmological applications.

Caveats and Future work: This paper has focused on testing a set of regularization and data augmentation schemes with a simple CNN design. Given that this approach is inspired by our understanding that cosmological fields are noisy on scales approaching the survey size, we anticipate that the random permutation layer, when applied at appropriate scales, could be synergized with more complex architectures. For instance, similar shuffling operations can be easily tested with Vision Transformer Hwang et al. (2023) or Graph-based Neural Networks Perraudin et al. (2019).

It is also interesting to test if this regularization and data augmentation method can be effective in other cosmological fields such as the overdensity field Villaescusa-Navarro et al. (2020, 2021b), 21-cm maps Novaes et al. (2023), and secondary anisotropies in CMB temperature or polarization maps Casas et al. (2022); Münchmeyer and Smith (2019).

A useful direction for future research is the optimization of the shuffling operation. One possibility is to introduce a hyperparameter that modulates the probability of performing the shuffle, which is a technique used in Swiderski et al. (2022). Our limited explorations are conservative in the sense that we did not tune hyperparameters, so there is room for improvement by tuning all the hyperparameters for different variations of the noise level. Additionally, exploring other randomness-incorporating structures, such as random shifting Zhao et al. (2017) or randomly wired CNNs Xie et al. (2019), could also offer valuable insights.

Another important direction for future research is in the interpretability of CNNs when applied to cosmological data. Despite their proven efficacy in regression problems, a persistent challenge with CNNs is the ’black box’ nature of their decision-making processes. Understanding how these networks derive their conclusions is important, not just for validating and improving accuracy, but also for gaining deeper insights into the underlying physics. Previous studies Matilla et al. (2020); Ntampaka and Vikhlinin (2022) show intriguing results by using saliency maps Simonyan et al. (2013); Springenberg et al. (2014) (see also a recent study You et al. (2023) using a ‘sum-of-parts’ approach to interpretability). Architectures designed to relate the neural network results to physical qualities like the N-point correlation function are also promising Miles et al. (2021); Gong et al. (2024). It will be interesting to see how the results change with the presence of random permutation layers.

We thank Shubh Agrawal, Pratik Chaudhari, Cyrille Doux, Rafael Gomes, Mike Jarvis, Amrut Nadgir, Shivam Pandey, Helen Qu, Dimitrios Tanoglidis, and Eric Wong for useful discussions. We are especially grateful to Tomek Kacprzak and his collaborators for generating and making available the DarkGrid simulations used in this work. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract No. DE-AC02-05CH11231. B.J. and M.G. are supported in part by the US Department of Energy grant DE-SC0007901.