Conditional Antibody Design as 3D Equivariant Graph Translation

An antibody is a Y-shaped protein (Figure 1) with two symmetric sets of chains, each composed of a heavy chain and a light chain (Kuroda et al., 2012). In each chain, there are some constant domains, and a variable domain ($V_{H}$/$V_{L}$) that has three Complementarity Determining Regions (CDRs). Antigen-binding sites occur on the variable domain where the interacting regions are mostly CDRs, especially CDR-H3. The remainder of the variable domain other than CDRs is structurally well conserved and often called framework region (Kuroda et al., 2012; Jin et al., 2021). Therefore, previous works usually formalize the antibody design problem as finding CDRs that fit into the given framework region (Shin et al., 2021; Akbar et al., 2022). As suggested by (Fischman & Ofran, 2018; Jin et al., 2021), we focus on generating CDRs in heavy chains since they contribute the most to antigen-binding affinity and are the most challenging ones to characterize. Nevertheless, in contrast to previous studies, we additionally incorporate the antigen and the light chain into the context in the form of antibody-antigen complexes, to better control the binding specificity of generated antibodies.

We represent each antibody-antigen complex as a graph of three spatially aggregated components, denoted as ${\mathcal{G}}=({\mathcal{V}}\coloneqq\{{\mathcal{V}}_{H},{\mathcal{V}}_{L},{\mathcal{V}}_{A}\},{\mathcal{E}}\coloneqq\{{\mathcal{E}}_{\text{in}},{\mathcal{E}}_{\text{ex}}\})$. Here, the components ${\mathcal{V}}_{H},{\mathcal{V}}_{L},{\mathcal{V}}_{A}$ correspond to the nodes (i.e. the residues) of the heavy chain, the light chain and the antigen, respectively; ${\mathcal{E}}_{\text{in}}$ and ${\mathcal{E}}_{\text{ex}}$ separately contain internal edges within each component and external edges across components. To be specific, each node in ${\mathcal{V}}$, i.e., $v_{i}=(h_{i},{\bm{Z}}_{i})$ is represented as a trainable feature embedding vector $h_{i}=s_{a_{i}}\in\mathbb{R}^{d_{a}}$ according to its amino acid type $a_{i}$ and a matrix of coordinates ${\bm{Z}}_{i}\in\mathbb{R}^{3\times m}$ consisting of $m$ backbone atoms. In our case, we set $m=4$ by choosing 4 backbone atoms $\{\text{N},\text{C}_{\alpha},\text{C},\text{O}\}$, where $\text{C}_{\alpha}$ denotes the alpha carbon of the residue and others refer to the atoms composing the peptide bond (Figure 1). We denote the residues in CDRs to be generated as $\mathcal{V}_{C}=\{v_{c_{1}},v_{c_{2}},...,v_{c_{n(c)}}\}$, which is a subset of $\mathcal{V}_{H}$. Since the information of each $v_{c_{i}}$ is unknown in the first place, we initialized its input feature with a mask vector and the coordinates according to the even distribution between the residue right before CDRs (namely, $v_{c_{1}-1}$) and the one right after CDRs (namely, $v_{c_{n(c)}+1}$). In our main experiments (§ 4), we select the 48 residues of the antigen closest to the antibody in terms of the $\text{C}_{\alpha}$ distance as the epitope like Jin et al. (2021). Instead of such hand-crafted residue selection, in Appendix L, we also incorporate the full antigen and let our method determine the epitope automatically, where the efficacy of our method is still exhibited.

To derive an effective translator, it is crucial to capture the 3D interactions of the residues in different chains. The message passing mechanism in $E(3)$-equivariant GNNs (Satorras et al., 2021; Huang et al., 2022) will fulfil this purpose. Particularly, we develop the Multi-channel Equivariant Attention Network (MEAN) to characterize the geometry and topology of the input antibody-antigen complex. Each layer of MEAN alternates between the two modules: internal context encoder and external interaction encoder, which is motivated by the biological insight that the external interactions between antibodies and antigens are different from those internal interactions within each heavy/light chain. After several layers of the message passing, the node representations and coordinates are transformed into the predictions by an output module. Notably, all modules are $E(3)$-equivariant.

