Dilated convolutional neural network for detecting extreme-mass-ratio
  inspirals
The detection of Extreme Mass Ratio Inspirals (EMRIs) is intricate due to
their complex waveforms, extended duration, and low signal-to-noise ratio
(SNR), making them more challenging to be identified compared to compact binary
coalescences. While matched filtering-based techniques are known for their
computational demands, existing deep learning-based methods primarily handle
time-domain data and are often constrained by data duration and SNR. In
addition, most existing work ignores time-delay interferometry (TDI) and
applies the long-wavelength approximation in detector response calculations,
thus limiting their ability to handle laser frequency noise. In this study, we
introduce DECODE, an end-to-end model focusing on EMRI signal detection by
sequence modeling in the frequency domain. Centered around a dilated causal
convolutional neural network, trained on synthetic data considering TDI-1.5
detector response, DECODE can efficiently process a year's worth of
multichannel TDI data with an SNR of around 50. We evaluate our model on 1-year
data with accumulated SNR ranging from 50 to 120 and achieve a true positive
rate of 96.3% at a false positive rate of 1%, keeping an inference time of less
than 0.01 seconds. With the visualization of three showcased EMRI signals for
interpretability and generalization, DECODE exhibits strong potential for
future space-based gravitational wave data analyses.
Tianyu Zhao, Yue Zhou, Ruijun Shi, Zhoujian Cao, Zhixiang Ren
# DECODE: DilatEd COnvolutional neural network for Detecting Extreme-mass-ratio inspirals

###### Abstract

The detection of Extreme Mass Ratio Inspirals (EMRIs) is intricate due to their complex waveforms, extended duration, and low signal-to-noise ratio (SNR), making them more challenging to be identified compared to compact binary coalescences. While matched filtering-based techniques are known for their computational demands, existing deep learning-based methods primarily handle time-domain data and are often constrained by data duration and SNR. In addition, most existing work ignores time-delay interferometry (TDI) and applies the long-wavelength approximation in detector response calculations, thus limiting their ability to handle laser frequency noise. In this study, we introduce DECODE, an end-to-end model focusing on EMRI signal detection by sequence modeling in the frequency domain. Centered around a dilated causal convolutional neural network, trained on synthetic data considering TDI-1.5 detector response, DECODE can efficiently process a year's worth of multichannel TDI data with an SNR of around 50. We evaluate our model on 1-year data with accumulated SNR ranging from 50 to 120 and achieve a true positive rate of 96.3% at a false positive rate of 1%, keeping an inference time of less than 0.01 seconds. With the visualization of three showcased EMRI signals for interpretability and generalization, DECODE exhibits strong potential for future space-based gravitational wave data analyses.

Gravitational Wave Deep Learning EMRI

## 1 Introduction

The groundbreaking detection of gravitational waves (GWs) in 2015, exemplified by the GW150914 event, has profoundly impacted the field of astrophysics [1]. Enabled by the Laser Interferemeter Gravitational Wave Observatory (LIGO) [2] and Virgo [3], this remarkable achievement unequivocally confirmed the existence of GWs, providing empirical validation of general relativity (GR) [4]. Beyond enriching our knowledge of the cosmos, this seminal discovery has ushered in a new era of astronomical observation [5]. With the spotlight now turning to space-based GW observatories [6, 7], the absence of terrestrial disturbances allows for a more dedicated exploration of the low-frequency
GWs [8]. This exciting pursuit carries the potential to reveal hitherto unobserved phenomena, offering profound insights into the nature of our universe [5].

Space-based GW detection, a largely unexplored domain, marks the next epoch in astrophysics [6]. Pioneering this exciting venture are projects such as the Laser Interferometer Space Antenna (LISA) [9] by the European Space Agency (ESA), with NASA's participation, and Asian projects including Japan's DECi-hertz Interferometer Gravitational wave Observatory (DECIGO) and B-DECIGO [10; 11], as well as China's Taiji [12] and TianQin [13] missions. Targeting the millimetric frequency band, these endeavors offer a novel perspective for the exploration of diverse astrophysical and cosmological phenomena through the detection of low-frequency GWs [6; 14; 15]. The scientific goals are broad, with the intent to shed light on the enigmas of massive black hole binaries (MBHBs), extreme-mass-ratio inspirals (EMRIs), continuous waves from galactic binaries (GBs), and the stochastic GW backgrounds produced by the early universe's myriad of unresolved sources [16].

