Revealing the 3D Cosmic Web through Gravitationally Constrained Neural Fields

Authors: Brandon Zhao, Aviad Levis, Liam Connor, Pratul P. Srinivasan, Katherine Bouman
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: computational imaging, signal processing, inverse problems, astrophysics, cosmology, neural fields, machine learning for physical sciences

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.

Decision: accept (poster)
Reviewer scores: 8, 3, 3, 6

Reviews and discussion:
<thread>
creation date:
1730267744690

thread_type:
official_review

<thread_object>
creator: reviewer_pr5t

<content>
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.

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. 
- **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.
- **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. 
- **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.
- **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.
- **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.

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?
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730285952415

thread_type:
official_review

<thread_object>
creator: reviewer_ggfb

<content>
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.

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
- make relevant improvement on the side of ML (unless I'm mistaking, it takes out-of-the-box methods)
- make a strong advance in the considered field
I 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:
 * sparsity prior: https://www.aanda.org/articles/aa/pdf/2021/05/aa39451-20.pdf, https://arxiv.org/pdf/1801.08945
 * wavelet: https://www.aanda.org/articles/aa/pdf/2006/21/aa2997-05.pdf
And I can imagine there are many others. This work does not provide any comparison to these methods.

The 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,

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. 
- 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 ?
- In addition, the authors say that they average over 3 lens planes. Why ? to which values of z does that correspond ?
- 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.
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700882731

thread_type:
official_review

<thread_object>
creator: reviewer_x8zd

<content>
summary:
The paper provides a new method for recovering the 3D cosmic web from galaxy weak lensing signal.
The proposed approach uses a coordinate-based neural field method with positional encoding. (Equation 6).
The reconstruction results are demonstrated in a series of experiments.

strengths:
* The paper proposes a new method for an important physics problem.
* The paper is well-written; I enjoyed reading it.
* The physics background is exceptionally well-presented, and even non-experts can understand the main scientific ideas.
* 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.
The 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.

Nonetheless, the proposed algorithm works well, and the physics application, estimating the 3D cosmic web, is a very important problem.
Since 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.

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.

* 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731469457487

thread_type:
official_review

<thread_object>
creator: reviewer_p3hc

<content>
summary:
***Late review: Apologies my review is in one day late. I will aim to engage as quickly as needed to makeup.***

This 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.

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). 

-- 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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731729101221

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Thank you for your helpful comments and questions. We present a “creative” (p3hC) and “novel approach” (pR5T) that attacks “an important physics problem” (x8ZD) with comparisons to baselines that “are well-presented, and highlight the advantage of the method in particular regimes” (pR5T). Reviewer x8ZD mentions that it is a “great paper” that is “exceptionally well-presented” so that “ even non-experts can understand the main scientific ideas.” As highlighted by reviewer pR5T, “presenting a paper like this while being faithful to the domain science [is] particularly challenging, and the authors do a good job in this respect.”

**Machine Learning Contributions (x8ZD, gGfb)**

We firmly believe that our work represents a significant contribution to the ML community. Below, we outline several reasons that we will clarify in the updated paper.

**Neural Fields for Science:** While neural field models with positional encoding have previously been applied to various problems, our work demonstrates how the framework can be incorporated with the underlying physics specific to our problem. We work with noisy measurements from a single viewpoint, and the cosmic web volumes we study have different statistical properties compared to natural scenes. As highlighted by pR5T, this unique context raises important ML questions: How does the implicit regularization in neural fields help estimate solutions to severely ill-posed, underconstrained inverse problems? And how well does positional encoding, effective in natural scenes, generalize to these new scientific contexts? Applications of neural fields was the topic of 2023 ICLR (https://sites.google.com/view/neural-fields) and 2024 ECCV (https://neural-fields-beyond-cams.github.io/) workshops, featuring invited talks that apply neural fields with positional encoding to different 3D reconstruction tasks. We believe that adapting neural fields to a new context with completely different physics represents a significant contribution, and gives value to the ML community.

**Introducing a New Problem to ML:** Additionally, one of the main challenges of interdisciplinary work is effective communication between fields, and we are glad to have expressed our results in a way that is understandable to the ML community (x8ZD, p3hC, pR5T), as this important problem could greatly benefit from their insights and contributions. We also wish to highlight that, to our knowledge, this is the first time that a comparison of 3D mass mapping methods has been done. We will make all code and datasets available so that this problem can be further studied by the ML community. 

**A Stepping Stone to Future Work:** We believe our work is a necessary stepping stone to future work that would take advantage of the neural representation we’ve presented. One direction we are actively pursuing is leveraging the neural representation to obtain efficient uncertainty quantification, which, as pR5T notes, is important for downstream science. While probabilistic approaches to 2D mass mapping are an active field of study, to our knowledge none have yet been proposed for 3D mass mapping, possibly due to computational intractability arising from the extra dimension. Even in the simple case of the Wiener filter, which has an analytic posterior, computing this posterior with a full covariance becomes intractable in 3D. In future work, we plan to build on the theory from [1], which demonstrates how neural representations can play an essential role in enabling efficient posterior estimation. The method we present serves as a solid foundation for further progress in this important direction.

**Wiener Filter Baseline (p3hC, gGfb)**

As gGfb highlights through various references, there have been recent developments in the field of 2D mass mapping. However, these methods are not directly applicable to 3D. Currently, there are no publicly available 3D mass mapping codes in working condition. While one code exists, it is non-functional, and the authors were unable to provide a working version. To the best of the authors' knowledge, the Wiener filter remains the only 3D method that has been applied to real cosmic shear data, with its most recent usage documented in 2018 [2]. Although code for this method was also not publicly available, we implemented our own transverse Wiener filter to establish a baseline.

In regards to p3hC’s question about Gaussian priors: On scales much larger than those we are targeting, the overdensity field of the universe can be described as Gaussian [3]. This coupled with the simplicity of the Gaussian prior makes it a popular choice for reconstructions of this type. However, non-Gaussianity present on small cosmic scales motivates our more general ML-based approach.

[1] https://arxiv.org/abs/2007.05864

[2] https://academic.oup.com/pasj/article/70/SP1/S26/4097646

[3] https://ned.ipac.caltech.edu/level5/March01/Coles/Coles4.html
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies the problem of reconstructing the 3D dark matter field from 2D observations. The method uses a coordinate-based neural field with positional encodings to represent the matter field which is then passed through a physics-based differentiable forward model to obtain a shear field which can be compared to observations.    

Strengths: This paper addresses an important problem in astronomy. The technique is a novel one in weak lensing, although the underlying neural field methods are not new, and is an interesting combination of ML methods with a physics-based simulator.  The evaluations are promising and show good results against reasonable baselines on simulated data. The paper is well-written and contains a good background on the application area for a general audience.

Weaknesses: The main the noted weakness of the paper is limited innovation the ML methodology which uses standard techniques, although they have not been applied to this physics problem before.  Another noted limitation of the method is limited ability to determine the impact of implicit priors in the reconstruction.

Conclusion:    Despite the range of scores, reviewers broadly agreed on the merits and weaknesses of the paper.  All reviewers actually agree that the paper is of high quality, but question whether ICLR is the right venue for publication since the paper does not make any new contributions to core ML methodology and is largely an application of ML to a new domain.  There is not a clear cut consensus here, and I appreciate the reviewers openness and flexibility during the discussion phase. In the end, I tend to believe that scientific applications papers are important to the field and figuring out how to deploy even existing methods in new domains is worthwhile to our community.