Disconnecting The Dots: Creating Leakage-Free Protein Datasets by Removal of Densely Connected Data Points

We thank the reviewer for their review. We address concerns/questions line-by-line below.

Re: Level of data exclusion
We appreciate the reviewer's concern about the level of data exclusion. While the reduction is significant, the final dataset size remains suitable for practical applications, and the number of points removed is much smaller than existing methods like GraphPart, particularly for low-to-intermediate difficulty levels. Additionally, our approach offers a robust way to assess model generalization, and training on the full dataset is always an option afterwards.

We also agree that a discussion of the data loss trade-off could be beneficial. In preliminary experiments, we monitored the proportion of the largest cluster in the data as a function of removed top connectors. While removing only a subset of the top connectors was sufficient to obtain well-distributed clusters, we also found that this proportion declined only after most connectors had been removed. This means that, while a stopping criterion could be introduced (such as monitoring the proportion of the largest cluster in the data), the number of additional data points retained would be minimal.

Re: Reliance on community detection
The reviewer is right to note that community detection can sometimes create arbitrary groupings. However, this should mostly happen with the Louvain version, which can sometimes create communities that are not well-defined or have unbalanced sizes. We use the Leiden version of the community detection algorithm, which aims to improve on the Louvain version by refining the community structure. It does so by ensuring that the communities it finds are more stable and well-defined, thus avoiding situations where communities overlap too much or are too fragmented.

Re: Additional biological tasks
We agree that the manuscript would benefit from including more tasks. We are currently working on adding two structure task predictions (contact and fold prediction, from the PEER benchmark), one localization task (also from PEER), and two retrieval, homology-based task (from the DGEB benchmark). We don’t have the complete data yet but will include them in the final results.

Re: Definition of similarity metrics
As we discuss in our response to reviewer TsLk , we agree that the similarity metrics should be defined in the main manuscript. We will add them in the final version. We used the following definitions:

• sequence identity: we used MMseqs fractional identity (“fid”) score, which measures the proportion of identical amino acid matches between two aligned sequences, normalized by the alignment length.
• TM score: a measure of global structural similarity widely used in protein structure comparison (e.g in AlphaFold). We use the definition introduced in the original paper (Zhang, Y., & Skolnick, J. (2004)) and its fast implementation by Foldseek.

Re: Complexity analysis
As we mention in our responses to reviewers MFnt and QpPD, we provide an estimate of the computational complexity of our method in Section E of the Appendix. We estimate the complexity to be quadratic in the number of samples, although we suspect the actual cost may be lower due to the sparse implementation of our method. We plan to provide a more refined analysis in the final version of the manuscript.

Re: Clarifying concepts
We acknowledge that the different terms used in our manuscript may be confusing, as they all refer to various notions of clustering. To clarify, we will standardize the terminology in the camera-ready version and provide clear definitions:

• Clusters: Groups of nodes in a graph with high internal connectivity relative to the rest of the network, typically defined by fixed criteria for local density or similarity. Cluster-based methods, such as k-means or hierarchical clustering, are used when the groups are expected to be well-separated and internally compact.
• Connected Component: A maximal subset of nodes where each node is reachable from any other node within the subset. Connected components are clusters that are strictly separated.
• Community: a subset of nodes more densely connected to each other than to nodes outside the subset, emphasizing modularity. Community-based methods are typically used when analyzing large-scale networks, such as social networks or protein interaction networks. They work well with graphs, focusing on connectivity patterns. Detecting communities allows us to identify groups of highly connected proteins within the large cluster found by hierarchical clustering or connected components.

In our manuscript, the resulting connected components from Algorithm 1 are the clusters we use to construct the new splits. We will revise the manuscript to use the term "cluster" consistently going forward.

Additionally, Algorithm 1 does not simply return the communities C with the top connector nodes v removed. We have observed that returning the connected components leads to slightly more clusters compared to returning communities. This occurs when the removal of v splits a community into two separate components. While this does not affect the total number of nodes removed, it results in a greater number of clusters to sample from when creating the splits, thereby enhancing the diversity of the clusters.

Re: Complexity analysis
As we mention in our response to reviewer MFnt, we provide an estimate of the computational complexity of our method in Section E of the Appendix. We estimate the complexity to be quadratic in the number of samples, although we suspect the actual cost may be lower due to the sparse implementation of our method. We plan to provide a more refined analysis in the final version of the manuscript.

