SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kai Yi *
University of New South Wales kai.yi@unsw.edu.au
Bingxin Zhou *
Shanghai Jiao Tong University  
bingxin.zhou@sjtu.edu.cn
Yiqing Shen
Johns Hopkins University  
yshen92@jhu.edu
Pietro Lio
University of Cambridge pl219@cam.ac.uk
Yu Guang Wang
Shanghai Jiao Tong University  
University of New South Wales  
yuguang.wang@sjtu.edu.cn
Inverse protein folding is challenging due to its inherent one-to-many mapping characteristic, where numerous possible amino acid sequences can fold into a single, identical protein backbone. This task involves not only identifying viable sequences but also representing the sheer diversity of potential solutions. However, existing discriminative models, such as transformer-based auto-regressive models, struggle to encapsulate the diverse range of plausible solutions. In contrast, diffusion probabilistic models, as an emerging genre of generative approaches, offer the potential to generate a diverse set of sequence candidates for determined protein backbones. We propose a novel graph denoising diffusion model for inverse protein folding, where a given protein backbone guides the diffusion process on the corresponding amino acid residue types. The model infers the joint distribution of amino acids conditioned on the nodes' physiochemical properties and local environment. Moreover, we utilize amino acid replacement matrices for the diffusion forward process, encoding the biologically meaningful prior knowledge of amino acids from their spatial and sequential neighbors as well as themselves, which reduces the sampling space of the generative process. Our model achieves state-of-the-art performance over a set of popular baseline methods in sequence recovery and exhibits great potential in generating diverse protein sequences for a determined protein backbone structure. The code is available on https://github.com/ykiiiiii/GraDe_IF.
Inverse protein folding, or inverse folding, aims to predict feasible amino acid (AA) sequences that can fold into a specified 3D protein structure [22]. The results from inverse folding can facilitate the design of novel proteins with desired structural and functional characteristics. These proteins can serve numerous applications, ranging from targeted drug delivery to enzyme design for both academic and industrial purposes [25, 31, 37, 52]. In this paper, we develop a diffusion model tailored for graph node denoising to obtain new AA sequences given a protein backbone.
Despite its importance, inverse folding remains challenging due to the immense sequence space to explore, coupled with the complexity of protein folding. On top of energy-based physical reasoning of a protein's folded state [1], recent advancements in deep learning yield significant progress in learning the mapping from protein structures to AA sequences directly. For example, discriminative models
*equal contribution.
Figure 1: Overview of GRADE-IF. In the diffusion process, the original amino acid is stochastically transitioned to other amino acids, leveraging BLOSUM with varied temperatures as the transition kernel. During the denoising generation phase, initial node features are randomly sampled across the 20 amino acids with a uniform distribution. This is followed by a gradual denoising process, conditional on the graph structure and protein secondary structure at different time points. We employ a roto-translation equivariant graph neural network as the denoising network.
formulate this problem as the prediction of the most likely sequence for a given protein structure [7, 10, 17, 50]. However, they have struggled to accurately capture the one-to-many mapping from the protein structure to non-unique AA sequences.
Due to their powerful learning ability, diffusion probabilistic models have gained increasing attention. They are capable of generating a diverse range of molecule outputs from a fixed set of conditions given the inherent stochastic nature. For example, Torsion Diffusion [19] learns the distribution of torsion angles of heavy atoms to simulate conformations for small molecules. Concurrently, SMCDIFF [43] enhances protein folding tasks by learning the stable scaffold distribution supporting a target motif with diffusion. Similarly, DIFFDOCK [5] adopts a generative approach to protein-ligand docking, creating a range of possible ligand binding poses for a target pocket structure.
Despite the widespread use of diffusion models, their comprehensive potential within the context of protein inverse folding remains relatively unexplored. Current methods in sequence design are primarily anchored in language models, encompassing Masked Language Models (MLMs) [25, 23] and autoregressive generative models [17, 26, 28]. By tokenizing AAs, MLMs formulate the sequence generation tasks as masked token enrichment. These models usually operate by drawing an initial sequence with a certain number of tokens masked as a specific schedule and then learning to predict the masked tokens from the given context. Intriguingly, this procedure can be viewed as a discrete diffusion-absorbing model when trained by a parameterized objective. Autoregressive models, conversely, can be perceived as deterministic diffusion processes [3]. It induces conditional distribution to each token, but the overall dependency along the entire AA sequence is recast via an independently-executed diffusion process.
On the contrary, diffusion probabilistic models employ an iterative prediction methodology that generates less noisy samples and demonstrates potential in capturing the diversity inherent in real data distributions. This unique characteristic further underscores the promising role diffusion models could play in advancing the field of protein sequence design. To bridge the gap, we make the first attempt at a diffusion model for inverse folding. We model the inverse problem as a denoising problem where the randomly assigned AA types in a protein (backbone) graph is recovered to the wild type. The protein graph which contains the spatial and biochemical information of all AAs is represented by equivariant graph neural networks, and diffusion process takes places on graph nodes. In real inverse folding tasks, the proposed model achieves SOTA recovery rate, improve  $4.2\%$  and  $5.4\%$  on recovery rate for single-chain proteins and short sequences, respectively, , especially for conserved region which has a biologically significance. Moreover, the predicted structure of generated sequence is identical to the structure of native sequence.
The preservation of the desired functionalities is achieved by innovatively conditioning the model on both secondary and third structures in the form of residue graphs and corresponding node features. The major contributions of this paper are three-fold. Firstly, we propose GRADE-IF, a diffusion model backed by roto-translation equivariant graph neural network for inverse folding. It stands out
from its counterparts for its ability to produce a wide array of diverse sequence candidates. Secondly, as a departure from conventional uniform noise in discrete diffusion models, we encode the prior knowledge of the response of AAs to evolutionary pressures by the utilization of Blocks Substitution Matrix as the translation kernel. Moreover, to accelerate the sampling process, we adopt Denoising Diffusion Implicit Model (DDIM) from its original continuous form to suit the discrete circumstances and back it with thorough theoretical analysis.
2.1 Residue Graph by Protein Backbone
A residue graph, denoted as  $\mathcal{G} = (X,A,E)$ , aims to delineate the geometric configuration of a protein. Specifically, every node stands for an AA within the protein. Correspondingly, each node is assigned a collection of meticulously curated node attributes  $X$  to reflect its physiochemical and topological attributes. The local environment of a given node is defined by its spatial neighbors, as determined by the  $k$ -nearest neighbor ( $k\mathrm{NN}$ ) algorithm. Consequently, each AA node is linked to a maximum of  $k$  other nodes within the graph, specifically those with the least Euclidean distance amongst all nodes within a  $30\AA$  contact region. The edge attributes, represented as  $E\in \mathbb{R}^{93}$ , illustrate the relationships between connected nodes. These relationships are determined through parameters such as inter-atomic distances, local N-C positions, and a sequential position encoding scheme. We detail the attribute construction in Appendix C.
2.2 Inverse Folding as a Denoising Problem
The objective of inverse folding is to engineer sequences that can fold to a pre-specified desired structure. We utilize the coordinates of  $C\alpha$  atoms to represent the 3D positions of AAs in Euclidean space, thereby embodying the protein backbone. Based on the naturally existing protein structures, our model is constructed to generate a protein's native sequence based on the coordinates of its backbone atoms. Formally we represent this problem as learning the conditional distribution  $p(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}^{\mathrm{pos}})$ . Given a protein of length  $n$  and a sequence of spatial coordinates  $\mathbf{X}^{\mathrm{pos}} = \{\pmb{x}_1^{\mathrm{pos}},\dots,\pmb{x}_i^{\mathrm{pos}},\dots,\pmb{x}_n^{\mathrm{pos}}\}$  representing each of the backbone  $C\alpha$  atoms in the structure, the target is to predict  $\mathbf{X}^{\mathrm{aa}} = \{\pmb{x}_1^{\mathrm{aa}},\dots,\pmb{x}_i^{\mathrm{aa}},\dots,\pmb{x}_n^{\mathrm{aa}}\}$ , the native sequence of AAs. This density is modeled in conjunction with the other AAs along the entire chain. Our model is trained by minimizing the negative log-likelihood of the generated AA sequence relative to the native wild-type sequence. Sequences can then be designed either by sampling or by identifying sequences that maximize the conditional probability given the desired secondary and tertiary structure.
2.3 Discrete Denoising Diffusion Probabilistic Models
Diffusion models belong to the class of generative models, where the training stage encompasses diffusion and denoising processes. The diffusion process  $q(\pmb{x}_1, \dots, \pmb{x}_T \mid \pmb{x}_0) = \prod_{t=1}^{T} q(\pmb{x}_t \mid \pmb{x}_{t-1})$  corrupts the original data  $\pmb{x}_0 \sim q(\pmb{x})$  into a series of latent variables  $\{\pmb{x}_1, \dots, \pmb{x}_T\}$ , with each carrying progressively higher levels of noise. Inversely, the denoising process  $p_\theta(\pmb{x}_0, \pmb{x}_1, \dots, \pmb{x}_T) = p(\pmb{x}_T) \prod_{t=1}^{T} p_\theta(\pmb{x}_{t-1} \mid \pmb{x}_t)$  gradually reduces the noise within these latent variables, steering them back towards the original data distribution. The iterative denoising procedure is driven by a differentiable operator, such as a trainable neural network.
While in theory there is no strict form for  $q(\pmb{x}_t \mid \pmb{x}_{t-1})$  to take, several conditions are required to be fulfilled by  $p_\theta$  for efficient sampling: (i) The diffusion kernel  $q(\pmb{x}_t \mid \pmb{x}_0)$  requires a closed form to sample noisy data at different time steps for parallel training. (ii) The kernel should possess a tractable formulation for the posterior  $q(\pmb{x}_{t-1} \mid \pmb{x}_t, \pmb{x}_0)$ . Consequently, the posterior  $p_\theta(\pmb{x}_{t-1} \mid \pmb{x}_t) = \int q(\pmb{x}_{t-1} \mid \pmb{x}_t, \pmb{x}_0) \, \mathrm{d}p_\theta(\pmb{x}_0 \mid \pmb{x}_t)$ , and  $\pmb{x}_0$  can be used as the target of the trainable neural network. (iii) The marginal distribution  $q(\pmb{x}_T)$  should be independent of  $\pmb{x}_0$ . This independence allows us to employ  $q(\pmb{x}_T)$  as a prior distribution for inference.
The aforementioned criteria are crucial for the development of suitable noise-adding modules and training pipelines. To satisfy these prerequisites, we follow the setting in previous work [3]. For categorical data  $\boldsymbol{x}_t \in \{1, \dots, K\}$ , the transition probabilities are calculated by the matrix  $[\boldsymbol{Q}_t]_{ij} = q(\boldsymbol{x}_t = j \mid \boldsymbol{x}_{t-1} = i)$ . Employing the transition matrix and on one-hot encoded categorical feature
$x_{t}$ , we can define the transitional kernel in the diffusion process by:
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) = \boldsymbol {x} _ {t - 1} \boldsymbol {Q} _ {t} \quad \text {a n d} \quad q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x}\right) = \boldsymbol {x} \bar {\boldsymbol {Q}} _ {t}, \tag {1}
$$
where  $\bar{Q}_t = Q_1 \dots Q_t$ . The Bayes rule yields that the posterior distribution can be calculated in closed form as  $q(x_{t-1} \mid x_t, x) \propto x_t Q_t^\top \odot x \bar{Q}_{t-1}$ . The generative probability can thus be determined using the transition kernel, the model output at time  $t$ , and the state of the process  $x_t$ . Through iterative sampling, we eventually produce the generated output  $x_0$ .
The prior distribution  $p(\pmb{x}_T)$  should be independent of the observation  $\pmb{x}_0$ . Consequently, the construction of the transition matrix necessitates the use of a noise schedule. The most straightforward and commonly utilized method is the uniform transition, which can be parameterized as  $Q_t = \alpha_t I + (1 - \alpha_t) \mathbf{1}_d \mathbf{1}_d^\top / d$  with  $I^\top$  be the transpose of the identity matrix  $I$ ,  $d$  refers to the number of amino acid types (i.e.,  $d = 20$ ) and  $\mathbf{1}_d$  denotes the one vector of dimension  $d$ . As  $t$  approaches infinity,  $\alpha$  undergoes a progressive decay until it reaches 0. Consequently, the distribution  $q(\pmb{x}_T)$  asymptotically approaches a uniform distribution, which is essentially independent of  $\pmb{x}$ .
3 Graph Denoising Diffusion for Inverse Protein Folding
In this section, we introduce a discrete graph denoising diffusion model for protein inverse folding, which utilizes a given graph  $\mathcal{G} = \{\pmb{X},\pmb{A},\pmb{E}\}$  with node feature  $\pmb{X}$  and edge feature  $\pmb{E}$  as the condition. Specifically, the node feature depicts the AA position, AA type, and the spatial and biochemical properties  $\pmb{X} = [X^{\mathrm{pos}},X^{\mathrm{aa}},X^{\mathrm{prop}}]$ . We define a diffusion process on the AA feature  $X^{\mathrm{aa}}$ , and denoise it conditioned on the graph structure  $\pmb{E}$  which is encoded by equivariant neural networks [35]. Moreover, we incorporate protein-specific prior knowledge, including an AA substitution scoring matrix and protein secondary structure during modeling. We also introduce a new acceleration algorithm for the discrete diffusion generative process based on a transition matrix.
3.1 Diffusion Process and Generative Denoising Process
Diffusion Process To capture the distribution of AA types, we independently add noise to each AA node of the protein. For any given node, the transition probabilities are defined by the matrix  $Q_{t}$ . With the predefined transition matrix, we can define the forward diffusion kernel by
$$
q \left(\boldsymbol {X} _ {t} ^ {\mathrm {a a}} \mid \boldsymbol {X} _ {t - 1} ^ {\mathrm {a a}}\right) = \boldsymbol {X} _ {t - 1} ^ {\mathrm {a a}} \boldsymbol {Q} _ {t} \quad \text {a n d} \quad q \left(\boldsymbol {X} _ {t} ^ {\mathrm {a a}} \mid \boldsymbol {X} ^ {\mathrm {a a}}\right) = \boldsymbol {X} ^ {\mathrm {a a}} \bar {\boldsymbol {Q}} _ {t},
$$
where  $\bar{\pmb{Q}}_t = \pmb{Q}_1\ldots \pmb{Q}_t$  is the transition probability matrix up to step  $t$ .
Training Denoising Networks The second component of the diffusion model is the denoising neural network  $f_{\theta}$ , parameterized by  $\theta$ . This network accepts a noisy input  $\mathcal{G}_t = (X_t, \mathbf{E})$ , where  $X_t$  is the concatenation of the noisy AA types and other AA properties including 20 one-hot encoded AA type and 15 geometry properties, such as SASA, normalized surface-aware node features, dihedral angles of backbone atoms, and 3D positions. It aims to predict the clean type of AA  $X^{\mathrm{aa}}$ , which allows us to model the underlying sequence diversity in the protein structure while maintaining their inherent structural constraints. To train  $f_{\theta}$ , we optimize the cross-entropy loss  $L$  between the predicted probabilities  $\hat{p}(X^{\mathrm{aa}})$  for each node's AA type.
Parameterized Generative Process A new AA sequence is generated through the reverse diffusion iterations on each node  $\mathbf{x}$ . The generative probability distribution  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t)$  is estimated from the predicted probability  $\hat{p} (\mathbf{x}^{\mathrm{aa}}|\mathbf{x}_t)$  by the neural networks. We marginalize over the network predictions to compute for generative distribution at each iteration:
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) \propto \sum_ {\hat {\boldsymbol {x}} ^ {\mathrm {a a}}} q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) \hat {p} _ {\theta} \left(\boldsymbol {x} ^ {\mathrm {a a}} \mid \boldsymbol {x} _ {t}\right), \tag {2}
$$
where the posterior
$$
q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) = \operatorname {C a t} \left(\boldsymbol {x} _ {t - 1} \left| \frac {\boldsymbol {x} _ {t} Q _ {t} ^ {\top} \odot \boldsymbol {x} ^ {\mathrm {a a}} \bar {Q} _ {t - 1}}{\boldsymbol {x} ^ {\mathrm {a a}} \bar {Q} _ {t} \boldsymbol {x} _ {t} ^ {\top}}\right)\right) \tag {3}
$$
can be calculated from the transition matrix, state of node feature at step  $t$  and AA type  $x^{\mathrm{aa}}$ . The  $x^{\mathrm{aa}}$  is the sample of the denoising network prediction  $\hat{p}(\boldsymbol{x}^{\mathrm{aa}})$ .
Figure 2: The middle two panels depict the transition probability of Leucine (L) from  $t = 0$  to  $T$ . Both the uniform and BLOSUM start as Dirichlet distributions and become uniform at time  $T$ . As shown in the two side figures, while the uniform matrix evenly disperses L's probability to other AAs over time, BLOSUM favors AAs similar to L.
3.2 Prior Distribution from Protein Observations
3.2.1 Markov Transition Matrices
The transition matrix serves as a guide for a discrete diffusion model, facilitating transitions between the states by providing the probability of moving from the current time step to the next. As it reflects the possibility from one AA type to another, this matrix plays a critical role in both the diffusion and generative processes. During the diffusion stage, the transition matrix is iteratively applied to the observed data, which evolves over time due to inherent noise. As diffusion time increases, the probability of the original AA type gradually decays, eventually converging towards a uniform distribution across all AA types. In the generative stage, the conditional probability  $p(\pmb{x}_{t-1}|\pmb{x}_t)$  is determined by both the model's prediction and the characteristics of the transition matrix  $Q$ , as described in Equation 2.
Given the biological specificity of AA substitutions, the transition probabilities between AAs are not uniformly distributed, making it illogical to define random directions for the generative or sampling process. As an alternative, the diffusion process could reflect evolutionary pressures by utilizing substitution scoring matrices that conserve protein functionality, structure, or stability in wild-type protein families. Formally, an AA substitution scoring matrix quantifies the rates at which various AAs in proteins are substituted by other AAs over time [44]. In this study, we employ the Blocks Substitution Matrix (BLOSUM) [12], which identifies conserved regions within proteins that are presumed to have greater functional relevance. Grounded in empirical observations of protein evolution, BLOSUM provides an estimate of the likelihood of substitutions between different AAs. We thus incorporate BLOSUM into both the diffusion and generative processes. Initially, the matrix is normalized into probabilities using the softmax function. Then, we use the normalized matrix  $\pmb{B}$  with different probability temperatures to control the noise scale of the diffusion process. Consequently, the transition matrix at time  $t$  is given by  $Q_{t} = B^{T}$ . By using this matrix to refine the transition probabilities, the generative space to be sampled is reduced effectively, thereby the model's predictions converge toward a meaningful subspace. See Figure 2 for a comparison of the transition matrix over time in random and BLOSUM cases.
3.2.2 Secondary Structure
Protein secondary structure refers to the local spatial arrangement of AA residues in a protein chain. The two most common types of protein secondary structure are alpha helices and beta sheets, which are stabilized by hydrogen bonds between backbone atoms. The secondary structure of a protein serves as a critical intermediary, bridging the gap between the AA sequence and the overall 3D conformation of the protein. In our study, we incorporate eight distinct types of secondary structures into AA nodes as conditions during the sampling process. This strategic approach effectively narrows down the exploration space of potential AA sequences. Specifically, we employ DSSP (Define Secondary Structure of Proteins) to predict the secondary structures of each AA and represent these structures using one-hot encoding. Our neural network takes the one-hot encoding as input and utilizes it to denoise the AA conditioned on it.
The imposition of motif conditions such as alpha helices and beta sheets on the search for AA sequences not only leads to a significant reduction in the sampling space of potential sequences, but also imparts biological implications for the generated protein sequence. By conditioning the sampling process of AA types on their corresponding secondary structure types, we guide the resulting protein
sequence towards acquiring not only the appropriate 3D structure with feasible thermal stability but also the capability to perform its intended function.
3.3 Equivariant Graph Denoising Network
Bio-molecules such as proteins and chemical compounds are structured in the 3-dimensional space, and it is vital for the model to predict the same binding complex no matter how the input proteins are positioned and oriented to encode a robust and expressive hidden representation. This property can be guaranteed by the rotation equivariance of the neural networks. A typical choice of such a network is an equivariant graph neural network [35]. We modify its SE(3)-equivariant neural layers to update representations for both nodes and edges, which reserves SO(3) rotation equivariance and E(3) translation invariance. At the  $l$ th layer, an Equivariant Graph Convolution (EGC) inputs a set of  $n$  hidden node embeddings  $\pmb{H}^{(l)} = \left\{\pmb{h}_1^{(l)}, \dots, \pmb{h}_n^{(l)}\right\}$  describing AA type and geometry properties, edge embedding  $\pmb{m}_{ij}^{(l)}$  with respect to connected nodes  $i$  and  $j$ ,  $\pmb{X}^{\mathrm{pos}} = \{\pmb{x}_1^{\mathrm{pos}}, \dots, \pmb{x}_n^{\mathrm{pos}}\}$  for node coordinates and  $t$  for time step embedding of diffusion model. The target of a modified EGC layer is to update hidden representations  $\pmb{H}^{(l + 1)}$  for nodes and  $\pmb{M}^{(l + 1)}$  for edges. Concisely,  $\pmb{H}^{(l + 1)}, \pmb{M}^{(l + 1)} = \operatorname{EGC}\left[\pmb{H}^{(l)}, \pmb{X}^{\mathrm{pos}}, \pmb{M}^{(l)}, t\right]$ . To achieve this, an EGC layer defines
$$
\boldsymbol {m} _ {i j} ^ {(l + 1)} = \phi_ {e} \left(\mathbf {h} _ {i} ^ {(l)}, \mathbf {h} _ {j} ^ {(l)}, \left\| \mathbf {x} _ {i} ^ {(l)} - \mathbf {x} _ {j} ^ {(l)} \right\| ^ {2}, \boldsymbol {m} _ {i j} ^ {(l)}\right)
$$
$$
\boldsymbol {x} _ {i} ^ {(l + 1)} = \mathbf {x} _ {i} ^ {(l)} + \frac {1}{n} \sum_ {j \neq i} \left(\mathbf {x} _ {i} ^ {(l)} - \mathbf {x} _ {j} ^ {(l)}\right) \phi_ {x} \left(\mathbf {m} _ {i j} ^ {(l + 1)}\right) \tag {4}
$$
$$
\boldsymbol {h} _ {i} ^ {(l + 1)} = \phi_ {h} \big (\mathbf {h} _ {i} ^ {(l)}, \sum_ {j \neq i} \mathbf {m} _ {i j} ^ {(l + 1)} \big),
$$