SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $I$  represents the identity matrix and  $J$  refers to the all-ones matrix. If the condition  $\forall i, j, \|X_i - \hat{Y}_j\| = \hat{D}_{ij}$  holds, the following equation exists
$$
H _ {m} F H _ {n} = H _ {m} (B - 2 E + C) H _ {n} = - 2 H _ {m} E H _ {n} = - 2 H _ {m} X ^ {\top} R Y H _ {n}. \tag {20}
$$
Table 1: Statistics of datasets.
Datasets	# Pairs of Proteins	# Residues per Protein	# Atoms per Protein
Training Set	4890	565.9 (±264.9)	4334.4 (±2028.3)
DB5.5 (Test)	24	428.4 (±132.0)	3308.0 (±1000.5)
VH-VL (Test)	68	230.1 (±5.4)	1749.9 (±53.7)
AB-AG (Test)	68	433.0 (±72.0)	3346.7 (±568.8)
By performing singular value decomposition (SVD) for  $H_{m}X^{\top}$  and  $YH_{n}$  respectively
$$
H _ {m} X ^ {\top} = U _ {X} \Sigma_ {X} V _ {X} ^ {\top}, \quad Y H _ {n} = U _ {Y} \Sigma_ {Y} V _ {Y} ^ {\top}, \tag {21}
$$
the rotation matrix can be solved as
$$
R = - \frac {1}{2} V _ {X} \Sigma_ {X} ^ {- 1} U _ {X} ^ {\top} H _ {m} F H _ {n} V _ {Y} \Sigma_ {Y} ^ {- 1} U _ {Y} ^ {\top}. \tag {22}
$$
Given the rotation matrix  $R$ , the translation vector  $\mathbf{t}$  can be gotten using
$$
H _ {m} F = H _ {m} (B - 2 E + C) = H _ {m} B - 2 H _ {m} X ^ {\top} R Y - 2 H _ {m} X ^ {\top} \mathbf {t} \mathbf {1} _ {n} ^ {\top}, \tag {23}
$$
$$
\mathbf {t} = - \frac {1}{2} V _ {X} \Sigma_ {X} ^ {- 1} U _ {X} ^ {\top} H _ {m} (F - B + 2 X ^ {\top} \hat {R} Y) \mathbf {1} _ {n}. \tag {24}
$$
By replacing random initialization with the above spectral initialization, the gradient descent will approach the global minima faster and better. We will empirically verify this conclusion through subsequent experiments.
Datasets We leverage Docking Benchmark 5.5 (DB5.5) [52], a gold standard dataset tailored for rigid docking. Following the experimental setting of EquiDock [25], the dataset is randomly partitioned into a test split of size 24. For a comprehensive comparison, we curate two datasets of antibodies (VH-VL) and antibody-antigen complexes (AB-AG) from Protein Data Bank (PDB) [3] and expect them to become new benchmarks. Specifically, the training set consists of 4,890 complexes containing at least one antibody chain, while the test set comprises 68 antibody-antigen complexes released after October 2022. The docking results of variable heavy-light chains and antigen-antibody can be separately evaluated. We direct the readers of interest to Appendix A for extraction criteria and detailed identifier lists. The statistics of the dataset are summarized in Table 1.
Evaluation Protocol We compare BiDock with two categories of representative methods, including three docking software ZDOCK [13], ClusPro [29], and HDOCK [58], and two deep learning models EquiDock [25], and AlphaFold-Multimer (Multimer for short [23]). To measure the quality of
Table 2: Quantitative results on protein docking. (bold: best; underline: runner-up)
Datasets	Metrics	ZDOCK	ClusPro	HDOCK	EquiDock	Multimer	BiDock
DB5.5	RMSD ↓	12.491±6.294	14.135±8.153	11.328±8.073	14.982±5.304	7.797±7.428	7.280±8.117
TM-score ↑	0.689±0.114	0.702±0.118	0.742±0.167	0.714±0.114	0.821±0.162	0.847±0.158
DockQ ↑	0.084±0.113	0.118±0.192	0.314±0.390	0.030±0.029	0.469±0.396	0.564±0.369
VH-VL	RMSD ↓	10.982±3.864	5.899±5.688	2.032±2.388	18.293±2.871	1.325±0.530	1.242±0.602
TM-score ↑	0.596±0.075	0.792±0.156	0.926±0.100	0.559±0.017	0.962±0.020	0.966±0.021
DockQ ↑	0.108±0.134	0.404±0.277	0.705±0.201	0.032±0.016	0.765±0.094	0.773±0.187
AB-AG	RMSD ↓	18.892±3.757	15.670±6.896	15.779±6.364	18.468±2.706	13.650±5.886	9.707±8.759
TM-score ↑	0.504±0.059	0.596±0.143	0.612±0.137	0.502±0.096	0.640±0.124	0.773±0.187
DockQ ↑	0.035±0.031	0.108±0.181	0.090±0.187	0.043±0.017	0.108±0.172	0.342±0.351
maxDockQ ↑	0.042±0.043	0.136±0.220	0.111±0.237	0.043±0.018	0.125±0.214	0.414±0.386
Figure 4: An intuitive comparison of DockQ metric on antibody-antigen docking. The box inside (a) violin plot represents 25-75 percentiles, and the median is shown by a white dot. Scatters in (b) scatter plot appear under the dashed diagonal line, indicating that BiDock outperforms Multimer on these complexes.
Table 3: Total inference time of different methods on antibody-antigen docking. (unit: hour)
Methods	Inference time
ZDOCK	72.61
ClusPro	87.27
HDOCK	20.40
EquiDock	0.60
Multimer	1.07
BiDock	1.47
predictions, we report universally accepted metrics Root Mean Square Deviation (RMSD), TM-score (Template Modeling score), and DockQ [6]. Please note that in the context of antibody-antigen docking, the original DockQ metric evaluates the docking performance by treating the entire antibody as a single entity to the antigen. Furthermore, we can assess the docking results separately for variable heavy/light chain (VH/VL) to the antigen and denote the maximum value as maxDockQ. Refer to Appendix B for details of experiments, including implementation and hyperparameters.
4.1 Main Results and Analysis
Docking Results Table 2 demonstrates that BiDock generally produces acceptable predictions on all datasets. Notably, BiDock achieves a significant improvement in performance on the more challenging antibody-antigen docking, with a  $234\%$  relative gain over the runner-up on the DockQ metric. These results highlight the effectiveness of our bi-level optimization, which effectively leverages multimodal information. We also observe that some established docking software, such as HDOCK, still provides reliable predictions. In contrast, deep learning methods may have some leeway in performance. For instance, the mean and deviation of RMSD evaluated from EquiDock are relatively large, indicating that some inappropriate SE(3) transformations are learned.
To be more intuitive, we select representative baselines, HDOCK and Multimer, along with BiDock to analyze their performance differences on antibody-antigen docking. The distribution of DockQ for each method is displayed using a violin plot, and a direct comparison between Multimer and BiDock is also drawn in Figure 4. It reveals that the distribution of DockQ for BiDock is more concentrated around high values, indicating a higher percentage of successful docking predictions. Overall, these results suggest that our proposed BiDock is a promising approach for rigid protein docking.
Visualization To visually demonstrate the superiority of our proposed BiDock, we chose the spike glycoprotein  $7F6Z$  as an example antigen. The spike glycoprotein is crucial for the entry of coronaviruses into host cells, making it a target of great interest for therapeutic intervention and vaccine development. In Figure 5, we align the ground truth structure of this protein with predictions from competitive baselines. Upon inspection, it is evident that predictions from baselines exhibit noticeable deviations from the ground truth structures. Such inaccuracies can significantly impact our understanding of the binding mechanisms and hinder the design of effective interventions. In contrast,
Figure 5: Structure comparison between predictions and the ground truth of protein complex  $7F6Z$ . The ground truth structures are represented in light gray, while predictions are colored cyan.
Table 4: Ablation studies on spectral initialization, bilevel optimization, and masked MSA loss. (number in parentheses: gradient descent steps; unit: thousand; bold: best)
RMSD	DockQ
w/o SI(1)	14.140±8.032	0.187±0.256
w/o SI(2)	11.191±8.754	0.290±0.328
w/o SI(4)	9.806±8.646	0.335±0.347
w/o Bi(2)	10.090±7.817	0.220±0.232
w/o MM(2)	9.821±8.688	0.330±0.350
BiDock(2)	9.707±8.759	0.342±0.351
Table 5: Impacts of different optimizers on antibody-antigen docking. (number in parentheses: gradient descent steps; unit: thousand; bold: best)
RMSD	DockQ
SGD(1)	9.939±8.525	0.333±0.341
SGD(2)	9.940±8.557	0.335±0.344
SGD(5)	9.914±8.556	0.337±0.345
SGD(10)	9.861±8.548	0.337±0.344
Adam(1)	9.780±8.773	0.341±0.350
Adam(2)	9.707±8.759	0.342±0.351
BiDock successfully captures the correct docking interface and accurately predicts the binding pose, highlighting its potential for providing biological insights and aiding drug design.
Computational Efficiency Table 3 shows the total inference time for antibody-antigen complexes. Traditional docking software, involving candidate sampling, ranking, and refinement steps, incurs substantial computational costs. Fortunately, deep learning methods provide a significant speed-up, which is particularly important in efficient screening. Although EquiDock is the fastest, its performance falls short of traditional software due to limitations in leveraging coevolution information and simple networks. On the other hand, Multimer and BiDock exhibit comparable inference times. Considering the performance improvement of our model, this trade-off is acceptable.
Effects of Spectral Initialization Recalling our utilization of spectral estimation to derive a numerical solution for initializing the gradient descent in the inner loop, we conduct ablation experiments to demonstrate the effectiveness of this spectral initialization. Table 4 presents the results on antibody-antigen docking, where the variant without spectral initialization is denoted as "w/o SI" and the number of gradient descent steps is listed in parentheses. It is seen that the "w/o SI" variants exhibit slower convergence and lower accuracy, even when we further increase the number of gradient descent steps. These findings further support the benefits of our proposed spectral initialization in accelerating the optimization process and achieving superior results.
Effects of Bi-level Optimization To justify the effectiveness of bi-level optimization, we use a two-stage strategy: training outer and inner loops separately. Specifically, we use cross-entropy and masked MSA losses to train cross-modal transformer. Based on the resulting distance map, we compute roto-translation transformation with spectral initialization. The results on antibody-antigen docking are presented in Table 4, where "w/o Bi" denotes the variant without bi-level optimization (results reported from the original paper [34]). Results support our contributions of employing bi-level optimization for end-to-end optimization, which customizes the parameter learning of the cross-modal transformer for rigid docking.
Effects of Masked MSA Loss In addition to the essential loss for the distance map, we introduce a masked MSA loss in the outer loop to supervise the learning of evolution representations. We also conduct an ablation study to investigate its significance, as shown in Table 4. The "w/o MM" refers to models without masked MSA loss. According to the results, it can be concluded that the masked MSA loss enables the cross-modal transformer to more effectively leverage the rich evolutionary information and seamlessly integrate it with the structure-modal information.
Effects of Optimizer As shown in Equation (15), we have the flexibility to choose different optimizers for the inner loop. To explore the impact of different optimizers on convergence speed and performance, we compare the effects of SGD and Adam optimizers. We vary the number of
gradient descent steps and evaluate the final performance, as displayed in Table 5. With Adam, BiDock achieves faster convergence and attains higher accuracy, indicating that a better optimizer can navigate the landscape and find better minima during the optimization of the inner loop.
In this study, we introduce BiDock, a novel model for rigid protein docking that tackles the challenge of accurately predicting the 3D structure of protein complexes from unbound states. By formulating this problem as a bi-level optimization, BiDock combines the advantages of integrating multimodal information by a cross-modal transformer in the outer loop and maintaining the rigidity of learning roto-translation transformation in the inner loop. Additionally, we derive a spectral initialization to expedite convergence. The maximum  $234\%$  relative performance improvement validates the effectiveness of BiDock in rigid protein docking.
Limitations and Broader Impact Despite the encouraging results, BiDock does not explicitly incorporate geometric constraints between residues when learning the distance map and does not account for potential atom clashes. Our future work will address these limitations and extend our framework to general proteins. Protein docking deepens our understanding of biological mechanisms and aids in the design of targeted interventions. This research may inspire the AI4Science community to pay more attention to the practical challenges in docking and promote further advancements in this field with significant real-world implications.
Acknowledgments We would like to extend our heartfelt appreciation to the anonymous reviewers who dedicated their valuable time and expertise to meticulously review this paper. Their constructive feedback significantly contributed to the refinement and enhancement of the quality of our work. Our gratitude also extends to our team members, whose insightful discussions and collaborative brainstorming sessions were instrumental in shaping this research.
We wish to express our sincere thanks to the Data Science team at BioMap, especially to Tingting Tang, Nachuan Shan, Shiro Hong, Di Wang, Dacheng Li, Chenrui Xu, Ke Wang, and Cheng Chen, for their exceptional support in data collection, processing, and analysis. Their efforts with complex biological datasets were indispensable to the success of our research. Furthermore, we also express our sincere thanks to the AI Infrastructure team at BioMap, especially to Ming Yang, Chuang Yu, Zedong Zheng, Zezhi Wang, Qian Wang, and Xiaoming Zhang, for their outstanding technical support and seamless collaboration. Their proactive approach and effective troubleshooting played a vital role in ensuring the smooth progress of our research.
Finally, thanks Aaron for proofreading the manuscript.
Author Contributions The author's contributions to this paper encompass a range of vital aspects. Ruijia Wang made substantial contributions to content, writing, and experiment design. Yiwu Sun took on a central role in implementing the bi-level optimization algorithm and conducting experiments. Yujie Luo and Shaochuan Li were crucial in developing the precursor algorithm, xTrimoDock [34], and also aiding in data preparation.
Le Song and Hui Li lead the antibody-antigen structure prediction project in BioMap. They made invaluable contributions by providing algorithmic insights and guidance to the project, influencing the overall conceptualization and direction of the research. Xingyi Cheng, Chuan Shi and Cheng Yang provided valuable guidance in organizing the writing of the paper.
This work was funded by BioMap. The intellectual property of BiDock is  $100\%$  owned by BioMap.
[1] Juhan Bae and Roger B Grosse. Delta-stn: Efficient bilevel optimization for neural networks using structured response jacobians. In NeurIPS, pages 21725-21737, 2020.
[2] Minkyung Baek, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, R Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using a three-track neural network. Science, 373(6557):871-876, 2021.
[3] Protein Data Bank. Protein data bank. Nature New Biol, 233:223, 1971.
[4] Jingxiao Bao, Xiao He, and John ZH Zhang. Deepbsp—a machine learning method for accurate prediction of protein-ligand docking structures. Journal of Chemical Information and Modeling, 61(5):2231-2240, 2021.
[5] Itzhack Y Bar-Itzhack. New method for extracting the quaternion from a rotation matrix. Journal of guidance, control, and dynamics, 23(6):1085-1087, 2000.
[6] Sankar Basu and Björn Wallner. Dockq: a quality measure for protein-protein docking models. PloS one, 11(8):e0161879, 2016.
[7] Jacek Biesiada, Aleksey Porollo, Prakash Velayutham, Michal Kouril, and Jaroslaw Meller. Survey of public domain software for docking simulations and virtual screening. Human genomics, 5(5):1-9, 2011.
[8] Ondrej Bohdal, Yongxin Yang, and Timothy Hospedales. Evograd: Efficient gradient-based meta-learning and hyperparameter optimization. In NeurIPS, pages 22234-22246, 2021.
[9] Jerome Bracken and James T McGill. Mathematical programs with optimization problems in the constraints. Operations research, 21(1):37-44, 1973.
A Details of Dataset
Background of Antibodies Antibodies are vital components of the immune system and are classified into various classes, including IgG, IgM, IgA, IgD, and IgE. Among them, IgG antibodies are the most abundant in the bloodstream and play a primary role in immune responses against pathogens.
As depicted in Figure 6, IgG antibodies exhibit a Y-shaped structure composed of two identical light chains and two identical heavy chains, where heavy chains provide structural stability. Each antibody chain is further divided into distinct regions. (1) The variable regions, referred to as the variable heavy (VH) and variable light (VL) regions, are lo
cated at the tips of the Y arms. These regions contribute to the specificity of antibodies in recognizing and binding to antigens. The VH and VL regions collaborate to form the fragment antigen-binding (Fab) region. (2) At the base of the Y structure, the constant regions, also known as the fragment crystallizable (Fc) region, are important in the effector functions of antibodies. The Fc region interacts with immune cells and triggers immune responses, such as the activation of complement proteins for pathogen destruction and the promotion of phagocytosis.
Given this background knowledge of antibodies, it becomes clear that antibody-antigen docking is fundamental in immune responses, therapeutic applications, vaccine development, and drug discovery. Therefore, our study places a particular emphasis on antibody-antigen docking, contributing to this field by curating a high-quality benchmark. This dataset will serve as a valuable resource for evaluating computational models in predicting antibody-antigen interactions, ultimately facilitating the development of novel therapeutics and immunological interventions.
Figure 6: Structure of an IgG antibody. The heavy chain is colored orange, while the light chain is blue.
Antibody-antigen Benchmark The training set comprises 4,890 complexes of antibody-antigen pairs, each consisting of proteins with a minimum of 30 residues. These complexes encompass three chains, including the light and heavy chains of the antibody, along with one antigen chain. All complexes were released before January 2022. Similarly, the test set consists of 68 antibody-antigen complexes with three chains, released after October 2022. Thus, we ensure that neither baselines nor our proposed model was trained using the test set and avoid data leakage.
In practical applications, obtaining the ground truth structures of antibody-antigen complexes poses significant challenges. Researchers often turn to existing folding models to predict them. To simulate real-world scenarios, we employ a specialized antibody model called xTrimoABFold [55, 56, 54] to predict the conformations of antibodies and AlphaFold2 [28] for antigens. Given these predicted structures as rigid structures, we construct training and test datasets essential for further analysis and investigation. Here are the PDB identifiers and their corresponding chain identifiers for the test set.
8cls:A,H,L; 8dlr:A,H,L; 8dfi:A,H,L; 8dfh:A,H,L; 8dcc:E,L,H; 8dad:H,L,B; 7zr8:A,H,L; 7zf8:H,L,E;
7xxl:C,A,B; 7xh8:A,B,C; 7x26:H,I,K; 7wsl:L,D,H; 7wsi:A,H,L; 7ws6:C,H,I; 7ws2:D,A,E;
7wrz:H,L,R; 7wrv:C,U,V; 7wro:H,L,R; 7wrl:A,B,R; 7wrj:R,A,B; 7wog:C,A,B; 7wlc:E,H,L;
7wef:C,E,L; 7wee:E,H,L; 7wed:E,H,L; 7wcr:A,a,b; 7wbz:A,H,L; 7urq:A,H,L; 7uaq:A,H,L;
7tty:A,H,L; 7ttx:A,H,L; 7ttm:A,H,L; 7tpj:B,H,L; 7tp4:Z,H,L; 7tp3:Z,H,L; 7tlz:A,B,J;
7the:A,B,C; 7tc9:B,H,L; 7t8w:L,H,D; 7t7b:A,H,L; 7t01:A,H,L; 7swp:A,H,L; 7su1:H,L,C;
7str:L,H,C; 7sem:B,F,C; 7sd5:A,H,L; 7sbu:A,H,L; 7sbg:H,L,C; 7sbd:H,L,C; 7sa6:A,H,L;
7s5p:A,H,L; 7rxp:L,H,A; 7rxi:L,H,A; 7rbu:B,H,L; 7qtk:D,B,C; 7n0a:C,A,B; 7lo8:Z,H,L;
7lo7:Z,H,L; 7kql:H,L,T; 7fjc:H,L,E; 7f7e:C,L,E; 7f6z:R,H,L; 7f6y:R,H,L; 7eng:L,H,B;
7ek0:R,H,L; 7ejz:R,H,L; 7ejy:R,H,L; 7e9p:B,L,H
B Details of Implementation
Baselines ZDOCK<sup>5</sup>, ClusPro<sup>6</sup>, and HDOCK<sup>7</sup> are user-friendly local packages suitable for automated experiments or web servers for manual submissions. We select the top-1 predicted structure from each of these methods for subsequent evaluation. For EquiDock<sup>8</sup> and Multimer<sup>9</sup>, we utilize their pretrained models available on GitHub for the inference. It is worth emphasizing that all methods except Multimer are designed for docking two chains. Therefore, during the evaluation, we employ a sequential docking strategy. This entails initially docking the light chain and heavy chain together, followed by treating them as a unified entity for docking with the antigen. And we calculate evaluation metrics using the tools USalign<sup>10</sup> and DockQ<sup>11</sup>.
MSA Extraction We utilize the heuristic approach described in [23] to pair sequences from per-chain multiple sequence alignments (MSAs). Initially, the per-chain MSA sequences are grouped based on species, with the species labels obtained from UniProt's idmapping  $^{12}$ . Within each specific species group, the sequences are paired together. We match the chain MSAs by minimizing the base-pair distance between the chains for prokaryotic species. While in terms of eukaryotic species, we order them based on sequence identity to the target sequence [62]. To reduce computational and memory costs, we employ the MSA clustering approach from AlphaFold2 [28]. We randomly select  $N_{cls} = 252$  sequences as the MSA cluster centers, with the primary protein sequence always set as the first cluster center. The remaining sequences are assigned to their closest cluster based on the Hamming distance.
Sequence-modal Input The sequence modality incorporates information derived from the primary sequence itself and co-evolutionary information obtained from MSAs. Following prior research [28, 23], we extract two types of features: type features  $F^{typ} \in \mathbb{R}^{N_{res} \times 21}$  and primary pair features  $F^{pp} \in \mathbb{R}^{N_{res} \times N_{res} \times 73}$  from the primary sequence, where  $N_{res}$  represents the number of residues. Regarding MSAs, we utilize cluster MSA features  $F^{msa} \in \mathbb{R}^{N_{cls} \times N_{res} \times 49}$ , where  $N_{cls}$  denotes the number of cluster centers. Specifically,
- The type feature  $F^{typ} \in \mathbb{R}^{N_{res} \times 21}$  comprises one-hot representations of the amino acid types, encompassing the 20 known amino acids and one additional category for unknown types.
- The primary pair feature  $F^{pp} \in \mathbb{R}^{N_{res} \times N_{res} \times 73}$  contains positional information within or across chains, including three components. (1) The relative positional feature of size  $[N_{res}, N_{res}, 66]$  represents the relative residue indices, which are clipped between  $[-32, 32]$ . The 66-th index is used to indicate cross-chain pairs. (2) The entity indicator of size  $[N_{res}, N_{res}, 1]$  identifies whether residues  $i$  and  $j$  originate from the same chain. (3) The relative index feature of size  $[N_{res}, N_{res}, 6]$  introduces the relative sym_id<sup>13</sup> indices clipped between  $[-2, 2]$ . The 6-th index is assigned to pairs where the two residues have different sym_ids.
- The cluster MSA feature  $F^{msa} \in \mathbb{R}^{N_{cls} \times N_{res} \times 49}$  consists of five components. (1) The one-hot representation of the amino acid types with size  $[N_{cls}, N_{res}, 23]$ , including 20 amino acids, one unknown type, one gap or missing residue, and one mask token as introduced in Section 3.1. (2) The amino acid distribution of size  $[N_{cls}, N_{res}, 23]$  represents the distribution of amino acid types within each MSA cluster. (3) The deletion indicator of size  $[N_{cls}, N_{res}, 1]$  indicates whether there is a deletion to the left of each residue. (4) The deletion value of size  $[N_{cls}, N_{res}, 1]$  is calculated using the formula  $\frac{2}{\pi}\arctan \frac{c}{3}$ , where  $c$  refers to the number of deletions to the left of each position.