Revealing the 3D Cosmic Web through Gravitationally Constrained Neural Fields
Authors: Brandon Zhao
Aviad Levis
Liam Connor
Pratul P. Srinivasan
Katherine Bouman
Venue: ICLR 2025 Poster (2025)
Keywords: computational imaging
signal processing
inverse problems
astrophysics
cosmology
neural fields
machine learning for physical sciences
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Weak gravitational lensing is the slight distortion of galaxy shapes caused primarily by the gravitational effects of dark matter in the universe. In our work, we seek to invert the weak lensing signal from 2D telescope images to reconstruct a 3D map of the universe's dark matter field. While inversion typically yields a 2D projection of the dark matter field, accurate 3D maps of the dark matter distribution are essential for localizing structures of interest and testing theories of our universe. However, 3D inversion poses significant challenges. First, unlike standard 3D reconstruction that relies on multiple viewpoints, in this case, images are only observed from a single viewpoint. This challenge can be partially addressed by observing how galaxy emitters throughout the volume are lensed. However, this leads to the second challenge: the shapes and exact locations of unlensed galaxies are unknown, and can only be estimated with a very large degree of uncertainty. This introduces an overwhelming amount of noise which nearly drowns out the lensing signal completely. Previous approaches tackle this by imposing strong assumptions about the structures in the volume. We instead propose a methodology using a gravitationally-constrained neural field to flexibly model the continuous matter distribution. We take an analysis-by-synthesis approach, optimizing the weights of the neural network through a fully differentiable physical forward model to reproduce the lensing signal present in image measurements. We showcase our method on simulations, including realistic simulated measurements of dark matter distributions that mimic data from upcoming telescope surveys. Our results show that our method can not only outperform previous methods, but importantly is also able to recover potentially surprising dark matter structures.

Published as a conference paper at ICLR 2025
REVEALING THE 3D COSMIC WEB THROUGH GRAVI- TATIONALLY CONSTRAINED NEURAL FIELDS
Brandon Zhao Aviad L
Liam Connor Pratul P. Srinivasan Katherine L. Bouman 1.6
1 'Department of Computing and Mathematical Sciences, California Institute of Technology
Department of Computer Science, University of Toronto
3 David A. Dunlap Department of Astronomy & Astrophysics, University of Toronto
arXiv:2504. 25....2.11 Dsstro-ph.Cor 21 Apr 2025
4 *Center for Astrophysics, Harvard & Smithsonian
5 GGogle DeepMind
6 Departments of Astronomy and Electrical Engineering, California Institute of Technology "byzhaoBcaltcee edu
ABSTRACT
Weak gravitational lensing is the slight distortion of galaxy shapes caused primar- ily by the gravitational effects of dark matter in the universe. In our work, seek to invert the weak lensing signal from 2D telescope images to reconstruct a 3D map of the universe's dark matter field. While inversion typically yields 2D pro- jection of the dark matter field, accurate 3D maps of the dark matter distribution are essential for localizing structures of interest and testing theories of our uni- verse. However, 3D inversion poses significant challenges. First, unlike standard 3D reconstruction that relies on multiple viewpoints, in this case, images are only observed from single viewpoint. This challenge can be partially addressed by observing how galaxy emitters throughout the volume are lensed. However, this leads to the second challenge: the shapes and exact locations of unlensed galaxies are unknown, and can only be estimated with a very large degree uncertainty. This introduces overwhelming amount of noise which nearly drowns out the lensing signal completely. Previous approaches tackle this by imposing strong assumptions about the structures in the volume. We instead propose a methodol- ogy using gravitationallyyconstiiii
neural field to flexibly model the contin- uous matter distribution. We take an sadllsss-b--yythhestt approach, optimizing the weights of the neural network through a fully differentiable physical forward model to reproduce the lensing signal present in image measurements. We show- case our method on simulations, including realistic simulated measurements of dark matter distributions that mimic data from upcoming telescope surveys. Our results show that our method can not only outperform previous methods, but im- portantly is also able to recover potentially surprising dark matter structures.
INTRODUCTION
According Einstein's theory of general relativity, massive objects in our universe bend their sur- rounding spacetime, causing nearby light to travel on curved trajectories. As a result, images of faraway galaxies as observed from Earth are warped due the continuous deflection of their light by intervening matter. This effect, called gravitational lensing, is a powerful probe of the distribu- tion of the massive structures underlying our universe, referred to colloquially as the cosmic web. Although the mass of visible matter can be directly measured by spectroscopic methods (Tinsley (1922). Roberts (1962)) dark matter does not appear to interact with electromagnetic radiation and thus cannot be observed directly. Nonetheless, gravitational lensing is an effect caused by the total matter distribution in the universe, both dark and luminous. Because dark matter constitutes around 85% of total matter by recent estimates (Aghanim et al.) (2020), gravitational lensing serves as an effective tool for estimating its distribution. In this work, we focus on the weak lensing regime, where light deflections are small to interpret for a single lensed galaxy. However, by observing how weak lensing subtly shears the images of galaxies in dense fields, we aim to reconstruct a map of the underlying matter.
Traditional weak lensing analysis is performed in 2D. Measured galaxy shears are interpolated to form single 2D shear map which is used to recover estimate of the 2D projected matter density

---

