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

TL;DR

GraviBERT presents a transformer-driven framework for inferring gravitational-wave source parameters directly from noisy time-domain waveforms. By combining an InceptionTime-inspired multi-scale CNN with a CLS-token transformer and a regression head, the model benefits from masked pretraining followed by supervised fine-tuning, yielding faster convergence and improved accuracy. The approach demonstrates strong domain adaptation across detectors and waveform models, achieving few-percent precision on intrinsic masses and distances and enabling real-time inference at high throughput. These results position GraviBERT as a step toward foundation-style models in gravitational-wave science, with potential to generalize to multi-messenger contexts and broader waveform families.

Abstract

We introduce GraviBERT, a novel deep learning framework for inference of gravitational-wave time series, which relies on an Inception-inspired multi-scale convolutional feature extractor combined with a transformer encoder and a suitable regression head. GraviBERT is trained in two stages: a BERT-style pretraining phase, in which the model learns to predict masked segments in feature space to capture universal patterns and physics, followed by supervised fine-tuning for accurate parameter estimation. This approach demonstrates impressive improvement across multiple metrics compared to training from scratch. On in-domain data, it reduces the mean absolute error for point-estimate parameter inference by up to $30\%$, and training convergence accelerates by up to a factor of six. Moreover, at low signal-to-noise ratio, the mean relative precision of the inferred masses and distances reaches the few-percent level, while the mean absolute error in the effective spin is about $10^{-3}$. For domain adaptation to new detector noise profiles, the pretrained model demonstrates remarkable efficiency, converging up to $15\times$ faster on small target datasets and reducing estimation errors by up to approximately $45\%$, indicating that it learns sufficient detector-agnostic representations. Cross-approximant transfer demonstrates comparable performance, achieving up to $44\%$ reductions in mean absolute error across all parameters and up to $15\times$ training speedups, with $R^2$ scores consistently exceeding 0.9 for mass parameters at signal-to-noise ratio 10, compared to 0.74 - 0.87 when training from scratch. Notably, GraviBERT works directly with noisy waveforms. The final regression head of the model can be adapted for a range of downstream tasks after pretraining, positioning it as a step towards foundation-style models in gravitational-wave and multi-messenger astronomy.

Figures (8)

• Figure 1: The ET's power spectral density (black, ET-B) overlaid with the frequency-domain representations of selected waveforms from our sample at fixed SNR = 30. Different waveforms intersect the PSD at different stages of their evolution. Our analysis uses waveforms with a low-frequency cutoff at $f_{\rm low}$ = 20 Hz, as explained in Section \ref{['sec:data']}.
• Figure 2: Architecture of our feature extractor ($\mathcal{F}$). The highlighted dashed box represents the main processing block, which is repeated $N_B$ times. IM stands for InceptionModule, and MP denotes MaxPool layer. $M$ is the number of modules stacked within each main block.
• Figure 3: Architecture of our InceptionModule. The bottleneck layer first reduces the channel dimensionality of the input $x$. This output is then processed by multiple parallel convolutions to extract multi-scale features. Concurrently, the original input is processed by a max-pooling pathway. The outputs of all pathways are concatenated and normalised.
• Figure 4: Overview of our model architecture.
• Figure 5: Illustration of the binary padding mask downsampling process. A MaxPool operation with a kernel size and stride of 2 is applied to the input mask. This ensures that a position in the downsampled mask is marked as padding (0) only if all corresponding input positions were also padding.