SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\phi_e, \phi_h$  are the edge and node propagation operations, respectively. The  $\phi_x$  is an additional operation that projects the vector edge embedding  $m_{ij}$  to a scalar. The modified EGC layer preserves equivariance to rotations and translations on the set of 3D node coordinates  $\mathbf{X}^{\mathrm{pos}}$  and performs invariance to permutations on the nodes set identical to any other GNNs.
3.4 DDIM Sampling Process
A significant drawback of diffusion models lies in the speed of generation process, which is typically characterized by numerous incremental steps and can be quite slow. Deterministic Denoising Implicit Models (DDIM) [39] are frequently utilized to counter this issue in continuous variable diffusion generative models. DDIM operates on a non-Markovian forward diffusion process, consistently conditioning on the input rather than the previous step. By setting the noise variance on each step to 0, the reverse generative process becomes entirely deterministic, given an initial prior sample.
Similarly, since we possess the closed form of generative probability  $p_{\theta}(\pmb{x}_{t-1}|\pmb{x}_t)$  in terms of a predicted  $\pmb{x}^{\mathrm{aa}}$  and the posterior distribution  $p(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}^{\mathrm{aa}})$ , we can also render the generative model deterministic by controlling the sampling temperature of  $p(\pmb{x}^{\mathrm{aa}}|\pmb{x}_t)$ . Consequently, we can define the multi-step generative process by
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} _ {t}\right) \propto \left(\sum_ {\hat {\boldsymbol {x}} ^ {\mathrm {a a}}} q \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) \hat {p} \left(\boldsymbol {x} ^ {\mathrm {a a}} \mid \boldsymbol {x} _ {t}\right)\right) ^ {T} \tag {5}
$$
where the temperature  $T$  controls whether it is deterministic or stochastic, and the multi-step posterior distribution is
$$
q \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) = \operatorname {C a t} \left(\boldsymbol {x} _ {t - k} \left| \frac {\boldsymbol {x} _ {t} Q _ {t} ^ {\top} \cdots Q _ {t - k} ^ {\top} \odot \boldsymbol {x} ^ {\mathrm {a a}} \bar {Q} _ {t - k}}{\boldsymbol {x} ^ {\mathrm {a a}} \bar {Q} _ {t} \boldsymbol {x} _ {t} ^ {\top}}\right)\right). \tag {6}
$$
We validate our GRADE-IF on recovering native protein sequences in CATH [30]. The performance is mainly compared with structure-aware SOTA models. The implementations for the main
Table 1: Recovery rate performance of CATH on zero-shot models.
Model	Perplexity ↓	Recovery Rate % ↑	CATH version
Short	Single-chain	All	Short	Single-chain	All	4.2	4.3
STRUCTGNN [18]	8.29	8.74	6.40	29.44	28.26	35.91	✓
GRAPHTRANS [18]	8.39	8.83	6.63	28.14	28.46	35.82	✓
GCA [41]	7.09	7.49	6.05	32.62	31.10	37.64	✓
GVP [20]	7.23	7.84	5.36	30.60	28.95	39.47	✓
GVP-large [17]	7.68	6.12	6.17	32.6	39.4	39.2	✓
ALPHADESIGN [9]	7.32	7.63	6.30	34.16	32.66	41.31	✓
ESM-IF1 [17]	8.18	6.33	6.44	31.3	38.5	38.3	✓
PROTEINMPNN [6]	6.21	6.68	4.57	36.35	34.43	49.87	✓
PIFOLD [10]	6.04	6.31	4.55	39.84	38.53	51.66	✓
GRADE-IF	5.49	6.21	4.35	45.27	42.77	52.21	✓
Figure 3: Recovery rate on core and surface residues and different secondary structure
algorithms (see Appendix D) at https://github.com/ykiiiii/GraDe_IF are programmed with PyTorch-Geometric (ver 2.2.0) and PyTorch (ver 1.12.1) and executed on an NVIDIA Tesla V100 GPU with 5,120 CUDA cores and 32GB HBM2 installed on an HPC cluster.
4.1 Experimental Protocol
Training Setup We employ CATH v4.2.0-based partitioning as conducted by GRAPHTRANS [18] and GVP [20]. Proteins are categorized based on CATH topology classification, leading to a division of 18,024 proteins for training, 608 for validation, and 1,120 for testing. To evaluate the generative quality of different proteins, we test our model across three distinct categories: short, single-chain, and all proteins. The short category includes proteins with sequence lengths shorter than 100. The single-chain category encompasses proteins composed of a single chain. In addition, the total time step of the diffusion model is configured as 500, adhering to a cosine schedule for noise [27]. For the denoising network, we implement six stacked EGNN blocks, each possessing a hidden dimension of 128. Our model undergoes training for default of 200 epochs, making use of the Adam optimizer. A batch size of 64 and a learning rate of 0.0005 are applied during training. Moreover, to prevent overfitting, we incorporate a dropout rate of 0.1 into our model's architecture.
Evaluation Metric The quality of recovered protein sequences is quantified by perplexity and recovery rate. The former measures how well the model's predicted AA probabilities match the actual AA at each position in the sequence. A lower perplexity indicates a better fit of the model to the data. The recovery rate assesses the model's ability to recover the correct AA sequence given the protein's 3D structure. It is typically computed as the proportion of AAs in the predicted sequence that matches the original sequence. A higher recovery rate indicates a better capability of the model to predict the original sequence from the structure.
4.2 Inverse Folding
Table 1 compares GRADE-IF's performance on recovering proteins in CATH, with the last two columns indicating the training dataset of each baseline method. To generate high-confidence
Table 2: Numerical comparison between generated sequence structure and the native structure.
Method	Success	TM score	avg pLDDT	avg RMSD
PIFOLD	85	0.80 ± 0.22	0.84 ± 0.15	1.67 ± 0.99
PROTEINMPNN	94	0.86 ± 0.16	0.89 ± 0.10	1.36 ± 0.81
GRADE-IF	94	0.86 ± 0.17	0.86 ± 0.08	1.47 ± 0.82
Recovery Rate: 0.612 RMSD: 0.915 avg. pLDDT: 0.841
Recovery Rate: 0.547 RMSD: 1.021 avg. pLDDT: 0.835
Figure 4: Folding comparison of GRADE-IF-generated sequences and the native protein (in nude).
Recovery Rate: 0.540 RMSD: 0.999 avg. pLDDT: 0.830
sequences, GRADE-IF integrates out uncertainties in the prior by approximating the probability  $p(\pmb{x}^{\mathrm{aa}}) \approx \sum_{i=1}^{N} p(\pmb{x}^{\mathrm{aa}} | \pmb{x}_T^i) p(\pmb{x}_T^i)$ . Notably, we observed an improvement of  $4.2\%$  and  $5.4\%$  in the recovery rate for single-chain proteins and short sequences, respectively. We also conducted evaluations on different datasets (Appendix E) and ablation conditions (Appendix F).
Upon subdividing the recovery performance based on buried and surface AAs, we find that the more conserved core residues exhibit a higher native sequence recovery rate. In contrast, the active surface AAs demonstrate a lower sequence recovery rate. Figure 3 examines AA conservation by Solvent Accessible Surface Area (SASA) (with  $\mathrm{SASA} < 0.25$  indicating internal AAs) and contact number (with the number of neighboring AAs within  $8\AA$  in 3D space) [11]. The recovery rate of internal residues significantly exceeds that of external residues across all three protein sequence classes, with the recovery rate increasing in conjunction with the contact number. We also present the recovery rate for different secondary structures, where we achieve high recovery for the majority of secondary structures, with the exception of a minor 5-turn helix structure that occurs infrequently.
4.3 Foldability
We extend our investigation to the foldability of sequences generated at various sequence recovery rates. We fold generated protein sequences (by GRADE-IF) with ALPHAFOLD2 and align them with the crystal structure to compare their closeness in Figure 4 (PDB ID: 3FKF). All generated sequences are nearly identical to the native one with their RMSD  $\sim 1\AA$  over 139 AAs, which is lower than the resolution of the crystal structure at  $2.2\AA$ . We have also folded the native sequence by ALPHAFOLD2, which yields an average pLDDT of 0.91. In comparison, the average pLDDT scores of the generated sequences are 0.835, underscoring the reliability of their folded structures. In conjunction with the evidence presented in Figure 3 which indicates our method's superior performance in generating more identical results within conserved regions, we confidently posit that GRADE-IF can generate biologically plausible novel sequences for given protein structures (See Appendix G).
The numerical investigation is reported in Table 2, where we pick the first 100 structures (ordered in alphabetical order by their PDB ID) from the test dataset and compare the performance of GRADE-IF with PROTEINMPNN and PiFOLD. We follow [48] and define the quality of a novel sequence by the TM score between its ALPHA FOLD2-folded structure and the native structure, with an above-0.5 score indicating the design is successful. Overall, our GRADE-IF exhibits high pLDDT, low RMSD, and high foldability. There are several proteins that face challenges with both GRADE-IF and baseline methods for folding with a high TM score, i.e., 1BCT, 1BHA, and 1CYU, whose structural determination is based on NMR, an experimental technique that analyzes protein structure in a buffer solution. Due to the presence of multiple structures for a single protein in NMR studies, it is reasonable for folding tools to assign low foldability scores.
Table 3: Numerical comparison on diversity and recovery rate
Method	low recovery rate	medium recovery rate	high recovery rate
diversity	recovery	diversity	recovery	diversity	recovery
PIFOLD	0.37	0.47	0.25	0.50	0.21	0.50
PROTEINMPNN	0.51	0.42	0.27	0.46	0.26	0.52
GRADE-IF	0.61	0.33	0.54	0.47	0.25	0.53
step:50   
avg.pLDDT:0.867   
TM score: 0.910   
diversity:0.259
step:10 avg.pLDDT:0.826 TM score:0.879 diversity:0.446
step: 5  
avg. pLDDT: 0.811  
TM score: 0.861  
diversity: 0.547
Figure 5: Folding structures of generated protein sequences with different steps.
step:1 avg.pLDDT:0.778 TM score: 0.820 diversity:0.618
4.4 Diversity
Given the one-to-many mapping relationship between a protein structure and its potential sequences, an inverse folding model must be capable of generating diverse protein sequences for a fixed backbone structure. To investigate the diversity of GRADE-IF in comparison to baseline methods, we define diversity as 1 - self-similarity, where higher diversity indicates a model's proficiency in generating distinct sequences. We employ PIFOLD, PROTEINMPNN, and GRADE-IF to generate new proteins for the test dataset at low, medium, and high recovery levels. We vary the temperature for the former two baseline methods within the range \{0.5, 0.1, 0.0001\} and adjust the sample step for GRADE-IF from \{1, 10, 50\}. The average performance from sampling 10 sequences is summarized in Table 3, revealing comparable results among the three models in both recovery rate and diversity. In general, increasing the probability temperature for PIFOLD and PROTEINMPNN (or decreasing the sample step in GRADE-IF) leads to a reduction in uncertainty and more reliable predictions, resulting in higher diversities and lower recovery rates for all three models. While all three methods demonstrate the ability to recover sequences at a similar level, particularly with high probability temperatures (or low sample steps), GRADE-IF produces significantly more diverse results when the step size is minimized. Our findings demonstrate that GRADE-IF can generate sequences with a  $60\%$  difference in each sample, whereas PIFOLD and PROTEINMPNN achieve diversity rates below  $50\%$ .
We next explore the foldability of these highly diverse protein sequences designed by GRADE-IF. Figure 5 compares the generated proteins (folded by ALPHAFOLD2) with the crystal structure (PDB ID: 3FKF). We vary the step size over a range of 1, 5, 10, 50, generating 10 sequences per step size to calculate the average pLDDT, TM score, and diversity. For simplicity, we visualize the first structure in the figure. Decreasing the step size results in the generation of more diverse sequences. Consequently, there is a slight reduction in both pLDDT and TM scores. However, they consistently remain at a considerably high level, with both metrics approaching 0.8. This reduction can, in part, be attributed to the increased diversity of the sequences, as ALPHAFOLD2 heavily relies on the MSA sequences. It is expected that more dissimilar sequences would produce a more diverse MSA. Remarkably, when step = 1, the sequence diversity exceeds 0.6, indicating that the generated sequences share an approximately 0.3 sequence similarity compared to the wild-type template protein sequence. This suggests the generation of protein sequences from a substantially distinct protein family when both pLDDT and TM scores continue to exhibit a high degree of confidence.
Deep Learning models for protein sequence design Self-supervised models have emerged as a pivotal tool in the field of computational biology, providing a robust method for training extensive
protein sequences for representation learning. These models are typically divided into two categories: structure-based generative models and sequence-based generative models. The former approaches protein design by formulating the problem of fixed-backbone protein design as a conditional sequence generation problem. They predict node labels, which represent AA types, with invariant or equivariant graph neural networks [2, 17, 18, 20, 40, 42]. Alternatively, the latter sequence-based generative models draw parallels between protein sequences and natural language processing. They employ attention-based methods to infer residue-wise relationships within the protein structure. These methods typically recover protein sequences autoregressively conditioned on the last inferred AA [25, 29, 36], or employing a BERT-style generative framework with masked language modeling objectives and enable the model to predict missing or masked parts of the protein sequence [23, 26, 33, 46].
Denoising Diffusion models The Diffusion Generative Model, initially introduced by Sohl-Dickstein et al. [38] and further developed by Ho et al. [13], has emerged as a potent instrument for a myriad of generative tasks in continuous time spaces. Its applications span diverse domains, from image synthesis [34] to audio generation [49], and it has also found utility in the creation of high-quality animations [14], the generation of realistic 3D objects [24], and drug design [5, 43]. Discrete adaptations of the diffusion model, on the other hand, have demonstrated efficacy in a variety of contexts, including but not limited to, text generation [3], image segmentation [15], and graph generation [16, 47]. Two distinct strategies have been proposed to establish a discrete variable diffusion process. The first approach involves the transformation of categorical data into a continuous space and then applying Gaussian diffusion [4, 16]. The alternative strategy is to define the diffusion process directly on the categorical data, an approach notably utilized in developing the D3PM model for text generation [3]. D3PM has been further extended to graph generation, facilitating the joint generation of node features and graph structure [47].
Deep learning approaches have striven to address a multitude of critical issues in bioengineering, such as protein folding, rigid-body docking, and property prediction. However, only a few methods have successfully generated diverse sequences for fixed backbones. In this study, we offered a viable solution by developing a denoising diffusion model to generate plausible protein sequences for a predetermined backbone structure. Our method, referred to as GRADE-IF, leverages substitution matrices for both diffusion and sampling processes, thereby exploring a practical search space for defining proteins. The iterative denoising process is predicated on the protein backbone revealing both the secondary and tertiary structure. The 3D geometry is analyzed by a modified equivariant graph neural network, which applies roto-translation equivariance to protein graphs without the necessity for intensive data augmentation. Given a protein backbone, our method successfully generated a diverse set of protein sequences, demonstrating a significant recovery rate. Importantly, these newly generated sequences are generally biologically meaningful, preserving more natural designs in the protein's conserved regions and demonstrating a high likelihood of folding back into a structure highly similar to the native protein. The design of novel proteins with desired structural and functional characteristics is of paramount importance in the biotechnology and pharmaceutical industries, where such proteins can serve diverse purposes, ranging from targeted drug delivery to enzyme design for industrial applications. Additionally, understanding how varied sequences can yield identical structures propels the exploration of protein folding principles, thereby helping to decipher the rules that govern protein folding and misfolding. Furthermore, resolving the inverse folding problem allows the identification of different sequences that fold into the same structure, shedding light on the evolutionary history of proteins by enhancing our understanding of how proteins have evolved and diversified over time while preserving their functions.
We thank Lirong Zheng for providing insightful discussions on molecular biology. Bingxin Zhou acknowledges support from the National Natural Science Foundation of China (62302291). Yu Guang Wang acknowledges support from the National Natural Science Foundation of China (62172370), Shanghai Key Projects (23440790200) and (2021SHZDZX0102). This project was undertaken with the assistance of computational resources and services from the National Computational Infrastructure (NCI), which is supported by the Australian Government.
[1] Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew J O'Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, P Douglas Renfrew, Vikram K Mulligan, Kalli Kappel, et al. The Rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation, 13(6):3031-3048, 2017.
[2] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019, 2022.
[3] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-17993, 2021.
[4] Ting Chen, Ruixiang ZHANG, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. In The Eleventh International Conference on Learning Representations, 2023.
[5] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference on Learning Representations, 2023.
[6] Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science, 378(6615):49-56, 2022.
[7] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Wang Yu, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. ProtTrans: Towards cracking the language of lives code through self-supervised deep learning and high performance computing. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
[8] Octavian-Eugen Ganea, Xinyuan Huang, Charlotte Bunne, Yatao Bian, Regina Barzilay, Tommi S Jaakkola, and Andreas Krause. Independent SE(3)-equivariant models for end-to-end rigid protein docking. In International Conference on Learning Representations, 2021.
[9] Zhangyang Gao, Cheng Tan, and Stan Z Li. Alphadesign: A graph protein design method and benchmark on alphafolddb. arXiv preprint arXiv:2202.01079, 2022.
[10] Zhangyang Gao, Cheng Tan, and Stan Z. Li. Pifold: Toward effective and efficient protein inverse folding. In International Conference on Learning Representations, 2023.
[11] Hai'e Gong, Haicang Zhang, Jianwei Zhu, Chao Wang, Shiwei Sun, Wei-Mou Zheng, and Dongbo Bu. Improving prediction of burial state of residues by exploiting correlation among residues. BMC bioinformatics, 18(3):165-175, 2017.
[12] Steven Henikoff and Jorja G Henikoff. Amino acid substitution matrices from protein blocks. Proceedings of the National Academy of Sciences, 89(22):10915-10919, 1992.
[13] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[48] Kevin E Wu, Kevin K Yang, Rianne van den Berg, James Y Zou, Alex X Lu, and Ava P Amini. Protein structure generation via folding diffusion. arXiv preprint arXiv:2209.15611, 2022.
[49] Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for text-to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 2023.
[50] Zaixiang Zheng, Yifan Deng, Dongyu Xue, Yi Zhou, Fei YE, and Quanquan Gu. Structure-informed language models are protein designers. In International Conference on Machine Learning, 2023.
[51] Bingxin Zhou, Outongyi Lv, Kai Yi, Xinye Xiong, Pan Tan, Liang Hong, and Yu Guang Wang. Accurate and definite mutational effect prediction with lightweight equivariant graph neural networks. arXiv:2304.08299, 2023.
[52] Bingxin Zhou, Lirong Zheng, Banghao Wu, Kai Yi, Bozitao Zhong, Pietro Lio, and Liang Hong. Conditional protein denoising diffusion generates programmable endonucleases. bioRxiv, pages 2023-08, 2023.
Broader Impact We have developed a generative model rooted in the diffusion denoising paradigm, specifically tailored to the context of protein inverse folding. As with any other generative models, it is capable of generating de novo content (protein sequences) under specified conditions (e.g., protein tertiary structure). While this method holds substantial potential for facilitating scientific research and biological discoveries, its misuse could pose potential risks to human society. For instance, in theory, it possesses the capacity to generate novel viral protein sequences with enhanced functionalities. To mitigate this potential risk, one approach could be to confine the training dataset for the model to proteins derived from prokaryotes and/or eukaryotes, thereby excluding viral proteins. Although this strategy may to some extent compromise the overall performance and generalizability of the trained model, it also curtails the risk of misuse of the model by limiting the understanding and analysis of viral protein construction.
Limitations The conditions imposed on the sampling process gently guide the generated protein sequences. However, in certain scenarios, stringent restrictions may be necessary to produce a functional protein. Secondary structure, as a living example, actively contributes to the protein's functionality. For instance, transmembrane  $\alpha$ -helices play essential roles in protein functions, such as passing ions or other molecules and transmitting a signal across the membrane. Moreover, the current zero-shot model is trained on a general protein database. For specific downstream applications, such as generating new sequences for a particular protein or protein family, it may necessitate the incorporation of auxiliary modules or the modification of training procedures to yield more fitting sequences.
B Non-Markovian Forward Process
We give the derivation of posterior distribution  $q(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}^{\mathrm{aa}})$  for generative process from step to step. The proof relies on the Bayes rule, Markov property, and the pre-defined transition matrix for AAs.
Proposition 1. For  $q(\pmb{x}_{t-1} \mid \pmb{x}_t, \pmb{x}^{\mathrm{aa}})$  defined in Eq 3, we have