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

Biwei Dai; Po-Wen Chang; Erwan Allys; Wahid Bhimji; Boris Bolliet; Thomas Borrett; Paolo Calafiura; Ragansu Chakkappai

astro-ph.CONov 6, 2026

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

Abstract

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

Panel of Figure 1 — Figure 1: Visual comparison of the simulation datasets at the same cosmology.

Panel of Figure 2 — 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	$\approx 10.5$
+ Cmbagent -One shot	BaseCNN	MCMC	No	$11.02$
+ GPT-5	BaseCNN	Grid	No	$11.42$
+ Gemini-2.5-Pro	Inception + InceptionSE	Grid	$D_{4}$	$11.72$
+ Humans	WST	Grid	$D_{4}$	$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	$N_{\mathrm{particle}}$	force resolution	$N_{\mathrm{step}}$	simulation code
$0<z<0.45$	$320\ h^{-1}\mathrm{Mpc}$	$960^{3}$	$0.03\ h^{-1}\mathrm{Mpc}$	adaptive	MP-Gadget [ 20 ]
$0.45<z<1.05$	$704\ h^{-1}\mathrm{Mpc}$	$2816^{3}$	$0.125\ h^{-1}\mathrm{Mpc}$	$60$	FastPM [ 21 ]
$1.05<z<2.72$	$1536\ h^{-1}\mathrm{Mpc}$	$1536^{3}$	$0.5\ h^{-1}\mathrm{Mpc}$	$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 (

\Omega_{m},S_{8}

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

Cosmology

See all 9 papers

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
COSMIO: A Benchmark for Cross-Survey Modality Imputation
Nov 3, 2026Dichang Zhang, Yixuan Shao, Jiaqi Deng +8Incomplete Data Imputation Learning with Missing Data
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
Competing with AI Scientists: Agent-Driven Approach to Astrophysics Research
Mar 18, 2026Thomas Borrett, Licong Xu, Andy Nilipour +7Scientific Workflow Automation Autonomous Research Agents
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

Apr 15, 2026J/K move · Enter open · S save

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

Biwei Dai, Po-Wen Chang, Wahid Bhimji, Paolo Calafiura+14

Institute for Advanced Study · Lawrence Berkeley National Laboratory · Université Paris-Saclay, CNRS/IN2P3, IJCLab+7

Uncertainty Quantification Distribution Shift Robustness Data-Efficient Learning Parameter Estimation

Open paper AskPDF arXiv

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.

Abstract

Cosmology

Explore similar work

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

COSMIO: A Benchmark for Cross-Survey Modality Imputation

Competing with AI Scientists: Agent-Driven Approach to Astrophysics Research