# Cross-Attention U-Net for Elastic Wavefield Decomposition in Anisotropic Media
## 2 Theoretical Framework
### Training  
To evaluate the effectiveness of the proposed cross-attention-based wavefield decomposition, a modified version of the BP 2007 anisotropic benchmark model, a widely used anisotropic subsurface model, was employed. Figure 2 shows the elastic properties of the medium: P-wave velocity (_Vp_), S-wave velocity (_Vs_), and Thomsen parameters (_e_, \(\delta\)). For training, the model domain was partitioned into 200 \(\times\) 200 subregions, yielding 60 distinct patches. To increase variability and improve generalization, 30 additional patches were generated using standard data augmentation techniques, including random flips and minor scaling. The patch size (200 \(\times\) 200) and the number of augmented samples (30) were empirically chosen to balance model performance and computational cost. These parameters were fixed during training for consistency but may be adjusted in future studies depending on the application. The process of extracting patches from the full model is illustrated in Figure 3, which provides an example of how the domain is segmented for training.  
For each patch, wave propagation is simulated under five different source locations, ensuring that the network is exposed to wavefields from multiple illumination angles. From each source simulation, 5 temporal snapshots of the propagating wavefield are extracted, leading to a total of 25 wavefield snapshots per patch. The snapshots are randomly selected from a time window excluding the initial 0.2 s of simulation, with at least 0.15 s between each selection to ensure sufficient temporal diversity. The ground truth labels are generated via the low-rank approximation method described in Section 2.2, allowing us to bypass explicit Christoffel equation solvers during training.  
Figure 2: Elastic properties of the modified BP 2007 anisotropic benchmark model. **(a)** P-wave velocity (_Vp_), **(b)** S-wave velocity (_Vs_), **(c)** Thomsen parameter \(\epsilon\), and **(d)** Thomsen parameter \(\delta\). These properties define the subsurface anisotropic characteristics used for wavefield decomposition.  
During training, two U-Net variants are compared--one with multi-head cross-attention integrated at the deepest encoder level and one without. Both variants utilize residual blocks in the encoder and decoder to facilitate gradient flow. The dataset is randomly partitioned into training, validation, and testing splits, with 64% of the data used for training, 16% for validation, and 20% for testing. All network weights are initialized by the Kaiming normal scheme [17] to stabilize the early training. Optimization is performed using the Adam optimizer [18], with an initial learning rate of 1.5 \(\times\) 10\({}^{-4}\). Training typically proceeds for up to 200 epochs with a patience-based early-stopping criterion set to 30 epochs. To further regularize the model, a minor physics-informed penalty (coefficient 0.01) is included within the loss function to encourage physically consistent separation. This term is particularly useful in noisy or highly heterogeneous scenarios, where purely data-driven separation may overfit spurious features.  
The training was conducted using an NVIDIA RTX 4090 Graphics Processing Unit (GPU; NVIDIA Corporation, Santa Clara, CA, USA) with CUDA (Compute Unified Device Architecture) acceleration. For the baseline U-Net model (without cross-attention), the total training time was approximately 4925 s, while the cross-attention U-Net required approximately 6487 s. This increase in training time reflects the additional computational complexity that is introduced by the attention mechanism.