SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruijia Wang $^{12*}$ † Yiwu Sun $^{1}$ † Yujie Luo $^{1}$ † Shaochuan Li $^{1}$ † Cheng Yang $^{2}$
Xingyi Cheng $^{1}$  Hui Li $^{1\ddagger}$  Chuan Shi $^{2\ddagger}$  Le Song $^{1\ddagger}$
<sup>1</sup> BioMap Research
$^{2}$  Beijing University of Posts and Telecommunications
{wangruijia, yangcheng, shichuan}@bupt.edu.cn,
{yiwu,luoyujie,shaochuan,xingyi,lihui,songle}@biomap.com
The structure of protein-protein complexes is critical for understanding binding dynamics, biological mechanisms, and intervention strategies. Rigid protein docking, a fundamental problem in this field, aims to predict the 3D structure of complexes from their unbound states without conformational changes. In this scenario, we have access to two types of valuable information: sequence-modal information, such as coevolutionary data obtained from multiple sequence alignments, and structure-modal information, including the 3D conformations of rigid structures. However, existing docking methods typically utilize single-modal information, resulting in suboptimal predictions. In this paper, we propose xTrimoBiDock $^{\alpha}$  (or BiDock for short) $^{4}$ , a novel rigid docking model that effectively integrates sequence-and structure-modal information through bi-level optimization. Specifically, a cross-modal transformer combines multimodal information to predict an inter-protein distance map. To achieve rigid docking, the roto-translation transformation is optimized to align the docked pose with the predicted distance map. In order to tackle this bi-level optimization problem, we unroll the gradient descent of the inner loop and further derive a better initialization for roto-translation transformation based on spectral estimation. Compared to baselines, BiDock achieves a promising result of a maximum  $234\%$  relative improvement in challenging antibody-antigen docking problem.
Protein-protein interactions (PPIs) are essential to the basic functioning of cells and larger biological systems. Due to their importance, elucidating such interactions up to atomic detail is necessary for understanding multicomponent complexes like ribosomes and discovering protein-based drugs, e.g., antibodies and peptides. However, the experimental golden standard for determining the structure of protein complexes, such as X-ray crystallography and cryo-EM, is extremely time-consuming.
Computational protein docking [49, 7, 57, 45] provides an alternative route to predict the 3D structures of complexes from unbound states. Here, we focus on the fundamental problem of rigid protein docking [25] where no deformations occur within any protein during the docking process.
*Work done during an internship at BioMap
† Contributed equally to this research. Each author's contribution is provided in Section 5.
$\ddagger$  Corresponding authors
$^{4}$  xTrimoBiDock is a member of BioMap's large-scale AI engine "xTrimo" series. "α" denotes the academic version, to distinguish it from the commercial product xTrimoBiDock.
This assumption is reasonable in many biological environments and stabilizes the prediction of natural structures. Therefore, all we need is an appropriate SE(3) transformation shown in Figure 1, i.e., roto-translation transformation, that places the ligand protein at the correct orientation and location concerning the receptor protein. In this context, two types of information are available. The first is sequence-modal information, such as the coevolutionary signals captured in multiple sequence alignments (MSAs). The second is structure-modal information, like 3D coordinates of rigid bodies and bond angles. Both types of information are indispensable for performing rigid protein docking.
Figure 1: Surface views of rigid protein docking. Keep receptor protein at a fixed location, and a rototranslation transformation  $T = (R, \mathbf{t})$  is predicted to place ligand protein at the correct docked pose.
Classical docking software [13, 19, 46, 29, 58] generally follow a three-step framework for predicting complex structures. Firstly, a large number of candidate structures are randomly sampled to explore the conformational space of the complex. Secondly, a scoring function is employed to evaluate and rank the sampled structures based on their compatibility with the binding interface. Finally, the top-ranked structures are refined using an energy model to improve their accuracy and eliminate steric clashes. Due to the evaluation of lots of candidates, these methods tend to be computationally expensive, particularly for high-throughput workflows.
Recently, deep learning has shown significant computational speed-up in this field. EquiDock [25] has emerged as a pioneering method for applying deep learning to rigid protein docking. However, it solely relies on structure-modal information, neglecting the valuable sequence-modal information in databases. This drawback hampers its ability to capture evolutionary constraints and exploit the intricate sequence-structure relationships. Building upon the success of AlphaFold2 [28], AlphaFold-Multimer [23] has been hailed as a breakthrough in directly folding complex structures from amino acid sequences, which considers the sequence-modal information but fails to effectively utilize the given rigid structures, leading to unnatural complexes in some cases.
To overcome the aforementioned limitations, we propose xTrimoBiDock (or BiDock for short), a novel rigid docking model that seamlessly integrates sequence- and structure-modal information through bi-level optimization. (i) To effectively utilize multimodal information, we introduce a cross-modal transformer that injects multimodal information into an inter-protein distance map. (ii) To satisfy rigid docking, we optimize a roto-translation transformation to minimize the difference between the docked pose of unbound proteins and the predicted inter-protein distance map. This framework naturally lends itself to a nested bi-level optimization paradigm, where the outer loop is the learning of the cross-modal transformer and the inner loop is dedicated to solving the roto-translation transformation. Inspired by gradient-based bi-level optimization [24, 32, 11], we unroll the gradient descent of the inner loop for approximation. Due to the ruggedness of the optimization landscape, a substantial number of iterations are required for convergence. Thus, we further derive a better initialization for the roto-translation transformation using spectral estimation. Extensive experiments conducted on diverse datasets and evaluation protocols validate the effectiveness of BiDock.
In summary, our contributions are three-fold:
- We effectively leverage sequence- and structure-modal information for rigid protein docking. By naturally integrating the fusion of multimodities and the docking of rigid bodies through bi-level optimization, we set up a new avenue for solving rigid protein docking.
- We solve the above bi-level optimization with unrolled gradient and spectral initialization. By unrolling the gradient descent of the inner loop and deriving a spectral estimation for initialization, BiDock enhances the convergence and controls the computational cost.
- Comprehensive experiments on three representative datasets demonstrate the effectiveness of the proposed model. Compared to state-of-the-art baselines, BiDock achieves the maximum  $234\%$  relative improvement in challenging but practical antibody-antigen docking.
Molecular Docking Molecular docking aims to characterize the binding poses between small molecule compounds and protein targets [42, 43, 63], playing an essential role in the discovery of effective and safe treatments for various diseases [48]. Traditional computational methods have been greatly enhanced by deep learning techniques [53], which offer increased expressive power in identifying, processing, and extrapolating complex patterns in molecular data [39, 31, 21, 35, 4, 38, 16]. It is important to note that these methods have primarily focused on molecular ligands and often assume the availability of known binding pockets, limiting their direct applicability to protein-protein docking. Protein-protein docking is more challenging due to larger, flexible proteins and the need to predict unknown binding interfaces.
Protein Docking Computational docking software [46, 47, 44, 57, 58, 45, 15] predict complex structures based on a framework of candidate sampling, ranking [37, 6, 22], and refinement [50]. These methods can be financially restrictive and time-consuming, as they often involve scoring and ranking thousands of candidate structures. Recently, deep learning has made significant contributions to structural biology [30, 17, 18, 36]. Notably, AlphaFold2 [28] and RoseTTAFold [2] have been employed to improve protein structure prediction from various angles [26, 40], such as integrating physics-based docking methods [29] or extending multiple sequence alignments [10]. Additionally, methods like AlphaFold-Multimer [23] and HSRN [27] have been developed to simultaneously fold and dock two proteins. Despite their remarkable achievements, these methods violate the rigidity of rigid docking and do not consider unbounded structures. EquiDock [25] is tailored for effective rigid docking but does not fully leverage the evolutionary information encoded in protein sequences, resulting in limited performance improvement compared to traditional docking software.
Bi-level Optimization Bi-level optimization has gained attention in the deep learning community for its ability to handle nested problem structures. It finds applications in various domains, including hyperparameter optimization [14] and metaknowledge extraction [24]. Traditional bi-level optimization methods rely on game theory [51] or mathematical programming [9], which may not scale well to large datasets or have strict mathematical requirements. Alternatively, gradient descent methods offer a promising solution, and they can be divided into explicit gradient update [60], explicit proxy update [1, 8], implicit function [41], and closed-form methods [59]. The first three are approximation methods suitable for general functions, while the last one is an accurate method specifically designed for certain functions. Recent surveys [33, 11] provide a more comprehensive review of these methods. In addition, the merging AI4Science directions, such as topology design [61] and protein representation learning [12], also present nested problem structures that can be compatible with bi-level optimization. By incorporating bi-level optimization into rigid protein docking, we anticipate that our work will have a significant impact at the intersection of these research areas.
3 BiDock Methodology
Multimodal Input The available information on the sequence modality mainly includes information inside the primary sequence itself and co-evolutionary information from MSAs. Following the existing work [28, 23], we extract 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}$  is the number of residues. In terms of MSAs, we leverage cluster MSA features  $F^{msa} \in \mathbb{R}^{N_{cls} \times N_{res} \times 49}$ , where  $N_{cls}$  is the number of cluster centers. For the structure modality, we extract angle features  $F^{ang} \in \mathbb{R}^{N_{res} \times 57}$  and pair features  $F^p \in \mathbb{R}^{N_{res} \times N_{res} \times 88}$  from the rigid protein structures. The angle features provide the orientation and position of each amino acid, while the pair features contain the distance information between amino acids. We present a brief introduction for features and give further details in Appendix B.
Architecture The illustration of the framework is presented in Figure 2. Given sequence-modal features  $\{F^{typ}, F^{msa}, F^{pp}\}$  and structure-modal features  $\{F^{ang}, F^p\}$ , the cross-modal transformer with parameters  $\phi$  transforms available features and integrates them to predict an inter-protein distance map  $\hat{D}$ . To maintain the rigid-body assumption, we define the objective function of learning rototranslation transformation  $(R, t)$  based on the distance map  $\hat{D}$  and coordinates of rigid proteins
Figure 2: The overall framework of BiDock. Taking the information from sequence and structure modality as input, the cross-modal transformer fuses them and predicts an inter-protein distance map. To achieve rigid docking, the roto-translation transformation is learned by minimizing the difference between the docked pose of rigid structures and the predicted distance map. We unroll the gradient descent and further derive a spectral initialization to address the formed bi-level optimization.
$\{X \in \mathbb{R}^{3 \times m}, Y \in \mathbb{R}^{3 \times n}\}$  as follows
$$
\mathcal {L} (R, \mathbf {t}, \phi) = \min  _ {R, \mathbf {t}} \frac {1}{m n} \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} \left(\left\| X _ {i} - R Y _ {j} - \mathbf {t} \right\| - \hat {D} _ {i j} (\phi)\right) ^ {2}, \tag {1}
$$
where  $X_{i}$  is the  $i$ -th column of  $X$  and  $Y_{j}$  is the  $j$ -th column of  $Y$ . Considering that optimizing roto-translation transformation  $(R, \mathbf{t})$  and parameters of cross-modal transformer  $\phi$  constitutes a nested structure, we reformulate the rigid protein docking as a bi-level optimization problem
$$
\phi^ {*} = \underset {\phi} {\operatorname {a r g m i n}} \mathcal {L} ^ {o u t} \left(R ^ {*} (\phi), \mathbf {t} ^ {*} (\phi), \phi\right) \quad \text {s . t .} R ^ {*} (\phi), \mathbf {t} ^ {*} (\phi) = \underset {R, \mathbf {t}} {\operatorname {a r g m i n}} \mathcal {L} (R, \mathbf {t}, \phi), \tag {2}
$$
where  $\mathcal{L}^{out}$  refers to the outer loss for the cross-modal transformer and its specific details will be introduced in the following subsections. To solve this bi-level optimization, we unroll the gradient descent of the inner loop to approximate  $(R^{*}(\phi),\mathbf{t}^{*}(\phi))$  by
$$
\begin{array}{l} R _ {t + 1} = \operatorname {O P T} \left(R _ {t}, \nabla_ {R} \mathcal {L} _ {t} \left(R _ {t}, \mathbf {t} _ {t}, \phi\right)\right), \\ \mathbf {t} _ {\text {一}} = \operatorname {O P T} \left(\mathbf {t} _ {\text {一}} \nabla_ {R} \mathcal {L} _ {t} \left(R _ {t}, \mathbf {t} _ {t}, \phi\right)\right) \end{array} \tag {3}
$$
$$
\mathbf {t} _ {t + 1} = \operatorname {O P T} \left(\mathbf {t} _ {t}, \nabla_ {\mathbf {t}} \mathcal {L} _ {t} \left(R _ {t}, \mathbf {t} _ {t}, \phi\right)\right),
$$
where OPT represents the optimization algorithm, such as the stochastic gradient descent (SGD). However, due to the complexity of the problem and the rugged optimization landscape, the gradient descent in the inner loop often requires a large number of iterations to converge and faces challenges in finding the global minima. Therefore, we further derive spectral initialization to provide a more favorable starting point for better convergence.
3.1 Cross-modal Transformer
Based on the multimodal features, the cross-modal transformer predicts an inter-protein distance map, as depicted in Figure 3. We first use multilayer perceptrons (MLPs) to capture the nonlinear relationship inside features and project them into the same space
$$
P ^ {p p} = \operatorname {M L P} \left(F ^ {p p}\right) \quad P ^ {p} = \operatorname {M L P} \left(F ^ {p}\right)
$$
$$
M ^ {t y p} = \operatorname {M L P} \left(F ^ {t y p}\right) \quad M ^ {m s a} = \operatorname {M L P} \left(F ^ {m s a}\right) \quad M ^ {a n g} = \operatorname {M L P} \left(F ^ {a n g}\right), \tag {4}
$$
where  $\{P^{pp}, P^p \in \mathbb{R}^{N_{res} \times N_{res} \times c_z}\}$  are transformed pair features,  $\{M^{typ}, M^{ang} \in \mathbb{R}^{N_{res} \times c_m}\}$  and  $M^{msa} \in \mathbb{R}^{N_{cls} \times N_{res} \times c_m}$  are transformed intra-sequence features. Then we add pair features as the input pair feature  $P$  of Evoformer [28]
$$
P = P ^ {p p} + P ^ {p}. \tag {5}
$$
Similarly, we integrate intra-sequence features as another input
$$
M = \left[ \left(M ^ {t y p} + M ^ {m s a}\right) \| M ^ {a n g} \right], \tag {6}
$$
Figure 3: Details on the architecture of cross-modal transformer that enables the interaction of different modalities and outputs updated representations. In particular, the pair representation is used to predict the inter-protein distance map.
where  $\parallel$  is the concatenation. Please note that the broadcast operation is used to make the dimensions consistent. Specifically,  $M^{typ}$  is broadcasted along the newly added first dimension during addition. Similarly, broadcasting is applied to the newly added first dimension of  $M^{ang}$  for concatenation. Through Evoformer, pair representation  $\hat{P}$  and evolution representation  $\hat{M}$  are obtained
$$
\hat {P}, \hat {M} = \operatorname {E V O F O R M E R} (P, M). \tag {7}
$$
Finally, we utilize pair representation  $\hat{P}$  to predict the inter-protein distance map. Here we define two loss functions to supervise the learning of distance map and evolution representation. Concerning the distance map, we can directly calculate the ground truth from rigid proteins. However, this ground truth is naturally noisy because of the experimental resolution. In light of successful practice in protein structure prediction [28, 23], we use a discretized distance map to replace the exact one. Specifically, distances are discretized into 64 bins ranging from 2 to  $22\AA$ . The prediction of discretized distances is converted into a classification problem by
$$
\bar {D} _ {i j} = \sigma \left(W \left(\hat {P} _ {i j} + \hat {P} _ {j i}\right)\right), \tag {8}
$$
where  $W$  is the learnable parameter and  $\sigma$  is the activation function. Then the cross-entropy loss averaged over all residue pairs is
$$
\mathcal {L} _ {\text {d i s t}} = - \frac {1}{m n} \sum_ {i, j} \sum_ {k = 1} ^ {6 4} \bar {G} _ {i j} ^ {k} \log \bar {D} _ {i j} ^ {k}, \tag {9}
$$
where  $\bar{G}_{ij}^{k}$  represents the  $k$ th-element of one-hot encoding of discretized actual distance. It is worth noting that the predicted distance map  $\hat{D}$  can be obtained by using the mean of each bin.
Inspired by the masked language model [20], we leverage the evolution representation to reconstruct masked MSA values. We consider 23 classes, including 20 common amino acid types, an unknown type, a gap token, and a mask token, and introduce the mask policy in Appendix B. Thus, a masked MSA loss can be defined as follows
$$
\bar {M} = \operatorname {S o f t m a x} (W \hat {M}), \tag {10}
$$
$$
\mathcal {L} _ {m s a} = - \frac {1}{N _ {m a s k}} \sum_ {i = 1} ^ {N _ {c l s}} \sum_ {j \in N _ {m a s k}} \sum_ {k = 1} ^ {2 3} A _ {i j} ^ {k} \log \left(\bar {M} _ {i j} ^ {k}\right), \tag {11}
$$
where  $N_{mask}$  denotes the number of masked tokens and  $A$  is the ground truth.
Similar to Equation (1), the SE(3) transformation optimized in the inner loop will exhibit differences from the ground truth complex, generating hypergradients through the cross-modal transformer
$$
\mathcal {L} ^ {i n} \left(R ^ {*} (\phi), \mathbf {t} ^ {*} (\phi)\right) = \frac {1}{m n} \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} \left(\| X _ {i} - R ^ {*} (\phi) Y _ {j} - \mathbf {t} ^ {*} (\phi) \| - D _ {i j}\right) ^ {2}, \tag {12}
$$
where  $D$  is the ground truth distance between amino acids. Overall, the outer loss for the cross-model transformer is
$$
\mathcal {L} ^ {\text {o u t}} = \lambda_ {1} \mathcal {L} _ {\text {d i s t}} + \lambda_ {2} \mathcal {L} _ {\text {m s a}} + \lambda_ {3} \mathcal {L} ^ {\text {i n}}, \tag {13}
$$
where hyperparameters  $\lambda_1, \lambda_2$  and  $\lambda_3$  balance the importance of different loss terms.
3.2 Unrolled Algorithm for Hypergradient
To solve the bi-level optimization formulated as Equation (2), hyper gradient  $\pmb{d}_{\phi}\mathcal{L}^{out}$  is required in the outer level and can be unrolled via the chain rule
$$
\boldsymbol {d} _ {\phi} \mathcal {L} ^ {o u t} = \frac {\partial \mathcal {L} ^ {o u t}}{\partial \left(R ^ {*}, \mathbf {t} ^ {*}\right)} \frac {\partial \left(R ^ {*} (\phi) , \mathbf {t} ^ {*} (\phi)\right)}{\partial \phi} + \frac {\partial \mathcal {L} ^ {o u t}}{\partial \phi}. \tag {14}
$$
One method to calculate it is approximating  $(R^{*}(\phi),\mathbf{t}^{*}(\phi))$  via optimizer
$$
\left(R _ {t}, \mathbf {t} _ {t}\right) = \operatorname {O P T} \left(R _ {t - 1}, \mathbf {t} _ {t - 1}, \phi\right), \quad t = 1, \dots , T, \tag {15}
$$
where  $T$  denotes the number of iterations. We explicitly calculate the gradients
$$
\tilde {D} _ {i j} = X _ {i} - R _ {t} Y _ {j} - \mathbf {t} _ {t},
$$
$$
\nabla_ {\mathbf {q}} \mathcal {L} _ {t} \left(R _ {t}, \mathbf {t} _ {t}, \phi\right) = \sum_ {i, j} \frac {2}{m n} \frac {\hat {D} _ {i j} - \| \tilde {D} _ {i j} \|}{\| \tilde {D} _ {i j} \|} \tilde {D} _ {i j} \odot Y _ {j} \cdot \frac {\partial R _ {t}}{\partial \mathbf {q} _ {t}}, \tag {16}
$$
$$
\nabla_ {\mathbf {t}} \mathcal {L} _ {t} (R _ {t}, \mathbf {t} _ {t}, \phi) = \sum_ {i, j} \frac {2}{m n} \frac {\hat {D} _ {i j} - \| \tilde {D} _ {i j} \|}{\| \tilde {D} _ {i j} \|} \tilde {D} _ {i j},
$$
where rotation matrix  $R$  and quaternion  $\mathbf{q}$  are converted to each other through the Bar-Itzhack algorithm [5], and  $\odot$  means Hadamard product. After  $T$  rounds of iterations,  $(R^{*}(\phi),\mathbf{t}^{*}(\phi))$  can be approximated as  $(R_T,t_T)$ . We can compute the hypergradient by substituting
$$
\frac {\partial \left(R ^ {*} (\phi) , \mathbf {t} ^ {*} (\phi)\right)}{\partial \phi} \approx - \gamma \frac {\partial^ {2} \mathcal {L} \left(R _ {T} , \mathbf {t} _ {T} , \phi\right)}{\partial \left(R _ {T} , \mathbf {t} _ {T}\right) ^ {\top} \partial \phi}. \tag {17}
$$
Due to the vast search space and the rugged landscape of the loss, the gradient descent algorithm with random initialization often requires a large number of iterations to converge and struggles to find the global minima. In the next subsection, we will derive a spectral initialization to enhance convergence.
3.3 Spectral Initialization
Recall that we intend to derive a good initialization  $(R_0, \mathbf{t}_0)$  for the gradient descent of Equation (1). To simplify it, we denote  $\hat{Y} = RY - \mathbf{t}\mathbf{1}_n^\top$ , where  $\mathbf{1}_n$  is the all-ones vector. Then substitute  $\hat{Y}$  into
$$
\left(\| X _ {i} - R Y _ {j} - \mathbf {t} \| - \hat {D} _ {i j}\right) ^ {2} = X _ {i} ^ {\top} X _ {i} - 2 X _ {i} ^ {\top} \hat {Y} _ {j} + \hat {Y} _ {j} ^ {\top} \hat {Y} _ {j} + \hat {D} _ {i j} ^ {2} - 2 \hat {D} _ {i j} \| X _ {i} - \hat {Y} _ {j} \|. \tag {18}
$$
According to these variable combinations, we define four variables and two centering matrices
$$
B = \left[ \begin{array}{c c c} X _ {1} ^ {\top} X _ {1} & \ldots & X _ {1} ^ {\top} X _ {1} \\ X _ {2} ^ {\top} X _ {2} & \ldots & X _ {2} ^ {\top} X _ {2} \\ \ldots & \ldots & \ldots \\ X _ {m} ^ {\top} X _ {m} & \ldots & X _ {m} ^ {\top} X _ {m} \end{array} \right] _ {m \times n} C = \left[ \begin{array}{c c c} \hat {Y} _ {1} ^ {\top} \hat {Y} _ {1} & \ldots & \hat {Y} _ {n} ^ {\top} \hat {Y} _ {n} \\ \hat {Y} _ {1} ^ {\top} \hat {Y} _ {1} & \ldots & \hat {Y} _ {n} ^ {\top} \hat {Y} _ {n} \\ \ldots & \ldots & \ldots \\ \hat {Y} _ {1} ^ {\top} \hat {Y} _ {1} & \ldots & \hat {Y} _ {n} ^ {\top} \hat {Y} _ {n} \end{array} \right] _ {m \times n}
$$
$$
E = X ^ {\top} \hat {Y} = \left[ \begin{array}{l l l} X _ {1} ^ {\top} \hat {Y} _ {1} & \dots & X _ {1} ^ {\top} \hat {Y} _ {n} \\ X _ {2} ^ {\top} \hat {Y} _ {1} & \dots & X _ {2} ^ {\top} \hat {Y} _ {n} \\ \dots & \dots & \dots \\ X _ {m} ^ {\top} \hat {Y} _ {1} & \dots & X _ {m} ^ {\top} \hat {Y} _ {n} \end{array} \right] _ {m \times n} \quad F = \left[ \begin{array}{l l l} \hat {D} _ {1 1} ^ {2} & \dots & \hat {D} _ {1 n} ^ {2} \\ \hat {D} _ {2 1} ^ {2} & \dots & \hat {D} _ {2 n} ^ {2} \\ \dots & \dots & \dots \\ \hat {D} _ {m 1} ^ {2} & \dots & \hat {D} _ {m n} ^ {2} \end{array} \right] _ {m \times n} \tag {19}
$$
$$
H _ {m} = I _ {m} - \frac {1}{m} J _ {m} H _ {n} = I _ {n} - \frac {1}{n} J _ {n},
$$