SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(4) For the antibody-antigen complex structure level, we select all available antibody-antigen protein complexes from SAbDab [16] at https://opig.stats.ox.ac.uk/webapps/newsabdab/sabdab/, leading to a dataset containing 9,823 structures. CDRs are identified using the antibody numbering program AbRSA [49]. Following the setting in [12], the chosen data points are divided into training and test data based on their release date and CDR sequence identity. To be explicit, the test split contains protein structures released after December 24, 2021, as well as structures with any CDR similar to those released after this date with a sequence identity higher than  $50\%$ . The antibodies in the test set are further clustered with a CDR sequence identity of  $50\%$  to eliminate duplicates, resulting in 20 antibody-antigen structures. The training and validation splits just include complexes not involved during the curation of the test split. After that, we randomly divided the remaining complexes with a ratio of  $90\%$  and  $10\%$  into training and validation sets.
Dataset Sequence Similarity. In our splitting strategy, the test set of SAbDab 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 pretraining datasets include antibody sequences similar to the test set. To navigate this issue, we have conducted a comprehensive analysis of the sequence similarity between the various pretraining data sources and the test set in SAbDab. This analysis includes general protein sequences from general protein-protein complexes from DIPS and antibodies from OAS. We have plotted the sequence similarity distributions in Figure 4 and present the statistical findings of this analysis in Table 4.
Table 4: The statistics of the sequence similarity between our splitted SAbDab test set and different pretraining datasets.
Dataset	Mean	Std.	Min.	25%	50%	75%	Max.
DIPS	0.188	0.036	0.000	0.183	0.198	0.208	0.429
OAS	0.246	0.017	0.200	0.235	0.243	0.254	0.401
(b)
Figure 4: The distributional plot of sequence similarity between different pre-training resources and the test set in SAbDab.
The statistical findings unequivocally demonstrate that the highest sequence similarity values across these three datasets are consistently below 0.5. This empirical evidence serves as a strong basis for our resolute assertion that neither DIPS, nor OAS encompasses sequences that exhibit significant similarity to the SAbDab test set. With this robust evidence at hand, we maintain our firm confidence that the hierarchical training paradigm we have employed is free from any potential data leakage concerns.
Implementation Details. HPT is implemented in PyTorch and PyTorch Geometric packages. For all four training stages, we leverage an Adam optimizer [50] with a weight decay of 1e-5. All experiments are run on multiple A100 GPUs, each with a memory storage of 80G.
(1) For ESM-2 in the single-protein sequence level training, we adopt a middle-size version, which has a parameter number of 150M, 30 layers, and a hidden dimension of 640. Besides, we append the ESM-2 with a three-layer perceptron to forecast the residue type for MLM.
(2) For the antibody sequence level training, we use a batch size of 2 to avoid out-of-memory error and 4 workers to load the data. The number of epochs is 100 and starting learning rate is 1e-5. Apart from that, we utilize a ReduceLROOnPlateau scheduler with a factor of 0.6, patience of 5 epochs, and a minimum learning rate of 1e-7.
(3) For protein-protein complex structure level training, we use a batch size of 32, 1000 epochs, and 4 works to speed up data loading. The starting learning rate is 1e-4, and a ReduceLROnPlateau scheduler is utilized to automatically adjust the learning rate with a factor of 0.6 and patience of 3 epochs. We adopt a distance threshold of  $8.0\mathring{\mathrm{A}}$  to determine the connection between different graph nodes (i.e., the alpha carbon of each residue). As for the loss weight balance, we set  $\lambda = 1$ .
(4) For the antibody-antigen complex structure level training, we also adopt the distance threshold of  $8.0\mathring{\mathrm{A}}$  to build the graph connection. For the random initialization of the CDR coordinates, we use a noise of  $\epsilon = 0.1$ . As for the other important hyperparameters, we use a grid search mechanism to find the optimal combination. Notably, the geometric neural networks used in the third and fourth levels are matched to each other. If we alter the setting of GGNNs in the antibody-antigen complex structure level training, we need to retrain it in the protein-protein complex structure level first. The entire hyperparameter search space is depicted in Table 5.
Reproduction of Baselines. Concerning the implementation of several baseline methods, we use the official repositories for conditional RefineGNN (https://github.com/wengong-jin/RefineGNN/), HERN (https://github.com/wengong-jin/abdockgen), DiffAb (https://github.com/luost26/diffab). To reproduce the performance of existing antibody-specific pretrained PLMs, we download the code from https://github.com/alchemab/antiberta for AntiBERTa and https://github.com/oxpig/AbLang for AbLang. In our comparison, we directly use their pretrained residue features as the input for GGNNs without any fine-tuning.
Table 5: Hyperparameters setup for HTP.
Hyperparameters Search Space	Symbol	Value
Training Setup
Epochs	-	[100, 500, 1000]
Batch size	-	[32, 64, 128]
Learning rate	-	[1e-4, 5e-5, 1e-6, 1e-7]
Warmup	-	[Yes, No]
Warmup epochs	-	[10, 20]
Loss Balance weight for Coordinates and Residue Types	λ	[0.1, 0.3, 0.5, 0.7]
GNN Architecture
Dropout rate	-	[0.1, 0.2]
Number of GNN layers	L	[2, 4, 6]
Tanh activation function	-	[Yes, No]
Coordinate Normalization	-	[Yes, No]
The hidden dimension of node representations	-	[320, 640]
The hidden dimension of edge representations	-	[16, 32, 64]
Code Availability. All relevant Python code to reproduce the results in our paper is stored in the GitHub repository at https://github.com/smiles724/HTP.
B Additional Results
B.1 Ablation Study
We investigate the effectiveness and necessity of each component of our HTP. As shown in Table 6, the elimination of the level of protein-protein complex structure level induces performance degradation, where RMSD increases from 2.06 to 2.49. Moreover, we implement a variant of HTP by replacing features obtained by pretrained PLMs with learnable embedding features, whose performance is worse than HTP. To be concise, AAR declines from 40.98 to 25.31, and RMSD increases from 2.06 to 2.65. In summary, our HTP brings significant relative improvements of  $78.56\%$  in AAR,  $41.97\%$  in RMSD, and  $2.94\%$  in TM-Score. This phenomenon strongly supports the superiority of our approach over existing naive co-design algorithms that are trained only on antibody-specific structure data.
Table 6: Effects of each module, where SPS stands for the single-protein sequence level, PPCS denotes the protein-protein complex structure level, and AS represents the antibody sequence level. The last row computes the relative improvements of HTP over the primitive baseline without any protein data augmentation.
SPS	AS	PPCS	SAbDab (CDR-H3)
AAR (%) ↑	RMSD ↓	TM-Score
1	X	X	X	22.95 ± 0.5	3.55 ± 0.01	0.9146 ± 0.003
2	✓	X	X	33.87 ± 0.8	2.77 ± 0.04	0.9450 ± 0.006
3	✓	✓	X	38.42 ± 1.6	2.49 ± 0.03	0.9538 ± 0.004
4	X	X	✓	25.31 ± 0.7	2.95 ± 0.02	0.9391 ± 0.005
5	✓	✓	✓	40.98 ± 1.5	2.06 ± 0.03	0.9621 ± 0.005
Imp.	-	-	-	78.56%	41.97%	2.94%
Antibody Design. The majority of old-school computational approaches for antibody design are based on sampling algorithms over handcrafted and statistical energy functions to iteratively modify protein sequences and structures [8, 9]. These physics-based algorithms are computationally expensive and prone to be stuck in local energy minimum, which triggers the adaptation of deep
learning in this sub-field. The initial researchers [4, 5] use pure PLMs to generate protein sequences but disregard the available antigen structures.
To circumvent this, Jin et al. [10] introduce RefineGNN, the first co-design architecture that aims to neutralize SARS-CoV-2. Later, HERN [11] is proposed as a more general version for paratope docking and design, which opens the door to produce antibodies given arbitrary antigen structures. Subsequent efforts are spent in either modifying the generative style or utilizing more advanced deep learning architectures such as diffusion denoise probabilistic models (DDPMs). For example, DiffAb [12] achieves atomic-resolution antibody design with SO(3)-equivariance, while MEAN [14] corrects the autoregressive manner with a full-shot one to prevent low efficiency and accumulated errors during inference.
Protein Sequence Modeling. Sequence-based protein representation learning is mainly inspired by the field of natural language processing. A large body of early work is focused on modeling individual protein families [51], solving problems such as functional nanobody design [5]. The success of this method then motivates the prospective trend to model large-scale databases of protein sequences by means of unsupervised learning. This line of study targets capturing the biochemical and co-evolutionary knowledge that underlies a large-scale protein sequence corpus by self-supervised pertaining. Thanks to them, a number of pertaining objects have been explored such as the next amino acid prediction [4, 24], masked language modeling (MLM) [51, 22], pairwise MLM [52], contrastive predictive coding [53], conditional generation [54], and position-specific scoring matrix prediction [55]. In addition, another line [56, 57] is based on multiple sequence alignment (MSA), leveraging sequences within a protein family to seize the conserved and variable regions of homologous sequences. Notably, some schemes for protein sequence modeling also seek to incorporate structural information in either the pretraining stage [58] or the finetuning stage [59].
The improvements in model scale and architecture are also crucial to the recent achievement of PLMs. Explicitly, Rao et al. [51] evaluate various PLMs in a panel of benchmarks and discover that multi-head attention outpaces the Potts model in contact prediction, even if using a single sequence for inference. Concurrently, Vig et al. [60] observe that specific attention heads of pretrained Transformers have straight correlations with protein contact. Others [24] investigate a variety of Transformer variants and demonstrate that large Transformers can obtain state-of-the-art features in various tasks. Apart from that, the latest ESM-2 [47] trains the largest PLM with 15B parameters and shows that as models are scaled, they learn information enabling the protein structure prediction at the resolution of individual atoms.
Protein Structure Learning. With the rapid advance of geometric deep learning, it has become increasingly attractive and challenging to represent and reason about the structures of macromolecules in 3D space. For the sake of encoding spatial information in protein structures including bond lengths and dihedral angles, numerous 3D geometric neural networks such as 3DCNN [61] or GNNs [42, 43, 62, 63] have been invented. They excel at capturing complex interactions between sets of amino acids [64] and attain pivotal Euclidean geometry, e.g., E(3) or SE(3)-equivariance and symmetry.
However, compared to protein sequences in databases like UniProt [65] or Pfam [66], the known structures in the PDB are scarce and hard to obtain. Therefore, it becomes urgent to develop structure-based mechanisms to efficiently learn protein representations with much fewer pretraining data. For instance, Hermosilla and Ropinski [17] use contrastive learning in terms of molecular substructures to help models understand protein structure similarity and functionality. Moreover, Chen et al. [67] propose a self-supervised framework that predicts angles and inter-residue distances. Additionally, Guo et al. [68] present a coordinate denoising score matching method. Wu et al. [21] put forward a novel prompt-based denoising conformation generative pretraining method based on the trajectories of molecular dynamics simulations. A recent attempt [34] makes a combination of both contrastive learning and self-prediction with more intriguing augmentation functions. Despite this progress, all of them are dealing with single-protein structures. No preceding studies have considered structure-based pretraining in the circumstance of multiple proteins. That is, how to pretrain on protein-protein complex, or more specifically, the antibody-antigen complex, remains unexplored.