In the spectrum of potential discoveries, EMRIs hold a unique position. These events, initiated when a compact stellar remnant spirals into a massive black hole (MBH), these events provide opportunities to investigate the MBH characteristics and the nature of the surrounding environments [17]. EMRIs emit low-frequency GWs throughout their extended inspiral phase, serving as a rich source of information for understanding system physical parameters and the MBH's spacetime geometry [18]. The successful detection and parameter estimation of EMRI signals could provide novel insights into the astrophysics of MBHs and the foundational principles of gravity [19; 20]. Traditional methods for EMRI detection, which include both time-domain and time-frequency-domain techniques, have been widely studied in prior research [21; 22; 23; 24; 25]. These strategies mainly employ matched filtering [21; 24] and the Short Time Fourier Transform [22; 23; 25]. However, the inherent complexities of EMRI signals present significant obstacles. Characterized by their complex waveform templates, high-dimensional parameter space, and multiple modes within a single waveform, EMRI signals require over \(\sim 10^{35}\) templates for matched filtering search [18], resulting in a computationally intensive and time-consuming procedure. An example of single EMRI in both the time and frequency domains can be seen in Figure 1, showcasing the aforementioned challenges of signal detection. Additionally, EMRI signals are typically faint and buried within detector and confusion noise, necessitating extended observation durations to achieve an adequate signal-to-noise ratio (SNR) for detection [18]. Time-frequency techniques, offering representations in both time and frequency domains, are frequently less sensitive than matched filtering, which limits their ability to identify weak signals [25]. Given these challenges, exploring alternative methods, such as deep learning, becomes crucial for potentially improving the efficiency of EMRI signal detection.

Deep learning, an advanced branch of machine learning, employs neural networks with multiple layers for different types of data. By facilitating the extraction of intricate patterns and representations from large datasets, it has played a crucial role in advancing various fields, from image recognition [26] to natural language processing [27]. Among the numerous architectures, the convolutional neural network (CNN) stands out for its proficiency in handling structured data, such as images and time series, by progressively learning features in a hierarchical manner. Starting with simple features like edges in the initial layers, they gradually combine these to recognize more complex patterns and structures in the deeper layers. This layered approach allows CNNs to automatically recognize and represent intricate details

Figure 1: **Visualization of a training data sample. This depicts an EMRI signal from the TDI-A channel spanning 1-year with an SNR of 70. (a), Time-domain representation of the TDI-A strain, showcasing both the combined data (signal + noise) and the signal. The signal’s amplitude is about 3 orders of magnitudes lower than the noise, which makes the detection challenging. (b), Welch PSD of the combined data and the signal, the signal contains lots of modes (peaks), with some reaching the noise level, highlighting the suitability of the frequency domain detection method. The designed detector noise PSD is also presented for reference.**
in the data, making them highly effective for tasks like object detection [28] and time-series classification [29]. In the area of GW data analysis, the potential of deep learning, especially CNNs, is becoming increasingly evident. A large amount of studies [30, 31, 32, 33, 34, 35, 36, 37] have demonstrated their effectiveness in ground-based GW detection. Beyond signal detection, deep learning methods have been applied to a variety of tasks, including parameter estimation [38, 39] and glitch classification [40, 41, 42]. However, the application of these methods to space-based GW detection is still in its early stages. While there have been some exploratory efforts, such as the adoption of MFCNN [32] to detect MBHB contaminated by confusion noise [43] and the application of dictionary learning to low-SNR space-based binary black hole (BBH) detection [44]. Notably, Zhang et al. [45] pioneered the detection of EMRIs using CNN, though without incorporating the time delay interferometry (TDI) technique. Therefore, further research is needed to harness the full capabilities of deep learning in space-based GW analysis.

In this paper, we introduce the DECODE (DilatEd COnvolutional neural network for Detecting Extreme-mass-ratio inspirals), an end-to-end model designed for detecting EMRI signals in the frequency domain with an SNR of around 50. As showed in Figure 2, the model incorporates dilated causal convolutional layers, which expand its receptive field, allowing it to efficiently process data covering an entire year in one pass. We trained our model using synthetic data that considers the TDI-1.5 detector response, accounting for unequal arm lengths. The results are promising: the DECODE detects EMRI signals with a 1-year accumulated SNR between 50 and 120, achieving a true positive rate (TPR) of 96.3% with a false positive rate (FPR) of 1%. Notably, our model can evaluate one batch of data samples within

