Gravitational-wave parameter estimation with machine-learning generated surrogate waveforms

Suyog Garg; Kipp Cannon

gr-qcAug 20, 2026

Gravitational-wave parameter estimation with machine-learning generated surrogate waveforms

Authors: Suyog Garg, Kipp Cannon

Organizations: Research Center for Early Universe, The University of Tokyo, Bunkyo-ku, Tokyo 113-0033 JAPAN

Abstract

The worldwide network of gravitational-wave detectors have detected more than 350 binary coalescence events till date. Future third-generation detectors, like Einstein telescope, are expected to detect orders-of-magnitude more signals from sources with more complicated characteristics, including eccentric orbits and high-mass ratio binaries. It is well-established that the computational cost of parameter estimation for signals from these kinds of sources will be extremely high. In particular, the process could be sped-up if generating theoretical waveform predictions, used for likelihood calculation becomes faster. Recently, various machine-learning techniques has been proposed to this end. In this work, we propose a two-stage deterministic conditional-autoencoder model for generating four-parameter SEOBNRv4 waveforms. The first-stage of the model generates amplitude and phase series of the waveform, while the second-stage calibrates the residual error in the predictions. Our model achieves a median mismatch of around $10^{-2}$ with the target polarization waveforms, while the calibrated amplitude/phase series achieve $10^{-6}$ level cosine distance error. We then propose a waveform conditioning step to enable use of these surrogate waveforms for downstream parameter estimation tasks. Finally, we perform extensive parameter estimation tests, with ML and EOB waveform injections and try to recover posterior estimates for the source parameters. We find that when ML waveforms are used to recover EOB target parameter estimates, the inferred posterior have some systematic bias. This inherent bias can be estimated and corrected for, and then importance reweighting of posterior samples can enable use of low-accuracy surrogate waveforms at low SNRs.

Figures & tables

Component	layers	Notes
XEncoder	pre-FC layers	None
	CNN layers	Conv1D (in=2,out=32,k=6,d=1) + GELU + MaxPool (k=4), Conv1D (in=32,out=32,k=6,d=1) + GELU + MaxPool (k=4)
	post-FC layers	Linear (16324,256) + GELU, Linear (256,256) + GELU, Linear (256,48) + GELU
		Encodes inputs shape $(2,8191)$ to latent dimension $d_{\rm X}=24$ .
KeyEncoder	pre-FC layers	Linear (8,500) + GELU, Linear (500,500) + GELU, Linear (500,8) + Identity
	CNN layers	None

Table 1: Details of the network architecture used in the waveform generation first-stage of the FlexCAEPhase model.

Layers	Configuration	Output shape
Input projection	Conv1D(in=6,out=128,k=7,p=3)	$128\times N$
Residual layers	$5\times$ [Conv1D(in=128,out=128,k=7) $\rightarrow$ GroupNorm(n=8) $\rightarrow$ GELU $\rightarrow$ Dropout(0) $\rightarrow$ Conv1D(in=128,out=128,k=7) $\rightarrow$ GroupNorm(n=8) $\rightarrow$ skip-connection $\rightarrow$ GELU ]	$128\times N$
Output projection	Conv1D(in=128,out=2,k=7,p=3); initialized to zero	$2\times N$

Table 2: Architecture of the residual calibration network used in the second-stage of the FlexCAEPhase model. The model takes a 6-channel input consisting of first-stage output amplitude/phase and the four normalized source parameters, and predicts two residual channels.

Figure 1: (Top) Distribution of the cosine distance of the ml model generated and calibrated amplitude and phase series outputs from the target. (Bottom) Distribution of the mismatch values for the polarization waveforms reconstructed from the calibrated ml model outputs. All distributions are shown for data in the test dataset.

Figure 2: An example overplot of the calibration result for the two-stage waveform generator model. The ml model outputs are normalized amplitude and phase.

Figure 3: Schematic showing waveform conditioning performed by applying a $\sin^{2}(t)$ taper at the start of the amplitude series output by the ml model.

Figure 4: An example overplot of the gw polarization strain with the waveform conditioning applied to the outputs of the two-stage waveform generation and calibration model.

