Stratified k-fold cross-validation

Most people meet k-fold cross-validation as a recipe. Cut the dataset into k folds, let each fold take a turn as the test set, train on the rest, and average the scores. The recipe says nothing about which samples land in which fold, so it is easy to assume that cutting the data into consecutive pieces is harmless. For a classifier, that assumption can break badly. A split that looks neutral can quietly make the evaluation meaningless, even when the model itself is perfectly good.

The lesson of this article is that the way data is cut into folds matters just as much as the model. The source uses the iris dataset to show the problem. Before you look at any output, try to reason about it yourself. Suppose the iris samples are arranged so that all samples of one class come first, then the next class, then the last class. If you cut that ordered sequence into three consecutive folds without shuffling or stratifying, what would each fold contain? Commit to an answer in your head before reading on.

The source loads the iris target labels with scikit-learn and prints them directly, which makes the ordering visible. The labels come in three consecutive blocks: the first third of the samples is class 0, the second third is class 1, and the last third is class 2. Nothing about the classifier is wrong here. The problem is in the split, because a standard k-fold split assigns rows by their position in the dataset, not by their class.

Now imagine three-fold cross-validation on this ordered data. The first fold is made only of class 0. In the first split, the test set therefore holds only class 0, and the training set holds only classes 1 and 2. The next two splits have the same mismatch: the classes in the training and test sets are different in every split. A classifier trained on two folds has never seen the class that the held-out fold contains, so it has no fair chance to predict it.

The source reports that three-fold accuracy on this arrangement would be zero, which is far below what a working classifier achieves on iris. The zero is a symptom of the split, not of the classifier. Keep this in mind when a result looks wrong: the place to look first is which samples ended up in which fold.

The cv parameter of cross_val_score can take a plain number, which sets the number of folds and hides the splitting rule behind it. It can also accept a cross-validation splitter, which is an object that controls how the data is split. Instead of accepting whatever rule sits behind a number, you build the splitter yourself and hand it over. That gives you finer control over the splitting step, and it makes the rule visible in your code.

The splitter used in the source is the KFold class, which lives in the model_selection module. You import it and instantiate it with the number of folds you want, set through the n_splits argument. The instance is a reusable object, so it can be passed to cross_val_score directly through the cv argument. The call then uses exactly the folds this object produces, rather than building its own.

Watch where the scores come from. Each score is produced by training on the samples in the other folds and testing on the held-out fold. The score list is the end product, so when it looks wrong, the cause is usually upstream, in which samples ended up in which fold. The zero output above is the result the source describes for unshuffled three-fold splitting on this ordered data. The model is not the place to look first.

A small example makes the divergence easy to see. Take nine samples with labels in three consecutive blocks: [0, 0, 0, 1, 1, 1, 2, 2, 2]. This dataset is invented for teaching, not taken from the source, but it has the same ordered structure as the iris example. Each class makes up one third of the data, so the overall class mix is even.

Read the diagram as a trace. Each fold is a consecutive run of positions, and because the labels are sorted, each run holds one class. No fold contains the full class mix, so no single test fold tells you how the classifier handles the other two classes. Trace this to the step that caused it: the position-based assignment, not the classifier.

Stratified k-fold cross-validation fixes the problem by changing what the split must preserve. The rule is that the proportions between classes are the same in each fold as they are in the whole dataset. The split is still into k folds, and each fold still takes a turn as the test set. What changes is that each fold is built to mirror the class composition of the full data, rather than whatever happens to sit in a block of rows.

Here $n_{c,f}$ is the number of samples of class $c$ in fold $f$, $n_f$ is the size of fold $f$, $n_c$ is the number of samples of class $c$ in the whole dataset, and $N$ is the total number of samples. The rule says that for each class, its share inside a fold should match its share in the full dataset. The approximately-equal sign is deliberate. Real datasets do not always divide evenly, so the goal is that each fold is representative, not that every fold has an identical count.

The source gives a simple ratio to make this concrete. If 90% of the samples belong to class A and 10% belong to class B, stratified cross-validation ensures that in each fold, 90% of the samples are class A and 10% are class B. The composition of every fold is anchored to the composition of the whole dataset, which is the property a fair test set needs.

The stratified column in the comparison is the generated nine-sample example, where each class is one third of the data, so every fold gets exactly one sample of each class. With real data the counts rarely divide this neatly, and each fold only needs to approximate the overall proportions.

The source recommends stratified k-fold cross-validation as the usual choice for evaluating a classifier, because it gives more reliable estimates of generalization performance. The reasoning is concrete. Suppose only 10% of the samples belong to class B. With standard k-fold, one fold can easily contain only samples of class A. Using that fold as the test set tells you little about overall performance, because the classifier was never asked to handle class B on that split.

Stratification matters most for imbalanced classification, where a standard split can produce a test fold that says little about the minority class. Scikit-learn uses stratified k-fold for classification by default. The stratified counterpart of the KFold splitter is available in the same model_selection module, so the choice is about which splitter you pass to cv, not about changing the rest of your code.

The source offers two remedies for the ordered-label failure. The first is stratifying the folds. The second is to shuffle the data before cutting the folds, so that the ordering of samples by label is removed. Shuffling is controlled through the shuffle parameter of KFold. Setting shuffle=True randomizes the order of the samples before the split. The source reports that shuffling before splitting gives a much better result than the unshuffled three-fold run, and it attributes that improvement to removing the label ordering.

Shuffling introduces randomness, and uncontrolled randomness makes results hard to reproduce. The source is direct about this: if you shuffle the data, you also need to fix random_state to get a reproducible shuffling. Without a fixed random_state, each run of cross_val_score would use a different split and so could give a different result. The source adds that this might not be a problem in every setting, but it can be surprising.

The same reasoning helps you reproduce someone else's results. If you know the number of folds, the shuffle setting, and the random_state, you have pinned down the split. Using the same number of folds is the detail the source emphasizes when describing reproduction. Recording those settings alongside a score is what makes the score meaningful to someone else.

Stratification is about class proportions, so it applies to classification. For regression, scikit-learn uses standard k-fold cross-validation by default. The source notes that one could try to make each fold representative of the different values a regression target takes, but this is not a commonly used strategy and would surprise most users. The practical lesson is to match the splitting strategy to the kind of prediction being evaluated.

The flowchart encodes the source's guidance, with one caveat. The source says the defaults of standard k-fold for regression and stratified k-fold for classification work well for most use cases. It also names a situation where people often reach for a different strategy: reproducing someone else's results with standard, non-stratified k-fold on a classification dataset. Whether that choice is harmful depends on how the samples are ordered, which is the part learners tend to miss.