Stratified K-Fold

Stratified k-fold is k-fold cross-validation's classification-specific variant — it guarantees every fold preserves roughly the same class proportions as the full dataset, exactly the discipline stratify=y already applies to a single train-test split.

Why Plain K-Fold Can Fail on Imbalanced Data

With plain K-Fold's random splitting, a rare class (say, 5% of the data) could easily end up with very few — or occasionally zero — examples in some folds' test sets purely by chance. A fold's precision/recall becomes meaningless (or undefined) if it contains no positive examples at all, and averaging across folds like this produces an unreliable, misleading final score.

Worked Example

A dataset with 100 samples: 90 class 0, 10 class 1 (a 90/10 imbalance). With \(k=5\):

Python Implementation

from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.linear_model import LogisticRegression

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

model = LogisticRegression(max_iter=1000)
scores = cross_val_score(model, X_train, y_train, cv=skf, scoring="f1")
print(scores.mean(), scores.std())

# Verify class balance is preserved per fold
import numpy as np
for fold_num, (train_idx, test_idx) in enumerate(skf.split(X_train, y_train)):
    test_labels = y_train.iloc[test_idx] if hasattr(y_train, "iloc") else y_train[test_idx]
    print(f"Fold {fold_num+1} class distribution:", np.bincount(test_labels))

Note: in practice, scikit-learn's cross_val_score and GridSearchCV automatically use stratified splitting by default for classification tasks — but it's important to know this is happening (and why), not just trust it silently.

Practical Use Cases

• Any classification cross-validation with meaningfully imbalanced classes
• The safe default choice for classification cross-validation generally — there's rarely a downside to stratifying, even on balanced data

Common Mistakes

• Using plain K-Fold on imbalanced classification data, risking folds with too few (or zero) minority-class examples.
• Applying Stratified K-Fold to a regression problem — stratification by class proportion is a classification-specific concept; regression has no discrete classes to preserve proportions of.

Interview Relevance

Q: "Why is stratified k-fold important for an imbalanced classification dataset?" Plain random folds can, by chance, produce test folds with very few or zero examples of the minority class — making per-fold precision/recall unreliable or undefined; stratified k-fold guarantees every fold mirrors the overall class distribution, producing consistent, trustworthy fold-level metrics.

Practice Question

A dataset has a target with 3 classes in proportions 70%/20%/10%. Explain what stratified k-fold with k=5 guarantees about each fold's composition.

Related ML Notes

• K-Fold Cross-Validation
• Imbalanced Data
• Train-Test Split in Python