Undersampling before or after Train/Test Split
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
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.
Related reading
- Understand Op Registration and Kernel Linking in TensorFlow
- Understanding cluster state update
- Understanding concept of Gaussian Mixture Models
- Understanding Cross Entropy `Loss`
- Understanding DynamicTreeCut algorithm for cutting a dendrogram
- Understanding multi-label classifier using confusion matrix
- Understanding CTC loss for speech recognition in Keras
- Understanding FeatureHasher, collisions and vector size trade-off
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.