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

arXiv: https://arxiv.org/abs/2501.10486 · HTML: https://mlsanity.com/paper/2501.10486/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 61 references that are in the map.
- First posted on arXiv 17 Jan 2025.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Quantifying Attention Flow in Transformers](https://mlsanity.com/paper/2005.00928/) (2020)
- [The Mythos of Model Interpretability](https://mlsanity.com/paper/1606.03490/) (2016)
- [Real-time gravitational-wave science with neural posterior estimation](https://mlsanity.com/paper/2106.12594/) (2021)
- [Bayesian parameter estimation using conditional variational autoencoders for gravitational-wave astronomy](https://mlsanity.com/paper/1909.06296/) (2019)
- [Complete parameter inference for GW150914 using deep learning](https://mlsanity.com/paper/2008.03312/) (2020)
- [Real-time gravitational-wave inference for binary neutron stars using machine learning](https://mlsanity.com/paper/2407.09602/) (2024)
- [Gravitational-wave parameter estimation with autoregressive neural network flows](https://mlsanity.com/paper/2002.07656/) (2020)
- [Learning Bayesian posteriors with neural networks for gravitational-wave inference](https://mlsanity.com/paper/1909.05966/) (2019)
- [Statistically-informed deep learning for gravitational wave parameter estimation](https://mlsanity.com/paper/1903.01998/) (2019)
- [Unsupervised Learning Architecture for Classifying the Transient Noise of Interferometric Gravitational-wave Detectors](https://mlsanity.com/paper/2111.10053/) (2021)
- [Training Process of Unsupervised Learning Architecture for Gravity Spy Dataset](https://mlsanity.com/paper/2208.03623/) (2022)

## Abstract (arXiv)

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.