GRIMM: Genetic stRatification for Inference in Molecular Modeling

TL;DR

GRIMM introduces a genetics-informed, cluster-based splitting framework to create realistic closed-set and open-set evaluations for enzyme function prediction. By clustering sequences by similarity and enforcing exclusive partitioning, GRIMM minimizes leakage and generates Test-1 (in-distribution) and Test-2 (pseudo-OOD) datasets that better reflect deployment conditions. The approach is demonstrated on Enzyme Commission annotations and shown to improve generalization to novel functions, with a freely available five-fold EC dataset and accompanying code. This framework is broadly applicable to any sequence-based labeling task and provides a rigorous, scalable benchmark for assessing model generalization beyond homology-dominated evaluation.

Abstract

The vast majority of biological sequences encode unknown functions and bear little resemblance to experimentally characterized proteins, limiting both our understanding of biology and our ability to harness functional potential for the bioeconomy. Predicting enzyme function from sequence remains a central challenge in computational biology, complicated by low sequence diversity and imbalanced label support in publicly available datasets. Models trained on these data can overestimate performance and fail to generalize. To address this, we introduce GRIMM (Genetic stRatification for Inference in Molecular Modeling), a benchmark for enzyme function prediction that employs genetic stratification: sequences are clustered by similarity and clusters are assigned exclusively to training, validation, or test sets. This ensures that sequences from the same cluster do not appear in multiple partitions. GRIMM produces multiple test sets: a closed-set test with the same label distribution as training (Test-1) and an open-set test containing novel labels (Test-2), serving as a realistic out-of-distribution proxy for discovering novel enzyme functions. While demonstrated on enzymes, this approach is generalizable to any sequence-based classification task where inputs can be clustered by similarity. By formalizing a splitting strategy often used implicitly, GRIMM provides a unified and reproducible framework for closed- and open-set evaluation. The method is lightweight, requiring only sequence clustering and label annotations, and can be adapted to different similarity thresholds, data scales, and biological tasks. GRIMM enables more realistic evaluation of functional prediction models on both familiar and unseen classes and establishes a benchmark that more faithfully assesses model performance and generalizability.

Figures (1)

• Figure 1: Decision-based flowchart for constructing stratified sequence datasets under joint classification and clustering constraints. In our use case, protein sequences are retrieved from SwissProt and grouped by enzyme class (EC number) and sequence-similarity clusters (UniRef50) to prevent cluster overlap across data partitions. Classes with sufficient cluster support are split into training, validation, and in-distribution test sets, while low-support classes are handled separately as "extras" or singletons ("orphans"). Singleton classes are used to construct an out-of-distribution evaluation set (Test-2), enabling controlled assessment of generalization beyond well-represented EC classes.