Re: Algorithm 2
We use Algorithm 2 to evaluate model performance at different difficulty levels. These often correspond to distinct biological problems, such as de novo design (high difficulty, requiring generalization) and protein optimization (low difficulty, requiring less generalization). However, we recognize that Algorithm 2 adds complexity to the main manuscript. It implements an existing method for distributing clusters across training, validation, and test sets to create multiple difficulty levels. This method, along with the ProteinNet dataset, was introduced five years ago and has been used in several protein benchmarks. We present a sparse, fast implementation to apply it to new datasets, including our DeepFRI splits. The core contribution of our work, however, is the removal of leakage through community detection, as described in Algorithm 1. To clarify this, we will move the section and figure discussing Algorithm 2 to the Appendix in the final version of the manuscript.

Re: Spearman correlation
We agree that Spearman correlations are not appropriate for that type of data, and will remove it in the final version.

We would like to thank the reviewer for their thoughtful comments and are happy that they found our work to be well written and of good quality. We discuss their concerns below:

Re: Scope
Our focus here is more an experimental investigation. For space reasons we chose not to overly formalize our description of the method.

Re: Goals when creating training and test sets
As we discuss in response to reviewer TsLk, proteins are not independent data points. They arise from evolutionary processes, often sharing high sequence similarity as a result. Random splits fail to account for this, frequently placing highly similar proteins across training, validation, and test sets. This overlap leads to performance metrics that do not reflect true generalization capabilities. This is a key challenge in biology because experimental labels are expensive and scarce, meaning researchers rely heavily on model predictions to annotate new proteins. Furthermore, well-studied proteins are the ones that are experimentally accessible – that is, only a small subset of existing proteins. On the other hand, proteins evolutionarily distant from the known protein universe are discovered every day.

As a result, annotated datasets typically only represent a fraction of the functional protein space. This motivates the need for benchmarking methods that enable generalization assessment to distant protein families. Two notable examples of label scarcity include:

• the human microbiome – many gut bacteria cannot be cultivated, and therefore cannot be labeled functionally;
• protein structures – historically, structure have been determined using X-ray crystallography, which works well for soluble proteins but not for membrane proteins. For a long time, the set of structures available for training protein structure prediction models was thus limited to a small subset of protein families, making it very hard for model to generalize to new structures. Significant improvements on the modeling side were made possible by the construction of structurally diverse datasets by CASP.

Re: Measuring vs achieving generalization
Yes, there is this tension in general. But even in instances where the desire is to achieve full generalization by training on all available data, it is still useful to understand the behavior of the model by using smaller subsets of the data to assess its generalization capabilities. Furthermore, in the biological context, data is often scarce. It is infer trend lines of how a give model improves as a function of the dataset size. Smaller subsets can be useful to infer how a model would improve as a function of the dataset size.

Re: Three-way splits
Our approach already uses a three-way splitting, with a validation set for hyperparameter tuning and a test set for model testing. Both the validation and test sets were constructed using our leakage-free clustering approach.

Re: Classification with a reject option
We appreciate the pedagogical value of this type of discussion but think they are out of the scope of this paper.

Re: Citations
We have now included references to CASP in the manuscript.

We thank the reviewer for their review. We are glad they believe our approach of detecting communities “neat and well-motivated”. We have updated the manuscript to fix the formatting errors. Below, we address the reviewer’s questions and concerns individually.

Re: Lack of evidence & relevance
In response to the reviewer’s comment about a lack of direct evidence, we acknowledge that the results section only indirectly demonstrated the benefit of our approach. To address this, we have made the following changes to more clearly show the advantages of our method:

• In Section A of the Appendix, we quantify leakage in the DeepFRI splits.
• In Section C of the Appendix, we compare model performance on the DeepFRI and our splits. We find that results are more sensitive to similarity thresholds with our splits, which is desirable for assessing model generalizability. We plan to incorporate both sections into the main body of the manuscript in the final version.

Re: Random and temporal splits
Proteins are not independent data points. They arise from evolutionary processes, often sharing high sequence similarity as a result. Random splits fail to account for this, frequently placing highly similar proteins across training, validation, and test sets. This overlap leads to inflated performance in expectation, as models can rely on memorizing patterns rather than generalizing to novel proteins. For tasks like protein function prediction, where researchers need reliable predictions for newly discovered, evolutionarily distant proteins, this is a significant limitation. Reported performance from random splits does not reflect the real-world generalization ability of models, making it harder for researchers to choose the most effective tools for annotating novel proteins.

Temporal splits face a similar issue: they also do not account for protein similarity, and can result in biologically close proteins being partitioned across dataset splits. While temporal splits might seem meaningful for applications where time order is relevant, they fail to address the key challenge in biology: generalizing to novel proteins that are highly dissimilar from known proteins. This is especially critical because experimental labels are expensive and scarce, meaning researchers rely heavily on model predictions to annotate new proteins. Evaluating models with similarity-based splits ensures that reported metrics realistically reflect their utility in these challenging scenarios.