Published as a conference paper at ICLR 2025
over the sky. Many methods for 2D density reconstruction have been proposed in the past (Kaiser & Squires [1933), Lanusse et al. (2016)) However, amalysis from these maps is limited; 2D mapping techniques do not provide any information about the mass of structures their distance from Earth.
We seek to tackle the problem of 3D mass mapping; that is, inverting the weak lensing signal from 2D telescope images to obtain a 3D reconstruction of the dark matter field. Detection and local- ization of peaks in the 3D structure of dark matter is important for the study of its fundamental properties, many of which are currently unknown. For example, an open question in physics is whether or not dark matter emits trace amounts of gamma ray radiation (Bringmann & Weniger (2012);; however, it difficult discriminate between potential weak dark matter emission and the overwhelming background of gamma rays without attenuating to regions with peaks in dark matter density. In addition, characterizing the exact nature of of the primordial matter field at times close to the Big Bang requires studying the 3D structure of the current universe further than just statis- ties from the 2D field (Bartolo et al. (2004)). Thus, studying the 3D structure of dark matter offers exciting avenues for answering fundamental questions about our universe.
Recovering the 3D matter density field from projected cosmic shear measurements poses several challenges. First, since galaxies are observed from single viewing angle, reconstructing the 3D distribution from shear data is inherently ill-posed, meaning many different dark matter configura- tions could produce the same observational data. While observing galaxy emitters throughout the volume helps to alleviate this issue, it introduces an even bigger challenge: the true unlensed shapes of visible sources are largely unknown. Because weak lensing describes the change between the ob- served shape and this "intrinsic" shape, this uncertainty introduces a nearly overwhelming amount of noise. For example, for single galaxy, the noise from galaxy's intrinsie shape can be on average 10 to 100 times the magnitude of the lensing shear signal.
We propose an approach that leverages the physics of gravitational lensing to recover a continuous 3D matter field represented using neural field. In particular, we represent the spatially varying matter density field with coordinate-basse entiable lensing forward model to reproduce a given set of observed shear measurements. We show on simplified simulation of dark matter that our method not only accurately recovers 3D matter structures, but also surpasses the baseline in localizing these structures at their correct distances from Earth, referred to as redshift. Crucially, we also show that our method can accurately reconstruct non- non-Gaussian features not present in the simplified simulation where the baseline method struggles, which will prove essential for constraining fundamental theories of structure formation in the uni- verse. We believe that our work serves as a flexible and powerful approach that shows promise for extension to upcoming weak lensing telescope surveys.
neural network that we optimize through a fully differ-
2 RELATED WORK
2.1 DARK MATTER MASS MAPPING
Dark matter mass mapping is the inverse problem of reconstructing the distribution of dark matter in the universe from observations of its gravitational lensing effects on images of galaxies. In weak lensing, which is the focus of this paper, the effects of lensing are very small and can be understood as small rotations or changes in the axis ratios of approximately elliptical galaxy shapes. These two effects are described by the complex components of the shear 7 a 2D field describing the spatial distribution of e elliptical shape changes in galaxy images across the sky.
It was discovered by Kaiser & Squires) (1993) that a dense shear field can be analytically inverted to obtain 2D projected mass map of the sky. However, in practice we are limited in our ability to measure the full shear field by the visibility and density of galaxies the sky as well as uncertainty in galaxies' unlensed shapes. Thus the problem of obtaining the 2D projected matter density is in itself difficult. Many methods for solving this 2D inversion have been developed: Kaiser & Squires (1993) propose smoothing of the observed shear field before directly applying the analytical inversion. Regularized approaches such as Gaussian prior (Hrrowitz et al. (2019) and sparsity (Lanusse) et al. (016)) have also been used. However, all these methods are fundamentally limited in that 2D
mass maps do not give information on the exact mass or redshift of detected structures. Only by measuring galaxy redshifts to reason in 3D can we estimate these quantities.
2

---

