Stratified Cross-Validation

Stratified cross-validation draws the folds so that each one carries roughly the same class proportions as the full dataset. Shuffled k-fold assigns rows at random and lets the class balance land where it falls; stratification constrains it. KFold does not shuffle by default.

On 1,000 samples split 900/100, stratified five-fold gives every fold about 180 of the majority class and 20 of the minority. Plain five-fold gives folds that average the same and vary around it.

How often the variation bites

The textbook warning is that a fold might contain almost no minority examples. On genuinely rare classes that is not a remote possibility.

import numpy as np
from sklearn.model_selection import KFold, StratifiedKFold

rng = np.random.default_rng(2)
n, positives = 500, 12
y = np.zeros(n, dtype=int)
y[rng.choice(n, size=positives, replace=False)] = 1
x = np.zeros((n, 1))

for name, splitter in (("KFold", KFold(5, shuffle=True, random_state=3)),
                       ("StratifiedKFold", StratifiedKFold(5, shuffle=True, random_state=3))):
    counts = [int(y[test].sum()) for _, test in splitter.split(x, y)]
    print(f"{name:16} positives per test fold {counts}  empty folds {counts.count(0)}")

empty = 0
for seed in range(200):
    counts = [int(y[test].sum()) for _, test in KFold(5, shuffle=True, random_state=seed).split(x, y)]
    empty += any(c == 0 for c in counts)
print(f"plain KFold over 200 seeds: at least one fold with no positives in {empty} of them")
# KFold            positives per test fold [2, 5, 2, 1, 2]  empty folds 0
# StratifiedKFold  positives per test fold [2, 2, 2, 3, 3]  empty folds 0
# plain KFold over 200 seeds: at least one fold with no positives in 56 of them

Twelve positives in 500 rows. On this particular shuffle plain k-fold happens to give every fold at least one positive, and stratified spreads them as evenly as twelve allows. Over 200 shuffles, though, plain k-fold produces at least one fold with no positives at all in 56 of them — better than one run in four.

A fold with no positives is not a slightly noisy fold. Recall is \(0/0\), precision is \(0/0\) as soon as the model predicts nothing there, and AUROC is undefined because there is no positive to rank. Depending on the metric and the zero_division convention, that fold contributes a zero, a one, or a nan that silently poisons the mean. The reported cross-validated score then reflects a convention rather than the model.

Stratification reduces missing-class folds when there are enough class members. scikit-learn uses it by default for binary and multiclass classification in cross_val_score. It addresses a practical scoring problem, but does not remove all uncertainty or override time and group constraints.

What it is and is not doing

Stratification controls one source of composition differences across folds. It changes which rows are used to train each model and can therefore change fitted models and their scores. It can also make the fold-score spread look smaller than the variability encountered in new samples. With twelve positives, stratified five-fold still estimates recall from only two or three positives per test fold.

Three cases where it needs care:

• Fewer class members than folds. Three positives cannot be spread across five folds; two folds will have none whatever the splitter does. scikit-learn warns rather than failing — see multiclass stratified CV.
• Grouped rows. Keeping class proportions and keeping a patient’s rows together are separate constraints, and satisfying both needs a splitter that does both.
• Regression. There is no class to stratify on. Binning the target and stratifying on the bin is a common workaround, and it is a modeling choice rather than a neutral default.

Stratifying on a feature is not inherently leakage. It defines a sampling design whose relevance depends on the population and variation the evaluation is meant to represent. Describe what was balanced, preserve time and group boundaries, and do not interpret constrained fold variability as all sources of deployment uncertainty.

Where it extends

The same idea appears under three other names: multiclass stratified CV preserves proportions across more than two classes, stratified shuffle split draws repeated random train/test pairs instead of a partition, and stratified group k-fold adds the constraint that a group’s rows stay together.

References

• scikit-learn API reference, StratifiedKFold
• scikit-learn User Guide, Cross-validation: evaluating estimator performance