Traditional methods (such as RefineGNN (Jin et al., 2021)) unravel the CDR sequence in an autoregressive way: generating one amino acid at one time. While such strategy is able to reduce the generation complexity, it inevitably incurs expensive computing and memory overhead, and will hinder the training owing to the vanishing gradient for long CDR sequences. It also acumulates errors during the inference stage. Here, thanks to the rich expressivity of MEAN, we progressively generate the CDRs over $T$ iterations ($T$ is much smaller than the length of the CDR sequences), and in each iteration, we predict the amino acid type and 3D coordinates of all the nodes in ${\mathcal{V}}_{C}$ at once. We call our scheme as full-shot decoding to distinguish it from previous autoregressive approaches.

To be specific, given the CDRs’ amino acid distribution and conformation $\{p_{i}^{(t)},\tilde{{\bm{Z}}}_{i}^{(t)}\}_{i\in{\mathcal{V}}_{C}}$ from iteration $t$, we first update the embeddings of all nodes: $h^{\prime}_{i}=\sum_{j=1}^{n_{a}}p_{i}^{(t)}(j)s_{j},\forall i\in{\mathcal{V}}_{C}$, where $p^{(t)}_{i}(j)$ returns the probability for the $j$-th class and $s_{j}$ is the corresponding learnable embedding as defined before. Such weighted strategy leads to less accumulated error during inference compared to the maximum selection counterpart. We then replace the CDRs with the new values $\{h^{\prime}_{i},\tilde{{\bm{Z}}}_{i}^{(t)}\}_{i\in{\mathcal{V}}_{C}}$, and denote the new graph as ${\mathcal{G}}^{(t+1)}$. The edges are also constructed dynamically according to the new graph. We update the next iteration as $\{p_{i}^{(t+1)},\tilde{{\bm{Z}}}_{i}^{(t+1)}\}_{i\in{\mathcal{V}}_{C}}=\mathrm{MEAN}({\mathcal{G}}^{(t+1)})$.

For sequence prediction, we exert supervision for each node at each iteration:

where $\ell_{ce}$ denotes the cross entropy between the predicted distribution $p^{(t)}_{i}$ and the true one $\hat{p}_{i}$.

For structure prediction, we only exert supervision on the output iteration. Since there are usually noises in the coordination data, we adopt the Huber loss (Huber, 1992) other than the common MSE loss to avoid numerical instability (further explanation can be found in Appendix F):

where $\hat{{\bm{Z}}}_{i}$ is the label. One benefit of the structure loss is that it conducts directly in the coordinate space, and it is still $E(3)$-invariant as our model is $E(3)$-equivariant. This is far more efficient than the loss function used in RefineGNN (Jin et al., 2021). To ensure invariance, RefineGNN should calculate over pairwise distance and angle other than coordinates, which is tedious but necessary since it can only perceive the input of node and edge features after certain invariant transformations. Finally, we balance the above two losses with $\lambda$ to form $\mathcal{L}=\mathcal{L}_{\mathrm{seq}}+\lambda\mathcal{L}_{\mathrm{struct}}$.

For quantitative evaluation, we employ Amino Acid Recovery (AAR), defined as the overlapping rate between the predicted 1D sequences and ground truths, and Root Mean Square Deviation (RMSD) regarding the 3D predicted structure of CDRs. Thanks to its inherent equivariance, our model can directly calculate RMSD of coordinates, unlike other baselines that resort to the Kabsch technique Kabsch (1976) to align the predicted and true coordinates prior to the RMSD computation. Our model requires each input complex to be complete (consisting of heavy chains, light chains, and antigens). Hence, we choose 3,127 complexes from the Structural Antibody Database (Dunbar et al., 2014, SAbDab) and remove other illegal datapoints that lack light chain or antigen. All selected complexes are renumbered under the IMGT scheme (Lefranc et al., 2003). As suggested by Jin et al. (2021), we split the dataset into training, validation, and test sets according to the clustering of CDRs to maintain the generalization test. In detail, for each type of CDRs, we first cluster the sequences via MMseqs2 (Steinegger & Söding, 2017) that assigns the antibodies with CDR sequence identity above 40% to the same cluster, where the BLOSUM62 substitution matrix (Henikoff & Henikoff, 1992) is adopted to calculate the sequence identity. The total numbers of clusters for CDR-H1, CDR-H2, and CDR-H3 are 765, 1093, and 1659, respectively. Then we split all clusters into training, validation, and test sets with a ratio of 8:1:1. We conduct 10-fold cross validation to obtain reliable results. Further details are provided in Appendix A.