Figure 2: **Comprehensive EMRI detection framework. (a), Depicts the entire EMRI detection process, from initial data preprocessing to the end-to-end DECODE model. (b), Highlights the mechanism of dilated causal convolution with dilation factors of \((1,2,4,8)\) and a kernel size of 2, emphasizing the exponential growth of the receptive field. (c), Detailed architecture of the residual block in DECODE, comprising two dilated causal convolutional layers, weight normalization, ReLU, and dropout layers. A \(1\times 1\) convolution is introduced to address any dimension discrepancies between the residual input and output.**
seconds. Visualizations of the model's intermediate outputs highlight its interpretable feature extraction process and its ability to generalize beyond GR. These findings emphasize the potential of DECODE in future space-based GW data analyses.

The remainder of this paper is organized as follows: Section 2 provides a detailed overview of the data generation procedure and outlines the architecture of our proposed model, the DECODE. In Section 3, we present the results of our EMRI detection experiments, demonstrating the effectiveness of our approach. Finally, Section 4 concludes the paper with a summary of our findings and a discussion on potential future work based on our findings.

## 2 Method

### EMRI Waveform Modeling

Detecting EMRIs has the potential to reveal key astrophysical insights, but modeling their waveform is challenging due to the delicate balance of strong-field GR and gravitational radiation dynamics. Accurately describing EMRIs demands a solution to the self-force problem, which considers the gravitational impact of the smaller compact object on its own motion within the powerful gravitational field of the central MBH [6]. Because the self-force problem is highly non-linear and defies analytical solutions, researchers have developed approximate waveform models, commonly referred to as kludge models [46, 47].

Two commonly used kludge models in EMRI modeling are the "analytic kludge" (AK) [46] model and the numerical kludge (NK) [47] model. The AK model relies on post-Newtonian expansions and perturbative calculations to evolve the orbital parameters and generate waveforms quickly. It provides computational efficiency but suffers from dephasing compared to more accurate models, leading to potential inaccuracies in parameter estimation. On the other hand, the NK model incorporates the orbital trajectory computed in curved space using Kerr geodesics and includes radiation reaction effects. Although more accurate, the NK model is computationally more expensive, making EMRI signal detection using this template highly formidable.

To address the limitations of both models, an argumented analytic kludge (AAK) [48, 49, 50] model has been proposed. The AAK model combines the computational efficiency of the AK model with improved phasing achieved through a mapping to Kerr geodesic frequencies and self-consistent post-Newtonian evolution. By incorporating self-force information and refining the phasing, the AAK model achieves higher waveform fidelity compared to the AK model while remaining computationally efficient. While its computational efficiency may not be adequate for matched filtering-based signal searches, it is suitable for producing training datasets for deep neural networks (DNNs).

Despite the advancements in kludge waveform modeling, challenges remain. Incorporating second-order self-force effects into the models and refining them for orbits approaching plunge are ongoing areas of research [6]. Nonetheless, these waveform models are crucial for accurately representing the dynamics of EMRIs and enabling the detection, parameter estimation, and data analysis of these elusive astrophysical sources.

### Data Curation

The process of curating training and testing datasets for the identification of EMRI signals using a DNN is a multi-step procedure consisting of signal generation, detector response simulation, and pre-processing.

Waveform GenerationThe first step involves the generation of signal templates. The AAK model used for generating these templates is based on [51]. The waveform, denoted as \(h(t)=h_{+}(t)-ih_{\times}(t)\), is typically characterized by 14 physical parameters. The parameter space used for sampling the training and testing dataset parameters in this study is detailed in Table 1. Here, \(M\) and \(a\) represent the mass and the spin parameter of the MBH respectively. The semi-latus rectum is denoted by \(p\), while \(e\) stands for orbital eccentricity, and \(\iota\) signifies the orbit's inclination angle from the equatorial plane. \(Y=\cos\iota\equiv L_{z}/\sqrt{L_{z}^{2}+Q}\), where \(Q\) is the Carter constant, and \(L_{z}\) is the \(z\) component of the specific angular momentum. The polar and azimuthal sky location angles are represented by \(\theta_{S}\), and \(\phi_{S}\), respectively. The orientation of the spin angular momentum vector of the MBH is described by the azimuthal and polar angles \(\theta_{K}\) and \(\phi_{K}\). These parameters are uniformly sampled for our dataset. It is important to note that \(\Phi_{\varphi,0},\Phi_{\theta,0},\Phi_{r,0}\), which represent the phase of azimuthal, polar, and radial modes, are all manually set to 0 respectively.