Figure 5: (Top) Distribution of the cosine distance for the conditioned amplitude and phase series outputs from the two-stage waveform generation and calibration model. (Bottom) Distribution of the mismatch values for the polarization waveforms reconstructed from conditioned outputs. All distributions are shown for data in the test dataset.

Figure 6: Contour plot of the mismatch values for $h_{+}$ reconstruction from conditioned wavefrom outputs of the two-stage waveform generation and calibration model. Full parameter space is used, i.e. $\chi_{\rm eff}\in[-0.99,0.99]$ .

Panel of Figure 7 — Figure 7: Mismatch values as a function of total mass (top left), chirp mass (top right), mass ratio (bottom left), and effective spin (bottom right) for conditioned waveform outputs from the two-stage waveform generation and calibration model, for all data in the test dataset. Full parameter space is used, i.e. $\chi_{\rm eff}\in[-0.99,0.99]$ .

Figure 8: Posterior distributions recovered for a ML2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=64.82\,\mathrm{M}_{\odot}$ , $m_{2}=49.75\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=-0.595$ , $\chi_{2}(z)=-0.08$ , $d_{L}=$400\text{\,}\mathrm{Mpc}$$ . snr of the injected signal is about $107.12$ . Other parameters are fixed at: $\theta_{\rm jn}=0.4,\psi=2.659,\phi_{c}=1.3,t_{c}=1126259642.413,\mathrm{ra}=1.375$ , and $\mathrm{dec}=-1.2108$ . The marginalized posterior histograms show the probability density for different posterior samples, alongwith the $1\sigma$ credible interval. The 2D joint posteriors show the 2D $1\sigma,2\sigma$ and $3\sigma$ credible regions. Each column header denotes the median value of the parameter and the $1\sigma$ credible range around it. — Figure 8: Posterior distributions recovered for a ML2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=64.82\,\mathrm{M}_{\odot}$ , $m_{2}=49.75\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=-0.595$ , $\chi_{2}(z)=-0.08$ , $d_{L}=$ 400\text{,}\mathrm{Mpc}$$ . snr of the injected signal is about $107.12$ . Other parameters are fixed at: $\theta_{\rm jn}=0.4,\psi=2.659,\phi_{c}=1.3,t_{c}=1126259642.413,\mathrm{ra}=1.375$ , and $\mathrm{dec}=-1.2108$ . The marginalized posterior histograms show the probability density for different posterior samples, alongwith the $1\sigma$ credible interval. The 2D joint posteriors show the 2D $1\sigma,2\sigma$ and $3\sigma$ credible regions. Each column header denotes the median value of the parameter and the $1\sigma$ credible range around it.

Figure 9: pp -plot for the ML2ML recovery experiment with 1045 injections using conditioned waveforms. The luminosity distance for the injections is set at $d_{L}=$400\text{\,}\mathrm{Mpc}$$ . The diagonal curve points that there is no systematic error in the pe set-up now. All other injection parameters values and other pe settings remain the same as in Figure 8 . — Figure 9: pp -plot for the ML2ML recovery experiment with 1045 injections using conditioned waveforms. The luminosity distance for the injections is set at $d_{L}=$ 400\text{,}\mathrm{Mpc}$$ . The diagonal curve points that there is no systematic error in the pe set-up now. All other injection parameters values and other pe settings remain the same as in Figure 8 .

Figure 10: Posterior distributions recovered for a EOB2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=74.09\,\mathrm{M}_{\odot}$ , $m_{2}=70.31\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=-0.32$ , $\chi_{2}(z)=0.31$ , $d_{L}=$400\text{\,}\mathrm{Mpc}$$ . snr of the injected signal is about $139.45$ . Other parameters and pe settings remain the same as in Figure 8 . — Figure 10: Posterior distributions recovered for a EOB2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=74.09\,\mathrm{M}_{\odot}$ , $m_{2}=70.31\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=-0.32$ , $\chi_{2}(z)=0.31$ , $d_{L}=$ 400\text{,}\mathrm{Mpc}$$ . snr of the injected signal is about $139.45$ . Other parameters and pe settings remain the same as in Figure 8 .

