Organizations: Department of Computer Science, Stony Brook University, USA · Department of Physics and Astronomy, Stony Brook University, USA · School of Computer Science, University of Technology Sydney, Australia · Futurewei Technologies, USA · Department of Computer Science, New Jersey Institute of Technology, USA
Large-scale astronomical surveys define a heterogeneous, partially observed multi-survey learning setting: each object is observed by only a subset of surveys, while survey observations differ in resolutions, wavelength coverage, noise characteristics, and footprints. This setting gives rise to a natural machine learning problem: given observations of an object in some surveys, how to infer its appearance in other surveys. We formalize this as the problem of cross-survey modality imputation and introduce COSMIO, a multi-survey benchmark of aligned observations obtained by forward-rendered shared astrophysical scenes across LSST, Euclid, and Roman under survey-specific instrument characteristics. The benchmark enables systematic evaluation of imputation across source-target survey combinations and missing-survey patterns. We evaluate a broad range of machine learning methods for imputation, which reveal a clear domain gap: methods developed for existing imputation settings transfer poorly to the survey setting. We identify rare-event imputation as a central challenge. This regime arises in the early stages of new surveys, when rare but scientifically valuable phenomena may be available only from prior surveys. We use strong gravitational lensing, where light from a background galaxy is gravitationally deflected by a foreground galaxy, as a representative rare-event setting. We evaluate model performance on zero-shot imputation of gravitational lenses and on downstream scientific tasks using the imputed lens observations. On our constructed evaluation set of galaxy-galaxy strong lenses, performance degrades under rare-event distribution shift, and this degradation further propagates to the downstream lens-detection task. Code: https://github.com/fw-ic/COSMIO.
Figures & tables
Figure 1: Scientific context for cross-survey rare-event imputation for early-stage observations. (a) Astronomical surveys cover partially overlapping regions of the sky. (b) The same underlying astrophysical object is observed with heterogeneous characteristics across surveys. (c) An imputation model is trained on aligned observations of common galaxies in the overlapping regions. (d) The trained model is applied to infer missing observations for rare events. (e) The reconstructed observations enable downstream scientific tasks.
Parameter
LSST
Euclid
Roman
Bands
g,r,i,z,y
VIS
F106(Y),F129(J),F158(H)
PSF (arcsec)
∼0.7′′
∼0.16′′
∼0.11′′
Pixel Scale (arcsec)
0.2′′
0.1′′
0.11′′
Depth
r∼27.5
VIS ∼24.5
H∼26 – 27
Noise
seeing + sky
readout
thermal + readout
Survey Area (deg 2 )
∼18,000
∼15,000
∼2,000 – 5,000
Table 1: Comparison of key observational parameters across LSST, Euclid, and Roman surveys.
Figure 2: COSMIO samples across three survey modalities, including both common galaxy fields and rare scenes such as strong-lensing systems. Additional dataset visualizations are provided in Appendix H .
Tasks
Method
RMSE ↓
PSNR ↑
CRPS ↓
Euclid→RomanLSST→RomanEuclid+LSST→Roman
PSF-Linear
0.403
33.758
/
Sérsic-FM
0.271
36.838
/
CodeBrain [ 38 ]
0.330
34.533
0.097
diffm4ri [ 39 ]
0.353
34.153
0.094
mmgan [ 33 ]
0.522
30.157
0.159
ummi [ 26 ]
0.449
33.567
0.134
Table 2: Performance comparison of cross-survey imputation methods on common galaxies. Results are grouped by target survey, with metrics averaged over all corresponding imputation tasks. Best results are in bold , and second-best results are underlined . CRPS is not reported for the two deterministic physical baselines. Detailed results for individual Roman-target imputation tasks are provided in Appendix F .
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Preprocessing
RMSE ↓
PSNR ↑
CRPS ↓
Padding
0.184
39.190
0.069
Rescaling
0.277
36.105
0.074
Appendix
Table 4: Padding versus common-grid rescaling for regular-galaxy Roman imputation, evaluated on the native target grid. Lower is better for RMSE and CRPS; higher is better for PSNR.
Method
Year
Core Idea
Source
CodeBrain [ 38 ]
2026
Quantized-code imputation
CVPR
DiffM4RI [ 39 ]
2026
Latent DiT inpainting
IEEE JBHI
MMGAN [ 33 ]
2020
GAN + GMM
IEEE TMI
UMMI [ 26 ]
2025
Mask-pretrained transformer
MICCAI
MMSyn [ 8 ]
2018
Modality-invariant latent
IEEE TMI
M2DN [ 21 ]
2024
Latent Conditional Diffusion
IEEE TMI
Appendix
Table 5: Overview of the brain-MRI and physically motivated imputation baselines.
Method
π=0.1
π=0.05
π=0.01
π=0.001
Asyn-Diff
0.6104
0.4747
0.2141
0.0411
Palette
0.5630
0.4180
0.1629
0.0339
M2DN
0.5388
0.4038
0.1750
0.0397
RePaint
0.5027
0.3546
0.1285
0.0288
Appendix
Table 6: Prior-adjusted AUPRC at target lens prevalence π (higher is better). Bold indicates the highest AUPRC in each column.
Method
AEE ↓
PSF-Linear
1.426
Sérsic-FM
1.460
CodeBrain [ 38 ]
1.424
m2dn [ 21 ]
0.822
RePaint [ 20 ]
0.848
Palette [ 29 ]
0.811
Appendix
Table 7: Mean Angular Extent Error (AEE, radians; lower is better) over the three Roman-imputation tasks. Task-level results are reported in Table 12 . Best and second-best values are bold and underlined, respectively.
Trial
0 shots
10 shots
20 shots
50 shots
100 shots
Trial 1
0.090
0.136
0.124
0.125
0.124
Trial 2
0.090
0.139
0.124
0.126
0.130
Appendix
Table 9: Asyn-Diff lens-imputation CRPS after common-galaxy pretraining followed by lens-only fine-tuning (lower is better). Each trial uses an independently sampled few-shot lens subset. The 0-shot column repeats the shared pretrained baseline without lens supervision.
Trial
0 shots
10 shots
20 shots
50 shots
100 shots
Trial 1
0.090
0.091
0.091
0.085
0.085
Trial 2
0.090
0.084
0.082
0.083
0.091
Appendix
Table 10: Asyn-Diff lens-imputation CRPS with importance-weighted joint training, implemented by weighted sampling with replacement: 25% probability mass on lenses and 75% on common galaxies (lower is better). Each trial uses an independently sampled lens subset. The 0-shot column repeats the shared common-galaxy baseline.
Task
Method
RMSE ↓
PSNR ↑
CRPS ↓
Euclid → Roman
PSF-Linear
0.256
36.33
/
Sérsic-FM
0.251
37.11
/
CodeBrain
0.331
34.50
0.098
diffm4ri
0.361
33.68
0.101
mmgan
0.591
28.99
0.165
ummi
0.406
34.48
0.129
Appendix
Table 11: Comparison of cross-survey imputation methods under deterministic (RMSE, PSNR) and probabilistic (CRPS) metrics for Roman-target imputation tasks. Lower is better for RMSE and CRPS; higher is better for PSNR. Best results are in bold, and second-best results are underlined. CRPS is not reported for the two deterministic physical baselines. A slash denotes an unavailable or unreported result.
Task
Method
RMSE ↓
PSNR ↑
CRPS ↓
AEE ↓
Euclid+LSST → Roman
PSF-Linear
0.715
31.88
/
1.595
Sérsic-FM
0.565
34.58
/
1.495
CodeBrain
0.413
34.56
0.119
1.400
m2dn
0.308
37.53
0.098
0.781
RePaint
0.362
36.29
0.101
0.786
Palette
0.235
39.05
0.085
0.790
Appendix
Table 12: Performance comparison of cross-survey imputation methods on rare galaxy-galaxy strong lensing. Results are reported for individual imputation tasks under distribution shift, where models trained on common galaxies are evaluated on rare events. Best results are in bold, and second-best results are underlined. CRPS is not reported for the two deterministic physical baselines.
Train source
Task
Method
AUROC
AUPRC
GT Roman
–
–
0.9626
0.9606
Imputed Roman
LSST+Euclid → Roman
asyn-diffusion
0.9056
0.8978
RePaint
0.8782
0.8623
Palette
0.8934
0.8832
m2dn
0.8794
0.8688
PSF-Linear
0.7892
0.7662
Appendix
Table 13: Downstream Roman lens detection performance.
Statistic
Asyn-Diff
Palette
M2DN
Sérsic-FM
Total flux
139.30
345.70
479.50
83.57
Aperture flux
41.39
32.95
77.25
69.97
Ellipticity
0.054
0.060
0.054
0.210
Second-moment size T
0.84
0.90
1.30
3.50
Appendix
Table 14: RMSE of photometric and low-order morphological statistics. Lower is better within each row; bold indicates the lowest error, including ties.
Figure 6: Examples of galaxy–galaxy strong-lensing systems in the COSMIO dataset.
Figure 7: Examples of regular galaxies in the COSMIO dataset.
Figure 8: Additional qualitative examples of strong-lensing imputation using asynchronous diffusion.
Figure 9: Additional qualitative examples of regular-galaxy imputation using asynchronous diffusion.
Weak gravitational lensing is a powerful cosmological probe that enables direct mapping of the matter distribution and precise constraints on the fundamental parameters governing structure formation in the universe. In recent years, machine learning has shown strong potential for extracting cosmological information from weak-lensing data beyond traditional summary statistics. However, progress has been limited by the lack of standardized benchmarks that capture the complexity of the dataset, simulator imperfections, and various systematic biases encountered in realistic analyses. To address this gap, we present the FAIR Universe Weak Lensing Machine Learning Uncertainty Challenge, one of the NeurIPS 2025 competitions, as a benchmark dataset and evaluation study. Our dataset reflects realistic weak-lensing observations, incorporating multiple systematic biases, and thereby providing a controlled testbed for robust cosmological inference under uncertainties. We summarize the main methodological advances developed by the winners of our challenge, including data augmentation with conditional optimal transport, an AI-agent-driven approach for parameter inference, ensemble learning, Vision Transformer-based architectures, and physics-informed feature integration. By unifying the benchmark dataset and the winning solutions, this work provides a community resource for developing and evaluating machine learning methods for weak lensing cosmology and, more broadly, for scientific inference under systematic uncertainties.
Biwei Dai, Po-Wen Chang, Erwan Allys +37
Institute for Advanced Study · Lawrence Berkeley National Laboratory · École normale supérieure – PSL +18
The upcoming decade of observational cosmology will be shaped by large sky surveys, such as the ground-based LSST at the Vera C. Rubin Observatory and the space-based Euclid mission. While they promise an unprecedented view of the Universe across depth, resolution, and wavelength, their differences in observational modality, sky coverage, point-spread function, and scanning cadence make joint analysis beneficial, but also challenging. To facilitate joint analysis, we introduce A(stronomical)S(urvey)-Bridge, a bidirectional generative model that translates between ground- and space-based observations. AS-Bridge learns a diffusion model that employs a stochastic Brownian Bridge process between the LSST and Euclid observations. The two surveys have overlapping sky regions, where we can explicitly model the conditional probabilistic distribution between them. We show that this formulation enables new scientific capabilities beyond single-survey analysis, including faithful probabilistic predictions of missing survey observations and inter-survey detection of rare events. These results establish the feasibility of inter-survey generative modeling. AS-Bridge is therefore well-positioned to serve as a complementary component of future LSST-Euclid joint data pipelines, enhancing the scientific return once data from both surveys become available. Data and code are available at https://github.com/ZHANG7DC/AS-Bridge.
Weak gravitational lensing, the correlated distortion of background galaxy shapes by foreground structures, is a powerful probe of the matter distribution in our universe and allows accurate constraints on the cosmological model. In recent years, high-order statistics and machine learning (ML) techniques have been applied to weak lensing data to extract the nonlinear information beyond traditional two-point analysis. However, these methods typically rely on cosmological simulations, which poses several challenges: simulations are computationally expensive, limiting most realistic setups to a low training data regime; inaccurate modeling of systematics in the simulations create distribution shifts that can bias cosmological parameter constraints; and varying simulation setups across studies make method comparison difficult. To address these difficulties, we present the first weak lensing benchmark dataset with several realistic systematics and launch the FAIR Universe Weak Lensing Machine Learning Uncertainty Challenge. The challenge focuses on measuring the fundamental properties of the universe from weak lensing data with limited training set and potential distribution shifts, while providing a standardized benchmark for rigorous comparison across methods. Organized in two phases, the challenge will bring together the physics and ML communities to advance the methodologies for handling systematic uncertainties, data efficiency, and distribution shifts in weak lensing analysis with ML, ultimately facilitating the deployment of ML approaches into upcoming weak lensing survey analysis.
Biwei Dai, Po-Wen Chang, Wahid Bhimji +15
Institute for Advanced Study · Lawrence Berkeley National Laboratory · Université Paris-Saclay, CNRS/IN2P3, IJCLab +7