Enhancing the reliability of machine learning for gravitational wave parameter estimation with attention-based models

TL;DR

The paper tackles the computational burden of Bayesian gravitational-wave parameter estimation by training two Vision Transformer–based models on spectrograms to estimate the effective spin $\chi_{\text{eff}}$ and chirp mass $\mathcal{M}$ from binary black hole signals. By leveraging attention maps, the authors verify that predictions rely on physically meaningful spectrogram regions and quantify how glitches bias estimates, showing that attention can flag unreliable results. An uncertainty-evaluation pipeline inspired by Monte Carlo ideas demonstrates 90% intervals broadly consistent with reference posterior estimates, with total inference time reduced to around six minutes. The approach offers a pathway to rapid, reliable GW parameter estimation and introduces a practical diagnostic tool for glitch robustness that could inform future training and automatic reliability checks in real data analyses.

Abstract

We introduce a technique to enhance the reliability of gravitational wave parameter estimation results produced by machine learning. We develop two independent machine learning models based on the Vision Transformer to estimate effective spin and chirp mass from spectrograms of gravitational wave signals from binary black hole mergers. To enhance the reliability of these models, we utilize attention maps to visualize the areas our models focus on when making predictions. This approach enables demonstrating that both models perform parameter estimation based on physically meaningful information. Furthermore, by leveraging these attention maps, we demonstrate a method to quantify the impact of glitches on parameter estimation. We show that as the models focus more on glitches, the parameter estimation results become more strongly biased. This suggests that attention maps could potentially be used to distinguish between cases where the results produced by the machine learning model are reliable and cases where they are not.

Figures (18)

• Figure 1: A predicted vs actual plot after 50 epochs for the effective spin estimation model. The model was trained for 50 epochs with a learning late of $1 \times 10^{-4}$ and a batch size of 16.
• Figure 2: A predicted vs actual plot after 50 epochs for the chirp mass estimation model. The model was trained for 50 epochs with a learning late of $5 \times 10^{-5}$ and a batch size of 16.
• Figure 3: Flow chart of uncertainty evaluation.
• Figure 4: Histogram of absolute error for the effective spin estimation model. The absolute errors are calculated from the effective spins estimated in Sec. \ref{['subsection:uncertainty']}-(4) and the estimated value in Sec. \ref{['subsection:uncertainty']}-(2).
• Figure 5: Histogram of absolute error for the chirp mass spin estimation model. The absolute errors are calculated from the chirp masses estimated in Sec. \ref{['subsection:uncertainty']}-(4) and the estimated value in Sec. \ref{['subsection:uncertainty']}-(2).