Figure 11: pp -plot for the EOB2ML pe experiment with 953 injections using conditioned waveforms. Some inherent systematic bias is seen. The luminosity distance for the injections is set at $d_{L}=$400\text{\,}\mathrm{Mpc}$$ . All other injection parameters values (apart from the four parameter shown) and other pe settings remain the same as in Figure 8 . — Figure 11: pp -plot for the EOB2ML pe experiment with 953 injections using conditioned waveforms. Some inherent systematic bias is seen. The luminosity distance for the injections is set at $d_{L}=$ 400\text{,}\mathrm{Mpc}$$ . All other injection parameters values (apart from the four parameter shown) and other pe settings remain the same as in Figure 8 .

Figure 12: Posterior distributions recovered for a EOB2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=72.12\,\mathrm{M}_{\odot}$ , $m_{2}=36.08\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=0.17$ , $\chi_{2}(z)=-0.69$ , $d_{L}=$4000\text{\,}\mathrm{Mpc}$$ . snr of the injected signal is about $9.95$ . Other parameters and pe settings remain the same as in Figure 8 . — Figure 12: Posterior distributions recovered for a EOB2ML pe run with conditioned waveform outputs. The injection parameters are $m_{1}=72.12\,\mathrm{M}_{\odot}$ , $m_{2}=36.08\,\mathrm{M}_{\odot}$ , $\chi_{1}(z)=0.17$ , $\chi_{2}(z)=-0.69$ , $d_{L}=$ 4000\text{,}\mathrm{Mpc}$$ . snr of the injected signal is about $9.95$ . Other parameters and pe settings remain the same as in Figure 8 .

Figure 13: pp -plot for the EOB2ML pe experiment with 511 injections using conditioned waveforms. The luminosity distance for the injections is set at $d_{L}=$4000\text{\,}\mathrm{Mpc}$$ . The snr of the signal scales inversely with the luminosity distance to the source. Thus, in comparison to Figure 11 , we find that there is some improvements in the inherent systematic bias at lower snr . All other injection parameters values (apart from the four parameter shown) and other pe settings remain the same as in Figure 8 . — Figure 13: pp -plot for the EOB2ML pe experiment with 511 injections using conditioned waveforms. The luminosity distance for the injections is set at $d_{L}=$ 4000\text{,}\mathrm{Mpc}$$ . The snr of the signal scales inversely with the luminosity distance to the source. Thus, in comparison to Figure 11 , we find that there is some improvements in the inherent systematic bias at lower snr . All other injection parameters values (apart from the four parameter shown) and other pe settings remain the same as in Figure 8 .

Panel of Figure 14 — Figure 14: Distribution of the ratio of the inferred posterior mode and median value for the primary (top) and secondary (bottom) mass parameters calculated from the EOB2ML pe runs results for injections described in Figure 11 .

Panel of Figure 15 — Figure 15: Distribution of the ratio of the inferred posterior mode and median value for the primary (top) and secondary (bottom) mass parameters calculated from the EOB2ML pe runs results for injections described in Figure 13 .

Panel of Figure 16 — Figure 16: Inferred posterior mode versus ratio of inferred to true value for the chirp mass, calculated from the distributions shown in Figure 14 for $d_{L}=$ 400\text{,}\mathrm{Mpc} $. EOB2ML pe case (top), and in Figure 15 for $d_{L}=$4000\text{\,}\mathrm{Mpc}$ (bottom). The linear regression fit lines are also plotted. The ratio with the inferred posterior median is shown for completeness sake.

Panel of Figure 17 — Figure 17: Biased (top) and bias-corrected (bottom) one-dimensional marginalized posteriors for the secondary mass for EOB2ML pe run for the injection described in Figure 10 .

Figure 18: Schematic sketch showing the expected trend in likelihood speedup gain, given by the ratio of number of likelihood evaluations required in an EOB2EOB pe run over the number of likelihood evaluations required when reweighting samples from the EOB2ML proposal distribution. At low snr , the EOB2ML proposal posterior is broad and may overlap significantly with the EOB2EOB target posterior, enabling efficient importance reweighting. At high snr , waveform-systematic bias becomes more significant relative to the statistical uncertainty, the effective sample fraction decreases, and the cost of likelihood evaluation increases. The reweighting strategy is useful only when it satisfies both a likelihood speedup condition, $\mathcal{S}_{\rm ML}>1$ , and an accuracy condition, such as $N_{\rm eff}\geq N_{\rm eff}^{\rm min}$ .

