FAIR Universe Weak Lensing ML Uncertainty Challenge: Dataset and Cosmological Inference Competition
Authors: Biwei Dai, Po-Wen Chang, Erwan Allys, Wahid Bhimji, Boris Bolliet, Thomas Borrett, Paolo Calafiura, Ragansu Chakkappai, Yuan-Tang Chou, Noé Dia, Sascha Diefenbacher, Ibrahim Elsharkawy, Steven Farrell, Sacha Guerrini, Isabelle Guyon, Chris Harris, Shih-Chieh Hsu, Wassim Kablan, François Lanusse, Celia Lecat, Julia Linhart, Yuezhou Ma, Benjamin Nachman, Shubhojit Naskar, Andy Nilipour, Laurence Perreault-Levasseur, Sébastien Pierre, Benjamin Remy, David Rousseau, Uroš Seljak, Haonan Shangguan, Sammy Nasser Sharief, Yuanxu Sun, Andreas Tersenov, Ihsan Ullah, Huikun Weng, Licong Xu, Justine Zeghal, Yulei Zhang, Hang Zhou
Organizations: Institute for Advanced Study · Lawrence Berkeley National Laboratory · École normale supérieure – PSL · University of Cambridge · Université Paris-Saclay, CNRS/IN2P3, IJCLab · ChaLearn · University of Washington · Université de Montréal · University of Toronto · CEA Paris-Saclay · François Arago Centre, Astroparticule et Cosmologie (APC) laboratory · Centre national de la recherche scientifique (CNRS) · New York University · Tsinghua University · SLAC National Accelerator Laboratory · Stanford University · Indian Institute of Technology Delhi · University of Chicago · NSF-Simons AI Institute for the Sky (SkAI) · University of California, Berkeley · University of Crete
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.
Figures & tables
Leaderboard
Participant
Score
Standardized MSE / coverage
(A) 1st
Cmbagent
11.703
0.103 / 0.700
(A) 2nd
Transatlantic Dream Team
11.654
0.104 / 0.709
(A) 3rd
Shubhojit
11.599
0.103 / 0.658
(B) 1st
Shubhojit
11.361
0.097 / 0.662
(B) 2nd
THUML
11.051
0.105 / 0.673
(B) 2nd
jagoncalves
11.037
0.107 / 0.668
Table 1: Final leaderboards evaluated on the holdout sets. Holdout set (A) contains new realizations of cosmological models seen in the training and public test data; holdout set (B) contains new realizations of unseen cosmological models. The final aggregate leaderboard uses the average of the scores on (A) and (B).
Figure 1: Visual comparison of the simulation datasets at the same cosmology.
Figure 2: The two model architectures are shown in the figure. The models exploit either residual feature skipping or dense feature concatenation to resolve deep maps.
Agents/LLMs
Model
Inference
Augmentation
Score
Cmbagent -P&C
ResNet-18/34/50
MCMC
Flips/rotations
≈10.5
+ Cmbagent -One shot
BaseCNN
MCMC
No
11.02
+ GPT-5
BaseCNN
Grid
No
11.42
+ Gemini-2.5-Pro
Inception + InceptionSE
Grid
D4
11.72
+ Humans
WST
Grid
D4
11.73
Table 2: Summary of model–generation pathways used during the architecture search, showing the models proposed under different Cmbagent /LLM configurations and their corresponding test-set performance. The left column indicates the source of each proposed architecture: some models were generated directly by Cmbagent (either through Planning & Control (P&C) or One-shot ), while in other cases a base model first proposed by Cmbagent was subsequently refined by external LLMs such as GPT-5 or Gemini [ 51 ] . In the inference method column, “MCMC” refers to sampling the posterior using a Markov Chain Monte Carlo algorithm. “Grid” means evaluating the likelihood on a fixed parameter grid and forming parameter estimates by weighting grid points according to their likelihood values. BaseCNN denotes our modified version of the starter-kit SimpleCNN , with slightly wider convolutional layers and a replacement of the large fully connected head by global average pooling followed by a small MLP head.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
redshift
box size
Nparticle
force resolution
Nstep
simulation code
0<z<0.45
320h−1Mpc
9603
0.03h−1Mpc
adaptive
MP-Gadget [ 20 ]
0.45<z<1.05
704h−1Mpc
28163
0.125h−1Mpc
60
FastPM [ 21 ]
1.05<z<2.72
1536h−1Mpc
15363
0.5h−1Mpc
15
FastPM [ 21 ]
Appendix
Table 3: N-body simulation parameters for different redshift ranges.
Figure 3: Example noiseless and noisy weak lensing convergence maps. The top panel and bottom panel show two different maps generated with different parameters of systematic effects and random seed, assuming the same cosmological parameters.
Figure 4: The 101 cosmological parameters ( Ωm,S8 ) in our suite of weak lensing simulations.
Figure 5: A comparison between the posterior distributions inferred by our baseline methods. The neural network-based approaches can estimate the parameters more accurately.
Figure 6: A summary of the number of submissions for the final evaluation in methodology categories.
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
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.
Dichang Zhang, Yixuan Shao, Jiaqi Deng +8
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 +2
We present an agent-driven approach to the construction of parameter inference pipelines for scientific data analysis. Our method leverages a multi-agent system, Cmbagent (the analysis system of the AI scientist Denario), in which specialized agents collaborate to generate research ideas, write and execute code, evaluate results, and iteratively refine the overall pipeline. As a case study, we apply this approach to the FAIR Universe Weak Lensing Uncertainty Challenge, a competition under time constraints focused on robust cosmological parameter inference with realistic observational uncertainties. While the fully autonomous exploration initially did not reach expert-level performance, the integration of human intervention enabled our agent-driven workflow to achieve a first-place result in the challenge. This demonstrates that semi-autonomous agentic systems can compete with, and in some cases surpass, expert solutions. We describe our workflow in detail, including both the autonomous and semi-autonomous exploration by Cmbagent. Our final inference pipeline utilizes parameter-efficient convolutional neural networks, likelihood calibration over a known parameter grid, and multiple regularization techniques. Our results suggest that agent-driven research workflows can provide a scalable framework to rapidly explore and construct pipelines for inference problems.
Thomas Borrett, Licong Xu, Andy Nilipour +7
Cavendish Astrophysics, University of Cambridge, Madingley Road, Cambridge CB3 0HA, UK · Kavli Institute for Cosmology, University of Cambridge, Madingley Road, Cambridge CB3 0HA, UK · Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge CB3 0HA, UK +4