SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhangyang Gao  $^{1,2,\dagger}$ , Cheng Tan  $^{1,2,\dagger}$ , Yijie Zhang  $^{3}$ , Xingran Chen  $^{4}$ , Lirong Wu  $^{1,2}$ , Stan Z. Li  $^{2,*}$ $^{1}$  Zhejiang University
$^{2}$  AI Lab, Research Center for Industries of the Future, Westlake University  $^{3}$  McGill University,  $^{4}$  University of Michigan
Protein inverse folding has attracted increasing attention in recent years. However, we observe that current methods are usually limited to the CATH dataset and the recovery metric. The lack of a unified framework for assembling and comparing different methods hinders the comprehensive investigation. In this paper, we propose ProteinInvBench, a new benchmark for protein design, which comprises extended protein design tasks, integrated models, and diverse evaluation metrics. We broaden the application of methods originally designed for single-chain protein design to new scenarios of multi-chain and de novo protein design. Recent impressive methods, including GraphTrans, StructGNN, GVP, GCA, AlphaDesign, ProteinMPNN, PiFold and KWDesign are integrated into our framework. In addition to the recovery, we also evaluate the confidence, diversity, sc-TM, efficiency, and robustness to thoroughly revisit current protein design approaches and inspire future work. As a result, we establish the first comprehensive benchmark for protein design, which is publicly available at https://github.com/A4Bio/OpenCPD.
Protein inverse folding is a fundamental problem in biology and has many applications in medicine, agriculture, and bioenergy [1-4]. It has thereby attracted increasing attention in both the machine learning and biology communities [5, 6]. Traditional physical-inspired methods suffer from the problem of expensive computation and unsatisfactory accuracy. Recently, deep learning methods have shown great potential in science areas [7-30], including protein inverse folding, by simplifying the process and improving accuracy [12, 31-62]. Among them, we observe that graph-based methods achieve state-of-the-art performance. However, previous methods are usually limited to the CATH dataset and the recovery metric. We emphasize that recovery is not the only important metric for protein design. Other metrics such as confidence, diversity, TM-score, efficiency, and robustness are also important for comprehensively revisiting current approaches. Also, the evaluation dataset should be further extended from CATH dataset to broader or more difficult cases to facilitate practical applications. All these challenges motivate us to establish ProteinInvBench, a unified benchmark for protein inverse folding, in which multiple tasks, models, and metrics are introduced and integrated.
ProteinInvBench extends the task of protein design from single-chain to multi-chain and de novo protein design. To our knowledge, many computational protein design methods [41, 46, 47, 50, 63] have only been evaluated on the outdated single-chain structure dataset CATH4.2. Furthermore, few studies [48, 56] have investigated protein design performance in multi-chain tasks, and even fewer have evaluated the performance comparison on de novo protein structures. To fill this knowledge gap,
$\dagger$  Equal Contribution, \*Corresponding Author.
we first benchmark open-sourced graph-based models on the latest CATH4.3 dataset and extend them to the case of multi-chain protein design. We then collect de novo protein structures with little or no similarity to existing structures from the CASP15 competition. Evaluating models on the CASP15 dataset allows us to gain a better understanding of the potential of AI models in designing de novo proteins and reveals that different models exhibit non-trivial differences in generalizability. We hope that the more complex and challenging tasks in ProteinInvBench will facilitate the development of protein design methods for practical applications.
ProteinInvBench also provides a range of metrics to comprehensively understand the strengths and weaknesses of each method. Beyond the recovery score that measures the percentage of exactly matched residues, we also evaluate the confidence, diversity, sc-TM, efficiency, and robustness. These metrics will be introduced in detail in the Sec.5. Notably, we use novel metrics such as confidence and sc-TM to measure the quality of designed sequences in an unsupervised and unbiased manner, respectively. With additional metrics, we encourage researchers to develop more robust, efficient models that can generate diverse proteins for higher success rates in wet experiments.
Tasks: Designing proteins over single-chain, multi-chain and de-novo structures
Models: Recent strong baselines that provide open-source training code
Metrics: The most comprehensive evaluation metrics to date
Figure 1: The framework of the proposed benchmark. The benchmark is organized incrementally from tasks to models, to metrics. We color contents in green and blue that are widely and partially considered by previous studies, respectively. Newly introduced contents are colored in pink.
Based on the constructed benchmark, we conduct extensive experiments to extend baseline models to new tasks and evaluate them on diverse metrics. All models are reproduced, integrated, and extended under a unified code framework. ProteinInvBench not only reproduces the reported results of baselines but also provides new insights into the detailed strengths and weaknesses of each method under different scenarios. To summarize, our contributions are as follows:
1. Tasks: We extend recent impressive models from single-chain protein design to the scenarios of multi-chain and de novo protein design.
2. **Models:** We integrate recent impressive models into a unified framework for efficiently reproducing and extending them to custom tasks.
3. Metrics: We incorporate new metrics such as confidence, sc-TM, and diversity for protein design, and integrate metrics including recovery, robustness, and efficiency to formulate a comprehensive evaluation system.
4. Benchmark: We establish the first comprehensive benchmark of protein design, providing insights into the strengths and weaknesses of different methods.
2 Overall Framework and Problem Definition
Overall Framework We present the overall framework in Figure 1, which contains three components: (1) Multiple datasets for the task of single-chain, multi-chain, and de-novo protein design. (2) Strong baselines are integrated into our unified framework, including GraphTrans, GVP, GCA, AlphaDesign, ProteinMPNN, PiFold, and KWDesign. (3) Diverse metrics are used to evaluate the designed proteins in a quantitative and comprehensive manner.
Problem Definition The protein inverse folding problem [64] aims to find the amino acids sequence  $S = \{s_i : 1 \leq i \leq n\}$  folding into the desired structure  $\mathcal{X} = \{X_i \in \mathbb{R}^{m,3} : 1 \leq i \leq n\}$ , where  $m$  is the maximum number of points belonging to the  $i$ -th residue,  $n$  is the number of residues and the natural proteins are composed by 20 types of amino acids, i.e.,  $1 \leq s_i \leq 20$  and  $s_i \in \mathbb{N}^+$ . Therefore, the protein inverse folding problem is usually formulated as a structure-to-sequence learning problem, where the goal is to learn a function  $\mathcal{F}_{\theta}$ :
$$
\mathcal {F} _ {\theta}: \mathcal {X} \mapsto \hat {\mathcal {S}}. \tag {1}
$$
The function  $\mathcal{F}_{\theta}$  is usually parameterized by a neural network, and the parameters  $\theta$  are learned by minimizing the cross-entropy loss, i.e.,  $\mathcal{L}(\mathcal{F}_{\theta}(\mathcal{X}),\mathcal{S}) = -\sum_{i = 1}^{n}\log s_{i}p(\hat{s}_{i}|\mathcal{X},\theta)$
CATH The CATH (Class, Architecture, Topology, Homology) [65] database is a comprehensive resource for protein structure classification that hierarchical group proteins based on their structural features. The database defines classes based on topological similarities, architectures based on the arrangement of secondary structure elements, topologies based on the connectivity of secondary structure elements, and homologous domains based on sequence similarity. Previous works such as GraphTrans [50], GVP [41], and GCA [46] have used the CATH4.2 version of the database as a benchmark for protein design, which splits the dataset by CATH topology classification and includes 18,024 proteins for training, 608 proteins for validation, and 1,120 proteins for testing. However, CATH4.2 is an outdated version and may not represent the current protein structure space. To address this, we use the newer version, CATH4.3 for benchmarking protein design and follow the data splitting protocol of ESMIF [53]. This results in a training set of 16,153 structures, a validation set of 1,457 structures, and a test set of 1,797 structures. Note that the curated CATH dataset contains only single-chain structures and does not consider the case of designing multi-chain proteins.
PDB The Protein Data Bank (PDB) [66] is a comprehensive database of 3D structural data for biological molecules. To study multi-chain protein design, we utilized a ProteinMPNN dataset derived from PDB assemblies with high resolution and less than 10,000 residues. 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,464), and test sets (1,539), ensuring that none of the chains from the target chain or biounits of the target chain were present in the other two sets. During each training epoch, we cycled through the sequence clusters and randomly selected a sequence member from each cluster. This dataset was used for the task of multi-chain protein design, expanding the comparison of computational protein design methods, as many previous methods were omitted in this task.
TS45 In addition to designing single- and multi-chain proteins, we also include a set of de novo proteins collected from the CASP15 competition to provide a more realistic assessment [67, 68]. The Critical Assessment of Protein Structure Prediction (CASP15), 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 CASP15, diverse protein targets are introduced, including FM (Free Modeling), TBM (Template-Based Modeling), TBM-easy, and TBM-hard proteins. There are 18 FM,  $25 + 2$  TBM (including 20 TBM-eazy, 5 TMB-hard, 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-easy targets are relatively easy TBM targets. The TBM-hard targets are more difficult TBM targets, with lower levels of sequence identity to known structures. We download the public TS-domains structures from CASP15 which consists of 45 structures, namely TS45. We use TS45 as a benchmark for de novo protein design, as the structures are less similar to known structures and were not determined prior to the construction of the training sets.
4 Baseline models
We ensemble recent strong protein design baselines under the unified framework. Currently, we support open-sourced graph methods, such as GraphTrans, StructGNN [50], GVP [41], GCA [46], AlphaDesign [47], ProteinMPNN [48], PiFold [63] and KWDesign [69], that we can access their codes and training scripts. StructGNN and GraphTrans [50] employ C-alpha for geometric features, while GCA [46] adds global attention. GVP [41] introduces a novel GNN layer for invariant and equivariant features. However, these methods suffer from poor inference efficiency due to autoregressive decoders. To overcome this, AlphaDesign [47] replaces the decoder with an iterative 1D CNN. Recent advancements include ProteinMPNN [48], which incorporates additional structural information, and PiFold [63], a combination of AlphaDesign and ProteinMPNN. KWDesign [69] is an ensemble model that utilizes PiFold to create a prompt template. It refines the template using pre-trained knowledge, including sequence pretraining (ESM-650M [70]) and structure pretraining models (ESMIF's encoder [53]).
New baselines [56, 71] will continue to be added in the future. We have not included ESMIF [53] in our benchmark since it lacks a training script and is challenging for us to train. We plan to add it to the benchmark once we successfully train the model. According to the generation scheme, these baselines can be categorized as autoregressive, iterative, and one-shot models.
Autoregressive models consider both sequential and structural dependencies by combining the structural encoder and autoregressive sequence decoder, such as GraphTrans [50], GVP [41], GCA [46], and ProteinMPNN [48]. The protein sequences are generated by:
$$
p (\hat {\mathcal {S}} | \mathcal {X}; \theta) = \prod_ {t = 1} ^ {n} p \left(\hat {s} _ {t} \mid \hat {s} _ {<   t}, \mathcal {X}; \theta\right). \tag {2}
$$
Autoregressive models have been criticized for being slow in generating long proteins [63].
Iterative models generate residues in parallel and iteratively refine the generated sequence (AlphaDesign [47] and KWDesign [69]):
$$
\hat {S} ^ {(0)} \sim p (\hat {S} | \mathcal {X}; \theta^ {(0)}), \hat {S} ^ {(t)} \sim p (\hat {S} | \hat {S} ^ {(t - 1)}, \mathcal {X}; \theta^ {(t)}), \tag {3}
$$
where  $t$  indicates the refinement step, affecting the inference time costs.  $\theta$  is a learnable parameter.
One-shot models generate the protein sequence in parallel, e.g., PiFold [63], which is quite efficient in generating long proteins, written as:
$$
\hat {\mathcal {S}} \sim p (\hat {\mathcal {S}} | \mathcal {X}; \theta). \tag {4}
$$
In this section, we introduce metrics that will be used for protein design evaluation, including recovery, confidence, diversity, and sc-TM. Previous researches [41, 46, 47, 50, 53, 63] mainly focus on improving recovery, while ignoring other metrics. However, we argue that recovery is not the only important metric for protein design. Other metrics introduced follows are also crucial for comprehensively revisiting current approaches, such as confidence, diversity, sc-TM, efficiency, and robustness.
Recovery Recovery is the primary metric of the ability of the designed protein to recover its original residues, defined as:
$$
\operatorname {R e c} = \frac {1}{n} \sum_ {i = 1} ^ {n} \mathbb {1} \left(\hat {s} _ {i} = s _ {i}\right) \tag {5}
$$
where  $\mathbb{1}(\cdot)$  is the indicator function,  $\hat{s}_i$  is the designed residue at the  $i$ -th position, and  $s_i$  is the corresponding reference residue. A high recovery rate indicates that the designed protein sequence is similar to the reference sequence, and it is therefore expected that the folded structure will also be similar to the reference structure. Since measuring structural similarity is computationally expensive, previous protein design methods have placed a great deal of emphasis on improving recovery. However, it is important to note that recovery itself is a proxy metric for measuring structural similarity. In other words, a higher recovery rate does not necessarily ensure a higher level of structural similarity. Moreover, a high recovery rate may result in low diversity.
**Confidence Calculating recovery requires access to the reference sequence, which is not always available in practice.** When the ground-truth sequence is unknown, measuring and ranking the quality of the designed sequence becomes more challenging. We introduce the confidence metric to address this problem, which is the average predictive probability of designed amino acids, defined as:
$$
\operatorname {C o n f} = \frac {1}{n} \sum_ {i = 1} ^ {n} p \left(\hat {s} _ {i}\right) \tag {6}
$$
Diversity To improve the success rate of protein design, it is important to explore a set of protein sequences rather than placing a bet on a single sequence. In this case, generating diverse sequences is crucial for exploring the reasonable protein sequence space. We define the pairwise diversity [72] as  $D_{ij} = \frac{\sum_{l=1}^{n} \mathbb{1}_{r_{i,l} \neq r_{j,l}}}{n}$ , where  $r_{i,l}$  indicates the  $l$ -th residue of the  $i$ -th designed sequence. The overall diversity score is
$$
\operatorname {D i v} = \sum_ {i, j} \frac {D _ {i , j}}{m ^ {2}} \tag {7}
$$
where  $i,j\in \{1,2,3,\dots ,m\}$  and  $m$  is the number of totally designed sequences. By default, we set  $m = 10$ . However, measuring diversity alone without combining it with other metrics may be misleading. For example, a high diversity indicates a low recovery rate, more likely to result in a low structural similarity.
sc-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 algorithms, such as AlphaFold [73], RoseTTAFold [74], OmegaFold [75] and ESMFold [70]. The protein folding algorithm itself has a certain inductive bias and will cause some prediction errors, which will affect the evaluation. To overcome the inductive bias, we introduce the self-consistent TM-score (sc-TM) metric:
$$
\operatorname {s c} - \operatorname {T M} = \operatorname {T M S c o r e} (f (\mathcal {S}), f (\mathcal {S})) \tag {8}
$$
where  $f$  is the protein folding algorithm and TMscore(.,.) is a widely used metric [76] for measuring protein structure similarity. Since the structures of the designed sequence and reference sequence are predicted by the same protein folding algorithm, the model's inductive bias is expected to be canceled out when calculating the TM-score. This approach results in a more robust metric, called the sc-TM, that is less affected by the inductive bias of the protein folding algorithm.
Robustness Robustness measures an algorithm's ability to maintain its original performance under geometric perturbations. It is a useful metric for assessing the stability and generalizability of an algorithm. We define robustness as:
$$
\operatorname {R o b} = \operatorname {R e c} ^ {\prime} - \operatorname {R e c} \tag {9}
$$
where  $\mathrm{Rec}$  and  $\mathrm{Rec}^{\prime}$  are the recovery after and before applying small Gaussian perturbations to the Cartesian coordinates of the structure, correspondingly. As the template protein structures may not be perfect, more robust methods are expected to be more suitable in real-world applications.
Efficiency Efficiency measures the computational resources and time required to design a set of proteins. This study reports the training time, evaluation time, and model parameters of different methods over the standard benchmarks. While efficiency may not be a crucial problem compared to the recovery and sc-TM, it is a useful metric for assessing the model's scalability and practicality.
6 Benchmarking Protein Design
In this section, we retrain baselines on the newly introduced datasets and evaluate them using diverse metrics, resulting in a comprehensive benchmark. The experiments are organized as follows:
1. Establish a basic benchmark within recovery and confidence. As emphasized by previous studies, the recovery rate and predictive confidence are the most important and straightforward metrics. We benchmark baselines over these metrics on CATH4.2, CATH4.3, PDB, and TS45 for the task of single-chain, multi-chain, and de novo protein design, respectively. These results could serve as the basic benchmark for future studies.
2. Measuring diversity and sc-TM for practical challenging tasks in protein design. We further extend the evaluation metrics to diversity and sc-TM. The diversity is opposite of recovery and is meaningless if we measure it alone. By examining the sequence diversity and structural sc-TM together, we could have a more comprehensive understanding of the designable protein space.
3. Assessing the robustness when input structures are not perfect. Although the model performs well on natural proteins, it may fail when the artificially designed structure is noisy. In this case, the robustness of the model is crucial for practical applications. We evaluate the robustness of different methods by applying geometric perturbations to the template protein structures during the evaluation phase.
4. Comparing the efficiency. Towards designing efficient, scalable, and generalizable algorithms, we evaluate the efficiency in terms of training time, evaluation time, and model parameters to facilitate the development of more efficient protein design methods.
6.1 Recovery and Confidence
In this section, we benchmark the recovery rate and confidence of different methods on the CATH4.2, CATH4.3, PDB, and TS45 datasets to address the problems of single-chain, multi-chain, and de novo protein design. By extending from CATH4.2 to the newer CATH4.3 and from single-chain to multi-chain to de novo, we have constructed the most comprehensive benchmark to date for protein design. All models are retrained and evaluated under the same code framework for a fair comparison. The hyperparameters used for training models are provided in the Appendix.
Table 1: Single-chain results. The best and suboptimal results are labeled with bold and underlined.
Model length	Confidence ↑	Recovery % ↑
L<100	100≤ L<300	300≤ L<500	Full	L<100	100≤ L<300	300≤ L<500	Full
CATH4.2	StructGNN	0.31	0.45	0.45	0.43	0.26	0.36	0.36	0.35
GraphTrans	0.31	0.43	0.43	0.43	0.25	0.35	0.35	0.34
GCA	0.34	0.46	0.47	0.45	0.27	0.38	0.38	0.37
GVP	0.40	0.52	0.53	0.51	0.28	0.40	0.41	0.39
AlphaDesign	0.36	0.49	0.49	0.47	0.33	0.43	0.44	0.42
ProteinMPNN	0.38	0.51	0.52	0.50	0.32	0.47	0.47	0.45
PiFold	0.44	0.58	0.60	0.57	0.39	0.53	0.56	0.52
KWDesign	0.50	0.68	0.72	0.67	0.44	0.62	0.66	0.61
CATH4.3	StructGNN	0.35	0.41	0.47	0.41	0.30	0.34	0.40	0.34
GraphTrans	0.37	0.42	0.48	0.42	0.29	0.34	0.39	0.34
GCA	0.38	0.43	0.49	0.43	0.32	0.36	0.41	0.36
GVP	0.45	0.51	0.55	0.50	0.33	0.38	0.45	0.38
AlphaDesign	0.41	0.48	0.53	0.47	0.37	0.43	0.47	0.42
ProteinMPNN	0.42	0.49	0.57	0.49	0.38	0.44	0.52	0.44
PiFold	0.47	0.56	0.64	0.55	0.43	0.52	0.59	0.51
KWDesign	0.58	0.68	0.76	0.67	0.51	0.61	0.69	0.60
Single-chain Results The results of single-chain protein design are shown in Tab. 1, where both CATH4.2 and CATH4.3 datasets are included. We present metrics for proteins of different sequence lengths. From Tab. 1, it could be concluded that:
1. KWDesign and PiFold are the best and second-best models. They consistently outperform all other models in terms of both confidence and recovery across all protein lengths in both CATH4.2 and CATH4.3 datasets. This highlights their effectiveness towards protein inverse folding.
2. Models perform better on longer proteins. This could be due to the increased complexity and information available for longer proteins, allowing the models to make more confident predictions.
3. CATH4.2 and CATH4.3 datasets show the same performance trend and very similar results, thereby validating the performance consistency of the different models. However, it also informs us that they are unable to provide complementary information for analyzing protein design methods. More diverse, complex, and challenging datasets are needed for further investigation.
4. The unsupervised confidence is highly correlated to supervised recovery. This discovery suggests that researchers can rank the quality of designed proteins based on confidence alone, without needing to access the ground truth.
Multi-chain Results To remedy the problem that CATH4.2 and CATH4.3 are highly consistent and do not bring complementary information, we extend the experiment to a multi-chain dataset. The corresponding results are presented in Table 2, showing that:
Table 2: Multi-chain results. The best and suboptimal results are labeled with bold and underlined.
Model length	Confidence ↑	Recovery % ↑
L < 100	100 ≤ L < 500	500 ≤ L < 1000	Full	L < 100	100 ≤ L < 500	500 ≤ L < 1000	Full
PDB	StructGNN	0.49	0.49	0.50	0.49	0.41	0.41	0.42	0.41
GraphTrans	0.48	0.47	0.48	0.48	0.40	0.39	0.40	0.40
GCA	0.45	0.45	0.46	0.45	0.41	0.41	0.42	0.41
GVP	0.51	0.53	0.55	0.54	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	0.60	0.63	0.61	0.54	0.58	0.60	0.58
KWDesign	0.65	0.71	0.74	0.71	0.59	0.66	0.67	0.66