AlphaFold Distillation for Improved Inverse Protein Folding
Authors: Igor Melnyk
Aurelie Lozano
Payel Das
Vijil Chenthamarakshan
Venue: Submitted to ICLR 2023 (2023)
Keywords: Inverse Protein Folding Design
Protein Design
Model Distillation
AlphaFold
Protein Folding

Abstract: Inverse protein folding, i.e., designing sequences that fold into a given three-dimensional structure, is  one of the fundamental design challenges in bio-engineering 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.

Under review as a conference paper at ICLR 2023
ALPHAFOLD DISTILLATION FOR IMPROVED INVERSE PROTEIN FOLDING
Anonymous authors Paper under double-blind review
ABSTRACT
Inverse protein folding, i.e., designing sequences that fold into a given three- dimensional structure, one of the fundamental design challenges in bio- engineering and drug discovery. Traditionally, inverse folding mainly involves learning from sequences that have an experimentally resolved structure. However, the known structures cover only 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 component of inverse folding approach offers the potential of sig- nificantly improving the inverse folding, as the folding model can provide feed- back any generated sequence in the form of the predicted protein structure a structural confidence metric. However, at present, these forward folding mod- els are still prohibitively slow be part the model optimization loop during training. In this work, we propose to perform knowledge distillation on the fold- ing model's confidence metrics, e.g.. pTTM or PLDDDT scores, to obtain a smaller, faster and end-to-end differentiable distilled model, which then can be included part of the structure consistency regularized inverse folding model training. More- over, regularization technique is general enough and can be applied in other design tasks, e.g., sequence- based protein infilling. Extensive experiments show clear benefit of our method 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.
INTRODUCTION
To date, 8 out of 10 top selling drugs are engineered proteins (Arnum 2022). For functional protein design, it is often pre-requisiee that the designed protein folds into specific three- dimensionl structure.Thr fundamental task of designing novel amino acid sequences that will fold into the given 3D protein structure 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 land- scape 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 physics-based based scoring function (Kuhlman al. 2003). recent years, deep generative models have been proposed to solve this task, which consist of learning mapping from protein structure to sequences (Jing et al. 2020 Cao et al.) 2021 Wu al. [02T Karimi al., 2020) Hsu et al. 2022 Fu & Sun 2022). These approaches frequently 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 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 re- cently in protein representation learning (Rives et al. 2021 Zhang al.] 2022). protein structure

---

