A Deep Learning Framework for Amplitude Generation of Generic EMRIs

Author(s)

Zeng, Yan-bo, Zhang, Jian-dong, Hu, Yi-Ming, Mei, Jianwei

Abstract

One of the main targets for space-borne gravitational wave detectors is the detection of Extreme Mass Ratio Inspirals (EMRIs). The data analysis of EMRIs requires waveform models that are both accurate and fast. The major challenge for the fast generation of such waveforms is the generation of the Teukolsky amplitudes for generic (eccentric and inclined) Kerr orbits. The requirement for the modeling of $\sim10^5$ harmonic modes across a four-dimensional parameter space makes traditional approaches, including direct computation or dense interpolation, computationally prohibitive. To overcome this issue, we introduce a convolutional encoder-decoder architecture for a fast and end-to-end global fitting of the Teukolsky amplitudes. We also adopt a transfer learning strategy to reduce the size of the training dataset, and the model is trained gradually from the simplest Schwarzschild circular orbits to generic Kerr orbits step by step. Within this framework, we obtain a surrogate model based on a semi-analytical Post-Newtonian dataset, and the full harmonic amplitudes can be generated within milliseconds, while the median mode-distribution error for generic orbits is approximately $\sim10^{-3}$. This result indicates that the framework is viable for constructing efficient waveform models for EMRIs.

Figures

Caption

The neural network employs an encoder-decoder architecture to predict the Teukolsky amplitude's modulus and phase in parallel branches. The encoder, a 10-layer residual MLP with Swish activations, maps the four orbital parameters $(a, p, e, x_I)$ to a latent vector. This vector is then upsampled to the full 4D mode-space dimensions using trilinear interpolation, which correspond to $(m, n, k)$ dimensions. A series of 10 residual CNN blocks, featuring 3D convolutions with anisotropic kernels and attention gates, refines the structural tensor. Independent output heads with physically-motivated activations (Softplus for the modulus and Tanh for the phase) produce the predictions, which are passed through a final layer that enforces physical constraints (e.g., $|m| \le \ell$).

Caption

Distribution of orbital parameters in the \textbf{training dataset}. The corner plot shows 1D marginalized histograms (diagonal) and 2D projected distributions (off-diagonal) for the spin $a$, semi-latus rectum $p$, eccentricity $e$, and inclination cosine $x_I$. The color map distinguishes the different orbital geometries as defined in Table~\ref{tab:orbit_classification}. The plot visualizes the dataset's stratified nature, with dense populations corresponding to specific classes like Schwarzschild (SC/SE), Kerr Equatorial (KEC/KEE), Kerr Inclined Circular (KIC), and Kerr Generic (KG) orbits.

Caption

Distribution of orbital parameters in the \textbf{validation dataset}. This dataset consists of randomly drawn samples that were held out from the training process. It covers all orbital classes, providing a robust test of the model's ability to generalize to unseen data.

Caption

The mode-distribution error ($\mathcal{M}_\text{amp}$) categorized by orbital geometry on a logarithmic scale. The shape of each violin shows the probability density of the error, while the inner box plot marks the median and interquartile range.

Caption

Comparison of the predicted and true log-magnitudes for the top 20 dominant modes of a representative Kerr Generic (KG) orbit. The specific orbital parameters are $(a, p, e, x_I) = (0.30, 16.26, 0.20, -0.50)$. The x-axis lists the mode indices $(\ell, m, n, k)$ for each of the 20 modes, ordered by their true amplitude. The y-axis shows the amplitude modulus, $|A_{\ell m n k}|$, on a logarithmic scale. Blue bars represent the true amplitudes from the PN dataset (Reference), while orange bars show the corresponding predictions from our surrogate model (NN). The Mean Absolute Percentage Error (MAPE) for this specific sample is 3.95\%.