Avoiding Data Leakage in Machine Learning Pipelines
Data leakage happens when information that would not be available at prediction time influences model training. The result is an evaluation score that looks excellent in development and collapses after deployment.
Leakage is often subtle because the model code itself can be correct. The mistake lives in how datasets, features, preprocessing, and time boundaries are constructed.
Split before learning from the data
A classic mistake is standardizing the full dataset and splitting afterward.
For a numeric feature, standardization estimates a mean and standard deviation. If those statistics include validation or test rows, the training pipeline has learned something about held-out data.
The correct order is:
- split raw examples into train, validation, and test sets;
- fit preprocessing only on the training set;
- apply the fitted transformation to validation and test sets;
- fit the model on transformed training data;
- use validation for model selection;
- evaluate once on the untouched test set.
This rule applies to more than scaling. Imputation values, vocabulary selection, dimensionality reduction, target encoding, feature selection, and learned embeddings can all leak information if fitted globally.
Keep preprocessing and the model in one pipeline
A practical way to avoid accidental refitting is to represent preprocessing as part of the model pipeline.
Conceptually:
training rows
-> fit imputer
-> fit scaler
-> fit feature selector
-> fit model
validation rows
-> use fitted imputer
-> use fitted scaler
-> use fitted selector
-> predictLibraries such as scikit-learn provide pipeline abstractions that automate this fit/transform discipline. When using cross-validation, ensure the preprocessing steps are inside the pipeline so each fold learns them only from that fold’s training partition.
Watch for target leakage
A feature can be present in the table and still be unavailable at real prediction time.
Suppose a model predicts whether a support ticket will breach its service-level target. Features such as resolution_time_minutes or closed_by_team may be highly predictive—but they are only known after the ticket is resolved.
Ask a timestamped question for every feature:
At the exact moment the model must produce a prediction, is this value already known?
If not, the feature leaks the future.
Respect time when the real task is temporal
Random splitting is inappropriate when production predicts future events from past events.
For a demand forecast, train on earlier dates and validate on later dates:
Jan -------- Jun | Jul -- Aug | Sep
training validation testThis better represents deployment, where future observations are unavailable during training.
Be careful with aggregate features. A “customer average order value” computed using the entire history can include purchases that occur after the row being predicted. Build such features using only data available up to that row’s prediction timestamp.
Keep related entities in the same split
Random row splitting can leak identity when multiple rows belong to one person, device, patient, account, document, or session.
Imagine classifying images when near-duplicate photographs from the same source appear in both training and test sets. The model may partly memorize source-specific patterns rather than generalize to new sources.
Use group-aware splitting when the real task requires generalization to unseen groups. The group key should reflect the dependency that would otherwise cross partitions.
Deduplicate before trusting the score
Exact and near duplicates can inflate evaluation metrics if copies land on both sides of a split.
Deduplication itself must be designed carefully. If the deduplication algorithm learns thresholds or representations from data, fit those decisions without using the held-out labels or future information.
At minimum, investigate repeated identifiers, identical rows, repeated text, and duplicated media hashes before finalizing the split.
Do not tune on the test set
The test set is not a second validation set.
If you repeatedly inspect test performance and change features, hyperparameters, thresholds, or preprocessing based on that result, test data has influenced model development. The final score becomes optimistic.
Use validation data or nested cross-validation for iteration. Reserve the test set for a small number of final evaluations after choices are frozen.
Leakage can happen outside the notebook
Production feature pipelines can accidentally differ from training pipelines.
Examples include:
- training reads a corrected historical table unavailable online;
- production computes a feature before an event, while training computes it after the event;
- a warehouse join includes records that were written later;
- labels and features use different timestamp semantics.
Document a feature availability time, not only an event time. A value generated at 10:05 cannot safely predict an event at 10:00 even if the underlying business event occurred earlier.
Use suspiciously strong performance as a debugging signal
Unexpectedly high accuracy should trigger investigation, especially on noisy real-world problems.
Useful checks include:
- train a simple baseline and compare the gap;
- inspect the most important features for post-outcome information;
- shuffle labels and confirm performance collapses toward chance where appropriate;
- recompute preprocessing independently inside each split or fold;
- test a stricter temporal or group split;
- examine duplicate and near-duplicate samples.
A strong metric is good only when the evaluation boundary matches deployment reality.
Common pitfalls
Fitting an imputer before cross-validation
Even an innocent median uses information from held-out folds. Put the imputer inside the cross-validation pipeline.
Encoding categories using all target labels
Target encoding is especially leakage-prone. Compute encodings using training-only information, with an out-of-fold strategy where appropriate.
Randomly splitting event histories
Rows from the future can teach the model patterns that would not have existed at the historical prediction point.
Selecting features after reading test results
Feature selection is model development. Once test results guide that choice, the test set is no longer independent.
Assuming anonymized IDs are harmless
An identifier can let a model memorize entities even when it contains no obvious personal meaning. Remove or group by identifiers according to the intended generalization task.
Design the split from deployment backward
Before choosing a model, write down who or what the model will predict, when the prediction happens, which information exists at that moment, and whether production must generalize to new time periods or new entities.
Then design data partitions to reproduce those constraints.
Preventing leakage is less about one library function than about preserving causality. The evaluation set must remain information the training process could not have known. When that boundary is real, model metrics become useful evidence rather than optimistic fiction.