Identifying and Mitigating Machine Learning Biases for the Gravitational Wave Detection Problem

TL;DR

This work targets the gap between traditional matched filtering and ML-based gravitational-wave detection by identifying 11 interconnected learning biases that hinder ML generalization in real detector noise. It introduces Sage, a two-stage architecture with a frontend multiscale feature extractor and a CBAM-ResNet backend, trained with a multi-objective loss and domain-informed data generation/augmentation strategies, including on-the-fly waveform synthesis and PSD recolouring. Sage demonstrates substantial gains over PyCBC and prior ML pipelines in injection tests (e.g., ~11.2% more detections at FAR=1/month in O3a noise and ~48.29% more than the previous ML baseline), and shows robustness to out-of-distribution PSDs and non-Gaussian glitches, while emphasizing the need for careful bias mitigation. The study also provides a thorough ablation and discusses deployment considerations for live data, highlighting practical implications for improving low-latency GW searches with interpretable ML approaches and sharing code for reproducibility.

Abstract

Matched filtering is a long-standing technique for the optimal detection of known signals in stationary Gaussian noise. However, it has known departures from optimality when operating on unknown signals in real noise and suffers from computational inefficiencies in its pursuit of near-optimality. A compelling alternative that has emerged in recent years to address this problem is deep learning. Although it has shown significant promise when applied to the search for gravitational waves (GWs) in detector noise, we demonstrate the existence of learning biases that hinder generalisation and lead to significant loss in detection sensitivity. Our work identifies the sources of a set of 11 interconnected biases present in the supervised learning of the GW detection problem and contributes mitigation tactics and training strategies to concurrently address them. In light of the identified biases, we demonstrate that existing detection sensitivity metrics are not reliable for machine-learning (ML) pipelines and discuss the trustworthiness of previous results. We use GW domain knowledge to build a bespoke ML based binary black hole search pipeline called Sage that addresses these biases. Via the injection study presented in the Machine Learning Gravitational-Wave Search Challenge, we show that Sage detects ~11.2% more signals than the benchmark PyCBC analysis at a false alarm rate of one per month in O3a noise. Moreover, we also show that it can detect ~48.29% more signals than the previous best-performing ML pipeline on the same dataset. We empirically prove that Sage can: [i] effectively handle out-of-distribution noise power spectral densities, [ii] strongly reject non-Gaussian transient noise artefacts, and [iii] achieve higher detection sensitivities using less data than network architectures of a similar size. All code and implementations are available at https://github.com/nnarenraju/sage.

Figures (24)

• Figure 1: Diagram showing the boundaries of the $(m_1, m_2)$ space coloured by chirp mass, $\mathcal{M}_c$. The solid lines show curves of constant chirp time calculated using equation \ref{['eqn:t_from_mc_Cutler_Flanagan']}.
• Figure 2: Flowchart of Sage methodology. The details in the dotted box are set once per run; the signal generation box is specific for generating the signal class; the common settings apply for both the signal and noise class. All boxes with solid lines are iterated over during the training process.
• Figure 3: Amplitude of estimated strain power as a function of frequency for different segments of O3a noise. We compare the estimated PSDs in the training (blue and orange curves) and testing (yellow and purple) datasets for the LIGO Hanford and Livingston detectors. Both plots have an $f_{\text{low}}$ of $15$ Hz at the noise low-frequency cut-off and an $f_{\text{high}}$ of $1024$ Hz at the Nyquist limit of the sampling rate. The solid and dashed black lines show the median PSD of all estimated PSDs in the testing and training dataset, respectively.
• Figure 4: 2D histogram of the joint distribution of component masses, $\textup{P}(m_1, m_2)$ used to generate the training dataset. $10^8$ points were uniformly sampled from $U(\tau_0, \tau_3)$; points that lie within the mass bounds were accepted ($\approx13.4\%$) and the rest were rejected.
• Figure 5: Time spans (in days) for the training, validation, and testing datasets taken from the O3a and O3b observing runs for the H1 and L1 detectors. The O3a noise is provided by the MLGWSC-1 DP5_MLGWSC1_2023 and is time-correlated between the H1 and L1. We downloaded the O3b noise from GWOSC, and it is not time-correlated.