SlowGuess/NeurIPS_2023_OCR huggingface dataset:
CDR Coordinates Initialization. How to initialize the positions of residues in CDRs is of great importance to the co-design problem, and there is no consensus among different approaches. HERN [11] tries two kinds of strategies. One strategy is to randomly initialize all coordinates by adding a small Gaussian noise around the center of the epitope as  $\mathbf{x}_i^{(0)} = \frac{1}{N_R}\sum_{j\in \mathcal{G}_R}\mathbf{x}_j + \epsilon ,\epsilon \sim \mathcal{N}(0,1)$ . The other is to predict the pairwise distance instantly  $\mathbf{D}^{(N_L + N_R)\times (N_L + N_R)}$  between paratope and epitope atoms
Table 1: Results of sequence and structure co-deign on SAbDab.
Model	CDR-H1	CDR-H2	CDR-H3
AAR (%)↑	RMSD ↓	TM-Score ↑	AAR (%)↑	RMSD ↓	TM-Score ↑	AAR (%)↑	RMSD ↓	TM-Score ↑
RAbD	20.63 ± 1.6	3.56 ± 0.05	0.9206 ± 0.007	27.80 ± 0.8	2.85 ± 0.09	0.9253 ± 0.010	21.73 ± 0.7	4.58 ± 0.13	0.8916 ± 0.012
C-RGNN	40.39 ± 3.2	1.98 ± 0.02	0.9380 ± 0.003	33.36 ± 1.7	1.32 ± 0.05	0.9507 ± 0.005	21.89 ± 1.5	3.59 ± 0.16	0.9187 ± 0.011
MEAN	43.80 ± 2.5	1.84 ± 0.04	0.9411 ± 0.008	37.18 ± 1.5	1.27 ± 0.04	0.9522 ± 0.007	22.56 ± 1.7	3.44 ± 0.18	0.9248 ± 0.009
HERN	48.42 ± 2.7	1.69 ± 0.04	0.9472 ± 0.005	41.53 ± 2.1	1.26 ± 0.03	0.9531 ± 0.006	25.73 ± 1.4	3.02 ± 0.11	0.9340 ± 0.004
DiffAb	52.82 ± 0.9	1.51 ± 0.01	0.9658 ± 0.001	45.95 ± 2.3	1.24 ± 0.01	0.9588 ± 0.002	27.04 ± 2.8	2.89 ± 0.15	0.9417 ± 0.008
HTP	81.33 ± 1.6	0.49 ± 0.02	0.9829 ± 0.006	67.77 ± 1.3	0.53 ± 0.02	0.9860 ± 0.004	40.98 ± 1.5	2.06 ± 0.03	0.9621 ± 0.005
Model	CDR-L1	CDR-L2	CDR-L3
AAR (%)↑	RMSD ↓	TM-Score ↑	AAR (%)↑	RMSD ↓	TM-Score ↑	AAR (%)↑	RMSD ↓	TM-Score ↑
RAbD	35.11 ± 1.0	1.88 ± 0.01	0.9458 ± 0.002	27.82 ± 0.6	1.35 ± 0.02	0.9611 ± 0.009	23.73 ± 0.5	2.14 ± 0.06	0.9247 ± 0.010
C-RGNN	41.44 ± 2.5	2.06 ± 0.02	0.9326 ± 0.008	36.71 ± 4.3	1.26 ± 0.01	0.9652 ± 0.006	33.80 ± 4.8	1.95 ± 0.06	0.9308 ± 0.008
MEAN	47.69 ± 2.3	1.87 ± 0.02	0.9461 ± 0.006	39.42 ± 3.5	1.24 ± 0.01	0.9647 ± 0.008	35.18 ± 2.6	1.84 ± 0.05	0.9370 ± 0.005
HERN	55.24 ± 2.7	1.63 ± 0.02	0.9502 ± 0.003	46.02 ± 4.1	1.18 ± 0.02	0.9712 ± 0.004	37.28 ± 4.1	1.77 ± 0.07	0.9389 ± 0.002
DiffAb	62.71 ± 1.2	1.48 ± 0.01	0.9637 ± 0.002	52.10 ± 3.6	1.11 ± 0.06	0.9780 ± 0.014	43.62 ± 2.6	1.65 ± 0.05	0.9447 ± 0.004
HTP	91.13 ± 1.0	0.67 ± 0.04	0.9869 ± 0.005	89.80 ± 0.8	1.03 ± 0.05	0.9846 ± 0.008	73.82 ± 1.1	0.78 ± 0.04	0.9704 ± 0.003
and reconstruct atom coordinates from this distance matrix. The results show that the former performs better. DiffAB [12] initializes them from the standard normal distribution as  $\mathbf{x}_i^{(0)}\sim \mathcal{N}(\mathbf{0},\mathbf{I}_3)$ . MEAN [14] leverage the even distribution between the residue right before CDRs and the one right after CDRs as the initial positions.
All of the above initialization mechanisms have corresponding drawbacks. To be specific, HERN insufficiently considers the context information since it only characterizes the shape of the epitope and ignores the incomplete antibody structure  $\mathcal{G}_{LR} - \mathcal{G}_{HC}$ . The initialization of normal distribution is only suitable in diffusion-based models. Moreover, our empirical experiments observe the least instability during training if an even distribution of MEAN is adopted, but residues right before or right after CDRs can be missing. Here, we propose another way to initialize the residue positions in CDRs. First, we follow MEAN and select the residue right before and right after CDRs. If both residues exist, we take the mean coordinates. Otherwise, we use only the existing one. After that, we add some little noise like HERN to separate nodes and introduce some randomization to prevent overfitting, which is proven to bring slight improvements in performance.
Loss Function. Two sorts of losses are used for supervision. First, a cross-entropy (CE) loss is employed for sequence prediction as  $\mathcal{L}_{seq} = \frac{1}{|\mathcal{V}_{HC}|}\sum_{v_i\in \mathcal{V}_{HC}}\mathrm{CE}(p_i,c_i)$ , where  $c_{i}$  is the ground truth residue type for each node. Apart from that, a common RMSD loss is utilized for structure prediction and we leverage the Huber loss [37] to avoid numerical instability as  $\mathcal{L}_{struct} = \text{Huber} \left( \mathbf{x}_{HC}^{(L)}, \mathbf{x}_{HC} \right)$ , where the latter is the ground truth coordinates of the target CDR. The total loss is a weighted sum of the above two as  $\mathcal{L} = \mathcal{L}_{seq} + \lambda \mathcal{L}_{struct}$ , where  $\lambda > 0$  is the balance hyperparameter.
We assess our HTP via two mainstream challenging tasks: sequence-structure co-design in Section 3.1, and antibody sequence design based on antibody backbones in Section 3.2. Our evaluation is conducted in the standard SAbDab database [16]. More experimental setting details and data descriptions are elucidated in Appendix A.
3.1 Sequence-structure Co-design
Task and Metrics. In this task, we remove the original CDR from the antibody-antigen complex in the test set and try to co-design both sequence and structure of the removed region. Here we set the length of the CDR to be identical to the length of the original CDR. But in practice, the lengths of CDRs can be variable. For quantitative evaluation, we adopt amino acid recovery (AAR) and root-mean-squared error (RMSD) regarding the 3D predicted structure of CDRs as the metric. AAR is defined as the overlapping rate between the predicted 1D sequences and the ground truths. We also take advantage of TM score [38] to calculate the global similarity between the predicted and ground truth antibody structures. It ranges from 0 to 1 and evaluates how well the CDRs fit into the frameworks.
Table 2: Results of the fix-backbone design task on SAbDab.
Model	CDR-H1	CDR-H2	CDR-H3
AAR (%)↑	Perplexity ↓	AAR (%)↑	Perplexity ↓	AAR (%)↑	Perplexity ↓
RosettaFix	36.29 ± 0.2	14.78 ± 0.01	37.70 ± 0.3	12.30 ± 0.02	28.13 ± 0.1	24.05 ± 0.14
Structured-TF	53.24 ± 3.2	8.61 ± 0.08	49.87 ± 1.2	10.27 ± 0.06	30.29 ± 0.4	19.65 ± 0.11
DiffAb	59.91 ± 1.2	6.44 ± 0.05	59.14 ± 1.8	6.92 ± 0.08	33.30 ± 0.5	16.84 ± 0.12
GVP-GNN	62.72 ± 1.5	4.08 ± 0.03	62.48 ± 1.7	4.77 ± 0.09	34.59 ± 0.6	15.79 ± 0.13
HTP	86.01 ± 1.1	1.71 ± 0.04	64.46 ± 1.3	3.29 ± 0.06	43.25 ± 1.2	9.01 ± 0.10
Model	CDR-L1	CDR-L2	CDR-L3
AAR (%)↑	Perplexity ↓	AAR (%)↑	Perplexity ↓	AAR (%)↑	Perplexity ↓
RosettaFix	35.42 ± 0.3	15.82 ± 0.01	36.76 ± 0.2	14.67 ± 0.01	32.17 ± 0.1	18.01 ± 0.00
Structured-TF	56.73 ± 3.1	7.63 ± 0.10	52.11 ± 1.8	8.93 ± 0.08	43.48 ± 0.7	13.88 ± 0.02
DiffAb	58.82 ± 1.6	6.89 ± 0.06	55.40 ± 1.2	7.16 ± 0.05	47.31 ± 0.5	10.60 ± 0.02
GVP-GNN	60.18 ± 1.4	5.48 ± 0.05	59.66 ± 1.5	6.48 ± 0.06	51.34 ± 0.6	8.27 ± 0.03
HTP	93.62 ± 1.5	1.09 ± 0.08	91.46 ± 1.7	1.58 ± 0.10	80.71 ± 1.1	2.66 ± 0.10
Baselines. We pick up a broad range of existing mechanisms for comparison. Rosetta Antibody Design (RAbD) [9] is an antibody design software based on Rosetta energy functions. RefineGNN [10] is an auto-regressive model that first consider 3D geometry for antibody design and is E(3)-invariant. However, its original version is merely conditioned on the framework region rather than the antigen and the remaining part of the antibody. We follow Jin et al. [11] and replace its encoder with a message-passing neural network (MPNN) encoder and use the attention layer to extract information from the antibody-antigen representation. We name this modified model C-RGNN to make a distinction from its original architecture. Hierarchical Equivariant Refinement (HERN) [11] is the abbreviation of Hierarchical Equivariant Refinement Network, which employs a hierarchical MPNN to predict atomic forces and refine a binding complex in an iterative and equivariant manner. Its autoregressive decoder progressively docks generated antibodies and builds a geometric representation of the binding interface to guide the next residue choice. Multichannel Equivariant Attention Network (MEAN) [14] adopts a multi-round progressive full-shot scheme instead of an autoregressive one to output both 1D sequences and 3D structures. DiffAb [12] is a diffusion-based mechanism that achieves state-of-the-art performance on antibody design recently. It consists of three diffusion processes for amino acid types, coordinates, and orientations, respectively.
Results and Analysis. We run each model three times with different random seeds and report the mean and standard deviation of each metric in Table 1, where metrics are labeled with  $\uparrow/\downarrow$  if higher/lower is better, respectively. It can be found that our model outperforms all baselines by a significant margin in terms of AAR, RMSD, and TM-Score. Specifically, HTP brings an improvement of  $53.97\%$ ,  $47.42\%$  and  $51.55\%$  in AAR and  $67.54\%$ ,  $57.25\%$  and  $29.75\%$  in RMSD over the state-of-the-art DiffAb in H1, H2 and H3, respectively. This implies that our HTP might have a higher success rate in designing new antibodies targeting the given antigen. Moreover, it can also be expected that the performance of our HTP can be further improved as the number of antibodies and antigens with solved 3D structures keeps increasing. In addition, the light chain generally has a higher AAR and a lower RMSD than the heavy chain. For example, our model achieves nearly  $90\%$  AAR in CDR-L1 and CDR-L2. This phenomenon is consistent with the fact that CDR in the heavy chain is much longer and more variant.
3.2 Fix-backbone Sequence Design
Task and Metrics. This problem is more straightforward than the previous co-design task. The backbone structure of CDRs is given and only requires the CDR sequence's design. The fixed-backbone design is a common setting in protein design [39, 40] with an alias of inverse folding. We rely on the metrics of AAR introduced in Section 3.1 to examine the generated antibodies. The metric RSMD is ruled out since the backbone structures are fixed. We also rely on perplexity (PPL) to understand the uncertainty of model predictions, which is common for evaluating language models in natural language processing. In short, the perplexity calculates the cross entropy loss and takes its exponent.
Table 3: Comparison with pretrained antibody-specific language models.
PLMs	CDR-H3
AAR (%)↑	RMSD ↓	TM-Score ↑
-	25.31 ± 0.7	2.95 ± 0.02	0.9391 ± 0.005
AbLang	28.46 ± 1.6	2.88 ± 0.17	0.9435 ± 0.009
AntiBERTa	34.52 ± 1.2	2.41 ± 0.13	0.9473 ± 0.006
HTP	40.98 ± 1.5	2.06 ± 0.03	0.9621 ± 0.005
Baselines. Apart from some baselines that have been listed in the sequence-structure co-design problem, we compare HTP with three additional approaches. RosettaFix [41] is a Rosetta-based software for computational protein sequence design. Structured Transformer [42] (Structured-TF) is an auto-regression generative model that is able to sample CDR sequence given the backbone structure. GVP-GNN [43] extends standard dense layers to operate on collections of Euclidean vectors and performs geometric and relational reasoning on efficient representations of macromolecules.
Results. Table 2 documents the result. It is clearly found that our model achieves the highest AAR and the lowest PPL compared to all the baseline algorithms. To be specific, HTP leads to an increase of  $37.13\%$ ,  $3.01\%$ , and  $18.47\%$  in AAR over the best GVP-GNN in H1, H2, and H3, separately, and  $28.76\%$ ,  $24.28\%$ , and  $49.90\%$  in L1, L2, and L3, respectively. This demonstrates that our model is also effective in capturing the conditional probability of sequences given backbone structures. Furthermore, we can also observe that CDR-H3 is the hardest segment in comparison to other regions as its AAR is usually lower than  $50\%$ . Meanwhile, the average AAR of the light chain is more than  $75\%$ , indicating that the light chain maintains less diversity than the heavy chain.
3.3 Discussion
Comparison with Antibody-specific Language Models. Recently, emerging efforts have been paid to train large antibody-specific language models. Studies have also demonstrated that these models can capture biologically relevant information and be generalized to various antibody-related applications, such as paratope position prediction. Here, we make a further investigation into the efficacy of these antibody-specific language models for sequence-structure co-design. To be precise, we abandon the pretraining stages of the single-protein and antibody sequence levels and directly leverage external antibody-specific language models. As shown in Table 3, AntiBERTa and AbLang provide biologically relevant information that is beneficial for the challenge of co-design with an increase of  $12.44\%$  and  $36.38\%$  in AAR. However, their improvements are much smaller than those of PLMs trained through the first two levels of tasks in HTP.
Up- and Downstream Protein. Recently, Wang et al. [44] proposed a joint sequence-structure recovery method based on RosettaFold to scaffold functional sites of proteins. They perform the inpainting and fix-backbone sequence design tasks without immediate up- and downstream protein visible. However, we discover that our HTP can generate adequate diversity without the need to mask neighboring residues, that is., up- and downstream proteins. This difference stems from the fact that our HTP considers the entire antigen and available antibody as the context to complete the masked CDR rather than only depending on the tiny context of up- and downstream protein.
Data Leakage of Language Models. It is undisputed that the evaluation of design methods that use PLMs should be stringent and ensure that the test data were not previously seen by those pretrained models. However, ESM-2 is trained on all protein sequences in the UniRef database (September 2021 version), while our test set includes sequence released after December 2021, as well as structures with any CDR similar to those released after this date (with sequence identity higher than  $50\%$ ). It is fairly possible that the training set of ESM-2 includes antibody sequences similar to the test set, leading to the intolerable data leakage problem.
Here, we conducted additional experiments aimed at assessing the contribution of ESM-2 in directly recovering CDR sequences. Specifically, we abandon the structural information and re-generate CDRs entirely based on sequential information. Towards this end, we first extract residue-level representations via (fixed-weight) ESM-2 and feed them to a three-layer perceptron to predict the masked CDR-H3, where no antigen sequences are given. The results show that this algorithm only
achieved an AAR of  $14.63\%$  to recover CDR-H3, much lower than all baseline methods such as RAdD  $(21.73\%)$  and our HTP  $(40.98\%)$ . This compellingly demonstrates that ESM-2 is not the primary driver of our favorable numerical outcomes and the experimental benefits brought forth by HTP are not due to data leakage. This also accords with ATUE's [45] findings that ESM models perform well in low-antibody-specificity-related tasks but can even bring negative impacts for high-antibody-specificity-related tasks. We have also provided adequate evidence in Appendix 4 to show that the other pretraining resources are free from any possible data leakage concern.
Antibodies are crucial immune proteins produced during an immune response to identify and neutralize the pathogen. Recently, several machine learning-based algorithms have been proposed to simultaneously design the sequences and structures of antibodies conditioned on the 3D structure of the antigen. This paper introduces a novel approach called the hierarchical training paradigm (HTP) to address the co-design problem. It leverages both geometric neural networks and large-scale protein language models and proposes four levels of training stages to efficiently exploit the evolutionary information encoded in the abundant protein sequences and complex binding structures. Extensive experiments confidently show that each stage of HTP significantly contributes to the improvements of the deep learning model's capacity in predicting more accurate antibody sequences and storing its corresponding 3D structures. Instead of focusing on the architecture side, our study hopes to shed light on how to better blend protein data of different modalities (i.e., one- and three-dimensions) and domains (i.e., antibodies, and non-antibodies) for tackling the sequence-structure co-design challenge, which has been long ignored by existing works.
5 Limitations and Future Work
In spite of the promising progress of our HTP, there is still some space left for future explorations. First, more abundant databases can be exploited in our framework. For example, AntiBodies Chemically Defined (ABCD) [46] is a large antibody sequence database that can be used to improve the capacity of protein language models at the second level. We do not use it in our work because our request for this database has not been approved by the authors so far. Secondly, we fix the language models during the last two levels of training (i.e., levels that need complex structure prediction) for simplicity and use them as the node feature initializer. It might be beneficial if both the PLM and the geometric encoder are tuned.
This work was supported by National Key R&D Program of China (No. 2022ZD0115100), National Natural Science Foundation of China Project (No. U21A20427), and Project (No. WU2022A009) from the Center of Synthetic Biology and Integrated Bioengineering of Westlake University. F.W. and S.L. led the research. F.W. contributed technical ideas and developed the method and performed analytics. S.L. provided evaluation and suggestions. The authors thank Professor Siqi Sun and Tao You from Shanghai Artificial Intelligence Lab (SAIL) for their efforts in processing the OAS data, and Professor Buyong Ma from Shanghai Jiaotong University for his comments on improving the quality of the paper.
[1] Gary W Litman, Jonathan P Rast, Michael J Shamblott, Robert N Haire, Michele Hulst, William Roess, Ronda T Litman, Kristin R Hinds-Frey, Anna Zilch, and Chris T Amemiya. Phylogenetic diversification of immunoglobulin genes and the antibody repertoire. Molecular biology and evolution, 10(1):60-72, 1993.
[2] Charles A Janeway, Paul Travers, Mark Walport, and Donald J Capra. Immunobiology. Taylor & Francis Group UK: Garland Science, 2001.
[3] Matthew IJ Raybould, Claire Marks, Konrad Krawczyk, Bruck Taddese, Jaroslaw Nowak, Alan P Lewis, Alexander Bujotzek, Jiye Shi, and Charlotte M Deane. Five computational
[55] Pascal Sturmfels, Jesse Vig, Ali Madani, and Nazneen Fatema Rajani. Profile prediction: An alignment-based pre-training task for protein sequence models. arXiv preprint arXiv:2012.00195, 2020.
[56] Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In International Conference on Machine Learning, pages 8844–8856. PMLR, 2021.
[57] Surojit Biswas, Grigory Khimulya, Ethan C Alley, Kevin M Esvelt, and George M Church. Low-n protein engineering with data-efficient deep learning. Nature methods, 18(4):389-396, 2021.
[58] Tristan Bepler and Bonnie Berger. Learning the protein language: Evolution, structure, and function. Cell systems, 12(6):654-669, 2021.
[59] Zichen Wang, Steven A Combs, Ryan Brand, Miguel Romero Calvo, Panpan Xu, George Price, Nataliya Golovach, Emmanuel O Salawu, Colby J Wise, Sri Priya Ponnapalli, et al. Lm-gvp: A generalizable deep learning framework for protein property prediction from sequence and structure. bioRxiv, 2021.
[60] Jesse Vig, Ali Madani, Lav R Varshney, Caiming Xiong, Richard Socher, and Nazneen Fatema Rajani. Bertology meets biology: interpreting attention in protein language models. arXiv preprint arXiv:2006.15222, 2020.
[61] José Jiménez, Stefan Doerr, Gerard Martínez-Rosell, Alexander S Rose, and Gianni De Fabritiis. Deepsite: protein-binding site predictor using 3d-convolutional neural networks. Bioinformatics, 33(19):3036–3042, 2017.
[62] Jiyu Cui, Fang Wu, Wen Zhang, Lifeng Yang, Jianbo Hu, Yin Fang, Peng Ye, Qiang Zhang, Xian Suo, Yiming Mo, et al. Direct prediction of gas adsorption via spatial atom interaction learning. Nature Communications, 14(1):7043, 2023.
[63] Fang Wu and Stan Z Li. Diffmd: A geometric diffusion model for molecular dynamics simulations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 5321-5329, 2023.
[64] Fang Wu, Siyuan Li, Lirong Wu, Stan Z Li, Dragomir Radev, and Qiang Zhang. Discovering the representation bottleneck of graph neural networks from multi-order interactions. arXiv preprint arXiv:2205.07266, 2022.
[65] UniProt Consortium. Uniprot: a hub for protein information. *Nucleic acids research*, 43(D1): D204–D212, 2015.
[66] Robert D Finn, Alex Bateman, Jody Clements, Penelope Coggill, Ruth Y Eberhardt, Sean R Eddy, Andreas Heger, Kirstie Hetherington, Liisa Holm, Jaina Mistry, et al. Pfam: the protein families database. Nucleic acids research, 42(D1):D222-D230, 2014.
[67] Can Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self-supervised learning. arXiv preprint arXiv:2204.04213, 2022.
[68] Yuzhi Guo, Jiaxiang Wu, Hehuan Ma, and Junzhou Huang. Self-supervised pre-training for protein embeddings using tertiary structures. 2022.
A Experimental setting
Data Descriptions for Different Stages of HTP. (1) For the single-protein sequence level, we employ the state-of-the-art ESM-2 [47], which outperforms all tested single-sequence protein language models across a wide range of structure prediction tasks and enables prediction of the atomic resolution structure. It is trained on 86 billion amino acids across 250 million protein sequences that span evolutionary diversity. Specifically, ESM uses UniRef50, September 2021 version. The training dataset was partitioned by randomly selecting  $0.5\%$  ( $\approx$  250,000) sequences to form the validation set. The training set has sequences removed via the procedure described in Hie et al. [48]. ESM-2 runs the MMseqs search to obtain the query and target databases. All train sequences that match a validation sequence with  $50\%$  sequence identity in this search are removed from the train set. The details of the ESM series can be found in https://github.com/facebookresearch/esm.
(2) For the antibody sequence level, we use the Observed Antibody Space database (OAS) [31] and its subsequent update [32] as the pretraining data. It currently contains more than one billion sequences, from more than 80 different studies that cover various immune states, organisms, and individuals, which can be downloaded from its official website at https://opig.stats.ox.ac.uk/webapps/oas/. We upload the processed paired data in https://pan.baidu.com/s/181B8gl9Maf0nnNPIw83ZzA?pwd=1212 (password: 1212) as well as the unpaired data in https://pan.baidu.com/s/161gU8fso6rz6-QGfNoCoHQ?pwd=96uF (password: 96uf).
(3) For the protein-protein complex structure level, we use the Database of Interacting Protein Structures (DIPS) [35]. It is a larger protein complex structure dataset than existing antibody-antigen complex structure datasets and is extracted from the Protein Data Bank [36]. We attain the database from Atom3d in Zendo https://zenodo.org/record/4911102, which is a collection of both novel and existing benchmark datasets spanning several key classes of biomolecules. Referring to Atom3d, we split protein complexes by sequence identity at  $30\%$ , resulting in train/validation/test sets with 87,303/31,050/15,268 instances.