What is stratified k-fold cross-validation and when is it necessary?
Stratified k-fold cross-validation preserves the class proportions in every fold, so each training and validation split represents the target population. It is especially important for imbalanced or small classification datasets, but it does not replace group-aware or time-aware splitting and does not fix class imbalance during training.
How to think about it
Short answer
Stratified k-fold cross-validation estimates model performance by splitting data into k folds while keeping roughly the same proportion of each class in every fold. It is most important for imbalanced or small classification datasets, but it is not a substitute for group-aware or time-aware splitting, and it does not rebalance the model during training.
Why plain k-fold can mislead
Suppose a fraud model has 10,000 transactions:
- 9,990 legitimate transactions, the negative class
- 10 fraudulent transactions, the positive class
With five-fold cross-validation, each validation fold contains 2,000 transactions on average. The average is two fraud cases per fold. The average is not a promise.
A plain random KFold can put zero fraud cases in one fold and four in another. A fold with no actual fraud cases cannot measure recall, because recall divides true positives by the number of actual positives. ROC AUC is also undefined when a validation fold contains only one class, because it needs both positive and negative examples to construct its curve.
Worse, a model that predicts “legitimate” for every transaction gets 100 percent accuracy on that zero-fraud fold. It has detected nothing. The metric still looks perfect.
That is the problem stratification addresses: not model quality, but whether each validation measurement is asking a sensible question.
What k-fold cross-validation does
A fold is one subset of the data. In k-fold cross-validation, the data is divided into k folds. The model trains on k - 1 folds and is evaluated on the remaining fold. The process repeats until every row has served as validation data once.
For five folds and 10,000 rows:
- Four folds, or 8,000 rows, train the model.
- One fold, or 2,000 rows, evaluates it.
- This repeats five times.
- The five scores are averaged, and their spread is reported.
The training sets overlap. These are not five independent test sets. They are five views of the same dataset, which is why cross-validation is useful for model comparison but should not be confused with a final untouched test set.
How stratification works
Stratification means allocating examples from each class separately before assembling the folds. If 0.1 percent of the full dataset is positive, each fold gets approximately 0.1 percent positive as well.
For the fraud example, stratified five-fold cross-validation can produce:
| Split | Negative | Positive | Total |
|---|---|---|---|
| Each validation fold | 1,998 | 2 | 2,000 |
| Each training fold | 7,992 | 8 | 8,000 |
The counts are exact here because 9,990 and 10 divide evenly across five folds. When they do not, folds differ by at most one example where possible.
This makes per-fold metrics more comparable. A fold with two fraud cases is still statistically fragile, but at least recall has a denominator. If the model catches one of those two cases, that fold’s recall is 50 percent. If it catches both, recall is 100 percent. The score is noisy, not meaningless.
Stratification preserves class counts. It does not select “easy” and “hard” examples intelligently, preserve every feature distribution, or guarantee that all demographic subgroups appear in equal proportions. It preserves the target label distribution and nothing more.
A practical scikit-learn implementation
For independent rows, I would make the splitter explicit:
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_validate
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
cv = StratifiedKFold(
n_splits=5,
shuffle=True,
random_state=42,
)
model = make_pipeline(
StandardScaler(),
LogisticRegression(max_iter=1000),
)
scores = cross_validate(
model,
X,
y,
cv=cv,
scoring={
"f1": "f1",
"roc_auc": "roc_auc",
},
)
print(scores["test_f1"].mean())
print(scores["test_roc_auc"].mean())
The result contains five F1 scores and five ROC AUC scores. I would inspect the mean and the variation across folds, not only the mean. With very few positive examples, a large spread is itself an important result.
shuffle=True matters when rows are independent but happen to be stored in an arbitrary or sorted order. random_state=42 makes the partition reproducible. The particular seed is not sacred; reproducibility is.
There is also a scikit-learn default worth knowing. When cross_val_score or cross_validate receives an integer such as cv=5, and the estimator is a classifier with a binary or multiclass target, scikit-learn automatically chooses StratifiedKFold. The default splitter does not shuffle. Passing the splitter explicitly makes the intended shuffling and seed visible instead of leaving an important assumption hidden.
Put preprocessing inside a pipeline. If you scale, impute, select features, or otherwise learn something from the data before cross-validation, the validation fold can influence that transformation. That is leakage: information from the supposed evaluation data enters training. The score is then too optimistic.
The same principle applies to a final train/test split: train_test_split(..., stratify=y) preserves class proportions in the holdout. See the related guide on train/test split.
When stratification is necessary
“Necessary” depends on the number of minority examples and the metric, not on a magical imbalance threshold.
| Situation | Sensible choice | Reason |
|---|---|---|
| Imbalanced binary classification | StratifiedKFold | Every fold is more likely to contain both classes |
| Imbalanced multiclass classification | StratifiedKFold | Rare classes are not accidentally omitted |
| Balanced, large, independent classification | KFold or StratifiedKFold | Either may give almost identical results |
| Repeated users, patients, or devices | StratifiedGroupKFold | Groups must not be split across training and validation |
| Time-ordered data | Chronological or forward-chaining split | Random folds would train on the future |
| Continuous regression target | Ordinary KFold or a regression-aware design | Regression has no discrete class label |
| Several labels per row | A multilabel-aware splitter | Ordinary stratification expects one target class per row |
For group data, avoiding leakage is more important than achieving perfect class proportions. If one patient has 20 records, putting some records in training and the rest in validation lets the model recognise the patient rather than generalise to a new patient. StratifiedGroupKFold tries to preserve class proportions while keeping groups intact, but exact balance may be impossible when groups have very different class compositions.
For time series, random stratification is usually wrong even when the class counts look attractive. A model predicting equipment failure in December should not train on labelled January failures merely because the rows were shuffled. Use a chronological evaluation that matches deployment.
For regression, binning a continuous target and then applying classification-style stratification can sometimes make folds cover similar target ranges. It is a heuristic, not standard stratified cross-validation. The bin boundaries can be arbitrary, and the split must still respect groups or time if those dependencies exist.
The senior-level nuance
Stratification is a split design, not an imbalance treatment.
In the fraud example, each training fold still contains only eight fraud cases. Stratification does not create more signal. Class-weighted loss, threshold tuning, undersampling, or oversampling address different problems. If you oversample or use SMOTE, do it inside each training fold. Oversampling the full dataset before cross-validation can place duplicates or synthetic relatives of training examples in validation, producing a contaminated score.
The metric also matters. Accuracy is a poor choice when 99.9 percent of transactions are legitimate. Depending on the business cost, I might report recall, precision, F1, balanced accuracy, or precision-recall AUC. Stratification makes those metrics computable and more comparable across folds; it does not make them automatically appropriate.
The smallest class limits the number of useful folds. If there are three positive examples and five folds, no splitter can place at least one positive in every validation fold. Reduce n_splits, gather more labelled data, or use a different evaluation design. With 10 positives and five folds, each fold has only two positives, so one mistake changes that fold’s recall by 50 percentage points. Repeated stratified cross-validation can show how sensitive the estimate is to the partition, but it does not manufacture more independent fraud cases.
A final trap is assuming that stratification preserves everything important. If a rare class is concentrated in one region, hospital, customer segment, or device type, label-only stratification may still produce unrepresentative folds. Stratify by the label when that is the main concern, but use groups, time boundaries, or carefully designed slices when those are part of the deployment problem.
Failure modes I would watch for
- A fold has no minority examples. The first symptom is a warning about the least-populated class,
nanfor ROC AUC, or fold scores that swing wildly between zero and one. - The split ignores groups. The first symptom is suspiciously excellent offline performance followed by a sharp drop when the model meets entirely new customers or patients.
- Preprocessing happens before cross-validation. The first symptom is a cross-validation score that looks unusually strong and then falls on a genuinely untouched test set.
- The split ignores time. The first symptom is a healthy offline score but poor performance immediately after launch, when the model faces later data.
What they’ll ask next
Does stratification fix class imbalance?
No. It only allocates rows so that folds have similar label proportions. Use class weights, resampling, threshold selection, or a different objective to change how the model learns. Those operations must be fitted separately inside each training fold.
How do you choose k when the positive class is rare?
Start with the smallest class count. If every fold needs a positive example, k cannot exceed that count. Then consider metric stability, training cost, and the fact that more folds give fewer minority examples per validation fold. Report the fold spread rather than hiding it behind one average.
Would you stratify if the data contains repeated patients or future events?
Not with ordinary StratifiedKFold. Use a group-aware splitter for repeated entities and a chronological or forward-chaining split for time-dependent data. Preventing leakage takes priority over perfectly matching class proportions.
Say this in the interview
“Stratified k-fold preserves each class’s proportion in every fold, which prevents imbalanced classification metrics from being undefined or dominated by arbitrary class counts; I use it for independent classification data, but switch to group- or time-aware splits when those dependencies matter.”