cs.LGSep 29, 2026
SaveReLaG: A Scalable Framework Generalizing Random Splits to Data with Latent Relations
Organizations: Department of Computer Science Université Laval Québec, QC, Canada · Department of Biochemistry, Microbiology and Bioinformatics Université Laval Québec, QC, Canada · Department of Molecular Medicine Université Laval Québec, QC, Canada
Abstract
Random splitting can yield non-independent train--test subsets when a dataset contains related samples, as is common in certain applications such as biochemical studies. This leads to overly optimistic generalization estimates. Here, we introduce ReLaG, a modality-agnostic framework that models sample relatedness through a hierarchical latent-variable process and infers groups of related samples using proximity graphs and community detection to produce independent train--test subsets. Across molecular and protein datasets, ReLaG matches existing relation-aware methods while scaling substantially better, enabling splits at previously impractical dataset sizes. We further introduce a label-free procedure that adapts the splitting resolution to production data, aligning evaluation with the intended deployment setting. ReLaG's inferred groups provide a cheap estimate of effective dataset size, enabling diversity-aware dataset scaling. ReLaG is open source and can be installed with pip install relag.
Figures & tables
Figure 1: Illustration of two generative processes that produce datasets containing related samples.
Figure 2: Two-level latent-variable model. Observed samples are generated conditionally on group-specific latent variables , which are drawn i.i.d. from a fixed population-level distribution . The objective is to infer this group structure from the observed samples.
Figure 3: ReLaG algorithm overview. (A) The algorithm requires a distance function that correlates with the probability that two samples are unrelated. (B) A proximity graph is built using this distance function. (C) Communities are identified, then graph edges that are unlikely to represent relationships are removed. Each community is an estimate of a group, i.e., a set of samples conditioned on the same latent variable . (D) The dataset is partitioned along the community boundaries.
Figure 4: Threshold selection intuition. A wrong threshold makes train communities closer or further away from each other than they are from production communities. The right threshold makes production communities integrate nicely in the community distribution.
Figure 5: Relation-aware splits yield evaluations that match production performance. In each of 100 repeats, we sample independently two synthetic peptide datasets (15,000 sequences, 1,000 families each) under the latent-variable model: an annotated dataset to split and a production dataset.
Figure 6: Performances on real data. Downstream MLP test performance of dataset-splitting strategies across eight datasets (six molecular and two protein datasets). Results are averaged over 10 split seeds, with error bars denoting standard deviations. We report AUROC for classification tasks and Pearson correlation coefficient for regression tasks.
Figure 7: Threshold selection characterizes the relatedness regime. Top: single-linkage distances from pure-training and pure-production communities to training communities (15th–85th percentile shaded). Bottom: KS statistic.
Figure 8: Scaling on PeptideAtlas subsets. Runtime as a function of dataset size for ReLaG and Hestia, the two fastest compared methods. Runs were limited to one day; values above the dashed line are extrapolated. Log–log fits indicate approximately linear scaling for ReLaG and quadratic scaling for Hestia ( for both).
Figure 9: Community structure captures an aspect of training-set diversity that is not reflected by the raw sample count and is associated with generalization performance.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | Modality | Task | Metric | Encoder | |
|---|---|---|---|---|---|
| cyp2c19_veith | molecule | 12,665 | classification | AUROC | ChemBERTa |
| ames | molecule | 7,278 | classification | AUROC | ChemBERTa |
| sr_are | molecule | 5,832 | classification | AUROC | ChemBERTa |
| lipophilicity | molecule | 4,200 | regression | PCC | ChemBERTa |
| pgp_broccatelli | molecule | 1,218 | classification | AUROC | ChemBERTa |
| caco2_wang | molecule | 910 | regression | PCC | ChemBERTa |
Table 1: Datasets used in the experiments. is the number of samples after filtering.
Figure 10: Scaling of ReLaG on molecular modality. Both axes are logarithmic. Exponents in the legend are fitted on the log–log space.
Figure 11: Threshold selection on large-scale and synthetic data. Upper panels: single-linkage distances from pure-training communities (train train) and from pure-production communities (train production) to training-containing communities, with shading denoting the 15th–85th percentile range. Lower panels: the corresponding Kolmogorov–Smirnov statistic. Dashed lines mark the selected threshold .
| Percentile | train train | train prod | prod train |
|---|---|---|---|
| 1 | 0.127 | 0.000 | |
| 5 | 0.500 | 0.516 | |
| 10 | 0.600 | 0.600 | |
| 25 | 0.650 | 0.650 | |
| 50 | 0.700 | 0.700 | |
| 75 | 0.750 | 0.750 |
Table 2: Nearest-train-neighbor distance on UniRef50: percentiles of the distance from a training sequence to its closest other training sequence (train train), and from a production benchmark sequence to its closest training sequence (train prod).
Figure 12: ProtSpaM distance against global alignment distance, on 39,380 pairs built from real UniRef50 seeds mutated at a sweep of rates (BLOSUM62 substitutions and indels) so the alignment axis is covered end to end. The two are monotonically related (Pearson , Spearman ).
| Dataset | Metric | Singleton | Community | |||
|---|---|---|---|---|---|---|
| caco2_wang | PCC | 0.6 | ||||
| lipophilicity | PCC | 0.4 | ||||
| sr_are | AUROC | 0.4 | ||||
| cyp2c19_veith | AUROC | 0.5 | ||||
| dbaasp | PCC | 0.3 | ||||
| ames | AUROC | 0.7 |
Table 3: Singleton information content. Comparison of models trained on singleton samples and equally sized samples from non-singleton communities, evaluated on the same validation set. is the performance difference (singleton community), with indicating greater information content in singletons. Mean over 5 splits 6 repeats. -values are from two-sided -tests. Rows are grouped by higher, lower, or indistinguishable singleton performance.
Figure 13: Family leakage on the synthetic dataset. Fraction of training samples whose family also appears in the test subset. Random: 93.8%, ReLaG: 0.62%, DataSAIL: 21.01%, Hestia: 0.00%, Oracle: 0.00%
Figure 14: Train–test distance across data-splitting methods. ECDFs of the minimum distance from each test sample to its nearest training sample, pooled within modality: (a) six small-molecule datasets and (b) two protein-sequence datasets. Distance is similarity, using Tanimoto similarity for molecules and global sequence identity ( Needleman & Wunsch, 1970 ) for proteins. Dashed lines indicate the splitting thresholds ( and , respectively); the legend reports the fraction of test samples closer to any train samples than the corresponding threshold.
Figure 15: Effect on test performance of dropping random samples or full communities on different datasets. corresponds to the number of communities remaining in the dataset.