Prevent Data Leakage in Protein Datasets with Grouped Splits

An embedding model can look useful when its test set contains close relatives of its training sequences. Before comparing ESM-2 embeddings with an amino acid composition baseline, decide what the model should generalize to: another measurement of a familiar protein, a new variant of the same parent, or a different family.

This tutorial creates a reproducible train/validation/test assignment that keeps supplied groups together. It follows the FASTA audit and prepares a shared evaluation split for the composition baseline and ESM-2 workflow.

Define what must stay together

Splitting rows randomly answers a different question from holding out entire families. The TAPE paper makes biologically relevant generalization part of its task-specific split design. For your own dataset, write the intended use case before choosing a grouping rule.

These are design choices, not interchangeable definitions of independence. For a sequence-only predictor, repeated measurements of the same sequence share the same input; keep them together unless your evaluation explicitly models another changing input. Do not silently average conflicting labels.

If two constraints matter, a tuple such as (sequence_cluster, donor) may fail: the same donor can occur in several clusters and therefore several tuples. One option is to join records connected by either constraint and use each connected component as a group. This can produce a very large group. Report that limitation rather than quietly breaking it apart.

Make the assignment explicit

The example uses ten invented groups containing 1 through 10 rows: 55 rows in total. It contains no real sequences, labels, or model performance results. Its sequence keys are placeholders used to demonstrate an overlap check.

Download split_demo.py. In a virtual environment, install the tested library version and run:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "scikit-learn==1.9.1"
python split_demo.py > split.csv

The example was tested with Python 3.14.7 and scikit-learn 1.9.1. The shell redirects the assignment table into split.csv; choose a new filename if you want to preserve an earlier file. The following summary appears on standard error:

train: rows=32 groups=6
valid: rows=12 groups=2
test: rows=11 groups=2
PASS: rows, groups, and exact-sequence keys checked

The first split reserves two groups for testing. The second reserves two of the eight remaining groups for validation. The result is 60/20/20 by group count, while row counts differ. In GroupShuffleSplit, fractional sizes apply to groups and the test-group count rounds upward. Small datasets may therefore depart from the intended proportions.

Here is the complete downloadable script:

"""Synthetic grouped split demo; no biological measurements or model scores."""
import csv
import sys
from itertools import combinations

import numpy as np
from sklearn.model_selection import GroupShuffleSplit


def grouped_split(groups):
    groups = np.asarray(groups)
    if groups.ndim != 1 or len(set(groups.tolist())) < 5:
        raise ValueError("Supply a one-dimensional array with at least five groups")
    rows = np.arange(len(groups))
    first = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=17)
    dev, test = next(first.split(rows, groups=groups))
    second = GroupShuffleSplit(n_splits=1, test_size=0.25, random_state=29)
    train_local, valid_local = next(second.split(dev, groups=groups[dev]))
    # The second split returns positions within dev, not original row numbers.
    return {"train": dev[train_local], "valid": dev[valid_local], "test": test}


def verify(parts, groups, sequence_keys):
    rows = [int(i) for indices in parts.values() for i in indices]
    if sorted(rows) != list(range(len(groups))):
        raise ValueError("Every row must occur exactly once")
    if len(sequence_keys) != len(groups):
        raise ValueError("Sequence keys and groups must have equal length")
    for left, right in combinations(parts, 2):
        for values, label in [(groups, "group"), (sequence_keys, "exact sequence")]:
            overlap = {values[i] for i in parts[left]} & {values[i] for i in parts[right]}
            if overlap:
                raise ValueError(f"{left}/{right}: overlapping {label}")


def main():
    # Invented groups of unequal sizes; keys stand in for canonical sequences.
    groups = np.repeat(np.arange(10), np.arange(1, 11))
    keys = [f"synthetic-sequence-{i}" for i in range(len(groups))]
    parts = grouped_split(groups)
    verify(parts, groups, keys)
    writer = csv.writer(sys.stdout, lineterminator="\n")
    writer.writerow(["sample_id", "group_id", "split"])
    for name, indices in parts.items():
        group_count = len(set(groups[indices]))
        print(f"{name}: rows={len(indices)} groups={group_count}", file=sys.stderr)
        for i in indices:
            writer.writerow([f"sample_{i:03d}", int(groups[i]), name])
    print("PASS: rows, groups, and exact-sequence keys checked", file=sys.stderr)


if __name__ == "__main__":
    main()

The second splitter returns positions within dev. Mapping those positions back through dev is essential; treating them as original row numbers can assign the wrong samples.

Replace the invented groups with real metadata

Keep one stable sample identifier per row. Supply a one-dimensional array of nonmissing group IDs with a consistent type, then call grouped_split(groups). The demonstration requires at least five groups to make its two-stage split practical; five groups do not make an evaluation statistically reliable.

Replace keys with actual canonical protein sequences, or their recorded cryptographic digests, in the same row order. Canonicalization must follow the input policy you documented during the FASTA audit. Run verify(parts, groups, keys) before fitting anything. It rejects missing or repeated row assignments, shared groups, and exact sequence keys crossing subsets. A passing check says nothing about near-identical sequences that have different keys.

For clustering, record the program, version, sequence-identity definition, alignment coverage requirement, and clustering mode. A cluster label is the output of that procedure. If your claim requires every test sequence to be below a stated similarity threshold to training data, perform a separate cross-subset comparison using the same defined identity and coverage policy. Do not infer that guarantee merely from different cluster labels.

Decide this policy without looking at model scores. Save the resulting sample_id,group_id,split table with the dataset version, file checksum, seeds, and grouping parameters. A seed alone cannot reconstruct a split after the input ordering or group annotations change.

Keep fitted preprocessing inside training

Use training data to fit the predictor and any learned scaling, imputation, PCA, or feature selection. Use validation data for model choices, and reserve test evaluation until those choices are fixed. The scikit-learn leakage guide explains why fitting transformations on the full dataset contaminates evaluation; a pipeline helps keep those operations within the training partition.

A fixed amino acid composition calculation does not learn dataset-wide statistics. Neither does independently encoding each sequence with an unchanged pretrained model in evaluation mode. However, fitting PCA on all resulting embeddings or fine-tuning the encoder on held-out labels changes that situation. Public pretraining may also have included evaluation sequences; a clean downstream split cannot establish pretraining independence.

For cross-validation, keep groups together inside the development set as well. The cross-validation guide describes grouped methods. GroupShuffleSplit does not balance class labels, so inspect class counts and target ranges in each subset. If the intended task includes extrapolation to future experiments, random group assignment may also be the wrong design; a time-based holdout needs its own policy.

Fix failed checks before fitting

Report what the split actually tests

For the next model comparison, publish the split manifest alongside row counts, group counts, label distributions, and the overlap checks performed. Compare the composition baseline and protein embeddings on the same assignments and state how each model was selected.

A lower score after grouping is not automatically a worse model. It may reflect a harder generalization question. Name that question, preserve the test set, and make the assignment reproducible before treating a score difference as progress.

The cover is conceptual artwork, not an experimental protein structure. Documentation checked September 24, 2026.