What Is Stratified Sampling in Machine Learning?
- 1 hour ago
- 28 min read
You build a classifier on a dataset with 500 fraud cases out of 100,000 transactions. You run a plain random 80/20 split. By chance, the test set ends up with only 40 fraud examples instead of the expected 100 — or worse, a tiny validation fold contains none at all. Your precision and recall numbers swing wildly between runs, and you can't tell whether a model change actually helped or whether you just got a luckier split. This is the exact problem stratified sampling exists to prevent.
TL;DR
Stratified sampling divides data into subgroups (strata) and samples independently from each one, most often to preserve class proportions across training, validation, test, or cross-validation partitions.
It stabilizes evaluation for classification problems, especially when a class is rare.
It does not balance classes — it preserves whatever distribution already exists in the data.
It's used constantly in scikit-learn workflows through stratify=y, StratifiedKFold, and related tools.
It cannot fix a biased source dataset, and it can be actively wrong for time-ordered or grouped data.
What Is Stratified Sampling in Machine Learning?
Stratified sampling is a method that divides a population or dataset into meaningful, non-overlapping subgroups called strata, then samples independently within each group. In machine learning, stratification usually means preserving class proportions across training, validation, test, or cross-validation partitions — improving evaluation stability without changing the underlying class balance.
Table of Contents
What Is Stratified Sampling?
Before any formula, here is the plain-English version: instead of drawing observations from the whole dataset at random, you first split the dataset into groups that share something meaningful, then sample from each group separately.
A population is the full set of items you care about. A sample is the subset you actually collect or select. A stratum (plural: strata) is one of the non-overlapping subgroups you divide the population into — for example, "customers who churned" and "customers who stayed." The variable used to form these groups is the stratification variable. The sampling frame is the list or dataset from which the sample is actually drawn, which may or may not match the true population exactly.
Good strata are mutually exclusive — every observation belongs to exactly one stratum — and collectively exhaustive — every observation is assigned to some stratum, with none left out (NIST, n.d.). Within each stratum, you still sample independently, usually with simple random sampling. The goal is for each stratum to be internally similar with respect to the variable you care about, while being meaningfully different from the other strata. A dataset of hospital patients grouped by "diagnosed" and "not diagnosed" is a simple, faithful example: each group is internally consistent on the outcome that matters, and the two groups differ from each other in exactly the way you want to preserve.
Stratified Sampling in Statistics vs Machine Learning
The phrase "stratified sampling" gets used for three related but distinct activities, and mixing them up leads to overconfident claims about how representative a dataset actually is.
Statistical stratified sampling — selecting observations from a larger population or sampling frame to build a dataset in the first place.
Stratified dataset splitting — dividing an existing dataset into training, validation, and test partitions while trying to preserve a target distribution, typically via stratify=y in scikit-learn's train_test_split.
Stratified cross-validation — constructing cross-validation folds, via tools such as StratifiedKFold, that preserve class proportions as closely as possible across folds.
Scikit-learn's documentation is explicit that its stratified splitters address a model-evaluation concern, not a survey-representativeness concern: "Stratification on the class label solves an engineering problem rather than a statistical one" (scikit-learn, n.d.). In other words, StratifiedKFold and StratifiedGroupKFold make your evaluation folds resemble each other; they do nothing to correct whatever bias already exists in how the dataset itself was collected.
Concept | What it operates on | Typical goal | Common tool |
Statistical stratified sampling | A population or sampling frame | Build a representative dataset | Survey design, manual sampling plans |
Stratified dataset splitting | An existing dataset | Preserve class proportions across train/test | train_test_split(stratify=y) |
Stratified cross-validation | An existing dataset | Preserve class proportions across CV folds | StratifiedKFold, StratifiedShuffleSplit |
The practical implication: a perfectly stratified train-test split does not make a biased dataset representative of the real-world population you'll deploy against. It only makes your training and test partitions consistent with each other.
Why Stratification Matters in Machine Learning
In classification tasks, especially those involving fraud detection, disease screening, equipment failure, or customer churn, one class is often rare. A plain random split can, by chance, under-represent that minority class in the test set or, in the worst case for a very small class, exclude it from a fold entirely. Stratification keeps the minority-class proportion consistent across partitions, which brings several practical benefits:
More stable evaluation. Metrics don't swing wildly between runs just because of which rows landed in the test set.
Comparability between partitions. Training, validation, and test sets reflect the same class mix, so a performance drop is more likely to reflect a real modeling issue rather than a lucky or unlucky split.
Better diagnosis of precision, recall, F1, sensitivity, specificity, ROC-AUC, and PR-AUC. These metrics depend on having enough minority-class examples in every partition to compute meaningfully.
Reduced risk of a class vanishing from a small validation or test set entirely, which would make some metrics undefined.
It's worth stating plainly, because it's the single most misunderstood idea in this space: stratification usually preserves an existing distribution; it does not automatically balance classes. If your dataset is 95% negative and 5% positive, a stratified split gives you a 95/5 training set and a 95/5 test set — not a 50/50 one. Google's machine learning materials make a related point about why this matters: accuracy alone is a poor way to judge a model trained on an imbalanced dataset, because a model that always predicts the majority class can still score a high accuracy while being useless for the minority class (Google, n.d.). A representative split is not the same thing as a balanced split, and confusing the two is a recurring source of misplaced confidence in imbalanced-classification projects.
How Stratified Sampling Works Step by Step
At a conceptual level, stratified sampling — whether for a survey or a machine-learning split — follows the same sequence:
Define the target population or dataset.
Select an appropriate stratification variable (often the class label).
Form non-overlapping strata from that variable.
Measure each stratum's size.
Choose an allocation strategy (how many observations to draw from each stratum).
Randomly sample or split independently within each stratum.
Combine the selected observations from all strata into the final sample or partition.
Verify the resulting distributions match what was intended.
Apply design weights if the allocation was disproportionate and population-level inference is required.
Randomization still happens — it just happens within each stratum rather than across the whole dataset. The table below shows a small, internally consistent example of an 80/20 stratified split on a two-class dataset.
Class | Total count | Training (80%) | Test (20%) |
Negative | 8,000 | 6,400 | 1,600 |
Positive | 2,000 | 1,600 | 400 |
Total | 10,000 | 8,000 | 2,000 |
Both partitions preserve the original 80/20 class ratio. This is the behavior stratify=y is designed to guarantee.
The Mathematics Behind Stratified Sampling
The notation is simple once defined:
N: total population size
N_h: size of stratum h
n: desired total sample size
n_h: number sampled from stratum h
H: number of strata
Proportional allocation assigns each stratum a sample size proportional to its share of the population:
$$n_h = n \cdot \frac{N_h}{N}$$
Because sample sizes must be whole numbers, rounding is applied, and the individual rounded values are usually adjusted slightly so they sum exactly to n.
The stratified mean estimator combines each stratum's mean, weighted by its share of the population:
$$\bar{y}{st} = \sum{h=1}^{H} \frac{N_h}{N} \bar{y}_h$$
The intuition: a stratum that represents 60% of the population should contribute 60% of the weight to the overall estimate, regardless of how many observations you happened to sample from it.
When sampling is disproportionate — some strata deliberately over- or under-sampled relative to their true share — a basic design weight corrects for this when estimating population-level quantities:
$$w_h = \frac{N_h}{n_h}$$
Rows from an under-sampled stratum get a larger weight; rows from an over-sampled stratum get a smaller weight. Without this correction, disproportionate sampling silently distorts any population-level estimate computed from the sample.
Stratification can reduce sampling variance when strata are meaningfully related to the variable being estimated — that is, when observations within a stratum are more similar to each other than to observations in other strata. But this is a conditional benefit, not a guarantee. Poorly chosen strata, tiny strata, incorrect allocation, measurement error, or a stratification variable unrelated to the outcome can erase any variance advantage, and in some cases add unnecessary complexity without benefit.
Types of Stratified Sampling and Allocation
Different allocation strategies decide how many samples come from each stratum.
Proportional allocation: each stratum contributes according to its share of the population. This is what stratify=y effectively replicates in a train-test split.
Equal allocation: every stratum contributes the same number of observations, regardless of size. Useful when you specifically want equal statistical power per group, at the cost of over-representing small strata.
Disproportionate allocation: some strata are deliberately sampled at higher or lower rates than their population share, often to guarantee enough minority-class examples for analysis. Requires design weights for unbiased population-level inference.
Neyman allocation: allocates more observations to strata that are both larger and more variable, following the proportional relationship
$$n_h \propto N_h S_h$$
where S_h is the standard deviation of the variable of interest within stratum h. Neyman allocation minimizes overall sampling variance for a fixed total sample size, assuming sampling costs are equal across strata.
Cost-optimal allocation: accounts for the fact that sampling may cost more in some strata than others (for example, surveying a remote region), balancing variance reduction against budget.
Allocation method | Core idea | Main advantage | Main limitation | Typical use |
Proportional | Match stratum sample size to population share | Simple, matches natural distribution | Small strata may still be under-sampled in absolute terms | ML train-test splits, CV folds |
Equal | Same count from every stratum | Equal statistical power per group | Distorts overall proportions | Comparing subgroups directly |
Disproportionate | Deliberately over/under-sample specific strata | Guarantees minority-class coverage | Requires design weights for inference | Rare-event research, oversampled minority classes |
Neyman | Weight by size and variability | Minimizes variance for fixed budget | Requires knowing/estimating S_h in advance | Formal survey design |
Cost-optimal | Balance variance against sampling cost | Efficient use of a limited budget | More complex to plan | Field surveys, expensive data collection |
Most everyday machine-learning splitting uses proportional allocation almost exclusively — it's what stratify=y and StratifiedKFold implement by default. Neyman and cost-optimal allocation belong more to formal survey and experimental design, where the goal is inference about a broader population rather than model evaluation.
Worked Example: Stratifying an Imbalanced Classification Dataset
Consider a hypothetical dataset of 10,000 observations across two classes: 9,500 majority-class (label 0) and 500 minority-class (label 1) — a 95:5 imbalance typical of fraud or defect-detection problems.
An 80/20 stratified split produces:
Partition | Majority class (label 0) | Minority class (label 1) | Total | Minority % |
Training (80%) | 7,600 | 400 | 8,000 | 5.0% |
Test (20%) | 1,900 | 100 | 2,000 | 5.0% |
Both partitions preserve the original 5% minority share exactly, because the counts divide evenly in this constructed example.
Contrast this with a possible — not guaranteed — outcome from an unstratified random 80/20 split on the same data. Because 500 minority-class rows are a small fraction of 10,000, ordinary sampling variability could plausibly place, say, only 82 minority-class rows in the test set instead of the expected 100, or occasionally far fewer if the dataset were smaller still. This is a description of a possible result, not a typical guaranteed one, and the risk grows as the minority class shrinks or the dataset gets smaller. When it does happen, the practical effect is real: precision, recall, and F1 computed on a test set with fewer positive examples become noisier and less trustworthy for comparing candidate models.
Implementing a Stratified Train-Test Split in Python
Scikit-learn's train_test_split accepts a stratify argument that performs this split directly (scikit-learn, n.d.).
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
# Create an imbalanced binary classification dataset
X, y = make_classification(
n_samples=10000,
n_features=20,
n_informative=10,
n_classes=2,
weights=[0.95, 0.05],
random_state=42,
)
X_train, X_test, y_train, y_test = train_test_split(
X, y,
test_size=0.2,
stratify=y,
random_state=42,
)
def class_counts(y_array, label):
counts = np.bincount(y_array)
proportions = counts / counts.sum()
print(f"{label}: counts={counts}, proportions={np.round(proportions, 3)}")
class_counts(y_train, "Training set")
class_counts(y_test, "Test set")Example output (exact counts depend on the generated dataset and its random seed, but proportions should closely track the original 95/5 split):
Training set: counts=[7600 400], proportions=[0.95 0.05]
Test set: counts=[1900 100], proportions=[0.95 0.05]Each argument matters here. stratify=y tells the function to split in a stratified fashion using y as the class labels (scikit-learn, n.d.). random_state=42 makes the split reproducible across runs. test_size=0.2 sets the proportion held out for testing. When stratify=None (the default), the split is a plain random shuffle with no guarantee about class proportions in either partition.
One practical error worth anticipating: if a class has fewer members than the number of partitions it needs to be divided across — for instance, a class with only one example being split into both a training and a test set — scikit-learn will raise an error rather than silently producing an empty or single-member group.
Stratified Cross-Validation
Standard K-fold cross-validation divides the dataset into K folds, training on K−1 folds and validating on the remaining one, repeated K times. Plain K-fold does not consider class labels when forming folds, so a fold can end up with a skewed class distribution purely by chance — a bigger risk as K grows relative to the size of the rarest class.
Stratified K-fold cross-validation, via StratifiedKFold, is a variation of KFold that returns stratified folds, made by preserving the percentage of samples for each class in a binary or multiclass classification setting (scikit-learn, n.d.). This keeps fold-to-fold class distribution consistent, which matters directly for the reliability of any classification metric computed per fold.
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.linear_model import LogisticRegression
X, y = make_classification(
n_samples=2000,
n_features=15,
weights=[0.9, 0.1],
random_state=42,
)
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
model = LogisticRegression(max_iter=1000)
scores = cross_val_score(model, X, y, cv=skf, scoring="f1")
print("Per-fold F1 scores:", np.round(scores, 3))
print("Mean F1:", round(scores.mean(), 3))The number of folds interacts directly with the size of the rarest class: if a class has only 8 members and you request 10 folds, some folds cannot receive even one example of that class, which will raise an error or produce degenerate folds. shuffle=True combined with random_state gives a reproducible shuffle before the split; without shuffling, the fold construction follows the order the data was given in.
Related tools worth knowing:
RepeatedStratifiedKFold repeats the entire stratified K-fold procedure multiple times with different randomization, producing a more stable performance estimate at the cost of proportionally more computation.
StratifiedShuffleSplit generates a specified number of independent stratified train/test splits, useful when you want repeated holdout evaluation rather than exhaustive fold coverage; scikit-learn again frames this as "an engineering problem rather than a statistical one" (scikit-learn, n.d.).
cross_val_score and cross_validate run a model across the folds produced by any of these splitters and aggregate scores.
It's worth restating plainly: stratified cross-validation does not eliminate dataset bias. If the underlying data collection process is skewed relative to the real-world deployment population, every stratified fold will faithfully reproduce that same skew.
Grouped Data and StratifiedGroupKFold
Many datasets have multiple rows belonging to the same underlying entity — a patient with several visits, a customer with several transactions, a device with many sensor readings, or a document split into many paragraphs. If a random or ordinary stratified split places some of that entity's rows in training and others in the test set, the model can effectively "see" a version of the test entity during training, inflating apparent performance. This is data leakage through grouping.
StratifiedGroupKFold addresses this by combining StratifiedKFold and GroupKFold: it attempts to create folds that preserve the percentage of samples from each class as much as possible, given the constraint of non-overlapping groups between splits (scikit-learn, n.d.). The distinction from GroupKFold is instructive: GroupKFold tries to balance the number of distinct groups per fold, while StratifiedGroupKFold tries to balance the class proportions per fold, subject to groups never being split across folds. Because both goals can conflict — a group might contain many minority-class rows that can't be evenly divided without breaking group integrity — perfect stratification is not always mathematically achievable, and fold sizes may end up uneven.
A concise scenario: a hospital readmission model built from patient records, where each patient contributes multiple encounters. Grouping by patient ID and requesting StratifiedGroupKFold on the readmission label keeps every encounter for a given patient inside a single fold, while still trying to keep each fold's readmission rate close to the dataset's overall rate.
Time-Series Data: Why Ordinary Stratification Can Be Wrong
Random stratification assumes it's fine to shuffle rows freely as long as class proportions are preserved. For time-ordered data, that assumption breaks: shuffling lets information from the future leak into the training set, a problem often called temporal leakage. A model evaluated this way can appear to "predict" outcomes it would never actually have access to at prediction time in production.
Scikit-learn's TimeSeriesSplit is built for this case. It provides train/test indices to split time-ordered data, where other cross-validation methods are inappropriate, as they would lead to training on future data and evaluating on past data (scikit-learn, n.d.). In each split, the training set is a growing prefix of the data and the test set is the segment immediately following it — successive training sets are supersets of the ones before them, never overlapping with future test periods.
import numpy as np
from sklearn.model_selection import TimeSeriesSplit
X = np.arange(24).reshape(12, 2)
y = np.arange(12)
tscv = TimeSeriesSplit(n_splits=3, test_size=2, gap=1)
for i, (train_index, test_index) in enumerate(tscv.split(X)):
print(f"Fold {i}: train={train_index}, test={test_index}")Distribution drift — where the class or feature distribution genuinely shifts over time — is common in forecasting and backtesting contexts. Monitoring class prevalence across time windows is still worthwhile, but the correct response to an imbalance in a time-series problem is not to randomly rearrange the data merely to obtain matching proportions; chronology takes priority, and any stratification-style adjustment should happen within each fixed time window, not across them.
Regression and Continuous Targets
Stratification is most straightforward for categorical labels, where "preserve the proportions" has an obvious meaning. For a continuous regression target, there's no natural category to stratify on, so a common heuristic is to bin the continuous target into a small number of quantile-based groups and stratify on those bins instead.
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
rng = np.random.default_rng(42)
y_continuous = rng.exponential(scale=1000, size=2000)
# Heuristic: bin into quartiles for stratification purposes only
y_bins = pd.qcut(y_continuous, q=4, labels=False, duplicates="drop")
X = rng.normal(size=(2000, 5))
X_train, X_test, y_train, y_test = train_test_split(
X, y_continuous,
test_size=0.2,
stratify=y_bins,
random_state=42,
)This approach carries real caveats. Binning discards information about the exact target value. Results depend on where the bin boundaries fall, and quantile-based bins can produce duplicate edges when the target has many repeated or clustered values, which is why duplicates="drop" is included above. Very small bins can cause splitting errors for the same reason a very small class can in classification. Most importantly, stratifying a regression target this way is a practical heuristic, not something with the same statistical justification as stratifying a true categorical label — it may or may not improve generalization, and it must never use information about the target that would be unavailable at prediction time.
Multiclass and Multilabel Stratification
StratifiedKFold and stratify=y extend naturally to multiclass problems with more than two categories. The same constraints apply, only more acutely: if some classes are extremely rare, or the requested number of folds is too high relative to the smallest class, the split may fail outright or produce degenerate folds with very few examples of the rarest classes.
Multilabel classification, where each observation can belong to multiple labels simultaneously (a document tagged with several topics, for instance), introduces another layer of difficulty. Ordinary single-label stratification, applied to just one label at a time, does not guarantee that every combination of labels is proportionally represented across folds — with many possible label combinations, some rare combinations may end up entirely in one fold. Iterative stratification approaches for the multilabel case exist conceptually in the research literature as a way to approximately balance label combinations across folds, but they should be treated as a distinct, more specialized technique rather than something available directly as a built-in scikit-learn splitter.
Stratified Sampling vs Resampling for Class Imbalance
Stratified splitting and resampling techniques solve different problems, and conflating them is one of the more consequential mistakes in imbalanced-classification work.
Technique | Preserves or changes prevalence | Adds/removes observations | Affects splitting, training, or both | Main risk |
Stratified splitting | Preserves | Neither | Splitting only | False sense that the split "fixed" imbalance |
Random oversampling | Changes (increases minority share) | Adds (duplicates) | Training only | Overfitting to duplicated minority rows |
Random undersampling | Changes (increases minority share) | Removes majority rows | Training only | Loses potentially useful majority-class information |
SMOTE | Changes (increases minority share) | Adds synthetic rows | Training only | Synthetic points may not reflect real minority structure |
Class weighting | Preserves prevalence, changes loss | Neither | Training (loss function) | May not help if model can't separate classes at all |
Threshold adjustment | Preserves | Neither | Prediction/decision stage only | Chosen threshold may not generalize to new data |
Stratification is about splitting: making sure every partition reflects the same class mix. Oversampling, undersampling, SMOTE, class weighting, and threshold adjustment are about changing what the model sees during training (or how its output is interpreted), typically to help it pay more attention to a minority class. Combining them is standard practice: stratify the split first, then apply resampling only to the training portion.
The leakage rule here is strict and well documented: applying resampling before splitting, so that the entire dataset — including what will become test data — is resampled together, is a common pitfall, and it is equivalent to resampling both partitions with information the model shouldn't have (imbalanced-learn, n.d.). The resampling procedure might use information about samples in the dataset to generate or select synthetic or duplicated rows, meaning information from what should be held-out test rows can influence what appears in training, which is a textbook case of data leakage (imbalanced-learn, n.d.). SMOTE, undersampling, and oversampling must be fitted only within each training fold, never on the complete dataset before splitting or cross-validation.
Pipelines are the standard defense. Using a Pipeline (or imblearn's pipeline for resampling steps) ensures the resampling and preprocessing steps are refit correctly on each training fold during cross-validation, avoiding manual leakage entirely (imbalanced-learn, n.d.):
from imblearn.pipeline import Pipeline
from imblearn.over_sampling import SMOTE
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, cross_val_score
pipeline = Pipeline([
("smote", SMOTE(random_state=42)),
("classifier", RandomForestClassifier(random_state=42)),
])
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
scores = cross_val_score(pipeline, X, y, cv=cv, scoring="f1")Because SMOTE is inside the pipeline, it is refit independently on each training fold and never touches the corresponding validation fold.
Stratified Sampling Compared With Other Sampling Methods
Method | How it works | Primary purpose | Strength | Weakness | Common ML use |
Simple random sampling | Every observation has equal selection probability | General-purpose sampling | Simple, unbiased | Can under-represent rare groups | Baseline train-test split when classes are balanced |
Stratified sampling | Sample independently within predefined subgroups | Preserve subgroup proportions | Stable subgroup representation | Requires a known, reliable stratification variable | stratify=y, StratifiedKFold |
Cluster sampling | Randomly select some clusters, then observe all/some units within them | Reduce cost when a full frame is expensive | Efficient for geographically or logistically dispersed data | Higher variance if clusters are internally similar | Rarely used directly in ML splitting |
Systematic sampling | Select every k-th observation from an ordered list | Simple, evenly spread sampling | Easy to implement | Can introduce bias if data has periodic patterns | Occasional data auditing |
Convenience sampling | Use whatever data is easiest to obtain | Fast, low-cost data collection | Cheap and quick | High risk of unrepresentative bias | Early prototyping only, with caution |
Bootstrap sampling | Sample with replacement, same size as original | Estimate variability, build ensembles | Useful for confidence intervals, bagging | Some observations repeated, others omitted | Bagging, Random Forest, confidence intervals |
Oversampling | Duplicate or synthesize minority-class rows | Increase minority representation in training | Helps model attend to minority class | Can overfit to duplicated/synthetic patterns | Imbalanced classification training |
Undersampling | Remove majority-class rows | Reduce majority dominance in training | Faster training, simpler minority signal | Discards potentially useful data | Very large majority classes |
The most important distinction here is between stratified and cluster sampling, since the two are frequently confused. Stratified sampling selects observations from every stratum. Cluster sampling, by contrast, typically selects only some clusters at random and then observes units within those selected clusters — the unselected clusters contribute nothing to the sample at all. These are not interchangeable substitutes for each other.
Advantages of Stratified Sampling
Better subgroup coverage — every meaningful class or group is guaranteed some representation in each partition, assuming its size allows it.
Potentially lower sampling variance, when strata are genuinely related to the outcome being measured.
More stable model evaluation across repeated splits or folds, reducing noisy swings in reported metrics.
More interpretable subgroup analysis, since each partition reflects a known, consistent composition.
Protection against a class disappearing from a small validation or test partition.
Flexible allocation — proportional, equal, or disproportionate allocation can be chosen to fit the analytical goal.
Improved comparability between cross-validation folds, making fold-to-fold metric variation more attributable to the model rather than to sampling luck.
Each of these benefits is conditional: they hold most reliably when strata are large enough, the stratification variable is reliably measured, and the split design matches the deployment scenario.
Limitations and Disadvantages
Requires a known, reliable stratification variable available for every observation.
Depends entirely on data quality — a mislabeled class ruins the point of stratifying on it.
Adds design complexity compared with a plain random split.
Very small strata can make requested splits or fold counts impossible.
Rounding to whole-number sample sizes can push actual proportions slightly off target.
Candidate stratification variables can be correlated or overlapping, muddying which one to prioritize.
High-dimensional intersectional strata (for example, stratifying jointly by class, region, and age band at once) can produce strata too small to be useful.
Class priors can shift over time, so a split stratified on historical proportions may not reflect future deployment conditions.
Stratification cannot repair selection bias that already exists in how the source data was collected.
Development-time class proportions may not match production proportions, especially if deployment populations differ from the historical dataset.
Disproportionate sampling requires design weights to support valid population-level inference.
Preserving both class balance and group integrity simultaneously (as in StratifiedGroupKFold) is not always perfectly achievable.
Randomly rearranging time-series rows to satisfy stratification is fundamentally incompatible with preserving chronology.
A confidently "balanced-looking" stratified split can create misleading confidence when the underlying dataset itself is not representative of the real world.
Stratification is not a guarantee. It manages one specific source of evaluation noise — split-to-split variability in class proportions — without addressing data quality, collection bias, deployment drift, or leakage from other sources.
Common Mistakes and How to Avoid Them
Assuming stratification balances the classes. It preserves the existing ratio; use oversampling, undersampling, SMOTE, or class weighting if you actually need balance.
Using the test set to decide the stratification scheme. Any decision informed by test data compromises its ability to give an honest final estimate; make splitting decisions using training-accessible information only.
Stratifying on a leaked or post-outcome variable. A variable only known after the outcome occurred (e.g., "was refunded," when predicting fraud) will bias the split and the model.
Applying SMOTE before splitting. This resamples across what will become both partitions, leaking synthetic information into the evaluation set (imbalanced-learn, n.d.).
Ignoring patient, customer, or device groups. Ordinary stratification without grouping constraints can let the same entity appear in both training and test.
Randomly stratifying time-series data. This introduces temporal leakage; use TimeSeriesSplit or another chronology-respecting method instead.
Creating too many tiny intersectional strata. Stratifying on several variables at once can shrink strata below a workable size.
Using too many folds for a rare class. If n_splits exceeds the smallest class count, StratifiedKFold cannot construct valid folds.
Ignoring sample weights after disproportionate sampling. Skipping design weights biases any population-level estimate derived from the sample.
Believing a stratified split makes the source dataset representative. It only makes partitions consistent with each other, not consistent with the real world.
Evaluating only accuracy. On imbalanced classes, accuracy can look excellent while the model ignores the minority class entirely (Google, n.d.).
Failing to preserve a final untouched test set. Repeated peeking at the same test set erodes its value as an honest final check.
Reusing the same observations across nominally separate partitions. Duplicate or near-duplicate rows across splits inflate apparent performance.
Selecting strata only because they produce better-looking metrics. Choosing a stratification variable retroactively, based on which one flatters the results, undermines the whole point of a principled evaluation design.
Choosing a Good Stratification Variable
An effective stratification variable should be:
Known before sampling or splitting — available at split time, not derived from the outcome.
Reliable — consistently and accurately measured across the dataset.
Relevant to the outcome or evaluation objective.
Available for all relevant observations, without excessive missing data.
Defined consistently across the whole dataset and over time.
Divided into sufficiently large strata to support the intended split or fold count.
Ethically and legally appropriate to use for the purpose at hand.
Common candidates include the target class itself, geographic region, age band, acquisition source, device type, severity group, institution, or product category. Protected or sensitive attributes — such as those related to demographic characteristics — can be genuinely useful for fairness auditing, checking whether a model performs consistently across subgroups. However, collecting and using such attributes involves legal, ethical, privacy, and governance considerations that go beyond a purely technical decision, and this article does not offer legal advice on that point.
Real-World Machine-Learning Use Cases
Fraud detection. Strata: fraudulent vs. legitimate transactions. Stratification helps ensure evaluation folds always contain enough fraud examples to compute meaningful precision and recall. Danger: fraud patterns evolve, so historical class proportions may not match future fraud rates.
Medical diagnosis. Strata: diagnosed vs. not diagnosed. Helps stabilize sensitivity and specificity estimates in disease-screening models. Danger: prevalence in a screening dataset often differs sharply from prevalence in the deployment population, especially if the dataset was built from referred patients rather than the general public.
Manufacturing defects. Strata: defective vs. non-defective units. Helps preserve rare defect representation across CV folds. Danger: production lines change over time, so class priors drift.
Customer churn. Strata: churned vs. retained customers. Stabilizes retention-model evaluation. Danger: customer-level grouping matters — multiple records for one customer must stay in the same partition.
Cybersecurity intrusion detection. Strata: malicious vs. benign network events. Helps keep rare attack types represented across folds. Danger: attack patterns shift quickly, and stratified historical data may not reflect emerging threats.
Content moderation. Strata: policy-violating vs. compliant content. Preserves violation-rate consistency across evaluation partitions. Danger: content categories and thresholds change with evolving platform policy.
Credit-risk modeling. Strata: default vs. non-default. Helps stabilize default-rate evaluation, especially given how rare defaults typically are. Danger: macroeconomic conditions shift default rates, so historical strata may not match future risk.
Remote sensing. Strata: land-cover or crop-type classes. Helps rare land-cover types remain represented across training and validation regions. Danger: seasonal and regional variation can make one region's strata unrepresentative of another's.
Education analytics. Strata: at-risk vs. on-track student outcomes. Helps ensure minority "at-risk" cases are consistently represented. Danger: using outcome-adjacent variables (like final grades) as stratification variables can introduce leakage.
A/B testing or experimentation. Strata: pre-existing user segments (e.g., new vs. returning users). Helps balance segment composition between control and treatment groups. Danger: over-stratifying on many segments simultaneously in an experiment can shrink each cell below a size that supports reliable statistical conclusions.
Businesses evaluating AI in business contexts like these should treat stratification as one input into a broader evaluation strategy, not a stand-alone guarantee of model quality.
Best-Practice Workflow
Define the prediction unit (a single transaction, a single patient visit, a single customer).
Define the future deployment population and how it might differ from the historical dataset.
Check class and subgroup distributions before doing anything else.
Identify grouping and temporal constraints in the data.
Select a stratification variable that satisfies the criteria above.
Confirm each stratum has enough observations for the requested split or fold count.
Create a final untouched test set before any tuning begins.
Use appropriate cross-validation (stratified, grouped, or time-aware) on the remaining training data.
Put preprocessing and any resampling steps inside a pipeline.
Evaluate class-sensitive and subgroup-sensitive metrics, not accuracy alone.
Compare development and production distributions once the model is deployed.
Document random seeds, split logic, exclusions, weights, and known limitations.
A short pre-implementation checklist: Is the stratification variable known before split time? Are there grouped entities that need StratifiedGroupKFold? Is the data time-ordered, requiring TimeSeriesSplit instead? Is any stratum too small for the requested number of folds? Has a final test set been held out untouched?
When Should You Use Stratified Sampling?
Use stratification when:
Class proportions matter to evaluation quality.
Rare classes risk disappearing from small partitions.
Every meaningful subgroup needs guaranteed representation.
The stratification variable is known, reliable, and available at split time.
Grouping and chronology constraints don't override the need for class-proportion preservation.
Avoid or modify ordinary stratification when:
Observations are time-ordered — use TimeSeriesSplit instead.
Rows are grouped by entity — use StratifiedGroupKFold or GroupKFold.
Strata are extremely small — reconsider fold count or allocation strategy.
The stratification variable leaks the outcome.
Deployment prevalence is intentionally different from the source data's prevalence.
The source data is already badly biased in ways stratification cannot fix.
There's no meaningful stratification variable available.
Stratified sampling in machine learning is, at its core, a disciplined way of making sure your evaluation setup doesn't add its own noise on top of whatever signal your model already has to contend with — no more, and no less.
FAQ
What is stratified sampling in machine learning?
Stratified sampling in machine learning is a method of dividing a dataset into subgroups called strata — usually based on the class label — and sampling or splitting independently within each group. The goal is typically to preserve class proportions across training, validation, test, or cross-validation partitions, which stabilizes evaluation, especially when one class is rare relative to the others.
What is a stratum?
A stratum is one non-overlapping subgroup within a stratified sampling design. Strata are mutually exclusive (every observation belongs to exactly one) and collectively exhaustive (every observation belongs to some stratum). Ideally, observations within a stratum are similar with respect to the variable of interest, while different strata are meaningfully distinct from each other.
Why is stratified sampling used?
It's used to make training, validation, and test partitions — or cross-validation folds — reflect the same class distribution as the overall dataset. This reduces the risk of a rare class being under-represented or entirely absent from a partition, which would otherwise make evaluation metrics unstable or undefined.
Is stratified sampling the same as class balancing?
No. Stratified sampling preserves whatever class distribution already exists in the data; it does not change it. Class balancing techniques such as oversampling, undersampling, SMOTE, or class weighting actively change the effective class distribution the model trains on. These are complementary, not interchangeable, techniques.
What is the difference between random and stratified sampling?
Simple random sampling selects observations with equal probability from the whole dataset, without regard to subgroup membership, so class proportions in the resulting sample can vary by chance. Stratified sampling samples independently within predefined subgroups, which keeps subgroup proportions consistent between the sample and the full dataset.
What does stratify=y do in scikit-learn?
Passing stratify=y to train_test_split tells the function to split the data in a stratified fashion, using the array y as the class labels that define the strata (scikit-learn, n.d.). The resulting training and test sets will each preserve, as closely as integer counts allow, the same class proportions found in y.
When should a train-test split be stratified?
Stratify a train-test split whenever the target is categorical and class proportions matter to evaluation — which covers most classification problems, and especially imbalanced ones. It's less relevant for regression without binning, and inappropriate as a random reshuffle for time-ordered or grouped data.
Can stratification be used for regression?
Only as a heuristic. Since regression targets are continuous, a common approach bins the target into groups (often quantile-based) and stratifies on those bins. This can help preserve target-range coverage across splits, but it discards information, depends on bin boundary choices, and is not statistically justified in the same way stratifying a true categorical label is.
What happens when a class has very few observations?
If a class has fewer members than the number of partitions or folds it needs to be spread across, train_test_split or StratifiedKFold will raise an error rather than silently creating an empty or degenerate group. In such cases, reduce the number of folds, combine rare classes if appropriate, or collect more data for that class.
What is StratifiedKFold?
StratifiedKFold is a scikit-learn cross-validation splitter that returns folds made by preserving the percentage of samples for each class across every fold, in a binary or multiclass classification setting (scikit-learn, n.d.). It's the standard way to get stable, comparable performance estimates across cross-validation folds for classification problems.
Is stratified cross-validation appropriate for time-series data?
Generally, no. Standard stratified cross-validation shuffles or reorders data without respect to time, which can let future information leak into training folds. Time-ordered data should use a chronology-respecting method such as TimeSeriesSplit, which always trains on past data and evaluates on subsequent data (scikit-learn, n.d.).
How does stratification work with grouped data?
StratifiedGroupKFold tries to preserve class proportions across folds while guaranteeing that all rows belonging to the same group (such as the same patient or customer) stay within a single fold, preventing that entity from appearing in both training and validation (scikit-learn, n.d.). Because both goals can conflict, perfect stratification isn't always achievable when groups must stay intact.
Is SMOTE a form of stratified sampling?
No. SMOTE (Synthetic Minority Oversampling Technique) generates synthetic minority-class observations to change the training set's class balance. Stratified sampling doesn't create or remove observations at all — it only controls how existing observations are divided across partitions. They address different stages of a machine-learning workflow and are often used together, with stratification applied at splitting time and SMOTE applied only within the training fold.
Can stratification prevent data leakage?
Stratification by class label does not by itself prevent leakage, and can even coexist with leakage if grouping or temporal structure is ignored. Preventing leakage requires additional measures: StratifiedGroupKFold for grouped entities, TimeSeriesSplit for time-ordered data, and fitting any resampling or preprocessing steps only within training folds, ideally inside a pipeline (imbalanced-learn, n.d.).
What are the main disadvantages of stratified sampling?
It requires a known, reliable stratification variable; it can fail or behave poorly with very small strata; it adds design complexity; it doesn't correct bias already present in the source dataset; and it can create a false sense that a dataset is representative when it only makes internal partitions consistent with each other. It also may conflict with grouping and temporal constraints if applied without care.
Key Takeaways
Stratified sampling divides data into meaningful, non-overlapping strata and samples independently within each one — most often by class label in machine learning.
It preserves existing class proportions; it does not balance classes.
Statistical stratified sampling, stratified dataset splitting, and stratified cross-validation are related but distinct concepts, and conflating them leads to overconfident claims about representativeness.
stratify=y and StratifiedKFold stabilize evaluation but do not fix a biased source dataset.
Grouped data needs StratifiedGroupKFold or GroupKFold to avoid entity-level leakage.
Time-ordered data needs TimeSeriesSplit, not random stratification, to avoid temporal leakage.
Resampling techniques like SMOTE must be fitted only on training folds — never before splitting — to avoid leakage.
Disproportionate sampling requires design weights for valid population-level inference.
Representativeness has limits: stratification cannot repair a dataset that was collection-biased from the start.
Actionable Next Steps
Inspect your target's class counts and proportions before choosing a split strategy.
Check for grouped entities (patients, customers, devices) and temporal ordering in your data.
Select stratify=y, StratifiedGroupKFold, or TimeSeriesSplit based on what you find.
Confirm no class or group is too small for your intended number of partitions or folds.
Implement the split with a fixed random_state for reproducibility.
Wrap any resampling or preprocessing steps inside a pipeline to prevent leakage.
Evaluate class-sensitive metrics (precision, recall, F1, PR-AUC) alongside — not instead of — accuracy.
Document your split logic, seeds, exclusions, and any weighting decisions for future reference.
Glossary
Allocation: The strategy used to decide how many observations are drawn from each stratum (e.g., proportional, equal, Neyman).
Class imbalance: A situation where one class in a classification problem has substantially more examples than another.
Cluster sampling: A sampling method that randomly selects entire clusters and then observes units within those selected clusters only.
Cross-validation: A model-evaluation technique that repeatedly splits data into training and validation folds to produce a more reliable performance estimate.
Data leakage: Information from outside the legitimate training scope — such as future data, test data, or grouped entities — influencing model training or evaluation.
Design weight: A correction factor applied to disproportionately sampled strata to enable valid population-level estimates.
Disproportionate sampling: Deliberately sampling some strata at higher or lower rates than their true population share.
Imbalanced dataset: A dataset where class proportions are far from equal.
Minority class: The less frequent class in an imbalanced classification problem.
Neyman allocation: An allocation strategy that samples more heavily from strata that are both larger and more variable, minimizing overall sampling variance for a fixed budget.
Oversampling: Increasing the representation of a minority class, typically by duplicating or synthetically generating new examples.
Population: The full set of items or observations a study or model is ultimately concerned with.
Proportional allocation: Allocating sample size to each stratum in proportion to its share of the population.
Random sampling: Selecting observations so that each has an equal chance of being included, without regard to subgroup membership.
Sampling frame: The list or dataset actually available for drawing a sample, which may differ from the true population.
Sampling variance: The variability in an estimate caused by which particular observations happened to be sampled.
Simple random sampling: Sampling with no stratification, clustering, or systematic structure — every observation has equal selection probability.
Strata: The plural of stratum; the non-overlapping subgroups used in stratified sampling.
Stratification variable: The variable used to divide a population or dataset into strata.
Stratified cross-validation: Cross-validation where each fold is constructed to preserve class proportions as closely as possible.
Stratified sampling: A sampling method that divides a population into strata and samples independently within each one.
Stratum: A single non-overlapping subgroup within a stratified sampling design.
Undersampling: Reducing the representation of a majority class, typically by removing some of its observations.
Sources & References
NIST Computer Security Resource Center. "stratified sampling — Glossary." National Institute of Standards and Technology. n.d. Accessed July 22, 2026. https://csrc.nist.gov/glossary/term/stratified_sampling
scikit-learn developers. "StratifiedKFold — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedKFold.html
scikit-learn developers. "train_test_split — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html
scikit-learn developers. "StratifiedShuffleSplit — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedShuffleSplit.html
scikit-learn developers. "StratifiedGroupKFold — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedGroupKFold.html
scikit-learn developers. "3.1. Cross-validation: evaluating estimator performance — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/cross_validation.html
scikit-learn developers. "TimeSeriesSplit — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.TimeSeriesSplit.html
scikit-learn developers. "Common pitfalls and recommended practices — scikit-learn documentation." scikit-learn.org. n.d. Accessed July 22, 2026. https://scikit-learn.org/stable/common_pitfalls.html
imbalanced-learn developers. "9. Common pitfalls and recommended practices — imbalanced-learn documentation." imbalanced-learn.org. n.d. Accessed July 22, 2026. https://imbalanced-learn.org/stable/common_pitfalls.html
Google for Developers. "Datasets: Class-imbalanced datasets — Machine Learning Crash Course." developers.google.com. n.d. Accessed July 22, 2026. https://developers.google.com/machine-learning/crash-course/overfitting/imbalanced-datasets
Google for Developers. "FAQPage structured data — Google Search Central." developers.google.com. n.d. Accessed July 22, 2026. https://developers.google.com/search/docs/appearance/structured-data/faqpage
Schema.org. "BlogPosting." schema.org. n.d. Accessed July 22, 2026. https://schema.org/BlogPosting
Schema.org. "FAQPage." schema.org. n.d. Accessed July 22, 2026. https://schema.org/FAQPage
Southern, Matt G. "Google Drops FAQ Rich Results From Search." Search Engine Journal. May 10, 2026. Accessed July 22, 2026. https://www.searchenginejournal.com/google-drops-faq-rich-results-from-search/574429/
