SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
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} \Big | \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).
$$
Proof. By Bayes rules, we can expand the original equation  $q(\pmb{x}_{t-1} \mid \pmb{x}_t, \pmb{x}^{\mathrm{aa}})$  to
$$
q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) = \frac {q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1} , \boldsymbol {x} ^ {\mathrm {a a}}\right) q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}{q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)} = \frac {q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}{q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}.
$$
As pre-defined diffusion process, we get  $q(\pmb{x}_t \mid \pmb{x}^{\mathrm{aa}}) = \pmb{x}^{\mathrm{aa}}\bar{Q}_t$ , and  $q(\pmb{x}_{t - 1} \mid \pmb{x}^{\mathrm{aa}}) = \pmb{x}^{\mathrm{aa}}\bar{Q}_{t - 1}$ .
For the term of  $q(\pmb{x}_t \mid \pmb{x}_{t-1}, \pmb{x}^{\mathrm{aa}})$  by Bayes rule and Markov property, we have
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}, \boldsymbol {x} ^ {\mathrm {a a}}\right) = q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) \propto q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) \pi \left(\boldsymbol {x} _ {t}\right) \propto \boldsymbol {x} _ {t} Q _ {t} ^ {\top} \odot \pi \left(\boldsymbol {x} _ {t}\right)
$$
where the normalizing constant is  $\sum_{\pmb{x}_{t - 1}}\pmb{x}_tQ_t^\top \odot \pi (\pmb{x}_t) = (\pmb{x}_t\sum_{\pmb{x}_{t - 1}}Q_t^\top)\odot \pi (\pmb{x}_t) = \pmb{x}_t\odot \pi (\pmb{x}_t)$
Then  $q(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}, \boldsymbol{x}^{\mathrm{aa}}) = \frac{\boldsymbol{x}_t Q_t^\top}{\boldsymbol{x}_t}$ , and the posterior distribution is:
$$
q \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} ^ {\mathrm {a a}}, \boldsymbol {x} _ {t}\right) = \operatorname {C a t} \left(\boldsymbol {x} _ {t - 1} \Big | \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).
$$
The following gives the derivation for the discrete DDIM which accelerates the generative process.
Proposition 2. For  $q(\pmb{x}_{t - k} \mid \pmb{x}_t, \pmb{x}^{\mathrm{aa}})$  defined in Eq 6,
$$
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} \Big | \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).
$$
Proof. By Bayes rules, we can expand the original equation  $q(\boldsymbol{x}_{t - k} \mid \boldsymbol{x}_t, \boldsymbol{x}^{\mathrm{aa}})$  to
$$
q \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} _ {t}, \boldsymbol {x} ^ {\mathrm {a a}}\right) = \frac {q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - k} , \boldsymbol {x} ^ {\mathrm {a a}}\right) q \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}{q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)} = \frac {q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - k}\right) q \left(\boldsymbol {x} _ {t - k} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}{q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} ^ {\mathrm {a a}}\right)}.
$$
As pre-defined diffusion process, we get  $q(\pmb{x}_t \mid \pmb{x}^{\mathrm{aa}}) = \pmb{x}^{\mathrm{aa}}\bar{Q}_t$ , and  $q(\pmb{x}_{t - 1} \mid \pmb{x}^{\mathrm{aa}}) = \pmb{x}^{\mathrm{aa}}\bar{Q}_{t - k}$ .
Similarly with  $q(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-1}, \boldsymbol{x}^{\mathrm{aa}})$  in Proposition 1,  $q(\boldsymbol{x}_t \mid \boldsymbol{x}_{t-k}, \boldsymbol{x}^{\mathrm{aa}}) = \frac{\boldsymbol{x}_t Q_t^\top \cdots Q_{t-k}^\top}{\boldsymbol{x}_t}$  and the posterior is
C Graph Representation of Folded Proteins
The geometry of proteins suggests higher-level structures and topological relationships, which are vital to protein functionality. For a given protein, we create a  $k$ -nearest neighbor ( $k\mathrm{NN}$ ) graph  $\mathcal{G} = (X, E)$  to describe its physiochemical and geometric properties with nodes representing AAs by  $X \in \mathbb{R}^{39}$  node attributes with 20-dim AA type encoder, 16-dim AA properties, and 3-dim AA positions. The undirected edge connections are formulated via a  $k\mathrm{NN}$ -graph with cutoff. In other words, each node is connected to up to  $k$  other nodes in the graph that has the smallest Euclidean distance over other nodes and the distance is smaller than a certain cutoff ( $e.g., 30\AA$ ). Edge attributes are defined for connected node pairs. For instance, if node  $i$  and  $j$  are connected to each other, their relationship will be described by  $\pmb{E}_{ij} = \pmb{E}_{ji} \in \mathbb{R}^{93}$ .
The AA types are one-hot encoded to 20 binary values by  $X^{\mathrm{aa}}$ . On top of it, the properties of AAs and AAs' local environment are described by  $X^{\mathrm{prop}}$ , including the normalized crystallographic B-factor, solvent-accessible surface area (SASA), normalized surface-aware node features, dihedral angles of backbone atoms, and 3D positions. SASA measures the level of exposure of an AA to solvent in a protein by a scalar value, which provides an important indicator of active sites of proteins to locate whether a residue is on the surface of the protein. Both B-factor and SASA are standardized with AA-wise mean and standard deviation on the associate attribute. Surface-aware features [8] of an AA is non-linear projections to the weighted average distance of the central AA to its one-hop neighbors  $i' \in \mathcal{N}_i$ , i.e.,
$$
\rho \left(\mathbf {x} _ {i}; \lambda\right) = \frac {\left\| \sum_ {i ^ {\prime} \in \mathcal {N} _ {i}} w _ {i , i ^ {\prime} , \lambda} \left(\boldsymbol {X} ^ {\text {p o s} , i} - \boldsymbol {X} ^ {\text {p o s} , i ^ {\prime}}\right) \right\|}{\sum_ {i ^ {\prime} \in \mathcal {N} _ {i}} w _ {i , i ^ {\prime} , \lambda} \left\| \boldsymbol {X} ^ {\text {p o s} , i} - \boldsymbol {X} ^ {\text {p o s} , i ^ {\prime}} \right\|},
$$
where the weights are defined by
$$
w _ {i, i ^ {\prime}, \lambda} = \frac {\exp \left(- \left\| \boldsymbol {X} _ {\mathrm {p o s} , i} - \boldsymbol {X} _ {\mathrm {p o s} , i ^ {\prime}} \right\| ^ {2} / \lambda\right)}{\sum_ {i ^ {\prime} \in \mathcal {N} _ {i}} \exp \left(- \left\| \boldsymbol {X} _ {\mathrm {p o s} , i} - \boldsymbol {X} _ {\mathrm {p o s} , i ^ {\prime}} \right\| ^ {2} / \lambda\right)}
$$
with  $\lambda \in \{1,2,5,10,30\}$ . The  $X^{\mathrm{pos},i} \in \mathbb{R}^3$  denotes the 3D coordinates of the  $i$ th residue, which is represented by the position of  $\alpha$ -carbon. We also use the backbone atom positions to define the spatial conformation of each AA in the protein chain with trigonometric values of dihedral angles  $\{\sin, \cos\} \circ \{\phi_i, \psi_i, \omega_i\}$ .
Edge attributes  $E \in \mathbb{R}^{93}$ , on the other hand, include kernel-based distances, relative spatial positions, and relative sequential distances for pairwise distance characterization. For two connected residues  $i$  and  $j$ , the kernel-based distance between them is projected by Gaussian radial basis functions (RBF) of  $\exp \left\{\frac{\|\pmb{x}_j - \pmb{x}_i\|^2}{2\sigma_r^2}\right\}$  with  $r = 1,2,\dots,R$ . A total number of 15 distinct distance-based features are created with  $\sigma_r = \{1.5^k \mid k = 0,1,2,\dots,14\}$ . Next, local frames [8] are created from the corresponding residues' heavy atoms positions to define 12 relative positions. They represent local fine-grained relations between AAs and the rigid property of how the two residues interact with each other. Finally, the residues' sequential relationship is encoded with 66 binary features by their relative position  $d_{i,j} = |s_i - s_j|$ , where  $s_i$  and  $s_j$  are the absolute positions of the two nodes in the AA chain [51]. We further define a binary contact signal [18] to indicate whether two residues contact in the space, i.e., the Euclidean distance  $\| C\alpha_i - C\alpha_j \| < 8$ .
D Training and Inference
In this section, we elucidate the training and inference methodologies implemented in the diffusion generative model. As shown in Algorithm 1, training commences with a random sampling of a time scale  $t$  from a uniform distribution between 1 and  $T$ . Subsequently, we calculate the noise posterior and integrate noise as dictated by its respective distribution. We then utilize an equivariant graph neural network for denoising predictions, using both the noisy amino acid and other properties as node features, and leveraging the graph structure for geometric information. This results in the model outputting the denoised amino acid type. Ultimately, the cross-entropy loss is computed between the predicted and original amino acid types, providing a parameter for optimizing the neural network.
Algorithm 1 Training
1: Input: A graph  $\mathcal{G} = \{\pmb {X},\pmb {E}\}$
2: Sample  $t\sim \mathcal{U}(1,T)$
3: Compute  $q(\mathbf{X}_t|\mathbf{X}^{\mathrm{aa}}) = \mathbf{X}^{\mathrm{aa}}\bar{Q}_t$
4: Sample noisy  $\mathbf{X}_t \sim q(\mathbf{X}_t | \mathbf{X}^{\mathrm{aa}})$
5: Forward pass:  $\hat{p}(\mathbf{X}^{\mathrm{aa}}) = f_{\theta}(\mathbf{X}_t, \mathbf{E}, t, ss)$
6: Compute cross-entropy loss:  $L = L_{\mathrm{CE}}(\hat{p}(\mathbf{X}^{\mathrm{aa}}), \mathbf{X})$
7: Compute the gradient and optimize denoise network  $f_{\theta}$
Upon completing the training, we are capable of sampling data using the neural network and the posterior distribution  $p(\pmb{x}_{t-1}|\pmb{x}_t, \pmb{x}^{\mathrm{aa}})$ . As delineated in the algorithm, we initially sample an amino acid uniformly from 20 classes, then employ our neural network to denoise  $X^{\mathrm{aa}}$  from time  $t$ . From here, we can calculate the forward probability utilizing the model output and the posterior distribution. Through iterative processing, the ultimate model sample closely approximates the original data distribution. We can also partially generate amino acids given some known AA within the structure. Analogous to an inpainting task, at each step, we can manually adjust the prediction at known positions to the known amino acid type, subsequently introducing noise, as illustrated in the Algorithm 4.
Algorithm 2 Sampling (DDPM)
1: Sample from uniformly prior  $X_{T}\sim p(\pmb{X}_{T})$
2: for  $t$  in  $\{T, T - 1, \dots, 1\}$  do
3: Predict  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t)$  by neural network  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t) = f_\theta(\mathbf{X}_t, \mathbf{E}, t, ss)$
4: Compute  $p_{\theta}(\boldsymbol{X}_{t-1}|\boldsymbol{X}_t) = \sum_{\hat{\boldsymbol{X}}^{\mathrm{aa}}} q(\boldsymbol{X}_{t-1}|\boldsymbol{X}_t, \hat{\boldsymbol{X}}^{\mathrm{aa}})\hat{p}(\boldsymbol{X}^{\mathrm{aa}}|\boldsymbol{X}_t)$
5: Sample  $X_{t - 1}\sim p_{\theta}(X_{t - 1}|X_t)$
6: end for
7: Sample  $\pmb{X}^{\mathrm{aa}}\sim p_{\theta}(\pmb{X}^{\mathrm{aa}}|\pmb{X}_1)$
Algorithm 3 Sampling (DDIM)
1: Sample from uniformly prior  $X_{T}\sim p(X_{T})$
2: for  $t$  in  $\{T, T - k, \dots, 1\}$  do
3: Predict  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t)$  by neural network  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t) = f_\theta(\mathbf{X}_t, \mathbf{E}, t, ss)$
4: Compute  $p_{\theta}(\boldsymbol{X}_{t - k}|\boldsymbol{X}_t) = \sum_{\hat{\boldsymbol{X}}^{\mathrm{aa}}}q(\boldsymbol{X}_{t - k}|\boldsymbol{X}_t,\hat{\boldsymbol{X}}^{\mathrm{aa}})\hat{p} (\boldsymbol{X}^{\mathrm{aa}}|\boldsymbol{X}_t)$
5: Sample  $X_{t - k}\sim p_{\theta}(X_{t - k}|X_t)$
6: end for
7: Sample  $\mathbf{X}^{\mathrm{aa}} \sim p_{\theta}(\mathbf{X}^{\mathrm{aa}} | \mathbf{X}_1)$
Algorithm 4 Partial Sampling
1: Input Mask  $M$  indicate which position is fixed
2: Sample from uniformly prior  $X_{T}\sim p(\pmb{X}_{T})$
3: for  $t$  in  $\{T, T - k, \dots, 1\}$  do
4: Predict  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t)$  by neural network  $\hat{p}(\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t) = f_\theta(\mathbf{X}_t, \mathbf{E}, t, ss)$
5: Compute  $\hat{p} (\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t) = p(\mathbf{X}^{\mathrm{aa}})\odot M + \hat{p} (\mathbf{X}^{\mathrm{aa}}|\mathbf{X}_t)\odot (1 - M)$
6: Compute  $p_{\theta}(\boldsymbol{X}_{t - k}|\boldsymbol{X}_t) = \sum_{\hat{\boldsymbol{X}}^{\mathrm{aa}}}q(\boldsymbol{X}_{t - k}|\boldsymbol{X}_t,\hat{\boldsymbol{X}}^{\mathrm{aa}})\hat{p} (\boldsymbol{X}^{\mathrm{aa}}|\boldsymbol{X}_t)$
7: Sample  $\mathbf{X}_{t - k} \sim p_{\theta}(\mathbf{X}_{t - k} | \mathbf{X}_t)$
8: end for
9: Sample  $\mathbf{X}^{\mathrm{aa}} \sim p_{\theta}(\mathbf{X}^{\mathrm{aa}} | \mathbf{X}_1)$
E Inverse Folding Performance on TS50 and T500
In addition to the CATH dataset, we also evaluated our model using the TS50 and T500 datasets. These datasets were introduced by DenseCPD [32], encompassing 9888 structures for training, and two distinct test datasets comprising 50 (TS50) and 500 (T500) test datasets, respectively. The same preprocessing steps applied to the CATH dataset were utilized here. The denoising network comprises six sequentially arranged EGNN blocks, each boasting a hidden dimension of 256. Our model's performance, outlined in Table 4, achieved an accuracy of  $61.22\%$  on T500, and  $56.32\%$  on TS50, respectively.
Table 4: Recovery rate performance of TS50 and T500 on zero-shot models.
Model	TS50	T500
Perplexity ↓	Recovery ↑	Perplexity ↓	Recovery ↑
STRUCTGNN [18]	5.40	43.89	4.98	45.69
GRAPHTRANS [18]	5.60	42.20	5.16	44.66
GVP [20]	4.71	44.14	4.20	49.14
GCA [41]	5.09	47.02	4.72	47.74
ALPHADESIGN [9]	5.25	48.36	4.93	49.23
PROTEINMPNN [6]	3.93	54.43	3.53	58.08
PIFOLD [10]	3.86	58.72	3.44	60.42
GRADE-IF(ours)	3.71	56.32	3.23	61.22
We conducted ablation studies to assess the impact of various factors on our model's performance. These elements encompassed the selection of the transition matrix (uniform versus BLOSUM), the integration of secondary structure embeddings in the denoising procedure, and the function of the equivariant neural network. As demonstrated in Figure 6, incorporating equivariance into the denoising neural network substantially enhances the model's performance. Given that the placement of protein structures in space can be arbitrary, considering symmetry in the denoising neural network helps to mitigate disturbances. Moreover, we found that including secondary structure as auxiliary information lessens uncertainty and improves recovery. Lastly, utilizing the BLOSUM matrix as the noise transition matrix boosted the recovery rate by  $2\%$ , highlighting the benefits of infusing biological information into the diffusion and generative processes. This approach reduces sample variance and substantially benefits overall model performance.
In our sampling procedure, we accelerate the original DDPM sampling algorithm, which takes every step in the reverse sampling process, by implementing the discrete DDIM as per Equation 6. This discrete DDIM allows us to skip every  $k$  steps, resulting in a speed-up of the original DDPM by a factor of  $k$ . We conducted an ablation study on the impact of speed and recovery rate by trying different skip steps: 1, 2, 5, 10, 20, 25, 50, and 100. We compare the recovery rates achieved by these different steps. Our results revealed that the recovery rate performance decays as the number of
Figure 6: Recovery rate with the different selection of the transition matrix, whether considering equivariance and secondary structure.
skipped steps increases. The best performance is achieved when skipping a single step, resulting in a recovery rate of  $52.21\%$ , but at a speed of 100 times slower than when skipping 100 steps, which yields a recovery rate of  $47.66\%$ .
Figure 7: Comparison of the distribution of mutational prior (BLOSUM replacement matrix) and sampling results.
We further compare the diversity of GRADE-IF with PIFOLD and PROTEINMPPNN in Figure 8. For a given backbone, we generate 100 sequences with a self-similarity less than  $50\%$  and employ t-SNE [45] for projection into a 2-dimensional space. At the same level of diversity, GRADE-IF encompasses the wild-type sequence, whereas the other two methods fail to include the wild-type within their sample region. Furthermore, inspiring at a recovery rate threshold of  $45\%$  for this protein, GRADE-IF manages to generate a substantial number of samples, whereas the other two methods revert to deterministic results. This further substantiates the superiority of our model in terms of achieving sequence diversity and a high recovery rate concurrently.
Figure 8: t-SNE of the generated sequences of GRADE-IF compared to PIFOLD and PROTEINMPNN.
G Additional Folding Results
We further analyzed the generated sequences by comparing different protein folding predictions. We consider the crystal structures of three native proteins with PDB IDs: 1ud9 (A chain), 2rem (B chain), 3drn (B chain), which we randomly choose from CATH dataset. For each structure, we generated three sequences from the diffusion model and used ALPHAFOLD 2 [21] to predict the respective structures. As shown in Figure 9, these predictions (in purple) were then compared with the structures of the native protein sequences (in nude). We can observe that the RMSD for all cases is lower than the preparation accuracy of the wet experiment. The results demonstrate that our model-generated sequences retain the core structure, indicating their fidelity to the original structures.
Figure 9: Folding comparisons between native sequence and generated sequence