🔭 astrophysics

GIGA-Lens 2.0: Strong-Lens Modeling on Multiple GPU Nodes

The paper introduces GIGA-Lens 2.0, a major upgrade to the GPU-accelerated Bayesian strong-lensing framework that enables distributed computing across up to 512 A100 GPUs, significantly improving performance for modeling both simulated and real astronomical systems.

Original authors: Xiaosheng Huang, Linus Upson, Nicolas Ratier-Werbin, Harry Lu, Sean Xu, Elden Yap, Evan Odell, Ansel Parke, Harsh Ambardekar, Saul Baltasar, Nestor Demeure, Bradley Richardson, Andi Gu, Yuan-Ming Hsu

Published 2026-06-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiaosheng Huang, Linus Upson, Nicolas Ratier-Werbin, Harry Lu, Sean Xu, Elden Yap, Evan Odell, Ansel Parke, Harsh Ambardekar, Saul Baltasar, Nestor Demeure, Bradley Richardson, Andi Gu, Yuan-Ming Hsu, Junyi Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Solving a Cosmic Puzzle Faster

Imagine trying to solve a massive, 3D jigsaw puzzle where the pieces are distorted by gravity, and you have to figure out exactly what the picture looked like before it was warped. This is what astronomers do when they study strong gravitational lensing. They look at how massive objects (like galaxy clusters) bend light from objects behind them to learn about the universe's invisible ingredients, like dark matter and dark energy.

For a long time, solving these puzzles was like trying to finish a 10,000-piece puzzle by hand, one piece at a time, in a single room. It was accurate, but it took forever.

GIGA-Lens 2.0 is a major upgrade to the software used to solve these puzzles. The authors have turned that single-room effort into a massive factory with 512 super-computers (GPUs) working together. They successfully ran this system on 128 different computer nodes, allowing them to solve these cosmic puzzles much faster and with greater precision than ever before.

The Three-Step Assembly Line

The software doesn't just guess; it follows a strict three-step process to find the best answer. Think of it like a team of detectives narrowing down a suspect:

The "Best Guess" (MAP): First, the computer makes a quick, rough guess at where the pieces fit. In the new version, they can send different "guessers" to different computers to try many starting points at once, ensuring they don't get stuck on a wrong answer.
The "Map Maker" (SVI): Next, the team builds a rough map of where the real answer is likely to be. This is the hardest part to coordinate because all the computers need to share their notes to update the map. The authors solved this by creating a new communication system that lets 512 computers talk to each other instantly, even if they are in different buildings. They also swapped out an old "calculator" (the Adam optimizer) for a smarter one called AdaBelief, which is less likely to get confused by "noisy" data and makes the map more accurate.
The "Deep Dive" (HMC): Finally, the team explores the map in detail to find the exact solution. They use a method called Hamiltonian Monte Carlo to sample the possibilities. In the past, this step was slow and sometimes got stuck. The new version includes a "warm-up" phase that helps the computer tune its steps before diving in, ensuring it finds the true answer quickly.

