# GraphEPN: A Deep Learning Framework for B-Cell Epitope Prediction Leveraging Graph Neural Networks
## 2 Materials and Methods
### Model Architecture
#### 2.3.2 Pre-Trained VQ-VAE Module  
To learn discrete latent representations from protein structures, we employ a VQ-VAE. The VQ-VAE model comprises an encoder, a vector quantization layer, and a decoder (Figure 1b) and is trained in an unsupervised manner. The encoder, which is based on a graph convolutional network (GCN), processes node features from the protein graph [19].  
Figure 1: Schematic view of the GraphEPN architecture. (**a**) Overall framework: The 3D structure of a protein (PDB ID: 1 OTU_A) is represented as a graph, where nodes correspond to amino acid residues. The VQ-VAE encodes and quantizes node features into discrete embeddings, which are passed to a graph transformer for epitope prediction. (**b**) VQ-VAE Module: The encoder extracts latent features and maps them to the nearest codebook vectors, generating discrete representations, while the decoder reconstructs original features. (**c**) Graph Transformer Architecture: The model applies graph attention networks (GAT) to capture residue interactions, followed by residual connections and feed-forward layers.  
The input for each node is \(x_{i}\in\mathbb{R}^{d}\), and the output is the latent representation \(z_{i}\in\mathbb{R}^{d_{z}}\). Each GCN layer updates node features by aggregating information from neighboring nodes:  
\[\mathrm{h}_{i}^{(l)}=\sigma\left(\sum_{(j\in\mathcal{N}(i))}\mathrm{W}^{(l)} \mathrm{h}_{j}^{(l-1)}+\mathrm{b}^{(l)}\right) \tag{8}\]  
where \(\mathrm{h}_{i}^{(l)}\) is the representation of node \(i\) at the all-th layer, initialized as \(\mathrm{h}_{i}^{(0)}=\mathrm{x}_{i}\). Here, \(\mathcal{N}(i)\) denotes the set of neighboring nodes of \(i\), \(\mathrm{W}^{(l)}\) is a trainable weight matrix, and \(\mathrm{b}^{(l)}\) is a bias term.  
The continuous representation \(z_{i}\) produced by the encoder is passed to the vector quantization layer, which uses a discrete codebook for quantization. Through nearest neighbor search, it is mapped to the closest vector in the codebook \(e_{i}\), expressed as follows:  
\[e_{i}=\mathrm{argmin}_{\mathbb{C}_{x}}\|z_{i}-C_{k}\|_{2} \tag{9}\]  
The quantization loss \(\mathcal{L}_{vq}\) ensures alignment between the continuous and discrete spaces, while a commitment loss \(\mathcal{L}_{commit}\) promotes stable updates for the codebook:  
\[\mathcal{L}_{vq}=\sum_{i}\lVert\mathrm{sg}[z_{i}]-e_{i}\rVert_{2}^{2}, \mathcal{L}_{commit}=\sum_{i}\lVert z_{i}-\mathrm{sg}[e_{i}]\rVert_{2}^{2} \tag{10}\]  
where \(\mathrm{sg}[\cdot]\) denotes the stop-gradient operation to prevent gradients from being propagated to the encoder.  
The decoder, which is based on the GCN architecture, decodes the discrete representation \(e_{i}\) back to the original node feature \(\hat{\mathrm{x}}_{i}\). The reconstruction loss \(\mathcal{L}_{recon}\) is expressed as follows:  
\[\mathcal{L}_{recon}=\sum_{i}\lVert\hat{\mathrm{x}}_{i}-x_{i}\rVert_{2}^{2} \tag{11}\]  
To enhance robustness, a masking mechanism is introduced during training, with a masking loss \(\mathcal{L}_{mask}\).  
Finally, the total loss function of the VQ-VAE model is as follows:  
\[\mathcal{L}_{VQ-VAE}=\mathcal{L}_{recon}+\beta\mathcal{L}_{vq}+\alpha\mathcal{ L}_{commit}+\lambda\mathcal{L}_{mask} \tag{12}\]  
where \(\beta\), \(\alpha\), and \(\lambda\) are hyperparameters used to balance the contributions of each loss term.