COSMIO: A Benchmark for Cross-Survey Modality Imputation

Dichang Zhang; Yixuan Shao; Jiaqi Deng; Jiali Cui; Haotian Yin; Yuanpeng Liu; Xiang Gao; Zhiqiang Lao

astro-ph.IMNov 3, 2026

COSMIO: A Benchmark for Cross-Survey Modality Imputation

Authors: Dichang Zhang, Yixuan Shao, Jiaqi Deng, Jiali Cui, Haotian Yin, Yuanpeng Liu, Xiang Gao, Zhiqiang Lao,

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 ·

Abstract

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)	$\sim 0.7^{\prime\prime}$	$\sim 0.16^{\prime\prime}$	$\sim 0.11^{\prime\prime}$
Pixel Scale (arcsec)	$0.2^{\prime\prime}$	$0.1^{\prime\prime}$	$0.11^{\prime\prime}$
Depth	$r\sim 27.5$	VIS $\sim 24.5$	$H\sim 26$ – $27$
Noise	seeing $+$ sky	readout	thermal $+$ readout
Survey Area (deg 2 )	$\sim 18{,}000$	$\sim 15{,}000$	$\sim 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 $\downarrow$	PSNR $\uparrow$	CRPS $\downarrow$
$\text{Euclid}\rightarrow\text{Roman}$ $\text{LSST}\rightarrow\text{Roman}$ $\text{Euclid+LSST}\rightarrow\text{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 $\downarrow$	PSNR $\uparrow$	CRPS $\downarrow$
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	$\pi=0.1$	$\pi=0.05$	$\pi=0.01$	$\pi=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

\pi

(higher is better). Bold indicates the highest AUPRC in each column.

Method	AEE $\downarrow$
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 $\downarrow$	PSNR $\uparrow$	CRPS $\downarrow$
Euclid $\rightarrow$ 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 $\downarrow$	PSNR $\uparrow$	CRPS $\downarrow$	AEE $\downarrow$
Euclid+LSST $\rightarrow$ 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 $\rightarrow$ 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.

Incomplete Data Imputation

Incomplete data imputation estimates unobserved values while accounting for data structure, missingness patterns, and uncertainty.

See all 80 papers

FAIR Universe Weak Lensing ML Uncertainty Challenge: Dataset and Cosmological Inference Competition
Nov 6, 2026Biwei Dai, Po-Wen Chang, Erwan Allys +37Cosmology Scientific ML
Institute for Advanced Study · Lawrence Berkeley National Laboratory · École normale supérieure – PSL +18
AS-Bridge: A Bidirectional Generative Framework Bridging Next-Generation Astronomical Surveys
Mar 12, 2026Dichang Zhang, Yixuan Shao, Simon Birrer +1Conditional Generative Modeling Image-to-Image Translation
Stony Brook University, Stony Brook, NY, USA
FAIR Universe Weak Lensing ML Uncertainty Challenge: Handling Uncertainties and Distribution Shifts for Precision Cosmology
Apr 15, 2026Biwei Dai, Po-Wen Chang, Wahid Bhimji +15Uncertainty Quantification Distribution Shift Robustness
Institute for Advanced Study · Lawrence Berkeley National Laboratory · Université Paris-Saclay, CNRS/IN2P3, IJCLab +7

Nov 6, 2026J/K move · Enter open · S save

FAIR Universe Weak Lensing ML Uncertainty Challenge: Dataset and Cosmological Inference Competition

Biwei Dai, Po-Wen Chang, Erwan Allys, Wahid Bhimji+36

Institute for Advanced Study · Lawrence Berkeley National Laboratory · École normale supérieure – PSL+18

Cosmology Scientific ML Parameter Estimation Robust Estimation

Open paper AskPDF arXiv

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.

Abstract

Incomplete Data Imputation

Explore similar work

FAIR Universe Weak Lensing ML Uncertainty Challenge: Dataset and Cosmological Inference Competition

AS-Bridge: A Bidirectional Generative Framework Bridging Next-Generation Astronomical Surveys

FAIR Universe Weak Lensing ML Uncertainty Challenge: Handling Uncertainties and Distribution Shifts for Precision Cosmology