Feature Selection
Machine Learning
Data Preprocessing
Train-Test Split
Model Evaluation

Should Feature Selection be done before Train-Test Split or after?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Feature selection is a critical step in the machine learning process, often determining the success and efficiency of a model. One frequent question in this context is whether feature selection should be performed before or after a train-test split. Understanding the implications of this decision is essential for building robust models.

Understanding Train-Test Split

Train-test splitting is a foundational concept in machine learning, essential for evaluating model performance. By splitting the dataset into training and test sets, we ensure that the model is tested on unseen data. This process helps to simulate real-world scenarios where the model has to make predictions on data it has not encountered before.

Importance of Feature Selection

Feature selection aims to identify the most relevant features for the task at hand, reducing dimensionality, improving model performance, and increasing interpretability. Having fewer features reduces the risk of overfitting and potentially improves model generalization.

Arguments for Feature Selection Before Train-Test Split

  1. Bias Introduction: Performing feature selection before the train-test split may introduce bias into the model. It can lead to overly optimistic performance results since the feature selection could be based on information from the entire dataset.
  2. Information Leakage: Selecting features based on the full dataset (before splitting) can leak information from the test set into the training process. This means the test set is no longer independent, compromising the evaluation's validity.
  3. Optimal Feature Set Identification: Early feature selection might rely on biassed feature importance measures influenced by the entirety of the data, which does not generalize well to truly unseen data.

Arguments for Feature Selection After Train-Test Split

  1. Prevention of Information Leakage: Conducting feature selection after the train-test split ensures that the test set remains unseen and unbiased. This maintains the integrity of the test set's role as a proxy for new, unseen data.
  2. Realistic Model Evaluation: By selecting features using only the training data, you mimic the real-world scenario where predictions are made using a model trained on available data, without peeking into future or external data.
  3. Robustness to Overfitting: Feature selection derived from the training data alone reduces the risk of tailoring the feature set too closely to specific dataset idiosyncrasies, thus enhancing model generalization.

Examples and Practical Guidelines

Example Scenario

Consider a dataset with 30 features. If feature selection is conducted on the entire dataset (before splitting), the process might choose features that best explain the data's variance. However, some of these features may align coincidentally with patterns in the training data, not necessarily representing trends generalizable to new data.

Conversely, by splitting the data first and then performing feature selection on the training set, the selected features are more likely to capture the true underlying relationships.

Practical Implementation

  1. Split the dataset: Consistently begin with a train-test split to ensure an unbiased evaluation.
  2. Perform feature selection: Use techniques like Recursive Feature Elimination (RFE) or Lasso regularization on the training data.
  3. Evaluate model: Use the test set, untouched during feature selection, to gauge model performance.

Summary Table

ApproachAdvantagesDisadvantages
Feature Selection Before Split- Utilizes full dataset statistics - Potentially finds globally optimal features- Information leakage - Unrealistic evaluation
Feature Selection After Split- No information leakage - Realistic evaluation - Reduces overfitting risk- May miss out on global trends

Additional Considerations

Cross-Validation

Cross-validation can further help in understanding the performance of a model. It's recommended to incorporate feature selection within each cross-validation fold rather than applying it globally before cross-validation. This ensures that each fold separately assesses the importance of features.

Domain Knowledge

While statistical measures are effective, incorporating domain knowledge can guide the feature selection process. Subject matter expertise can often highlight features critical to maintaining model interpretability.

Feature Engineering

Feature engineering is complementary to feature selection. New features that better capture underlying patterns may reduce the number of features needed, thus impacting the subsequent feature selection step.


In conclusion, performing feature selection after the train-test split is generally recommended to prevent bias and information leakage, ensuring a robust and realistic evaluation of model performance. By adhering to this guideline, you enhance your model's ability to generalize to unseen data, a crucial aspect of deploying machine learning solutions.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track 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.

Practice ML system design

All Rights Reserved.