P ROTEIN STRUCTURE GENERATION VIA FOLDING DIFFUSION
Anonymous authors
Paper under double-blind review
A BSTRACT
The ability to computationally generate novel yet physically foldable protein structures could lead to new biological discoveries and new treatments targeting yet incurable diseases.
Despite recent advances in protein structure prediction, directly generating diverse, novel protein structures from neural networks remains difficult.
In this work, we present a new diffusion-based generative model that designs protein backbone structures via a procedure that mirrors the native folding process.
We describe protein backbone structure as a series of consecutive angles capturing the relative orientation of the constituent amino acid residues, and generate new structures by denoising from a random, unfolded state towards a stable folded structure.
Not only does this mirror how proteins biologically twist into energetically favorable conformations, the inherent shift and rotational invariance of this representation crucially alleviates the need for complex equivariant networks.
We train a denoising diffusion probabilistic model with a simple transformer backbone and demonstrate that our resulting model unconditionally generates highly realistic protein structures with complexity and structural patterns akin to those of naturally-occurring proteins.
As a useful resource, we release the first open-source codebase and trained models for protein structure diffusion.
I NTRODUCTION
Proteins are critical for life, playing a role in almost every biological process, from relaying signals across neurons (Zhou et al., 2017) to recognizing microscopic invaders and subsequently activating the immune response (Mariuzza et al., 1987), from producing energy for cells (Bonora et al., 2012) to transporting molecules along cellular highways (Dominguez & Holmes, 2011).
Misbehaving proteins, on the other hand, cause some of the most challenging ailments in human healthcare, including Alzheimer’s disease, Parkinson’s disease, Huntington’s disease, and cystic fibrosis (Chaudhuri & Paul, 2006).
As a consequence of their importance, proteins have been extensively studied as a therapeutic medium (Kamionka, 2011; Dimitrov, 2012) and constitute a rapidly growing segment of approved therapies (H Tobin et al., 2014).
Thus, the ability to computationally generate novel yet physically foldable protein structures could open the door to discovering novel ways to harness cellular pathways and eventually lead to new treatments targeting yet incurable diseases.
Many works have tackled the problem of generating new protein structures, but have generally run into challenges with creating diverse yet realistic folds.
Traditional approaches typically apply heuristics to assemble fragments of experimentally profiled proteins into structures (Schenkelberg & Bystroff, 2016; Holm & Sander, 1991).
This approach is limited by the boundaries of expert knowledge and available data.
More recently, deep generative models have been proposed.
However, due to the incredibly complex structure of proteins, these commonly do not directly generate protein structures, but rather constraints (such as pairwise distance between residues) that are heavily post-processed to obtain structures (Anand et al., 2019; Lee & Kim, 2022).
Not only does this add complexity to the design pipeline, but noise in these predicted constraints can also be compounded during post-processing, resulting in unrealistic structures – that is, if the constraints are at all satisfiable to begin with.
Other generative models rely on complex equivariant network architectures or loss functions to learn to generate a 3D point cloud that describes a protein structure (Anand & Achim, 2022; Trippe et al., 2022; Luo et al., 2022; Eguchi et al., 2022).
Such equivariant architectures can ensure that the probability density from which the protein structures are sampled is invariant under translation and rotation.
However, translation- and rotation-equivariant architectures are often also symmetric under reflection, leading to violations of fundamental structural properties of proteins like chirality (Trippe et al., 2022).
Intuitively, this point cloud formulation is also quite detached from how proteins biologically fold – by twisting to adopt energetically favorable configurations (ˇSali et al., 1994; Englander et al., 2007).
Inspired by the in vivo protein folding process, we introduce a generative model that acts on the interresidue angles in protein backbones instead of on Cartesian atom coordinates (Figure 1).
This treats each residue as an independent reference frame, thus shifting the equivariance requirements from the neural network to the coordinate system itself.
This angular representation has been successfully used by some protein structure prediction works (Gao et al., 2017; AlQuraishi, 2019; Chowdhury et al., 2022).
For generation, we use a denoising diffusion probabilistic model (diffusion model, for brevity)
(Ho et al., 2020; Sohl-Dickstein et al., 2015) with a vanilla transformer parameterization without any equivariance constraints.
Diffusion models train a neural network to start from noise and iteratively “denoise” it to generate data samples.
Such models have been highly successful in a wide range of data modalities from images (Saharia et al., 2022; Rombach et al., 2022) to audio (Rouard & Hadjeres, 2021; Kong et al., 2021), and are easier to train with better modal coverage than methods like generative adversarial networks (GANs) (Dhariwal & Nichol, 2021; Nichol & Dhariwal, 2021).
We present a suite of validations to quantitatively demonstrate that unconditional sampling from our model directly generates realistic protein backbones – from recapitulating the natural distribution of protein inter-residue angles, to producing overall structures with appropriate arrangements of multiple structural building block motifs.
We show that our generated backbones are diverse and designable, and are thus biologically plausible protein structures.
Our work demonstrates the power of biologically-inspired problem formulations and represents an important step towards accelerating the development of new proteins and protein-based therapies.
R ELATED WORK 2. G ENERATING NEW PROTEIN STRUCTURES
Many generative deep learning architectures have been applied to the task of generating novel protein structures.
Anand et al.
(2019) train a GAN to sample pairwise distance matrices that describe protein backbone arrangements.
However, these pairwise distance matrices must be corrected, refined, and converted into realizable backbones via two independent post-processing steps, the Alternating Direction Method of Multipliers (ADMM) and Rosetta.
Crucially, inconsistencies in these predicted constraints can render them unsatisfiable or lead to significant errors when reconstructing the final protein structure.
Sabban & Markovsky (2020) use a long short-term memory (LSTM)
GAN to generate the ( ϕ, ψ ) dihedral angles.
However, their network relies on downstream postprocessing to filter, refine, and fold predicted 3D structures, partly due to the fact that these two dihedrals do not sufficiently specify backbone structure.
Eguchi et al.
(2022) propose a variational auto-encoder with equivariant losses to generate 3D coordinates for protein backbones.
However, their work only targets immunoglobulin proteins and also requires refinement through Rosetta.
Nondeep learning methods have also been explored.
Schenkelberg & Bystroff (2016) apply heuristics to ensembles of similar sequences to make relatively small perturbations to known protein structures, while Holm & Sander (1991) use a database search to find and assemble existing protein fragments that might fit a new scaffold structure.
These approaches’ reliance on known proteins and handengineered heuristics limit them to relatively small deviations from naturally-occurring proteins.
2.1. D IFFUSION MODELS FOR PROTEIN STRUCTURE GENERATION
Several recent works have proposed extending diffusion models towards generating protein structures.
These predominantly perform diffusion on the 3D Cartesian coordinates of the residues themselves.
For example, Trippe et al.
(2022) use an E(3)-equivariant graph neural network to model the coordinates of protein residues.
Anand & Achim (2022) adopt a hybrid approach where they train an equivariant transformer with invariant point attention (Jumper et al., 2021); this model generates the 3D coordinates of C α atoms, the amino acid sequence, and the angles defining the orientation of side chains.
Another recent work by Luo et al.
(2022) performs diffusion for generating antibody fragments’ structure and sequence by modeling 3D coordinates using an equivariant neural network.
Note that these prior works all use some form of equivariance to translation, rotation, and/or reflection due to their formulation of diffusion on Cartesian coordinates.
Another method,
N
Cα
C
O
R
N (cid:31) (cid:30) (cid:29)
Cα (cid:28) (cid:27) (cid:28) (cid:26) (cid:28) (cid:25) ...
R
[Figure]
ProteinSGM, (Lee & Kim, 2022) implements a score-based diffusion model that generates imagelike square matrices describing pairwise angles and distances between all residues in an amino acid chain.
However, this set of values is highly over-constrained, and must be used as a set of input constraints for Rosetta’s folding algorithm
(Yang et al., 2020), which in turn produces the final folded output.
This is a similar approach to Anand et al.
(2019), and is likewise subject to the aforementioned concerns regarding complexity, satisfiability, and cleanliness of predicted constraints.
Our work instead uses a minimal set of angles required to specify a protein backbone, and thus directly generates structures without relying on additional methods for refinement.
Unfortunately, none of these prior works have publicly-available code, model weights, or generated examples at the time of this writing.
Thus, our ability to perform direct qualitative and quantitative comparisons is limited.
2. D IFFUSION MODELS FOR SMALL MOLECULES
A related line of work focuses on creating and modeling small molecules, typically in the context of drug design, using similar generative approaches.
These small molecules average 44 atoms in size (Jing et al., 2022).
Compared to proteins, which average several hundred residues and thousands of atoms (Tiessen et al., 2012), the relatively small size of small molecules makes them easier to model.
The E(3) Equivariant Diffusion Model (Hoogeboom et al., 2022) uses an equivariant transformer to design small molecules by diffusing on their coordinates in Euclidean space.
Other works have explored torsional diffusion, i.e., modelling the angles that specify a small molecule, to sample from the space of energetically favorable molecular conformations (Jing et al., 2022).
This work still requires an SE (3) -equivariant model as the input to their model is still a 3D point cloud.
In contrast, our problem formulation allows us to work entirely in terms of relative angles.
M ETHOD 3. S IMPLIFIED FRAMING OF PROTEIN BACKBONES USING INTERNAL ANGLES
Proteins are variable-length chains of amino acid residues.
There are a total of 20 canonical amino acids, all of which share the same three-atom N − C α − C backbone, but have varying side chains attached to the C α atom (typically denoted R , see illustration in Figure 1).
These residues assemble to form polymer chains typically hundreds of residues long (Tiessen et al., 2012).
These chains of amino acids fold into 3D structures, taking on a shape that largely determines the protein’s functions.
These folded structures can be described on four levels: primary structure, which simply captures the linear sequence of amino acids; secondary structure, which describes the local arrangement of amino acids and includes structural motifs like α -helices and β -sheets; tertiary structure, which
[Table]
describes the full spatial arrangement of all residues; and quaternary structure, which describes how multiple different amino acid chains come together to form larger complexes (Sun et al., 2004).
We propose a simplified framing of protein backbones that follows the biological intuition of protein folding while removing the need for complex equivariant networks.
Rather than viewing a protein backbone of length N amino acids as a cloud of 3D coordinates (i.e., x ∈ R N × 3 if modeling only
C α atoms, or x ∈ R 3 N × 3 for a full backbone) as prior works have done, we view it as a sequence of six internal, consecutive angles x ∈ [ − π, π ) ( N − 1) × 6 .
That is, each vector of six angles describes the relative position of all backbone atoms in the next residue given the position of the current residue.
These six angles are defined precisely in Table 1 and illustrated in Figure 1.
These internal angles can be easily computed using trigonometry, and subsequently converted back to 3D Cartesian coordinates by iteratively adding atoms to the protein backbone as described in Parsons et al.
(2005), keeping bond distances fixed to average lengths (see Appendix A.1, Figure S1).
This internal angle formulation has several key advantages.
Most importantly, since each residue forms its own independent reference frame, there is no need to use an equivariant neural network.
No matter how the protein is rotated or shifted, the angle of the next residue given the current residue never changes.
This allows us to use a simple transformer as the backbone architecture; in fact, we demonstrate that our model fails when substituting our shift- and rotation-invariant internal angle representation with Cartesian coordinates, keeping all other design choices identical (see Appendix
C.1, Figure S4).
This internal angle formulation also closely mimics how proteins actually fold by twisting into more energetically stable conformations.
3. D ENOISING DIFFUSION PROBABILISTIC MODELS
Denoising diffusion probabilistic models (or diffusion models, for short) leverage a Markov process q ( x t | x t − 1 ) to corrupt a data sample x 0 over T discrete timesteps until it is indistinguishable from noise at x T .
A diffusion model p ξ ( x t − 1 | x t ) parameterized by ξ is trained to reverse this forward noising process, “denoising” pure noise towards samples that appear drawn from the native data distribution (Sohl-Dickstein et al., 2015).
Diffusion models were first shown to achieve good generative performance by Ho et al.
; we adapt this framework for generating protein backbones, introducing necessary modifications to work with periodic angular values.
We modify the standard Markov forward noising process that adds noise at each discrete timestep t to sample from a wrapped normal instead of a standard normal (Jing et al., 2022):
[Equation]
(cid:19) where β t ∈ (0 , 1) T t =1 are set by a variance schedule.
We use the cosine variance schedule (Nichol &
Dhariwal, 2021) with T = 1000 timesteps: β t = clip (cid:18) 1 − [Equation] α t ¯ α
[Equation]
1 , 0 .(cid:19)
[Equation]
[Equation]
f
[Equation]
t
[Equation] (cid:18) [Equation] s 1 + s · π (cid:19) where s = 8 × 10 − 3 is a small constant for numerical stability.
We train our model for p ξ ( x t − 1 | x t ) with the simplified loss proposed by Ho et al.
(2020), using a neural network nn ξ ( x t , t ) that predicts the noise ϵ ∼ N (0 , I ) present at a given timestep (rather than the denoised mean values themselves).
To handle the periodic nature of angular values, we introduce a function to “wrap” values within the range [ − π, π ) : w ( x ) =
(( x + π ) mod 2 π ) − π .
We use w to wrap a smooth L1 loss (Girshick, 2015)
L w , which behaves like L1 loss when error is high, and like an L2 loss when error is low; we set the transition between these two regimes at β L = 0 .
1 π .
While this loss is not as well-motivated as torsional losses used by Jing et al.
(2022), we find that it achieves strong empirical results.
[Equation]
During training, timesteps are sampled uniformly t ∼ U (0 , T ) .
We normalize all angles in the training set to be zero mean by subtracting their element-wise angular mean µ ; validation and test sets are shifted by this same offset.
Figure 1 illustrates this overall training process, including our previously described internal angle framing.
The internal angles describing the folded chain x 0 are corrupted until they become indistinguishable from random angles, which results in a disordered mass of residues at x T ; we sample points along this diffusion process to train our model nn ξ .
Once trained, the reverse process of sampling from p ξ also requires modifications to account for the periodic nature of angles, as described in Algorithm 1.
The variance of this reverse process is given by σ t = (cid:113) 1 − ¯ α t − 1 1 − ¯ α t · β t .
1: x T ∼ w ( N (0 , I ))
▷ Sample from a wrapped Gaussian 2: for t = T, . . .
, 1 do 3:
z = N (0 , I ) if t > 1 else z =4: x t − 1 = w (cid:16) 1 √ α t (cid:16) x t − 1 − α t √ 1 − ¯ α t nn ξ ( x t , t ) (cid:17) + σ t z (cid:17)
▷ Wrap sampled values about [ − π, π ) 5: end for 6: return w ( x 0 + µ ) ▷ Un-shift generated values by original mean shift
This sampling process can be intuitively described as refining internal angles from an unfolded state towards a folded state.
As this is akin to how proteins fold in vivo , we name our method FoldingDiff.
3. M ODELING AND DATASET
For our reverse (denoising) model p ξ ( x t , t ) , we adopt a vanilla bidirectional transformer architecture (Vaswani et al., 2017) with relative positional embeddings as described in Shaw et al.
Our six-dimensional input is linearly upscaled to the model’s embedding dimension ( d = 384 ).
To incorporate the timestep t into this model, we generate random Fourier feature embeddings (Tancik et al., 2020) as done in Song et al.
and add these embeddings to each upscaled input.
To convert the transformer’s final per-position representations to our six outputs, we apply a regression head consisting of a densely connected layer, followed by GELU activation (Hendrycks & Gimpel, 2016), layer normalization, and finally a fully connected layer outputting our six values.
We train this network with the AdamW optimizer (Loshchilov & Hutter, 2019) over 10,000 epochs, with a learning rate that linearly scales from 0 to 5 × 10 − 5 over 1,000 epochs, and back to 0 over the final 9,000 epochs.
Validation loss appears to plateau after ≈ 1,400 epochs; additional training does not improve validation loss, but appears to lead to a poorer diversity of generated structures.
We thus take a model checkpoint at 1,484 epochs for all subsequent analyses.
We train our model on the CATH dataset, which provides a “de-duplicated” set of proteins spanning a wide range of functions where no two chains share more than 40% sequence identity over 60% overlap (Sillitoe et al., 2015).
We exclude any chains with fewer than 40 residues.
Chains longer than 128 residues are randomly cropped to a 128-residue window at each epoch.
Using a random 80/10/10 training/validation/test split, we have 24,316 training backbones, 3,039 validation backbones, and 3,040 test backbones.
Expanding this training set is a target for future work.
E XPERIMENTS 4. G ENERATING PROTEIN INTERNAL ANGLES
After training our model, we check that it is able to recapitulate the correct marginal distributions of dihedral and bond angles in proteins.
We unconditionally generate 10 backbone chains each for every length from 50 to 128, which results in a total of 780 generated backbone chains.
We plot the distributions of all six angles, aggregated across these 780 structures, and compare each distribution to that of test set structures less than 128 residues in length (Figures 2, S9).
We observe that, across all angles, the generated distribution almost exactly recapitulates the test distribution.
This is true both for angles that are nearly Gaussian with low variance ( ω, θ 1 , θ 2 , θ 3 ) as well as for angles with highly complex, high-variance distributions ( ϕ, ψ ) .
Compared to similar plots generated from other protein diffusion methods (e.g., Figure 1 in Anand & Achim (2022), reproduced with permission in Figure S10), we qualitatively observe that our method produces a much tighter distribution that more closely matches the natural distribution of bond angles.
[Figure]
n
[Figure]
, KL
[Figure]
Figure S9 shows the cumulative distribution function (CDF) corresponding to these histograms.
However, looking at individual distributions of angles alone does not capture the fact that these angles are not independently distributed, but rather exhibit significant correlations.
A Ramachandran plot, which plots the frequency of co-occurrence between the dihedrals ( ϕ, ψ ) , is commonly used to illustrate these correlations between angles (Ramachandran & Sasisekharan, 1968).
Figure 3 shows the Ramachandran plot for chains with fewer than 128 residues in the test set, as well as that for our 780 generated structures.
The Ramachandran plot for natural structures (Figure 3a) contains three major concentrated regions corresponding to right-handed α helices, left-handed α helices, and β sheets.
All three of these regions are recapitulated in our generated structures (Figure 3b).
In other words, FoldingDiff is able to generate all three major secondary structure elements in protein backbones.
Furthermore, we see that our model correctly learns that right-handed α helices are much more common than left-handed α helices (Cintas, 2002).
Prior works that use equivariant networks, such as Trippe et al.
(2022), cannot differentiate between these two types of helices due to network equivariance to reflection.
This concretely demonstrates that our internal angle formulation leads to improved handling of chirality (i.e., the asymmetric nature of proteins) in generated backbones. [Figure]
3 (radians)       ( r a d i a n s ) helix, LH helix, RH sheet (b) Ramachandran plot, generated backbones      3 (radians)       ( r a d i a n s )
4. A NALYZING GENERATED STRUCTURES
We have shown that our model generates realistic distributions of angles and that our generated joint distributions capture secondary structure elements.
We now demonstrate that the overall structures specified by these angles are biologically reasonable.
Recall that naturally occurring protein structures contain multiple secondary structure elements.
We use P-SEA (Labesse et al., 1997) to count the number of secondary structure elements in each test-set backbone of fewer than 128 residues, and plot the frequency of α / β co-occurrence counts in Figure 4a.
Figure 4b repeats this analysis for our generated structures, which frequently contain multiple secondary structure elements just as naturally-occurring proteins do.
FoldingDiff thus appears to generate rich structural information.
This is a nontrivial task, as an autoregressive baseline collapses into a failure mode of endlessly generating α helices (Appendix C.2, Figures S5, S6)
[Figure]
t
[Figure] Beyond demonstrating that FoldingDiff’s generated backbones contain reasonable structural motifs [Figure] it is also important to show that they are designable – meaning that we can find a sequence of amino acids that can fold into the designed backbone structure.
After all, a novel protein structure is not useful if we cannot physically realize it.
Previous works evaluate this in silico by predicting possible amino acids that fold into a generated backbone and checking whether the predicted structure for these sequences matches the original backbone (see Appendix B).
Following this general procedure, for a generated structure s , we use the ESM-IF1 inverse folding model (Hsu et al., 2022) to generate 8 different amino acid sequences.
We then use OmegaFold (Wu et al., 2022) to predict the 3D structures ˆ s 1 , . .
, ˆ s 8 corresponding to each of these sequences.
We use TMalign (Zhang & Skolnick, 2005), which evaluates structural similarity between backbones, to score each of these 8 structures against the original structure s .
The maximum score max i ∈ [1 , 8] TMalign( s, ˆ s i ) is the self-consistency TM (scTM) score.
A scTM score of ≥ 0.5 is considered to be in the same fold, and thus is “designable.”
We repeat this process for each of our 780 generated backbones.
[Figure]
arrows are illustrated in Figure S12.
Overall, we find that 111 of our 780 structures (14.2%) are designable with an scTM score ≥ 0 .(Figure 5a), which is higher than the value of 11.8% reported by Trippe et al. This designability is highly consistent across random seeds (Table S1).
Our improved designability comes from a significantly higher proportion of short structures ( ≤ 70 residues) being designable (57/ for ours compared to 36/210 in Trippe et al.
p = 0 .
014 , Chi-square test).
There is no significant difference in designability for longer structures of 71-128 residues (54/570 for ours compared to 51/570 in Trippe et al.
p = 0 .
76 , Chi-square test).
While ProteinSGM (Lee & Kim, 2022) reports an even higher designability proportion of 50.5%, this metric is not directly comparable, as ProteinSGM generates sets of constraints that are then used to fold the final structure with Rosetta, rather than generating backbone structures directly as with our approach or that of Trippe et al. (2022).
It is unsurprising that a structure produced by a dedicated protein folding tool produces greater designability.
The authors themselves note that this Rosetta “post-processing” significantly improves the viability of their structures.
To additionally contextualize our scTM scores, we benchmark against a naive baseline structure sampling method that preserves the overall distribution of protein bond angles but destroys their positional relationships; FoldingDiff significantly outperforms this baseline (see Appendix C.3, Figure S7). Notably, even natural sequences do not have perfect designability; only 80% of experimentally obtained structures have an scTM ≥ 0 . 5 (Figure S7).
We additionally evaluate the similarity of each generated backbone to any training backbone by taking the maximum TM score across the entire training set.
The maximum training TM-score is significantly correlated with scTM score (Spearman’s r = 0 .
79 , p = 1 . 2 × 10 − 165 , Figure 5b), indicating that structures more similar to the training set tend to be more designable.
However, this does not suggest that we are merely memorizing the training set; doing so would result in a distribution of training TM scores near 1.0, which is not what we observe.
We note that ProteinSGM reports a distribution of training set TMscores much closer to 1.0; this suggests a greater degree of memorization and may contribute to their high reported scTM designability ratio.
Selected examples of our generated backbones and corresponding OmegaFold predictions of various lengths are visualized using PyMOL (Schr¨odinger, LLC, 2015) in Figure 6.
Interestingly, we find that of our 111 designable backbones, only 4 contain β sheets as annotated by P-SEA.
Conversely, of our 669 backbones with scTM < 0.5, 545 contain β sheets.
This suggests that generated structures with β sheets may be less designable ( p < 1 .
× 10 − 5 , Chi-square test).
It is unclear whether this is due to our model generating poor backbones in this condition, or because ESM-IF1 and
OmegaFold struggle with β sheets.
We additionally cluster our designable backbones and observe a large diversity of structures (Figure S14) comparable to that of natural structures (Figure S15).
This suggests that our model is not simply generating small variants of a handful of core structures, which prior works appear to do (Figure S16).
[Figure]
C ONCLUSION
In this work, we present a novel parameterization of protein backbone structures that allows for simplified generative modeling.
By considering each residue to be its own reference frame, we describe a protein using the resulting relative internal angle representation.
We show that a vanilla transformer can then be used to build a diffusion model that generates high-quality, biologically plausible, diverse protein structures.
These generated backbones better respect protein chirality and exhibit greater designability compared to prior works that use equivariance assumptions.
While we demonstrate promising results with our model, there are several limitations to our work.
Although formulating a protein as a series of angles enables us to use simpler models without equivariance mechanisms, this framing allows accumulated errors to significantly alter the overall generated structure.
This takes two forms: small errors iteratively compounding into larger errors over many residues, and single large errors drastically modifying the structure.
Additionally, some generated structures exhibit collisions where the generated structure crosses through itself.
Future work could explore methods to avoid these pitfalls using geometrically-informed architectures such as those used in Wu et al.
Our generated structures are still of relatively short lengths (up to 128 residues) compared to natural proteins which typically have several hundred residues.
We also do not handle multi-chain complexes or ligand interactions, and are only able to generate static structures that do not capture the dynamic nature of proteins.
Additional future work could incorporate amino acid sequence generation in parallel with structure generation, along with guided generation using protein function or domain annotations.
In summary, our work provides an important step in using biologically-inspired problem formulations for generative protein design.
A I NTERNAL ANGLE FORMULATION OF PROTEIN BACKBONES A. C HOICE OF ANGLES FOR REPRESENTATION
A protein backbone structure can be fully specified by a total of 9 values per residue: 3 bond distances, 3 bond angles, and 3 dihedral torsional angles.
The three bond angles and dihedrals are described in Table 1, and the three bond distances correspond to N i
→ Cα i , Cα i
→ C i , and
C i → N i +1 where i denotes residue index.
These values enable a protein backbone to be losslessly converted from Cartesian to internal angle representation, and vice versa.
To determine which subset of values to use to frame proteins in our model, we take a set of experimentally profiled proteins and translate their coordinates from Cartesian to internal angles and distances and back, measuring the
TM score between the initial and reconstructed structures.
When excluding an angle or distance, we fix all corresponding values to the mean.
The reconstruction TM scores of various combinations of values is illustrated in Figure S1.
Of these 9 values, the three bond distances are the least important for reliably reconstructing a structure from Cartesian coordinates to the inter-residue representation and back; they can usually be replaced with constant average values without much impact on the recovered structure.
In comparison, removing even two bond angles with relatively little variance ( θ 2 , θ 3 ) results in a large loss in reconstruction TM score (third bar).
Removing all bond angles and retaining only dihedrals ( ϕ, ψ, ω ) results in only about half of proteins being able to be reconstructed (last bar).
Thus, we model the three dihedrals and the three bond angles (second bar in Figure S1); this simplifies our prediction problem to use only periodic angular values (instead of a mixture of angular and real values) without a substantial loss in the accuracy of described structures.
Future work might include additional modeling of these real-valued bond distances.
d i h e d r a l s , a n g l e s , a n d d i s t a n c e s d i
h e d r a l s a n d a n g l e s d i h e d r a l s , 1 , a n d d i s t a n c e s d i h e d r a l s , d i h
e d r a l s , d i s t a n c e s d i
h e d r a l s o n l y 0.
T M - s c o r e o f r e c o n s t r u c t i o n
[Figure] One challenge with converting between a N -residue set of Cartesian coordinates to a set of N − 1 angles between consecutive residues is that the latter representation does not capture the first residue’s information.
To solve this, we use a fixed set of coordinates to “seed” generation of Cartesian coordinates, using the N − 1 specified angles to build out from this fixed point.
For all generations, this fixed point is extracted from the coordinates of the N − C α − C atoms in first residue on the
N-terminus of the PDB structure 1CRN (Teeter, 1984).
A. E FFECT OF LENGTH ON STRUCTURE RECONSTRUCTION
One of the primary concerns of using an angle-based formulation is that small errors can propagate across many residues to culminate in a large difference in overall structure. To try and quantify the effect of this, perform two evaluations: the “lossiness” of the representation itself, and the ability of the model to learn long range angle dependencies. We start by evaluating the accuracy of our representation itself over different structure lengths.
To do this, we sample 5000 structures of varying length from the CATH dataset.
For each, we compare the original 3D coordinate representation x c , and the 3D coordinates obtained after converting the structure to angles and back to coordinates ˆ x c using the TM score algorithm, i.e., TMscore ( x c , ˆ x c ) .
We find that longer structures have reconstruction, i.e., are lossier when using our representation (Figure S2).
However, even at our maximum considered length of 128 residues, structures still retain a TMscore ≈ 0 . 9 , which is well above the accepted similarity threshold of 0.5 denoting the same fold.
This indicates that while our representation itself is somewhat lossy, the losses are insignificant.
Even at longer sequences up to 512 residues in length, the reconstructed structures still share a TMscore similarity much greater than 0.5 (graph not shown), which suggests that our method could effectively scale up .
[Figure]
f r e c o n s t r u c t i o n
Next, we evaluate our model’s ability to successfully reconstruct sequences of varying length.
For each structure in our held-out test set ( n = 3040 ), we add t = 750 timesteps of noise to that structure’s angles (recall that we use a total of T = 1000 timesteps).
This adds a significant amount of noise to corrupt the structure, while retaining a hint of the target, true structure as a weak “guide” to what the model should ultimately diffuse towards.
We then apply our trained model to these mostlynoised examples, running them for the requisite 750 iterations to fully denoise and reconstruct the angles.
Afterwards, we assess reconstruction accuracy by taking the TMscore between the structure specified by the reconstructed angles, and the true structure specified by the ground truth denoised angles.
Figure S3 illustrates the relationship between test set structure length and this reconstruction
TM score.
We find that there is no significant correlation between length and reconstruction TM score (Spearman’s correlation r = − 0 .
0024 , p = 0 . 89 ).
[Figure]
All in all, these results indicate that (1) our representation itself is indeed noisier with longer structures, but not to a degree that would significantly impact the overall structures, and (2) our model is capable of learning robust relationships between angles regardless of structure length.
B A DDITIONAL NOTES ON SELF - CONSISTENCY TM SCORE
Our scTM evaluation pipeline is similar to previous evaluations done by Trippe et al.
and
Lee & Kim (2022), except that we use OmegaFold instead of AlphaFold (Jumper et al., 2021).
OmegaFold is designed to be used without multiple sequence alignments (MSAs), and performs similarly to AlphaFold while generalizing better to orphan proteins that may not have such evoluationary annotations (Wu et al., 2022).
Furthermore, given that prior works use AlphaFold without
MSA information in their evaluation pipelines (Trippe et al., 2022; Lee & Kim, 2022), OmegaFold appears to be a more appropriate method for scTM evaluation.
Trippe et al.
(2022) use ProteinMPNN (Dauparas et al., 2022) for inverse folding, whereas Lee & Kim (2022) uses ESM-IF1, as we do.
We empirically found that ESM-IF1 has better performance than ProteinMPNN.
Both these prior works generate 8 candidate amino acid sequences for each generated structure, as we do here.
We perform an “ablation” of our internal angle representation by replacing our framing of proteins as a series of inter-residue internal angles with a simple Cartesian representation of C α coordinates x ∈ R N × 3 .
Notably, this Cartesian representation is no longer rotation or shift invariant.
We train a denoising diffusion model using this Cartesian representation, using the same variance schedule, transformer backbone architecture, and loss function, but sampling from a standard Gaussian and with all usages of our wrapping function w removed.
This represents the same modelling approach as our main diffusion model, with only our internal angle formulation removed.
To evaluate the quality of this Cartesian-based diffusion model’s generated structures, we calculate the pairwise distances between all C α atoms in its generated structures and compare these with distance matrices calculated for real proteins and for our internal angle diffusion model’s generations.
For a real protein, this produces a pattern that reveals close proximity between pairwise residues where the protein is folded inwards to produce a compact, coherent structure (Figure S4a).
However, similarly visualizing the C α pairwise distances in the Cartesian model’s generated structures yields no significant proximity or patterns between any residues (Figure S4b).
This suggests that the ablated Cartesian model cannot learn to generate meaningful structure.
Our internal angle model, on the other hand, produces a visualization that is very similar to that of real proteins (Figure S4c).
Simply put, our model’s performance drastically degrades when we change only how inputs are represented.
This demonstrates the importance and effectiveness of our internal angle formulation.
[Figure]
Residue index
R e s i d u e i n d e x
C
C d i s t a n c e , Å
Residue index
[Figure]
s i d u e i n d e x
[Figure]
Residue index
R e s i d u e i n d e x
C
C d i s t a n c e , Å C. A UTOREGRESSIVE BASELINE FOR ANGLE GENERATION
As a baseline method for our generative diffusion model, we also implemented an autoregressive (AR) transformer
f AR that predicts the next set of six angles in a backbone structure (the same angles as we use in our main model, described in Figure 1 and Table 1) given all prior angles.
Architecturally, this model consists of the same transformer backbone as used in FoldingDiff combined with the same multi-layer perceptron to convert per-token embeddings to angle outputs, though it is trained using absolute positional embeddings rather than relative embeddings as we saw this produce an improved loss on validation set.
The total length of the sequence is encoded using random Fourier feature embeddings, just as time was encoded in our primary model, and this embedding is similarly added to each position in the sequence of angles.
Autoregressive training is done by asking the model to predict the i th set of six angles given all prior angles, using masking to hide the i th angle and onwards. We use the same wrapped smooth L1 loss as our main FoldingDiff model to handle the fact that these angle predictions exist in the range [ − π, π ) ; specifically: L w ( x ( i ) , f AR ( x (0 ,...,i − 1) )) where we use superscripts to indicate positional indexing. This approach is conceptually similar to causal language modeling (Radford et al., 2019), with the important difference that the inputs and outputs are continuous values, rather than discrete tokens (or probabilities over them).
This model is trained using the exact same data set and data splits as our main FoldingDiff model with the same preprocessing and normalization.
We train f AR using the AdamW optimizer (Loshchilov & Hutter, 2019) with weight decay set to 0.01.
We use a batch size of 256 over 10, epochs, where we linearly scale the learning rate from 0 to 5 × 10 − 5 over the first 1,000 epochs, and back to 0 over the remaining 9,000 epochs.
To generate structures from f AR , we “seed” the autoregressive model with 4 sets of 6 angles taken from the corresponding first 4 angle sets in a randomly-chosen naturally occurring protein structure.
This serves as a random, but biologically realistic “prompt” for the model to begin generation.
We then supply a fixed length value and repeated run f AR to obtain the next i -th set of angles, appending each prediction to the existing i − 1 values in order to predict the i + 1 set of angles.
We repeat this until we reach our desired structure length.
We use the above procedure to generate 10 structures for each structure length l ∈ [50 , 128) each with a different set of seed angles.
Examples of generated structures of varying lengths are illustrated in Figure S5.
All structures generated by this autoregressive approach appear to consist of singular α helix chains of varying length.
This is confirmed by running P-SEA on the generated structures, which detects only α helices across all 780 generated structures (Figure S6a).
Such behavior where autoregressive models generate a single looped pattern endlessly has been observed in language models as well, where it is often worked around using a temperature parameter to sample different words, thus breaking out of the loop.
In our continuous regime, we try to do something similar by adding small amounts of noise to each angle during generation, but are unable to meaningfully deviate from the aforementioned failure mode of collapse to endless α helices.
We additionally evaluate the AR model’s generated structures’ designability using the same ESMIF1 and OmegaFold pipeline used in our primary results.
We find that 223 of the 780 generated structures are designable, leading to an overall designability ratio of 0.29.
Interestingly, we observe that even when generating only α helices, as is the case here, longer helices are still less designable than shorter sequences.
This may be due to nature’s
propensity towards shorter coils in the range of 9-17 residues (Qin et al., 2013) – in other words, the AR model’s generated long coils are not frequently found in nature.
Importantly, however, these designable sequences being singular coils greatly limits their potential usefulness.
C. B ASELINE METHOD FOR CONTEXTUALIZING SC TM SCORES To contextualize the distribution of scTM scores (Figure 5a), we implement a naive angle generation baseline.
We take our test dataset, and concatenate all examples into a matrix of ˆ x ∈ [ − π, π ) ˆ N × 6 , where ˆ N denotes the total number of angle sets in our test dataset, aggregating across all individual chains.
To generate a backbone structure of length l , we simply sample l indices from U (0 , ˆ N ) .
This creates a chain that perfectly matches the natural distribution of protein internal angles, while additional structures generated all angles (see Figure S6a of diversity complexity indicates modeling
[Figure]
of the 780 structures being designable with scTM ≥ 0.5,
though this increase is of questionable value due to the lack of diversity in generated sequences. also perfectly reproducing the pairwise correlations, i.e., of dihedrals in a Ramachandran plot, but critically loses the correct ordering of these angles.
We randomly generate 780 such structures ( samples for each integer value of l ∈ [50 , 128) ).
This is the same distribution of lengths as the generated set in our main analysis.
For each of these, we perform the same scTM evaluation as before.
The distribution of scTM scores for these randomly-sampled structures compared to that of FoldingDiff’s generated backbones is shown in Figure S7.
We observe that this random protein generation method produces significantly poorer scTM scores than our model does ( p = 4 .
4 × 10 − 60 , Mann-Whitney test).
This suggests that our model is not simply learning the overall distribution of angles, but is learning the correct ordering of angles that comprise a folded protein structure.
We additionally take the structures specified by these randomly generated angles and analyze them for secondary structure counts using P-SEA, as we did for Figure 4.
The resulting 2D histogram is illustrated in Figure S8, and represents a completely different distribution of secondary structures
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
N o r m a li z [Figure] scTM scores, naive baseline vs. generated structures
[Figure]
[Figure] present than natural sequences, or that of FoldingDiff’s generations.
This further suggests that this random angle baseline cannot reasonably generate natural-appearing structures.
It is clear that scTM scores and secondary structure analyses cannot be satisfied using a trivial solution. D A DDITIONAL S UPPLEMENTARY F IGURES AND T ABLES
C u m u l a t
i v e f r e q u e n c y
CDF
TestSampled      3 0.
CDF      3 0.
CDF 0.
2.5 Angle (rad) 0.
C u m u l a t
i v e f r e q u e n c y 1 CDF 0.
3.0 Angle (rad) 0.
2 CDF 0.
3.0 Angle (rad) 0.
3 CDF the work of Anand & Achim (2022) exhibit a much larger spread than the natural distribution of angles, whereas our work matches much more tightly (Figures 2, S9). [Table]
[Figure]
[Figure]
a t e d ( ou r s )
O m e g a F o l d
M a x t r a i n i ng
T M
[Figure]
[Figure]
These authors similarly cluster unconditionally generated backbones with scTM ≥ 0.5 using 1 −
TMscore ( x, y ) as a distance metric.
Compared to our identical evaluation, illustrated in Figure S14, we notice that this clustering has a few dark blocks of nearly 0 distance, or a TM score of nearly 1.
This suggests that among these designable backbones, many are actually minor variants of a core structure; in actuality, though this work claims to produce 90 designable structures, there seem to be fewer unique structures in this set due to many being near-duplicates.