GraviBERT: Transformer-based inference for gravitational-wave time series

TL;DR

GraviBERT is a transformer-based deep learning framework for gravitational wave inference that employs staged training, enabling transfer learning, domain adaptation, and improved accuracy across detector configurations and waveform models.

Contribution

Introduces GraviBERT with a novel staged training process, demonstrating transfer learning, domain adaptation, and multi-task capabilities for gravitational wave data analysis.

Findings

Pretraining reduces MAE by up to 31% and accelerates convergence by 6.6 times.

Pretrained models converge up to 15 times faster on new detector noise profiles.

Cross-waveform transfer reduces MAE by up to 44% and achieves high R^2 scores for mass parameters.

Abstract

We introduce GraviBERT, a novel deep learning framework for gravitational wave inference, built on a multi-scale feature extractor with a transformer encoder and a suitable regression head. A key novelty of GraviBERT is its staged training: a BERT-style self-supervised pretraining phase to learn transferable representations, followed by supervised fine-tuning on labeled data. GraviBERT demonstrates consistent transfer learning across detector configurations and waveform models. On in-domain data, pretraining reduces the MAE by up to $31\%$ and accelerates convergence by $\sim 6.6 \times$, with mean relative precision for point estimates reaching the few-percent level and MAE in effective spin of $\sim 10^{-3}$ at SNR = 10. For domain adaptation to new detector noise profiles, the pretrained model converges up to $15\times$ faster on small target datasets and reduces estimation errors by up to $\sim 47\%$, demonstrating detector-agnostic learning. Cross-waveform approximant transfer achieves up to $44\%$ MAE reductions and up to $15\times$ training speedups, with $R^2$ scores consistently exceeding $0.9$ for mass parameters at SNR = 10 compared to $0.74$ - $0.87$ when training from scratch. GraviBERT works directly with noisy waveforms, and in its current form quantifies predictive uncertainty through MC dropouts. After pretraining, the regression head could be adapted to multiple downstream inference tasks in gravitational-wave astronomy.