TDI ResponseThe next stage involves simulating the detector's response to these signals. The specific detector configurations utilized in this study are detailed in Table 2. For the breathing arm length, we employed the TDI-1.5 technique, which yielded the GW strain of TDI A and E channels, denoted as \(h_{A}(t)\) and \(h_{E}(t)\), respectively. A detailed derivation of this technique can be found in Ref. [52]. Their CUDA-based implementation, enable us to calculate the
response cost in seconds. The signal is then rescaled according to the desired SNR using the formula:

\[\mathrm{SNR}^{2}=(h_{A}\mid h_{A})+(h_{E}\mid h_{E}). \tag{1}\]

Here, the inner product \((a\mid b)\) is defined as:

\[(a\mid b)=2\int_{f_{min}}^{f_{max}}\frac{\tilde{a}^{*}(f)\tilde{b}(f)+\tilde{a} (f)\tilde{b}^{*}(f)}{S_{n}(f)}\ \mathrm{d}f. \tag{2}\]

In this equation, \(f_{min}=\frac{1}{\text{Duration}}\simeq 3.17\times 10^{-8}\,\mathrm{Hz}\) and \(f_{max}=\frac{1}{2\cdot\text{Calence}}=\frac{1}{30}\mathrm{Hz}\). \(\tilde{a}(f)\) and \(\tilde{b}(f)\) represent the frequency domain signals, and the superscript \(*\) denotes the complex conjugate. \(S_{n}(f)\) is the one side noise power spectral density (PSD), which will be specified later.

Noise GenerationThe third step introduces noise to the signal. This noise, \(n(t)\), is modeled as a colored Gaussian noise with a PSD defined by

\[\mathrm{S}_{n}(f)=16\sin^{2}(\omega L)\left(P_{\mathrm{oms}}(f)+(3+\cos(2 \omega L))P_{\mathrm{acc}}(f)\right)\, \tag{3}\]

