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.
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:
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:
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
StratifiedKFoldfor classification - use
KFoldfor regression
Example with cross-validation scoring:
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
- '
KFoldsplits rows without preserving class proportions.' - '
StratifiedKFoldpreserves target-class distribution across folds.' - For classification, especially imbalanced classification,
StratifiedKFoldis usually the better choice. - For regression,
KFoldis usually the correct default. - Always think about data ordering and use shuffling when appropriate.
Related reading
- Streaming large training and test files into Tensorflow's DNNClassifier
- String Distance Matrix in Python
- String Matching Using Recurrent Neural Networks
- String Matching Using Recurrent Neural Networks
- Structuring dataset for OpenAI's GPT-3 fine tuning
- Submitting Assignment on Coursera ML in Octave
- Suboptimal convergence in PyTorch compared to TensorFlow when using Adam optimizer
- ''super'' object has no attribute ''__sklearn_tags__''
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.