ai_ml / ml foundations / 02_train_val_test_split.md

Train / validation / test splits

6 interview angles 5 min read source

Train / validation / test splits

Splitting looks trivial and is where most “amazing offline, useless in production” stories begin. The rule underneath everything here: the test set must resemble the situation the model will actually face.

The three sets, and what each is for

Set Used for How often you may look
Train fitting parameters constantly
Validation choosing hyperparameters, early stopping, model selection every experiment
Test one final unbiased estimate as close to once as you can manage

Typical split is 60/20/20, or 80/10/10 when data is plentiful. With very large datasets the validation and test sets don’t need to grow proportionally — you need enough samples for a tight confidence interval, not a fixed percentage.

Every decision you make using the validation set leaks a little information into it. After a hundred experiments, validation performance is optimistic. That’s the entire reason a separate test set exists, and why looking at it repeatedly destroys its value.

Random splits are wrong more often than you’d think

from sklearn.model_selection import train_test_split

# Fine only when rows are i.i.d.
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42)

Random splitting assumes rows are independent and identically distributed. Reach for something else when:

The data has time order. Predicting the future from the past means the test set must come after the train set. A random split lets the model learn from Wednesday to predict Tuesday, which it will never get to do in production.

from sklearn.model_selection import TimeSeriesSplit

for train_idx, val_idx in TimeSeriesSplit(n_splits=5).split(X):
    ...   # each fold trains on the past, validates on the future

Add an embargo gap between train and validation when the label itself takes time to materialise — if you’re predicting 30-day churn, the last 30 days of training data have labels that depend on the validation window.

Rows are grouped. Multiple rows per user, per patient, per document. If the same user appears in train and test, the model can memorise the user rather than learn the pattern.

from sklearn.model_selection import GroupKFold, StratifiedGroupKFold

for tr, va in GroupKFold(n_splits=5).split(X, y, groups=user_ids):
    ...

Classes are imbalanced. With 1% positives and a small test set, a random split can leave you with a handful of positives and a metric that’s mostly noise.

X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)

There are near-duplicates. Scraped datasets, product catalogues and document corpora are full of them. A duplicate straddling the split is memorisation dressed up as generalisation. Deduplicate — by hash for exact matches, by embedding similarity or MinHash for near ones — before splitting.

Cross-validation

When data is limited, a single split wastes it and gives a high-variance estimate. K-fold trains K times, each time holding out a different fold.

from sklearn.model_selection import cross_val_score, StratifiedKFold

scores = cross_val_score(
    model, X, y,
    cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=42),
    scoring="average_precision",
)
print(f"{scores.mean():.3f} +/- {scores.std():.3f}")

Report the standard deviation, not just the mean. If fold-to-fold variance is large, your estimate is fragile and any comparison between two models at that scale is meaningless.

Variant Use
KFold plain i.i.d. data
StratifiedKFold classification, preserves class balance per fold
GroupKFold grouped rows
StratifiedGroupKFold both at once
TimeSeriesSplit temporal data
LeaveOneOut tiny datasets; expensive and high variance

Cross-validation is standard for classical ML. For deep learning it’s usually too expensive — a single held-out validation set is the norm.

Nested cross-validation

If you tune hyperparameters with cross-validation and then report that same cross-validated score, it’s optimistic — you selected on it. Nested CV separates the two loops:

from sklearn.model_selection import GridSearchCV, cross_val_score

inner = GridSearchCV(model, param_grid, cv=3, scoring="average_precision")
scores = cross_val_score(inner, X, y, cv=5)   # outer loop = honest estimate

Expensive, and correspondingly rare in practice — but knowing why it exists is the interview point.

Preprocessing belongs inside the split

The most common leakage bug in real code:

# WRONG - the scaler saw the test set's mean and variance
X_scaled = StandardScaler().fit_transform(X)
X_tr, X_te = train_test_split(X_scaled, ...)

# RIGHT - fit on train only, apply to test
X_tr, X_te, y_tr, y_te = train_test_split(X, y, ...)
scaler = StandardScaler().fit(X_tr)
X_tr, X_te = scaler.transform(X_tr), scaler.transform(X_te)

The same applies to imputation, target encoding, feature selection, PCA and oversampling. Use a Pipeline so it’s structurally impossible to get wrong:

from sklearn.pipeline import Pipeline

pipe = Pipeline([
    ("scale", StandardScaler()),
    ("model", LogisticRegression()),
])
cross_val_score(pipe, X, y, cv=5)   # scaler refits inside every fold

This is the single strongest argument for using Pipeline at all, and it’s worth saying so explicitly in an interview.

SMOTE and other resampling must also live inside the fold, and must only touch training data — never the validation fold. imblearn.pipeline.Pipeline handles this; the sklearn one does not.

When your test set stops being representative

Even a correct split decays. Production distribution drifts away from the frozen test set, so offline metrics stay flat while real performance degrades. Two habits mitigate it:

  • Maintain a rolling recent-data holdout alongside the fixed test set.
  • Track online metrics, not just offline ones. See ../15_mlops_llmops/.

Interview angle

  • “How would you split data for a churn model?” — temporally, with a gap. Train on data up to T, validate on T to T+k, test after that, and make sure the label window doesn’t overlap the training window. Also group by user so one customer can’t span sets.
  • “Why can’t I just use train_test_split everywhere?” — it assumes i.i.d. rows. Time order, grouping, class imbalance and near-duplicates each break that assumption in a way that inflates offline metrics.
  • “You scaled features before splitting. What’s wrong?” — the scaler’s mean and variance were computed using test data, so information leaked into training. Fit transforms on train only, ideally via a Pipeline so it happens per fold automatically.
  • “Why keep a test set separate from validation?” — every hyperparameter choice made on validation biases it upward. The test set is your only unbiased estimate, and it stops being unbiased the moment you start iterating against it.
  • “Model scores 0.95 offline and 0.6 in production. Where do you look?” — leakage first (a feature unavailable or differently distributed at prediction time), then split methodology (temporal or group leakage, duplicates), then genuine distribution drift, then training/serving skew in the feature computation itself.
  • “When would you use cross-validation over a single split?” — limited data, or when you need a variance estimate on the metric. It’s standard in classical ML and usually too expensive for deep learning.