ML Notes
Master data splitting strategies including train-test split, stratification, time-series splits, cross-validation, and preventing data leakage in ML pipelines.
The train-test split is perhaps the most fundamental concept in applied machine learning — and one of the most frequently misunderstood. The core principle is simple: you must evaluate your model on data it has never seen during training. If you test on training data, you are measuring memorization, not generalization. This lesson covers proper splitting strategies, common pitfalls, and advanced techniques for different data types.
Why Splitting Data Is Essential
Imagine studying for an exam by memorizing the exact questions and answers. You would score perfectly on those specific questions but fail any new question you have not memorized. Machine learning models can do the same thing — perfectly fit training data while being useless on new data. This is overfitting, and the train-test split is our primary defense against it.
By holding out a portion of data that the model never sees during training, we get an honest estimate of how the model will perform in the real world. This held-out test set simulates the future data the model will encounter in production.
Basic Train-Test Split
The simplest approach splits data into two groups: a larger training set (typically 70-80 percent) for fitting the model and a smaller test set (20-30 percent) for evaluation.
The random_state parameter ensures reproducibility — the same split every time you run the code. Without it, each run produces a different split, making debugging and comparison impossible.
Stratified Splitting for Classification
When classes are imbalanced, random splitting might produce a test set with very few minority examples — or even zero examples of a rare class. Stratification preserves the class ratio in both splits.
# Stratified split — preserves class proportions
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
# Verify proportions are preserved
train_ratio = np.bincount(y_train) / len(y_train)
test_ratio = np.bincount(y_test) / len(y_test)
print(f"Original class ratio: {np.bincount(y) / len(y)}")
print(f"Train class ratio: {train_ratio}")
print(f"Test class ratio: {test_ratio}")
# All three should be very similarAlways use stratify=y for classification problems, especially when any class represents less than 20 percent of the data.
Three-Way Split: Train, Validation, and Test
In practice, you often need THREE sets: training (fit the model), validation (tune hyperparameters and make modeling decisions), and test (final unbiased evaluation, touched exactly once).
# First split: separate test set (final evaluation)
X_temp, X_test, y_temp, y_test = train_test_split(
X, y, test_size=0.15, random_state=42, stratify=y
)
# Second split: separate validation from training
X_train, X_val, y_train, y_val = train_test_split(
X_temp, y_temp, test_size=0.18, random_state=42, stratify=y_temp
)
# Results in approximately 70% train, 15% validation, 15% test
print(f"Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")The validation set is where you compare different models, tune hyperparameters, and decide which features to use. The test set is sacred — use it only once at the very end to report your model's expected real-world performance.
Time-Series Splitting
For temporal data, random splitting violates causality — you would be training on future data to predict the past. Instead, use chronological splits where the training set comes before the test set in time.
Data Leakage: The Silent Killer
Data leakage occurs when information from the test set inadvertently influences training. This makes your evaluation metrics unrealistically optimistic — the model appears to perform well in testing but fails in production.
Common leakage sources include fitting preprocessing (scaling, encoding, imputation) on the full dataset before splitting, including features that are proxies for the target variable, and not respecting temporal ordering in time-series data.
# WRONG: Leakage! Scaler sees test data statistics
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X) # Fits on ALL data including test
X_train, X_test = train_test_split(X_scaled, ...)
# CORRECT: Fit only on training data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train) # Fit on train ONLY
X_test_scaled = scaler.transform(X_test) # Transform test with train statsCross-Validation: Better Than a Single Split
A single train-test split can give misleading results depending on which examples land in which set. Cross-validation provides a more robust estimate by performing multiple splits:
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier
model = RandomForestClassifier(n_estimators=100, random_state=42)
# 5-fold cross-validation: train on 4 folds, test on 1, rotate 5 times
scores = cross_val_score(model, X, y, cv=5, scoring='accuracy')
print(f"CV Accuracy: {scores.mean():.4f} (+/- {scores.std():.4f})")
# The mean gives expected performance, std gives uncertaintyChoosing the Right Split Ratio
The ideal split ratio depends on your dataset size. With 100,000+ samples, 80/20 or even 90/10 works well. With 1,000 samples, consider using cross-validation instead of a fixed split to maximize both training data and evaluation reliability. With very small datasets (under 500 samples), use leave-one-out cross-validation or repeated stratified k-fold.
Real-World Applications and Industry Patterns
In industry, train-test split is applied differently across domains. E-commerce companies engineer features around customer behavior patterns (session duration, cart abandonment rate, browse-to-buy ratio). Financial institutions create risk indicators from transaction patterns (velocity of spending, unusual amounts, geographic anomalies). Healthcare systems derive clinical features from raw measurements (body mass index from height and weight, estimated glomerular filtration rate from creatinine and age). Each domain has established feature engineering patterns that experienced practitioners know to apply immediately, giving them a significant head start on any new project in their domain.
The most impactful feature engineering often comes from understanding the business context deeply. Talk to domain experts, read industry literature, and study how the target variable is influenced by real-world factors. Technical skill in pandas and numpy is necessary but not sufficient — the creative insight about which features to engineer comes from understanding the problem domain at a fundamental level.
Split Size Guidelines by Dataset Size
For datasets with more than 100,000 samples, a simple 80/20 split provides sufficient test examples for reliable evaluation. For 10,000-100,000 samples, consider 70/15/15 (train/validation/test) when you need hyperparameter tuning. For 1,000-10,000 samples, cross-validation becomes essential since individual splits have high variance. For fewer than 1,000 samples, use repeated stratified k-fold cross-validation (5 or 10 folds, repeated 3-5 times) to get stable performance estimates without wasting precious training data on a large held-out set. These guidelines balance the competing needs of having enough data to train well and enough held-out data to evaluate reliably.
Key Takeaways
Train-test splitting ensures honest evaluation of model generalization. Always stratify for classification, respect temporal ordering for time-series, and use three-way splits when tuning hyperparameters. Guard against data leakage by fitting all preprocessing exclusively on training data. When dataset size permits, cross-validation provides more reliable performance estimates than a single split. These practices are non-negotiable foundations of rigorous machine learning — skip them and your real-world model performance will disappoint.
Exam Focus
Revise definitions, diagrams, examples, and short-answer points for Train-Test Split.
Interview Use
Prepare one clear explanation, one practical example, and one common mistake for this Machine Learning topic.
Search Terms
machine-learning, machine learning, machine, learning, data, preprocessing, train, test
Related Machine Learning Topics