We perform fine-grained validation on designing CDR-H3 that binds to a given antigen. In addition to AAR and RMSD, we further adopt TM-score (Zhang & Skolnick, 2004; Xu & Zhang, 2010), which calculates the global similarity between two protein structures and ranges from 0 to 1, to evaluate how well the CDRs fit into the frameworks. We use the official implementation²²2https://zhanggroup.org/TM-score/ to calculate TM-score. For autoregressive baselines, we follow Jin et al. (2021) to alleviate the affects of accumulated errors by generating 10,000 CDR-H3s and select the top 100 candidates with the lowest PPL for evaluation. Notably, our MEAN reduces accumulated errors significantly by the progressive full-shot decoding strategy, therefore can directly generate the candidates. We also include RosettaAD (Adolf-Bryfogle et al., 2018), a widely used baseline of conventional approach, for comparison. We benchmark all methods with the 60 diverse complexes carefully selected by Adolf-Bryfogle et al. (2018) (RAbD). The training is still conducted on the SAbDab dataset used in § 4.1, but we eliminate all antibodies from SAbDab whose CDR-H3s share the same cluster as those in RAbD to avoid any potential data leakage. We then divide the remainder into the training and validation sets by a ratio of 9:1. The numbers of clusters/antibodies are 1,443/2,638 for training and 160/339 for validation.

It is crucial to optimize various properties like binding affinity of antibodies for therapeutic purposes. This can be formulated as a search problem over the intrinsic space of generative models. In our case, we jointly optimize the sequence and structure of CDR-H3 to improve the binding affinity of any given antibody-antigen complex. For evaluation, we employ the geometric network from Shan et al. (2022) to predict the change in binding energy ($\Delta\Delta$G) after optimization. Particularly, we leverage the official checkpoint³³3https://github.com/HeliXonProtein/binding-ddg-predictor that is trained on the Structural Kinetic and Energetic database of Mutant Protein Interactions V2.0 (Jankauskaitė et al., 2019, SKEMPI V2.0). $\Delta\Delta$G is calibrated under the unit of kcal/mol, and lower $\Delta\Delta$G indicates better binding affinity. Since all the methods model backbone structures only, we use Rosetta (Alford et al., 2017) to do sidechain packing before affinity prediction. To ensure the expected generalizability, we select a total of 53 antibodies from its training set (i.e. SKEMPI V2.0) for affinity optimization. Besides, we split SAbDab (pre-processed by the strategy in § 4.1) into training and validation sets in a ratio of 9:1 for pretraining the model.

Following Jin et al. (2021), we exploit the Iterative Target Augmentation (Yang et al., 2020, ITA) algorithm to tackle the optimization problem. Since the original algorithm is designed for discrete properties, we adapt it for compatibility with our affinity scorer of continuous values. Please refer to Appendix B for detailed description of the adaption. During the process, we discard any unrealistic candidate with PPL above 10 in accordance with Jin et al. (2021). It is observed that our model learns the constraints of net charge, motifs, and repeating amino acids (Jin et al., 2021) implicitly, therefore we do not need to impose them on our model explicitly.

Results As shown in Table 4.3 (middle), our MEAN models achieve obvious progress towards discovering antibodies with better binding affinity. This further validates the advantage of explicitly modeling the interface with MEAN. Moreover, we provide the predicted $\Delta\Delta$G of mutating the CDRs to random sequences, denoted as Random, for better interpretation of the results. We provide further interpretation in Appendix C. We also provide a visualization example in Figure 3 (B), which indicates our MEAN does produce a novel CDR-H3 sequence/structure, with improved affinity.

We further provide a possible pipeline to utilize our model in scenarios when the binding complex is unknown. Specifically, for antigens from RAbD (Adolf-Bryfogle et al., 2018), we aim at generating binding antibodies with high affinity. To this end, we first select an antibody from the database and remove its CDR-H3, then we use HDOCK (Yan et al., 2017) to dock it to the target antigen to obtain a template of antibody-antigen complex. With this template, we employ our model to generate antigen-binding CDR-H3s in the same way as § 4.2. To alleviate the risk of docking inaccuracy, we compose 10 such templates for each antigen and retain the highest scoring one in the subsequent generation. We first refine the generated structure with OpenMM (Eastman et al., 2017) and Rosetta (Alford et al., 2017), and then use the energy functions in Rosetta to measure the binding affinity. The comparison of the affinity distribution between the generated antibodies by our method, those by C-RefineGNN, and the original ones in RAbD is shown in Figure 4 (B). Obviously, the antibodies designed by our MEAN exhibit higher predicted binding affinity. We also present a tightly binding example in Figure 4 (A).