SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wengong Jin, Siranush Sarzikova, Xun Chen, Nir Hacohen, Caroline Uhler
{wjin,sarkizov,xun,nhacohen,cuhler}@broadinstitute.org
Protein-ligand binding prediction is a fundamental problem in AI-driven drug discovery. Previous work focused on supervised learning methods for small molecules where binding affinity data is abundant, but it is hard to apply the same strategy to other ligand classes like antibodies where labelled data is limited. In this paper, we explore unsupervised approaches and reformulate binding energy prediction as a generative modeling task. Specifically, we train an energy-based model on a set of unlabelled protein-ligand complexes using SE(3) denoising score matching (DSM) and interpret its log-likelihood as binding affinity. Our key contribution is a new equivariant rotation prediction network for SE(3) DSM called Neural Euler's Rotation Equations (NERE). It predicts a rotation by modeling the force and torque between protein and ligand atoms, where the force is defined as the gradient of an energy function with respect to atom coordinates. Using two protein-ligand and antibody-antigen binding affinity prediction benchmarks, we show that NERE outperforms all unsupervised baselines (physics-based potentials and protein language models) in both cases and surpasses supervised baselines in the antibody case.
One of the challenges in drug discovery is to design ligands (small molecules or antibodies) with high binding affinity to a target protein. In recent years, many protein-ligand binding prediction models have been developed for small molecules [13, 19, 43]. These models are typically trained on a large set of crystal structures labeled with experimental binding data. However, it is difficult to apply this supervised learning approach to other ligand classes like antibodies where binding affinity data is limited. For example, the largest binding affinity dataset for antibodies [31] has only 566 data points. Transfer learning from small molecules to antibodies is also hard as their structures are very different.
Given this data scarcity challenge, we aim to develop a general unsupervised binding energy prediction framework for small molecules and antibodies. The basic idea is to learn an energy-based model (EBM) [16] of protein-ligand complexes by maximizing the log-likelihood of crystal structures in the Protein Data Bank (PDB). This generative modeling approach is motivated by recent advances of protein language models (PLMs) which show that the likelihood of protein sequences is correlated with protein mutation effects [20, 11]. However, PLMs are not applicable to small molecules and they model the likelihood of protein sequences rather than structures. We suppose the likelihood of protein complex structures is a better predictor of binding affinity, which is determined by the relative orientation between a protein and a ligand. Indeed, our method shares the same spirit as traditional physics-based energy functions [1] designed by experts so that crystal structures have low binding energy (i.e., high likelihood). Our key departure is to use neural networks to learn a more expressive energy function in a data-driven manner. Importantly, our model is applicable to large-scale virtual screening because it does not require expensive molecular dynamic simulations [21].
Specifically, we train EBMs using SE(3) denoising score matching (DSM) [33, 15] as maximum likelihood estimation is hard. For each training example, we create a perturbed protein-ligand complex by randomly rotating/translating a ligand and ask our model to reconstruct the rotation/translation noise. The main challenge of SE(3) DSM in our context is how to predict rotation noise from the score of an EBM in an equivariant manner. Motivated by physics, we develop a new equivariant rotation prediction network called Neural Euler's Rotation Equation (NERE). The key observation is that the score of an EBM equals the force of each atom, which further defines the torque between a protein and a ligand. Using Euler's rotation equation, we convert this torque into the angular velocity of a ligand and its corresponding rotation matrix to calculate SE(3) DSM loss. Importantly, NERE is compatible with any SE(3)-invariant networks and guarantees equivariance without any further requirement on network architectures.
We evaluate NERE on two protein-ligand and antibody-antigen binding affinity benchmarks from PDBBind [36] and Structural Antibody Database (SAbDab) [31]. We compare NERE with unsupervised physics-based models like MM/GBSA [21], protein language models (ESM-1v [20] and ESM-IF [11]), and a variety of supervised models regressed on experimental binding affinity data. To simulate real-world virtual screening scenarios, we consider two settings where input complexes are crystallized or predicted by a docking software. NERE outperforms all unsupervised baselines across all settings and surpasses supervised models in the antibody setting, which highlights the benefit of unsupervised learning when binding affinity data is limited.
Protein-ligand binding models fall into two categories: supervised and unsupervised learning. Supervised models are trained on binding affinity data from PDBBind [36, 29, 3, 32, 40, 19, 43]. They typically represent protein-ligand 3D complex represented as a geometric graph and encode it into a vector representation using a neural network for affinity prediction. Unsupervised models are either based on physics or statistical potentials. For example, molecular mechanics generalized Born surface area (MM/GBSA) [21] calculates binding energy based on expensive molecular dynamics. Our work is closely related to statistical potentials like DrugScore $_{2018}$  [6]. It learns an energy score for each atom pair independently based on their atom types and distance alone, which ignores the overall molecular context. Our key departure is to use neural networks to learn a more expressive energy function powered by context-aware atom representations.
Antibody-antigen binding models are traditionally based on scoring functions in protein docking programs [25, 1, 10, 39, 2]. They are typically weighted combinations of physical interaction terms whose weights are tuned so that crystal structures have lower energy than docked poses. Recently, protein language models like ESM-1v [20] and ESM-IF [11] have been successful in predicting mutation effects for general proteins and we evaluate their performance for antibodies in this paper. Due to scarcity of binding affinity data, very few work has explored supervised learning for antibody binding affinity prediction. Existing models [23] are based on simple hand-crafted features using binding affinity data from SAbDab. For fair comparison, we implement our own supervised baselines using more advanced neural architectures [38, 28].
Denoising score matching (DSM) is a powerful technique for training score-based generative models [33] and has been applied to protein structure prediction [41] and protein-ligand docking [5]. Our departure from standard DSM is two fold. First, standard DSM is based on Gaussian noise while our method uses rigid transformation noise based on an isotropic Gaussian distribution for SO(3) rotation group [15]. Second, we parameterize the score as the gradient of an energy-based model rather than the direct output of a neural network. We choose this parameterization scheme because we need to predict energy and score at the same time.
Equivariant rotation prediction. Motivated by molecular docking, there has been a growing interest in predicting rotations with equivariant neural networks. For instance, Ganea et al. [9], Stärk et al. [35] developed a graph matching network that first predicts a set of key points between a protein and a ligand, and then uses the Kabsch algorithm [14] to construct a rigid transformation. Corso et al. [5] predicted rotations based on tensor field networks [37] and spherical harmonics. Our method (NERE) has two key differences with prior work. First, it does not require key point alignment and directly outputs an rotation matrix. Second, NERE does not require spherical harmonics and can be plugged into any SE(3)-invariant encoder architectures for equivariant rotation prediction.
Figure 1: The architecture of  $E(A, X)$  for small molecules (left) and antibodies (right). For small molecules, the input is a binding pocket and a molecular graph encoded by a message passing network (MPN). For antibodies, the input contains antibody CDR residues and epitope residues, where residue features  $A$  come from an ESM2 protein language model.
3 Unsupervised Binding Energy Prediction
A protein-ligand (or an antibody-antigen) complex is a geometric object that involves a protein and a ligand (a small molecule or an antibody). It is denoted as a tuple  $(\mathbf{A},\mathbf{X})$ , with atom features  $\mathbf{A} = [a_{1},\dots ,a_{n}]$  and atom coordinates  $\mathbf{X} = [\pmb{x}_1,\dots ,\pmb{x}_n]$  (column-wise concatenation). The key property of a complex is its binding energy  $E(\mathbf{A},\mathbf{X})$ . A lower binding energy means a ligand binds more strongly to a protein. In this section, we describe how to parameterize  $E(\mathbf{A},\mathbf{X})$  and design proper training objectives to infer the true binding energy function from a list of crystal structures without binding affinity labels.
3.1 EBM Architecture
An energy function  $E(\mathbf{A}, \mathbf{X})$  is composed of a protein encoder and an output layer. The encoder is a frame averaging neural network [28] that learns a SE(3)-invariant representation  $\mathbf{h}_i$  for each atom. We choose this architecture because of its simplicity (detailed in the appendix). The output layer  $\phi_o$  is a feed-forward neural network with one hidden layer. It predicts the interaction energy  $\phi_o(\mathbf{h}_i, \mathbf{h}_j)$  for each pair of atoms. Finally, we define  $E(\mathbf{A}, \mathbf{X})$  as the sum of pairwise interaction energies:
$$
E (\boldsymbol {A}, \boldsymbol {X}) = \sum_ {i, j: d _ {i j} <   d} \phi_ {o} \left(\boldsymbol {h} _ {i}, \boldsymbol {h} _ {j}\right) \tag {1}
$$
Since atomic interaction vanishes beyond certain distance, we only consider atom pairs with distance  $d_{ij} < d$ . So far, we have described our method in generic terms. We now specify the input features and preprocessing steps tailored to small molecules and antibodies.
Small molecules. When the ligand is a small molecule, the input to our model is a protein-ligand complex where the protein is cropped to its binding pocket (residues within  $10\AA$  from the ligand). The model architecture is illustrated in Figure 1a. On the protein side, each atom in the binding pocket is represented by a one-hot encoding of its atom name  $(C_{\alpha}, C_{\beta}, N, O, \text{etc.})$ . We consider all backbone and side-chain atoms. On the ligand side, the atom features  $\mathbf{A}$  are learned by a message passing network (MPN) [42] based on a ligand molecular graph. The MPN and energy function are optimized jointly during training.
Antibodies. When the ligand is an antibody, the input to our model is an antibody-antigen binding interface. The binding interface is composed of residues in the antibody complementarity-determining region (CDR) of both heavy and light chains and top 50 antigen residues (epitope) closest to the CDR. The model architecture is depicted in Figure 1b. Different from the small molecule case, we only consider  $C_{\alpha}$  atoms of each antibody/antigen residue. Each residue is represented by a 2560 dimensional pre-trained ESM-2 embedding [18]. For computational efficiency, we freeze ESM-2 weights during training.
3.2 Training EBMs with Denoising Score Matching
Given that our training set does not have binding affinity labels, we need to design an unsupervised training objective different from a supervised regression loss. Our key hypothesis is that we can infer the true binding energy function (up to affine equivalence) by maximizing the likelihood of crystal
Algorithm 1 Training procedure (single data point)
Require: A training complex  $(A,X)$
1: Sample a noise level  $\sigma$
2: Sample rotation vector  $\omega \sim \mathcal{N}_{SO(3)}$  with variance  $\sigma^2$
3: Sample translation vector  $\pmb{t} \sim \mathcal{N}(0, \sigma^2 \pmb{I})$ .
4: Perturb the coordinates  $\bar{X}$  by applying rigid transformation  $(\omega, t)$  to the original complex.
5: Compute the score of energy function  $(\tilde{\omega},\tilde{t})$  based on its gradient (force)  $-\partial E / \partial \tilde{X}$ .
6: Minimize DSM objective  $\ell_{\mathrm{dsm}}$
structures in our training set. The motivation of our hypothesis is that a crystal structure is the lowest energy state of a protein-ligand complex. The maximum likelihood objective seeks to minimize the energy of crystal structures since the likelihood of a complex is  $p(\mathbf{A}, \mathbf{X}) \propto \exp(-E(\mathbf{A}, \mathbf{X}))$ .
While maximum likelihood estimation (MLE) is difficult for EBMs due to marginalization, recent works [33, 34] has successfully trained EBMs using denoising score matching (DSM) and proved that DSM is a good approximation of MLE. In standard DSM, we create a perturbed complex by adding Gaussian noise to ligand atom coordinates, i.e.,  $\tilde{\pmb{X}} = \pmb {X} + \pmb {\epsilon}$ , where  $\epsilon \sim p(\epsilon) = \mathcal{N}(0,\sigma^2\pmb {I})$ . DSM objective tries to match the score of our model  $-\partial E / \partial \tilde{X}$  and the score of the noise distribution  $\nabla_{\epsilon}\log p(\epsilon) = -\epsilon /\sigma^{2}$ :
$$
\ell_ {\mathrm {g}} = \mathbb {E} \left[ \| \partial E (\boldsymbol {A}, \tilde {\boldsymbol {X}}) / \partial \tilde {\boldsymbol {X}} - \epsilon / \sigma^ {2} \| ^ {2} \right] \tag {2}
$$
Intuitively,  $\ell_{\mathrm{g}}$  forces the gradient to be zero when the input complex is a crystal structure  $(\epsilon = 0)$ . As a result, a crystal structure pose will be at the local minima of an EBM under the DSM objective.
3.3 Perturbing a Complex with Rigid Transformation Noise
Nevertheless, adding Gaussian noise is not ideal for protein-ligand binding because it may create nonsensical conformations that violate physical constraints (e.g., an aromatic ring must be planar). A better solution is to create a perturbed complex  $(\mathbf{A},\tilde{\mathbf{X}})$  via random ligand rotation and translation, similar to molecular docking. To construct a random rotation, we sample a rotation vector  $\omega$  from  $\mathcal{N}_{SO(3)}$ , an isotropic Gaussian distribution over  $SO(3)$  rotation group [15] with variance  $\sigma^2$ . Each  $\omega \sim \mathcal{N}_{SO(3)}$  has the form  $\omega = \theta \hat{\omega}$ , where  $\hat{\omega}$  is a vector sampled uniformly from a unit sphere and  $\theta \in [0,\pi]$  is a rotation angle with density
$$
f (\theta) = \frac {1 - \cos \theta}{\pi} \sum_ {l = 0} ^ {\infty} (2 l + 1) e ^ {- l (l + 1) \sigma^ {2}} \frac {\sin ((l + 1 / 2) \theta)}{\sin (\theta / 2)} \tag {3}
$$
Likewise, we sample a random translation vector  $\pmb{t}$  from a normal distribution  $t\sim \mathcal{N}(0,\sigma^2 I)$ . Finally, we apply this rigid transformation to the ligand and compute its perturbed coordinates  $\tilde{\pmb{X}} = \pmb{R}_{\omega}\pmb {X} + t$  where  $\pmb{R}_{\omega}$  is the rotation matrix given by the rotation vector  $\omega = (\omega_{x},\omega_{y},\omega_{z})$ .
$$
\boldsymbol {R} _ {\omega} = \exp (\boldsymbol {W} _ {\omega}), \boldsymbol {W} _ {\omega} = \left( \begin{array}{c c c} 0 & - \omega_ {z} & \omega_ {y} \\ \omega_ {z} & 0 & - \omega_ {x} \\ - \omega_ {y} & \omega_ {x} & 0 \end{array} \right). \tag {4}
$$
Here exp means matrix exponentiation and  $W_{\omega}$  is an infinitesimal rotation matrix. Since  $W_{\omega}$  is a skew symmetric matrix, its matrix exponential has the following closed form
$$
\boldsymbol {R} _ {\omega} = \exp (\boldsymbol {W} _ {\omega}) = \boldsymbol {I} + c _ {1} \boldsymbol {W} _ {\omega} + c _ {2} \boldsymbol {W} _ {\omega} ^ {2} \tag {5}
$$
$$
c _ {1} = \frac {\sin \| \boldsymbol {\omega} \|}{\| \boldsymbol {\omega} \|}, \quad c _ {2} = \frac {1 - \cos \| \boldsymbol {\omega} \|}{\| \boldsymbol {\omega} \| ^ {2}} \tag {6}
$$
Moreover, we do not need to explicitly compute the matrix exponential  $R_{\omega}$  since  $W_{\omega}$  is the linear mapping of cross product, i.e.  $\boldsymbol{\omega} \times \boldsymbol{r} = \boldsymbol{W}_{\omega}\boldsymbol{r}$ . Therefore, applying a rotation matrix only involves cross product operations that are very efficient:
$$
\boldsymbol {R} _ {\omega} \boldsymbol {x} _ {i} = \boldsymbol {x} _ {i} + c _ {1} \boldsymbol {\omega} \times \boldsymbol {x} _ {i} + c _ {2} \boldsymbol {\omega} \times (\boldsymbol {\omega} \times \boldsymbol {x} _ {i}) \tag {7}
$$
Figure 2: NERE predicts a rotation in three steps. It first calculates pairwise residue force based on the gradient  $\pmb{f}_i = (-\partial E / \partial \pmb{x})^\top$ . It then predicts the corresponding torque and angular velocity by solving Euler's rotation equation. Lastly, it converts the angular velocity vector to a rotation via matrix exponential map.
3.4 Training EBMs with SE(3) Denoising Score Matching
Under this perturbation scheme, DSM aims to match the model score  $-\partial E / \partial \hat{X}$  and the score of rotation and translation noise  $\nabla_{\omega} \log p(\omega)$ ,  $\nabla_t \log p(t)$ . Our SE(3) DSM objective is a sum of two losses:  $\ell_{\mathrm{dsm}} = \ell_t + \ell_r$ , where  $\ell_t$  and  $\ell_r$  correspond to a translation and a rotation DSM loss. The translation part is straightforward since  $t$  follows a normal distribution and  $\nabla_t \log p(t) = -t / \sigma^2$ :
$$
\ell_ {t} = \mathbb {E} \left[ \| \tilde {\boldsymbol {t}} - \nabla_ {\boldsymbol {t}} \log p (\boldsymbol {t}) \| ^ {2} \right], \quad \tilde {\boldsymbol {t}} = - \frac {1}{n} \sum_ {i} \partial E / \partial \tilde {\boldsymbol {x}} _ {i} \tag {8}
$$
The rotation part of DSM is more complicated. As  $\hat{\omega}$  is sampled from a uniform distribution over a sphere (whose density is constant), the density and score of  $p(\omega)$  is
$$
p (\omega) \propto f (\theta), \quad \nabla_ {\omega} \log p (\omega) = \nabla_ {\theta} \log f (\theta) \cdot \hat {\omega} \tag {9}
$$
In practice, we calculate the density and score by precomputing truncated infinite series in  $f(\theta)$ . However, the main challenge is that the model score  $-\partial E / \partial \tilde{\pmb{X}} \in \mathbb{R}^{n \times 3}$  is defined over atom coordinates, which is not directly comparable with  $\nabla_{\omega} \log p(\omega) \in \mathbb{R}^3$  as they have different dimensions. To address this issue, we need to map  $-\partial E / \partial \tilde{\pmb{X}}$  to a rotation vector  $\tilde{\omega}$  via a rotation predictor  $F$  and perform score matching in the rotation space:
$$
\ell_ {r} = \mathbb {E} \left[ \| \tilde {\boldsymbol {\omega}} - \nabla_ {\boldsymbol {\omega}} \log p (\boldsymbol {\omega}) \| ^ {2} \right], \quad \tilde {\boldsymbol {\omega}} = F (- \partial E / \partial \tilde {\boldsymbol {X}}) \tag {10}
$$
The rotation predictor  $F$  is based on Neural Euler's Rotation Equation, which will be described in the next section. The overall training procedure is summarized in Algorithm 1 and Figure 2.
4 Neural Euler's Rotation Equation
In this section, we present Neural Euler's Rotation Equation (NERE), which can be used in SE(3) DSM to infer rotational scores from the gradient  $-\partial E / \partial \tilde{X}$ . To motivate our method, let us first review some basic concepts related to Euler's Rotation Equation.
4.1 Euler's Rotation Equations
In classical mechanics, Euler's rotation equation is a first-order ordinary differential equation that describes the rotation of a rigid body. Suppose a ligand rotates around its center mass  $\mu$  with angular velocity  $\omega$ . Euler's rotation equation in an inertial reference frame is defined as
$$
\boldsymbol {I} _ {N} \frac {d \boldsymbol {\omega}}{d t} = \boldsymbol {\tau}, \quad \boldsymbol {\tau} = \sum_ {i} \left(\boldsymbol {x} _ {i} - \boldsymbol {\mu}\right) \times \boldsymbol {f} _ {i} \tag {11}
$$
where  $\pmb{I}_N \in \mathbb{R}^{3 \times 3}$  is the inertia matrix of a ligand,  $\pmb{\tau}$  is the torque it received, and  $\pmb{f}_i$  is the force applied to a ligand atom  $i$ . The inertia matrix describes the mass distribution of a ligand and the
torque needed for a desired angular acceleration, which is defined as
$$
\boldsymbol {I} _ {N} = \sum_ {i} \| \boldsymbol {x} _ {i} - \boldsymbol {\mu} \| ^ {2} \boldsymbol {I} - (\boldsymbol {x} _ {i} - \boldsymbol {\mu}) (\boldsymbol {x} _ {i} - \boldsymbol {\mu}) ^ {\top} \tag {12}
$$
For a short period of time  $\Delta t$ , we can approximate the new angular velocity  $\omega_{t = \Delta t}$  by
$$
\frac {d \boldsymbol {\omega}}{d t} \approx \frac {\boldsymbol {\omega} _ {t = \Delta t} - \boldsymbol {\omega} _ {t = 0}}{\Delta t} = \boldsymbol {I} _ {N} ^ {- 1} \tau \tag {13}
$$
Since we assume the system is in an inertial reference frame  $(\omega_{t=0} = 0)$ , we have  $\omega_{t=\Delta t} = I_N^{-1}\tau\Delta t$  (we set  $\Delta t = 0.1$ ). We note that calculating the inverse  $I_N^{-1}$  is cheap because it is a  $3\times 3$  matrix.
4.2 Applying NERE to SE(3) DSM
In SE(3) DSM, our goal is to map the gradient  $-\partial E / \partial x_{i}$  to a rotation vector  $\omega$ . In physics, the gradient  $-\partial E / \partial x_{i}$  is the force  $\pmb{f}_{i}$  of atom  $i$ . Therefore, the rotation predictor  $F$  in Eq.(10) is a simple application of the Euler's Rotation Equation
$$
\boldsymbol {\omega} = F (- \partial E / \partial \boldsymbol {X}) = \boldsymbol {I} _ {N} ^ {- 1} \boldsymbol {\tau} \Delta t, \quad \boldsymbol {\tau} = \sum_ {i} \left(\boldsymbol {x} _ {i} - \boldsymbol {\mu}\right) \times (- \partial E / \partial \boldsymbol {x} _ {i}) \tag {14}
$$
A nice property of NERE is that it is equivariant under SO(3) rotation group because it is derived from physics. We formally state this proposition as follows and its proof in the appendix.
Proposition 1. Suppose we rotate a ligand so that its new coordinates become  $\pmb{x}_i' = \pmb{R}\pmb{x}_i$ . The new force  $\pmb{f}'$ , torque  $\pmb{\tau}'$ , inertia matrix  $\pmb{I}_N'$ , and angular velocity  $\omega'$  for the rotated complex are
$$
\boldsymbol {f} _ {i} ^ {\prime} = \boldsymbol {R} \boldsymbol {f} _ {i}, \tau^ {\prime} = \boldsymbol {R} \tau , \boldsymbol {I} _ {N} ^ {\prime} = \boldsymbol {R} \boldsymbol {I} _ {N} \boldsymbol {R} ^ {\top}, \omega^ {\prime} = \boldsymbol {R} \omega
$$
In other words, NERE is equivariant under SO(3) rotation group.
Once we establish SO(3) equivariance, it is easy to satisfy SE(3) equivariance by first placing the rotation center at the origin  $(\pmb{x}_i\gets \pmb{x}_i - \pmb{\mu})$ , applying the predicted rotation via NERE, and then adding  $\pmb{\mu}$  back to each atom.
We evaluate our model on two drug discovery applications: protein-ligand binding and antibody-antigen binding affinity prediction. The experimental setup is described as follows.
5.1 Protein-Ligand Binding
Data. Our training data has 5237 protein-ligand complexes from the refined subset of PDBbind v2020 database [36]. Our test set has 285 complexes from the PDBbind core set with binding affinity labels converted into log scale. Our validation set has 357 complexes randomly sampled from PDBbind by Stärk et al. [35] after excluding all test cases. The final training set has 4806 complexes (without binding affinity labels) after removing all ligands overlapping with the test set.
Metric. We report the Pearson correlation coefficient between true binding affinity and predicted energy  $E(\mathbf{A},\mathbf{X})$ . We do not report root mean square error (RMSE) because our model does not predict absolute affinity values. In fact, shifting  $E(\mathbf{A},\mathbf{X})$  by any constant will be equally optimal under the DSM objective. We run our model with five different random seeds and report their average.
Baselines. We consider three sets of baselines for comparison: