Stratified Group K-Fold
Stratified Group K-Fold splits data under two constraints at once: every sample sharing a group — a patient, a user, a session, a document — lands in exactly one fold, and each fold holds roughly the class proportions of the whole set.
Either constraint alone is a standard tool. Stratified k-fold balances classes and ignores groups. Group k-fold keeps groups intact and ignores classes. The combination exists because datasets with repeated measurements are often also imbalanced, and dropping either constraint breaks the evaluation in a different way.
What a split group actually costs
The usual warning is that splitting a group across folds “leaks information.” Here is the size of it, on a problem constructed so that the honest answer is known.
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, StratifiedGroupKFold, cross_val_score
rng = np.random.default_rng(3)
n_groups, per_group = 120, 8
groups = np.repeat(np.arange(n_groups), per_group)
label = rng.integers(0, 2, size=n_groups)[groups]
signature = rng.normal(size=(n_groups, 6))[groups]
x = signature + rng.normal(scale=0.25, size=signature.shape)
"""Each group has its own fingerprint and one label. Nothing generalizes across groups:
the label is assigned at random, so an honest score is 0.5."""
model = RandomForestClassifier(n_estimators=100, random_state=0)
ignores = cross_val_score(model, x, label, cv=StratifiedKFold(5, shuffle=True, random_state=0)).mean()
respects = cross_val_score(model, x, label, groups=groups,
cv=StratifiedGroupKFold(5, shuffle=True, random_state=0)).mean()
print(f"StratifiedKFold (groups split across folds) {ignores:.4f}")
print(f"StratifiedGroupKFold (groups kept together) {respects:.4f}")
print(f"the truth 0.5000")
rates = [round(float(label[test].mean()), 3) for _, test in
StratifiedGroupKFold(5, shuffle=True, random_state=0).split(x, label, groups)]
print(f"positive rate per fold under grouping {rates} whole set {label.mean():.3f}")
# StratifiedKFold (groups split across folds) 0.9198
# StratifiedGroupKFold (groups kept together) 0.5323
# the truth 0.5000
# positive rate per fold under grouping [0.5, 0.5, 0.5, 0.458, 0.5] whole set 0.492
There is nothing to learn in this data. Group labels are coin flips and the features carry only a per-group fingerprint plus noise, so the expected accuracy of any rule fitted on one set of groups is 0.5 on new, independent groups. Stratified k-fold reports 0.9198 anyway. It is not wrong about what it measured: seven of a group’s eight rows sat in training, so at validation time the model was recognizing a fingerprint it had already been told the answer for. That is memorization scored as generalization — leakage in its most ordinary form.
Keeping groups whole removes the shortcut and the score falls to 0.5323, close to the expected 0.5. The gap that remains is within what chance variation can produce: a single 5-fold estimate of a chance-level quantity scatters around 0.5, so 0.5323 is best compared with that variation rather than read as a small real effect.
The practical reading of 0.9198 versus 0.5323 is that group leakage does not shade a result — it can manufacture one. A model that looks deployable under the wrong split can be worthless under the right one.
The stratification is approximate
The second constraint is weaker than it sounds. Group membership is indivisible, so the fold sizes and their class composition are whatever the available groups allow. scikit-learn’s user guide is explicit that the method tries to preserve the class distribution while keeping each group in one split, and that the greedy assignment “is suboptimal in a sense that it might produce imbalanced splits even if perfect stratification is possible.”
The last line of the run shows it: against an overall positive rate of 0.492, four folds land on 0.5 and one on 0.458. That case is the hard one by construction — every group here carries a single label, so a fold’s composition is decided entirely by which whole groups it receives. Where groups are internally mixed, the approximation is much closer.
So the guarantee to rely on is the group one. Treat the class balance as a best effort, and check it rather than assume it — the per-fold rates take one line, as above.
Choosing the grouping column
The method is only as good as the groups array, and that array is a modeling decision, not a lookup. The question it answers is: at deployment, what will be genuinely new?
- New patients, not new scans. If the system will read images from people it has never seen, the group is the patient. Grouping by scan leaves the same patient on both sides.
- Watch for grouping under another name. Near-duplicate records, augmented copies of one image, several rows from one survey response and translations of one document all form groups that no ID column marks.
- Nest the strictest unit. When rows belong to a session inside a user inside a clinic, group by whichever level will be unseen in production — usually the outermost.
Getting this wrong is not caught by any validator. The split will run, the folds will look balanced and the score will be high.
When a simpler split is the better one
Stratified Group K-Fold is the right choice when groups exist and the class balance is skewed enough that plain group splitting would produce folds with too few positives to score.
With no relevant grouping, stratified k-fold may suffice. When grouped class distributions are already similar, plain group k-fold avoids an unnecessary stratification constraint. Time order can add another requirement rather than replace grouping: evaluating future observations from unseen patients may require both chronological separation and patient isolation. Choose the split from what is new at deployment, as discussed under cross-validation.
Combining time-based splits with group constraints requires a cross-validation design that matches both the prediction time and the intended population.
References
- scikit-learn User Guide, Cross-validation: evaluating estimator performance
- scikit-learn API reference, StratifiedGroupKFold
Discover more from Insightful Data Lab
Subscribe to get the latest posts sent to your email.
