SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yangtian Zhang $^{1,2*}$  Zuobai Zhang $^{1,2*}$  Bozitao Zhong $^{1,2}$  Sanchit Misra $^{3}$  Jian Tang $^{1,4,5\dagger}$
*equal contribution †corresponding author  
<sup>1</sup>Mila - Québec AI Institute <sup>2</sup>Université de Montréal <sup>3</sup>Intel
$^{4}$ HEC Montréal  $^{5}$ CIFAR AI Research Chair
contacts: <yangtian.zhang, zuobai.zhang> @ mila.quebec, jian.tang@hec.ca
Proteins play a critical role in carrying out biological functions, and their 3D structures are essential in determining their functions. Accurately predicting the conformation of protein side-chains given their backbones is important for applications in protein structure prediction, design and protein-protein interactions. Traditional methods are computationally intensive and have limited accuracy, while existing machine learning methods treat the problem as a regression task and overlook the restrictions imposed by the constant covalent bond lengths and angles. In this work, we present DiffPack, a torsional diffusion model that learns the joint distribution of side-chain torsional angles, the only degrees of freedom in side-chain packing, by diffusing and denoising on the torsional space. To avoid issues arising from simultaneous perturbation of all four torsional angles, we propose autoregressively generating the four torsional angles from  $\chi_{1}$  to  $\chi_{4}$  and training diffusion models for each torsional angle. We evaluate the method on several benchmarks for protein side-chain packing and show that our method achieves improvements of  $11.9\%$  and  $13.5\%$  in angle accuracy on CASP13 and CASP14, respectively, with a significantly smaller model size ( $60\times$  fewer parameters). Additionally, we show the effectiveness of our method in enhancing side-chain predictions in the AlphaFold2 model. Code is available at https://github.com/DeepGraphLearning/DiffPack.
Proteins are crucial for performing a diverse range of biological functions, such as catalysis, signaling, and structural support. Their three-dimensional structures, determined by amino acid arrangement, are crucial for their function. Specifically, amino acid side-chains play a critical role in the stability and specificity of protein structures by forming hydrogen bonds, hydrophobic interactions, and other non-covalent interactions with other side-chains or the protein backbone. Therefore, accurately predicting protein side-chain conformation is an essential problem in protein structure prediction [16, 15, 11], design [48, 54, 14, 65] and protein-protein interactions [64, 24].
Despite recent advancements in deep learning models inspired by AlphaFold2 for predicting the positions of protein backbone atoms [32, 4], predicting the conformation of protein side-chains remains a challenging problem due to the complex interactions between side chains. In this work, we focus on the problem of predicting side-chain conformation with fixed backbone structure, a.k.a., protein side-chain packing. Traditional methods for side-chain prediction rely on techniques such as rotamer libraries, energy functions, and Monte Carlo sampling [27, 73, 37, 3, 1, 35, 71, 7]. However, these methods are computationally intensive and struggle to accurately capture the complex energy landscape of protein side chains due to their reliance on search heuristics and discrete sampling.
Several machine learning methods have been proposed for side-chain prediction, including DL-Packer [46], AttnPacker [45], and others [47, 69, 70, 73, 39, 71]. DLPacker, the first deep learning-based model, employs a 3D convolution network to learn the density map of side-chain atoms, but it lacks the ability to capture rotation equivariance due to its 3D-CNN structure. In contrast, AttnPacker, the current state-of-the-art model, directly predicts side-chain atom coordinates using Tensor Field Network and SE(3)-Transformer, ensuring rotation equivariance in side chain packing. However, it does not consider the restrictions imposed by covalent bond lengths and angles, leading to inefficient training and unnatural bond lengths during generation. Furthermore, previous methods that treat side-chain packing as a regression problem assume a single ground-truth side-chain structure and overlook the fact that proteins can fold into diverse structures under different environmental factors, resulting in a distribution of side-chain conformations.
In this study, we depart from the standard practice of focusing on atom-level coordinates in Cartesian space as in prior research [45, 46]. Instead, we introduce DiffPack, a torsional diffusion model that models the exact degree of freedom in side-chain packing, the joint distribution of four torsional angles. By perturbing and denoising in the torsional space, we use an SE(3)-invariant network to learn the gradient field for the joint distribution of torsional angles. This result in a much smaller conformation space of side-chain, thereby capturing the intricate energy landscape of protein side chains. Despite its effectiveness, a direct joint diffusion process on the four torsion angles could result in steric clashes and accumulative coordinate displacement, which complicates the denoising process. To address this, we propose an autoregressive diffusion process and train separate diffusion models to generate the four torsion angles from  $\chi_{1}$  to  $\chi_{4}$  in an autoregressive manner. During training, each diffusion model only requires perturbation on its corresponding torsional angle using a teacher-forcing strategy, preserving the protein structure and avoiding the aforementioned issues. To improve the capacity of our model, we further introduce three schemes in sampling for consistently improving the inference results: multi-round sampling, annealed temperature sampling, and confidence models.
We evaluate our method on several benchmarks for protein side-chain packing and compare it with existing state-of-the-art methods. Our results demonstrate that DiffPack outperforms existing state-of-the-art approaches, achieving remarkable improvements of  $11.9\%$  and  $13.5\%$  in angle accuracy on CASP13 and CASP14, respectively. Remarkably, these performance gains are achieved with a significantly smaller model size, approximately 60 times fewer parameters, highlighting the potential of autoregressive diffusion models in protein structure prediction. Furthermore, we showcase the effectiveness of our method in enhancing the accuracy of side-chain predictions in the AlphaFold2 model, indicating its complementary nature to existing approaches.
Protein. Proteins are composed of a sequence of residues (amino acids), each containing an alpha carbon  $(\mathrm{C}_{\alpha})$  atom bounded to an amino group  $(- \mathrm{NH}_2)$ , a carboxyl group  $(- \mathrm{COOH})$  and a side-chain  $(- \mathrm{R})$  that identifies the residue type. Peptide bonds link consecutive residues through a dehydration synthesis process. The backbone of a protein consists of the  $\mathrm{C}_{\alpha}$  atom and the connected nitrogen, carbon, and oxygen atoms. We use  $S = [s_1, s_2, \dots, s_n]$  to denote the sequence of a protein with  $n$  residues, where  $s_i \in \{0, \dots, 19\}$  denotes the type of the  $i$ -th residue.
Protein Conformation. Physical interactions between residues make a protein fold into its native 3D structure, a.k.a., conformation, which determines its biologically functional activity. We use  $\mathcal{X} = [x_1,x_2,\dots,x_n]$  to denote the structure of a protein, where  $x_{i}$  denotes the set of atom coordinates belonging to the  $i$ -th residue. The backbone structure  $x_{i}^{(\mathrm{bb})}$  of the  $i$ -th residue is a subset consisting of backbone atoms, i.e.,  $\mathrm{C}_{\alpha}$ ,  $\mathrm{N}$ ,  $\mathrm{C}$ ,  $\mathrm{O}$ , while the side-chain consists of the remaining atoms  $x_{i}^{(\mathrm{sc})} = x_{i} \setminus x_{i}^{(\mathrm{bb})}$ . The backbone conformation  $\mathcal{X}^{(\mathrm{bb})}$  and side-chain conformation  $\mathcal{X}^{(\mathrm{sc})}$  of the protein are defined as the set of backbones and side-chains of all residues, respectively.
Protein Side-Chain Packing. Protein side-chain packing (PSCP) problem aims to predict the 3D coordinates  $\mathcal{X}^{(\mathrm{sc})}$  of side-chain atoms given backbone conformations  $\mathcal{X}^{(\mathrm{bb})}$  and protein sequence  $S$ . That is, we aim to model the conditional distribution  $p(\mathcal{X}^{(\mathrm{sc})}|\mathcal{X}^{(\mathrm{bb})}, S)$ .
$p(\chi_1)$
$p(\chi_2\mid \chi_1)$
Figure 1: Overview of DiffPack. Given a protein sequence and backbone structure, we aim to model the conditional distribution of side-chain conformation. (A) Distribution of side-chain conformation is modeled through diffusion process in torsion space  $\mathbb{T}^m$ . An SE(3)-invariant network is used to learn the torus force field (torsion score). (B) Four torsion angles are generated autoregressively across all residues.
$p(\chi_3\mid \chi_1,\chi_2)$
$p(\chi_4\mid \chi_1,\chi_2,\chi_3)$
In this paper, we introduce an autoregressive diffusion model DiffPack to predict the side-chain conformation distribution in torsion space. We address the issue of overparameterization by introducing a torsion-based formulation of side-chain conformation in Section 3.1. Then we give a formulation of the torsional diffusion model in Section 3.2. However, directly applying the diffusion model encounters challenges in learning the joint distribution, which we address by introducing an autoregressive-style model in Section 3.3. We then provide details of the model architecture in Section 3.4, followed by an explanation of the inference procedure in Section 3.5.
3.1 Modeling Side-Chain Conformations with Torsional Angles
Previous methods [45, 46] model the side chain conformation as a series of three-dimensional coordinates in the Cartesian space. However, this approach does not take into account the restrictions imposed by constant covalent bond lengths and angles on the side chain's degree of freedom, resulting in inefficient training and unnatural bond lengths during generation.
To overcome this overparameterization issue, we propose modeling side chain conformation in torsion space. As illustrated in Figure 2, torsional angles directly dictate protein side-chain conformation by determining the twist between two neighboring planes. Table 6 lists the atom groups to define corresponding the neighboring plane for each residue type. The number of torsional angles varies across different residues, with a maximum of four  $(\chi_{1},\chi_{2},\chi_{3},\chi_{4})$ . Modeling side chains in torsion space reduces the number of variables by approximately one third and restricts degrees of freedom [6]. Formally, we transform the problem of side-chain packing to modeling the distribution of torsional angles:
$$
p \left(\mathcal {X} ^ {\left(\mathrm {s c}\right)} \mid \mathcal {X} ^ {\left(\mathrm {b b}\right)}, \mathcal {S}\right) \Leftrightarrow p \left(\chi_ {1}, \chi_ {2}, \chi_ {3}, \chi_ {4} \mid \mathcal {X} ^ {\left(\mathrm {b b}\right)}, \mathcal {S}\right), \tag {1}
$$
Figure 2: Illustration of four torsional angles.
where  $\chi_i \in [0, 2\pi)^{m_i}$  is a vector of the  $i$ -th torsional angles of all residues in the protein and we use  $m_i$  to denote the number of residues with  $\chi_i$ . The space of possible side chain conformation is, therefore, reduced to an  $m$ -dimension sub-manifold  $\mathcal{M} \subset \mathbb{R}^{3n}$  with  $m = m_1 + m_2 + m_3 + m_4$ .
3.2 Diffusion Models on Torsional Space
Denoising diffusion probabilistic models are generative models that learn the data distribution via a forward diffusion process and a reverse generation process [23]. We follow [30] to define diffusion models on the torsional space with the continuous score-based framework in [57]. For simplicity, we omit the condition and aggregate torsional angles on all residues as a torsional vector  $\chi \in [0,2\pi)^{m}$ . Starting with the data distribution as  $p_0(\chi)$ , the forward diffusion process is modeled by a stochastic differential equation (SDE):
$$
d \boldsymbol {\chi} = \boldsymbol {f} (\boldsymbol {\chi}, t) d t + g (t) d \boldsymbol {w}, \tag {2}
$$
where  $\boldsymbol{w}$  is the Wiener process on the torsion space and  $f(\chi, t)$ ,  $g(t)$  are the drift coefficient and diffusion coefficient, respectively. Here we adopt Variance-Exploding SDE where  $f(\chi, t) = 0$  and  $g(t)$  is exponentially decayed with  $t$ . With sufficiently large  $T$ , the distribution  $p_T(\chi)$  approaches a uniform distribution over the torsion space. The reverse generation process samples from the prior and generates samples from the data distribution  $p_0(\chi)$  via approximately solving the reverse SDE:
$$
d \boldsymbol {\chi} = \left[ \boldsymbol {f} \left(\boldsymbol {\chi} _ {t}, t\right) - g ^ {2} (t) \nabla_ {\boldsymbol {\chi}} \log p _ {t} \left(\boldsymbol {\chi} _ {t}\right) \right] d t + g (t) d \boldsymbol {w}, \tag {3}
$$
where a neural network is learned to fit the score  $\nabla_{\chi}\log p_t(\chi_t)$  of the diffused data [23, 57]. Inspired by [30], we convert the torsional angles into 3D atom coordinates and define the score network on Euclidean space, enabling it to explicitly learn the interatomic interactions.
The training process involves sampling from the perturbation kernel of the forward diffusion and computing its score to train the score network. Given the equivalence  $(\chi_1, \ldots, \chi_m) \sim (\chi_1 + 2\pi, \ldots, \chi_m) \ldots \sim (\chi_1, \ldots, \chi_m + 2\pi)$  in torsion space, the perturbation kernel is a wrapped Gaussian distribution on  $\mathbb{R}^m$ . This means that any  $\chi, \chi' \in [0, 2\pi)^m$ , the perturbation kernel is proportional to the sum of exponential terms, which depends on the distance between  $\chi$  and  $\chi'$ :
$$
p _ {t \mid 0} \left(\chi^ {\prime} \mid \chi\right) \propto \sum_ {\boldsymbol {d} \in \mathbb {Z} ^ {m}} \exp \left(- \frac {\| \boldsymbol {\chi} - \boldsymbol {\chi} ^ {\prime} + 2 \pi \boldsymbol {d} \| ^ {2}}{2 \sigma^ {2} (t)}\right), \tag {4}
$$
In order to sample from the perturbation kernel, we sample from the corresponding unwrapped isotropic normal and take element-wise mod  $2\pi$ . The kernel's scores are pre-computed using a numerical approximation. During training, we uniformly sampled a time step  $t$  and the denoising score matching loss is minimized:
$$
J _ {\mathrm {D S M}} (\theta) = \mathbb {E} _ {t} \left[ \lambda (t) \mathbb {E} _ {\boldsymbol {\chi} _ {0} \sim p _ {0}, \boldsymbol {\chi} _ {t} \sim p _ {t | 0} (\cdot | \boldsymbol {\chi} _ {0})} \left[ \| \boldsymbol {s} (\boldsymbol {\chi} _ {t}, t) - \nabla_ {t} \log p _ {t | 0} (\boldsymbol {\chi} _ {t} | \boldsymbol {\chi} _ {0}) \| ^ {2} \right] \right], \tag {5}
$$
where the weight factors  $\lambda(t) = 1 / \mathbb{E}_{\boldsymbol{\chi} \sim p_{t|0}(\cdot | 0)}\left[\|\nabla_{\boldsymbol{\chi}} \log p_{t|0}(\boldsymbol{\chi}|\mathbf{0})\|^2\right]$  are also pre-computed.
3.3 Autoregressive Diffusion Models
The direct approach to torsional diffusion models introduces noise on all torsional angles simultaneously and poses two significant challenges:
1. Cumulative coordinate displacement: Adding noises to torsional angles is equivalent with rotating side-chain atoms in and beyond the corresponding atom groups. For instance, if  $\chi_{1}$  is rotated, it affects the coordinates of atoms in the  $\chi_{2},\chi_{3}$ , and  $\chi_{4}$  groups. This cumulative effect complicates the denoising process of the latter three angles. Similar issues arise after rotating  $\chi_{2}$  and  $\chi_{3}$ . The effect of rotating  $\chi_{1}$  is illustrated in Figure 3.
2. Excessive steric clash: Adding noises to all torsional angles may damage protein structures and complicate the denoising process in our score network. In Appendix C, Figure 8 compares the number of atom clashes observed when using noise schemes in vanilla diffusion models and autoregressive-style ones to show the straightforward application of diffusion models results in significantly more steric clashes.
Figure 3: Effects of rotating  $\chi_{1}$ .
To address the aforementioned issues, we propose an autoregressive diffusion model over the torsion space. We factorize the joint distribution of the four torsional angles into separate conditional
distributions. Specifically, we have
$$
p \left(\chi_ {1}, \chi_ {2}, \chi_ {3}, \chi_ {4}\right) = p \left(\chi_ {1}\right) \cdot p \left(\chi_ {2} \mid \chi_ {1}\right) \cdot p \left(\chi_ {3} \mid \chi_ {1, 2}\right) \cdot p \left(\chi_ {4} \mid \chi_ {1, 2, 3}\right). \tag {6}
$$
This allows us to model the side-chain packing problem as a step-by-step generation of torsional angles: first, we predict the first torsional angles  $\chi_{1}$  for all residues given the protein backbone structure; next, based on the backbone and the generated  $\chi_{1}$ , we predict the second torsional angles  $\chi_{2}$ ; and so on for  $\chi_{3}$  and  $\chi_{4}$ .
We train a separate score network for each distribution on the right-hand side using the torsional diffusion model from Section 3.2. To train the autoregressive model, we use a teacher-forcing strategy. Specifically, when modeling  $p(\chi_i|\chi_{1,\dots,i - 1})$ , we assume that  $\chi_{1,\dots,i - 1}$  are ground truth and keep these angles fixed. We then rotate  $\chi_{i}$  by sampling noises from the perturbation kernel in (4) and discard all atoms belonging to the  $\chi_{i + 1,\dots,4}$  groups to remove the dependency on following torsional angles. This approach eliminates the cumulative effects of diffusion process on  $\chi_{i}$  and preserves the overall structure of the molecule, thus overcoming the aforementioned challenges. The generation process is illustrated in Figure 1 and the training procedure is described in Algorithm 1.
3.4 Model Architecture
To model  $p(\chi_i|\chi_{1,\dots,i-1})$ , we utilize a score network constructed using the 3D coordinates  $\boldsymbol{s}(\mathcal{X}_t,t)$  of backbone atoms and atoms in the  $\chi_{1,\dots,i-1}$  groups. The score network's output represents the noise on torsional angles, which should be SE(3)-invariant w.r.t. the conformation  $\mathcal{X}_t$ . Unlike previous methods [45] operating in Cartesian coordinates that require an SE(3)-equivariant score network, our method only requires SE(3)-invariance, providing greater flexibility in designing the model architecture. To ensure this invariance, we employ GearNet-Edge [75], a state-of-the-art protein structure encoder. This involves constructing a multi-relational graph with atoms as nodes, where node features consist of one-hot encoding for atom types, corresponding residue types, and time step embeddings. Edges are added based on chemical bond and 3D spatial information, determining their type. To learn representations for each node, we perform relational message passing between them [50]. We denote the edge between nodes  $i$  and  $j$  with type  $r$  as  $(i,j,r)$  and set of relations as  $\mathcal{R}$ . We use  $\boldsymbol{h}_i^{(l)}$  to denote the hidden representation of node  $i$  at layer  $l$ . Then, message passing can be written as:
$$
\left. \boldsymbol {h} _ {i} ^ {(l)} = \boldsymbol {h} _ {i} ^ {(l - 1)} + \sigma \left(\operatorname {B a t c h N o r m} \left(\sum_ {r \in \mathcal {R}} \boldsymbol {W} _ {r} \sum_ {j \in \mathcal {N} _ {r} (i)} \left(\boldsymbol {h} _ {j} ^ {(l - 1)} + \operatorname {L i n e a r} \left(\boldsymbol {m} _ {(i, j, r)} ^ {(l)}\right)\right)\right)\right), \right. \tag {7}
$$
where  $W_{r}$  is the learnable weight matrix for relation type  $r$ ,  $\mathcal{N}_r(j)$  is the neighbor set of  $j$  with relation type  $r$ , and  $\sigma(\cdot)$  is the activation function. Edge representations  $m^{(l)}(i,j,r)$  are obtained through edge message passing. We use  $e$  as the abbreviation of the edge  $(i,j,r)$ . Two edges  $e_1$  and  $e_2$  are connected if they share a common end node, with the connection type determined by the discretized angle between them. The edge message passing layer can be written as:
$$
\boldsymbol {m} _ {e _ {1}} ^ {(l)} = \sigma \left(\operatorname {B a t c h N o r m} \left(\sum_ {r \in \mathcal {R} ^ {\prime}} \boldsymbol {W} _ {r} ^ {\prime} \sum_ {e _ {2} \in \mathcal {N} _ {r} ^ {\prime} (e _ {1})} \boldsymbol {m} _ {e _ {2}} ^ {(l - 1)}\right)\right), \tag {8}
$$
where  $\mathcal{R}'$  is the set of relation types between edges and  $\mathcal{N}_r'(e_1)$  is the neighbor set of  $e_1$  with relation type  $r$ . After obtaining the hidden representations of all atoms at layer  $L$ , we compute the residue representations by taking the mean of the representations of its constituent atoms. The residue representations are then fed into an MLP for score prediction. The details of architectural components are summarized in Appendix D.
3.5 Inference
After completing the training, we adopt the common practice of autoregressive and diffusion models for inference, as described in Algorithm 2. We generate four torsional angles step by step as described in Section 3.3. When sampling  $\chi_{i}$  based on the predicted torsional angles  $\chi_{1,\dots,i-1}$ , we begin by sampling a random angle from the uniform prior and then discretize and solve the reverse diffusion. At each time step, we generate atoms in the  $\chi_{1,\dots,i}$  group and use our learned score network for denoising. We also discover several simple techniques that significantly improve our performance, including multi-round sampling, annealed temperature sampling and confidence models.
Predictor-corrector sampling. After discretizing the reverse diffusion SDE, we perform multiple Langevin steps subsequent to each denoising step. This hybrid method, which mixes the denoising
process with Langevin dynamics, is called predictor-corrector sampling, as suggested by [57]. This approach can be seen as introducing an equilibration process to stabilize  $p_t$ , and it is demonstrated to be effective in diffusion process.
Annealed temperature sampling. When designing a generative model, two critical aspects to consider are quality and diversity. The diffusion model often suffers from overdispersion, which prioritizes diversity over sampling quality. However, in the context of side chain packing, quality is more important than diversity. Directly using the standard reverse sampling process may lead to undesirable structures. Following [28], we utilize an annealed temperature sampling scheme to mitigate this issue. Specifically, We modify the reverse SDE by adding an annealed weight  $\lambda_{t}$  to the score function (details in Appendix E):
$$
d \boldsymbol {\chi} = - \lambda_ {t} \frac {d \sigma^ {2} (t)}{d t} \nabla_ {\boldsymbol {\chi}} \log p _ {t} (\boldsymbol {\chi} _ {t}) d t + \sqrt {\frac {d \sigma^ {2} (t)}{d t}} d \boldsymbol {w}, \text {w h e r e} \lambda_ {t} = \frac {\sigma_ {\max } ^ {2}}{T \sigma_ {\max } ^ {2} - (T - 1) \sigma^ {2} (t)}. \tag {9}
$$
The above modification results in a reverse process approximating the low temperature sampling process, where ideally decreasing temperture  $T$  lead to a sampling process biased towards quality.
Confidence model. As per common practice in protein structure prediction [32], we train a confidence model to select the best prediction among multiple conformations sampled from our model. The model architecture we use is the same as that used in diffusion models, which takes the entire protein conformation as input and outputs representations for each residue. We train the model on the same dataset, using the residue representations to predict the negative residue-level RMSD of our sampled conformations. When testing, we rank all conformations based on our confidence model.
3.6 Handling Symmetric Issues
Torsional angles generally exhibit a periodicity of  $2\pi$ . However, certain rigid side-chain structures possess a  $\pi$ -rotation-symmetry, meaning that rotating the torsional angle by  $\pi$  does not yield distinct physical structures, as demonstrated in Figure 4. For instance, in the case of a tyrosine residue, the phenolic portion of its side chain ensures that a  $\pi$  rotation of its  $\chi_{2}$  torsion angles only alters the internal atom name without affecting the actual structure.