We present a physicsinspired transformer model that predicts the nonlinear dynamics of higherorder wave modes emitted by quasicircular spinning nonprecessing binary black hole mergers. The model forecasts the waveform evolution from the premerger phase through the ringdown starting with an input timeseries spanning  t in 5000textrmM 100textrmM . The merger event defined as the peak amplitude of waveforms that include the l  m  2 modes occurs at  t  0textrmM . The transformer then generates predictions over the time range  t in 100textrmM 130textrmM . We produced training evaluation and test sets using the NRHybSur3dq8 model considering a signal manifold defined by mass ratios  q in 1 8  spin components  sz12 in 0.8 0.8  modes up to l leq 4 including the 55 mode but excluding the 40 and 41 modes and inclination angles theta in 0 pi. We trained the model on 14440761 waveforms completing the training in 15 hours using 16 NVIDIA A100 GPUs in the Delta supercomputer. We used 4 H100 GPUs in the DeltaAI supercomputer to compute within 7 hours the overlap between ground truth and predicted waveforms using a test set of 840000 waveforms finding that the mean and median overlaps over the test set are 0.996 and 0.997 respectively. Additionally we conducted interpretability studies to elucidate the waveform features utilized by our transformer model to produce accurate predictions. The scientific software used for this work is released with this manuscript.
['General Relativity and Quantum Cosmology', 'Astrophysics - Instrumentation and Methods for Astrophysics', 'Computer Science - Artificial Intelligence', '68T10', '85-08', '83C35', '83C57', 'I.2']
AI forecasting of higherorder wave modes of spinning binary black hole mergers
{'AI forecasting of higher-order wave modes of spinning binary black hole mergers': 'Victoria Tiki ab , Kiet Pham c , and E. A. Huerta abd a Data Science and Learning Division, Argonne National Laboratory, Lemont, Illinois 60439, USA b Department of Physics, University of Illinois Urbana-Champaign, Urbana, Illinois 61801, USA c School of Physics and Astronomy, University of Minnesota, 55455 Minnesota, USA d Department of Computer Science, University of Chicago, Chicago, Illinois 60637, USA', 'Abstract': 'We present a physics-inspired transformer model that predicts the non-linear dynamics of higher-order wave modes emitted by quasi-circular, spinning, non-precessing binary black hole mergers. The model forecasts the waveform evolution from the pre-merger phase through to the ringdown, starting with an input time-series spanning t ∈ [ -5000 M , -100 M ) . The merger event, defined as the peak amplitude of waveforms that include the ℓ = | m | = 2 modes, occurs at t = 0 M. The transformer then generates predictions over the time range t ∈ [ -100 M , 130 M ] . Training, validation, and test sets were generated using the NRHybSur3dq8 model, considering a signal manifold defined by mass ratios q ∈ [1 , 8] ; spin components s z { 1 , 2 } ∈ [ -0 . 8 , 0 . 8] ; modes up to ℓ ≤ 4 , including the (5 , 5) mode but excluding the (4 , 0) and (4 , 1) modes; and inclination angles θ ∈ [0 , π ] . We trained the model on 14,440,761 waveforms, completing the training in 15 hours using 16 NVIDIA A100 GPUs in the Delta supercomputer. To evaluate the model, we utilized 4 H100 GPUs in the DeltaAI supercomputer to compute, within 7 hours, the overlap between ground truth and predicted waveforms on a test set of 840,000 waveforms. The mean and median overlaps achieved were 0.996 and 0.997, respectively. Additionally, we conducted interpretability studies to elucidate the waveform features that our transformer model utilizes to produce accurate predictions. The scientific software used for this work is released with this manuscript.  \nKeywords: AI, Transformers, Higher-order gravitational wave modes, Black Hole Mergers', '1. Introduction': 'Artificial Intelligence (AI) is being explored and applied in earnest to address computational grand challenges in gravitational wave astrophysics [13]. Contemporary AI applications cover the entire spectrum of gravitational wave data science, ranging from data quality [4-8], detection [9-17], parameter estimation [18-24], denoising [6, 25-27], detection [28-32] and forecasting [33-35] of sources with potential electromagnetic counterparts, among many others. Ref. [36] provides a comprehensive compilation of recent AI advances in gravitational wave astrophysics.  \nIn this article we study whether AI can learn and predict the physics of higher-order wave modes emitted by the merger of quasi-circular, spinning, non-precessing binary black hole mergers. This work is inspired by previous studies demonstrating that AI can forecast the time-series signals of the ℓ = | m | = 2 mode of quasi-circular, non-spinning binary black hole mergers [37], and of quasi-circular, spinning, non-precessing binary black hole mergers [38]. In practice, these two studies [37, 38] demonstrated that AI can learn the key features that define the strongly non-linear, highly dynamical behaviour of binary black hole mergers, and then predict the time-series evolution of the pre-merger, merger and ringdown phases. This paper explores the use of transformers to learn the non-linear properties of higher-order wave modes and discusses methods for training these models on supercomputers with millions of waveforms.  \nTo gain a basic understanding of the type of physics that physics-inspired AI models need to master in the context of higher-order wave modes, we present in Figure 1 a comparison between former studies and this work. The left column of Figure 1 displays samples of waveforms used in previous work [37, 38], which have a very smooth evolution throughout inspiral, merger and ringdown. Although learning the physics of these waveforms poses challenges, these signals are relatively simple compared to the significantly more complex waveforms analyzed in this work, as illustrated in the right panel of Figure 1. In this case, the amplitude evolution is no longer monotonic up until merger, at t = 0 M, followed by a smooth decay. Rather, these waveforms show a complex structure due to the inclusion of higher-order wave modes, which  \nalso shift the location of the amplitude peak. For any set of parameters ( q, s z 1 , s z 2 ) , the inclination angle, θ , modifies the physics of these waveforms significantly. For this reason, we densely sample the signal manifold with over 14 million waveforms to capture the dynamics of these higher-order wave modes. Our transformer model is designed to learn and forecast these intricate physical dynamics.  \nFigure 1: Normalized amplitude | h | / | h | max of gravitational waves around merger for different binary black hole configurations. Notice the impact of higher-order wave modes for inclination angles θ = 0 ( left column) and θ = π/ 4 (right column).  \n<!-- image -->  \nAs we demonstrate below, the algorithms we introduce in this work enable physics-inspired AI to learn and predict the nonlinear, highly dynamical  \nevolution of higher-order wave modes with state-of-the-art accuracy. Key results of this work encompass:  \n- 1. A novel transformer model that is trained using the higher-order wave modes ℓ ≤ 4 and (5 , 5) , except for (4 , 0) and (4 , 1) .\n- 2. A new transformer architecture that takes as input data the two gravitational wave polarizations ( h + , h × ) , as given by h ( t ) = h + -ih × , and forecasts time-series data for ( h + , h × ) from the pre-merger phase through to ringdown.\n- 3. A distributed training method that leverages 4 nodes, equivalent to 16 NVIDIA A100 GPUs, in the Delta supercomputer to train a transformer model with optimal forecasting performance within 15 hours using a training dataset of 14,440,761 million waveforms-the largest dataset used to date to train AI models for gravitational wave physics applications.\n- 4. For inference, we used 4 NVIDIA H100 GPUs to quantify, within 7 hours, the performance of our transformer model using a test set of 840,000 waveforms. We found that the mean and median overlap between ground truth and predicted signals are 0.996 and 0.997, respectively, over the parameter space under consideration.\n- 5. Open source scientific software to reproduce the key results of this article, available in GitHub [39].  \nThis paper is organized as follows. Section 2 describes the data used to train, validate and test our transformer model. It also provides a description of the physics-inspired design of our transformer model. We then discuss the training approach, and present inference results in Section 3. We summarize our key findings in Section 4. Future directions of work are outlined in Section 5. The numerical relativity waveforms used to train the NRHybSur3dq8 surrogate model use geometric units in which G = c = 1 , and M represents the total mass of the binary black-hole merger spacetime. Thus, the waveform datasets we use throughout this work follow the same convention, and we measure time in units of M, with M = 1 M ⊙ = 4 . 93 × 10 -6 s.', '2.1. Data': "We produced three independent datasets of modeled waveforms using the numerical relativity (NR) surrogate model NRHybSur3dq8 [40]. These wave-  \nforms describe the gravitational emission from quasi-circular, spinning, nonprecessing binary black hole mergers throughout the inspiral, merger, and ringdown phases. The datasets are constrained within the valid parameter space for this NR surrogate model, i.e., mass-ratios q ≤ 8 and individual spins ( | s z i | ≤ 0 . 8 for i = { 1 , 2 } ). The gravitational wave strain time-series, h , can be expressed as a sum of spin-weighted spherical harmonic modes, h lm , on the 2-sphere [41]  \nh ( t, θ, ϕ ) = ∑ l =2 m = l ∑ m = -l h lm ( t ) -2 Y lm ( θ, ϕ ) , (1)  \nwhere -2 Y lm are the spin-weight-2 spherical harmonics, θ is the inclination angle between the orbital angular momentum of the binary and line of sight to the detector, and ϕ is the initial binary phase, which we set to zero. We include the higher-order wave modes ℓ ≤ 4 and ( ℓ, m ) = (5 , 5) , except for (4,0) and (4,1). The waveforms cover the timespan t = [ -5000 M , 130 M ] .  \nFigure 2: Sampling of the signal manifold defined by mass-ratios q ∈ [1 , 8] , individual spins s z { 1 , 2 } ∈ [ -0 . 8 , 0 . 8] , and inclination angle θ ∈ [0 , π ] to produce three independent datasets to train, validate and test our transformer models. The training set is produced using step sizes of { δq, δs z i , δθ } = { 0 . 1 , 0 . 1 , π/ 29 } . The validation and test sets are generated by alternately sampling values between the training set values.  \n<!-- image -->  \nTraining dataset. The training dataset consists of 14,440,761 waveforms, generated by sampling the mass ratio q ∈ [1 , 8] in steps of δq = 0 . 1 ; individual  \nspins s z i ∈ [ -0 . 8 , 0 . 8] in steps of δs = 0 . 1 ; and inclination angle θ ∈ [0 , π ] in steps of δθ = π/ 29 .  \nTest and validation datasets. Each of these sets consists of 840,000 waveforms generated by sampling values spaced evenly between the training set values. Figure 2 illustrates the sampling strategy for generating the training, validation, and test sets, demonstrating that these datasets are independent and do not overlap.  \nFigure 3: Gravitational wave sequence and separation into encoder input, decoder input, and decoder target for training. Note that the decoder target complements the encoder input exactly, reflecting the transformer's intended function of completing the signal. The decoder input sequence is shifted to the left by ∆ t = 1 M with respect to the decoder target sequence and thus overlaps the encoder input by one timestep.  \n<!-- image -->  \nSampling and encoder-decoder split. We reduce the temporal resolution of the data by subsampling the original input sequences, selecting every other timestep. This approach reduces computational complexity. The modeled waveform is then split up such that the transformer's encoder module receives  \nthe early inspiral sequence consisting of time steps [ -5000 M , -100 M ] . The decoder module receives the inspiral, merger and ringdown segments made up of time steps [ -101 M , 129 M ] . The target waveform then consists of the gravitational wave segment [ -100 M , 130 M ] , i.e. the decoder target is onetime step ahead of the decoder input, as shown in Figure 3.", '2.2. Architecture': 'Our transformer model is based on the original architecture introduced in Ref. [42], which leverages attention mechanisms to weigh different segments of the input sequence and improve performance in sequence prediction tasks. We have modified the model to incorporate specific properties of gravitational waves. Figure 4 schematically illustrates these modifications. Below, we describe the functionality of each component of the transformer as it processes the gravitational wave sequences. We represent the amplitudes of the input sequences for the encoder and decoder modules by x p , where p indicates the position within the sequence of length n . Note that the sequence length n varies for encoder and decoder inputs, denoted as n enc and n dec , respectively.', '2.2.1. Embedding': "The transformer architecture is not intrinsically sensitive to the order of tokens within the input sequence. However, for sequence prediction, information about token order is critical. This issue is addressed by introducing non-trainable positional encoding which embeds the input sequence into a higher-dimensional space. Through positional encoding, the input sequence x p is transformed into a sequence of d embed dimensional vectors each taking the form X p = ( x p , PE ( p, 1) , PE ( p, 2) , ..., P E ( p, d embed -1 )) . This transformation carries information about the original token's position. We represent this transformed set of vectors as a matrix X ∈ R n × d embed . The embedding function PE , where p is the position of any token within the original sequence and i is the dimension within the positional encoding vector, takes the form:  \nPE ( p, 2 i ) = sin ( p 10000 2 i/d embed ) , (2)  \nPE ( p, 2 i +1) = cos ( p 10000 2 i/d embed ) . (3)  \nAs demonstrated in Ref. [42], this type of sinusoidal positional encoding enables the model to interpolate and extrapolate positional encodings  \nFigure 4: Schematic representation of transformer modules. The input time-series data for ( h + , h × ) is processed concurrently to generate corresponding predictions ( h + , h × ) . This approach ensures the model maintains physical consistency in forecasting both time-series data  \n<!-- image -->  \nfor unseen positions, providing better generalization than other embedding schemes. Building upon previous transformer models [38] for gravitational wave forecasting, this work facilitates the processing of complex-valued waveforms: Each input sequence is split into real and imaginary components, corresponding to plus, h + , and cross, h × , polarizations, and independently embedded using the procedure outlined above. These embeddings are concatenated along the embedding dimension, effectively doubling the input's dimensionality for subsequent processing by the encoder and decoder modules.", '2.2.2. Encoder': "The multi-head self-attention mechanism is designed to allow each position in the input sequence to attend to all positions within the same sequence. For each of the h (= 10) attention heads, the computation reads  \nAttention ( Q j , K j , V j ) = softmax ( Q j ( K j ) T √ d k ) V j . (4)  \nHere, Q j ∈ R n enc × d q , K j ∈ R n enc × d k , and V j ∈ R n enc × d v denote the query, key, and value matrices, respectively, defined as linear transformations of the input matrix X :  \nQ j = XW j Q K j = XW j K V j = XW j V , (5)  \nwith weight matrices W j Q ∈ R d embed × d q , W j K ∈ R d embed × d k , and W j V ∈ R d embed × d v . In this work, d k = d q = d v = d embed / h .  \nQK T functions as a measure of similarity between keys and queries. The softmax activation operates on the rows of QK T and returns weights corresponding to values V . The output of each attention head is concatenated along the last dimension and then linearly transformed:  \nAttention ( Q,K,V ) = concat ( Attention ( Q 1 , K 1 , V 1 ) , · · · , Attention ( Q h , K h , V h ) ) W 0 , (6)  \nwith W 0 ∈ R d embed × d embed , enabling the model to integrate information across different representational subspaces. The resulting attention matrix has the same shape as the input X . After the attention mechanism, each  \nposition's output undergoes processing by a position-wise feed-forward network, which applies two linear transformations with ReLU activations. Each sublayer (self-attention and feed-forward network) is enclosed by a residual connection, followed by layer normalization. This configuration aids in stabilizing the training process of deep networks.", '2.2.3. Decoder': "The decoder receives both the output from the encoder as well as the decoder input, which has been embedded in the same way as the encoder input. First, a masked multihead self-attention mechanism is applied to the decoder input only. This mechanism prevents each position in the decoder input from attending to subsequent positions in the sequence, preserving causality and enabling autoregressive predictions during inference. The operation is defined as:  \nMaskedAttention ( Q j , K j , V j ) = softmax ( Q j ( K j ) T + M √ d k ) V j , (7)  \nM is a mask matrix that applies negative infinity to positions not to be attended to, ensuring that the softmax operation assigns them zero weight. As before, Q j , K j , and V j are the query, key, and value matrices derived from linear transformations on the decoder's input. The dimensionalities of these matrices are the same as those of the encoder self-attention matrices with the exception of the sequence length n , which varies for encoder inputs and decoder inputs. Outputs from each attention head are again concatenated and linearly transformed. These transformations leave the shape of the decoder input unchanged.  \nFollowing the self-attention layer, the cross-attention layer allows the decoder to focus on relevant positions in the encoder output sequence. The queries Q stem from the previous decoder layer (denoted as Y ∈ R n dec × d embed ), while the keys K and values V are derived from the encoder output X :  \nCrossAttention ( Q j , K j , V j ) = softmax ( Q j ( K j ) T √ d k ) V j , (8)  \nwith  \nQ j = Y W j Q K j = XW j K V j = XW j V . (9)  \nThis operation preserves the shape of the queries, in this case the decoder input, ensuring the sequence length of the decoder output matches that of the decoder input. Again, QK T serves as a measure of similarity between keys and queries. The array returned by the softmax activation function is of shape R n dec × n enc and serves as an array of weights for the value array V . Figure 8 presents a visualization of these weights for a choice of input waveform and attention head. The subsequent processing after application of the attention mechanisms mirrors that of the encoder module: each position's output is fed to a feed-forward network, residual connections and layer norms are applied after both attention layers as well as the feed forward network to stabilize training. Finally, a 1D convolution is applied to the decoder output to produce predictions for the two time-series ( h + , h × ) polarizations, respectively.", '2.2.4. h × mask': 'Empirically, we find that separating waveforms with null cross-polarization leads to slightly higher reconstruction accuracy. To facilitate this, we apply a mask to identify and filter waveforms with inclination angle θ = π/ 2 , as these exhibit null cross-polarization. This mask allows us to modify the final convolutional layer and feed-forward layers, which are specifically separated for these waveforms to better capture their unique characteristics. The remaining transformations in the model are consistent across all waveforms, ensuring uniform processing while accommodating the distinct features of the null cross-polarization cases.', '3. Training and inference': "For our transformer model, we used an embedding dimension of 160 (this includes both the contributions from the real and the imaginary part), a feedforward dimension of 80, and 10 attention heads. This model was trained on the NCSA Delta system, using 4 quad NVIDIA A100 GPU nodes. We employed a batch size of 16 and trained the model on 14,440,761 waveforms using the Adam optimizer with mean squared error (MSE) loss metric. The initial learning rate was set to 0.001, with decay applied during training. Training was distributed across all GPUs using PyTorch's DistributedDataParallel (DDP) framework. The model was trained for 30 epochs and reached convergence in around 15 hours.  \nFigure 5: During training, a single decoder output sequence is generated. The attention mechanism uses an attention mask to ensure that each predicted timestep within the decoder output is based on the entire encoder input sequence and the preceding elements in the decoder input. During inference, autoregressive predictions are generated sequentially: at each step, the model produces an output sequence, from which only the last timestep is taken. This output is concatenated with the existing decoder input to form a new input for subsequent predictions. This iterative process continues until the desired sequence is fully generated.  \n<!-- image -->  \nInference on 840,000 test waveforms was completed within 7 hours using 1 node on the DeltaAI system, i.e., 4 NVIDIA H100 GPUs. During training, the model processes entire sequences at once. Inference, in contrast, involves sequential autoregressive generation, where each output is used to predict the next timestep. The functional differences between training and inference are illustrated in Figure 5.", '4. Results and Discussion': "We evaluated the performance of the transformer model on a test set of 840,000 waveforms that covers the signal manifold under consideration. We used the overlap, O , as the figure of merit to quantify the accuracy of forecasted waveforms. This metric is equivalent to the cosine similarity between the predicted and target waveforms. Specifically:  \nO ( h true , h pred ) = max t c ( ⟨ h true , h pred [ t c ] ⟩ √ ⟨ h true , h true ⟩⟨ h pred [ t c ] , h pred [ t c ] ⟩ ) , (10)  \nwhere ⟨ h true , h pred ⟩ denotes the inner product between the complex-valued waveforms h true and h pred , and h pred [ t c ] indicates that the waveform has been time-shifted. A selection of waveform comparisons is depicted in Figure 6, demonstrating the agreement between predicted and target waveforms for various parameter values. The figure includes two randomly selected waveforms, an example of a low overlap waveform, and a waveform with θ = π/ 2 , where the cross polarization h × = 0 .  \nTo provide a more detailed assessment of the model's performance, the top panel in Figure 7 presents a histogram of the overlap values across the test set. Notably, none of the overlap values fall below 0.85, only 0 . 003% of test set waveforms fall below 0.9, while 7 . 15% fall below 0.99, indicating that the model achieves high accuracy in most cases. The mean and median overlaps are 0.996 and 0.997, respectively.  \nThe bottom panels in Figure 7 display overlaps as a function of merger parameters, specifically mass ratio q , inclination angle θ , and effective spin s z eff , defined as [43]:  \ns z eff = qs z 1 + s z 2 1 + q . (11)  \n=0.996  \nFigure 6: True and predicted waveforms for selected ( q, s z 1 , s z 2 , θ ) values. The first two rows represent random samples, the third a predicted waveform with one of the lowest overlap values in the test set, and the last a random sample with θ = π/ 2 .  \n<!-- image -->  \n.  \nAreas of comparatively poor reconstruction accuracy are those at the edge of the signal manifold, i.e., both low mass ratio, q , and low inclination angle, θ , as well as those of both high mass ratio and inclination angle around π/ 2 , as observed in the bottom left panel. Additionally, we observe that low spin values of the heavier object, s z 1 , generally result in slightly lower reconstruction accuracy, with these regions visible as distinctive green bands in the bottom right panel.  \nFigure 7: Top panel: Histogram of overlaps. Bottom panels: Overlap O as a function of merger parameters. Point opacity is inversely proportional to overlap value, with overlaps close to 1 rendered fully transparent, to highlight poor overlaps that might otherwise be obscured. A slight random jitter has been added for the same reason.  \n<!-- image -->  \nIn Figure 8, we provide visualizations of the attention weights derived from the two decoder attention modules. These visualizations aid in forming an intuitive understanding of how the model attends to different parts of the input sequences for a given input waveform and choice of attention head. We present visualizations of attention weight matrices for all other attention heads in Appendix A.  \nFigure 8: Visualization of self-attention (left) and cross-attention (right) weights for a random choice of attention head and input waveform. The raw attention weights have been transposed to obtain the visual representation achieved here. To aid visualization, the corresponding input sequences have been plotted alongside the axes.  \n<!-- image -->  \nAs detailed in the model section, the attention weights are computed using the dot product of queries and keys, QK T , which measures the similarity between them. The softmax function then transforms these similarities into a probability distribution, resulting in a matrix of attention weights with shape R n dec × n enc in the case of the cross-attention module or with shape R n dec × n dec in the case of the self-attention module. These weights are used to weigh the contribution of the values V for the next representation. For visualization purposes, we transpose these attention matrices, as shown in Figure 8.  \nIn both cross- and self-attention plots, we include the corresponding input sequences alongside the x- and y-axes to help visualize the model's focus. For self-attention, the plotted decoder input sequence represents both queries (x-axis) and keys (y-axis). For cross-attention, the plotted decoder input sequence corresponds to the queries (x-axis), while the encoder input sequence represents the keys (y-axis). Although the attention modules operate on embeddings or intermediate representations rather than the raw  \nsequences shown, the plotted sequences provide a useful reference for interpreting the attention patterns. Note that each decoder timestep t dec in the cross-attention plot (i.e., each column in the transposed weight matrix corresponding to some decoder timestep t dec ) contains contextual information about how each encoder timestep relates to the current decoder timestep, which is needed for the transformer to make the t dec + 1 prediction in the decoder output. Due to the softmax operation, all values in each column sum up to 1.  \nThe self-attention heatmap illustrates how each position in the decoder sequence attends to all other positions in the same sequence. For the selfattention mechanism in the decoder, each timestep can attend to all timesteps up to and including itself, but not to future timesteps, due to the causal attention mask. As a consequence of this property, there are no nonzero elements below the diagonal of the transposed weight matrix. As exemplified in Figure 8, the diagonal elements frequently tend to dominate in these selfattention weight plots, indicating that local context is especially important to the self-attention mechanism when making predictions.  \nThe cross-attention heatmap reveals the model's focus on different parts of the encoder output sequence while decoding each position of the output sequence. The attention is more diffused than in the self-attention case, as the model is integrating information from a broader range of encoder timesteps to inform each prediction in the decoder output.  \nThese results, along with those presented in Appendix A, provide insight into how the attention heads prioritize different timesteps in the input signal during the forecasting process. The patterns shown in these visualizations change as we consider different input signals, highlighting the model's sensitivity to variations in the input features. Interactive visualizations for different input signals are available at our GitHub repository [39].", '5. Summary and conclusions': 'Wehave introduced a physics-inspired transformer model capable of learning and forecasting the nonlinear dynamics of higher-order wave modes that describe the physics of quasi-circular, spinning, non-precessing binary black hole mergers.  \nWe developed algorithms to train AI models using over 14 million waveforms with distributed training, harnessing 16 NVIDIA A100 GPUs in the Delta supercomputer. This approach reduced time-to-solution to 15 hours.  \nFurthermore, we used distributed inference to quantify the performance of our transformer model using 4 NVIDIA H100 GPUs in the DeltaAI supercomputer, finding that over a test set of 840,000 waveforms, our transformer model predicts waveforms whose mean and median overlaps with ground truth signals are 0.996 and 0.997, respectively. These results indicate that the transformer model learned the physics of these sources to a high level of accuracy. There are some regions, however, in which the AI model may be improved in future work, namely: low mass-ratios with anti-aligned spins, and high mass-ratio binaries that describe head-on mergers, i.e., with inclination angles θ = π/ 2 . Nonetheless, our results indicate that these samples represent a small set, since only 7.15% of test cases have overlaps below 0.99.  \nThe work presented here, in terms of AI models, distributed training, and inference, paves the way for future studies that may consider the signal manifold of higher-order wave modes of spinning, precessing binary black hole mergers. This high dimensional signal manifold serves as a sandbox to stress test the capabilities of physics-inspired AI models and exascale computing to learn complex, non-linear dynamics of compact binary sources. This study should be pursued in future work.', 'Acknowledgements': 'This work was supported by Laboratory Directed Research and Development (LDRD) funding from Argonne National Laboratory, provided by the Director, Office of Science, of the U.S. Department of Energy under Contract No. DE-AC02-06CH11357, and by the Diaspora project of the U.S. Department of Energy, Office of Science, Advanced Scientific Computing Research, under contract number DE-AC02-06CH11357. KP was in part supported by the NSF grant NRT-1922512. The work used resources of the Argonne Leadership Computing Facility, a DOE Office of Science User Facility supported under Contract DE-AC02-06CH11357. EAH acknowledge support from National Science Foundation (NSF) award OAC-2209892. This research also used the Delta advanced computing and data resources, which is supported by the National Science Foundation (award OAC 2005572) and the State of Illinois. Delta is a joint effort of the University of Illinois Urbana-Champaign and its National Center for Supercomputing Applications. This research used the DeltaAI advanced computing and data resource, which is supported by the National Science Foundation (award OAC 2320345) and the State of Illinois. DeltaAI is a joint effort of the University of Illinois Urbana-Champaign  \nand its National Center for Supercomputing Applications.', 'Appendix A. Attention head visualizations': 'Here we present visualizations of self-attention and cross-attention weights for all attention heads, except for those presented in Figure 8. These results indicate that attention heads gather information across the input time span of input signals to produce accurate predictions.  \nFigure A.9: Visualization of decoder self-attention weights for all 10 attention heads, excluding the one already presented in Figure 8, and choice of input waveform.  \n<!-- image -->  \nimesteps  \nDecoder T  \nimesteps  \nDecoder T  \nimesteps  \nDecoder T  \nFigure A.10: Visualization of cross-attention weights for all attention heads, excluding the one already presented in Figure 8, and choice of input waveform.  \n<!-- image -->', 'References': "- [1] E. Cuoco, J. Powell, M. Cavaglià, K. Ackley, M. Bejger, C. Chatterjee, M. Coughlin, S. Coughlin, P. Easter, R. Essick, H. Gabbard, T. Gebhard, S. Ghosh, L. Haegel, A. Iess, D. Keitel, Z. Marka, S. Marka, F. Morawski, T. Nguyen, R. Ormiston, M. Puerrer, M. Razzano, K. Staats, G. Vajente, and D. Williams, 'Enhancing Gravitational-Wave Science with Machine Learning,' Mach. Learn. Sci. Tech. , vol. 2, no. 1, p. 011002, 2021.\n- [2] E. A. Huerta and Z. Zhao, 'Advances in machine and deep learning for modeling and real-time detection of multi-messenger sources,' in Handbook of Gravitational Wave Astronomy (C. Bambi, S. Katsanevas, and K. D. Kokkotas, eds.), pp. 1-27, Singapore: Springer Singapore, 2020.\n- [3] E. A. Huerta, G. Allen, I. Andreoni, J. M. Antelis, E. Bachelet, G. B. Berriman, F. B. Bianco, R. Biswas, M. Carrasco Kind, K. Chard, M. Cho, P. S. Cowperthwaite, Z. B. Etienne, M. Fishbach, F. Forster, D. George, T. Gibbs, M. Graham, W. Gropp, R. Gruendl, A. Gupta, R. Haas, S. Habib, E. Jennings, M. W. G. Johnson, E. Katsavounidis, D. S. Katz, A. Khan, V. Kindratenko, W. T. C. Kramer, X. Liu, A. Mahabal, Z. Marka, K. McHenry, J. M. Miller, C. Moreno, M. S. Neubauer, S. Oberlin, A. R. Olivas, D. Petravick, A. Rebei, S. Rosofsky, M. Ruiz, A. Saxton, B. F. Schutz, A. Schwing, E. Seidel, S. L. Shapiro, H. Shen, Y. Shen, L. P. Singer, B. M. Sipocz, L. Sun, J. Towns, A. Tsokaros, W. Wei, J. Wells, T. J. Williams, J. Xiong, and Z. Zhao, 'Enabling real-time multi-messenger astrophysics discoveries with deep learning,' Nature Reviews Physics , vol. 1, pp. 600-608, Oct. 2019.\n- [4] D. George, H. Shen, and E. Huerta, 'Glitch Classification and Clustering for LIGO with Deep Transfer Learning,' in NiPS Summer School 2017 , 11 2017.\n- [5] D. George, H. Shen, and E. A. Huerta, 'Classification and unsupervised clustering of LIGO data with Deep Transfer Learning,' Phys. Rev. D , vol. 97, p. 101501, May 2018.\n- [6] H. Shen, D. George, E. A. Huerta, and Z. Zhao, 'Denoising gravitational waves with enhanced deep recurrent denoising auto-encoders,' in  \nICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pp. 3237-3241, IEEE, 2019.  \n- [7] S. Bahaadini, N. Rohani, S. Coughlin, M. Zevin, V. Kalogera, and A. K. Katsaggelos, 'Deep multi-view models for glitch classification,' in 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pp. 2931-2935, IEEE, 2017.\n- [8] S. Coughlin, S. Bahaadini, N. Rohani, M. Zevin, O. Patane, M. Harandi, C. Jackson, V. Noroozi, S. Allen, J. Areeda, M. Coughlin, P. Ruiz, C. P. L. Berry, K. Crowston, A. K. Katsaggelos, A. Lundgren, C. Østerlund, J. R. Smith, L. Trouille, and V. Kalogera, 'Classifying the unknown: Discovering novel gravitational-wave detector glitches using similarity learning,' Phys. Rev. D , vol. 99, p. 082002, Apr. 2019.\n- [9] D. George and E. A. Huerta, 'Deep Learning for real-time gravitational wave detection and parameter estimation: Results with Advanced LIGO data,' Physics Letters B , vol. 778, pp. 64-70, Mar. 2018.\n- [10] D. George and E. A. Huerta, 'Deep neural networks to enable real-time multimessenger astrophysics,' Phys. Rev. D , vol. 97, p. 044039, Feb. 2018.\n- [11] D. George and E. A. Huerta, 'Deep Learning for Real-time Gravitational Wave Detection and Parameter Estimation with LIGO Data,' in NiPS Summer School 2017 , 11 2017.\n- [12] H. Gabbard, M. Williams, F. Hayes, and C. Messenger, 'Matching Matched Filtering with Deep Networks for Gravitational-Wave Astronomy,' Physical Review Letters , vol. 120, p. 141103, Apr. 2018.\n- [13] E. A. Huerta, A. Khan, X. Huang, M. Tian, M. Levental, R. Chard, W. Wei, M. Heflin, D. S. Katz, V. Kindratenko, D. Mu, B. Blaiszik, and I. Foster, 'Accelerated, scalable and reproducible AI-driven gravitational wave detection,' Nature Astronomy , vol. 5, pp. 1062-1068, July 2021.\n- [14] W. Wei, A. Khan, E. A. Huerta, X. Huang, and M. Tian, 'Deep learning ensemble for real-time gravitational wave detection of spinning binary black hole mergers,' Physics Letters B , vol. 812, p. 136029, Jan. 2021.  \n[43] L. Blanchet, 'Gravitational Radiation from Post-Newtonian Sources and Inspiralling Compact Binaries,' Living Rev. Rel. , vol. 17, p. 2, 2014."}