07 - Data Science, AI & Machine Learning

How to avoid data leakage when training a machine learning model in Python

Learn how data leakage happens in machine learning workflows, how to fix it with scikit-learn pipelines, and how to validate the repair.

4 min read / An Introduction to Statistical Learning with Applications in Python (ISLP)

How to avoid data leakage when training a machine learning model in Python

Data leakage is one of the easiest machine learning bugs to miss because the model appears to be working better than expected. The validation score looks strong, the notebook runs cleanly, and the charts make the result feel trustworthy. The problem is that information from the test set, future rows, or the target variable has slipped into the training process. A model trained this way is not learning a general pattern. It is getting help from information it would never have in production.

This is a natural topic to pair with An Introduction to Statistical Learning with Applications in Python (ISLP) because statistical learning is not just about choosing a model. It is about estimating how well a model will behave on unseen data. If the validation setup is contaminated, every downstream decision becomes weaker: feature selection, model comparison, hyperparameter tuning, and the business decision that follows the score. The practical lesson is simple but important: split first, fit every learned transformation only on the training data, and evaluate only on data the model has not influenced.

What leakage looks like

A common mistake is scaling or imputing the full dataset before the train-test split. The scaler calculates means and standard deviations using rows that should be held out for evaluation. That seems harmless, but it lets the training process see distribution information from the test set. The same mistake happens with target encoding, feature selection, missing-value imputation, outlier removal, PCA, and any transformation that learns from data.

# Leaky: the scaler sees the full dataset before the split.
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

X_scaled = StandardScaler().fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=42, stratify=y
)

model = LogisticRegression(max_iter=1000)
model.fit(X_train, y_train)
pred = model.predict(X_test)
print(accuracy_score(y_test, pred))

The bug is not that scaling is bad. Scaling is often necessary. The bug is where the scaling is fitted. fit_transform estimates parameters from the data. If it runs before the split, the test set has already influenced those parameters. A clean workflow treats the test set as locked away until the final evaluation.

How to rectify it

Use a pipeline so every learned preprocessing step is fitted only inside the training fold. The pipeline keeps the order honest: split the data, fit preprocessing on the training rows, transform the test rows using training-only parameters, then score the model. This is the safest default because it makes the correct workflow harder to accidentally break.

from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

pipe = Pipeline([
    ("scale", StandardScaler()),
    ("model", LogisticRegression(max_iter=1000)),
])

pipe.fit(X_train, y_train)
pred = pipe.predict(X_test)
print(accuracy_score(y_test, pred))

For mixed numeric and categorical data, use ColumnTransformer inside the same pipeline. This prevents a second common leak: fitting encoders or imputers outside the model workflow and accidentally using the whole dataset to learn category lists, medians, or missing-value patterns.

from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder, StandardScaler

numeric_features = ["age", "income", "visits"]
categorical_features = ["region", "plan"]

preprocess = ColumnTransformer([
    ("num", Pipeline([
        ("impute", SimpleImputer(strategy="median")),
        ("scale", StandardScaler()),
    ]), numeric_features),
    ("cat", Pipeline([
        ("impute", SimpleImputer(strategy="most_frequent")),
        ("onehot", OneHotEncoder(handle_unknown="ignore")),
    ]), categorical_features),
])

pipe = Pipeline([
    ("preprocess", preprocess),
    ("model", LogisticRegression(max_iter=1000)),
])

How to test the workflow

Cross-validation should also wrap the full pipeline, not only the model. If preprocessing happens before cross_val_score, every fold can leak information into every other fold. When the pipeline is passed into cross-validation, scikit-learn refits preprocessing separately inside each training fold.

from sklearn.model_selection import cross_val_score, StratifiedKFold

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
scores = cross_val_score(pipe, X, y, cv=cv, scoring="accuracy")
print(scores.mean(), scores.std())

Also scan for feature leakage. Any column created after the outcome is known can make the model look magical. Examples include days_until_cancelled when predicting cancellation, final invoice totals when predicting purchase likelihood, or a status field updated after the event. A practical review question is: would this feature exist, with this exact value, at the moment the prediction is made? If the answer is no, remove it or rebuild it from information available at prediction time.

What to ask next

Once the leakage is fixed, the model score may drop. That is not a failure. It is a more honest estimate. You can open An Introduction to Statistical Learning with Applications in Python (ISLP) in Vidyora and ask how train-test splits, cross-validation, bias-variance tradeoffs, and model assessment connect to your dataset. The article solves the immediate bug; the book helps you understand why the fix works.

machine learningdata leakagepythonmodel validation07 - Data Science, AI & Machine Learning