Under review as a conference paper at ICLR 2023
Training
Inference
Structure Consistency (SC) score
Final loss
Predicted Protein Sequence
XX
loss
Input 3D Structure
Inverse Foldiigg Model
Predicted Protein Sequence
Proposed AFDistil Model
Inverse Folding Model
Higher Fidelity, Slow
Predicted
TM LOOT
3D Structure
AlphaFold
Fidelity,
Input
2TM, PLOOT
3D Structure
Lower Fidelity, Fast (Proposed)
pTM,
Overview of the proposed system. The traditional inverse protein folding (designing sequences that fold into given 3D structures) is augmented by our proposed AFDistill model to increase the diversity of generated sequences while maintaining consisteney 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 compute metrics such as TM or LDDT and, finally, regularize the original training loss (usually cross-entropy (CF)). 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 maintain the comparable accuracy of the AlphaFoll-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.
Figure I:
prediction from sequences (Jumper et al. 2021 Back et al., 2021b), as well as in conditional pro- tein sequence generation (Das et al. 202 AAnihhhhenko et al., 2021). These lines of works have largely benefited by learning from millions of available protein sequences (that may or may not have 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 exam- ple, it has been observed that structural and functional aspects emerge from representation learned on broad protein sequence data (Rives et al. [021]. In contrast, inverse protein folding has mainly focused on learning from sequences that do have an experimentally resolved structure. Those re- ported 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 terms of amino acid recovery (Hsu al.) 2022). However, such large-scale training from scratch is computa- tionally 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 ob- jective that consists of regular sequence reconstruction loss, augmented with an additional structure consistency loss (SC) (see Fig. [orr the system overview). The straightffrward way of implementing this would be to use forward protein folding models, e.g., AlphaFold, estimate the structure from generated sequence, compare it with ground truth and compute TM score to regularize the training. However, challenge in using Alphafold (or similar) directly the computational cost associated with its inference(see Fig. (2) as well as the need of ground truth reference structure. Internal confi- dence structure metrics from the forward folding model can 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 consis- tency (SC) loss). The main properties of AFDistill model are that is fast, accurate and end-to-end differentiable. (ii) Perform extensive evaluations, where the results standard meuuture--uide quence design benchmarks show that our proposed system outperforms existing basclines in terms of lower perplexity and higher amino acid recovery, while maintaining closeness original protein
2

---

Under review as a conference paper at ICLR 2023
AFDistill built 42K
AF Disisill built PLDDT bal 6GM
163
25c
255
1.0
1.0
niernnce tme (recondrs
123
10s
0.5
Predcztd 81201
0.8
104
0.7
Prediccdd
2014
0.3
0.0
0.025
2011
101
0.3
Hipharedd
Opentelo
Ornggar EShiold
0.6
p.aOT
Figure 2: Average inference time for protein sequence of length 500 by our AFDistill model ver- sus the alternatives (left plot). Note that AFDistill fast performance holds for longer sequences: sequence length 1024- inference time 0.028s, sequence length 2048 inference time 0.035s. The timings for AlphaFold and OpenFold (Ahdritz et al. (222I) exclude the MSA searching time, which can vary from a few minutes to few hours. Note that the values for HelixFold (Fang et al. 2022). OmegaFold (Wu et al. 2022) and ESMFold (Lin et al. 2022) are taken from their respective pub- lished results. As can be seen, the inference time of the current alternatives is too slow, which makes them impractical to be included as part of the model optimization loop. On the other hand, AFDistill is fast, accurate and end-to-end differentiable. The middle plot shows true and AFDistill predicted TM scores TM distillation datasets (Pearson's correlation is 0.77) (see Section (3 for details). The right plot shows similar scatter plot of the (averaged) true and the predicted pLDDT values on the PLDDT distillation dataset (Pearson's 's correlation is 0.76).
structure, More interestingly, we improve diversity in the designed sequences, one of the main goals of protein design. As result of trade-off between sequence recovery 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 'sequence capacity' exceeding 1023 (Tian & Best, 2017]. Note that our regularization technique is not limited to the inverse folding design and, as we show, can applied to other applications, such as sequence-band observe performance improvement over the baseline. (iii) Finally, the estimated structure consis- tency 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. (021I) which would benefit from structural consistency estimation of the designed protein sequence, or as a cheap surrogate of AlphaFold that provides scoring of given protein, reflecting its structural content.
protein infilling, where we also
2 RELATED WORK
Forward Protein Folding. Several powerful computational approaches have recently been pro- posed for the task of forward folding, namely predicting the structure of a protein given its sequence. In particular, AlphaFold (Jumper al., 2021 l) makes use of multiple sequence alignments (MSAs) along with pairwise features. Building on similar ideas, RoseTTAFoll (Back et al. (222aa) is three- track 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 repro- duction 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. There- fore, MSA-free approaches such as OmegaFold (Wu et al., 2022), HelixFold (Fang et al. [0022) and ESMFold (Lin et al.) (022) have been subsequently proposed. These architectures make use of large pretrained language models to produce predictions that are more accurate than AlphaFold and Fold 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 de- sired 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., (022) introduced a deep neural network architecture that explicitly models side-chain conformers in strueture-based context. Instead of making use of
3

---

Under review as a conference paper at ICLR 2023
Training
Proposed AFDssill Model
PTM. PLDDT
Logits
CE loss
Claspes
Discretizer
AlphaFold
Protein Sequence
Inference
Input Sequence
Proposed
AFDissill
Structure Consistency (SC) score
3D structure
3D structure
pTM, PLDDT
Figure 3: Distillation overview. Top diagram shows the training of AFDistill. The scores from AlphaFold's confidence estimation are denoted as pTM and pLDDT, while the scores which are computed using ground truth and the AlphaFold's predicted 3D structures are denoted TM and LDDT. These values are then discretized and treated as class labels during cross-entropy (CE) train- ing. Note that the training based TM/LDTT is limited since the number of known ground truth structures is small. The bottom diagram shows the inference stage of AFDistill, where for each protein sequence it estimates pTT and pL.DTT scores.
trRosetta, (Jendrusch et al. 2021) presented a framework that incorporates the structure prediction algorithm AlphaFold (Jumper et al. 2021) into optimization loop to generate protein sequences. However, as noted above, the inference of AlphaFold is expensive due to MSA search, hence em- ploying 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) 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 student model that 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.
ALPHAFOLD DISTILL
Knowledge distillation (Hinton al. 2013) transfers knowledge from 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 they are harder to collect or unavailable, but rather the model's predictions (PTM/pL.DTT) and the hard labels, TM/LDDT scores, computed based AlphaFold's predicted 3D structures.
3.1 SCORES TO DISTILL
TM-score (Template Modeling score) Zhhang & Skolnick 2004), the mean distance between struc- turally aligned C., atoms scaled by length-dependent distance parameter. LDDT (Local Distance Difference Test) (Mariani et al. (203)) is the average of four fractions computed using distances be- tween all pairs of atoms based on four tolerance thresholds (0.5A, 1A. 2A, 4A) within 15A inclusion radius. The range of both metrics is (0,1), and the higher values represent more similar structures.
pTM and PLDDT are the AlphaFold-predicte sponding to the reconstructed 3D protein structure, which represent model's 's confidence of the es- timated structure. PLDDT is local per-residue score (pLDDT chain is another score that simply averages per-residue pDDDTss across the chain), while p7M is global confidence metric for assess- ing 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
metrics for a given input protein sequence, corre-
3.2 DATA
Using Release (January 2022) of AlphaFold Protein Structure Database (aardd et al.) 2021). collected set of 907.578 predicted structures. Each of these predicted structures contains 31>
4

---

Under review as a conference paper at ICLR 2023
Releuse 3 (January 2022)
Releuse 4 (July 2022)
Name
Size
Name
Size
Original
907,578
Original
214,687,408
TM 42K
42,605
PLDDT balanced 1M
1,000,000
'M augmented 86K
86,811
PLDDT balanced 10M
10,000,000
PTM synthetic IM
1,244,788
PLDDT balanced 60M
66,736,124
42K
42,605
PLDDT IM
905,850
Table Statistics from January 2022 (left side) and July 2022 (right size) releases of the AlphaFold database. For the earlier release, we created multiple datasets for pT and pLDDT estimation, while for the later, larger release we curated datasets only for pLDDT estimation.
coordinates of all the residue atoms as well as the per-resiude pL.DIT confidence scores. To avoid any data leakage to the downstream aplications, 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, created PLDDT 1M dataset (see Table D 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.
12.5
TM 42K
LDDT 42K
86K
IM
synthetic
PLDOT balanced
Density
7.5
PLDOT balanced 10M
balanced 6OM
5.0
0.0
0.0
0.2
1.0
TM scores
(P]LDDT scores
Figure 4: Distribution of the (p)TM/(P)LDDT scores in various datasets used in AFDistill training.
We also created datasets which are based on the true TM and LDDT values using the predicted Al- phaFold structures. Specifically, using the PDB-to- PDB-to-UniProd mapping, selected subset of sam- ples with matching ground truth PDB sequences and 3D structures, resulting in 42,605 structures. We denote these datasets 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 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 T'M-score, and keeping the low and medium range TM values. p'TM 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 also shown in Fig. 4 where both TM augmented 86K and pTM synthetic 1M datasets less skewed, covering lower (p)TM values better.
additional TM-based
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-pI.DDT values (also called pLDDT chain). The resulting dataset contains 60M sequences, for which we additionally created 10M and IM versions, see Fig. for their density.
3.3 MODEL
AFDistill model is based on ProtBert (Elnaggar al. 2020). Transformer BERT model (420M parameters) pretrained on a large corpus of protein sequences using masked language modeling. For task modify ProtBert head by setting the vocabulary size to 50, corresponding discretizing pTM/pLDDT in range (0,1). For pTM (scalar) the output corresponds to the first (CLS) token of the output sequence, while for DLDDT (sequence) the predictions are made for each residue position.
5

---

Under review as a conference paper at ICLR 2023
Training data
Validation CB loss
Training data
Validation CE loss
TM 42K
1.10
LDDT 42K
3.39
augmented
2.12
pLDDT IM
3.24
M synthetic IM
2.55
pl.Di>T balanced 1M
2.63
pLDDT bulanced 10M
2.43
PL.DTT balanced 60M
2.40
Table Validation loss of AFDistill on datasets from Table |
(For more details, see Tables and
2003,
diain
074
2TM (TV 43caug) 243
(TM (TM2 081
KTM (TM-
Figure 5: Examples of 3D protein structures from the dataset, corresponding to high, medium, and low actual TM scores (top row in legend), well as AFDistill predictions, trained on TM 42K (middle row) and TM augmented 86K (bottom row).
3.4 DISTILLATION EXPERIMENTAL RESULTS
In this section we present evaluation results of the model after training 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 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. Fig. 2 we show scatter plots of the true PTM scores and pL.DDT values the entire validation 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, Fig. and 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 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).
shown Table [ We
We
Uniprot ADA7SOVSN2
AGAS3REBU7
Unigret ACASWSKA3S
2.0
2.8
2.0
pLDDT
PLOTT
PLDDT AishaFald
Presicted IPLDUT bal S0M)
Presictidn (pLDDT bal IM1
Amino
index
Figure 6: Dataset examples of the per-residue predictions for two AFDistill models (blue and green lines), build on DDDD balanced 1M and 60M datasets, versus the AlphaFold predictions (red line).
INVERSE PROTEIN FOLDING DESIGN
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 for the protein sequence infilling as a means to novel antibody generation. The overall framework is presented in Fig. [ (following the green line in the diagram), where the traditional inverse folding model is regularized by our SC

---

Under review as a conference paper at ICLR 2023
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
== LCE aLsc,
where LC = Σ1 Lee (s,, (8,,8) 8,) is the CE loss, 8, is the ground truth and & is the generated protein N sequence, Lsc Di SC((s)) is the structure consistency loss, N is the number of training sequences, and ar is the weighting scalar for the SC loss, in our experiment it is set to 1.
4.1 METRICS
To measure the quality of the prediction designs, we compute the following set of sequence eval- uation 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 recov- cry computed for all pairwise comparisons
the set. While in general the recovery and diversity tend to be inversely correlated, i.c., higher recovery leads to lower diversity, and vice versa, interested in models that achieve high recovery rates and be able to maintain high protein sequence diversity. Perplexity measures the likelihood of given sequence, lower values mean better perfor- mance. Finally, for structure evaluation, we use TM-score as well as the structure consistency (SC) score, which the AFDistill's output (pTM/pLDDTT for given input.
4.2 RESULTS
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 ct al. 2013) dataset curated by (Jin al.) 2021) and focus on infilling CDR-H3 loop. The dataset has 3896 training, 403 validation and 437 test sequences.
Sowenne mezics
.324
82
5.24
Stoucuue metitis
(21.21
(16.3)
157.31
3.5
42K
134 sug 2AK
PTM Synth JM
1DOT 42K
IM
OLDUT 1M
PLBOT SOM
Figure 7: Evaluation results of GVP trained with SC regularization. The horizontal x-axis shows various datasets for AFDistill pretraining, the left vertical y-axis shows sequence metrics (recovery and diversity gains), while the right y-axis shows structure metrics (TM and SC scores). Blue and orange solid bars show recovery and diversity gains (top number percentage, bottom actual value) of SC-regularized GVP over the vanilla GVP baseline. Solid olive bar shows the predicted SC (pTM or PLDDT, depending on AFDistill model), while the purple bar is the test set TM score (structures predicted by AlphaFold), Overlaid dashed cyan bar is the TM score of the baseline GVP. We can see that overall TM 42K and TM augmented 86K pretrained AFDistill achieve the best overall performance, with high diversity and moderate improvement in sequence and structure recovery.
GVP Geometric Vector Perceptron GNNs (GVP) (Jing et al. 2020 is the inverse folding model, that for 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 geomet- ric 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, Ta- ble 8 for additional results). Baseline GVP with regularization achieves 38.6 in recovery, 15.1 in diversity and 0.79 in TM score on the test set. can be seen that there a consistent improve- ment in sequence recovery gain (1-3%) over the original GVP and significant diversity gain (up to
7

---

Under review as a conference paper at ICLR 2023
augmented
POB ID laal. chain
POB ID laal. chain
Figure 8: Comparison between baseline GVP (left) and SC regularized GVP (right), where AFDistill was pre-trained on TM augmented 86K dataset. The tables display 15 generated protein sequences from each model, top red row the ground truth. Green cell with the * indicates amino acid identity compared to ground truth (top red row), while blue cell shows novelty. The shade of blue cell represents the frequency of the amino acid in that column (darker + more frequent, lighter > rare). Therefore, the method with high recovery and diversity rates will have many green and light blue cells. Bottom plots show AlphaFold estimated structures (green) and the ground truth (red). 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 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 p'Dsd SC scores had overall better influence performance as compared to pI.DDT-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-struce structure relationship and provides more performance boost to the inverse folding model.
the model
1.
S/N
Dvera STMiUUDOT
Orggenze mettee
200
sore
Structure metbrie
Ghara
52
135
43K
TH aug 9M
syrth 2M
L9BT 43K
PLOOT
SUDOT
SOM
Figure 9: Evaluation results Graph Transformer model trained with SC score regularization.
To further illustrate the effect of recovery and diversity, show in Fig. protein sequences and AlphaFold -ewerrreedttddd 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
8

---

Under review as a conference paper at ICLR 2023
Recovary gain Diveriily gain pTMiSLDOT
0.3%
2.85
2.00
Structue mupricc
TM score
Ind (TH c B5]
Sequerce
LAX
HL.
5.9
TM
wuy
5TM synth
LOOT
PLOOT
PLUOT hal
1LOOT bat
Figure 10: Evaluation results of protein infilling model trained with SC regularization.
recovery 40.8 and diversity is 11.2 for GVP, while for GVP+SC it is 42.8 and 22.6, respectively, confirming that GVP+SC achieves higher recovery and diversity. The bottom plots show AlphaFold estimated structures (green) and the ground truth (red). It can be seen that for GVP+SC, the high sequence diversity still results in very accurate reconstructions (for this example, average TM score is 0.95), while GVP alone shows more inconsistencies, marked with blue arrows, (TM score 0.92).
Graph Transformer We evaluated the effect of SC score Graph Transformer WW al. 202 another inverse folding model, which seeks to improve standard GNNs to represent the protein 3D structure. Graph Transformer applies 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 for additional results). Baseline model with no regularization has 25.2 in recovery, 72.2 diversity and 0.81 in TM score on the test set. As compared to GVP (Fig. 77) we can see that for this model, the recovery and diversity gains upon SC regularization are smaller. We also see that TM score regularized model (TM 42K and TM augmented 86K pretraining) is slightly higher compared to pLDDT-baat -baeed models.
Protein Infilling Our proposed structure consistency regularization quite general and not lim- ited 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 given structure, infilling focuses on filling specific regions of protein conditioned on a se- quence/structure template. The
regions (CDRs) of an antibody pro- tein 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 con- ditioned on a fixed framework region. We focus on CDR-H3 and use a baseline pretrained protein model ProtBERT (Elnaggar et al. (020)) finetuned on the infilling dataset, and use ProIBERT+SC as an alternative (finetuned with SC regularization). The CDR-H3 masked and the objective is to reconstruct it using the rest of the protein sequence as a template. The results are shown in Fig. IO (see also Appendix, Table for additional results). Baseline model achieves 41.5 in recovery, 14.5 in diversity, and 0.80 in TM score the test set. Similar for the other applications, 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.
an
CONCLUSION
In this work we introduce AFDistill, distillation model based on AlphaFold, which for given protein sequence estimates its structural consistency (SC: p'TM or PLDDT) score. We provide ex- perimental results to showcase the efficiency and efficacy of the AFDistill model in high-quality protein sequence design, when used together with graph neural net based inverse folding model or large protein language model for sequence infilling. Our AFDistill model is small and accurate enough that it can be efficiently used for regularizing structural consistency in protein optimiza- tion tasks, maintaining sequence and structural integrity, while introducing diversity and variability in the generated proteins.

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
For the sake 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.
REFERENCES
Gustaf Ahdritz, Nazim Bouatta, Sachin Kadyan, Qinghui Xia, William Gerecke, and Mo- hammed AlQuraishi. OpenFold, 11 2021. htRt:::s///ttttmmm cttp//aaaaaaaaaarrrr......o/////
openfold
Namrata Anand, Raphael Eguchi, Irimpan Mathews, Carla Perez, Alexander Derry, Russ Altman, and Po-Ssu Huang. Protein sequence design with learned potential. Nature communi- cations, 13(1) 2022.
Ivan Anishchenko, Samuel Pellock, Tamuka M Chidyausiku, Theresa A Ramelot, Sergey Ovehin- nikov, Jingzhou Hao, Khushboo Bafna, Christoffer Norn, Alex Kang, Asim K Bera, et al. De novo protein design by deep network hallucination. Nature, 600(7889):547-522 2021.
Patricia Van Arnum. Biologics, 2022.
Pharma
Pulse: Top-Selling
Small Molecules and
URL
https: www. www.tcc..tticeeeeeeeeeooo fearrg/feaurrs////
pharma pu 1 se top-top---- se Ling- smal
-mool wwwlecules-biollln.
Minkyung Back, Frank DiMaio, Ivan Anishchenko. Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, and R Dustin Schaeffer. Accurate prediction of pro- tein structures and interactions using a three-track neural network. Science, 373(6557) 2021a.
Minkyung Back, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov. Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using three-track neural network. Science, 373(6557):871 876, 2021b.
Yue Cao, Payel Das, Vijil Chenthamarakshan. Pin- Yu Chen, Igor Melnyk, and Yang Shen. Fold2seq: A joint sequence (Id)-fold (3d) embedding- based generative model for protein design. In Inter- national Conference on Machine Learning, pp. 1261-1271. PMLR, 2021.
Payel Das, Tom Sercu, Kahini Wadhawan, Inkit Padhi, Sebastian Gehrmann, Flaviu Cipcigan, Vijil Chenthamarakshan, Hendrik Strobelt, Cicero Dos Santos, Pin-Yo Chen, et al. Accelerated an- timicrobial discovery via deep generative models and molecular dynamics simulations. Nature Biomedical Engineering, 5(6):613 0at::613-2233 623, 2021.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel, et al. Robust deep learning- based protein sequence design using proteinmpnn. Science, 378(6615):49956 2022.
James Dunbar, Konrad Krawezyk, Jinwoo Leem, Terry Baker, Angelika Fuchs, Guy Georges, Jiye Shi, and Charlotte M. Deane. SAbDab: the structural antibody database. Nucleic Acids Research, 42(D1), 11 2013.
Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. ProtTrans: towards crack- ing the language of life's code through self-supervised deep learning and high performance com- puting, arXiv preprint arXiv:2007.022. 06225, 2020.
Xiaomin Fang, Fan Wang, Lihang Liu, Jingzhou He, Dayong Lin, Yingfei Xiang, Xiaonan Zhang, Hua Wu, Hui Li, and Le Song. Helixxold-Single: Msa-free protein structure prediction by using protein language model as an alternative. arXiv preprint 2xXxx:2207.13221 2022.
Tianfan Fu and Jimeng Sun. Sipf: Sampling method for inverse protein folding. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 378-388, 2022.
10

---

Under review as a conference paper at ICLR 2023
Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. In NIPS Deep Learning and Representation Learning Workshop, 2015. URL [RLLtt:///aaxi org/absa 1503. 02531
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie. Tom Sereu, Adam Lerer, and Alexan- der Rives. Learning inverse folding from millions of predicted structures. bioRxiv, 2022.
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. Advances neural information processing systems, 32, 2019.
Michael Jendrusch, Jan O Korbel, and S Kashif Sadiq. AlphaDesign: A de novo protein design framework based on AlphaFold. bioRxiv, 2021.
Wengong Jin, Jeremy Wohlwend, Regina Barzilay, and Tommi Jaakkola. Iterative refinement graph neural network for antibody sequence-structure co-design. arXiv preprint arXiv:2110.04000 2021.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.00 0/10/ 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, and Anna Potapenko. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873) 2021.
Mostafa Karimi, Shaowen Zhu, Yue Cao, and Yang Shen. De novo protein design for novel folds using guided conditional wasserstein generative adversarial networks. Journal of chemical infor- mation and modeling, 60(12):5667-5681, 2020.
Brian Kuhlman, Gautam Dantas, Gregory C Treton, Gabriele Varani, Barry I. Stoddard, and David Baker. Design of a novel globular protein fold with atomic-level accuracy. science, 302(5649): 1364-1368, 2003.
Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision. pp. 2980-2988, 2017.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos San- tos Costa, Maryam Fazel-Zarande Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv, 2022.
Valerio Mariani, Marco Biasini, Alessandro Barbato, and Torsten Schwede. IDDT: local superposition--free score for comparing protein structures and models using distance difference tests. Bioinformatics, 29(21):2722 27288, 2013.
Lewis Moffat, Joe G Greener, and David Jones. Using AlphaFold for rapid and accurate fixed backbone protein design. bioRxiv, 2021.
Christoffer Norn, Basile IM Wicky, David Juergens, Sirui Liu, David Kim, Brian Koepnick, Ivan Anishchenko, David Baker, and Sergey Ovchinnikov. Protein sequence design by explicit energy landscape optimization. bioRxiv, 2020.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Out, C Lawrence Zitnick, Jerry Ma, al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15):20016231 118, 2021.
Pengfei Tian and Robert B Best. How many protein sequences fold to a given structure? a coevolu- tionary analysis. Biophysical journal, 113()):119-17300 2017.
Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, Augustin Zidek, Tim Green, Kathryn Tunyasuvunakool, Stig Petersen, John Jumper, Ellen Clancy, Richard Green, Ankur Vora, Mira Lutfi, Michael Figurnov, Andrew Cowie, Nicole Hobbs, Pushmeet Kohli, Gerard Kleywegt,

---

Under review as a conference paper at ICLR 2023
Ewan Birney, Demis Hassabis, and Sameer Velankar. AlphaFold Protein Structure Database: mas- sively expanding the structural coverage of protein-sequence space with sigh-accuracy models. Nucleic Acids Research, (OKDI):D439-444, 11 2021.
Ruidong Wu, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan
Wu, Qi Xie, Bonnie Berger, Jianzhu Ma, and Jian Peng. High-resolution de novo structure pre- diction from primary sequence. bioRriv, 2022. doi: 10. 00.1100/202200021150000999999999 URL https:
h/www.bionn www.bbiioiiii..................e wwt.
Zhanghao Wu, Paras Jain, Matthew Wright, Azalia Mirhoseini, Joseph E Gonzalez, and Ion Stoica. Representing long-range context for graph neural networks with global attention. Advances in Neural Information Processing Systems, 34: 34:13266-13279 2021.
Jianyi Yang. Ivan Anishchenko, Hahnbeom Park, Zhenling Peng, Sergey Ovchinnikov, and David Baker. Improved protein structure prediction using predicted interresidue orientations. Proceed- ings of the National Academy of Sciences,
2020.
Yang Zhang and Jeffrey Skolnick. Scoring function for automated assessment of protein structure template quality. Proteins: Structure, Function, and Bioinformatics, 57(4):702-770, 2004.
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining, arXiv preprint 2:Xi::2003.06/25, 2022.
12

---

Under review as a conference paper at ICLR 2023
A BACKGROUND ON PROTEIN DESIGN
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, T-Phenylalanine, W-Tryphtophan, Y-Tyrosine, D-Aspartic Acid, E-Glutamic Acid, R- Arginine, H-Histidine, K-Lysine, S-Serine, T-Threonine, C-Cystene, M-Methionine, N-Asparagine, Q D-Glutamine). Each amino acid has the same core structure (backbone), consisting of alpha carbon atom C.. connected an amino group H2, 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 Co. always different and determines the amino acid, including its chem- ical and mechanical properties. Amino acids are attached to each other by covalent bond, known as peptide bond (carboxyl group COOH one amino acid and the amino group N H2 of the other amino acid combine, releasing water molecule H2O and create peptide bond). In this work, as is commonly done, we define protein 3D structure specified only by the Ca atoms of amino acids.
The protein inverse folding task is draw sequence from the true distribution of n-length quences of amino acids Y E (1...., 20}, conditioned on a fixed protein structure, such that the designed protein folds into that structure. The protein structure graph G = (V, E) with node features V = {uI..... EN J. describing each residue and edge features E {eiy] capturing relationships between them. Thus, the final conditional distribution are interested in modeling is: P(Y[X) = p(y....., Nn X),, which is known as computational protein design task.
be represented
attributed
Protein structures are intrinsically dynamic and each structure thus possess high designability, i.c. the total number of amino acid sequences that can fold 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, need to learn "soft" function that can model this high designability associated with a protein structure, i.c generating diverse sequences for a given protein structure.
B ALPHAFOLD MODEL OVERVIEW
schematic overview of AlphaFold model is shown in Fig. I1 which takes as input protein sequence and produces output, among others, the predicted 3D structure, well as the confidence estimates of its prediction, pTM and pLDDT, which measure the estimated confidence how well the predicted and ground truth structures match.
13

---

Under review as a conference paper at ICLR 2023
MSA
PLDDT Head
Input Protein Sequence
EvoFormer Module
Structure Module
pTM Head
Template
3D Structure
Recycling
Figure 11: Overview of the inference stage in AlphaFold model. Given an input protein sequence, first the search is performed in genetic database find similar sequences and construct multiple sequence alignments (MSA). Then a structure database search is done to find similar 3D structures and construct templates. The MSA and templates fed into EvoFormer module, whose output is then sent to the Structure module, which is finally completed with the multiple output heads. The 3D structure head generates predicted 3D protein structure, while pLDDT and pTM heads estimate the confidence of the computed structure. Optionally, the generated structure together with the intermediate states are recycled and sent back to update/correet MSA and template representations for further processing and improvement.
C AFDISTILL TRAINING
Tables 4 show the validation performance of AFDistill trained on each of the (p)TM-base -assed and (pysDDDTbased datasets, respectively. Table 5 shows results (p)LDDT chain-based datasets. Note that (p)LDDT chain the dataset, similar to (p)TM datasets, where for each sequence we associate single scalar, in this case the average of all the per-reside (p)LDDT values.
14

---

Under review as a conference paper at ICLR 2023
Training
Validation
Data
Weighted sampling
Focal loss (7)
CE loss
1.33
1.37
TM 42K
+
1.0
1.16
3.0
1.10
+
10.0
1.29
2.12
1.0
2.15
TM augmented 86K
+
3.0
2.19
+
10.0
2.25
2.90
1.0
2.75
рTM synthetic 1M
+
3,0
2.55
10.0
3.20
Table Validation CE loss for the AFDistill model trained on each of the (p)TM-basedd datasets. To address data imbalance during training, employed weighted sampling for minibatch generation to so that the TM-scores cover their range (0.1) close to uniform distribution. Moreover, we also used Focal loss (Lin al. 2017) in place of the standard cross-entropy (CE) loss (the evaluation is still done using CE loss across all the training setups). Based on the validation loss, we see that the AFDistill model trained on TM 42K dataset performed the best, followed by the dataset with augmentations, and the synthetic performed the worst. We also see that weighted sampling and focal loss do help addressing the data imbalance problem, although for TM augmented 86K, the balanced augmentation seemed to help better and the best performance was for the case when no weighted sampling is applied and the traditional CE loss is used. As shown in Section 4 the valida- tion performance the distillation data may not always indicate the performance on the downstream applications, where in particular we observed that the Distill model, trained on TM augmented 86K dataset, overall performed better than TM 42K, while having slightly worse validation CE loss.
Training
Validation
Data
Weighted sampling
Focal loss (r)
CE loss
3.47
1.0
3.44
LDDT 42K
3.0
3.42
10.0
3.39
3.27
1.0
3,28
PLDDT IM
3.0
3.25
10.0
3.24
Table Validation CE loss for AFDistill trained on each of the (p).DDT-bbsed datasets. We see that weighted sampling coupled with Focal loss, performed the best.
15

---

Under review as a conference paper at ICLR 2023
Training
Validation
Data
Weighted Sampling
Focal loss (7)
CE loss
3.69
1.0
3.57
LDDT chain 42K
3.0
3.63
10.0
3.59
3.29
1.0
3.36
PLDDT chain 1M
3.0
3.30
10.0
3.31
PLDDT chain balanced 1M
2.45
PLDDT chain balanced 10M
2.24
pLDDT chain balanced 60M
2.21
Table 5: Validation CE loss for the AFDistill model trained on each of the (p)LDDT chain-based datasets. LD)LDDT chain is the dataset, similar to (p)TM datasets, where for each sequence associate a single scalar, in this case the average of all the per-reside (D)LDDT values. Similar as before, that the use of weighted sampling coupled with Focal loss helps in boosting the model performance. We also see that increasing the scale of data (which is already balanced) improves the performance even further.
16

---

Under review as a conference paper at ICLR 2023
D GVP TRAINING DETAILS
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. 12 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.
TM-42K
TM augmented 86K
pTM synthetic 1M
None
0.00
2.50
0.92
0.89
0.55
2.59
0.85
0.90
0.00
0.90
0.70
0.70
2.70
0.70
0.49
score
0.55
score
2.53
score
0.05
score
score
score
0C
SCK
100K 1.5K
2006
04
50C
1,0K
150c
200K
0K
50K
100K
150k
200K
04
50C
100K
1.50K
200c
Taiiiig steps
Traning Stope
Prining steps
Training steas
Figure 12: Example of the training progress (on CATH 4.2 dataset) of the GVP model regularized by the structure consistency (SC) score computed by the AFDistill model pre-trained on different datasets. Each plot shows the results for one of the Distill pre-training datasets, where the blue line represents the SC score computed by the AFDistill model (in this case generating pTTM value), while the orange line shows the actual TM score computed between the ground truth structure and the AlphaFold's estimated 3D structures for GVP-generated protein sequences. The last plot on the right shows the original, unregularized GVP training, where SC score computed but never applicd as part of the loss. can be seen that SC correlates well with the TM score for TM 42K, while for others (TM augmented 86K and PTTM 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 can clearly identify the ill ill-generated protein sequences early in the training (lower SC scores) and recognize well-defined sequences later during the training (higher SC scores).
D.1 EFFECT OF USING AFDISTILL TRAINED FROM SCRATCH
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 I 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 11, which confirms the benefit of common practice of fine-tuning the pretrained models as opposed to starting from random models weights.
D.2 EFFECT OF STRUCTURE CONSISTENCY (SC) SCORE ON GVP PERFORMANCE
For protein design (e.g., using GVP as a base model) the objective is CE + SC (cross-entropy + AFDistill structure consistency score). In Fig. [4 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, 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//LDDT, 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
17

---

Under review as a conference paper at ICLR 2023
across training egochs
0.8-
38-
64
0.44
3A
07
22-
3.0
00
(AFDittil trained augmented 96K) across training epochs
0.6
3.6
6.5-
0.2
22
02
400 0. 0.5 60
055
1.0
3.0
0.0 0. 1.0 05
2.5
Figure 13: Estimated pTM versus TM score (based on AlphaFold structure prediction) for GVP- generated protein sequences regularized by SC score. The top row shows results for SC computed by AFDistill model trained on TM 42K, while the bottom row for AFDistill trained TM augmented 86K. The columns in each row correspond to the progress as GVP trains. Note that the top row corresponds to the first left plot Fig. [2] while the bottom row corresponds to the second plot in Fig. I2 It can be observed that in the earlier stages of GVP training, the generated protein sequences are of poor quality, reflected pTM and TM scores that are spread across the (0,1) range. On the other hand, the training progresses, the generated sequences are getting better and the pTM/TM score is concentrated more in the upper range. Another observation is that for AFDistill trained on TM 42K dataset, the predicted and true TM score better aligned across the diagonal (compare with orange and blue lines on the left plot in Fig. [2 while for AFDistill trained on TM augmented 86K dataset, PTM tends underestimate sequence scoring tool, which fairly accurately measures the structural consistency of the generated protein sequences. Combined with the fact that fast and end-to-end differentiable, shows its potential for many of the protein optimization problems.
TM score. These plots show that AFDistill viable
50
Test Recovery
10.8
37.5
Test Diversity
40
37.5
30
22.8
28.2
20
20
15.1
10
CE+1'SC
CE+2'CC
CE+4'5C
CE48ISC
CE+16TS
CE+32'5C
CE+1+SC
CE+2'SC
CE+415C
CE+8/SC
CE+16'S
CE+32'SC
GVP Training Regime
Figure 14: The effect of Structure Consistency (SC) loss on the performance of GVP. Left panel shows the amino acid recovery rate and the right panel shows the diveristy rate on the test set of CATH dataset. The horizontal y-axis shows the different choices of objective function during train- ing: CE is the cross-entropy loss, SC is the Structure Consistency score computed by AFDistill.
it is CE+SC, corresponding to the orange bar in both panels), we get a full benefit, i.c. the accurate recovery and high diversity of the generated protein sequences.
18

---

Under review as a conference paper at ICLR 2023
E ADDITIONAL PERFORMANCE COMPARISONS ON GVP
E.I ESM-IF
In this Section we compare GVP performance under different training scenarios and present the
results in T The first row is the recovery rate of the original GVP model as reported in Jing
al. (2020). The following three rows are the results presented in the work of Hssu al.] (2022)
(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 + AlphaFold2- generated 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) 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 AlphaFoldd-ggeneraee 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.
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 ( 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 consis- tency of the original model while introducing diversity into the generated sequences.
E.2 PROTEINMPNN
In this section we present additional experimental results on ProteinMPNN model Dau- paras et al.] (2022) We used the publicly available code repository and data from htt.ghhhwbb.oo/daaaaaaa//ootiiie
We compared the results of original unmodified train- ing of MPNN to the training using SC regularization (AFDistill model trained on TM aug 86K dataset). We evaluated both training regimes amino acid recovery, protein sequence diver- sity 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 gener- ated 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 diver- sity it also leads the decrease in amino acid recovery rate. SC regularization, the other hand, promotes diverse sequences and maintains similar recovery rate.
19

---

Under review as a conference paper at ICLR 2023
Model
Recovery
Change
ling al. 2026
40.2
GVP
42.2
Hsu 2022
VV AlphaFold2 data
Hsu al. 2022
-3.6 (-8.5%)
38.6
(our experiment)
GVP SC experiment)
39.6
+1.0 +1.06+))6)
Table 6: Comparison of amino acid recovery rate of protein sequences generated by GVP the test split of CATH dataset. The lighter and darker highlighted blocks represent the results from three experiments: original results from GVP authors, the results from ESM authors, and the results from our experiments. A small difference between the values in first, second and forth rows can be attributed some discrepancies in experimental settings see that simple data augmentation baseline results in 3.6 (or 8.5%) drop of recovery relative to the unaugmented GVP. On the other hand, the use of SC regularization leads 1.0 (or 2.6%) gain in recovery, signaling the benefit of the proposed distillation approach.
well model initialization. We can
Recovery
Diversity
Perplexity
ProteinMPNN
ProteinMPNN
ProteinMPNN
ProteinMPNN
ProteinMPNN
ProteinMPNN
+SC
+SC
+SC
Backbone Noise 0.02
43.7
43.5
22.5
24.3
5.1
5.1
Backbone Noise 0.1
39.8
40.0
28.1
30.4
5.3
5.4
Backbone Noise 0.2
36.5
36.2
31.3
34.4
5.8
5.8
Backbone Noise 0.3
33.3
33.1
33.0
37.8
6.2
Table Evaluation results ProteinMPNN trained with and without SC regularization (AFDistill trained on TM aug 86K dataset). Structure Consistency (SC) regularization maintains recovery and
perplexity rates while improving the diversity of the generated protein sequences. Backbone noise part of ProteinMPNN model and it can be seen that the increase in noise leads less accurate amino acid recovery, while it improves sequence diversity. SC regularization, on the other hand, leads to diverse sequences, while maintaining same recovery rates.
20

---

Under review as a conference paper at ICLR 2023
F AFDISTILL EVALUATION ON DOWNSTREAM APPLICATIONS
Finally, in Tables 8 and 10 we show detailed results for GVP and Graph Transformer inverse fold- ing task 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).
21

---

Under review as a conference paper at ICLR 2023
Distill model
GVP
Training
Validation
PTM/ PLDDT
TM
Data
Recovery
Diversity
Perplexity
Weighted sampling
Focal loss
Score
CE loss
38.6
15.1
6.1
0.80
0.79
+
1.37
36.8
22.2
6.3
0.78
TM
+
1.0
1.16
37.6
21.1
6.0
0.87
42K
+
3.0
1.10
39.6
21.1
5.9
0.84
0.84
+
10.0
1.29
37.9
18.4
6.0
0.80
2.12
38.3
22.2
5.9
0.73
TM augmented
+
1.0
2.15
39.8
22.1
5.8
0.78
0.85
86K
3.0
2.19
37.8
19.8
6.1
0.73
10.0
2.25
38.5
21.2
5.9
0.72
2.90
38.8
21.4
5.8
0.73
TM synthetic 1M
1.0
2.55
39.1
22.5
5.9
0.77
0.81
+
3,0
2.75
39.0
21.9
5.8
0.74
10.0
3.20
39.0
22.0
5.9
0.69
3,47
39.0
18.9
5.8
0.74
0.78
LDDT
1.0
3.44
38.7
22.5
5.8
0.73
42K
+
3.0
3,42
38.9
21.2
5.8
0.73
+
10.0
3.39
38.5
22.3
5.9
0.72
3.27
39.3
16.5
5.9
0.76
0.79
PLDDT IM
+
1.0
3.28
38.8
15.5
5.9
0.72
+
3.0
3.25
38.9
18.2
5.8
0.78
+
10.0
3.24
38.4
16.2
6.0
0.73
3.69
38.8
20.0
5.8
0.74
LDDT chain 42K
+
1.0
3,57
39.3
16.3
5.8
0.79
0.78
+
3.0
3.63
38.9
15.9
5.9
0.72
+
10.0
3.59
37.9
23.2
6.0
0.73
3.29
39.4
17.4
5.8
0.78
PLDDT chain IM
1.0
3.36
38.7
16.3
5.8
0.76
+
3.0
3,30
39.6
18.3
5.7
0.79
0.77
+
10.0
3,31
38.2
20.1
6.0
0.76
pLDDT balanced IM
2.63
39.1
17.1
5.8
0.75
0.82
pLDDT
2.43
39.3
17.7
5.9
0.73
balanced 10M
PLDDT
2.40
39.8
17.5
5.9
0.74
0.81
balanced 60M
PLDDT chain
2.45
38.6
16.6
5.9
0.73
balanced IM
pL.DTT chain
2.24
39.1
17.8
5.8
0.73
balanced 10M
PLDDT chain balanced 60M
2.21
39.7
17.9
5.9
0.74
0.82
Table 8: Evaluation results of GVP inverse folding task, trained without (top row) and with SC regularization (all other rows). The table combines the choices for AFDistill pretraining and showing their validation accuracy, as well as the corresponding performance on the downstream application. We select the best performance for each experiment based on the highest recovery rate (marked in bold).
22

---

Under review as a conference paper at ICLR 2023
Distill model
Graph Transformer
Training
Validation
PTM/ PLDDT
TM
Data
Recovery
Diversity
Perplexity
Weighted sampling
Focal loss
score
CE loss
25.2
72.2
7.2
0.80
0.81
+
1.37
24.1
74,4
7.4
0.81
TM
+
1.0
1.16
25.2
73.2
7.2
0.86
42K
+
3.0
1.10
25.8
74.5
7.2
0.83
0.80
+
10.0
1.29
24.9
73.9
7.3
0.81
2.12
25.0
73.3
7.1
0.78
TM augmented
+
1.0
2.15
25.9
74.0
7.1
0.80
0.82
86K
3.0
2.19
24.9
73.4
7.3
0.76
10.0
2.25
24.8
73.4
7.2
0.79
2.90
25.3
73.2
7.1
0.72
TM synthetic 1M
1.0
2.55
25.5
73.9
7.2
0.79
0.75
+
3,0
2.75
25.2
73.5
7.2
0.77
10.0
3.20
24.9
74.2
7.2
0.76
3,47
25.4
73.2
7.1
0.75
LDDT
1.0
3.44
25.7
74.2
7.1
0.72
0.76
42K
+
3.0
3,42
25.5
74.4
7.2
0.73
+
10.0
3.39
25.3
22.3
7.2
0.72
3.27
25.6
73.4
0.79
PLDDT
+
1.0
3.28
25.4
74.1
7.2
0.78
IM
+
3.0
3.25
25.6
74.3
7.1
0.79
0.79
+
10.0
3.24
25.4
74.0
7.1
0.77
3.69
25.3
74.1
7.2
0.76
LDDT chain 42K
+
1.0
3,57
25.8
74.3
7.1
0.75
0.80
+
3.0
3.63
25.5
74.2
7.1
0.77
+
10.0
3.59
25.6
74.1
7.2
0.76
3.29
25.3
74.3
7.1
0.78
PLDDT chain IM
1.0
3.36
25.2
74.1
7.1
0.76
+
3.0
3.30
25.6
74.4
7.2
0.79
0.81
+
10.0
3,31
25.3
74.3
7.1
0.77
pLDDT balanced IM
2.63
25.8
74.2
7.1
0.70
pLDDT
2.43
25.7
74.5
7.1
0.73
balanced 10M
PLDDT
2.40
26.0
74.2
7.2
0.74
0.78
balanced 60M
PLDDT chain
2.45
25.7
74.3
7.1
0.72
balanced IM
pL.DTT chain
2.24
25.9
74.4
7.1
0.74
balanced 10M
PLDDT chain balanced 60M
2.21
25.9
74.5
7.2
0.72
0.77
Table 9: Evaluation results of Graph Transformer inverse folding task, trained without (top row) and with SC regularization (all other rows). The table combines all the choices for AFDistill pretraining and showing their validation accuracy, well as the corresponding performance on the downstream application. We select the best performance for each experiment based on the highest recovery rate (marked in bold).
23

---

Under review as a conference paper at ICLR 2023
Distill model
CDR Infill
Training
Validation
PTM/ PLDDT
TM
Data
Recovery
Diversity
Perplexity
Weighted sampling
Focal loss
score
CE loss
41.5
14.5
6.8
0.80
0.85
+
1.37
41.9
15.7
6.5
0.81
TM
+
1.0
1.16
42.4
16.6
6.3
0.77
0.85
42K
+
3.0
1.10
41.7
14.6
6.7
0.78
+
10.0
1.29
40.8
14.4
6.6
0.79
2.12
42.8
15.5
6.5
0.78
0.84
TM augmented
+
1.0
2.15
41.6
14.8
6.6
0.74
86K
3.0
2.19
41.3
14.6
6.7
0.76
10.0
2.25
40.9
15.4
6.8
0.79
2.90
41.8
16.0
6.6
0.79
TM synthetic 1M
1.0
2.55
41.9
15.9
6.7
0.79
0.84
+
3,0
2.75
41.3
16.1
6.6
0.77
10.0
3.20
40.9
16.2
6.7
0.78
3,47
41.3
15.1
6.5
0.83
LDDT
1.0
3.44
40.3
15.5
6.7
0.84
42K
+
3.0
3,42
40.8
14.4
6.8
0.81
+
10.0
3.39
41.9
14.9
6.6
0.79
0.84
3.27
41.8
15.4
6.3
0.85
0.83
PLDDT
+
1.0
3.28
40.7
14.3
6.5
0.85
IM
+
3.0
3.25
41.7
17.2
6.5
0.84
+
10.0
3.24
41.6
16.1
6.6
0.85
3.69
40.8
15.1
6.7
0.77
LDDT chain 42K
+
1.0
3,57
40.9
15.7
6.6
0.85
+
3.0
3.63
41.7
15.2
6.9
0.84
0.85
+
10.0
3.59
41.6
15.2
6.8
0.83
3.29
40.5
16.1
6.6
0.81
PLDDT chain IM
1.0
3.36
40.8
17.1
6.5
0.88
+
3.0
3.30
41.0
15.0
6.5
0.85
+
10.0
3,31
41.8
15.4
6.3
0.87
0.85
pLDDT balanced IM
2.63
42.1
15.8
6.4
0.75
0.83
pLDDT
2.43
42.0
14.9
7.0
0.76
balanced 10M
PLDDT
2.40
42.1
16.5
6.3
0.73
balanced 60M
PLDDT chain
2.45
41.1
18.0
6.1
0.75
balanced IM
pL.DTT chain
2.24
41.3
17.0
6.7
0.74
balanced 10M
PLDDT chain balanced 60M
2.21
41.9
17.5
6.3
0.73
0.83
Table 10: Evaluation results of Protein Infilling task, trained without (top row) and with SC regular- ization (all other rows). The table combines all the choices for AFDistill pretraining and showing their validation accuracy. as well as the corresponding performance on the downstream application. We select the best performance for each experiment based on the highest recovery rate (marked in bold).
24


Reviews:
[{"review_id": "Y_AgqjjiS1", "reviewer": "Reviewer_jZLG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposed a knowledge distillation module for protein inverse folding, i.e., generating a protein sequence by its 3D structure. The proposed system recovers more diverse amino acid chains with a lower perplexity and higher recovery rate efficiently.\n \n\n \n", "summary_of_the_review": "The submission is not sufficient to be published in ICLR.", "strength_and_weaknesses": "**Strengths:**\n\n1. The paper overall is easy to follow, and the model construction is clearly introduced.\n\n2. According to the experimental results reported in the main paper, the main claimed contributions of “efficiently interpreting the protein sequence” and “appliable to other tasks (protein infilling)” are supported with evidence (although this evidence is flawed, see weaknesses).\n \n**Weaknesses:**\n\n1. The system design can be better motivated. In particular, why it is essential to infer protein structure from the predicted sequence for an inverse folding problem? Does it follow some convention or some specific empirical meaning?\n\n2. It is questionable whether the comparison of 'the efficiency' of the proposed method against baselines is fundamentally fair. It seems the only (in)direct comparison is in the left plot of figure 2, where the inference time is compared against different folding models. However, if the reviewer understands correctly, the distillation model itself requires running AlphaFold2 to obtain the predicted labels (which are later used for training). In other words, Figure 2 actually compares a fraction of the proposed model with the entire competitor model.\n\n3. The comparison of computational efficiency is incomplete. While the authors report the inference speed with a chain of 500 amino acids, it is not clear if the proposed method also performs swift inference speed in large-scale protein sequences. As the chain length can go beyond thousands of amino acids, the practical importance of the proposed model is questionable. While the authors claimed “truncated proteins up to length 500 to reduce computational complexity” (above Figure 4), other language models (such as ESM-IF1) work well on amino acid chains of up to 1024 tokens.\n\n4. More experimental results should be included (in the main text) to support the merit of the proposed method. For instance, as the designed module mainly solves the inverse folding problem, at least some specific baselines should be compared, such as ESM-1v and ESM-IF1. Meanwhile, the designed computational blocks should be investigated with detailed ablation studies. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SJlPjn30CO", "reviewer": "Reviewer_LQja", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper mainly tackles the problem of protein design through training of an inverse folding model. The procedure involves using a protein folding model such as AlphaFold as a consistency regularizer along with the sequence generation task. To make such learning feasible in terms of computational efficiency and training time, it distills AlphaFold folding model into a smaller, more efficient, but similarly-accurate BERT-like transformer model.", "summary_of_the_review": "The paper proposes a simple consistency regularization method for protein modeling, with the most focus on protein design. Such a simple setup achieves encouraging results which significantly improve over recent baselines on important tasks and aspects. Conditioned on a satisfying discussion around my questions above I find the paper a clearly well-performed application paper on an important topic of relatively-wide interest.\n\n-------- post rebuttal ---------\nI read the other reviews and the authors' response. I believe the paper has clear merits for publication as the approach is reasonable and novel, the diversity of baselines are increased with the new SC mechanism while mostly maintaining the recovery rate, and the efficiency of distilling model for AF2 scores.   ", "strength_and_weaknesses": "+ The main idea is a simplistic distillation of the AlphaFold model to be used as a consistency regularization for other protein modeling tasks.\n+ The papers shows that such as simple idea works in drastically improving efficiency of the folding model and the cosistency regulariztion is clearly effective for multiple tasks.\n+ The correlation of the distilled model's scores seem to be acceptably high with respect to the original AF's scores.\n+ The general performance on protein design seems to significantly improve over the recent baselines in terms of ground-truth sequence likelihood and amino-acids recovery.\n+ Different types of inverse folding backbones is used in conjunction with the consistency loss to demonstrate the universality of its usefulness.\n+ Interestingly, the diversity of the obtained well-structured sequences is improved over the baselines.\n+ Results on (structure-preserving) sequence infilling is also encouraging.\n\n\nSome questions:\n- how does the distilled model perform if trained from scratch and not from a pre-trained language model? \n- can the fact that the model is pre-trained as a self-supervised language model suggest that it might have memorized some sequences that it has to generate on the validation set?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "SIaVIrcRVSl", "reviewer": "Reviewer_QCkw", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tries to incorporate forward folding/structure prediction models, i.g., AlphaFold 2 in particular, to provide informative feedback for inverse folding models during training. One of the challenges here is forward folding models are usually too large and computationally expensive to be integrated into the learning framework. To tackle this, this paper proposes to distill AF2 by a ProtBert, a Transformer based language model, in terms of AF2’s confidence metrics, i.e., pLDDT and pTM scores. After distillation, the distilled score model is used to provide signals regarding structural consistency for learning inverse folding models. Experiments on fixed backbone design (CATH 4.2) and antibody CDR infilling (SabDab) show the effectiveness of the proposed approach. \n\n", "summary_of_the_review": "-", "strength_and_weaknesses": "Overall, the motivation for incorporating folding models to improve the structural fidelity of inverse folding models is straightforward, and resorting to knowledge distillation makes sense. However, my major concerns lie in both the model and evaluation parts. Please refer to the questions for more details.\n\nQuestions:\n\n1. The baseline models, GVP, are weak. It remains unclear whether the proposed idea could bring consistent improvements when strong inverse folding models are considered, such as protein mpnn?\n2.  Since additional 907k AF2 predicted proteins are used to build the distilled model, a simple baseline in which all of these predicted data are directly used for inverse folding training should be carefully considered and compared. What’s more, does the proposed approach, which requires significant additional overhead in terms of both distillation and regularized training, complement with a simple data augmentation approach? This important question remains unclear. \n3. Comparisons to the state of the arts methods are missing, e.g., ESM-IF (Hsu+ ICML 2022).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Opxh-o3m5CP", "reviewer": "Reviewer_Y6Az", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper focuses on inverse protein folding, a fundamental and challenging problem in protein science. The authors propose a new regularization term by including the trained AlphaFold model and its predicted structural confidence metric. To improve the efficiency of the model optimization loop, they distill the AlphaFold into a smaller model. They perform extensive experiments with different protein design models to show generality of the regularization trick. Experimental results show that the model with regularization achieves better sequence recovery rate and large improvement in protein diversity, while preserving the structural consistency.", "summary_of_the_review": "A method distilling Alphafold for inverse folding, which is an important topic for computational protein design. The idea of distilling the advanced protein structure prediction model for inverse folding might be interesting to the specific \"ML for protein\" audience, while generally the ML models and methods are all just existing works. Experiments are sound and comprehensive.\n\nQuestions:\n1. It’s interesting to see that such a regularization term can increase sequence diversity. Could the authors provide some explanations and in-depth analysis of this phenomenon? \n\nMinor points:\n1. In the caption of Fig.3. Line 5: TM/LDTT -> TM/LDDT\n2. In Fig. 7,9,10. The cyan bar is overlapped with the purple bar, which makes the color similar to the blue bar. Since the baseline value is the same for all methods in one figure, I suggest using a dashed red line rather than a cyan bar.", "strength_and_weaknesses": "Strength:\n1. Address an important problem with a well-motivated method.\n2. Simple techniques with impressive results, especially in protein diversity.\n\nWeakness:\n1. Lack of technical novelty. The models and distillation methods are adapted from existing works without non-trivial modifications.\n2. Lack of related work on knowledge distillation. \n3. The presentation is not friendly to readers that are unfamiliar with inverse protein design.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject