KW-Design: Pushing the Limit of Protein Design via Knowledge Refinement
Authors: Zhangyang Gao
Cheng Tan
Xingran Chen
Yijie Zhang
Jun Xia
Siyuan Li
Stan Z. Li
Venue: ICLR 2024 poster (2024)
Keywords: Protein Design
Graph
Finetuning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Recent studies have shown competitive performance in protein inverse folding, while most of them disregard the importance of predictive confidence, fail to cover the vast protein space, and do not incorporate common protein knowledge. Given the great success of pretrained models on diverse protein-related tasks and the fact that recovery is highly correlated with confidence, we wonder whether this knowledge can push the limits of protein design further. As a solution, we propose a knowledge-aware module that refines low-quality residues. We also introduce a memory-retrieval mechanism to save more than 50\% of the training time. We extensively evaluate our proposed method on the CATH, TS50, TS500, and PDB datasets and our results show that our KW-Design method outperforms the previous PiFold method by approximately 9\% on the CATH dataset. KW-Design is the first method that achieves 60+\% recovery on all these benchmarks. We also provide additional analysis to demonstrate the effectiveness of our proposed method. The code is publicly available via \href{https://github.com/A4Bio/ProteinInvBench}{GitHub}.

Published as a conference paper at ICLR 2024
KW-DESIGN: PUSHING THE LIMIT OF PROTEIN DE- SIGN VIA KNOWLEDGE REFINEMENT
Zhangyang Gao 1,2,t, Cheng Tan Xingran Chen 3, Yijie Zhang 1 Jun Xia 2, Siyuan Li Stan Z. Li 2,x
Zhejiang University; Research Center for Industries of the Future, Westlake University
3 University of Michigan, USA; McGill University, Canada
gaazhangyang, tancheng, Stan. 2Q. Li} (@westlakee eew.aaaeeeeeeee... en
ABSTRACT
Recent studies have shown competitive performance in protein inverse folding, while most of them disregard the importance of predictive confidence, fail to cover the vast protein space, and do not incorporate common protein knowledge. Given the great success of pretrained models on diverse protein- reetted tasks and the fact that recovery is highly correlated with confidence, we wonder whether this knowledge can push the limits of protein design further. As a solution, we propose a knowledge-awee aware module that refines low-quality residues. We also introduce a memory-retriewe mechanism to save more than 50% of the training time. We extensively evaluate our proposed method on the CATH, TS50, TS500, and PDB datasets and our results show that our K KW-Design method outperforms the previ- PiFold method by approximately 9% on the CATH dataset. KW KW-Design is the first method that achieves 60+% recovery on all these benchmarks. We also pro- vide additional analysis demonstrate the effectiveness of our proposed method. The code is publicly available via GitHub
INTRODUCTION
Protein sequences, which are linear chains of amino acids, play a crucial role in determining the structure and function of cells and organisms. In recent years, there has been significant interest
in designing protein sequences that can fold into desired structures (Pabol 1983). Deep learning
models (Li al. 2014 Wu al. [202]] Pearce & Zhang [022]] 2202] Ovchinnikov & Huang 2021 Ding
it al. [022] Gao et al. 2020 2022a| Dauparas et al.
2022
Ingraham et al. 2019 |iing et al.
2020
Tan et al. [022cc] Hsu al., 2022; O'Connell et 2018
Wang et al.) 2018 Qi & Zhang
2020
Strokach et al. |020 Chen et |019| et al. 2020a
Huang et al. [017 Anand et al.
2022
Strokach & Kim| 2022
[i & Koehl, 2014] Greener | [018]]Karimi Karimi et al. 2020 Anishchenko
et al. 2021 Cao et al.) 2021 Liu et al. 2022 McPartlon et al. [2022] Huang et al.) 2022 Dumortier
jet al. 2022 Li et al. 2022a, Maguire et al. 2021: Li et al. (022b) have made significant progress
in this area. However, many of these methods either ignore the importance of predictive confidence, fail to cover the vast protein space, or lack considcration of common protein knowledge. We argue that the absence of common protein knowledge limits the generalizability of protein design models, and propose a confidence-awarr module that refines low-quality residues using pretrained features.
Previous protein design methods have not fully utilized pre- dictive confidence, which is the maximum probability of a residue. Using PiFold (Gao et al.; 2023) as our base- line, a recent competitive protein design model, we observed significant differences in confidence distributions between positive and negative residues, as shown in Fig. П This finding inspired us propose confidence-awee that automatically identifies low-quality residues and iter- atively refines them reduce prediction errors. However, we encountered a challenge: the recovery rate of PiFold plateaued at around 52% regardless of how many refining layers (PiGNNs) were added. We hypothesized that this was due to the model being trapped in a local optimum based on the current training set. Escaping the local minimum would require additional inductive bias from other teacher models.
module
Figure 1: Confidence distributions of pos- itive and negative residues designed by Pi- FoldGao al. (2023). Positive residues are identical to native residues and vice versa.
Equal Contribution, "Corresponding Author.

---

Published as a conference paper at ICLR 2024
To escape the local minimum and improve the performance of our protein design model, we suggest leveraging pretrained teacher models. These models have made significant progress on a variety of downstream tasks (Zang et al.] 2022b Meter et al.) 202I; Zhang et al.. 2022a Chen et al., 2023 by learning common knowledge across a vast protein space. structural knowledge (Zhang et al. 20226 |Hsu al.] 2002)) help to learn expressive protein features, while the sequential knowl- edge (Meie et al. 2021) can aid in designing rational proteins. In this study, we investigate three pretrained models, i.c., ESM (Meier al.) 2021 Lin et al. (022), ESM-IF (Hsu al. (022), and GearNet (Zhang et al., (022b), to extract structural and sequential embeddings as prior knowledge that can enhance our refinement module. As structure-in and sequence- out task, the structural- based protein design can benefit from the multimodal knowledge and automatically revise residues
that violate common sense.
To boost protein design, we propose confidence-aware refining model that leverages multimodal knowledge. However, we face several challenges: (1) how to adaptively fuse multimodal pretrained knowledge based on predictive confidence, (2) how to develop more effective refining technologies, and (3) how to efficiently tune the model with large-scale pretrained parameters. Firstly, we propose multimodal fusion module that combines the knowledge of structure pretraining, sequence pre- training. and history predictions. Predictive confidence is used to control the combination through gated attention, enabling the model to adaptively select proper knowledge. Secondly, we suggest using virtual MSA and recycling technologies to improve the recovery. Thirdly, we introduce a seemory-retreeall mechanism that caches the intermediate results of modules. This mechanism en- ables the model to retrieve historical embeddings without performing a forward pass, resulting in more than 50% training time savings.
We call our method KW-Design, a refining method that considers multimodal knowledge as well as predictive confidence. We evaluate our method on four benchmark datasets: CATH, TS50, TS500 and PDB, and observe significant improvements across all settings. For example, KW-Desigg is the first method to achieve 60+% recovery on all these datasets. On the CATH dataset, we observe 9.1 111% 1% improvement compared to the previous PiFold method. We also conduct extensive ablation studies to demonstrate how the wwwlledg--rffnemeet module works and how the memory-retrieval mechanism saves training time. Additionally, we provide further analysis to demonstrate the supe- riority of our proposed method. Overall, our results demonstrate the effectiveness of our approach in improving the recovery of protein sequence design. In the appendix, we show that the recovery is highly correlated with the structural similarity and discuss that under the same level of structural similarity, our method can generate more novel protein sequences than baselines.
2 RELATED WORK
Recently, Al algorithms have evolved rapidly in many fields (Gao et al. 2022d Cao et al. 2022 Tan
ct al. 2022a Li el al. 2022c He et al. 2020 Stärk et al. 2022), where the protein folding pro broblem em
(Jumper et 2021 Wu et al. 2022 Lin al. 2022 Mirdita et al. 2022 Wang et al. 2022 Li
ct al. 2022d) that has troubled humans for decades has been nearly solved. Its inverse problem-
structure-based protein design- - is receiving increasing attention.
Problem definition The structure-based protein design aims to find the amino acids sequence S {s+ :1 S i < n} folding into the desired structure X' {x; R 1 S n), where n is the number of residues and the natural proteins are composed by 20 types of amino acids, i.c., 1 I s; < 20 and s; E N Formally, that is to learn function Fo:
Fø : X S.
Because homologous proteins always share similar structures (Pearson & Sierk 2005), the problem itself is underdetermined, i.e., the valid amino acid sequence may not be unique (Gao et al. 2020).
MLP-based models MLP is used to predict the probability of 20 amino acids for each residue, and various methods are mainly difficult in feature construction. These methods are commonly evaluated on the TS50, which contains 50 native structures. For example, SPIN {Li et al.) 2014) achieves 30% recovery on TS50 by using torsion angles (ó and v³), sequence profiles, and energy profiles. Through adding backbone angles (0 and r), local contact number, and neighborhood dis- tance, SPIN2 (O'Connell al., 2018) improves the recovery to 34% Wang's 's model (Wang al.,
2

---

Published as a conference paper at ICLR 2024
2018) suggests using backbone dihedrals (Ø, 2b and w), the solvent accessible surface area of back- bone atoms (Ca, NN C, and O), secondary structure types (helix, sheet, loop), Ca - Ca distance and unit direction vectors of C. Ca. C, N and C. C and achieves 33% recovery. The MLP method enjoys a high inference speed but suffers from a low recovery rate because the structural information is not sufficiently considered.
CNN-based models
These methods use 2D CNN or 3D CNN extract protein features (Torg
& Altman| 2017 Boomsma & Frellsen 2017 Weiler et al. 2018 Zhang et al. 2020a Qi & ZZhag
2020 Chen al. 2019) and commonly evaluated
the TS50 and TS500. SPROF (Chen
et al. 2019) adopts 2D CNN to learn residue representations from the distance matrix and achieves
a 40.25% recovery on TS500. 3D CNN-based methods, such as ProDCoNN (Zhang al.. 2020a) and DenseCPD QQ & Zhang 2020), extract residue features from the atom distribution in three- dimensional grid box. For each residue, after being translated and rotated standard position, the atomic distribution is fed to the model to learn translation- and rotation- potation-invariant features. ProD- CoNN (Zhang et al. 2020a) designs nine-layer 3D CNN with multi-scale convolution kernels and achieves 42.2% recovery on TS500. DenseCPD (Qi & Zhang (2220) uses the DensetNet architec- ture (Huang et al.) (017) to boost the recovery to 55.53% on TS500. Recent works (Anand et al.) 2022) have also explored the potential of deep models to generalize to de novo proteins. Despite the improved recovery achieved by the 3D CNN models, their inference is slow, probably because they require separate preprocessing and prediction for each residue.
Graph-based models These methods use k--N graph to represent the 3D structure and employ
graph neural networks (Defferrard et al., 2016 Kipf & Welling 2016 Veličković et al. 201 7 Zhou
jet al. 2020. Zhang et al. [2000b] Gao al. 2022b Tan et al., 2022b Gao al. (022c) extract residue features while considering structural constraints. The protein graph encodes residue infor- mation and pairwise interactions as the node and edge features, respectively. GraphTrans (Ingraham itt al. 2219)) uses the graph attention encoder and autoregressive decoder for protein design. GVP (Jing et al. 2020) proposes geometric vector perceptrons to learn from both scalar and vector fea- tures. GCA (Tan et al. (202c) introduces global graph attention for learning contextual features. In addition, ProteinSolver (Strokach et al., 2020) is developed for scenarios where partial sequences are known while not reporting results on standard benchmarks. Recently, AlphaDesign (Gao et al.) 2022a), ProteinMPNN (Dauparas et al. 2022) and Inverse Folding (Hsu et al.) (022) achieve dra- matic improvements. Compared to CNN methods, graph models do not require rotating each residue separately as in CNN, thus improving the training efficiency. Compared to MLP methods, the well-
exploited structural information helps GNN obtain higher recovery.
METHOD
3.1
OVERALL FRAMEWORK
The framework of our KW-Design model is illustrated in Figure 2 comprising initial design model (Foto)), and L confidence-aware knowledge-tuning modules (oam, fgrzs,... ""(,,,)))) where g( and g(*)) are learnable parameters. Each knowledge tuning process consists of knowledge extractor (pretrained ESMIF, ESM2-650M, or GearNet), a confidence predictor, multimodal fusion layer, and a refinement module. Note that pretrained protein models are frozen during optimizing and not included in (0)) To simplify the notation, we write FØc = fe(m) ... foc Fg where g(k) = (100 (166) and denotes concatenation operation.
For the 1-th knowledge-tuning module, we denote the protein structure as a E Rm.3 the residue embedding as (!)) E R=.dd and the predicted probabilities as p(²) E R=.21, Formally, we have h0 = = ... faca, o Fgto: (x) and (')) = Predi ct (1(")) Here, o denotes the operation of compositing functions, is the number of residues, d is the embedding size, 21 is the number of amino acids plus special token of [mask], and Predi ct (-TT) is linear layer followed by the softmax activation using a temperature value of T.i.., Softmax(11) (enftmma(tnner(())) near(.)/T) The overall
objective of our KW-Design model is to minimize the loss function L with respect the learnable parameters ((0) (()) (L):
g()),(1))
min 000, (00) C(far) ... Fart Egcor (x), s)
(2)
3

---

Published as a conference paper at ICLR 2024
sequence
XL
Siructure
Sequence
Structure
Sequence
Structure
Knowledge
Sequence
Design
Design
Tuning
Design
Fuse
Tuning
Canfidence
(c) Krnowledge-tunigg
(a) Baseline
(b) Fiiee-tuiing
Retrieval
Fgen
fan
S pa
f jax
B
Model
Design
Tuning
Tuning
Tuning
Full fine-tuning
fx
team-
leam
lear
step1
retrieval
learn
leam
learn
refinmment
Step-by step
step2
retrivva
step3
retrieval-
learn
(d) Comparing the workflow
Figure Comparison various models. (a) The baseline trains from scratch without using pretrained knowl- edge. (b) The fine-tuning model refines the output of the baseline. (c) The knowledge-tuning model fuses multimodal pretained knowledge and the confidence to enhance the refinement module. (d) The proposed KW- Design model introduces memory bank to speed up the training process skipping the forward pass of well-tuned modules.
Here, s is the reference sequence, and x is the protein structure.
Under the assumptions that if C(Fgre (x), < C(Fgur, s), then C(fan- (Foc)) C(fare+ (Faww (z)), s), indicating that stronger design model leads to better results with the same fine-tuning module, we simplify the objective as:
min L(farr (x(x-))), s), (Th)
s), s.t.,(t--1) 0(h-1)
min [(F) 1)k-1)
(3)
Note that by optimizing the fine-tuning modules fa(v); Sarz2 ... Sares sequentially. Therefore, we conduct module-wise training strategy, where the parameters g(k-1) are frozen when optimizing ((k) In the appendix A we show the module- training strategy is helpful in reducing the computational cost by avoiding redundant forward passes of large pretrained models. The optimal embedding in the memory bank automatically determined by the early stopping (patience=5), where the indicator is the average predictive confidence of the whole sequence. We will introduce the details of the refining technique, knowledge-tunie tuning module and memory bank in the following sections.
= =oo k 1-, 1}. Eq. suggests the problem could be solved
3.2 REFINING TECHNIQUE
Recycling Process Given the initial residue embedding ()) = F Fotos (æ), our KW-Design applies sequence of knowledge-tuning modules to update the residue embedding:
1000
fa0t2) (1+1) faa+1 (1+1)
15(e) (L)
(0)
(1)
(4)
where L is the maximum number of refinement modules. The predictive probability is obtained by:
0 (1
Virtual MSA To capture diverse protein knowledge, we sample a set of protein sequences (s.) Multinomial((()))0 < m) from the predicted probabilities by setting T 0.05 in 5 Any residues with confidence score below 0.99 are replaced with [mask), which helps in constructing high-quality prompt template. We call this set of sequences the virtual multiple

---

Published as a conference paper at ICLR 2024
sequence alignment (MSA) and feed them into pretrained models to obtain the residue embeddings:
h (1), seq
Feeg ( (8(0.t)
(6)
(1) h 3d
= F33 (3())
(7)
where a: is the 3d coordinates residues, and Fseg and F3d are sequence and structure pretrained models, respectively. The sequential embedding (1).5 E Rnndd captures the primary sequence knowledge using ESM2-650M, while the structural embedding had E Rrndds captures the structure knowledge using ESMIF or GearNet. Together, these features are combined as unified embedding 21)) through the fusion module:
(0),
(.amm (h (0 Fuse ( Jt=1 {had JE= (0
11=1;
(8)
which can be further converted as h (1+1) and p(l=1);
h (++1) Refi Re (2 pp (41) Predict (h+))))))))))))))))))))))))
where the Refine (1+1) module is the 10-layer PiGNNs, described later in Section 3.3.
Confidence-aware updating We define the confidence vector of a sequence s as the correspond- ing predictive probability, written as Cg:
Ca = (p..s P2,52: PP....
(10)
Note that Pt.e; represents the predicted probability of the i-th amino acid and the predicted residue type is s,. Because some residues are harder to design than others, they may benefit more from refinement. Considering this, we introduce confidence-aware gated attention mechanism that up- dates the pre- and post-refinement embeddings based on the predictive confidence of each residue. This allows us to focus more on difficult residues during refinement and improve overall design performance:
(1+1) h (++1) (MLP) (c (1+1) (c()t) h0 o (MLP2(cc)) C2 (211))
(11)
where Multinomial( and s' (u are sampled from multimodal distri- butions. is the sigmoid function, @ is element-wise multiplication.
3.3 ANOWLEDGE-TUNINE MODULE
The mhowledge-tuning module updates the residue embeddings of well-tuned models to generate more rational protein sequences. As shown in Figure 2 the knowledge-tuning module includes knowledge extractor, a confidence predictor, a fusion layer, and tuning layer.
Knowledge extractor & Confidence predictor As introduced in Sec S.c3.2 the pretrained knowl- edge extractors (F(() and F(()) extract embeddings from virtual MSAs. The confidence predictor (Predict (1+1)) outputs the predictive probability ((+1) based on the residue embedding h (!+1),
Fusion layer The fusion layer combines the sequential and structural embeddings with a confi- dence score to obtain unified embedding. Specifically, the structural and sequential MSA embed- dings are fused using a confidence-aware gated layer:
Embed( s ():) (h + MLP, 1 (h (0. O (MLPss(( (cD)))
(12)
Refinement module The refinement module is learnable graph neural network (GNN) that takes
(!)) input node features and el input edge features. Given the protein graph G, node features
(1)) and edge features e!!), PiGNNs the refinement module for updating embeddings:
h (++1) ((+1),e((t))) (l+1) P (iGNNs
G)
(13)
Note that the initial node features 20) and edge features (0) extracted from PiFold,
5

---

Published as a conference paper at ICLR 2024
3.4 MEMORY RETRIEVAL
From 3 we know that 8(0) (0))........ o(0)
are frozen parameters when optimizing ((++1)
Algorithm Memory Net Framework
Therefore, we can use a memory bank M to Usage: Retrieve embedding from the memory
store and retrieve the intermediate embeddings of bank without the forward pass.
the l-th design model Pocc (a)) for speeding up the
Input: A hatch of inputs H = 11 |ng):
process of optimizing (b(-1)) As shown in Alg the protein embedding (t)) can be retrieved from the memory bank MM without the need for forward pass, provided that the following condi-
Output: A batch of outputs H(
h1 h [|| IIng |n!t".
1: Stepl: Debatch input data
tions are satisficd: (1) the embedding h) is al- ready stored in M! and (2) the saved embed-
2:
(h hn" .... _n DeBatch(H(.))
3:
4: Step2: Retrieve embeddings
dings are consistently obtained from an optimal model Foe.. While the first condition is straight- forward, the second condition requires the algo- rithm to automatically determine the optimal gi(t) and freeze fao0 to ensure that the memorized em- beddings are consistent. To determine the optimal 10:
5: for i E (0,66)
6:
if Ih ((+1) E (M and faco is early stopped
h ((+1)) [M]+
7:
8:
else
9:
h ((+1) = Refine
g(0), we use the average confidence score over the
11: Step3: Batch output
validation set as an indicator and apply the early
12: Save {11 (1+1) hg (4+1)
hy 2 to M
stopping operation to determine the optimal (()).
13: Return h 1h2
(1+1) (111)
||ng (+1)
with the patience of 5.
EXPERIMENTS
We evaluate the performance of KW-Design on multiple datasets, including CATH4.2, CATH4.3, TS50, TS500, and PDB. We also conduct systematic studies to answer the following questions:
Performance (Q1): Can KW-Design achieve state-of-the- accuracy on real-world datasets?
Refining technology (Q2): How much can models gain from different refinement techniques?
Knowledge tuning (Q3): Which pretrained knowledge helpful in improving protein design,
and how much of a speed boost can the memory bank bring?
More analysis (Q4): How does the KW-Design make a difference on the basis of PiFold?
4.1 PERFORMANCE ON CATH (Q1)
Objective & Setting We demonstrate the effectiveness of KW-Design on the widely used CATH (Oreggo et al. 1997) dataset. To provide a comprehensive comparison, we conduct experiments on both CATH4.2 and CATH4.3, The CATH4.2 dataset consists of 18,024 proteins for training, 608 proteins for validation, and 1,120 proteins for testing, following the same data splitting as Graph- Trans (Ingraham et al. 2019), GVP (Ting et al.) 2020p. and PiFold (Gao et al.) 2023). The CATH4.3 dataset includes 16, 153 structures for the training set, 1,457 for the validation set, and 1,797 for the test set, following the same data splitting as ESMIF (Hsu et al. 2022). The model is trained up to 20 epochs using the Adam optimizer on an NVIDIA V10O. The batch size and learning rate used for training are 32 and 0.001, respectively. To evaluate the generative quality, we report perplexity and median recovery scores on short-chain, single-chain, and all-chain settings.
Baselines We compare KW-Design with recent graph models, including StructGNN, StructTrans 2022c). GVP (Jing et al. 2020), Gap-large, AlphaDesign 2022). ProteinMPNN (Dauparas et al. 2022), and Pifold Gao et al. (202)), as most of them are open-source. To provide a head-to-head comparison with
IIgrraham et al. 2019), GCA (Tan et al.
Gao et al. 2022a), ESM-IF (Hsu al.
ESMIF, we retrain our model on the CATH4.3 dataset following the same data splitting as ESMIF,
while we do not utilize AF2DB for training the model.
Results & Analysis From Table I we conclude that KW-Design consistently achieves state-of- the-art performance different settings. Specifically, we observe the following: (1) KW-Design is the first model to exceed 60% recovery on both CATH4.2 and CATH4.3, demonstrating its superior ability in generating protein structures. (2) On the full CATH4.2 dataset, KW- KW-Design achieves a per- plexity of 3.46 and a recovery of 60.77%, outperforming the previous state-of of-te-ofthhean the-art model PiFold by 23.95% and 9.11%, respectively. Furthermore, KW-Design achieves recovery improvement of

---

Published as a conference paper at ICLR 2024
Table 1: Results comparison the CATH dataset. All baselines reproduced under the code frame- work, except ones marked with t. copy results GVP-large and ESM-IF from their manuscripts (Hsu et al. 2022). The best and suboptimal results are labeled with bold and underline.
Perplexity Single-chain
Recovery % 1
CATH version
Model
Short
All
Short
Single-chain
All
4.2
4.3
StructGNN
8.29
8.74
6.40
29.44
28.26
35,91
GraphTrans
8.39
8.83
6.63
28. 14
28.46
35.82
GCA
7.09
7,49
6.05
32.62
31.10
37.64
GVP
7.23
7.84
5.36
30.60
28.95
39.47
GVP-large
7.68
6.12
6.17
32.6
39.4
39.2
AlphaDesign
7.32
7.63
6.30
34.16
32.66
41.31
ESM-IFF
8.18
6,33
6,44
31.3
38.5
38.3
ESM-IF (AF2DB)
6.05
4.00
4.01
38.1
51.5
51.6
ProteinMPNN
6.21
6,68
4.61
36.35
34.43
45.96
PiFold
6.04
6.31
4.55
39.84
38.53
51.66
KW-Design (Ours)
5.48
5.16
3.46
44.66
45.45
60.77
KW-Design (Ours)
5.47
5.23
3.49
43.86
45.95
60.38
4.82% and 6.92% the short and single-chain settings, respectively. (3) KW-Desigl also achieves similar improvements when extending to the CATH4.3 dataset, further validating its effectiveness and generalizability. Overall, the strong performance suggests that K KW-Design could be valuable tool for advancing protein engineering and drug design.
4.2 PERFORMANCE ON TS50 AND TS500 (Q1)
To provide a more comprehensive evaluation and demonstrate the general- izability of KW-Design, we also evaluate it two standard protein benchmarks, TS50 and TS500. These datasets contain 50 and 500 proteins, respectively, and are widely used for evaluation.
Objective & Setting
Table 2: Results TS50 and TS500. All baselines reproduced under the code framework.
TS50
TS500
Model
Perplexity |
Recovery f
Worst
Perplexity J
Recovery t
Worst
StructGNN
5.40
43.89
26.92
4.98
45.69
0.05
GraphTrans
5.60
42,20
29.22
5.16
44.66
0.03
GVP
4.71
44.14
33.73
4.20
49. 14
0.09
GCA
5.09
47.02
28.87
4.72
47.74
0.03
AlphaDesign
5.25
48.36
32.31
4.93
49.23
0.03
ProteinMPNN
3,93
54.43
37.24
3.53
58.08
0.03
PiFold
3.86
58.72
37.93
3.44
60.42
KW-Design (Ours)
3.10
62.79
39.31
2.86
69.19
0.02
Results & Analysis Experimental results are shown Table where KW-Design significantly outperforms previous baselines all benchmarks. We observe that: (1) On the TS50 dataset, KW-Design achieves a perplexity of 3.10 and recovery rate of 62.79%, outperforming the previ- state-of-the-art model, PiFold, by 19.69% and 4.07%. (2) On the TS500 dataset, KW-Design achieves a perplexity of 2.86 and a recovery rate of 69.19%, outperforming PiFold by 16.86% and 8.77%. (3) Notably, KW-Design is the first model exceed 60% and 65% recovery on the TS50 and TS500 benchmarks, respectively.
4.3 PERFORMANCE ON PDB (Q1)
Objective & Setting In addition, we expand the evaluation to include multi-chain protein design using the dataset curated by ProteinMPNN, The dataset was preprocessed by clustering sequences at 30% identity, resulting in 25.361 clusters. Following ProteinMPNN's setup, we divided the clusters randomly into training (23,358), validation (1, (1,464), and test sets (1,539), ensuring that none of the chains from the target chain or biounits the target chain were present in the other two sets. During each training epoch, cycled through the sequence clusters and randomly selected sequence

---

Published as a conference paper at ICLR 2024
member from each cluster. All baselines retrained and evaluated under the PDB dataset. We report the median values of confidence scores and recovery rates for the designed sequences.
Table 3: Multi-chain results on the PDB dataset.
Model
Confidence f
Recovery
length
L < 100
100 L 500
500
L 1000
Full
L L
100
100 S < 500
500 S < 1000
Full
StructGNN
0.49
0.49
0.50
0.49
0.41
0.41
0.42
0.41
GraphTrans
0.48
0.47
0.39
0.40
0.40
GCA
0.45
0.45
0..66
11.45
0.41
0.41
01.42
0.41
GVP
0.51
0.53
0.44
0.42
0.45
0.43
AlphaDesign
0.52
0.53
0.54
0.53
0.48
0.49
0.50
0.49
ProteinMPNN
0.54
0.56
0,58
0.57
0.52
0.53
0.55
0.53
PiFold
0.56
KWDesign
0.65
0.71
0.74
0.71
0.59
0.66
0.67
0.66
Results & Analysis From Table 3 we conclude that: (1) The predictive confidence score, an unsupervised metric, exhibits strong correlation with the recovery. As result, we leverage this metric for two purposes: firstly, fuse pre-trained knowledge, and secondly, to construct high- quality sequence prompts. (2) The longer the protein, the higher the recovery. Given that the protein length in the PDB dataset is generally larger than that the CATH dataset, all models achieve higher recovery rates the PDB dataset. (3) KW-Design achieves the best performance all protein lengths, demonstrating its effectiveness in multi-chain protein design.
4.4 REFINING TECINOLOGY (Q2)
Objective & Setting We conduct ablation studies investigate the effects of virtual MSA, cycling, and the confidence-aware tuning module. We follow the same experimental setting as in Section 4.1 and report the results the CATH dataset. Specifically, vary the number of virtual MSA and recycling times from to 3, and remove the confidence-aware tuning module by replac- ing the confidence score with constant value of 1.0. We also compare the training time with and without using the memory bank.
Conlig
Perplexity Single-chain
Recovery Single-chain
Training (per epoch)
w/o confidence
MSA
Recycle
Short
All
Short
All
w/o memory
w memory
3.54
$.39
3.59
42.58
42.74
$8.39
20min
70 min
5.52
3.52
44.72
44.19
59.72
40min
5.46
5.17
3.48
43.91
44.16
601.4
60min
210min
3.54
5.39
3.59
42.58
42.74
58.39
20min
70 min
5.55
5.42
3.56
42.72
42.16
58.62
33min
83min
5.57
5.42
3.56
42.94
43.23
58.71
45min
95min
2
5.49
5,22
3.49
44.76
45.90
39.84
65min
165min
5.48
5.16
3.46
44.66
45.45
61..7
100min
250min
5,50
43.88
44.08
59.64
Table 4: Ablation of refining technology. "w/o confidence means" replacing the confidence score as a con- sistent value 1.0. "w memory" and "w/o memory" indicate whether using the memory-retrieval mechanism or not. The training time is measured an NVIDIA V100.
Results and Analysis Ablation studies about MSA, recycling and confidence embedding are pre- sented Table We conclude that:(1) Recycling has a more significant impact on performance than virtual MSA. When increasing the recycling from I to 3 while keeping the number of virtual MSAs constant the recovery rate on the full dataset improves by 1.95%, from 58.39% 60.34%. In contrast, increasing the number of virtual MSAs from 1 to 3 only results in 0.32% improvement. (2) The confidence-aware tuning module makes a non-trivial improvement by 0.78%, 1.37%, and 1.13% on the short, single-chain, and full datasets, respectively. (3) Increasing the number of virtual MSAs and recycling times leads to higher computational overhead during training. Therefore, we introduce memory retrieval mechanism that saves more than 50% of the training time in all cases. (4) Based on the above analysis, we conclude that the importance order the three components is follows: recycling confidence-awee tuning module virtual MSA.
8

---

Published as a conference paper at ICLR 2024
4.5 PRETRAIN KNOWLEDGE(Q3)
Objective & Setting By deleting the corresponding pre-trained embedding in Eq we investi- gate how much performance gain model can achieve from different pre-trained models, including ESM22-650M (Meier et al. 2021] Lin et al., 2022), ESM-IF (Hsu et al. (002)). and GearNet (Zhang
ct al. 2022b). The experimental settings are kept the same as in Section
Table 5: Ablation study multimodal knowledge. We investigate how much performance gain the model can obtain from different pre pre-trained models, where ESM+ESMIF provides the best performance.
Perplexity Single-chain
Recovery % Single-chain
Model
Short
All
Short
All
PiFold
6.04
6.31
4.55
39.84
38.53
$1.66
KW-Design (GearNet)
6.66
6.89
4.96
38.72
38.02
50.43
KW-Design (ESM2-650M)
6.05
5.29
3,90
43.32
46.30
57.38
KW- Desige (ESMIF)
6,15
6.51
4.18
38.79
39.71
54.52
KW-Design (ESMIF-ESM2-6SOM)
5.48
5.16
3,46
44.66
45.45
60.77
Results & Analysis We present results Table We observe that: (1) The knowledge of ESMIF and ESM pre-trained models contributes to improving the performance, with ESM providing a larger improvement than ESMIF. Specifically, the knowledge of ESMIF and ESM results in 2.86% and 5.72% improvement, respectively. In contrast, the knowledge of GearNet does not contribute to the improvement. (2) The best recovery rate is achieved when the model combines the knowledge of ESM and ESMIF, resulting in 9.1 1% improvement. Notably, the improvement is not linear, as the combination of ESMIF and ESM provides larger improvement than the sum of their individual contributions (i.e., 9.11% >> 2.86% + 5. 5.72%). These results highlight the importance of selecting the appropriate pre-trained models for protein structure refinement and demonstrate the effectiveness of combining multiple sources of knowledge to achieve better performance.
4.6 MORE ANALYSIS(Q4)
Recovery States We randomly selected 10 proteins from the CATH4.2 dataset test set and de- signed their sequences using PiFold. The sequences recovery states of the designed residues were visualized in Fig. Our results show that KW-Design tends make more positive corrections than negative corrections, with positive corrections occur- ring mostly in adjacent locations to initially positive residues. This suggests that the model learns the local consistency of the protein structure and can automatically correct incorrectly designed residues that violate this consistency.
then refined using K KW-Design, and the
FUsidbes
PoethNeeg
Figure Recovery states. The light blue cell indicates negative residue designed by PiFold and differs from the reference one, while light orange cell indicates positive residue matches reference one. The darker blue cells indicate where KW-Design wrongly converted positive residues into negative ones, while the darker orange cells indicate where KW-Design corrected negative residues to positive oncs.
Due space limitations, have included additional analysis and details in the appendix, which will help readers to better understand the proposed method.
CONCLUSION&LIMTTATON
We propose KW-Design, novel method that iteratively refines low-confidence residues using common protein knowledge extracted from pretrained models. KW-Design is the first model that achieves 60+% recovery CATH4.2, CATH4.3, TS50, TS500, and PDB, demonstrating its effec- tiveness and generalizability. However, the proposed method has not yet been verified through wet experiments in real applications, and this will be direction for future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was supported by National Science and Technology Major Project (No. 0022D00111011 National Natural Science Foundation of China Project (No. U21A20427), the Center of Synthetic Biology and Integrated Bioengineering of Westlake University and Integrated Bioengineering of Westlake University Project (No. WU2022A009) and the Westlake University Industries of the Future Research Funding Project (No. WU2023C019).
REFERENCES
Namrata Anand, Raphael Eguchi, Irimpan Mathews, Carla P Perez, Alexander Derry, Russ B Altman, and Po-Ssu Huang. Protein sequence design with a learned potential. Nature communications, 13(1):1-11, 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-5555 2021.
Wouter Boomsma and Jes Frellsen. Spherical convolutions and their application in molecular mod- elling. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- http:///.w....i....paaaeee/ee//a wttppppppp/..////.......e. [LL38t2766fecal005300e6556562222
sociates, Inc., 2017. URL /
Hanqun Cao, Cheng Tan, Zhangyang Gao, Guangyong Chen, Pheng-Ann Heng, and Stan Li. A survey on generative diffusion model. arXiv preprint arXiv:2209.02644 2022.
Yue Cao, Payel Das, Vijil Chenthamarakshan, Pin- Yu Chen, Igor Melnyk, and Yang Shen. Fold2seq:
A joint sequence (Id)-fold (3d) embedding- generative model for protein design. In International Conference on Machine Learning, pp. 1261-1271. PMLR, 2021.
Sheng Chen, Zhe Sun, Lihua Lin, Zifeng Liu, Xun Liu, Yutian Chong, Yutong Lu, Huiying Zhao, and Yuedong Yang. To improve protein sequence profile prediction through image captioning pairwise residue distance map. Journal chemical information and modeling, 60(1):391-39) 2019.
Zhiyuan Chen, Zuobai Zhang, and Jian Tang. Data-efficient protein 3d geometric pretraining via refinediff. arXiv preprint arXiv:2302. 10888, 2023.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning based protein sequence design using proteinmpan. bioRxiv, 2022.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks graphs with fast localized spectral filtering. Advances neural information processing systems, 3852, 2016.
Wenze Ding, Kenta Nakai, and Haipeng Gong. Protein design via deep learning. Briefings in bioinformatics, 1303) 2022.
Baldwin Dumortier, Antoine Liurkus, Clément Carré, and Gabriel Krouk. Petribert: Augmenting bert with tridimensional encoding for inverse protein folding and design. bioRxiv, 2022.
Wenhao Gao, Sai Pooja Mahajan, Jeremias Sulam, and Jeffrey J Gray. Deep learning in protein structural modeling and design. Patterns, 101442 2020.
Zhangyang Gao, Cheng Tan, Stan Li, et al. Alphadesign: A graph protein design method and benchmark alphafolddb. arXiv preprint arXiv:2202.01079, 2022a.
Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z Li. Cosp: Co-supervised pretraining of pocket and ligand. arXiv preprint Xix::2206..1 12241, 2022b.
10

---

Published as a conference paper at ICLR 2024
Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z Li. Semiretro: Semi-template framework boosts deep retrosynthesis prediction. arXiv preprint arXiv:2202.082205 2022c
Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z Li. Simvp: Simpler yet better video prediction. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3170-3180, 2022d.
Zhangyang Gao, Cheng Tan, and Stan Z. Li. Pifold: Toward effective and efficient protein in-
verse folding. In International Conference Learning Representations, 2023. URL https:
//openreview. net./ wtt.:rrrmmmmm....ttt/t
Joe G Greener, Lewis Moffat, and David T Jones. Design of metalloproteins and novel protein folds using variational autoencoders. Scientific reports, 8(1):1-12, 2018.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp.
2020.
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie. Tom Sercu, Adam Lerer, and Alexan- der Rives. Learning inverse folding from millions of predicted structures. bioRxiv, 2022.
Bin Huang, Tingwe Fan, Kaiyue Wang, Haicang Zhang, Chungong Yu, Shuyu Nic, Yangshuo Qi, Wei-Mou Zheng, Jian Han, Zheng Fan, et al. Accurate and efficient protein sequence design through learning concise local environment of residues. bioRxiv, 2022.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700 2017.
John Ingraham, Vikas K Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. 2019.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786-9801. PMLR, 2022.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint (XXXx:0009.01411 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 589, 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 information and modeling, 60(12):5667-566 2020.
Lisa N Kinch, R Dustin Schaeffer, Andriy Kryshtafovych, and Nick V Grishin. Target classification in the 14th round of the critical assessment of protein structure prediction (casp14). Proteins: Structure. Function, and Bioinformatics, 89(12):1618-666) 2021.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: www.::10990022... 2016.
Alex J Li, Mindren Lu, Israel Tilahum Desta, Vikram Sundar, Gevorg Grigoryan, and Amy Keat- ing. Neural network-derived potts models for baseed protein design using backbone atomic coordinates and tertiary motifs. bioRxiv, 2022a.
Alex J Li. Vikram Sundar, Gevorg Grigoryan, and Amy E Keating, Terminator: A neural framework for structure-base based protein design using tertiary repeating motifs. arXiv preprint arXiv:2204. 13048, 2022b.
11

---

Published as a conference paper at ICLR 2024
Jie Li and Patrice Kochl. 3d representations of amino acids applications to protein sequence com- parison and classification. Computational and structural biotechnology journal, 11(18):47-58, 2014.
Siyuan Li, Zedong Wang. Zicheng Liu, Di Wu, and Stan Li. Openmixup: Open mixup toolbox and benchmark for visual representation learning. arXiv preprint arXivv2209.044 .0851,, 2022c.
Zhixiu Li, Yuedong Yang. Eshel Faraggi, Jian Zhan, and Yaoqi Zhou. Direct prediction of profiles of sequences compatible with protein structure by neural networks with fragment-based local and energy-based nonlocal profiles. Proteins: Structure, Function, and Bioinformatics, 82(10): 2573, 2014.
Ziyao Li, Xuyang Liu, Weijie Chen, Fan Shen, Hangrui Bi, Guolin Ke, and Linfeng Zhang. Uni- fold: An open-source platform for developing protein folding models beyond alphafold. bioRxiv, 2022d.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos San- tos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences the scale of evolution enable accurate structure prediction. bioRxiv, 2022.
Yufeng Liu, Lu Zhang, Weilun Wang, Min Zhu, Chenchen Wang, Fudong Li, Jiahai Zhang, Houqiang Li, Quan Chen, and Haiyan Liu. Rotamer-free protein sequence design based on deep learning and self-consistency. 2022.
Jack B Maguire, Daniele Grattarola, Vikram Khipple Mulligan, Eugene Klyshko, and Hans Melo. Xenet: Using a new graph convolution to accelerate the timeline for protein design quantum computers. PLoS computational biology, 17(9)ee009033,, 2021.
Matt McPartlon, Ben Lai, and Jinbo Xu. A deep se (3)-equivariant model for learning inverse protein folding. bioRxiv, 2022.
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. bioRxiv, 2021. doi: 10. 10.1101/2021.00 .07 09 450648. URL http::://// www. www.:iiooxivvv.ooom www.biorxii..rrg//cnteett
10.1101/2021 07. 09.44448
Milot Mirdita, Konstantin Schütze, Yoshitaka Moriwaki, Lim Heo, Sergey Ovchinnikov, and Martin Steinegger. Colabfold: making protein folding accessible to all. Nature Methods, pp. 1-4, 2022.
James O'Connell, Zhixiu Li, Jack Hanson, Rhys Heffernan, James Lyons, Kuldip Paliwal, Abdollah Dehzangi, Yuedong Yang, and Yaoqi Zhou. Spin2: Predicting sequence profiles from protein structures using deep neural networks. Proteins: Structure, Function, and Bioinformatics, 86(6): 629-633, 2018.
Christine A Orengo, Alex D Michie, Susan Jones, David T Jones, Mark B Swindells, and Janet M Thorton. Cath-a hierarchic classification of protein domain structures. Structure, 508) 1109, 1997.
Sergey Ovchinnikov and Po-Ssu Huang. Structure- based protein design with deep learning. Current opinion in chemical biology, 65:136114, 2021.
Carl Pabo. Molecular technology: designing proteins and peptides. Nature, 301(58977:200-200, 1983.
Robin Pearce and Yang Zhang. Deep learning techniques have significantly impacted protein struc- ture prediction and protein design. Current opinion in structural biology, 68:194-207, 2021.
William R Pearson and Michael L Sierk. The limits of protein sequence comparison? opinion in structural biology, 2(3):254-260. 2005.
Current
Oscar Pina and Verónica Vilaplana. Liser-wise training for self-supervised learning on graphs. arXiv preprint 03ti:23009005003 2023.
12

---

Published as a conference paper at ICLR 2024
Yifei Qi and John ZH Zhang. Denseepd: improving the accuracy of tational protein sequence design with densenet. Journal of chemical information and modeling, 6063):2245-1252, 2020.
compu-
Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek, Alexander WR Nelson, Alex Bridgland, et al. Protein structure prediction using multiple deep neural networks in the 13th critical assessment of protein structure prediction (casp13). Proteins: Structure, Function, and Bioinformatics, 87(12): 141-1148, 2019.
Hannes Stärk, Octavian Ganea, Lagnajit Pattanaik, Regina Barzilay, and Tommi Jaakkola. Equibind: Geometric deep learning for drug binding structure prediction. In International Conference on PMLR, 2022.
Machine Learning, pp.
Alexey Strokach and Philip M Kim. Deep generative modeling for protein design. Current opinion in structural biology, 72:226-236, 2022.
Alexey Strokach, David Becerra, Carles Corbi- Verge, Albert Perez-Riba, and Philip M Kim. Fast and flexible protein design using deep graph neural networks. Cell Systems, 11(4):402-411, 2020.
Cheng Tan, Zhangyang Gao, and Stan Li. Simvp: Towards simple yet powerful spatiotemporal predictive learning. arXiv preprint uTXiv:2211.12509, 2022a.
Cheng Tan, Zhangyang Gao, and Stan L.i Target-aware molecular graph generation. arXiv preprint arXiv:2202.0482, 2022b.
Cheng Tan, Zhangyang Gao, Jun Xia, and Stan Li. Generative de novo protein design with global context. arv preprint arXiv:2204. 11tt222244.077,,,. 2022c.
Wen Torng and Russ B Altman. 3d deep convolutional neural networks for amino acid environment similarity analysis. BMC bioinformatics, 18(1):1-23, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1710.1110. 2017.
Guoxia Wang, Xiaomin Fang, Zhihua Wu, Yiqun Liu, Yang Xue, Yingfei Xiang, Dianhai Yu, Fan Wang, and Yanjun Ma. Helixfold: An efficient implementation of alphafold2 using paddlepaddle. arXiv preprint 2rXiv:2207 200.552777.04477, 2022.
Jingxue Wang, Huali Cao, John ZH Zhang, and Yifci Qi. Computational protein design with deep learning neural networks. Scientific reports, 8(1):1-9, 2018.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma. and Taco S Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018, URL
http://s....pppeee//////1/ pdf
(18844045200caaa waw.e00980001111111101----1
Ruidong Wu, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, et al. High-resolution de novo structure prediction from primary sequence. BioRxiv, 2022.
Zachary Wu, Kadina Johnston, Frances H Amoll, and Kevin K Yang. Protein sequence design with deep generative models. Current opinion in chemical biology, 65:18-27, 2021.
Ningyu Zhang, Zhen Bi, Xiaozhuan Liang, Siyuan Cheng, Haosen Hong, Shumin Deng, Jiazhang Lian, Qiang Zhang, and Huajun Chen. Ontoprotein: Protein pretraining with gene ontology embedding. arXiv preprint arXiv:2201.11144, 2022a.
Yang Zhang and Jeffrey Skolnick. Tm-align: a protein structure alignment algorithm based on the tm-score. Nucleic acids research, 33(7):2302-2309 2005.
Yuan Zhang, Yang Chen, Chenran Wang. Chun-Chao Lo, Xiuwen Liu, Wei Wu, and Jinfeng Zhang. Prodconn: Protein design using a convolutional neural network. Proteins: Structure, Function, and Bioinformatics, 88(7)8119-22) 2020a.
13

---

Published as a conference paper at ICLR 2024
Ziwei Zhang, Peng Cui, and Wenwu Zhu. Deep learning graphs: A survey. IEEE Transactions on Knowledge and Data Engineering, 2020b.
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. arXiv preprint arXiv:2203.06125 2022b.
Zaixiang Zheng, Yifan Deng, Dongyu Xue, Yi Ziou, Fei Ye, and Quanquan Gu. Stracture-informed language models are protein designers. bioRxiv, pp. 2023-02 2023.
Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. Graph neural networks: A review of methods and applica- tions. Al Open, 1:57-81, 2020.
14

---

Published as a conference paper at ICLR 2024
A MORE ANALYSIS
Module-wise training Regarding to the module-wise training strategy, we employ a stagewise
fashion for sequence refinement. This is represented by the equation s(() = Laer (8 (s(1-1); 11)))) where s()) represents the refined sequence generated by the l-th refinement module facs, and (1- denotes the pretrained embedding extracted from (l--1), The module-wise training could signifi- cantly reduce the computational cost and helpful for alleviating oversmoothing:
1. Stantard: For given protein, during each epoch, the parameters of layers / and I are updated, resulting in a change in 8( Due to the large size of the pretrained model (ESM2-650M),, performing a forward pass to obtain the updated version of x(() consumes significant GPU memory and increases the time cost. Consequently, conducting end-to-end training would significantly increase the computational cost.
2. Ours: By fixing g() before training g(+1), can ensure that s()) is retrained without
the need for updating (0) This approach allows to initialize g()) in the first epoch and
reuse it in subsequent epochs, effectively avoiding redundant forward passes of the large pretrained model.
3. The overall KWDesign is x 10-layer GNN model, where each refienment module consists of 10 GNN layers. It well known that GNN models suffer from the issue of oversmooth- ness when training in an end-to-end fashion. To overcome this, we adopt module-wise training to ensure that the pre-placed module serves as good initialization for the subse- quent module. As discovered by (Pina & Vilaplanal 2023), the layer-wise training promotes the node features to be uncorrelated at each single layer to alleviate oversmoothing.
Distribution comparison Fig. shows the confidence distributions of positive and negative residues generated by PiFold and KW-Design on the CATH4.2 test set. Pos- itive residues tend towards a confidence of 1.0, while negative residues have mostly below 0.6 confidence, in- dicated by different colors. Our results demonstrate that KW-Design produces positive residues with higher confi- dence compared to PiFold, while also reducing the num- ber of negative residues. This suggests that KW-Design can convert low-confidence positive residues to high- confidence ones and correct negative residues as positive ones.
Ros
Mag
PostOal Neg10un
02
0.4
0.8
Figure 4: Confidence distributions.
Compare structures In Fig 5 we use ESMFold(Lin et al. 2022) to generate protein structures from designed sequences and compare the designed proteins of PiFold and KW-Design against the reference ones. We observe that the designed structures of KW-Design are more similar to the reference ones than those of PiFold. Specifically, KW-Design achieves 15.9% 35.3%, and 60% improvement in structural mean squared eтror (MSE) on the la73, 1a81, and lacl proteins, respectively. These results demonstrate that KW-Design can generate proteins that are structurally more similar to the reference ones compared to PiFold.
(3) POUID 1873
POSILL 1681
PIFodd MSE. 0.50
PIFad MSE:
Figure Comparing the designed proteins. The green structures are reference ones, while the gray and purple structures are designed by PiFold and KW-Design, respectively. We use red circles to highlight the regions where KW-Design produces more similar structures to the reference ones than PiFold.
15

---

Published as a conference paper at ICLR 2024
B DISCUSSION ABOUT DIVERSITY
TS45 In addition to designing single- and multi-chain proteins, we also include a set of de novo proteins collected from the CASPI5 competition to provide a more realistic assessment (Senior et al. 2019 Kinch et al., 2021). The Critical Assessment of Protein Structure Prediction (CASPI5), which took place from May through August 2022, was held after the release dates of CATH4.3 (July 1, 2019) and PDB (August 2, 2021). In CASPI5, diverse protein targets are introduced, including FM (Free Modeling), TBM (Template-Be Baste Modeling), TBM-easy. easy, and TBM-hard proteins. There are 18 FM, 25+2 TBM (including 20 TBM-eazy, 5 TMBB-ard, 2 FM/TBM). The FM targets have no homology to any known protein structure, making them particularly suitable for de novo protein design. The TBM targets have some homology to known protein structures, while the TBM-casy targets are relatively easy TBM targets. The TBM-hard targets are more difficult TBM targets, with lower levels of sequence identity known structures. We download the public TS-domains struc- tures from CASP15 which consists of 45 structures, namely TS45. We use TS45 as a benchmark for de novo protein design, as the structures less similar known structures and were not determined prior to the construction of the training sets.
Diversity Definition To improve the success rate of protein design, it is important to explore set of protein sequences rather than placing bet on single sequence. In this case, generating diverse sequences crucial for exploring the reasonable protein sequence space. define the pairwise
diversity (Jain et al. 2022) as Di Σi= huaaaa where Ti indicates the I-th residue of the i-th
designed sequence. The overall diversity score is
D&J m
Div= Σ Lj
(14)
where i.j (1, m and m is the number of totally designed sequences. By default, set m = 10. However, measuring diversity alone without combining it with other metrics may be misleading. For example, high diversity indicates low recovery rate, more likely to result in low structural similarity.
Experiments & Analysis We benchmark the diversity on TS45 dataset using models pre-trained
on CATH4.3. As discovered by previous research (Hsu et al.) 2022 Dauparas et al. (022), the sampling temperature affects diversity. Denote the temperature as T, predicted prohability vector p E 1R,,200 we sample new sequences from the distribution of Mu Mult nomia a (so ft ft.aa We vary the temperature from 0.0 to 0.5 and plot the trends of recovery and diversity in Fig. 66] Under the same sampling temperature, high recovery leads to decreased diversity. However, at the same level of recovery, stronger models have higher diversity. This phenomenon can be attributed to the fact that stronger models, such KWDesign, PiFold, and ProteinMPNN, exhibit relatively high confidence levels for a larger number of residues. Even after applying smoothing to the probability distribution, the recovery remains consistently high, while noticeable improvement is observed in terms of diversity.
0.9
0.5
Diversity
0.8
0.7
0.4 Recovery
0.6
StructGNN
0.5
GraphTrans
GCA
GVP
0.3
0.2
AlphaDesign
0.2
ProteinM
Pifeld
0.1
0.1
KWDesign
0.0
0.1
0.2
0.3
0.4
0.5
Figure 6: The trends of recovery and diversity.
16

---

Published as a conference paper at ICLR 2024
C MORE COMPARISON
Compare to LMDesign In parallel with our work, we are observing another exciting project called LMDesign (Zheng al. 2023), which was recently published at ICML as an oral presentation. LMDesign aims to use the pre-trained ESM model to improve protein design. However, there are several differences between our Showledge-Desigeee and LMDesign.
More comprehensive: We enhance protein design by fusing multimodal knowledge from pre-trained models, including both structural and sequential information, while LMDesign only single-modal information. Our experiments demonstrate that combining these modalities leads to nontrivial improvements, as shown in Table
More efficient: We introduce the memory-retrieval mechanism to save more than 50% of the training time, while LMDesign does not this mechanism.
Novel modules: We introduce confidence-aware recycling techniques as well virtual MSA to boost the model performance.
More effective:
Overall, model outperforms LMDessggn by 5.12% the CATH4.2
dataset.
Table Results comparison on the CATH dataset. baselines are reproduced under the same code frame- work, except ones marked with t. We copy results GVP-large and ESM-IF from their manuscripts (Hsu
ct al. 2022) The best and suboptimal results are labeled with bold and underline.
Perplexity +
Recovery f
CATH
version
Model
Short
cingle-chain
All
Short
Single-chain chain
All
4.2
4.3
StrucGGNN
8.29
8.74
6.40
29.44
28.26
35,91
V
GraphTrans
8.39
8.83
6.63
28.14
28.46
35.82
GCA
7.09
7.49
6.05
32.62
31.10
37.64
GVP
7.84
30.60
28.95
39.47
GVP-large"
7.68
6.12
6.17
32.6
39.4
39.2
AlphaDesign
7.32
7.63
6.30
34.16
32.66
41.31
ESM-IFT
8.18
6.33
6.44
31.3
38.5
38.3
ProteinMPNNN
4.61
36.35
34.43
45.96
PiFold
6.04
6.31
4.55
39.84
38.53
51.66
LMDesign
6.77
6.46
4.52
37.88
42.47
55.65
Knowledge-Design (Ours)
5.48
5.16
3.46
44.66
45,45
60.77
17

---

Published as a conference paper at ICLR 2024
se-TM The structural similarity is the ultimate standard for measuring the quality of the designed sequence. However, the structures of designed protein sequences needed to be predicted by other al- gorithms, such AlphaFold (Jumper al. (02T), OmegaFold (Wu al. 2022) and ESMFold (Lin) et al., 2022). The protein folding algorithm itself has certain inductive bias and will cause some prediction errors, which will affect the evaluation. To overcome the inductive bias, introduce the self-consistent TM-score (sc-TM) metric:
SC-TM (MSc
(15)
where f is the protein folding algorithm and TMscore(;.) is widely used metric (Zhang & Skolnick, 2005) for measuring protein structure similarity. Since the structures of the designed sequence and reference sequence inductive bias is expected to be canceled out when calculating the TM-score. This approach results in more robust metric, called the sc-TM, that is less affected by the inductive bias of the protein folding algorithm.
predicted by the same protein folding algorithm, the model's
se-TM Results On the CASP15 dataset, we investigate the recovery and sc-TM metries as the temperature increases. To compute sc- Sc.TM we utilize AlphaFold to predict protein structures from sequences. According 7 and Fig.6 we observe that a slight increase in temperature from O to 0.1 is beneficial in significantly enhancing diversity while maintaining good recovery and sc- TM. Specifically, when the sampling temperature set to 0.0 and 0.1, KWDesign outperforms the baselines, achieving the highest sc- sc-orr.
63
Figure The statistics recovery and sc-TM with increasing temperature.
18

---

Published as a conference paper at ICLR 2024
Do pretrained language models prefer to correct disordered regions? Thanks to the insightful comments of Reviewer 4EsM, we investigate the preferred regions that pretrained language models tend to correction. Expeinments follow five steps:
Select PDBs containing diverse protein structures, run the pretrained PiFold get the de- signed sequence Sps
2. Further use the pretrained KWDesign to refine Spi. resulting in the optimized sequence Skw. All the used PiFold and KWDesign are pretrained on CATH4.3.
3. For each residue, we record whether PiFold has successfully recovered the residue type (goodps) or not (badps). Similarly, we record goodkw and badkw.
4. We list four states for each residues and set different color for each state: goodgi (red,, badpi - (gdx goodpi + (Whtee( bad + goodsn(green)
5. We show the colored proteins with Mol and present them in Fig g
Our finding is that pretraining knowledge corrects for both disordered and ordered regions. However, it seems that the correction more pronounced for regions with well-defined secondary structures (a-helix and B-sheet).
1kllA00
1ud9A00
2d9eA00
3fkfA00
5nckA01
17gsA001
Figure 8: The colored proteins.
19


Reviews:
[{"review_id": "88FISKPfdb", "reviewer": "Reviewer_Mme4", "summary": "This work proposes a refining method, called KW-design, to refine the generated protein sequences from previous methods, such as PiFold. In specific, KW-design is an iterative process that uses pretrained sequence and structure models to extract multi-modal features and a confidence-aware gated layer to fuse these features. It also applies a memory retrieval mechanism to save the training time.", "questions": "See my comments in the above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The idea of iteratively refining generated sequences from protein structures based on a confidence-aware gated fusion layer that fuses features from multi-modal feature extractors looks novel to me. \n- Experiments across multiple datasets show the effectiveness of the proposed method. \n- Ablation studies are well conducted to show the impact of main components.", "weaknesses": "My main concern is on the clarification of the paper. I feel like a large revision is needed to make the presentation more clear to the readers. \n- The method section is not very clear. In particular, Section 3.2 and 3.3 have a large portion of redundancy. For instance, I think “Knowledge extractor & Confidence predictor” in Section 3.3 seems to just repeat the content in “Virtual MSA” in Section 3.2. “Fuse layer” can be merged to “Virtual MSA” to make the presentation more concise. Besides, I don’t know how Eq (13) is related to Eqs (8) and (9). It looks like they are both about how to update the embeddings, but how does “PiGNN” process $e^{(l)}$ to make Eq. (13) consistent to Eqs (8) and (9)?\n- In “Confidence-aware updating”, why do you need two MLPs to process the same input (one is the negative sign of another) for the gates? Any justification on Eq. (11)?\n- In Algorithm 1, I’m not sure how the optimal $\\phi^{(l)}$ is determined and how the “patience of 5” is applied. Also, in line 12 (Algorithm 1), what does $n$ mean in $h_n^{(l+1)}$?\n- I wonder why not split Table 1 into two subtables: one for CATH 4.2 and another for CATH 4.3?\n- In Table 2, what does the “Worst” metric mean? Why does the proposed method perform worse than all other methods?\n- Minor typo issue. “settings Specifically” -> “settings. Specifically” on page 6. “3d CNN” -> “3D CNN” on page 3.", "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": "", "extra_scores": {}}, {"review_id": "V52MzjJqxT", "reviewer": "Reviewer_KJsj", "summary": "The paper introduces KW-Design, a novel method for protein design that iteratively refines low-confidence residues using knowledge extracted from pretrained models. The approach incorporates a multimodal fusion module, virtual MSA, recycling technologies, and a memory-retrieval mechanism to enhance performance and efficiency. The method demonstrates substantial improvements across various benchmarks, achieving over 60% recovery on datasets such as CATH4.2, CATH4.3, TS50, TS500, and PDB.", "questions": "1. Is there any fine-tuning done on the pretrained language model used in your approach?\n2. In Section 4.3, Table 3 claims that “the predictive confidence score, an unsupervised metric, exhibits a strong correlation with recovery.” Could the authors provide a more detailed analysis, perhaps including the Spearman correlation between these two values?\n3. Regarding the virtual MSA, does the sequence order affect the resulting residue embedding? If not, what criteria are used to determine the sequence order?\n4. In Section 3.3, the initial node and edge features are extracted using PiFold. Is PiFold fine-tuned or kept fixed during this process?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written and exhibits a clear and logical structure.\n2. The proposed method effectively leverages the knowledge from pretrained protein sequence/structure models, resulting in notable benefits for protein sequence design.\n3. The paper includes a thorough ablation study, examining different components of the models such as recycling, virtual MSA numbers, and the pretrained model. This is crucial for gaining a deep understanding of the proposed methodology.", "weaknesses": "1. The code associated with the paper is currently unavailable.\n2. Given that the model relies on pretrained models, some of which have been trained on the test set utilized, there is a potential risk of data leakage.\n3. The paper predominantly employs perplexity and recovery as metrics for evaluating the designed sequences. However, there is a chance that the designed proteins may not be soluble or may not fold correctly to the given backbone. It would be beneficial for the authors to incorporate additional metrics (e.g., scTM score, solubility) in their evaluation.", "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": "", "extra_scores": {}}, {"review_id": "XaE3cSudbJ", "reviewer": "Reviewer_4EsM", "summary": "The paper addresses the 'inverse protein folding'  problem of designing a protein sequence that folds into a particular shape. They achieve strong results on standard benchmarks.  It draws on a number of interesting ideas, from graph neural networks to leveraging embeddings of pretraining language models, to a 'recycling' approach that updates predictions based on the current uncertainty over those predictions.", "questions": "Can you please clarify my questions regarding train-test splits and data leakage from the 'knowledge' module?\n\nCan you please explain why you use such a non-standard training procedure?\n\nI have raised my score to weak accept. We had an extensive back and forth regarding train-test splits and I appreciate the attention that the authors have devoted to this important topic.\n\n\n==Comments after author's response==\nMy principal concerns about the paper were (1) data leakage and (2) the unconventional presentation of a multi-layer neural network architecture and stagewise training procedure.\n\nRegarding (1), the overall challenge with this research field is that the primary application of inverse folding models is to find sequences that fold into de-novo designed protein structures. However, to validate models on offline natural data, the sequence-structure pairs in the test set aren't de-novo, particularly if they were seen in pretraining. The authors created a new test set at my request, which is small, but establishes good community standards for how to think about  these issues in the future.\n\nRegarding (2), please update the paper to use more standard terminology for a multi-layer model. When describing the stagewise training procedure, please make it clear in updated versions of the manuscript that this was done for computational reasons, and the superiority of it in terms of achieving a better model is not evaluated empirically.\n\nFinally, note that in future work an alternative to your stagewise training approach would be to use gradient checkpointing, which reduces memory overhead (at the expense of additional forward passes).", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper achieves really strong empirical results on a panel of common benchmarking setups. The results will definitely be of interest to the community.", "weaknesses": "The paper's explanation of the model is extremely difficult to understand because it is not using standard terms for deep neural networks. For example, 'knowledge' is used in a vague way, I guess, to mean leveraging pretraining? If I understand correctly, the composition of functions in section 3.1 is just describing a multi-layer neural network, but uses very verbose and non-standard notation. Further, why are the layers trained in a stagewise fashion in section 3.4 instead of standard back-propagation? It was confusing to me why you chose this, since it is much more complex to implement. Does it provide better performance?\n\nThe paper's experiments are quite careful about train-test splits, using a number of clustering-based splits that are well-established in the literature. I'm concerned about data leakage, however, from the pretrained language models used to help guide predictions for low-confidence residues in the 'knowledge' module. There models were trained across vast numbers of proteins and likely do not follow the same train-test splits as for the structure -> sequence benchmarks. As a result. The exact target sequence for structure -> sequence design may have appeared in the LM pretraining data. This may explain why the paper's method is able to increase the per-residue recovery rate so dramatically. The paper also mentions that the ESM-IF train-test split is not compatible with some of the other benchmarking setups, yet ESM-IF embeddings are used here.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)