# Cross-Attention U-Net for Elastic Wavefield Decomposition in Anisotropic Media
## 2 Theoretical Framework
### Network Architecture  
The proposed model is based on a U-Net architecture [15], enhanced with cross-attention mechanisms to improve vector decomposition while preserving inter-component relationships. Unlike conventional approaches that concatenate displacement components into a single input, this model processes them separately in the encoder and integrates cross-attention in the decoder to facilitate feature exchange. This structure ensures that the vector decomposition remains physically meaningful while improving the separation process. The model takes a seismic wavefield snapshot as input, where the horizontal and vertical displacement components, \(u_{x}\) and \(u_{z}\), serve as the input channels. Each displacement component contains contributions from both qP and qS waves, which must be properly separated to ensure accurate vector decomposition. Since this study focuses on 2D wavefield decomposition, the analysis is restricted to qP and qS waves, without considering separate shear wave polarizations such as qSH and qSV. The network is designed to output two channels, \(u_{x,qS}\) and \(u_{z,qS}\), corresponding to the qS wave components of the horizontal and vertical displacements. The qP components, \(u_{x,qP}\) and \(u_{z,qP}\), are then inferred using the physical constraint:  
\[u_{x,qP}=u_{x}-u_{x,qS},\quad u_{z,qP}=u_{z}-u_{z,qS}, \tag{8}\]  
Empirical results indicate that learning 5-wave components directly and computing P-wave components as residuals leads to improved separation performance, making this a preferred approach for training. As illustrated in Figure 1, the network consists of three main components: an encoder, a cross-attention module, and a decoder. The encoder follows a dual-branch structure, where the two displacement components are processed independently through separate encoding paths. Each encoding path consists of multiple down-sampling stages, applying residual blocks [16] for feature extraction while progressively reducing the spatial resolution through pooling layers. This design allows the network to learn independent representations for each component while preserving their structural integrity. Unlike conventional U-Netmodels, which process displacement components jointly, this approach ensures that the unique polarization characteristics in \(u_{x}\) and \(u_{z}\) are retained during feature extraction. To facilitate feature exchange between the encoded displacement components, cross-attention is introduced before the decoder. Instead of simply concatenating feature maps  
Figure 1: Architecture of the proposed cross-attention-based U-Net model for seismic wavefield decomposition. The network consists of three main components: an encoder, a cross-attention module, and a decoder. The encoder follows a dual-path structure, where horizontal (\(u_{x}\)) and vertical (\(u_{z}\)) displacement components are processed separately through independent encoding paths. The cross-attention modules facilitate feature exchange between the two paths, ensuring improved vector decomposition while preserving inter-component relationships. The decoder reconstructs the qS wave components (\(u_{x,qS}\), \(u_{z,qS}\)), while the qP wave components (\(u_{x,qP}\), \(u_{z,qp}\)) are inferred using a physical constraint. before decoding, cross-attention selectively integrates relevant information from one component into the other, ensuring that the decomposition process remains consistent with the underlying physical relationships between wave modes. The cross-attention mechanism enables dynamic feature interaction and reinforces inter-component dependencies, which are essential for robust wave mode separation, particularly in anisotropic media. Following cross-attention, the decoder mirrors the hierarchical structure of the encoders. For each up-sampling step, transposed convolutions restore spatial resolution, and feature maps from the corresponding encoder layers are concatenated through skip connections to retain fine-scale structural details [15].