Gravitational-Wave Astronomy

Gravitational-wave astronomy uses ripples in spacetime to detect cosmic events, measure compact objects, and test gravity.

See all 20 papers

Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms
Jul 27, 2026Michael Pürrer, Ashwin Girish, Lucy M. Thomas +2Neural Surrogate Modeling Gravitational-Wave Astronomy
Department of Physics, East Hall, University of Rhode Island, Kingston, RI 02881, USA · Institute for AI & Computational Research, Carothers Library, University of Rhode Island, Kingston, RI 02881, USA · 3TAPIR, California Institute of Technology, Pasadena, California 91125, USA +2
Discovery of Interpretable Surrogates via Agentic AI: Application to Gravitational Waves
May 11, 2026Tousif Islam, Digvijay Wadekar, Tejaswi Venumadhav +4Surrogate Modeling Gravitational-Wave Astronomy
Kavli Institute for Theoretical Physics, University of California Santa Barbara, Kohn Hall, Lagoon Rd, Santa Barbara, CA 93106, USA. · Center for Gravitational Physics, University of Texas at Austin, Austin, TX 78712, USA. · Department of Physics, University of California at Santa Barbara, Santa Barbara, CA 93106, USA. +5
Binary Black Hole Parameter Estimation with Hybrid CNN-Transformer Neural Networks
Jun 11, 2026Panagiotis N. Sakellariou, Spiros V. Georgakopoulos, Sotiris Tasoulis +1Gravitational-Wave Astronomy Hybrid CNN-Transformer Architectures
Department of Mathematics, University of Thessaly, 3rd km. Old National Road Lamia-Athens, Lamia, 35100, Greece · Department of Computer Science and Biomedical Informatics, University of Thessaly, 35100, Greece

Jul 27, 2026J/K move · Enter open · S save

Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms

Michael Pürrer, Ashwin Girish, Lucy M. Thomas, Scott E. Field+1

Department of Physics, East Hall, University of Rhode Island, Kingston, RI 02881, USA · Institute for AI & Computational Research, Carothers Library, University of Rhode Island, Kingston, RI 02881, USA · 3TAPIR, California Institute of Technology, Pasadena, California 91125, USA+2

Neural Surrogate Modeling Gravitational-Wave Astronomy Differentiable Programming

Open paper AskPDF arXiv

We present a neural network surrogate model that emulates the NRSur7dq4 gravitational waveform model for precessing binary black hole mergers. The surrogate decomposes the waveform into constituent quantities and trains an independent multilayer perceptron (MLP) for each. We validate the surrogate against NRSur7dq4 on 10,000 waveforms spanning its full parameter space (

1 \leq q \leq 4

|χ_{A,B}| \leq 0.8

). For representative total masses between 60 and 300

M_\odot

, median sky-averaged frequency-domain mismatches range from

8.0 \times 10^{-5}

1.7 \times 10^{-4}

, with 95th percentiles below

10^{-3}

. On an NVIDIA L40S GPU the JAX surrogate evaluates a single waveform in about 1 ms end-to-end, roughly 10 times faster than the LALSimulation C implementation of NRSur7dq4, and sustains about 140 times the LALSimulation throughput at batch size 64, making it well suited for both low-latency parameter-estimation samplers and large-scale waveform generation. The full NRSur7dq4 NN waveform-to-likelihood pipeline is implemented in JAX and is differentiable. This is the first neural-network surrogate of a precessing numerical-relativity waveform model to combine validated NR-faithful accuracy with a fully differentiable, GPU-accelerated inference pipeline, enabling gradient-based inference approaches via automatic differentiation including Fisher information matrices, GPU-accelerated nested sampling, gradient-based MCMC and importance sampling.

Abstract

Gravitational-Wave Astronomy

Explore similar work

Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms

Discovery of Interpretable Surrogates via Agentic AI: Application to Gravitational Waves

Binary Black Hole Parameter Estimation with Hybrid CNN-Transformer Neural Networks