Re: Similarity metric
As we discussed in our response to reviewer MFnt, we report results for two standard similarity metrics for proteins:

• sequence identity: we used MMseqs fractional identity (“fid”) score, which measures the proportion of identical amino acid matches between two aligned sequences, normalized by the alignment length.
• TM score: a measure of global structural similarity widely used in protein structure comparison (e.g in AlphaFold). We use the definition introduced in the original paper (Zhang, Y., & Skolnick, J. (2004)) and its fast implementation by Foldseek.

Re: Splitting methods and leakage
The common approach for creating similarity-based splits involves clustering proteins at a set similarity threshold, then distributing these clusters into training, validation, and test sets. CD-HIT and MMseqs are widely used for this purpose. For example, CD-HIT was used to create the DeepFRI splits, while MMseqs was employed to generate the TAPE, ProteinFlow, and ProteinShake benchmarks. MMseqs also helped create the dataset for the inverse folding model ProteinMPNN. However, neither method ensures strict cluster separation by default, allowing for pairs of proteins with similarity higher than the specified threshold across clusters. As shown in Section A of the Appendix, this leads to significant leakage in the DeepFRI splits. While MMseqs offers a “connected components” mode to ensure strict separation, it produces a single large cluster, making it difficult to create balanced and diverse splits.

Re: Concept of leakage
We agree that a definition of leakage is necessary. Following previous work, we define leakage as the presence of proteins in the validation (or test) set that are more similar to the training set than they should be, given the specified similarity threshold. In Section A of the Appendix, where we discuss leakage in the DeepFRI splits, we quantify it by calculating the proportion of such proteins in the validation (or test) set. We plan to move this section to the main body of the manuscript in its final version.

We thank the reviewer for their constructive feedback. We are pleased that the reviewer finds our data splitting approach to be “justified and sound” and our experimental evaluation extensive. We attempt to address remaining concerns/questions line-by-line below.

Re: Venue Fit
While we acknowledge that our manuscript is most valuable to practitioners of machine learning in bioinformatics, we believe the most recent revision of our manuscript is a good fit for ICLR for three reasons:

1. The bioinformatic machine learning community is large and growing, especially in the wake of AlphaFold2. For instance, at ICLR 2024, there were 19 papers with “protein” in the title, 23 papers with “molecule” or “molecular”, and another 2 with “biological”. One of these introduces a benchmark for DNA language models on 6 biological tasks (“BEND: Benchmarking DNA Language Models on biologically meaningful tasks”). Previous protein-related ICLR papers have been very successful with hundreds of citations (e.g. "Transformer protein language models are unsupervised structure learners" in 2021, “Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking" in 2022, or “Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem" in 2023 ). Furthermore, the official ICLR call for papers includes “applications in biology”. In general, as ML/AI in science has grown, biology and in particular molecular biology have been one of the key application areas.
2. The problem of similarity between the training, validation, and test set is crucial to estimating model generalizability in biology. The approach and the datasets we introduce should be of interest to many machine learning practitioners in biology, precisely because biological problems introduce new challenges that are not present to the same degree in vision or natural language
3. Although we have not explored applications in other fields of machine learning, the approach could find interesting applications elsewhere, particularly in areas where there is a well-justified similarity metric, for example in chemistry and materials science.

Re: Scaling
With regards to whether our community-based approach scales well to large datasets, we provide an estimate of the computational complexity of our method in Section E of the Appendix. We estimate the complexity to be quadratic in the number of samples, although we suspect the actual cost may be lower due to the sparse implementation of our method. We plan to provide a more refined analysis in the final version of the manuscript. Additionally, our choice of dataset was driven by the task at hand and the availability of labels, rather than dataset size. The exemplar dataset we selected happened to be relatively small, which is a common scenario in biology. However, due to the sparse implementation of our method, it should also be applicable to larger datasets.

Re: Qualitative results
We recognize that the original manuscript focused primarily on the comparative performance of models on the new sequence- and structure-based splits. In Section C of the Appendix, we now compare model performance on the DeepFRI and our splits. We observe that results are more sensitive to similarity thresholds with our splits, which is desirable for assessing model generalizability. We plan to move this section to the main body of the manuscript in the final version.

Re: Choice of similarity metric
You are right that our rewording might have been misleading. We meant that the choice of different similarity metrics is independent of our splitting strategy, not that they would have no effect on the outcome - they certainly would. Furthermore, we already report results for two commonly used protein similarity metrics: the sequence identity (for sequences), and the TM score (for structures). Figure 4 in the main paper presents results on the new splits using both similarity metrics. Additionally, Section B of the Appendix displays the cluster distribution for both metrics, both before and after applying our method.