Architecture Acceleration of Machine Learning Gravitational Waveform Surrogate Models
Abstract
Accurate and computationally efficient waveform models of compact binary coalescences are a fundamental requirement for gravitational wave data analysis. This work presents differentiable and hardware-accelerated implementations of the machine-learning surrogate models mlgw and mlgw-bns within the JAX ecosystem. The proposed framework enables JIT compilation, vector paralellization, native GPU acceleration, and automatic differentiation, while preserving the original surrogate training infrastructure and its applicability to arbitrary non-precessing waveform approximants. Benchmarks are performed in the binary black hole case with a newly trained surrogate of the SEOBNRv5HM approximant; on CPU, they show speed-ups in waveform evaluation time with respect to the original mlgw implementation that exceed one order of magnitude; further, when exploiting GPU acceleration and large-batch vectorization these reach approximately two orders of magnitude. The benchmarks performed in the case of binary neutron stars with a TEOBResumSPA surrogate model achieve similar gains on GPU. Our JAX-based surrogate modeling framework can be integrated into Bayesian parameter estimation pipelines through GWgpu-jax, an open-source package that connects JAX-compatible gravitational waveform generators to JAX-native sampling algorithms. We demonstrate this capability with a nested-sampling pipeline built upon BlackJax-NS, performing parameter estimation analyses of the GW150914 and GW170817 events with mlgw and mlgw-bns. The analyses require approximately twelve and seventeen minutes, respectively, on a single GPU and yield posterior distributions consistent with those reported by the LIGO-Virgo-KAGRA Collaboration. Beyond nested sampling, JAX automatic differentiation provides efficient gradient and Hessian evaluations of the waveform models, enabling integration with gradient-based Bayesian inference methods.
Related
- DANSur_HM: Modularly incorporating higher modes in a deep learning based gravitational-wave surrogate
- Gravitational-wave parameter estimation with machine-learning generated surrogate waveforms
- Line-of-Sight Acceleration-driven waveform corrections for unequal-mass compact binary coalescences
- Discovering gravitational waveform distortions from lensing: A deep dive into GW231123
- Reconstructing the gravitational waveform from its probe limit
BibTeX
@misc{mascioli2026architecture,
title = {{Architecture Acceleration of Machine Learning Gravitational Waveform Surrogate Models}},
author = {Adriano Frattale Mascioli and Lorenzo Piccari and Saulo Albuquerque and Gabriele Demasi and Giulia Capurri and Massimo Lenti and Angelo Ricciardone and Barbara Patricelli and Gianluca M. Guidi and Giulia Stratta and Walter Del Pozzo and Francesco Pannarale},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.31118},
eprint = {2609.31118},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.31118},
}
RIS
TY - UNPB
TI - Architecture Acceleration of Machine Learning Gravitational Waveform Surrogate Models
AU - Adriano Frattale Mascioli
AU - Lorenzo Piccari
AU - Saulo Albuquerque
AU - Gabriele Demasi
AU - Giulia Capurri
AU - Massimo Lenti
AU - Angelo Ricciardone
AU - Barbara Patricelli
AU - Gianluca M. Guidi
AU - Giulia Stratta
AU - Walter Del Pozzo
AU - Francesco Pannarale
PY - 2026
DA - 2026/09/25
PB - arXiv
DO - 10.48550/arxiv.2609.31118
UR - https://arxiv.org/abs/2609.31118
AB - Accurate and computationally efficient waveform models of compact binary coalescences are a fundamental requirement for gravitational wave data analysis. This work presents differentiable and hardware-accelerated implementations of the machine-learning surrogate models mlgw and mlgw-bns within the JAX ecosystem. The proposed framework enables JIT compilation, vector paralellization, native GPU acceleration, and automatic differentiation, while preserving the original surrogate training infrastructure and its applicability to arbitrary non-precessing waveform approximants. Benchmarks are performed in the binary black hole case with a newly trained surrogate of the SEOBNRv5HM approximant; on CPU, they show speed-ups in waveform evaluation time with respect to the original mlgw implementation that exceed one order of magnitude; further, when exploiting GPU acceleration and large-batch vectorization these reach approximately two orders of magnitude. The benchmarks performed in the case of binary neutron stars with a TEOBResumSPA surrogate model achieve similar gains on GPU. Our JAX-based surrogate modeling framework can be integrated into Bayesian parameter estimation pipelines through GWgpu-jax, an open-source package that connects JAX-compatible gravitational waveform generators to JAX-native sampling algorithms. We demonstrate this capability with a nested-sampling pipeline built upon BlackJax-NS, performing parameter estimation analyses of the GW150914 and GW170817 events with mlgw and mlgw-bns. The analyses require approximately twelve and seventeen minutes, respectively, on a single GPU and yield posterior distributions consistent with those reported by the LIGO-Virgo-KAGRA Collaboration. Beyond nested sampling, JAX automatic differentiation provides efficient gradient and Hessian evaluations of the waveform models, enabling integration with gradient-based Bayesian inference methods.
ER -