GPU-R2D2

PI: Prof. Yves Wiaux

Executive Summary

Modern radio telescopes such as the Square Kilometre Array (SKA) will produce unprecedented volumes of data, requiring new computational approaches to transform raw measurements into high-fidelity images of the Universe. Traditional imaging methods struggle to scale to the ultra- high resolutions and dynamic ranges required by next-generation surveys. Using DiRAC’s GPU resources, we have advanced a new deep-learning-based framework called R2D2 (‘ R esidual-to- R esidual D NN series for high- D ynamic range imaging”), designed to replace the most computationally expensive components of state-of-the-art radio imaging algorithms with learned models. In 2025, our DiRAC allocation enabled three major advances:

1. Solving the Data Bottleneck for Large-Scale Imaging

Scaling deep learning to large image sizes introduces two fundamental challenges:

1. Data availability for training at high resolution and dynamic range

2. Model scalability to fit and process large-scale images

To address the first challenge, using DiRAC resources, we developed a novel large-scale radio interferometric simulation pipeline capable of generating diverse, high-dynamic-range training datasets. This pipeline is not limited to standard monochromatic imaging. It is fully extensible and can generate:

  • Spherical (wide-field) radio images, enabling training beyond the small-field approximation
  • Wideband, multi-frequency datasets, supporting spectral imaging across different observ- ing bands

By enabling scalable simulation across geometries, and frequencies, this framework estab- lishes a general foundation for next-generation AI-driven radio imaging. This work forms the basis of a forthcoming dataset and simulation methodology paper, enabled directly by DiRAC’s large-scale parallel simulation capability.

2. Architectural Scalability: Ultra-High Resolution (4096 × 4096)

The massive parallelization of DiRAC allowed us to break the ”memory wall” of 4k imaging. The combination of our new data simulation pipeline and the novel transformer-based DNN ar- chitecture establishes a viable pathway for the next generation of ultra-large-scale radio surveys (e.g., SKA, ngVLA, and DSA-2000), as envisioned in the project’s original proposal. To address the ”prohibitive memory footprint” of standard convolutional architectures, which prevents scal- ing to the resolutions required by next-generation surveys, we developed Swin-U-WDSR. This hybrid architecture integrates Swin-style Transformer blocks into U-WDSR DNN [1]. By lever- aging the long-range dependency modeling of Transformers, we capture global structures while preserving local emission continuity, drastically reducing GPU memory requirements. Bench- marking against state-of-the-art algorithms (multiscale CLEAN, AIRI, and uSARA) confirms that R2D2 achieves superior precision and data fidelity at a fraction of the computational time. This work is currently in draft for a second major publication.

3. Cross-Disciplinary Translation: JC-R2D2 for MRI

We have successfully adapted the R2D2 paradigm, originally engineered for radio interferome- try (RI), to medical imaging. We introduce the Jointly Calibrated R2D2 (JC-R2D2) framework for fast, scalable, and calibrated image reconstruction from accelerated non-Cartesian k-space acquisitions in Magnetic Resonance Imaging (MRI). JC-R2D2 enables the joint estimation of images and coil sensitivity maps, correcting physical model errors on-the-fly, resulting in signif- icant increase of GPU memory requirements compared to R2D2. Training such models at scale required DiRAC’s high-memory GPU infrastructure. Future work will extend JC-R2D2 to 3D volumetric and 4D dynamic MRI, where memory demands grow dramatically, making continued access to extreme-scale compute essential.

References

[1] A. Aghabiglou, C. S. Chu, C. Tang, A. Dabbech, and Y. Wiaux, “Toward a robust r2d2 paradigm for radio-interferometric imaging: Revisiting deep neural network training and architecture,” The Astrophysical Journal Supplement Series , vol. 280, no. 2, p. 63, 2025.