SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fang Wu*
Stan Z. Li
Westlake University
Hangzhou, China
Therapeutic antibodies are an essential and rapidly expanding drug modality. The binding specificity between antibodies and antigens is decided by complementarity-determining regions (CDRs) at the tips of these Y-shaped proteins. In this paper, we propose a hierarchical training paradigm (HTP) for the antibody sequence-structure co-design. HTP consists of four levels of training stages, each corresponding to a specific protein modality within a particular protein domain. Through carefully crafted tasks in different stages, HTP seamlessly and effectively integrates geometric graph neural networks (GNNs) with large-scale protein language models to excavate evolutionary information from not only geometric structures but also vast antibody and non-antibody sequence databases, which determines ligand binding pose and strength. Empirical experiments show that HTP sets the new state-of-the-art performance in the co-design problem as well as the fix-backbone design. Our research offers a hopeful path to unleash the potential of deep generative architectures and seeks to illuminate the way forward for the antibody sequence and structure co-design challenge.
Antibodies, known as immunoglobulins (Ig), are large Y-shaped proteins that the immune system uses to identify and neutralize foreign objects such as pathogenic bacteria and viruses [1, 2]. They recognize a unique molecule of the pathogen, called an antigen. As illustrated in Figure 1 (a), each tip of the "Y" of an antibody contains a paratope that is specific for one particular epitope on an antigen, allowing these two structures to bind together with precision. Notably, the binding specificity of antibodies is largely determined by their complementarity-determining regions (CDRs). Consequently, unremitting efforts have been made to automate the creation of CDR subsequences with desired constraints of binding affinity, stability, and synthesizability [3]. However, the search space is vast, with up to  $20^{L}$  possible combinations for a  $L$ -length CDR sequence. This makes it infeasible to solve the protein structures and then examine their corresponding binding properties via experimental approaches. As a remedy, a group of computational antibody design mechanisms has been introduced to accelerate this filtering process.
Some prior studies [4, 5] prefer to generate only 1D sequences, which has been considered suboptimal because it lacks valuable geometric information. Meanwhile, since the target structure for antibodies is rarely given as a prerequisite [6], more attention has been paid to co-designing the sequence and structure. One conventional line of research [7-9] resorts to sampling protein sequences and structures on the complex energy landscape constructed by physical and chemical principles. But it is found to be time-exhausted and vulnerable to being trapped in local energy optima. Another line [10-12] relies on deep generative models to simultaneously design antibodies' sequences and structures. They take advantage of the most advanced geometric deep learning (DL) techniques and can seize
*Corresponding Authors, emails: fw2359@columbia.edu.
(b)
Figure 1: (a) Schematic structure of an antibody bonded with an antigen (figure modified from Wikipedia). (b) The workflow overview of our hierarchical training paradigm (HTP).
higher-order interactions among residues directly from the data [13]. Their divergence mainly lies in their generative manner. For instance, early works [10, 11] adopt an iterative fashion, while successors [14, 15] employ full-shot generation. Most utilize the traditional translation framework, while some [12] leverage the diffusion denoise probabilistic model.
Despite this fruitful progress, the efficacy of existing co-design methods is predominantly limited by the small number of antibody structures. The Structural Antibody Database (SAbDab)[16] and RAbD[9] are two widely used datasets in the field. After eliminating structures without antigens and removing duplicates, SAbDab comprises only a few thousand complex structures, whereas RAbD consists of 60 complex structures. These numbers are orders of magnitude lower than the data sizes that can inspire major breakthroughs in DL areas [17, 18]. Consequently, deep-generative models fail to benefit from large amounts of 3D antibody-antigen complex structures and are of limited sizes, or otherwise, overfitting may occur.
To address this issue, in this paper, we propose a hierarchical training paradigm (HTP), a novel unified prototype to exploit multiple biological data resources, and aim at fully releasing the potential of geometric graph neural networks (GGNNs) [19, 20] for the sequence and structure co-design problem. Explicitly, HTP consists of four distinct training ranks: single-protein sequence level, antibody sequence level, protein-protein complex structure level, and antibody-antigen complex structure level. These steps are itemized as the data abundance declines, but the task specificity increases, as depicted in Figure 1 (b). Alongside them, we present various pretraining objectives for the first three levels to mine correlated evolutionary patterns, which benefit the co-design task in the final stage. Specifically, we first pretrain the protein language models (PLMs) on tremendous single-protein sequences to obtain general representations and fine-tune them on antibody sequences to capture more condensed semantics. For 3D geometry, we invented a pocket-painting task to exploit the protein-protein complex structures and simulate the CDR generation process. After that, we combine the marvelously pretrained PLMs and the well-pretrained GGNNs to co-design the expected antibodies. Our study provides a promising road to excavate the power of existing deep generative architectures and hopes to shed light on the future development of antibody sequence and structure co-design. The contributions of our work can be summarized as follows.
- First, we equip GGNNs with large-scale PLMs to bridge the gap between protein databases of different modalities (i.e., 1D sequence, and 3D structure) for the antibody co-design problem.
- Second, we design four distinct levels of training tasks to hierarchically incorporate protein data of different domains (i.e., antibodies, and non-antibodies) for the antibody co-design challenge. HTP breaks the traditional co-design routine that separates proteins of different domains and extends antibody-antigen complex structures to broader databases.
- Comprehensive experiments have been conducted to indicate that each stage of HTP significantly contributes to the improvements in the capacity of the DL model to predict more accurate antibody sequences and restore its corresponding 3D structures. To be explicit, HTP brings a rise of  $78.56\%$  in the amino acid recovery rate (AAR) and a decline of  $41.97\%$  in structure prediction error for the sequence-structure co-design. It also leads to an average increase of  $26.92\%$  in AAR for the fixed-backbone design problem.
Figure 2: The illustration of the first three stages of our hierarchical training mechanism for antibody sequence-structure co-design. In level I, a Transformer-based language model is trained by masked language modeling on a large number of single protein sequences to extract general-purposed representations. In level II, the Transformer-based PLM is then further fine-tuned on specific antibody sequences from databases like OAS or ABCD. CDRs are all masked and require recovery, with other framework regions reserved. In level III, GGNNs are asked to predict both the sequence and structure of the pseudo-CDR on protein-protein complex structure databases. The residue features are based on PLMs gained in the previous step.
This section is organized as follows. Subsection 2.1 describes the background knowledge and mathematical formulation of the co-design problem. Subsection 2.4 introduces the backbone architecture to encode the geometric structure of protein-protein complexes as well as antibody-antigen complexes in the 3D space. Subsection 2.2 to 2.6 concentrate on explaining the four different levels of our HTP.
2.1 Preliminary
Background An antibody is a Y-shaped protein with two symmetric sets of chains, each consisting of a heavy chain and a light chain. Each chain has one variable domain (VH/VL) and some constant domains. The variable domain can be further divided into a framework region (FW) and three CDRs. Notably, CDRs in heavy chains contribute the most to the antigen-binding affinity and are the most challenging to characterize.
Notations and Task Formulation We represent each antibody-antigen complex as a heterogeneous graph  $\mathcal{G}_{LR}$ . It is made up of two spatially aggregated components, i.e., the antibody and antigen denoted as  $\mathcal{G}_L = \{\mathcal{V}_L,\mathcal{E}_L\}$  and  $\mathcal{G}_R = \{\mathcal{V}_R,\mathcal{E}_R\}$ , respectively.  $\mathcal{G}_L$  and  $\mathcal{G}_R$  use residues as nodes with numbers of  $N_{L}$  and  $N_{R}$  separately. The node locations  $\mathbf{x}_L\in \mathbb{R}^{N_L\times 3}$  and  $\mathbf{x}_R\in \mathbb{R}^{N_R\times 3}$  are defined as their corresponding  $\alpha$ -carbon coordinate, and are associated with the initial  $\psi_h$ -dimension roto-translation invariant features  $\mathbf{h}_L\in \mathbb{R}^{N_L\times \psi_h}$  and  $\mathbf{h}_R\in \mathbb{R}^{N_R\times \psi_h}$  (e.g. residue types, electronegativity). CDRs are subgraphs of  $\mathcal{G}_L$  and can be divided into  $\mathcal{G}_{HC} = \{\mathcal{V}_{HC},\mathcal{E}_{HC}\}$  and  $\mathcal{G}_{LC} = \{\mathcal{V}_{LC},\mathcal{E}_{LC}\}$ , which belong to the heavy chain and the light chain, respectively. We assume that  $\mathcal{G}_{HC}$  and  $\mathcal{G}_{LC}$  have  $N_{HC}$  and  $N_{LC}$  residues. Besides, it is worth noting that the distance scales between the internal and external interactions are very different. Based on this fact, we strictly distinguish the interaction within and across two graphs  $\mathcal{G}_R$  and  $\mathcal{G}_L$  as and  $\mathcal{E}_L\cup \mathcal{E}_R$  and  $\mathcal{E}_{LR}$ , individually. This implementation avoids the underutilization of cross-graph edges' information due to implicit positional relationships between the antibody and the antigen [21].
In this work, we hypothesize that the antigen structure and the antibody framework are known, aiming to design CDR with more efficacy. Formally, our goal is to jointly model the distribution of CDR in the heavy chain given the structure of the remaining antibody-antigen complex as  $p(\mathcal{G}_{HC}|\mathcal{G}_{LR} - \mathcal{G}_{HC})$  or in the light chain as  $p(\mathcal{G}_{LC}|\mathcal{G}_{LR} - \mathcal{G}_{LC})$ . Since the heavy chain plays a more critical role in determining antigen binding affinity, we take  $\mathcal{G}_{HC}$  as the target example in the following content, but design both heavy and light chains in the experiment section 3.1.
2.2 Single-protein Sequence Level
The idea that biological function and structures are recorded in the statistics of protein sequences selected through evolution has a long history [22]. Unobserved variables that determine the fitness of a protein, such as structure, function, and stability, leave a mark on the distribution of the natural sequence observed [23]. To uncover that information, a group of PLMs has been developed at the scale of evolution, including the series of ESM [22] and ProtTrans [24]. They are capable of capturing information about secondary and tertiary structures and can be generalized across a broad range of downstream applications. Recent studies [25] also demonstrate that equipping GGNNs with pretrained language models can end up with a stronger capacity. Accordingly, we adopt an ESM-2 with 150M parameters to extract per-residue representations, denoted as  $\mathbf{h}_L' \in \mathbb{R}^{N_L \times \psi_{PLM}}$  and  $\mathbf{h}_R' \in \mathbb{R}^{N_R \times \psi_{PLM}}$ , and use them as input node features. Here  $\psi_{PLM} = 640$  and we use  $\Phi$  to denote the trainable parameter set of the language model.
Noteworthy, ESM-2 supplies plenty of options with different model sizes, ranging from 8M to 15B, and a persistent improvement has been found as the model scale increases. However, the trajectory of improvement becomes relatively smooth after the scale reaches  $10^{8}$ . Therefore, for the sake of computational efficiency, we select the 150M version of ESM-2, which performs comparably with the 650M parameters of the ESM-1b model [22]. As declared by Wu et al. [26], incompatibility exists between the experimental structure and its original amino acid sequence (i.e., FASTA sequence). For simplicity, we obey Wu et al. [26]'s mode, which uses the fragmentary sequence directly as the substitute for the integral amino acid sequence and forwards it to the PLMs.
2.3 Antibody Sequence Level
Though PLMs have achieved great progress within protein informatics, their protein representations are generally pursued. Noticeably, the residue distribution of antibodies is significantly different from non-antibodies. Over the past decades, billions of antibodies have been sequenced [27], which enables the training of a language model specifically for the antibody domain [28]. Several researchers have recognized this problem and present models such as AntiBERTa [29] and AbLang [30] to decipher the biology of disease and the discovery of novel therapeutic antibodies. Nonetheless, those pretrained antibody language models have some intrinsic flaws and may not be suitable to apply in our sequence-structure co-design problem immediately.
First and foremost, they are directly pretrained on antibody sequence datasets and fail to exploit the vast amounts of all protein sequences. Second, the existing pretrained antibody language models are all on a small scale. For example, AntiBERTa consists of 12 layers with 86M parameters. Last but not least, both AntiBERTa and AbLang regard the heavy and light chains individually, and AbLang even trains two different models for them. This approach is precluded from encoding a comprehensive and integral representation of the whole antibody. More importantly, they are pretrained on antibodies and cannot be perfectly generalized to extract representations of antigens, which is necessary for analyzing the interactions between antibodies and antigens.
To avoid their drawbacks, we chose to fine-tune ESM on available antibody sequence datasets and considered both antibody and antigen sequences. Following AbLang [30], we leverage the Observed Antibody Space database (OAS) [31] and subsequent update [32]. OAS is a project that collects and annotates immune repertoires for use in large-scale analysis. It contains over one billion sequences, from over 80 different studies. These repertoires cover diverse immune states, organisms, and individuals. Additionally, Olsen et al. [30] has observed in OAS that approximately  $80\%$  sequences lack more than one residue at the N-terminus, nearly  $43\%$  of them are missing the first 15 positions, and about  $1\%$  contain at least one ambiguous residue for each sequence. Remarkably, OAS contains both unpaired and paired antibody sequences and we utilize only paired ones. To better align with our
co-design target, we implement masked language modeling (MLM) and mask residues in all CDRs (i.e., VH, and VL) simultaneously to increase the task difficulty.
2.4 Geometric Graph Neural Networks
To capture 3D interactions of residues in different chains, we adopt a variant of equivariant graph neural network (EGNN) [33] to act on this heterogeneous 3D antibody-antigen graph. The architecture has several key improvements. First, it consists of both the intra- and inter- message-passing schemes to distinguish interactions within the same graph and interactions between different counterparts. Second, it only updates the coordinates of residues in CDR, i.e., the part that is to be designed, while the positions of other parts are maintained as unchangeable. Last, it uses features from the pretrained PLMs as the initial node state rather than a randomized one. Here, all modules are E(3)-equivariant.
The  $l$ -th layer of our backbone is formally defined as the following:
$$
\mathbf {m} _ {j \rightarrow i} = \phi_ {e} \left(\mathbf {h} _ {i} ^ {(l)}, \mathbf {h} _ {j} ^ {(l)}, d \left(\mathbf {x} _ {i} ^ {(l)}, \mathbf {x} _ {j} ^ {(l)}\right)\right), \forall e _ {i j} \in \mathcal {E} _ {L} \cup \mathcal {E} _ {R}, \tag {1}
$$
$$
\boldsymbol {\mu} _ {j \rightarrow i} = a _ {j \rightarrow i} \mathbf {h} _ {j} ^ {(l)} \cdot \phi_ {d} \left(d \left(\mathbf {x} _ {i} ^ {(l)}, \mathbf {x} _ {j} ^ {(l)}\right)\right), \forall e _ {i j} \in \mathcal {E} _ {L R}, \tag {2}
$$
$$
\mathbf {h} _ {i} ^ {(l + 1)} = \phi_ {h} \left(\mathbf {h} _ {i} ^ {(l)}, \sum_ {j} \mathbf {m} _ {j \rightarrow i}, \sum_ {j ^ {\prime}} \boldsymbol {\mu} _ {j ^ {\prime} \rightarrow i}\right), \tag {3}
$$
where  $d(.,.)$  is the Euclidean distance function.  $\phi_e$  is the edge operation and  $\phi_h$  denotes the node operation that aggregates the intra-graph messages  $\mathbf{m}_i = \sum_j\mathbf{m}_{j\rightarrow i}$  and the cross-graph message  $\pmb{\mu}_i = \sum_{j'}\pmb{\mu}_{j'\rightarrow i}$  as well as the node embeddings  $\mathbf{h}_i^{(l)}$  to acquire the updated node embedding  $\mathbf{h}_i^{(l + 1)}$ .  $\phi_d$  operates on the inter-atomic distances.  $\phi_e$ ,  $\phi_h$  and  $\phi_d$  are all multi-layer perceptrons (MLPs). Besides that,  $a_{j\rightarrow i}$  is an attention weight with trainable MLPs  $\phi^q$  and  $\phi^k$ , and takes the following form as:
$$
a _ {j \rightarrow i} = \frac {\exp \left(\left\langle \phi^ {q} \left(\mathbf {h} _ {i} ^ {(l)}\right) , \phi^ {k} \left(\mathbf {h} _ {j} ^ {(l)}\right)\right\rangle\right)}{\sum_ {j ^ {\prime}} \exp \left(\left\langle \phi^ {q} \left(\mathbf {h} _ {i} ^ {(l)}\right) , \phi^ {k} \left(\mathbf {h} _ {j ^ {\prime}} ^ {(l)}\right)\right\rangle\right)}. \tag {4}
$$
As for coordinate iterations, we note that residues located in CDRs (i.e.,  $\mathcal{G}_{HC}$ ) are the sole constituent that needs spatial transformation. On the contrary, the position of the remaining piece (i.e.,  $\mathcal{G}_{LR} - \mathcal{G}_{HC}$ ) is ascertainable. If we change the coordinate of  $\mathcal{G}_{LR} - \mathcal{G}_{HC}$ , its conformation can be disorganized and irrational from physical or biochemical perspectives. Therefore, it is reasonable to stabilize  $\mathcal{G}_{LR} - \mathcal{G}_{HC}$  in each layer and simply alter  $\mathcal{G}_{HC}$ . Mathematically,
$$
\mathbf {x} _ {i} ^ {(l + 1)} = \left\{ \begin{array}{l l} \mathbf {x} _ {i} ^ {(l)} + \frac {1}{| \mathcal {N} _ {i} |} \sum_ {i \in \mathcal {N} _ {i}} \left(\mathbf {x} _ {j} ^ {(l)} - \mathbf {x} _ {j} ^ {(l)}\right) \phi_ {x} (i, j), & \text {i f} \mathbf {x} _ {i} \in \mathcal {G} _ {H C} \\ \mathbf {x} _ {i} ^ {(l)}, & \text {o t h e r w i s e} \end{array} , \right. \tag {5}
$$
where  $\mathcal{N}_i$  denotes the neighbors of node  $i$  and we take the mean aggregation to update the coordinate for each movable node.  $\phi_{x}$  varies according to whether the edge  $e_{ij}$  represent intra-graph connectivity or cross-graph connectivity. In particular,  $\phi_{x} = \phi_{m}(\mathbf{m}_{i\rightarrow j})$  if  $e_{ij}\in \mathcal{E}_L^{(t)}\cup \mathcal{E}_R^{(t)}$ . Otherwise,  $\phi_{x} = \phi_{\mu}\left(\pmb{\mu}_{i\rightarrow j}\right)$  when  $e_{ij}\in \mathcal{E}_{LR}^{(t)}$ , where  $\phi_{m}$  and  $\phi_{\mu}$  are two different functions to deal with different types of messages. Then,  $\phi_{x}$  is left multiplied with  $\mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)}$  to keep the direction information. Equation 5 takes as input the edge embedding  $\mathbf{m}_{i\rightarrow j}$  or  $\pmb{\mu}_{i\rightarrow j}$  as a weight to sum all relative coordinate  $\mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)}$  and output the renewed coordinates  $\mathbf{x}_i^{(l + 1)}$ .
To summarize, the  $l$ -th layer of our architecture  $(l \in [L])$  takes as input the set of atom embeddings  $\left\{\mathbf{h}_L^{(l)}, \mathbf{h}_R^{(l)}\right\}$ , and 3D coordinates  $\left\{\mathbf{x}_L^{(l)}, \mathbf{x}_R^{(l)}\right\}$ . Then it outputs a transformation on  $\left\{\mathbf{h}_L^{(l)}, \mathbf{h}_R^{(l)}\right\}$  as well as coordinates of residues on the CDRs, that is,  $\mathbf{x}_{HC}^{(l)}$ . Concisely,  $\mathbf{h}_L^{(l+1)}, \mathbf{x}_{HC}^{(l+1)}, \mathbf{h}_R^{(l+1)} = \mathrm{GGNN}^{(l)}\left(\mathbf{h}_L^{(l)}, \mathbf{x}_L^{(l)}, \mathbf{h}_R^{(l)}, \mathbf{x}_R^{(l)}\right)$ , while the coordinates of other non-CDR parts remain the same as in the last layer as  $\mathbf{x}_L^{(l+1)} \cup \mathbf{x}_R^{(l+1)} \backslash \mathbf{x}_{HC}^{(l+1)} = \mathbf{x}_L^{(l)} \cup \mathbf{x}_R^{(l)} \backslash \mathbf{x}_{HC}^{(l)}$ . We assign  $\Theta$  as the trainable parameter set of the whole GGNN architecture.
Figure 3: The final stage of our hierarchical training mechanism for antibody sequence-structure co-design. In level IV, both the pretrained language model and geometric graph neural networks are employed to implement the co-design task with the parameters of the pretrained language model fixed. At last, the in silico antibody is experimentally validated via high-throughput binding quantification.
2.5 Protein-protein Complex Structure Level
Apart from empowering PLMs with tremendous protein sequences, plenty of protein-protein complex structures stay uncultivated, which can be harnessed to promote the geometric backbone architecture. Despite the emergence of several structure-based pretraining methods in the biological domain [21, 34], it is not trivial to apply them to our antibody design problem because all previous studies focus on single protein structures. In order to enable GGNNs to encode the general docking pattern between multiple proteins, we propose the pocket inpainting task.
Specifically, we used the Database of Interacting Protein Structures (DIPS) [35]. DIPS is a much larger protein complex structure dataset than existing antibody-antigen complex structure datasets and is derived from the Protein Data Bank (PDB) [36]. In DIPS, each complex  $\mathcal{G}_{LR}$  has two sub-units  $\mathcal{G}_L$  and  $\mathcal{G}_R$ . We calculate the distance between each amino acid of these two substructures and select residues whose minimum distance to the counterpart substructure is less than the threshold  $\epsilon_P = 8\AA$  as the pocket  $\mathcal{G}_P$ . Mathematically, the pocket nodes  $\mathcal{V}_P$  is written as follows:
$$
\mathcal {V} _ {P} = \left\{v _ {L, i} \mid \min  _ {j = 1} ^ {N _ {L}} d \left(\mathbf {x} _ {L i}, \mathbf {x} _ {R, j}\right) <   \epsilon_ {P} \right\} \cup \left\{v _ {R, i} \mid \min  _ {j = 1} ^ {N _ {R}} d \left(\mathbf {x} _ {R, i}, \mathbf {x} _ {L, j}\right) <   \epsilon_ {P} \right\}, \tag {6}
$$
then our target is to retrieve  $\mathcal{G}_P$  given the leftover  $\mathcal{G}_{LR} - \mathcal{G}_P$ . Here, we follow the official split based on PDB sequence clustering at a  $30\%$  sequence identity level to ensure little contamination between sets. It results in train/val/test of 87,303/31,050/15,268 complex samples.
2.6 Antibody-antigen Complex Structures Level
After the preceding levels of preparation, it is time to co-design the antibody. Given an antigen  $\mathcal{G}_R$  and a fractional antibody  $\mathcal{G}_L - \mathcal{G}_{HC}$ , we first employ the well-trained PLM  $\Phi$  to attain per-residue node features  $\mathbf{h}_R^{(0)}$  and  $\mathbf{h}_L^{(0)}$ , where the nodes in the unknown part  $\mathcal{G}_{HC}$  are tagged as a masked token. Then both features and coordinates are fed into the well-trained geometric encoder to unravel the CDR sequences and structures in a concurrent way, i.e.,  $\mathbf{h}_L^{(L)}, \mathbf{x}_{HC}^{(L)}, \mathbf{h}_R^{(L)} = \mathrm{EGNN}_{\Theta}\left(\mathbf{h}_L^{(0)}, \mathbf{x}_L^{(0)}, \mathbf{h}_R^{(0)}, \mathbf{x}_R^{(0)}\right)$ . Eventually, we use an MLP  $\phi_o$  and a Softmax operator as the classifier to output the probability distribution of residue types as  $p_i = \text{Softmax}\left(\phi_o\left(\mathbf{h}_i^{(L)}\right)\right) \in \mathbb{R}^{20}$  for  $v_i \in \mathcal{V}_{HC}$ .