Stratified Sampling for ML
Why plain random sampling can mislead you
A simple random split of data into train/validation/test sets is unbiased in expectation, but any single draw can, by chance, produce a split where the class distribution (or another important subgroup distribution) differs meaningfully between sets — especially with imbalanced classes or small datasets. A dataset that's 5% positive overall might, by unlucky random draw, end up 3% positive in train and 8% positive in test, which can distort both training (the model sees a slightly different base rate than deployment) and evaluation (performance metrics reflect the quirks of that particular split, not the underlying task).
What stratified sampling does
Stratified sampling divides the dataset into subgroups (strata) based on a key variable — most commonly the target class, but it can be any categorical attribute worth preserving proportionally — and samples from each stratum independently, so that the proportion of each stratum in every resulting subset matches its proportion in the full dataset. For classification, this guarantees train, validation, and test sets all have (approximately) the same class balance as the original data, removing one source of noisy variance from your evaluation.
from sklearn.model_selection import train_test_split, StratifiedKFold
import numpy as np
# Simple stratified train/test split
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, stratify=y, random_state=42
)
# Stratified K-fold cross-validation: every fold preserves class proportions
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
for train_idx, val_idx in skf.split(X, y):
X_tr, X_val = X[train_idx], X[val_idx]
y_tr, y_val = y[train_idx], y[val_idx]
# train and evaluate within this fold
class_balance = np.bincount(y_val) / len(y_val)
print("validation fold class balance:", class_balance)Beyond class labels: stratifying on other variables
Stratification isn't limited to the target label. In practice, it's common to stratify on variables that are known to correlate strongly with outcome but aren't the target itself — for example, stratifying a medical dataset by both diagnosis *and* hospital site, or stratifying a marketing dataset by both conversion label and customer acquisition channel. Combined (multi-variable) stratification keeps proportions balanced across the cross-product of these variables, though with enough strata and too little data, some strata can become too small to split sensibly — a practical limit on how many variables you can stratify on simultaneously.
Stratification matters more as datasets get smaller or more imbalanced
| Sampling method | Preserves class balance? | Best suited for |
|---|---|---|
| Simple random split | Approximately, with enough data | Large, balanced datasets |
| Stratified split | Exactly (within rounding) | Imbalanced classes, small-to-medium datasets |
| Stratified K-fold CV | Exactly, per fold | Model selection under imbalance |
| Group-aware split (GroupKFold) | N/A — preserves group integrity instead | Data with non-independent groups (e.g. same patient/user across rows) |
Stratification doesn't fix leakage from grouped data
What's Next?
Stratification is one ingredient in a sound train/validation/test strategy; the broader question of how to size and structure those splits — and what can go wrong even with a technically correct split — is worth examining as its own topic.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.