\begin{table}
\begin{tabular}{l r r} \hline \hline
**Parameter** & \multicolumn{2}{c}{**Lower bound**} & \multicolumn{1}{c}{**Upper bound**} \\ \hline \hline \(\log_{10}(M/M_{\odot})\) & \(5\) & \(8\) \\ \(a\) & \(10^{-3}\) & \(0.99\) \\ \(e_{0}\) & \(10^{-3}\) & \(0.8\) \\ \(p_{0}/M\) & \(15\) & \(25\) \\ \(Y_{0}\) & \(-1\) & \(1\) \\ SNR & \(50\) & \(120\) \\ \(\theta_{S}\) & \(0\) & \(\pi\) \\ \(\phi_{S}\) & \(0\) & \(2\pi\) \\ \(\theta_{K}\) & \(0\) & \(\pi\) \\ \(\phi_{K}\) & \(0\) & \(2\pi\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Summary of parameter setups in EMRI signal simulation.

\begin{table}
\begin{tabular}{l r} \hline \hline
**Parameter** & \multicolumn{1}{c}{**Configuration**} \\ \hline \hline Size of training dataset & \(5000\) \\ Size of testing dataset & \(1000\) \\ Cadence & \(15\,\mathrm{s}\) \\ Duration & \(1\,\mathrm{year}\) \\ Re-sampled data length \(N\) & \(1024/2048/4096\) \\ \hline Arm length \(L\) & \(2.5\times 10^{9}\,\mathrm{m}\) \\ Detector orbit & \(1\)st order Keplerian orbit \\ TDI & \({}^{*}\,\mathrm{TDI}\)-1.5 \\ Acceleration noise \(A_{\mathrm{acc}}\) & \(3\,\mathrm{fm}/\sqrt{\mathrm{Hz}}\) \\ OMS noise \(A_{\mathrm{oms}}\) & \(15\,\mathrm{pm}/\mathrm{Hz}\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Summary of configurations of training and testing dataset.
with

\[\begin{split} P_{\text{oms}}(f)&=A_{\text{oms}}^{2} \left[1+\left(\frac{2\,\mathrm{mHz}}{f}\right)^{4}\right]\left(\frac{2\pi f}{c} \right)^{2}\,,\\ P_{\text{acc}}(f)&=A_{\text{acc}}^{2}\left[1+ \left(\frac{0.4\,\mathrm{mHz}}{f}\right)^{2}\right]\\ &\cdot\left[1+\left(\frac{f}{8\,\mathrm{mHz}}\right)^{4}\right] \left(\frac{1}{2\pi fc}\right)^{2}\,.\end{split} \tag{4}\]

Where \(A_{\text{acc}}\) and \(A_{\text{oms}}\) are the noise budget of test mass acceleration noise and readout noise coming from the optical metrology system (OMS), \(L\) is the arm length of LISA detector, and \(c\) is the speed of light. Then the signal is injected into the noise, resulting in the synthetic data Figure 1 is a showcase of the training data in the time and frequency domain.

Whitening and PSD EstimationIn the final stage of data curation, the data undergoes several pre-processing steps to prepare it for input into the DNN. The first of these steps is whitening, which serves to remove the frequency-dependent variations in the noise. This process allows the DNN to concentrate on the underlying signal patterns, simplifying the learning task and enhancing the network's ability to detect subtle patterns in the data, thereby improving the overall performance of the EMRI signal identification. Following whitening, the PSD of the data is estimated using Welch's method. The data then undergoes sub-sampling, where it is re-sampled onto a log-uniform frequency grid. This step is aimed at reducing the computational load of subsequent analyses by decreasing the number of data points. 3 different grid density is selected as listed in Table 2. The final pre-processing step is standardization, which ensures that all input features are on a uniform scale, a fundamental requirement for most deep learning algorithms. This step is crucial in enhancing the learning efficiency of the neural network and improving the overall performance of the model.

### Decode

In this work, we introduce the DECODE, a novel architecture for sequence modeling tasks, as illustrated in Figure 2. The DECODE is inspired by the TCN architecture [53], which has been shown to outperform traditional recurrent architectures across a diverse range of tasks and datasets. The DECODE architecture leverages the strengths of convolutional networks, which have been proven to be highly effective for sequence modeling. It incorporates dilated convolutions, which are a powerful tool for capturing long-range dependencies in sequence data. The causal nature of the DECODE ensures that the model's output at each step is conditioned on all previous steps, making it suitable for tasks that require an understanding of sequential dependencies. While the TCN and other sequence modeling architectures have predominantly been applied to time series data, the DECODE stands out in its application to frequency domain data. Detecting EMRI in the time domain presents challenges due to the extended duration of the signals and their low SNR. As illustrated in Figure 0(a), the amplitude of the signal is typically three orders of magnitude lower than the noise, and the data spans a full year. However, as shown in Figure 0(b), in the frequency domain, the signal's PSD

Figure 3: **EMRI detection performance across SNR and \(N\). All sub-plots depict receiver operating characteristic (ROC) curves for distinct input sample lengths \(N\) within specific SNR ranges, presented on a logarithmic scale. Each line style signifies the balance between TPR and FPR for a given sample length, with the area beneath each curve representing the model’s efficacy. A reference yellow dashed line indicates the random prediction. The use of logarithmic scales enhances the visibility of performance difference, especially at lower FPR levels. (a), Evaluation for \(\mathrm{SNR}\in[50,120]\). (b), Evaluation for \(\mathrm{SNR}\in[70,170]\). (c), Evaluation for \(\mathrm{SNR}\in[100,240]\).**
has lots of peaks, with some even reaching the noise level. Despite this shift from time to frequency domain, the core principles of sequence modeling remain applicable. The DECODE effectively exploits these principles, achieving notable performance in EMRI signal detection.

Causal Sequence ModelingThe DECODE framework is designed for sequence modeling, with a focus on maintaining causality throughout its structure. Central to DECODE's design are two fundamental principles. Firstly, the architecture ensures that the output sequence's length aligns with the input sequence. This alignment is achieved via a 1D-convolutional network design, where each hidden layer matches the length of the input layer. To maintain this length consistency, zero padding of length (\(\text{kernel size}-1\)) is applied. Following this, the architecture emphasizes the causality of the sequence. This is achieved by using causal convolutions, which ensure that the output at a particular time step is convolved only with preceding elements in the previous layer.

Dilated ConvolutionIncorporated into the DECODE architecture, dilated convolutions play a pivotal role in capturing long-range dependencies in sequence data. Drawing inspiration from the WaveNet [54], the DECODE employs dilated convolutions to exponentially expand the receptive field without a significant increase in computational complexity or number of parameters. We provide an illustration in Figure 1(b), More formally, for a 1-D sequence input \(\mathbf{x}\in\mathbb{R}^{n}\) and a filter \(f:\{0,...,k-1\}\longrightarrow\mathbb{R}\), the dilated convolution operation \(F\) on element \(s\) of the sequence is defined as:

\[F(s)=\left(\mathbf{x}\ast_{d}f\right)(s)=\sum_{i=0}^{k-1}f(i)\cdot\mathbf{x}_ {s-d\cdot i}\,, \tag{5}\]

where \(d\) is the dilation factor, \(k\) is the filter size (i.e. kernel size), and \(s-d\cdot i\) accounts for the direction of the past. When \(d=1\), a dilated convolution reduces to a regular convolution. By employing larger dilations, the receptive field of a DECODE is effectively expanded, allowing it to capture long-range dependencies within the sequence data more effectively.

Residual ConnectionsResidual connections, another key feature of the DECODE architecture, are designed to facilitate the training of deep networks. These connections, introduced by He et al. [55], allow the gradient to flow directly through the network, mitigating the problem of vanishing gradients that often jeopardize deep networks. In the DECODE, a residual block is composed of two dilated causal convolutional layers, with a residual connection skipping over them. If we denote the input to the residual block as \(\mathbf{x}\), the output of the block, \(\mathbf{y}\), can be computed as:

\[\mathbf{y}=\mathrm{Activation}(\mathbf{x}+\mathcal{F}(\mathbf{x}))\,, \tag{6}\]

where \(\mathcal{F}(\mathbf{x})\) represents the transformations performed by the dilated causal convolutional layers. This design choice has been shown to improve the performance of deep networks and is a key component of the DECODE architecture. The residual block used in the DECODE model is illustrated in Figure 1(c). Each block comprises two layers of dilated causal convolution, followed by the rectified linear unit (ReLU) activation function. Weight normalization [56] and dropout [57] are incorporated after each dilated convolution within the residual block.

Loss FunctionIn our DECODE model, the output of the residual block has a shape of \((H,N)\), where \(H\) represents the hidden size of our model and \(N\) is the length of the input sequence. The last column of this output is then passed through a linear layer to generate the predicted probability for EMRI signal detection. To train the model, we use the cross-entropy loss, a common choice for classification tasks. One of the advantages of using the cross-entropy loss is its ability to accelerate convergence during training, especially when compared to other loss functions like mean squared error [58]. The cross-entropy loss for a binary classification problem is given by:

\[\mathcal{L}=-\frac{1}{n}\sum_{i=1}^{n}y_{i}\log(\mathcal{P}_{i})+(1-y_{i})\log (1-\mathcal{P}_{i})\,. \tag{7}\]

In this equation, \(y_{i}\) denotes the actual label, while \(\mathcal{P}_{i}\) is the predicted probability for the \(i\)-th sample, with \(n\) representing the total number of samples in the training dataset. The cross-entropy loss quantifies the divergence between the actual label and the predicted probability.

### Implementation Detail

For waveform generation of training data, we employed FastEMRIWaveform2[50, 51] for EMRI signal creation and lisa-on-gpu3[52] for GPU-accelerated detector response simulations, which includes TDI. We also integrated
additional functionalities from the SciPy library. Our DECODE architecture consists of 10 residual blocks, each with a kernel size of 3 and a hidden size of 128. Developed using the PyTorch framework, known for its computational efficiency and speed, computations were performed on a high-performance computing cluster equipped with NVIDIA Tesla V100 GPUs. The training utilized the Adam optimizer with a learning rate of \(2\times 10^{-4}\) and a batch size of 64.

## 3 Results

### EMRI Detection Proficiency

Receiver operating characteristic (ROC) curve and the area under the curve (AUC) are essential tools for evaluating the performance of models in binary classification tasks. In the context of our study, where the task is to detect EMRI signals buried in noise, these tools provide valuable insights. The ROC curve, which plots the TPR against the FPR, offers a visual representation of the model's performance across various threshold settings. The AUC, on the other hand, provides a single, overall measure of the model's performance across all thresholds. A model with perfect discrimination has an AUC of 1, while a model performing no better than random guessing has an AUC of 0.5.

In our research, we employ ROC curves as the primary benchmark to quantify the performance of the DECODE. Our test dataset used here is generated like the training datasets, i.e. the waveform parameters are uniformly distributed as shown in Table 1 but with different SNR range. As depicted in Figures 2(a) to 2(c), we show three separate ROC curves, with each corresponding to a unique input sample length fed into the DECODE. For the specified input lengths of \(N=(1024,2048,4096)\), the SNR ranges are set at \([50,120]\), \([70,170]\), and \([100,240]\). The associated AUC values, detailed within the figures, offer quantitative insight into the model's sensitivity in detecting EMRI signals. For clarity in visual representation, especially at lower FPR values, Figure 3 adopt a logarithmic scale for their axes.

It's noteworthy that our test dataset comprises signals with a duration of 1 year, achieving twice the SNR compared to the 3-month data scenario presented in Ref. [45]. While their study tested models on datasets with \(\mathrm{SNR}\in[50,120]\), we evaluated ours on datasets with \(\mathrm{SNR}\in[100,240]\). Both datasets, when rescaled for a 1-year duration, maintain equivalent SNR values, implying consistent signal amplitudes. Impressively, our model attains a TPR of 97.5% at a FPR of 1% as showcased in Figure 2(c).

One significant advantage of deep learning methods over matched filtering-based approaches is their speed. Once trained, the model can be rapidly deployed for inference. In our tests, conducted on a single NVIDIA Tesla V100 GPU, our model processed 2000 data samples in approximately 4 seconds, amounting to less than \(10^{-2}\) seconds per sample.

### EMRI Detection Efficacy

In Figure 4, we provide a detailed examination of the DECODE's performance across different physical parameters. Figure 3(a) illustrates the relationship between TPR and SNR. The sub-figure clearly demonstrates that as the SNR increases, the TPR increases correspondingly, particularly at the specified FPR thresholds of 0.10 and 0.01.

Figure 4: **Detection capability of DECODE across various parameters. (a), Illustrates the TPR as a function of SNR, highlighting the model’s capability to detect signals with varying strengths. (b) Showcases the TPR plotted against the relative amplitude \(\mathcal{A}\) (defined in eq. (8)), emphasizing the model’s ability to detect power excesses in the frequency domain and detect signals even when they are submerged within the noise. (c) Explores the TPR in relation to the spin parameter \(a\), keeping the MBH mass consistent at \(10^{6}M_{\odot}\). This sub-figure is evaluated at three distinct SNR levels: 50, 70, and 100, shedding light on the relationship between spin parameters and detection capabilities.**
To gain a deeper understanding of the sensitivity of our model, we introduce the relative amplitude, denoted as \(\mathcal{A}\). It is defined as:

\[\mathcal{A}=\max_{i\in A,E}\sqrt{\frac{S_{h}^{i}(f)}{S_{n}(f)}}\,, \tag{8}\]

where \(S_{h}^{i}\) represents the Welch PSD of waveform \(h_{i}\). This metric effectively captures the signal's amplitude in the frequency domain. Figure 3(b) plots the TPR against the relative amplitude, at FPRs of 0.1 and 0.01, this sub-figure presents the model's proficiency in discerning power exceeds in the frequency domain. Notably, the DECODE can also detect signals that are entirely submerged within the noise.

In Figure 3(c), we evaluate the DECODE's sensitivity to varying spin parameters, while keeping the MBH mass constant at \(10^{6}M_{\odot}\). The evaluation, performed at SNR levels of 50, 70, and 100 and FPR thresholds of 0.1 and 0.01, indicates that the model's detection performance is mainly influenced by the SNR. In contrast, the spin parameter appears to have a limited effect on detection, suggesting that the spin parameter contribution to the overall strength of the EMRI signal is relatively minor.

### Interpretability

CNN-based models are powerful tools for pattern recognition and prediction. Their unique architecture and operational mechanism make them inherently interpretable, a feature that is particularly valuable in interdisciplinary research. CNN-based models learn hierarchical patterns in the data through their convolutional layers, with each layer extracting a set of high-level features from the input data. These features are then used by subsequent layers to understand more complex patterns. This transparent process of feature extraction can be visualized, providing insight into how the network interprets the data and makes predictions.

The activation maps, often used in the context of neural networks, provide a visual representation of the features that the model identifies and emphasizes during its processing. Essentially, they capture the output values or "activations" from various layers or blocks within the network when presented with an input. These maps offer insights into which parts of the input data the model finds significant or relevant for a particular task. In the case of the DECODE, the activation maps generated at the output of each residual block reveal how the model processes and interprets the frequency-domain data of EMRI signals.

The activation maps illuminate the interpretability of the DECODE. By analyzing the outputs of multiple residual blocks, the processes of feature extraction are made transparent. Figure 5 provides a detailed visualization of these maps, demonstrating the ability of the DECODE to distinguish EMRI signals from noise. Specifically, panel **i** of each sub-figure depicts activation maps for inputs with an EMRI signal, while panel **iii** depicts the corresponding frequency domain data. These maps emphasize activated neurons in regions that correspond to the frequency components of the signal. In contrast, panel **ii** of each sub-figure depicts diminished activations for noise-only samples. The corresponding frequency domain data for these samples is presented in panel **iv**, validating the model's ability at identifying EMRI signals.

### Generalization Ability

Generalization ability is the capacity of a model trained on a specific dataset to perform well on new, untrained data. It indicates how well a model can extrapolate from its training data to make accurate predictions on unknown data. In practical applications, a model will frequently be presented with data that differs from its training set, so this ability is crucial. A model that generalizes well is robust and flexible, ensuring that it does not simply memorize the training data but rather understands inherent patterns and relationships.

In Figure. 4(b) and 4(c), we provide evidence of the generalization capabilities of our model. Even though the model was only trained on AAK waveform datasets, it identified the AK waveform accurately during evaluation with the output probability equal to 1, demonstrating its ability to generalize across various waveform templates. In contrast, the model's successful detection of the XSPECG waveform [59, 60], which was formulated using the KRZ metric, demonstrates its generalization ability with respect to various gravitational theories. These results demonstrate the generalization ability of the model, suggesting that it is capable of handling scenarios beyond its training datasets.

## 4 Conclusion and Discussion

The detection of EMRIs in gravitational wave astronomy presents a formidable challenge. In this paper we introduce the DECODE, a state-of-the-art end-to-end DNN model designed for the detection of EMRI signals in the frequency
Figure 5: **Interpretability and generalization ability showcase.** This figure provides an in-depth visualization of the intermediate outputs from each residual block, demonstrating the model’s capability for feature extraction within the frequency domain and it’s generalization ability to different waveform templates and gravitational theories. For each sub-figure, panels **i** and **ii** represent the intermediate results corresponding to the input data samples shown in panels **iii** and **iv**. In contrast to the faint activations in panel **ii**, the noticeable activated neurons in panel **i** indicate the extraction of essential characteristics when a signal is present in the input. (a), AAK waveform. (b), AK waveform. (c), XSPEC waveform.
domain. By leveraging dilated causal convolutional layers, the DECODE efficiently processes year-long data. Our evaluations on synthetic datasets have revealed the model's robustness and efficiency, achieving remarkable detection rates at varied SNR levels. Furthermore, the model's rapid inference capabilities and its ability to generalize beyond its training parameters but there is still room for future advancement.

The precision of the EMRI detection model is intrinsically related to the precision of the training data. While our current training dataset employs the TDI-1.5 detector response, future developments could benefit from the incorporation of more sophisticated simulations, such as the TDI-2.0 technique. This would provide a more accurate simulation of the detector's response, potentially enhancing the model's applicability in the actual world.

Our current approach primarily focuses on the amplitude information of the EMRI signals. However, the phase information, which has been largely wasted in this research, holds considerable potential. By integrating phase-related features into the model, we could capture more intricate patterns and details of the EMRI signals. This may lead to improved detection rates and lower false alarm rates.

In conclusion, DECODE is a step forward in EMRI detection. Even though there are avenues for improvement, its foundational accomplishments demonstrate its potential as a tool for future space-based GW data analyses.

## 5 Acknowledgments

The research was supported by the Peng Cheng Laboratory and by Peng Cheng Laboratory Cloud-Brain. This work was also supported in part by the National Key Research and Development Program of China Grant No. 2021YFC2203001 and in part by the NSFC (No. 11920101003 and No. 12021003). Z.C was supported by the "Interdisciplinary Research Funds of Beijing Normal University" and CAS Project for Young Scientists in Basic Research YSBR-006.
