Undersampling before or after Train/Test Split
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
Undersampling should be applied after the train/test split, not before it. More precisely, it should be applied only to the training data, and if you are using cross-validation, it should happen separately inside each training fold.
The reason is data leakage. If you undersample before splitting, information from the full dataset influences what ends up in training and test data, and the test distribution stops representing the real problem you actually care about.
Why Splitting Comes First
The test set is supposed to simulate unseen production data. That means it should preserve the natural class imbalance of the original dataset unless you have a very unusual evaluation goal.
If you undersample first, you are changing the population before creating the evaluation set. That makes the test set easier and less realistic, which can produce misleading performance numbers.
So the safe order is:
- split the data
- undersample only the training portion
- train the model on the resampled training data
- evaluate on the untouched test data
Correct Workflow with scikit-learn and imbalanced-learn
Here is a simple example using RandomUnderSampler.
The training set becomes balanced, but the test set keeps the original imbalance. That is exactly what you want for honest evaluation.
Why Undersampling Before the Split Is Wrong
This is the pattern to avoid:
This leaks information from the full dataset into the resampling step and produces a test set whose class balance no longer matches the original problem.
The result is often an overly optimistic estimate of model quality.
Cross-Validation Needs the Same Discipline
The same rule applies when using cross-validation. Resampling should happen inside each training fold, not once on the entire dataset before the folds are created.
A good way to do that is with an imblearn pipeline.
This ensures the sampler sees only the training portion of each fold.
Keep Evaluation Metrics Appropriate
Because the test set remains imbalanced, plain accuracy is often a weak metric. In class-imbalance problems, metrics such as precision, recall, F1, PR AUC, or ROC AUC are usually more informative.
That is not a reason to rebalance the test set. It is a reason to choose metrics that reflect the real problem.
Common Pitfalls
- Undersampling before the train/test split and leaking information into evaluation.
- Balancing the test set and then believing the measured performance reflects production behavior.
- Applying the sampler once before cross-validation instead of inside each fold.
- Using only accuracy on a strongly imbalanced problem.
- Forgetting stratification during the original split and accidentally distorting the class ratios even before resampling.
Summary
- Undersample after the train/test split, not before.
- Apply undersampling only to the training data.
- Keep the test set untouched so it reflects the real class distribution.
- In cross-validation, resample inside each training fold through a pipeline.
- Use evaluation metrics suited to imbalanced classification rather than trying to make the test set artificially balanced.

