SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shiwei Liu1,2*
Tian Zhu $^{1,2*}$
Milong Ren $^{1,2}$
Chungong  $\mathbf{Y}\mathbf{u}^{1,2}$
Dongbo Bu $^{1,2,3}$
Haicang Zhang $^{1,2\dagger}$
1. Institute of Computing Technology, Chinese Academy of Sciences, Beijing 100080, China.
2. University of Chinese Academy of Sciences, Beijing 100080, China.
3. Central China Research Institute for Artificial Intelligence Technologies, Zhengzhou 450046, China
Many crucial biological processes rely on networks of protein-protein interactions. Predicting the effect of amino acid mutations on protein-protein binding is vital in protein engineering and therapeutic discovery. However, the scarcity of annotated experimental data on binding energy poses a significant challenge for developing computational approaches, particularly deep learning-based methods. In this work, we propose SidechainDiff, a representation learning-based approach that leverages unlabelled experimental protein structures. SidechainDiff utilizes a Riemannian diffusion model to learn the generative process of side-chain conformations and can also give the structural context representations of mutations on the protein-protein interface. Leveraging the learned representations, we achieve state-of-the-art performance in predicting the mutational effects on protein-protein binding. Furthermore, SidechainDiff is the first diffusion-based generative model for sidechains, distinguishing it from prior efforts that have predominantly focused on generating protein backbone structures.
Many crucial biological processes, including membrane transport and cell signaling, are likely regulated by intricate networks of protein-protein interactions rather than single proteins acting independently [Calakos et al., 1994; Jones and Thornton, 1996; Levskaya et al., 2009; Schwikowski et al., 2000]. One representative example is the interaction between the spike protein of the SARS-CoV-2 virus and its receptor protein ACE2 on the surface of human cells, which is crucial for the virus to invade target cells [Lan et al., 2020]. Meanwhile, specific human antibodies, an essential type of protein in the immune system, can prevent the virus entry by competitively binding to the spike protein [Shan et al., 2022].
In protein engineering and therapeutic discovery, it is a crucial step to induce amino acid mutations on the protein-protein interface to modulate binding affinity [T O'Neil and Hoess, 1995]. For example, to enhance the efficacy of an antibody against a virus, one practical approach is to introduce amino acid mutations and filter the resulting antibody sequences to increase binding affinity and specificity to the target viral protein [Mason et al., 2021]. However, the variant space that can be explored
$^{\dagger}$ Correspondence should be addressed to H.Z. (zhanghaicang@ict.ac.cn)
using experimental assays is very limited, and developing an effective high-throughput screening can require a significant experimental effort.
Computational methods have been developed to predict the mutational effect on binding affinity measured by the change in binding free energy (i.e.,  $\Delta \Delta G$ ). Traditional computational approaches mostly used physical energy features, such as van der Waals and solvation energy, in combination with statistical models to predict  $\Delta \Delta G$  [Schymkowitz et al., 2005, Meireles et al., 2010, Alford et al., 2017, Barlow et al., 2018]. The limited model capacity and bias inherent in human-engineered features in these methods hinder their ability to characterize the complex mutational landscape of binding energy. Despite recent breakthroughs in protein modeling with deep learning [Jumper et al., 2021, Baek et al., 2021, Watson et al., 2023], developing deep learning-based models to predict mutational effects on protein-protein binding remains challenging due to the scarcity of labeled experimental data [Jankauskaite et al., 2019].
Recent studies have investigated various self-supervised learning strategies on protein structures and sequences [Liu et al., 2021, Meier et al., 2021, Hsu et al., 2022, Luo et al., 2023] to ease the data scarcity issue. Among them, GeoPPI [Liu et al., 2021] and RDE [Luo et al., 2023] have focused on optimizing protein side-chain conformations during pre-training, as the side-chain conformations on the protein-protein interface play a critical role in determining binding energy. In protein-protein interactions, the side-chain conformations of amino acids at the interface may exhibit significant variations when comparing the same receptor with different ligands. These side-chain conformations can be more effectively characterized using probability density. Notably, RDE explores a flow-based model to estimate the uncertainty in side-chain conformations and leverages the learned representations to achieve state-of-the-art performance in predicting  $\Delta \Delta G$ . However, flow-based models possess inherent limitations as they require specialized architectures to construct accurate bijective transformations in probability density [Papamakarios et al., 2021], which results in increased costs associated with model design and implementation.
To address the above limitations, we propose SidechainDiff, a representation learning framework for the protein-protein interfaces. It employs a Riemannian diffusion model to learn the generative process of side-chain conformations and the representation of the structural context of amino acids. To the best of our knowledge, SidechainDiff is the first diffusion probabilistic model for side-chain modeling, whereas previous methods have only focused on protein backbone structures. Furthermore, we leverage the learned representations and neural networks to predict  $\Delta \Delta G$ .
2.1 Protein side-chain conformation prediction
Accurate side-chain modeling is essential in understanding the biological functions of proteins. There are two primary categories in predicting protein side-chain conformations: end-to-end full-atom structure prediction methods and side-chain packing. AlphaFold and RoseTTAFold [Jumper et al., 2021, Baek et al., 2021] are two representative methods that simultaneously generate side-chain conformations and backbone structures. In scenarios like structure refinement and protein design, side-packing becomes pivotal. The objective is to determine the conformations of protein side-chains while having their backbone structures already defined. Traditional methods, including Rosetta [Leman et al., 2020] and SCWRL4 [Krivov et al., 2009] operate by minimizing the energy function across a pre-defined rotamer library. Recent methods for side-chain packing often employ deep learning models, such as 3D convolution networks and graph attention networks [McPartlon and Xu, 2023, Misiura et al., 2022, Xu et al., 2022].
Our model distinguishes itself from previous methods for side-chain modelling in two key aspects. First, it's capable of generating a distribution of side-chain conformations rather than a single conformation. Second, we emphasize side-chain modeling specifically for mutation sites within the protein-protein interface, leveraging the known structural context of these sites.
2.2 Prediction of mutational effects on protein-protein binding
Methods for predicting  $\Delta \Delta G$  can be broadly classified into three categories: biophysical methods, statistical methods, and deep learning-based approaches. Biophysical methods offer a robust means of elucidating the molecular mechanisms governing protein-protein interactions and the impact of
mutations on these interactions [Schymkowitz et al., 2005, Alford et al., 2017, Steinbrecher et al., 2017]. These methods directly integrate information on protein structures and key biophysical properties, such as solvent accessibility, electrostatic potential, and hydrogen bonding patterns. Statistical methods tailor statistical models for the protein properties such as evolutionary conservation and geometric characteristics [Li et al., 2016, Geng et al., 2019, Zhang et al., 2020].
Deep learning-based methods can be categorized into sequence-based methods and structure-based methods. Sequence-based methods primarily either focus on the evolutionary history, multiple sequence alignments (MSAs) in most cases, of the target protein [Hopf et al., 2017], Riesselman et al., 2018, Frazer et al., 2021] or act as protein language models (PLMs) trained on large amounts of protein sequences [Meier et al., 2021, Notin et al., 2022]. Structure-based methods can be categorized into end-to-end methods and pre-training-based methods. The end-to-end methods extract features from protein complexes and directly train a neural network model on them [Shan et al., 2022]. To mitigate overfitting caused by data sparsity, an alternative approach is to learn representations by pre-training a feature extraction network on unlabeled structures [Liu et al., 2021, Luo et al., 2023]. Among them, RDE-Network [Luo et al., 2023] utilizes normalizing flows in torus space to estimate the density of amino acid side-chain conformations and leverages the learned representation to predict  $\Delta \Delta G$ .
2.3 Diffusion probabilistic models
Diffusion Probabilistic Models (DPMs) are generative models to transform a sample from a tractable noise distribution towards a desired data distribution with a gradual denoising process [Sohl-Dickstein et al., 2015, Kingma et al., 2021, Dhariwal and Nichol, 2021]. DPMs have achieved impressive results in generating various data types, including images [Ho et al., 2020, Nichol and Dhariwal, 2021], waveforms, and discrete data like text [Hoogeboom et al., 2021]. DPMs-based autoencoders have also been proposed to facilitate representation learning for image data [Preechakul et al., 2022, Zhang et al., 2022].
Inspired by these progresses, DPMs have also been explored in protein modeling, including de novo protein design [Watson et al., 2023, Ingraham et al., Yim et al., 2023], motif-scaffolding [Trippe et al., 2023], and molecular dynamics [Arts et al., 2023, Wu and Li, 2023]. While the existing methods utilize DPMs for generating protein backbones, a research gap remains in modeling side-chain conformations, which play a critical role in protein-protein binding. The studies by Jing et al. [2022] and Corso et al. [2023] employ diffusion models for generating torsion angles in the context of small molecular design. In contrast, our research is centered on modeling the torsional angles of protein side-chains. Furthermore, our approach distinguishes itself from their models in terms of how we construct the diffusion process.
Our methods have two main components. We first propose SidechainDiff (Figure 1), which is a diffusion probabilistic model for protein side-chains and can also give the structural context representations of mutations on the protein-protein interface. We then propose DiffAffinity that utilize the learned representations to predict  $\Delta \Delta G$ . We organize this section as follows: Section 3.1 presents the preliminaries and notations that are used throughout the paper and formally defines the problem. Section 3.2 and 3.3 describe SidechainDiff and DiffAffinity, respectively. Section 3.4 describes key details in model training.
3.1 Preliminaries and notations
A single-chain protein structure is composed of multiple residues (amino acids). Each amino acid shares a common central carbon atom, called the alpha  $(\alpha)$  carbon. The alpha carbon is bonded to an amino group (-NH2), a carboxyl group (-COOH), and a hydrogen atom, forming the amino acid's backbone. Each amino acid has a distinct atom or group bonded to the alpha carbon, known as the side-chain. A residue's side-chain structure is usually termed a side-chain conformation. The side-chain conformation varies across different amino acids, and it depends on the type of amino acid and the overall protein structure. Our main work focuses on building a generative process of side-chain conformations given the protein backbone structures.
Figure 1: The overall architecture of SidechainDiff.
Multiple protein chains bind with each other and form a protein complex. Here we only consider protein complexes with two chains just for notation simplicity and note that our model works for protein complexes with more chains. For a protein complex with  $n$  residues, we denote the two chains as chain  $A$  and  $B$ , and the set of its residues as  $\{1, \dots, n\}$ . The properties of the  $i$ -th amino acid include the amino acid type  $a_i \in \{1, \dots, 20\}$ , the orientation  $\mathbf{R}_i \in SO(3)$ , the 3D coordinates  $\mathbf{x}_i \in \mathbb{R}^3$ , and the side-chain torsion angles  $\chi_i = (\chi_i^{(k)})_{k=1}^4$  where  $\chi_i^{(k)}$  belongs to  $[0, 2\pi)$ .
The bond lengths and bond angles in side-chains are usually considered fixed and given backbone structures, we can solely use  $\chi_{i}$  to reconstruct the coordinates of each atom in a side-chain. Thus, a side-chain conformation is also called a rotamer, which is parameterized with  $\chi_{i}$ . The number of torsional angles ranges from 0 to 4, depending on the residue type. Here, we place a rotamer in a 4-dimensional torus space  $\mathbb{T}^4 = (\mathbb{S}^1)^4$  where  $\mathbb{S}^1$  denotes a unit circle.
3.2 SidechainDiff: side-chain diffusion probabilistic model
SidechainDiff utilizes a conditional Riemannian diffusion model in  $\mathbb{T}^4$  to build the generative process of the four rotamers and estimate their joint distributions (Figure ①). The generative process is conditioned on the structural context of mutations on the protein-protein interface, which is encoded using a SE(3)-invariant neural network. The learned conditional vectors serve as a natural representation of the structural context.
Conditional diffusion probabilistic model on  $\mathbb{T}^4$  We adopt the continuous score-based framework on compact Riemannian manifolds [De Bortoli et al., 2022] to construct the generative process for rotamers within  $\mathbb{T}^4$ . Our primary adaptation involves conditioning our diffusion model on vectors that are concurrently learned via an encoder network.
Let the random variable  $\mathbf{X}$  denote the states of rotamers. And let  $(\mathbf{X}_t)_{t\in [0,T]}$  and  $(\mathbf{Y}_t)_{t\in [0,T]} = (\mathbf{X}_{T - t})_{t\in [0,T]}$  denote the diffusion process and associated reverse diffusion process in  $\mathbb{T}^4$  , respectively. The stochastic differential equation (SDE) and reverse SDE [De Bortoli et al., 2022, Song et al., 2021] can be defined as follows:
$$
\mathrm {d} \mathbf {X} _ {t} = \mathrm {d} \mathbf {B} _ {t} ^ {\mathcal {M}}. \tag {1}
$$
$$
\mathrm {d} \mathbf {Y} _ {t} = \nabla \log p _ {T - t} (\mathbf {Y} _ {t} | \mathbf {Z}) \mathrm {d} t + \mathrm {d} \mathbf {B} _ {t} ^ {\mathcal {M}}. \tag {2}
$$
Here,  $\mathbf{B}_t^{\mathcal{M}}$  represents the Brownian motion on  $\mathbb{T}^4$ , which is approximated with a Reodesic Random Walk (GRW) on  $\mathbb{T}^4$  (Algorithm 1). The score  $\nabla \log p_{T - t}(\mathbf{Y}_t|\mathbf{Z})$  in Equation 2 is estimated with a
score network  $s_{\theta}(\mathbf{X}, t, \mathbf{Z})$ . We introduce the conditional vector  $\mathbf{Z}$  in the score, parameterized with an encoder network that takes the structural context of mutation sites as inputs. We will delve into these components separately in the following subsections.
Geodesic Random Walk in  $\mathbb{T}^4$  When approximating the Brownian motion on the compact Riemannian manifolds using GRWs, a crucial step is defining the projection map and the exponential map that establishes the connection between the ambient Euclidean space and the associated tangent space, as discussed in previous studies [Jorgensen, 1975; De Bortoli et al., 2022]. In the context of  $\mathbb{T}^4$ , which is constructed as the quadruple product of  $\mathbb{S}^1$  within the ambient Euclidean space  $\mathbb{R}^2$ , it is adequate to focus on modeling Brownian motion on  $\mathbb{S}^1$ , defined as follows:
$$
\mathrm {d} \mathbf {x} _ {t} = \mathbf {K} \cdot \mathrm {d} \mathbf {B} _ {t}, \quad \mathbf {K} = \left( \begin{array}{c c} 0 & - 1 \\ 1 & 0 \end{array} \right). \tag {3}
$$
Here,  $\mathbf{B}_t$  represents the Brownian motion on the real line in the ambient Euclidean space  $\mathbb{R}^2$ , and  $\mathbf{K}$  denotes the diffusion coefficient.
Similarly, the projection map  $\mathrm{Proj}_{\mathbf{Y}_k^\gamma}$  within the tangent space  $T_{\mathbf{Y}_k}(\mathbb{T}^4)$  and the exponential map  $\operatorname{Exp}$  on  $\mathbb{T}^4$  can be derived by applying the Cartesian product of the maps designed for  $\mathbb{S}^1$ . To be more precise, these two maps are defined as follows:
$$
\exp_ {\mu} (\mathbf {v}) = \cos (\| \mathbf {v} \|) \boldsymbol {\mu} + \sin (\| \mathbf {v} \|) \frac {\mathbf {v}}{\| \mathbf {v} \|}. \tag {4}
$$
$$
\operatorname {p r o j} _ {\boldsymbol {\mu}} (\mathbf {z}) = \mathbf {z} - \frac {<   \boldsymbol {\mu} , \mathbf {z} >}{\| \boldsymbol {\mu} \| ^ {2}} \boldsymbol {\mu}. \tag {5}
$$
Here,  $\pmb {\mu}\in \mathbb{S}^1$ $\mathbf{v}\in T_{\pmb{\mu}}(\mathbb{S}^{1})$  , and  $\mathbf{z}\in \mathbb{R}^2$
Algorithm 1 Geodesic Random Walk (GRW) during the sampling phase
Require:  $T, N, Y_0^\gamma, P, \{a_i, \mathbf{R}_i, \mathbf{x}_i, \boldsymbol{\chi}_i\}_{i=1}^{127}$   
1:  $\gamma = T/N$   
2:  $\mathbf{Z} = \operatorname{Enc}_{\phi^*}(\{a_i, \mathbf{R}_i, \mathbf{x}_i, \boldsymbol{\chi}_i\}_{i=1}^{127})$   
3: for  $k \in \{0, \dots, N-1\}$  do  
4:  $\bar{\mathbf{R}}_{k+1} \sim \mathrm{N}(0, \mathrm{Id})$   
5:  $\mathbf{R}_{k+1} = \operatorname{Proj}_{\mathbf{Y}_k^\gamma}(\bar{\mathbf{R}}_{k+1})$   
6:  $\mathbf{W}_{k+1} = -\gamma s_\theta(\mathbf{Y}_k^\gamma, k\gamma, \mathbf{Z}) + \sqrt{\gamma} \mathbf{R}_{k+1}$   
7:  $\mathbf{Y}_{k+1}^\gamma = \operatorname{Exp}_{\mathbf{Y}_k^\gamma}(W_{k+1})$   
8: end for  
9: return  $\{\mathbf{Y}_k^\gamma\}_{k=0}^N$
Score network The architecture of the score network  $s_{\theta}(\mathbf{X}, t, \mathbf{Z})$  is implemented as a three-layer Multilayer Perceptron (MLP) with 512 units in each layer. It takes as inputs the conditional vector from the encoder network representing the structural context of the mutations, the sampled rotamers, and the timestep  $t$ .
Conditional encoder network The encoder network (Figure S1) takes as inputs both single features, consisting of amino acid type, backbone dihedral angles, and local atom coordinates for each amino acid, as well as pair features such as pair distances and relative positions between two amino acids. An SE3-invariant IPA network [Jumper et al., 2021] is utilized to obtain the conditional vector  $\mathbf{Z}$  for the structural context of the mutations. Further information regarding the encoder can be found in App. A.1
Training objectives The parameters  $\theta$  and  $\phi$  in the score network and conditional encoder network are trained simultaneously. We adopt the implicit loss  $l_t^{\mathrm{im}}(s_t)$  in De Bortoli et al. [2022], defined as follows:
$$
l _ {t} ^ {\mathrm {i m}} (s _ {t}) = \int_ {\mathcal {M}} \left\{\frac {1}{2} \| s _ {\theta} (\mathbf {X} _ {t}, t, \mathbf {Z}) \| ^ {2} + \operatorname {d i v} (s _ {t}) (\mathbf {X} _ {t}, t, \mathbf {Z}) \right\} \mathrm {d} \mathbb {P} _ {t} (\mathbf {X} _ {t}). \tag {6}
$$
3.3 DiffAffinity: mutational effect predictor
For the mutational effect predictor, DiffAffinity, we imitate the structure of the mutated protein by setting the side-chain conformation as an empty set and altering the amino acids at the mutation
sites. Subsequently, we use the pre-trained conditional encoder  $\mathrm{Enc}_{\phi^*}$  from SidechainDiff to encode wild-type and mutant structures into distinct hidden representations. These hidden representations, following additional IPA-like transformer network and max aggregation, are fed into a Multilayer Perceptron (MLP) of only one layer for the prediction of  $\Delta \Delta G$ . To assess the quality of our hidden representation, we also utilize the linear combination of SidechainDiff's hidden representations to predict  $\Delta \Delta G$ . More model details can be found in App. A.2
3.4 Model training
We train SidechainDiff with the refined experimental protein structures on the database of Protein Data Bank-REDO [Joosten et al., 2014]. To fairly benchmark the performance of SidechainDiff, we follow the same protocol in RDE [Luo et al., 2023] to preprocess the data and split the train and test sets. We train DiffAffinity with the SKEMPI2 dataset [Jankauskaite et al., 2019]. Due to limited samples, we perform a strategy of three-fold cross-validation to train the model and benchmark the performance. Specifically, we split the SKEMPI2 dataset into three folds based on structure, ensuring that each fold contains unique protein complexes not present in the other folds. Two of the folds are utilized for training and validation, whereas the remaining fold is reserved for testing purposes. More training settings are shown in AppA.
First, we assess DiffAffinity's performance in predicting  $\Delta \Delta G$  using the SKEMPI2 dataset (Section 4.1) and the latest SARS-CoV-2 dataset (Section 4.2). Subsequently, we demonstrate DiffAffinity's effectiveness in optimizing antibodies, with a focus on SARS-CoV-2 as an illustrative case (Section 4.3). Additionally, we evaluate the prediction accuracy of SidechainDiff in side-chain packing.
4.1 Prediction of the mutational effects on binding affinity
DiffAffinity leverages the learned representation from SidechainDiff to predict mutational effects on protein-protein binding. Here, we benchmark the performance on the widely used SKEMPI2 dataset.
Baseline models As mentioned in section 2.2 the models for mutational effects prediction can be categorized into energy-based, sequence-based, unsupervised, end-to-end, and pre-trained models. We have selected several representative models from each category. The selected models include FoldX [Schymkowitz et al., 2005], Rosetta [Alford et al., 2017], flex ddG [Barlow et al., 2018], ESM-1v [Meier et al., 2021], ESM-IF [Hsu et al., 2022], ESM2 [Lin et al., 2023], ESM2*, DDGPred [Shan et al., 2022], and RDE-Network [Luo et al., 2023]. More details of baselines are shown in App. A.4
For ablation studies, we have also developed two variants named DiffAffinity* and DiffAffinity-Linear. DiffAffinity* employs the same model architecture as DiffAffinity but without the learned representations from SidechainDiff and it serves as an end-to-end prediction method. And DiffAffinity-Linear is a direct linear projection of SidechainDiff's hidden representations to  $\Delta \Delta G$ .
Evaluation metrics The performance of our model is evaluated using a variety of metrics. These include traditional prediction accuracy metrics such as root mean squared error (RMSE) and mean absolute error (MAE). Given that the scalar values of experimental  $\Delta \Delta G$  are not standardized, we also employ Pearson and Spearman correlation coefficients to evaluate the models' predictive capability. To assess our model's ability to identify positive or negative mutations affecting protein-protein binding, we use the AUROC and AUPRC metrics. Mutations are classified based on the sign of  $\Delta \Delta G$ , where positive values indicate positive impacts on binding and negative values indicate negative impacts. AUROC measures the overall discriminative power, while AUPRC accounts for imbalanced datasets.
Following the RDE-network ([Luo et al. 2023]), we also evaluate the prediction performance for each complex separately. We group mutations by complex structure, excluding groups with fewer than ten mutation data points, and derive two metrics: average per-structure Pearson correlation coefficient per-structure Pearson and average per-structure Spearman correlation coefficient per-structure Spearman.
Experimental results Our methods, DiffAffinity and DiffAffinity*, outperform all baseline models across almost all metrics and achieve state-of-the-art performance on the SKEMPI2 dataset (Table I).