A LPHA F OLD D ISTILLATION FOR I MPROVED I NVERSE P ROTEIN F OLDING
Anonymous authors
Paper under double-blind review
A BSTRACT
Inverse protein folding, i.e., designing sequences that fold into a given threedimensional structure, is one of the fundamental design challenges in bioengineering and drug discovery.
Traditionally, inverse folding mainly involves learning from sequences that have an experimentally resolved structure.
However, the known structures cover only a tiny space of the protein sequences, imposing limitations on the model learning.
Recently proposed forward folding models, e.g., AlphaFold, offer unprecedented opportunity for accurate estimation of the structure given a protein sequence.
Naturally, incorporating a forward folding model as a component of an inverse folding approach offers the potential of significantly improving the inverse folding, as the folding model can provide a feedback on any generated sequence in the form of the predicted protein structure or a structural confidence metric.
However, at present, these forward folding models are still prohibitively slow to be a part of the model optimization loop during training.
In this work, we propose to perform knowledge distillation on the folding model’s confidence metrics, e.g., pTM or pLDDT scores, to obtain a smaller, faster and end-to-end differentiable distilled model, which then can be included as part of the structure consistency regularized inverse folding model training.
Moreover, our regularization technique is general enough and can be applied in other design tasks, e.g., sequence-based protein infilling.
Extensive experiments show a clear benefit of our method over the non-regularized baselines.
E.g., in inverse folding design problems we observe up to 3% improvement in sequence recovery and up to 45% improvement in protein diversity, while still preserving structural consistency of the generated sequences.
I NTRODUCTION
To date, 8 out of 10 top selling drugs are engineered proteins (Arnum, 2022).
For functional protein design, it is often a pre-requisite that the designed protein folds into a specific three-dimensional structure.
The fundamental task of designing novel amino acid sequences that will fold into the given 3D protein structure is named inverse protein folding.
Inverse protein folding is therefore a central challenge in bio-engineering and drug discovery.
Computationally, inverse protein folding can be formulated as exploring the protein sequence landscape for a given protein backbone to find a combination of amino acids that supports a property (e.g. structural consistency).
This task, computational protein design, has been traditionally handled by learning to optimize amino acid sequences against a physics-based scoring function (Kuhlman et al., 2003).
In recent years, deep generative models have been proposed to solve this task, which consist of learning a mapping from protein structure to sequences (Jing et al., 2020; Cao et al., 2021;
Wu et al., 2021; Karimi et al., 2020; Hsu et al., 2022; Fu & Sun, 2022).
These approaches frequently use high amino acid recovery with respect to the ground truth sequence (corresponding to the input structure) as one success criterion.
Other success criteria are high TM score (reflecting structural consistency) and low perplexity (measuring likelihood to the training/natural sequence distribution).
However, such criteria solely ignore the practical purpose of inverse protein folding, i.e., to design novel and diverse sequences that fold to the desired structure and thus exhibit novel functions.
In parallel to machine learning advances in inverse folding, notable progresses have been made recently in protein representation learning (Rives et al., 2021;
Zhang et al., 2022), protein structure
[Table]
[Figure]
with the given structure.
One way of doing this (red line) would be to use forward protein folding models, e.g., AlphaFold, estimate structure from generated sequence, compare it with the ground truth to compute metrics such as TM or LDDT and, finally, regularize the original training loss (usually cross-entropy (CE)).
However, inference through folding models is slow (see Fig. 2), making them impractical to be a part of the optimization loop.
Alternatively, bypassing structure estimation, the folding model’s internal confidence metrics, such as pTM or pLDDT can be used instead (blue line).
This results in lower fidelity solutions that are still slow.
Instead, in this work, we propose to distill the confidence metrics of AlphaFold into a smaller, faster, and differentiable model (referred as AFDistill).
AFDistill is trained to maintain the comparable accuracy of the AlphaFold-estimated pTM/pLDDT, which now can be seamlessly used as part of the training loop (green line).
The inference of the improved inverse folding model remains unmodified and is shown on the right side of the figure.
prediction from sequences (Jumper et al., 2021;
Baek et al., 2021b), as well as in conditional protein sequence generation (Das et al., 2021; Anishchenko et al., 2021).
These lines of works have largely benefited by learning from millions of available protein sequences (that may or may not have a resolved structure) in a self/un-supervised pre-training paradigm.
Such large-scale pre-training has immensely improved the information content and task performance of the learned model.
For example, it has been observed that structural and functional aspects emerge from a representation learned on broad protein sequence data (Rives et al., 2021).
In contrast, inverse protein folding has mainly focused on learning from sequences that do have an experimentally resolved structure.
Those reported structures cover only less than 0.1% of the known space of protein sequences, limiting the learning of the inverse folding model.
In this direction, a recent work has trained an inverse folding model from scratch on millions of AlphaFold-predicted protein structures (in addition to tens of thousands of experimentally resolved structures) and shown performance improvement in terms of amino acid recovery (Hsu et al., 2022).
However, such large-scale training from scratch is computationally expensive.
A more efficient alternative would be to use the guidance of an already available forward folding model pre-trained on large-scale data in training the inverse folding model.
In this work we construct a framework where the inverse folding model is trained using a loss objective that consists of regular sequence reconstruction loss, augmented with an additional structure consistency loss (SC) (see Fig.
1 for the system overview).
The straightforward way of implementing this would be to use forward protein folding models, e.g., AlphaFold, to estimate the structure from generated sequence, compare it with ground truth and compute TM score to regularize the training.
However, a challenge in using Alphafold (or similar) directly is the computational cost associated with its inference(see Fig. 2), as well as the need of ground truth reference structure.
Internal confidence structure metrics from the forward folding model can be used instead.
However, that approach is still slow for the in-the-loop inverse folding model optimization.
To address this, in our work we: (i) Perform knowledge distillation on AlphaFold and include the resulting model, AFDistill (frozen), as part of the regularized training of the inverse folding model (we term this loss structure consistency (SC) loss).
The main properties of AFDistill model are that it is fast, accurate and end-to-end differentiable.
(ii) Perform extensive evaluations, where the results on standard structure-guided sequence design benchmarks show that our proposed system outperforms existing baselines in terms of lower perplexity and higher amino acid recovery, while maintaining closeness to original protein structure.
More interestingly, we improve diversity in the designed sequences, one of the main goals of protein design.
As a result of a trade-off between sequence recovery vs structure recovery our regularized model yields better sequence diversity for a given structure, consistent with the fact that even small (35-40 amino acid) protein fold holds a ‘sequence capacity’ exceeding 10 23 (Tian & Best, 2017).
Note that our regularization technique is not limited to the inverse folding design and, as we show, can be applied to other applications, such as sequence-based protein infilling, where we also observe performance improvement over the baseline.
(iii) Finally, the estimated structure consistency metric can either be used as part of the regularization of an inverse folding or infilling, during any other protein optimization tasks (e.g., (Moffat et al., 2021)) which would benefit from structural consistency estimation of the designed protein sequence, or as a cheap surrogate of AlphaFold that provides scoring of a given protein, reflecting its structural content.
R ELATED W ORK
Forward Protein Folding.
Several powerful computational approaches have recently been proposed for the task of forward folding, namely predicting the structure of a protein given its sequence.
In particular, AlphaFold (Jumper et al., 2021) makes use of multiple sequence alignments (MSAs) along with pairwise features.
Building on similar ideas, RoseTTAFold (Baek et al., 2021a) is a threetrack network combining information at the sequence level, the 2D distance map level, and the 3D coordinate level.
On the implementation side, OpenFold (Ahdritz et al., 2021) is a PyTorch reproduction of Alphafold.
However, MSAs of homologous proteins (e.g., antibodies, orphan proteins) are not always available and the MSA search adds significant overhead in inference time.
Therefore, MSA-free approaches such as OmegaFold (Wu et al., 2022), HelixFold (Fang et al., 2022), and ESMFold (Lin et al., 2022) have been subsequently proposed.
These architectures make use of large pretrained language models to produce predictions that are more accurate than AlphaFold and
RoseTTAFold given a single sequence as input, and on par when given full MSAs as input.
Inverse Protein Folding.
Several algorithms have also been proposed recently for the inverse protein folding problem, namely the task of identifying amino acid sequences that fold to a desired structure.
(Norn et al., 2020) proposed a deep learning approach that optimizes over the entire folding landscape by backpropagating gradients through the trRosetta structure prediction network (Yang et al., 2020).
(Anand et al., 2022) introduced a deep neural network architecture that explicitly models side-chain conformers in a structure-based context.
Instead of making use of
[Figure]
trRosetta, (Jendrusch et al., 2021) presented a framework that incorporates the structure prediction algorithm AlphaFold (Jumper et al., 2021) into an optimization loop to generate protein sequences.
However, as noted above, the inference of AlphaFold is expensive due to MSA search, hence employing AlphaFold in an optimization loop for protein sequence generation is very cumbersome.
Though faster, the inference of MSA-free approaches (e.g. OmegaFold, HelixFold and EMSFold) is still too slow for their use in an optimization loop.
In this work, we propose knowledge distillation from the forward folding algorithm AlphaFold, and build a student model that is small, practical and accurate enough.
We show that the distilled model can be efficiently used within the inverse folding model optimization loop and improve quality of designed protein sequences.
A LPHA F OLD D ISTILL
Knowledge distillation (Hinton et al., 2015) transfers knowledge from a large complex model, in our case AlphaFold, to a smaller one, here this is the AFDistill model (see Fig. 3).
Traditionally, the distillation would be done using soft labels, which are probabilities from AlphaFold model, and hard labels, the ground truth classes.
However, in our case we do not use the probabilities as they are harder to collect or unavailable, but rather the model’s predictions (pTM/pLDDT) and the hard labels, TM/LDDT scores, computed based on AlphaFold’s predicted 3D structures.
3. S CORES TO D ISTILL
TM-score (Template Modeling score) (Zhang & Skolnick, 2004), is the mean distance between structurally aligned C α atoms scaled by a length-dependent distance parameter.
LDDT (Local Distance
Difference Test) (Mariani et al., 2013) is the average of four fractions computed using distances between all pairs of atoms based on four tolerance thresholds (0.5 ˚A, 1 ˚A, 2 ˚A, 4 ˚A) within 15 ˚A inclusion radius.
The range of both metrics is (0,1), and the higher values represent more similar structures.
pTM and pLDDT are the AlphaFold-predicted metrics for a given input protein sequence, corresponding to the reconstructed 3D protein structure, which represent model’s confidence of the estimated structure.
pLDDT is a local per-residue score (pLDDT chain is another score that simply averages per-residue pLDDTs across the chain), while pTM is a global confidence metric for assessing the overall chain reconstruction.
In this work we interpret these metrics as the quality or validity of the sequence for the purpose of downstream applications (see Section 4).
3. D ATA
Using Release 3 (January 2022) of AlphaFold Protein Structure Database (Varadi et al., 2021), we collected a set of 907,578 predicted structures.
Each of these predicted structures contains 3D
[Table]
coordinates of all the residue atoms as well as the per-resiude pLDDT confidence scores.
To avoid any data leakage to the downstream aplications
[Figure]
we first filtered out the structures that are part of the validation and test splits of CATH 4.2 dataset (discussed in Section 4).
Then, using the remaining structures
[Figure]
we created our pLDDT 1M dataset (see Table 1)
[Figure]
where each protein sequence is paired with the sequence of per-residue pLDDTs.
We also truncated proteins up to length 500 to reduce computational complexity of AFDistill training
[Figure]
AFDistill training.
We also created datasets which are based on the true TM and LDDT values using the predicted AlphaFold structures.
Specifically, using the PDB-to-UniProt mapping, we selected a subset of samples with matching ground truth PDB sequences and 3D structures, resulting in 42,605 structures.
We denote these datasets as TM 42K and LDDT 42K (see Table 1).
In Fig.
4 we show the score density distribution of each dataset.
As can be seen, the TM 42K and LDDT 42K are very skewed to the upper range of the values.
To mitigate this data imbalance, we curated two additional TM-based datasets.
TM augmented 86K was obtained by augmenting TM 42K with a set of perturbed original protein sequences, estimating their structures with AlphaFold, computing corresponding TM-score, and keeping the low and medium range TM values.
pTM synthetic 1M was obtained by generating random synthetic protein sequences and feeding them to AFDistill (pre-trained on TM 42K data), to generate additional data samples and collect lower-range pTM values.
The distribution of the scores for these additional datasets is also shown in Fig.
4, where both TM augmented 86K and pTM synthetic 1M datasets are less skewed, covering lower (p)TM values better.
Finally, using Release 4 (July 2022), containing over 214M predicted structures, we plotted its distribution density of pLDDT values and observed similar high skewness towards upper range.
To fix this, we rebalanced the data by filtering out samples with upper-range mean-pLDDT values (also called pLDDT chain).
The resulting dataset contains 60M sequences, for which we additionally created 10M and 1M versions, see Fig.
4 for their density.
3. M ODEL
AFDistill model is based on ProtBert (Elnaggar et al., 2020), a Transformer BERT model (420M parameters) pretrained on a large corpus of protein sequences using masked language modeling.
For our task we modify ProtBert head by setting the vocabulary size to 50, corresponding to discretizing pTM/pLDDT in range (0,1).
For pTM (scalar) the output corresponds to the first ⟨
⟩ token of the output sequence,
(sequence) predictions are made for each residue position.
[Figure] [Figure]
3. D ISTILLATION E XPERIMENTAL R ESULTS
In this section we present evaluation results of the model after training it on the presented datasets.
We note that to further improve the data imbalance problem, during training we employed weighted sampling in the minibatch generation as well as used Focal loss (Lin et al., 2017) in place of the traditional cross-entropy loss.
The results for (p)TM-based datasets are shown in Table 2.
We see that AFDistill trained on TM 42K dataset performed the best, followed by the dataset with the augmentations, and the synthetic data.
For the (p)LDDT-based datasets, we observe that increasing the scale, coupled with the data balancing, improves the validation performance.
In Fig.
2 we show scatter plots of the true vs pTM scores and pLDDT values on the entire validation set.
We see a clear diagonal pattern in both plots, where the predicted and true values match.
There are also some number of incorrect predictions (reflected along the off-diagonal), where we see that for the true scores in the upper range, the predicted scores are lower, indicating that AFDistill tends to underestimate them.
Finally, in Fig.
5 and 6, we show a few examples of the data samples together with the corresponding AFDistill predictions.
Fig.
2 and Fig.
12, 13 (in Appendix) also show plots of SC (pTM or pLDDT) versus TM score, indicating that AFDistill is a viable tool for regularizing protein representation to enforce structural consistency or structural scoring of protein sequences, reflecting its overall composition and naturalness (in terms of plausible folded structure).
[Figure]
I NVERSE P ROTEIN F OLDING D ESIGN
In this section we demonstrate the benefit of applying AFDistill as a structure consistency (SC) score for solving the task of inverse protein folding, as well as for the protein sequence infilling as a means to novel antibody generation.
The overall framework is presented in Fig.
1 (following the green line in the diagram), where the traditional inverse folding model is regularized by our SC score.
Specifically, during training, the generated protein is fed into AFDistill, for which it predicts pTM or pLDDT score, and combined with the original CE training objective results in
[Equation]
where L CE = (cid:80) N 1 L CE ( s i , ˆ s i ) is the CE loss, s i is the ground truth and
ˆ s i is the generated protein sequence, L SC
=
(cid:80) Ni =1 (1 − SC (ˆ s i )) is the structure consistency loss, N is the number of training sequences, and α is the weighting scalar for the SC loss, in our experiment it is set to 1.
4. M ETRICS
To measure the quality of the prediction designs, we compute the following set of sequence evaluation metrics.
Recovery (range (0, 100) higher is better) is the average number of exact matches between predicted and ground truth sequences, normalized by the length of the alignment.
Diversity (range (0, 100) higher is better) of a predicted protein set is the complement of the average recovery computed for all pairwise comparisons in the set.
While in general the recovery and diversity tend to be inversely correlated, i.e., higher recovery leads to lower diversity, and vice versa, we are interested in models that achieve high recovery rates and be able to maintain high protein sequence diversity.
Perplexity measures the likelihood of a given sequence, lower values mean better performance.
Finally, for structure evaluation, we use TM-score as well as the structure consistency (SC) score, which is the AFDistill’s output (pTM/pLDDT) for a given input.
4. R ESULTS
We present experimental results for several recently proposed deep generative models for protein sequence design accounting for 3D structural constraints.
For the inverse folding tasks we use
CATH 4.2 dataset, curated by (Ingraham et al., 2019).
The training, validation, and test sets have 18204, 608, and 1120 structures, respectively.
While for protein infilling we used SabDab (Dunbar et al., 2013) dataset curated by (Jin et al., 2021) and focus on infilling CDR-H3 loop.
The dataset has 3896 training, 403 validation and 437 test sequences.
GVP Geometric Vector Perceptron GNNs (GVP) (Jing et al., 2020) is the inverse folding model, that for a given target backbone structure, represented as a graph over the residues, replaces dense layers in a GNN by simpler layers, called GVP layers, directly leveraging both scalar and geometric features.
This allows for the embedding of geometric information at nodes and edges without reducing such information to scalars that may not fully capture complex geometry.
The results of augmenting GVP training with SC score regularization are shown in Fig.
7 (see also Appendix, Table 8 for additional results).
Baseline GVP with no regularization achieves 38.6 in recovery, 15.
in diversity and 0.79 in TM score on the test set.
It can be seen that there is a consistent improvement in sequence recovery gain (1-3%) over the original GVP and significant diversity gain (up to
[Figure]
SC-regularized GVP, while having high sequence diversity, still results in accurate reconstructions, while GVP alone has more inconsistencies, marked with arrows.
45%) of the generated protein sequences, when we employ SC regularization.
At the same time the estimated structure (using AlphaFold) remains close to the original as measured by the high
TM score.
We also observed that pTM-based SC scores had overall better influence on the model performance as compared to pLDDT-based ones.
It should be further noted, that the validation performance of AFDistill on the distillation data is not always reflective of the performance on the downstream applications, as AFDistill trained on TM augmented 86K overall performs better than
TM 42K, while having slightly worse validation CE loss (Table 2).
This observation indicates that the augmented models might be less biased by the teacher model, hence enables more generalized representation learning of sequence–structure relationship and provides more performance boost to the inverse folding model.
To further illustrate the effect of recovery and diversity, we show in Fig.
8 protein sequences and
AlphaFold-generated 3D structures of GVP and GVP+SC models, where the latter model achieves higher diversity of the sequence while retaining accurate structure of the original protein.
Here, the
Graph Transformer We evaluated the effect of SC score on Graph Transformer (Wu et al., 2021), another inverse folding model, which seeks to improve standard GNNs to represent the protein 3D structure.
Graph Transformer applies a permutation-invariant transformer module after GNN module to better represent the long-range pair-wise interactions between the graph nodes.
The results of augmenting Graph Transformer training with SC score regularization are shown in Fig.(see also Appendix, Table 9 for additional results).
Baseline model with no regularization has in recovery, 72. diversity and 0.81 in score on set.
As compared to GVP (Fig. 7) we can see that for this model, the recovery and diversity gains upon SC regularization are smaller.
We also see that TM score of regularized model (TM 42K and TM augmented 86K pretraining) is slightly higher as compared to pLDDT-based models.
Protein Infilling Our proposed structure consistency regularization is quite general and not limited to the inverse folding task.
Here we show its application on protein infilling task.
Recall, that while the inverse folding task considers generating the entire protein sequence, conditioned on a given structure, infilling focuses on filling specific regions of a protein conditioned on a sequence/structure template.
The complementarity-determining regions (CDRs) of an antibody protein are of particular interest as they determine the antigen binding affinity and specificity.
We follow the framework of (Jin et al., 2021) which formulates the problem as generation of the CDRs conditioned on a fixed framework region.
We focus on CDR-H3 and use a baseline pretrained protein model ProtBERT (Elnaggar et al., 2020) finetuned on the infilling dataset, and use ProtBERT+SC as an alternative (finetuned with SC regularization).
The CDR-H3 is masked and the objective is to reconstruct it using the rest of the protein sequence as a template.
The results are shown in Fig.(see also Appendix, Table 10 for additional results).
Baseline model achieves 41.5 in recovery, 14.
in diversity, and 0.80 in TM score on the test set.
Similar as for the other applications, we see an improvement in the sequence recovery and even bigger gain in diversity, while using the AFDistill pretrained on TM 42K and TM augmented 86K, together with the pLDDT balanced datasets.
TM score shows that the resulting 3D structure remains close to the original, confirming the benefit of using SC for training regularization.
C ONCLUSION
In this work we introduce AFDistill, a distillation model based on AlphaFold, which for a given protein sequence estimates its structural consistency (SC: pTM or pLDDT) score.
We provide experimental results to showcase the efficiency and efficacy of the AFDistill model in high-quality protein sequence design, when used together with a graph neural net based inverse folding model or large protein language model for sequence infilling.
Our AFDistill model is small and accurate enough so that it can be efficiently used for regularizing structural consistency in protein optimization tasks, maintaining sequence and structural integrity, while introducing diversity and variability in the generated proteins.
R EPRODUCIBILITY S TATEMENT
For the sake of reproducibility, we use publicly available datasets and describe the data in detail in
Section 3.2.
We provide training details in the Appendix.
The code will be publicly released upon acceptance
A B ACKGROUND ON P ROTEIN D ESIGN
A protein is a linear chain of variable length made up of twenty amino acids, also called residues.
These are denoted by 20 characters (A-Alanine, G-Glycine, I-Isoleucine, L-Leucine, P-Proline,
V-Valine, F-Phenylalanine, W-Tryphtophan, Y-Tyrosine, D-Aspartic Acid, E-Glutamic Acid, RArginine, H-Histidine, K-Lysine, S-Serine, T-Threonine, C-Cystene, M-Methionine, N-Asparagine,
Q-Glutamine).
Each amino acid has the same core structure (backbone), consisting of alpha carbon atom
C α , connected to an amino group NH 2 , carboxyl group COOH and hydrogen atom H .
The backbone is identical in all amino acids, while the variable group, called side chain, which is also attached to alpha carbon C α , is always different and determines the amino acid, including its chemical and mechanical properties.
Amino acids are attached to each other by a covalent bond, known as peptide bond (carboxyl group COOH of one amino acid and the amino group NH 2 of the other amino acid combine, releasing water molecule H 2 O and create a peptide bond).
In this work, as is commonly done, we define protein 3D structure specified only by the C α atoms of amino acids.
The protein inverse folding task is to draw a sequence from the true distribution of n -length sequences of amino acids Y ∈ { 1 , . .
, 20 } , conditioned on a fixed protein structure, such that the designed protein folds into that structure.
The protein structure can be represented as an attributed graph G =
( V, E ) with node features V = { v 1 , . . .
, v N } , describing each residue and edge features
E = { e ij } , capturing relationships between them.
Thus, the final conditional distribution we are interested in modeling is: P ( Y | X ) =
p ( y i , . .
, y n | X ) , which is known as computational protein design task.
Protein structures are intrinsically dynamic and each structure thus possess high designability, i.e. the total number of amino acid sequences that can fold to a target protein structure is high, without losing stability of the structure.
The highly designable structures always enjoy beneficial properties such as higher thermodynamic stability, mutational stability, fast folding, functional robustness, etc.
Therefore, we need to learn a “soft” function that can model this high designability associated with a protein structure, i.e. generating diverse sequences for a given protein structure.
B A LPHA F OLD M ODEL O VERVIEW
A schematic overview of AlphaFold model is shown in Fig.
11, which it takes as input a protein sequence and produces as output, among others, the predicted 3D structure, as well as the confidence estimates of its prediction, pTM and pLDDT, which measure the estimated confidence of how well the predicted and ground truth structures match.
[Figure]
C AFD ISTILL T RAINING
Tables 3, 4 show the validation performance of AFDistill trained on each of the (p)TM-based and (p)LDDT-based datasets, respectively.
Table 5 shows results on (p)LDDT chain-based datasets.
Note that (p)LDDT chain is the dataset, similar to (p)TM datasets, where for each sequence we associate a single scalar, in this case the average of all the per-reside (p)LDDT values.
[Table] [Table]
[Table]
D GVP T RAINING D ETAILS
An example of GVP training progress regularized by the structure consistency (SC) score computed by the AFDistill model (pre-trained on various (p)TM-based datasets) is shown in Fig.
This figure shows that although SC score may be less accurate on the absolute scale, on the relative scale we can see it accurately detecting decays and improvements in the sequence quality as the
GVP trains.
Similarly, in Fig.
13 we show scatter plots of estimated pTM versus true TM scorefor
GVP-generated protein sequences regularized by SC score
[Figure]
GVP-generated protein sequences.
The last plot on the right shows the original, unregularized GVP training, where SC score was computed but never applied as part of the loss.
It can be seen that SC correlates well with the TM score for TM 42K, while for others (TM augmented 86K and pTM synthetic 1M datasets) it tends to underestimate true
TM score.
Therefore, SC score may be less accurate on the absolute scale, while on the relative scale we can see that it can accurately detect decays and improvements in the sequence quality as the GVP trains.
And the latter is of particular importance for SC to be a regularization loss during training, since it can clearly identify the ill-generated protein sequences early in the training (lower
SC scores)
and recognize well-defined sequences later during the training (higher SC scores).
D. E FFECT O F U SING AFD ISTILL T RAINED F ROM S CRATCH
We also experimented with AFDistill models trained from scratch (as opposed to starting from pre-trained ProtBert), but observed worse performance.
As an example, we trained AFDistill from scratch on TM42K dataset.
The validation CE loss during distillation was 1.5 (versus 1.1 when using pre-trained ProtBert model).
Moreover, training of AFDistill model from scratch takes longer (3 days vs 1 day).
When regularizing GVP with AFDistill from scratch, we get similar recovery rate (39.4 vs 39.6) but lower sequence diversity (15.9 vs 21.1), which confirms the benefit of common practice of fine-tuning the pretrained models as opposed to starting from random models weights.
D. E FFECT OF S TRUCTURE C ONSISTENCY
S
O N GVP P ERFORMANCE
For protein design (e.g., using GVP as a base model) the objective is CE + SC (cross-entropy +
AFDistill structure consistency score).
In Fig.
14 we present the effect of SC magnitude on the GVP performance on the test set of CATH dataset.
As can be seen, when only the CE term is present (the blue left most bar in both panels, representing the original GVP), the model is encouraged to recover the specific ground truth protein sequence for a given 3D structure, and this promotes model accuracy, and high amino acid recovery rate, while also resulting in low diversity.
On the other hand, when only the SC term is present (the pink right most bar, reprenting CE+32*SC, i.e., when SC completely dominates and CE can be ignored), this results in poor and degenerated protein sequences.
This is expected, since AFDistill alone cannot guide GVP which sequence it should generate to match the given input 3D structure.
Recall, that AFDistill has no information about the structure, and since many of the relevant protein sequences can have high pTM/pLDDT, all of them could be good candidates, and this promotes high diversity and low recovery.
Consequently, when both CE and SC terms are present and when appropriate balance between them is found (in our case
[Figure]
[Figure]
E A DDITIONAL P ERFORMANCE C OMPARISONS ON GVP E. ESM-IF
In this Section we compare GVP performance under different training scenarios and present the results in Table 6.
The first row is the recovery rate of the original GVP model as reported in Jing et al.
The following three rows are the results presented in the work of Hsu et al.
(ESM-IF).
Their evaluation showed that the vanilla GVP achieved a slightly higher recovery rate of 42.2.
GVP+AlphaFold2 represents the GVP trained on augmented dataset (CATH
+ AlphaFold2generated structure/sequence pairs).
Interestingly, this simple data augmentation baseline showed worse performance as compared to the original GVP, and the authors had to significantly increase
GVP capacity (from 1M to 21M)to get any benefit from the data augmentation.
Moreover, note that such a data augmentation idea can also serve as the baseline for our approach of AFDistill regularization, since AFDistill was trained on AlphaFold2-generated data and it can be thought of as a compressed representation of that data.
The last two rows show our evaluation results of the vanilla GVP, achieving slightly lower base recovery rate of 38.6, while this same GVP but trained with AFDistill regularization achieves a small boost in recovery (39.6) and significant increase in the sequence diversity (as we showed in
Fig. 7).
Therefore, comparing data augmentation and model distillation for the task of protein design, we see that for the original GVP model, AFDistill offers a clear advantage (compare third and last rows in Table 6), providing a modest boost in recovery, while significantly increasing diversity of the generated sequences.
Moreover, the distillation overhead is amortized, as we train AFDistill once and use it in many downstream applications.
The data augmentation would require additional computational cost in every downstream application.
On the other hand, our AFDistill regularization was able to improve the original GVP (with modest recovery and significant diversity gains).
Moreover, the distillation overhead is amortized - train once use everywhere.
We paid extra cost to train AFDistill but then applied it in many downstream tasks with little overhead.
Data augmentation would require additional significant cost in every application.
Finally, AFDistill offers unique advantage - it maintains recovery and structural consistency of the original model while introducing diversity into the generated sequences.
E. P ROTEIN MPNN
In this section we present additional experimental results on ProteinMPNN model Dauparas et al. We used the publicly available code repository and data from https://github.com/dauparas/ProteinMPNN. We compared the results of original unmodified training of ProteinMPNN to the training using SC regularization (AFDistill model trained on TM aug 86K dataset). We evaluated both training regimes on amino acid recovery, protein sequence diversity and perplexity.
We also varied ProteinMPNN internal parameter, which adds noise to the input backbone protein structure.
Table 7 presents the results.
We can see that Structure Consistency (SC) regularization maintains recovery and perplexity rates while improving the diversity of the generated protein sequences.
Backbone noise, which is a part of ProteinMPNN model, can also be seen as a form of regularization, however while the increase in noise leads to improved sequence diversity it also leads to the decrease in amino acid recovery rate.
SC regularization, on the other hand, promotes diverse sequences and maintains similar recovery rate.
[Table]
[Table]
generated protein sequences.
Backbone noise is part of ProteinMPNN model and it can be seen that the increase in noise leads to less accurate amino acid recovery, while it improves sequence diversity.
SC regularization, on the other hand, leads to more diverse sequences, while maintaining same recovery rates.
F AFD ISTILL E VALUATION ON D OWNSTREAM A PPLICATIONS
Finally, in Tables 8 9, and 10 we show detailed results for GVP and Graph Transformer inverse folding task as well as protein infilling task.
The table combines all the choices for AFDistill pretraining, showing their validation accuracy, and presents the corresponding performance on the downstream application without (top row in each table) and with SC regularization (all the following rows).
[Table]
PLDDT1M .
+ 1.
+ 3.
+ 10.
LDDT chain 42K – – 3.
+ 1.
+ 3.
+ 10.
[Table]
1M – – 3.
+ 1.
+ 3.
+ 10.
[Table]
balanced 1M – – 2.
[Table]
balanced 10M – – 2.
pLDDT balanced 60M – – 2. pLDDT chain
[Table]
[Table] [Table]
pLDDT1M – – 3.
+ 1.
+ 3.
+ 10.
LDDT chain 42K – – 3.
+ 1.
+ 3.
+ 10.
pLDDT chain 1M – – 3.
+ 1.
+ 3.
+ 10.
pLDDT balanced 1M – – 2.
pLDDT balanced 10M – – 2.
pLDDT balanced 60M – – 2.
pLDDT chain balanced 1M – – 2.
pLDDT chain balanced
[Table]
[Table] [Table]
pLDDT1M – – 3.
+ 1.
+ 3.
+ 10.
LDDT chain 42K – – 3.
+ 1.
+ 3.
+ 10.
pLDDT chain 1M – – 3.
+ 1.
+ 3.
+ 10.
pLDDT balanced 1M – – 2.
pLDDT balanced 10M – – 2.
pLDDT balanced 60M – – 2.
pLDDT chain balanced 1M – – 2.
pLDDT chain balanced 10M – –
[Table]