Why Speed Matters (It's Not Just About Being Fast)

You might think, "Why do we need it to be faster? Can't we just wait?" The paper argues that speed is actually about accuracy and safety, not just saving time.

The "Double-Check" Analogy: Imagine you are a judge deciding a complex case. If you only have one day to review the evidence, you might make a quick decision. But if you have a team of 128 judges, you can review the same evidence 100 times under different assumptions to see if your conclusion holds up.
Testing for System Errors: Because GIGA-Lens 2.0 is so fast, scientists can run the same cosmic puzzle through the system multiple times with slightly different rules. This helps them spot "systematic errors"—mistakes that happen because of how the model is built, not because of the data itself.
Handling the "Big Ones": Some cosmic puzzles are huge. The paper mentions a specific system called the Cosmic Carousel Lens, which is like a galaxy cluster. Modeling this used to be impossible because the data was too big for a single computer's memory. With the new multi-node system, they can now tackle these giant, complex systems that were previously out of reach.

The Results: A Five-Fold Leap

The authors tested their new system in two ways:

100 Simulated Systems: They created 100 fake cosmic puzzles and solved them all. Almost every single one converged to a perfect solution (a statistical score called $\hat{R} < 1.1$ ), proving the method works reliably.
A Real System: They applied it to a real galaxy lens called DESI J238.5690+04.7276.
- Speed: Using 8 computer nodes, they finished the job 5 times faster than using just one node.
- Precision: They achieved a level of precision ( $\hat{R} < 1.01$ ) that was previously impossible for real-world data, meaning their results are statistically rock-solid.

In Summary

GIGA-Lens 2.0 is like upgrading from a single bicycle to a high-speed train. It allows astronomers to:

Use 512 GPUs working in perfect sync.
Solve massive, complex puzzles (like galaxy clusters) that were too big for old computers.
Run multiple tests to ensure their conclusions about dark matter and the universe are correct, not just lucky guesses.

The paper concludes that this speed isn't just a convenience; it is essential for doing rigorous, high-quality science on the scale required by the next generation of telescopes.

Technical Summary of GIGA-Lens 2.0: Strong-Lens Modeling on Multiple GPU Nodes

Problem Statement
Strong gravitational lensing offers a powerful probe for cosmology, including constraints on dark matter substructure, the Hubble constant ( $H_0$ ), and dark energy. However, realizing this potential requires fully Bayesian inference with rigorous posterior sampling. As lens models become more flexible and datasets grow, ensuring robust posterior exploration and numerical convergence places stringent demands on computational performance and scalability. The original GIGA-Lens framework (GIGA-Lens 1.0) was limited to single-node GPU configurations, hindering its application to large sample sizes, complex systems (e.g., cluster-scale lenses), and joint analyses involving stellar dynamics. These limitations often force simplifying assumptions or restrict the exploration of systematic uncertainties.

Methodology
GIGA-Lens 2.0 introduces a major upgrade to the framework, enabling fully forward-modeling Bayesian strong-lens inference across multiple GPU nodes (up to 128 nodes or 512 A100 GPUs). The pipeline utilizes JAX's shard_map API to distribute computation while maintaining a unified programming model. The methodology addresses three distinct stages of the modeling pipeline:

Parallelization Strategies:
- MAP and HMC: Parallelization is achieved by distributing independent particles (for Multi-start MAP estimation) or independent chains (for Hamiltonian Monte Carlo) across devices.
- Stochastic Variational Inference (SVI): This stage is more complex as it requires coordinated gradient updates across devices to optimize a shared surrogate posterior. The implementation draws a fraction of samples on each device, computes local gradients via automatic differentiation, aggregates these gradients across all nodes, and applies an optimizer step. This necessitates cross-node communication, managed via JAX.
Algorithmic Improvements:
- Optimizer Selection: The authors identified that the default Adam optimizer often caused instability in SVI, characterized by sharp increases in loss before recovery. They replaced Adam with AdaBelief, which adapts step sizes based on gradient deviation, resulting in more stable convergence and a surrogate posterior closer to the true distribution.
- Mass Matrix Adaptation: For challenging systems where the SVI surrogate remains poor, GIGA-Lens 2.0 implements a new mass matrix adaptation step. This involves a longer burn-in phase (approx. 10 $\times$ typical) followed by fitting a new Gaussian surrogate to the final 90% of burn-in samples to better initialize HMC.
Convergence Diagnostics:
The framework employs the Gelman-Rubin metric ( $\hat{R}$ ) and Effective Sample Size (ESS) to ensure convergence. The implementation correctly accounts for cross-chain variance in ESS calculations.

Key Results

Scaling Performance: The pipeline demonstrates near-linear scaling. For large workloads, the parallelizable fraction ( $p$ ) reaches 0.983 for MAP, 0.979 for SVI, and 0.968 for HMC. Even at modest workloads, efficiency remains above 85%.
Simulated Systems: The authors modeled 100 simulated strong lens systems on 16 GPUs. With the new hyperparameters and algorithmic improvements, all but one system converged to $\hat{R} < 1.1$ ; the final system converged after applying the new mass matrix adaptation.
Real System Application: The framework was applied to a real system, DESI J238.5690+04.7276. Using eight nodes, the execution time saw a five-fold speed gain compared to a single node.
Resolution and Precision: Unlike previous applications of GIGA-Lens 1.0 on real HST data which used a supersampling factor of 1, GIGA-Lens 2.0 successfully modeled this system with supersample = 2. Furthermore, while previous work achieved $\hat{R} < 1.1$ , this work achieved $\hat{R} < 1.01$ for all parameters in the real system, indicating superior convergence.

Significance and Claims
The paper claims that GIGA-Lens 2.0 overcomes the primary bottleneck for next-generation wide-field imaging surveys: the ability to model lenses robustly and at scale. The significance of the speed gains is not merely efficiency but the ability to pursue rigor and accuracy. By enabling high-fidelity Bayesian inference at unprecedented scales, the framework allows researchers to:

Model complex, high-dimensional systems (e.g., the "Cosmic Carousel Lens") that were previously computationally prohibitive.
Systematically explore modeling assumptions and parameterizations to quantify systematic uncertainties, rather than relying on simplifying assumptions.
Achieve tighter convergence metrics ( $\hat{R} < 1.01$ ) and higher resolution modeling (supersample = 2) for real observational data.

The authors emphasize that the framework is now accessible, with improved documentation and reproducible setup instructions, making scalable, rigorous strong lens modeling powerful and available for large cosmological inference projects.