GRIMM: Genetic stRatification for Inference in Molecular Modeling

TL;DR

GRIMM introduces a clustering-based benchmark for enzyme function prediction that improves evaluation of model generalization to novel and known enzyme classes, addressing dataset biases and out-of-distribution challenges.

Contribution

It formalizes a genetic stratification strategy for creating realistic training, validation, and test splits, including open-set scenarios, applicable to sequence-based classification tasks.

Findings

Provides a reproducible framework for sequence clustering and data splitting.

Enables evaluation of models on both in-distribution and out-of-distribution enzyme functions.

Applicable to various biological sequence classification problems.

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.