Stratified Splitting

September 2026 · 8 min read

You reserve five test examples from a dataset of 20. Four of the 20 examples are positive, so you expect about one positive in test. A random split can still give you none. You then have no test evidence about whether the model catches positives.

Stratified splitting controls this allocation. It divides each class between the subsets in roughly the requested proportions. It does not change the class balance to 50/50, train a better classifier, or guarantee that the future population will match this dataset.

Allocate each class, then combine the rows

Our dataset has four positives and 16 negatives. A 25% test split can reserve one positive and four negatives. Training receives the remaining three positives and 12 negatives. Both subsets have the original 20% positive rate:

The counts are controlled, but which positive and which four negatives go to test are still randomized. In ordinary stratification, shuffle within each class, allocate the requested counts, and combine the class-specific selections. Integer rounding means exact proportions are not always possible.

Watch the class count change under random splitting

Each card below is a unique example. P marks a positive and N a negative. Switch to random splitting and draw several times. Then switch back to stratification: the identities can change while the class counts stay fixed.

For this dataset, about 28.2% of uniformly random five-row test subsets contain no positive. That is not a bug in random splitting. It is a consequence of drawing a small sample from an imbalanced collection.

A runnable stratified split

Pass the label vector to stratify. Scikit-learn then uses class membership when allocating rows. Keep the same row IDs for features and labels, as in the hold-out example:

import numpy as np
from sklearn.model_selection import train_test_split

row_id = np.arange(20)
y = np.array([1] * 4 + [0] * 16)

train, test = train_test_split(
    row_id, test_size=0.25, random_state=42, stratify=y
)

assert len(train) == 15 and len(test) == 5
assert y[train].sum() == 3
assert y[test].sum() == 1
assert set(train).isdisjoint(test)

For a three-way split, stratify both calls. After reserving test IDs, the second call splits only the development IDs, so its stratification labels must be y[dev], not the full label array. Verify counts in all three subsets, especially when a class has only a few examples.

Preserving proportions does not create evidence

Our stratified test set always contains one positive. That fixes the missing-class problem, but recall is still either 0% or 100% depending on one prediction. A stable class ratio is not a stable performance estimate.

With three positive examples and five folds, every validation fold cannot contain a positive. With just one example of a class, that example cannot be in both train and test. Standard stratified shuffling also has minimum class-count and subset-size requirements. Even when a call succeeds, extreme imbalance and rounding can leave a subset with no examples of a very rare class; inspect the actual counts.

Possible responses include collecting more examples, reducing the number of folds, or reporting that a class cannot yet be evaluated reliably. Do not duplicate the rare example before splitting: copies could land on opposite sides, producing an easier evaluation without adding independent evidence.

Know which problem stratification solves

Prioritize the evaluation question. For unseen customers, keep customer groups separate even if their sizes make the class ratios imperfect. StratifiedGroupKFold tries to preserve class proportions while respecting groups, but exact balance may be impossible. For future traffic, use a time split rather than mixing dates to obtain prettier proportions.

Stratification is different from balancing the training set

A 1%-positive dataset remains approximately 1% positive in each stratified subset. If you oversample positives or use class weights to help training, apply that decision within training data. Evaluate on a distribution appropriate to the intended use. Artificially balancing test data changes what aggregate metrics such as precision mean.

There is one further subtlety: stratification deliberately makes class proportions more similar across splits. Smaller variation in fold scores therefore need not mean that all sources of performance uncertainty are small. The scikit-learn guide explicitly discusses this limitation.

For this chapter, the useful habit is straightforward: inspect class counts before and after splitting, then check groups and time separately. Continue with group splitting to see why distinct row IDs can still hide shared information.

References and further reading

• Scikit-learn: StratifiedShuffleSplit — class-preserving shuffled splits and the sampling design.
• Scikit-learn: what stratification does and does not solve — rare classes, fold homogeneity, and limits on uncertainty interpretation.
• Scikit-learn: StratifiedGroupKFold — balancing labels while preserving group separation.