Published as a conference paper at ICLR 2025
In the problem of 3D mass mapping seek to recover not only the projected mass density but its full 3D distribution. This amounts to solving an ios-p single-view tomography along the line of sight. Thus, some form of regularization is necessary for reconstruction. Previous work in 3D reconstruction from weak lensing has focused on methods that incorporate very strong priors on the spatial structure of the universe. [Simon et al. (2009) proposed Wiener filter method that combines Gaussian prior with an inverse variance filter to regularize a reconstruction, VanderPlas et al. (201T) use an inverse variance filter as well, but regularize by truncating the singular value decomposition of their reconstruction. While both methods introduce heavy smoothing of their reconstructions along the line of sight, [eoonadd et al. (2014) proposed method to achieve higher resolution in the line of sight direction by using sparsity prior. However, this method induces strong constraint, i.e. that the universe can be represented as a sparse sum of dark matter "halos", and can be susceptible to false peak detections along the line of sight. In addition, the optimal amount of sparsity is not clear, and different values can result in significantly different reconstructions.
Out of the aforementioned 3D reconstruction methods, the Wiener filtering method that employs a Gaussian prior based on the power spectrum of the dark matter field (Simon et al. (2009)) has been the most widely used in real data (Oguri al.] (2018), Simon et al. (2012)). Assuming that the underlying dark matter field is Gaussian, this method is favorable in that it is provably the optimal reconstruction method. However, there is currently plethora of evidence pointing towards the mon-Gaussiag nature of the universe's s density field, especially at low redshift regions close to Earth (Bartolo et al. (2004)). We show in this work that our method can reconstruct non-Gaussian structures more accurately, which will become critical for studying non-Gaussian features of the universe.
2.2 COORDINATE- BASED NEURAL FIELDS
Coordinate- based Neural Fields (Xie et al. (0022)) also commonly referred to as coordinate-based neural representations, have gained popularity in computer vision and graphics across a wide range of inverse problems. These models parameterize continuous three-dimensional vector scalar field with the weights of a multi-layer perceptron (MLP) that takes a spatial coordinate as input and outputs the value of the field at that coordinate. Typical approaches using neural fields formulate inverse problems as an optimization, directly using an underlying physical forward model to recover single solution field that matches set of observed measurements.
For ill-posed inverse problems where there are potentially many solution fields which fit a set of observed measurements, neural fields have been shown to provide a good implicit prior. In particular, MLPs favor learning lower spatial frequencies when optimized with gradient descent (Rahaman et al. (211 ) an effect known as spectral bias. [Tancik et al. (2020) conducted a theoretical analysis which showed that by prepending a positional encoding layer one can modify the network's 's bandwidth in tunable fashion. Thus, positionally encoded MLP implicitly imposes a smoothness prior on the final reconstruction as it will favor representing certain spatial frequencies of a signal. Neural fields with positional encoding have been shown to provide impressive results on wide range of ill- pll-po inverse problems, such medical imaging (Shen et al. (2022)) cryo-electron microscopy (Zhong cct al. (2021)) refractive field estimation (Zhao et al. (2024)), and black hole tomography (Levis 2022 Levis et al. (2024). Recently, /Zhao et al. (2024) presented an approach to 3D dark
ee al. matter mass mapping that optimizes a neural field model to fit to simulated weak lensing image measurements and showed how it led superior performance over traditional voxel grid. While this method is relevant to a larger range of gravitational lensing regimes, as it doesn't 't require lensed galaxies to be modeled as ellipses, it struggles when there is noise on the intrinsic shape of galaxies. In contrast, our work focuses on the noise- coise-dominated weak lensing regime and introduces a method that takes advantage of fitting directly to the estimated shear of elliptical galaxies.
3 METHODS
3.1 WEAK LENSING MEASUREMENTS
As light travels towards us from faraway galaxies, its path is continuously perturbed by the dark matter it encounters along the line of sight. In weak lensing, the degree of this perturbation is rel- atively minor, and results in small changes in the observed shapes of galaxies. To quantify these small changes, approximately elliptical galaxy shapes are described using complex ellipticity pa-
3

---

Published as a conference paper at ICLR 2025
Exaggerated
Pint
Cbbs
Cobs
Realistic
0.5
Cint
eabs
kinemmnce uncertainty
uncertainty
0.02
Figure 1: Weak Lensing Measurements. (Left) As light travels through the universe, it is lensed by intervening matter structures, causing slight changes between the unlensed intrinsic shape ein and the observed shape Cobs of galaxies. Measurements of these faint shape changes is useful probe of the 3D structure of dark matter. (Right) To quantify the effects of cosmic shearing, it is useful to consider galaxy shapes in the complex elliptical domain, where the components of the complex ellipticity describe the axis ratio and orientation angle of a given ellipse (Figure adapted from Schneider et al.] (2006)). In kinematic weak lensing, the detectable cosmic shear signal is outweighed by our uncertainty in a galaxy's intrinsic shape by more than an order of magnitude. In traditional weak lensing the uncertainty is more than two orders of magnitude greater than the shear signal. The lensing effects of intervening matter can be described by a shear y, which is approximately additive in the ellipticity domain. In practice, we combine many shear measurements from a dense field of galaxies to obtain a coherent signal from the underlying matter distribution.
rameterization e. Given approximately elliptical galaxy with axis ratio and orientation angle Φ, we can define the magnitude and phase of its ellipticity e as in Schneider et al. (2006):
1 1-P
lel
Le 20 20
1+r
Weak lensing measurements relate a galaxy's observed ellipticity cobs to its intrinsic ellipticity ennt which describes what its observed shape would be in the absence of lensing. In the regime of weak lensing, the lensing effects of intervening matter can be approximately described as a shear 7, which is additive in the ellipticity domain (Schneider et al. )2006):
Cobs
Cint %
An overview of shear measurements and the complex ellipticity can be found in Fig. In prac- tice, although the observed shape coss can be measured directly from an image, galaxy's 's intrinsic shape Cin much more difficult to estimate. In traditional weak lensing surveys, Cin is completely unknown (Schneider et al.) (2006)). More recently a technique called kinematic weak lensing (Huff) let al. (2013) has been proposed that leads an estimate of éim using spectral data. Nonetheless, even with kinematic estimates of the intrinsic shape, the noise on 7 will still typically be more than 10 times larger than the signal itself; this uncertainty is called the shape noise. To mitigate the effects of the shape noise, typical weak lensing surveys combine measurements from many galax- ies into galaxy catalog containing the sky position, redshift (distance), and observed and intrinsic shape estimates of millions of galaxies, which constitute the measurements for the 3D mass mapping problem.
3.2 FORWARD MODEL
The shear observed in a galaxy image arises from the cumulative lensing occurring along the line of sight between us and the galaxy. Because a uniform sheet of mass does not induce any gravitational shearing effect, the shear is measurement of the overdensity 8 = (p p)/p. where p is the continuous matter density field and p is mean density at a given distance from Earth.
To compute the shear from the matter overdensity field 5, first introduce the convergence (Schneider et al. (2006)), which for galaxy with sky position 0 at distance w is:
3462
n dw' (wwww- - w') 8(w' 0, w') a(w")
k(0, w) = Q8
(3)
20
4

---

Published as a conference paper at ICLR 2025
Deep Ensemble
Neural Density Field
Forward Model: P Projection Convolution
Measurement Loss
Bobs
iint
X = (0,.02 2)
ic (0)
(Lm)
Ymaas
M
Sshapa
Optimize via Gradient Descent
Figure 2: Proposed 3D Mass Mapping Pipeline. We model 3D matter overdensity field continuous function using a fully-connecter neural network. then differentiably compute cosmic shear measurements from given galaxy catalog through this overdensity field with physics-basee forward model to produce set of predicted shear measurements. Next, we solve for the weights of the neural network by minimizing data loss between the model and observed shear measurements plus a peysically-motivateee power spectrum regularization loss. We obtain a final reconstruction by taking the median from a deep ensemble of 100 independently initialized neural fields.
where Ho and Ωm are assumed cosmological constants and is the speed of light. The scale pa- rameter a(w) represents the known expansion of the universe that occurs as light travels towards us, and 8 is the matter overdensity. The 3D shear is obtained by convolving the convergence with complex kernel (Kaiser & Squires 1993)):
π {
1
700, w) = PK
de 0')n((', w), D(0) (0-)
(4)
Here, the angular position 0 01 + i0g is expressed in complex coordinates, and the asterisk represents complex conjugation.
Notably, the forward operators in Eqns. and are both linear with respect to the overdensity 6, so obtaining the overdensity from the shear measurements is linear inverse problem. Combining with Eqn. 2 gives the full formulation with the effects of the additive shape noise s:
(5) where Q and P the linear forward operators from Eqns. 3 and respectively, and éim is the estimated intrinsic ellipticity. Note that even in traditional weak lensing, where no estimate of the intrinsic shape available, we can ēint possible ellipses. In practice we calculate the forward model by discretizing the overdensity ò into of lens planes spaced throughout the volume. We then efficiently compute with summation, and the convolution P with the Fast Fourier Transform.
Eabs PQ1 +
~ N (Cint, (shape)s
and adjust shape account for the full range of
3.3 CONTINUOUS NEURAL REPRESENTATION
To solve the inverse problem of mass mapping, we take an by-ythsss approach. We model the continuous matter overdensity field ó with a neural field consisting of a coordinate-based Multi-Layer Perceptron (MLP), which is optimized to fit the observed measurements Cobs through the forward model described in Sec, 3.2 The MLP, parameterized by weights W, takes 3D spherical coordinates x (01, 02, w) as input and outputs scalar matter density:
(6) where (x) is positional encoding layer mapping each input coordinate to a set of Fourier basis functions with increasing frequencies. The positional encoding layer has been shown to control the bandwidth of the neural network's 's interpolation kernel (Tancik et al. (2020)), tunable via the frequency parameter T.. Higher levels of allow the MI.P to express higher spatial frequencies, while lower levels of 7, encourage learning lower spatial frequencies. In our experiments we use for angular coordinates and for the radial coordinate.
8(x) (F((((((x)),, Fx(x) = [sin(x), cos(x).......................... sin()) cos(2_
3.4 OPTIMIZATION
Because our network is fully continuous estimate of the matter overdensity field 8, we can simulate cosmic shear measurements that would have been observed had the galaxies in catalog been
5

---

Published as a conference paper at ICLR 2025
lensed by the density distribution represented by neural field. To do so, simply apply the forward operators P and Q detailed in Eqns. 3 to the neural overdensity estimate to obtain the modeled shear ÿ. We then optimize our network weights W to match the observed galaxy ellipticities Cobs the specified level of shape noise, minimizing a x² data loss.
We regularize our reconstruction with an additional loss term equal to the mean squared error be- tween the lens plane power spectra of the ground truth (P) and the reconstruction (P). Inclusion of this term helps us to prevent overfitting the noise and produce physically realistic reconstructions. The theoretical form of the power spectrum P can be analytically calculated and has been used to characterize the covariance of Gaussian prior in previous work; however, we expect the ground truth field to contain non-Gaussian features in practice.
The full loss function is:
(m) w
(mm) 2
((w) M
Σ
A N
(W (())
Shupe
where and index over M galaxy catalog entries and N lens planes, and E is the frequency vector of a lensplane's ' 2D Fourier transform. A diagram of our model's optimization is shown in Fig. 2
3.5 NEURAL DEEP ENSEMBLE
Because the forward operator Q is weighted line-of-sight projection, solving the inverse problem of 3D mass mapping is akin to computing a tomography from a single line of sight. This, in combi- nation with the large degree of noise in the lensing signal, makes the inversion ill-posed; there exist infinitely many mass fields that could produce the same set of shear observations. In particular, two neural networks with different sets of randomly initialized weights could potentially converge to dif- ferent reconstructions, both of which fit a given observed galaxy catalog equally well. We find that constructing an ensemble estimator by taking the median values over many neural reconstructions gives robust reconstruction results and mitigates artifacts from individual samples"
In our experiments, we use deep ensemble of 100 fully connected MLPs each with 4 layers, where each layer is 256 units wide with the ReLU activation function. The network output is passed through the sigmoid activation function followed by a final single node linear layer. We optimize the network weights by minimizing the loss in Eqn. / via gradient descent using the Adam optimizer (Kingmal (2014)), with exponential learning rate decay from le-4 to 5e-6 over 100K iterations. All code is implemented in JAX and will be made publicly available.
4 EXPERIMENTS
4.1 COSMIC SHEAR SIMULATIONS
To simulate ground truth dark matter fields we used the low resolution Particle-Mesh N-body solver JaxPM (Initiative (202I)) The shear signal was simulated over a field of 5 x 5 with a resolution of 75 x 75 pixels. For the kinematic weak lensing experiment in Sec. 4.2 we simulate shape noise in line with current estimates for the instrument capabilities of the Roman Space Telescope Xu et al. (2023)) with a galaxy number density of 7iaal = 4 aremin -1 and a shape noise level of Ssape 0.035. For the traditional weak lensing survey in Sec. 4.3 we assume no estimation has been done on the intrinsic shape of a denser field of galaxies, corresponding to a shape noise level of shape 0.25 and galaxy number density of Ngal = 30 aremin 1 Finally, we assume a realistically distributed galaxy sample as in Leonard et al. (2014) of 360,000 galaxies for kinematic weak lensing and 2,700,000 galaxies in the traditional weak lensing experiment; more details are in the appendix.
Throughout this section, we will use the term redshift, which should be understood as interchange- able with distance in the radial direction; for reference, redshift = = corresponds to Earth's 's loca- tion, while redshift of 2 2 corresponds roughly to distance of 15 billion light years away in the present day. For our experiments, we divide the lensing volume into 18 equally spaced lens planes from redshift 2 = 0 to 2. Although simulated galaxies are lensed throughout the whole volume, the SNR of the reconstruction decreases with distance from the sensor, as most of the lensing is done by
Interestingly, although the individual members of the so-called deep ensemble do not correspond to sam- ples from Bayesian posterior, deep ensembles have been shown to yield predictive uncertainty estimates ( akkshminarayanan al. 2017)).

---

Published as conference paper at ICLR 2025
Ground Truth
GT (Blur)
Ours
GT (Blur)
Wiener
0.72
0.57
Region 0
3.13
Region
Region
1.58
1.90
Region
Region
Region
1.52
Region 2
1.96
Region Region
Region Region
1.47
Region 3
1.80
-0.99
Figure 3: Kinematic weak lensing result: We present 3D reconstruction results of realistic lensing volume derived from an N-body simulation. Results includes 12 lens planes from redshift = = 0 to average every 3 adjacent lens planes form
1. To concisely visualize reconstructions, lensing regions, increasing in distance away from the sensor. Ground truth (GT) lensing regions are shown on the far left. For visual comparison, we blur the ground truth volume with Gaussian filter in the lfne-of-sight direction that maximizes the cross-correlation with each reconstructed slice. We then visualize this optimally blurred ground truth volume for each reconstruction, along with the average line-of-sight standard deviation used for each region (top right corner). Our reconstruction corresponds well with GT volume with significantly lower radial blurring, especially at lower redshift, showing that our method is less susceptible to smearing along the line of sight.
matter in the front half. Thus, only present the reconstructed volume up redshift 1; this is a common practice in 3D mass mapping (Simon et al. 2009 Leonard et al. (0014).
4.2 LARGE-SCALE STRUCTURE RECOVERY WITH KINEMATIC WEAK LENSING
We first use our approach to perform 3D mass reconstruction realistic kinematic weak lensing measurements of simulated dark matter fields. Results are shown in Fig. We optimize a neural field ensemble as described in Sec. 3.5 taking the median value over each neural reconstruction at sampled points. As baseline, we compare against the Wiener filter described in [Simon et al. (209). This method has been the most widely used real data (Simon et al. (2012). (Oguni et al. 2018 and employs a Gaussian prior that requires knowledge of each lens plane power spectrum.
We present reconstruction results of our overdensity volume 12 lens planes from redshift to = = 1. To visualize our reconstruction, we average every 3 adjacent lens planes to form lensing regions. Due large amounts of shape noise, compare each reconstruction of the ground truth volume. For each reconstructed lens plane we blur the ground truth volume with Gaussian filter using a 2px in the transverse (angular) directions, and a in the radial = di- rection chosen to maximize cross-corred coos-correlltion with each reconstruction. This produces two volumes, optimally blurred to match our reconstruction and the Wiener filtered reconstruction. Additionally, we use the optimal radial blur as a metric for the amount of radial smearing in each reconstruction; a reconstruction with large amounts of smearing along the line of sight would correlate most highly with the ground truth blurred with wide Gaussian filter in the direction.
blurred version
quantitative analysis of reconstruction performance is shown Table Downstream applications of 3D mass mapping include the study of mass peaks (areas with high overdensity) and voids (areas with low overdensity); thus, it is important for our reconstruction methods to accurately discriminate between the two. To evaluate the reconstruction quality of both methods, cross-cones cross-correlation metric without mean subtraction, which harshly penalizes predicting overdensities of the wrong sign. We find that method outperforms the Wiener filter on both the unblurred and
the normalized

---

Published as conference paper at ICLR 2025
Ground Truth
GT (Blur)
Ours
GT (Blur)
Wiener
0.72
0.82
Region 0
2.51
Region
Region
1.78
2.06
Region
Region
Region
2.27
Region 2
2.78
Region Region
Region Region
2.78
Region 3
2.69
-0.99
Figure 4: Traditional weak lensing result: We perform 3D reconstruction simulated traditional weak lensing survey assuming no information on the intrinsic shapes of galaxies. As in Fig. 3 the ground truth volume and setup, dividing our volume into four redshift regions and generating blurred ground truth volumes for visual comparison. The optimal blur level for each region and method indicated in the top right corner of each blurred ground truth image. Our reconstruction has less radial blurring than the Wiener filter baseline used in previous weak lensing surveys, and has higher correlation with the ground truth volumes.
optimally blurred ground truth volumes. In addition, we find that the optimally blurred ground truth volume for our reconstruction corresponds Gaussian filter with significantly lower width the radial direction, suggesting that our method better localizes structures along the line of sight; this is especially evident at lower redshifts of the reconstruction.
Table 1: Quantitative Analysis: We report the normalized cross-correlation between the reconstruc- tions and the ground truth, both with and without Gaussian blurring. Cross-correlation values range from -1 1 to +1, where higher values signify better reconstruction performance.
Kinematic WL
Traditional WL
MNIST Reconstruction
CC
CC (Blur)
CC
CC (Blur)
CC
CC (Blur)
Ours
0.60
0.80
Ours
0.52
0.77
Ours
0.61
0.86
Wiener 0.39
0.70
Wiener
0.34
0.61
Wiener
0.52
0.68
4.3
LARGE-SCALE STRUCTURE RECOVERY WITH TRADITIONAL WEAK LENSING
Although kinematic weak lensing measurements show great promise for enabling accurate 3D dark matter reconstructions, kinematic surveys are still in the planning stage (Xu et al. (2023)). Tradi- tional weak lensing surveys do not estimate galaxies' intrinsic shapes, and instead use the fact that they are distributed isotropically and are circular on average. As such, they are subject to much higher levels of shape noise.
In this section we simulate shear measurements with shape noise the amount expected from tra- ditional weak lensing survey. Although traditional weak lensing surveys detect much denser field of galaxies than expected from kinematic surveys, the statistical effect of shape noise only decays of /ggal (Schneider al. (2006)); thus, the combined measure- ments still contain more noise in aggregate (roughly 2,6x more noise than kinematic). Quantitative results as described Sec. 4.2 can found Table [ We find that our method again outperforms the Wiener filter baseline and better localizes structures radially, especially at lower redshifts. Thus, we believe that our method could be powerful tool for analyzing existing weak lensing surveys.
with the galaxy density
4.4 RECONSTRUCTION OF NON-GAUSSIAN STRUCTURES
As our reconstruction in Sec. 4.2 was performed on low-resolution N- N-Body simulation, the ground truth volume lacks non-Gaussian features such sharp edges peaks, and was well-suited to
8

---

Published as conference paper at ICLR 2025
Ground Truth
GT (Blurred)
Mean Recon
GT (Blurred)
Wiener Filter
0.50
0.98
1.71
Plane 0
Plane 0
Plane 0
Plane 1
Plane 1
0.73
Plane 1
1.06
Plane 2
Plane 2
0.73
Plane 2
2.04
2
Plane 3
1.22
0.90
Plane
Plane
-0.26
Figure 5: Sensitivity to Non-Gaussianity: In this experiment, we perform reconstruction on toy lensing field of 4 lens planes, each with an MNIST digit through in if its corners. Due to the non-Gaussian nature of these images, the Wiener filter which uses Gaussian prior struggles to reconstruct the lens planes accurately. We blur the ground truth volume to match each reconstructed lensplane described in Sec. 4.2 and report the optimal blur (px) in each corner; the Wiener filter exhibits significantly more blurring in low redshifts.
the of Gaussian prior. However, recent observations Komatsu et al.) (2009) well state-of- the-art simulations Pillepich et al. 2018 reveal that some of the most fascinating and scientifically important structures in the universe are highly non-Gaussian. These structures are ill-suited for reconstruction with a Wiener filter, which uses a strong Gaussian prior.
In this experiment consider toy example of 4 lens planes, each with an MNIST digit 0 through 3 in single quadrant, at 4 equally spaced distances from redshift = O to 1. We simulate shear measurements with the same galaxy parameters as in Sec. 4.2, where galaxies throughout the volume are lensed only by these 4 lens planes. While the structures of each digit are apparent in both reconstructions, the Wiener filter suffers from heavy amounts of artifacts due the non-Gaussian nature of the true lens planes. In addition, our method is better able to localize the central locations of each digit in its reconstructions; for the Wiener filter, traces of the digit can be found prominently in the incorrect lens planes. As can be seen in Table | our method outperforms the Wiener filter quantitatively on both blurred and unblurred ground truth volumes.
CONCLUSION
This work presents method for recovering the 3D cosmic web using neural field constrained by galaxy shear measurements. First, we demonstrate accurate reconstruction results on simulations of dark matter that outperform the leading baseline method for real data. Second, we show how our method addresses key limitations of the baseline, particularly in capturing non-Gaussian structures within the dark matter field. These structures, which are among the most intriguing and scientifically significant, often exhibit highly non-Gaussian characteristics that are obscured by methods con- strained by Gaussian prior. Furthermore, the continuous and differentiable neural representation underlying model offers flexible framework for future advancements, such incorporating stronger physically motivated constraints, non din-discrtte galaxy distributions, and more sophisticated spatial priors. Our method offers flexible yet still powerful solution for 3D dark matter mass map- ping, potentially enabling insights into structure formation in the universe and the nature of dark matter itself, two fundamental questions at the heart of physics.
Although outputs from these simulations are publicly available, simulation of weak lensing shear measure- ments from these volumes prohibitively computationally expensive due exponentially greater partiele compared to the simplified N-Body simulations shown in this work.

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by NSF award 2048237, Google, a Carver Mead New Adventure Fund Award, Stanback innovation Fund Award, Amazon AI4Science Discovery Award, and a Caltech S2I Award. A.L. was supported by the Natural Sciences & Engineering Research Council Canada (NSERC). We thank Francois Lanusse for his helpful discussions regarding weak lensing and pro- viding us with guidance on the JaxPM code, Carolina Cuesto-Lazaro for her input on dark matter mapping applications, and Olivier Doré for helpful discussions about cosmology.
REFERENCES
Nabila Aghanim, Yashar Akrami, Mark Ashdown, Jonathan Aumont, Carlo Baccigalupi, Mario Ballardini, Anthony J Banday, RB Barreiro, Nicola Bartolo, S Basak, al. Planck 2018 results- vi. cosmological parameters. Astronomy & Asstrophyscc, 641:A6, 2020.
Nicola Bartolo, Eiichrro Komatsu, Sabino Matarrese, and Antonio Riotto. Non-gaussianity from inflation: Theory and observations. Physics Reports, 402(3-4):003-266 2004.
Torsten Bringmann and Christoph Weniger. Gamma ray signals from dark matter: Concepts, status and prospects. Physics of the Dark Universe, 1(1-2):194-217, 2012.
Benjamin Horowitz, Uros Seljak, and Grigor Aslanyan. Efficient optimal reconstruction of linear fields and bowe-s from cosmological data. Journal of Cosmology and Astroparticle Physics, 2019(10):035, 2019.
Eric M Huff, Elisabeth Krause, Tim Eifler, Xiao Fang, Matthew R George, and David Schlegel. Cosmic shear without shape noise. arXiv preprint arXiv:1311.148) 2013.
Differentiable Universe Initiative.
Jaxpm. httpa://https://gtt
com/
JaXPM/ 20211.
Nick Kaiser and Gordon Squires. Mapping the dark matter with weak gravitational lensing. Astra- physical Journal, Part I (ISSN 0004-637X), vol. 404, no. p. 441-450,, 404:441 1993.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
Eiichiro Komatsu, Niayesh Afshordi, Nicola Bartolo, Daniel Baumann, JR Bond, EI Buchbinder, Christian T Byrnes, Xingang Chen, Daniel JH Chung, Asantha Cooray, al. Non-gaussianity as probe of the physics of the primordial universe and the astrophysics of the low redshift universe. arXiv preprint 4rX.4:0902.4759, 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances neural information processing systems, 30, 2017.
Francois Lanusse, J-L. Starck, Adrienne Leonard, and Sandrine Pires. High resolution weak lensing mass mapping combining shear and flexion. Astronomy & Astrophysics, 591:A2, 2016.
Adrienne Leonard, François Lanusse, and Jean-Luc Starck. Glimpse: accurate 3d weak lensing reconstructions using sparsity. Monthly Notices of the Royal Astronomical Society, 440(2) 1294, 2014.
Aviad Levis, Pratul P Srinivasan, Andrew A Chael, Ren Ng, and Katherine L Bouman. Gravita- tionally lensed black hole emission tomography. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 19841-19850, 2022.
Aviad Levis, Andrew A Chael, Katherine L Bouman, Maciek Wielgus, and Pratul P Srinivasan. Orbital polarimetric tomography a flare near the sagittarius a supermassive black hole. Nature Astronomy, pp. 1-9, 2024.
10

---

Published as a conference paper at ICLR 2025
Masamune Oguri, Satoshi Miyazaki, Chiaki Hikage, Rachel Mandelbaum, Yousuke Utsumi, Hi- ronao Miyatake, Masahiro Takada, Robert Armstrong, James Bosch, Yutaka Komiyama, et al. Two-and three- diceensiinnii wide-field weak lensing mass maps from the hyper suprime-cam sub- aru strategic program s16a data. Publications of the Astronomical Society of Japan, 70(SP1):S26, 2018.
Annalisa Pillepich, Volker Springel, Dylan Nelson, Shy Genel, Jill Naiman, Rüdiger Pakmor, Lars Hernquist, Paul Torey, Mark Vogelsberger, Rainer Weinberger, et al. Simulating galaxy for- mation with the illustristng model. Monthly Notices of the Royal Astronomical Society, 473(3): 4077-4106, 2018
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International conference on machine learning, pp- 5301-5310. PMLR, 2019.
Morton S Roberts. The neutral hydrogen content of late-type spiral galaxies. Astronomical Journal, Vol. 67, p. 437-446 (1962), 67:437-446, 1962.
Peter Schneider, Christopher Kochanek, and Joachim Wambsganss. Gravitational lensing: strong, weak and micro: Fue advanced course 33, volume 33. Springer Science & Business Media, 2006.
Liyue Shen, John Pauly, and Lei Xing. Nerp: implicit neural representation learning with prior embedding for sparsely sampled image reconstruction. IEEE Transactions on Neural Networks and Learning Systems, 35(1):770-782, 2022.
P Simon, Catherine Heymans, Schrabback, Andy N Taylor, Meghan Gray, Ludovic Van Waer- beke, Christian Wolf, David Bacon, Marco Barden, A Böhm, et al. Spatial matter density mapping of the stages abell a901/2 supercluster field with 3d lensing. Monthly Notices of the Royal Astro- nomical Society, (19(())998--0166 2012.
Patrick Simon, AN Taylor, and Jan Hartlap. Unfolding the matter distribution using three- dimensional weak gravitational lensing. Monthly Notices of the Royal Astronomical Society, 399 (1):48 68, 2009.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
Beatrice M Tinsley. Galactic evolution. Astronomy and Astrophysics, Vol. 20, p. 383 (1972), 20: 383, 1972.
JT VanderPlas, AJ Connolly, Bhuvnesh Jain, and Mike Jarvis. Three-dimennional reconstruction of the density field: An svd approach weak-lensing tomography. The Astrophysical Journal, 727 (2):118, 2011.
Yiheng Xie, Towaki Takikawa, Shunsuke Saito, Or Litany, Shiqin Yan, Numair Khan, Federico Tombari, James Tompkin, Vincent Sitzmann, and Srinath Sridhar. Neural fields in visual comput- ing and beyond. In Computer Graphics Forum, volume 41, pp. 641-676. Wiley Online Library, 2022.
Jiachuan Xu, Tim Eifler, Eric Huff, RS Pranjal, Hung-Jin Huang, Spencer Everett, and Elisabeth Krause. Kinematic lensing with the roman space telescope. Monthly Notices of the Royal Astro- nomical Society, 519(2):2535-255 25511, 2023.
Brandon Zhao, Aviad Levis, Liam Connor, Pratul P Srinivasan, and Katherine L Bouman. Single view refractive index tomography with neural fields. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 25358-25367 2024.
Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis. Cryodrgn: reconstruction of heterogeneous cryo-em structures using neural networks. Nature methods, 1862):176-185, 2021.
11

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 N-BODY SIMULATION PARAMETERS
All measurements were simulated in a flat ACDM cosmology with Ω2m 0.3075, ΩĄ 0.6925, and Ho = 100 (ids/(h"MMpoo We assume that galaxy redshifts are sampled from a distribution with probability density function given by n() = exp (=(=/z)))). where we take 20 1/1.4, = and B 1.5.
used a multiple lensplane approximation of the shear, dividing the volume into 18 equally spaced (in comoving distance) lens planes from redshift 2 = 0 to 2. Reconstruction results are shown for the first lensplanes from redshift to 1. Radial distances to the center of each reconstructed lensplane are given in Table 2 as well as the redshift with respect to the cosmological parameters used in our experiments:
For simulated cosmic shear experiments (Figs
Table 2: Reconstructed Lensplane Redshifts and Distances
Lensplane
Comoving Distance (Mpe h h
Redshift
Present Distance (10' Light Years)
100
0.03
0.48
300
0.10
1.44
3
500
0.17
2.41
700
0.25
3.37
5
900
0.33
4.33
6
1100
0,41
5.30
1300
0.49
6.26
8
1500
0.58
7.22
1700
0.68
8.19
10
1900
0.78
9.15
2100
0.88
10.11
12
2300
1.00
11.07
The MNIST experiment (Fig. 5) was done with 4 equally spaced (in redshift) lensplanes. A table of distances center of each of these lensplanes can be found in Table
Table 3: Reconstructed Lensplane Redshifts and Distances
Lensplane
Comoving Distance (Mpe h 1)
Redshift
Present Distance (10" Light Years)
705
0.25
3.39
2
1319
0.5
6.35
3
1848
0.75
8.90
4
2303
1.0
11.09
A.2
NEURAL ENSEMBLE SAMPLES
In this section we visualize samples from the Neural Field ensemble used for reconstruction in Sec. 4.2 Due to the ill-posed nature of the reconstruction problem, individual samples can fit the measurement data but vary in structure or contain reconstruction artifacts. However, taking the ensemble mean or median yields more robust reconstruction that matches well with the ground truth volume.
In addition, we empirically find statistics from the ensemble can be used to perform uncertainty quantification. Results are shown in Fig. 7 Our ensemble can give well-calibrated uncertainty estimates for most reconstruction points, except for outliers at either tail.
A.3 ADDITIONAL RESULT FOR LARGE-SCALE STRUCTURE RECOVERY WITH KINEMATIC WEAK LENSING
In this section we show a reconstruction experiment with the same simulation parameters as pre- sented in Sec. 4.2 of the main paper, but simulated volume initialized with different random
12

---

Published as a conference paper at ICLR 2025
Table Cross correlation for the kinematic weak lensing experiment presented in Fig.
CC
CC (Blur)
Ours
0.52
0.75
Wiener Filter
0.39
0.64
Ground Truth
Ours (Median) Sample
Sample
Sample
Sample
0.72
Region
Region 1
Region 2
Region 3
0.69
Figure 6: Neural ensemble samples: we visualize samples from the neural ensemble used for re- construction in Sec. (.2) Individual samples may fit the measurement data equally well but vary in structure or contain artifacts due to the ill-posed nature the reconstruction. However, taking the ensemble median gives robust reconstruction which matches well with the ground truth volume.
seed. We show similar findings as in the previous experiment; our method correlates more strongly with both blurred and unblurred ground truth volumes. Moreover, our method exhibits smearing in the radial direction of reconstructed structures at low redshifts. Reconstruction results are shown Fig. and quantitative results in Table
A.4 FULL RECONSTRUCTION VOLUME OF LARGE-SCALE STRUCTURE RECOVERY WITH KINEMATIC WEAK LENSING
In this section we present reconstruction results corresponding to Sec. 4.2 of the main paper. While in Fig. 3 we average lensplanes and optimal blur levels for sets of adjacent lensplanes for visual brevity, we present the same results in Fig. 9 without averaging.
13

---

Published as a conference paper at ICLR 2025
Deep Ensemble Calibration
0.45
0.45
Prediction Errors
Proportion of Test Points
0.40
y=x
0.40
0.30
0.30
Absolute
0.25
0.25
Proportion
0.15
0.15
0.05
0.05
0.00
0.00
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.40
0.45
Ensemble Std
Figure 7: Ensemble Uncertainty Estimates. Although our current ensemble representation doesn't t mathematically correspond to a Bayesian posterior, find that our current method already provides meaningful uncertainty estimates. We present an uncertainty calibration plot for the reconstruction ensemble corresponding Fig. 3 There is a strong correlation between the ensemble standard deviation and reconstruction error, except for few outlier regions. In regions with high ensemble variance we find our model is underconfident. In regions with low ensemble variance our uncertainty estimates are slightly overconfident, but the error is still well within 2 standard deviations.
Ground Truth
GT (Blur)
Ours
GT (Blur) 3.10
Wiener
0.57
Region 0
1.71
Region
Region
Region 1
Region 1
2.23
Region 1
2.80
Region 2
Region 2
1.20
3.16
Region
Region 3
2.45
1.74
Region
Region
00.88
Figure 8: Large-scale structure recovery for an additional simulated kinematic weak lensing survey. We present reconstruction results for mock kinematic lensing survey with the same parameters as in the main paper, only where the N-Body simulation is initialized with a different random seed. Our method again correlates more strongly with both blurred and unblurred ground truth volumes, and is subject to significantly less radial smearing at low redshifts.
14

---

Published as a conference paper at ICLR 2025
Ground Truth
GT (Blur) Ours (Median)
GT (Blur)
Wiener
0.19
Region 0
4.00
Region
Region
Region 1
Region 1
0.89
Region 1
4.00
Region 2
0.78
Region 3 Region 2
1.39
Region
Region 3
1.59
1.68
Region
0.72
Region 4
Region 4
1.68
Region 4
2.02
Region 5
Region 5
1.50
Region 5
1.97
Region 6
Region 6
1.25
Region 6
1.63
Region 7
Region 7
1.62
Region 7
2.09
0.69
Region 8
Region 8
1.72
Region 8
2.10
Region 9
Region 9
1.33
Region 9
1.77
Region 11 Region 10
Region 10
1.31
Region 10
1.68
Region 11
1.80
Region 11
1.88
Figure 9: Full Reconstruction volume for Kinematic WL. Reconstruction results in Fig, are shown
with sets of adjacent lensplanes averaged. This figure presents the full set 12 lensplanes without averaging.
15


Reviews:
[{"review_id": "e0trP7PN38", "reviewer": "Reviewer_p3hC", "summary": "***Late review: Apologies my review is in one day late. I will aim to engage as quickly as needed to makeup.***\n\nThis very interesting paper considers a very important application, although using some well-established tools in machine learning and computer vision. The paper focuses on the difficult weak-lensing regime, where deflections are too weak to be interpretable reliably for a single galaxy. The idea is that attempting to understand how weak lensing alters (or shears) images of galaxies in dense fields, one could reconstruct a density map of the underlying matter. This is a difficult problem as the matter density is usually analyzed in 2D: the measured shear map is used to recover a projected 2D estimate. The paper aims to solve the inverse problem i.e. to obtain a 3D reconstruction of the dark matter field from 2D images. From a modeling perspective, this becomes challenging as galaxies are observed from a single view. Further, unlensed shapes of known visible sources are themselves not fully understood, introducing a large amount of uncertainty. The main idea of the paper is to use the underlying physics of gravitational lensing to recover a continuous 3D field, with a focus on also capturing non-Gaussian features, which are not easy with traditional methods, which use strong priors. The idea is to use coordinate-based neural fields augmented with the underlying physics (clearly described in figures 1 and 2), with the total loss function described in equation 7. The experiments, over simulated data, are able to show that the method has promise in both reconstructing the 3D matter field and also non-Gaussian features.", "questions": "-- Could the authors elaborate on the training protocols and how difficult it was to get stabilized training? From the writeup, it is not clear how easy it is to get the model to work.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-- The paper attacks a problem of quite some importance. While the machine learning used is standard, its combination with the underlying physics allows the authors to obtain a promising solution, although it is only tested on simulations (standard for this area of physics at the moment). \n\n-- The paper is well-written, well-structured and easy to follow. While I am not an expert in this specific application, I have worked in some adjacent applications. I think the results support the underlying proposal and motivation well, and are creative in combining standard ML tools with physics-based models.", "weaknesses": "-- More of a question: Could the authors elaborate in the paper why are strong Gaussian priors usually made? This will not be obvious to a standard ML audience. I see that it could make sense from the perspective of filtering. But I am not sure I understand why this should drop out from the physics models. I would expect the features to be highly non-Gaussian. I find this confusing whenever I venture to look at papers in this area.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6wv5IlMdDF", "reviewer": "Reviewer_x8ZD", "summary": "The paper provides a new method for recovering the 3D cosmic web from galaxy weak lensing signal.\nThe proposed approach uses a coordinate-based neural field method with positional encoding. (Equation 6).\nThe reconstruction results are demonstrated in a series of experiments.", "questions": "* I would like to know some more details about the number of CPUs/GPUs used for training, memory requirements, and the training and inference time of the proposed method.\n\n* I'm also curious about the next steps. After the 3D cosmic web estimation step is done, what are the next science questions that can be answered with the new results?", "limitations": "", "rating": 3, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* The paper proposes a new method for an important physics problem.\n* The paper is well-written; I enjoyed reading it.\n* The physics background is exceptionally well-presented, and even non-experts can understand the main scientific ideas.\n* The proposed method works better than the leading baseline.", "weaknesses": "I think this is a great paper, but in my opinion, ICLR might not be the best venue for this.\nThe new machine learning contributions are quite limited: neural field models and positional encoding have been around for some time, and I didn't find significantly new machine learning contributions in the paper.\n\nNonetheless, the proposed algorithm works well, and the physics application, estimating the 3D cosmic web, is a very important problem.\nSince the most significant contributions of the paper are in astrophysics, I think an astrophysics or cosmology journal would be a better venue for this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k7NIlXnEsR", "reviewer": "Reviewer_gGfb", "summary": "The article considers the task of reconstructing a 3D map of the universe from weak gravitational lensing obtained in 2D signals. The methods consider first to encode the overdensities using a neural field based on the spherical coordinate and its Fourier decomposition.  The target density field is used for two purposes: (i) to match the power spectrum of the dataset under consideration and (ii) to match the shear observed (after using the transformation from overdensity to shear) in the data. The method is tested on simulated data mimicking the instrumental observations from an N-body simulation. The results show that the method seems better than using Wiener filter for the same task.", "questions": "- it might be useful to related the various red-shift to a number of years, it is done for z=0 and z=2 but not for z=1. \n- The caption of Fig. 3 says: \"results includes 12 lens planes from redshift z=0 to z=1\", can the authors specify the precise value of the considered z ?\n- In addition, the authors say that they average over 3 lens planes. Why ? to which values of z does that correspond ?\n- The ground truth images are blurred with a factor that maximize the cross-correlation with the reconstructed picture. It seems a bit weird to do so, it might bias the reader to believe that a reconstruction is better than what it should be.\n- The cross-correlation between the reconstructed signal and the true one seems quite small... while we can agree that it is even worse using the Wiener filter, I have the impression that it is working very-well.", "limitations": "", "rating": 3, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors apply ML techniques to a dataset that is different from the usual benchmarks.", "weaknesses": "My overall impression is that we have just an application of ML methods to a cosmological problem, which neither\n- make relevant improvement on the side of ML (unless I'm mistaking, it takes out-of-the-box methods)\n- make a strong advance in the considered field\nI would like the authors to comment on how their work is particularly timely. I'm also thinking that the Wiener filter is not very modern. By quickly looking on the web, we can find more elaborated techniques:\n * sparsity prior: https://www.aanda.org/articles/aa/pdf/2021/05/aa39451-20.pdf, https://arxiv.org/pdf/1801.08945\n * wavelet: https://www.aanda.org/articles/aa/pdf/2006/21/aa2997-05.pdf\nAnd I can imagine there are many others. This work does not provide any comparison to these methods.\n\nThe last section that investigate non-Gaussian structure is a bit of a mystery to me. Does the application of the method to MNIST (whatever it means) should be considered of a test of something precise ? In addition,", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KWKwFclFtF", "reviewer": "Reviewer_pR5T", "summary": "The paper presents a method for reconstructing the underlying 3D dark matter distribution from weak gravitational lensing observations (specifically, shape distortions of galaxies). This is an important problem in cosmology, and is of contemporary relevance given the number of astronomical surveys which will be measuring galaxy shapes. The overall approach is to model the dark matter distribution implicitly through a neural network (NeRF-like or implicit neural representation), and use a differentiable forward model that models the observed shear field. The authors compare the method to a few traditional approaches for mass mapping, finding it to especially perform well in reconstructing non-Gaussian structures.", "questions": "- The authors use an ensembling approach, taking the median to get a point estimate on the reconstructed mass map. Much of the ongoing literature/effort is focused on _probabilistic_ approaches to mass mapping. The authors mention in a footnote that ensembling can lead to good predictive uncertainties -- can it be expected that this can be used to yield a calibrated distribution over plausible mass maps? If not, are there alternative approaches to this incorporating neural fields, e.g. variational or Bayesian approaches?\n- What is the role of bandwidth $L$ and the induced spectral bias in the study? How does the induced prior compare with that assumed in traditional and model ML-based approaches?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Novel approach and good technical execution:** Using an implicit representation to model the underlying dark matter field along with a differentiable forward model is a very neat but challenging idea, and the authors set it up and executed it very well. \n- **Paper straddling AI and fundamental physics:** This is a bit of a meta point, but AI + physics (in particular a sub-field like cosmology) is not particularly mainstream at AI conference so far, in contrast to e.g. AI + bio. This makes presenting a paper like this while being faithful to the domain science particularly challenging, and the authors do a good job in this respect.\n- **Simulation realism:** The authors use a realistic set of tools from the cosmology literature, e.g. JaxPM in the differentiable forward model, making it a contribution beyond a simple proof-of-principle. \n- **Baseline comparisons:** Comparisons with baselines are well-presented, and highlight the advantage of the method in particular regimes (e.g. in recovering non-Gaussian structures).", "weaknesses": "- **Context in current literature and state of the field:** There exists substantial literature in neural approaches to mass mapping; e.g., https://arxiv.org/abs/2201.05561. While the authors compare to traditional approaches like Kaiser-Squires (KS) and Wiener filter, they do not contextualize their work in the more recent ML-based approaches to weak lensing mass mapping. Another example is https://arxiv.org/abs/2206.14820, which uses an implicit neural representation to perform *strong* lensing reconstruction; high-level comparison with existing literature could be improved.\n- **Hyperparameter choices:** The authors mention specific hyperparameter choices, e.g. bandwidth of positional encodings L=2 and 5, without further justification. This should be expected to have a significant impact on the results, as it controls the spectral biases of the implicit representation, and should be further expanded upon, possibly including ablations.\n- **Role of induced prior:** Traditional approaches typically have a regularization mechanism, often explicit, which allows for mass reconstruction (since the inverse problem is fundamentally ill-posed). In this study, the authors mention that \"neural fields have been shown to provide a good implicit prior\"; while true in the real-world setting of scene reconstruction, scientific problems are inherently different in nature, and it is not clear a-priori whether the implicit prior induced by the neural field is a good one. The authors approach this empirically, however further understanding of the role of induced priors is necessary for downstream science from the mass maps.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)