scikit-learn
machine learning
StratifiedKFold
KFold
cross-validation

StratifiedKFold vs KFold in scikit-learn

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

KFold and StratifiedKFold both split data into k folds, but they do not preserve the target distribution in the same way. The difference matters most in classification, especially when classes are imbalanced, because a bad fold split can make model evaluation far less reliable than it looks.

What plain KFold does

KFold divides the dataset into k folds and rotates which fold is used for validation:

python
1from sklearn.model_selection import KFold
2import numpy as np
3
4X = np.arange(20).reshape(10, 2)
5y = np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1])
6
7kf = KFold(n_splits=5, shuffle=True, random_state=42)
8
9for train_idx, test_idx in kf.split(X):
10    print("test labels:", y[test_idx])

KFold does not care about class balance. It only splits rows. That is fine for regression and for classification when classes are already balanced and randomly distributed. It becomes risky when the target distribution is skewed.

What StratifiedKFold adds

StratifiedKFold tries to keep the class proportions in each fold similar to the full dataset:

python
1from sklearn.model_selection import StratifiedKFold
2import numpy as np
3
4X = np.arange(20).reshape(10, 2)
5y = np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1])
6
7skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
8
9for train_idx, test_idx in skf.split(X, y):
10    print("test labels:", y[test_idx])

This is usually the better choice for classification because each fold stays more representative of the original class distribution. That gives metrics such as accuracy, precision, recall, and ROC AUC a fairer evaluation context.

Why stratification matters in imbalanced classification

Imagine a binary classification dataset where 95 percent of the samples belong to class 0 and only 5 percent belong to class 1. With plain KFold, some folds may contain very few minority-class examples or even none at all. That creates problems:

  • the model may never be tested meaningfully on the minority class
  • metrics can fluctuate wildly from fold to fold
  • some scoring functions become unstable or invalid

With StratifiedKFold, each fold is more likely to reflect the original imbalance rather than accidentally exaggerating it.

This does not "fix" class imbalance, but it does make the validation split more faithful.

When KFold is still the right choice

KFold is not worse in general. It is the right baseline for problems where there is no class label distribution to preserve.

Typical cases include:

  • regression
  • unsupervised workflows
  • classification tasks where stratification is not relevant or not possible

If your target is continuous, StratifiedKFold is not the default tool because it expects class labels, not arbitrary numeric targets.

For regression, use KFold or RepeatedKFold, or build a careful custom binning approach only if you have a specific reason to approximate stratification.

Example decision rule

A practical rule of thumb is:

  • use StratifiedKFold for classification
  • use KFold for regression

Example with cross-validation scoring:

python
1from sklearn.datasets import load_breast_cancer
2from sklearn.linear_model import LogisticRegression
3from sklearn.model_selection import cross_val_score, StratifiedKFold
4
5X, y = load_breast_cancer(return_X_y=True)
6model = LogisticRegression(max_iter=5000)
7cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
8
9scores = cross_val_score(model, X, y, cv=cv, scoring="accuracy")
10print(scores.mean())

This is the sort of setup you want for ordinary supervised classification in scikit-learn.

Common Pitfalls

The biggest mistake is using KFold on a strongly imbalanced classification problem and trusting the resulting metrics without checking fold composition.

Another common issue is assuming StratifiedKFold balances the features. It does not. It only preserves the target distribution.

People also forget to shuffle when the dataset may be ordered by class or source. Without shuffling, both KFold and StratifiedKFold can produce misleading splits if the rows are structured in a non-random order.

Finally, StratifiedKFold is not a general regression splitter. If the target is continuous, plain stratification by label is not the right abstraction.

Summary

  • 'KFold splits rows without preserving class proportions.'
  • 'StratifiedKFold preserves target-class distribution across folds.'
  • For classification, especially imbalanced classification, StratifiedKFold is usually the better choice.
  • For regression, KFold is usually the correct default.
  • Always think about data ordering and use shuffling when appropriate.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.