Perform Chi-2 feature selection on TF and TFIDF vectors
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
Chi-squared feature selection is a practical way to reduce vocabulary size in text classification models built on sparse bag-of-words features. It works with both term-frequency and TF-IDF matrices as long as the features are non-negative, and the main engineering concern is fitting the selector inside the training pipeline so you do not leak label information from the evaluation set.
What Chi-Squared Is Measuring
The chi-squared score asks whether a feature's distribution is independent of the class label. In text terms, a token gets a high score when it appears disproportionately in one class and not in others.
This makes chi-squared a good fit for sparse count-like features such as:
- raw term frequency vectors
- TF-IDF vectors
- binary bag-of-words indicators
It is not a good fit for features that can be negative or for dense embeddings where the semantics are very different.
Use It With Either TF or TF-IDF
Both raw term frequency and TF-IDF are compatible with chi2 in scikit-learn. The best choice depends on the corpus rather than on a universal rule.
TF can work well when raw frequency carries strong class signal. TF-IDF can help when common words need to be downweighted more aggressively.
Keep Selection Inside the Pipeline
The most important operational rule is to fit the vectorizer and selector only on training data. If you run feature selection on the full dataset before the split, you leak label information from the test set.
Using Pipeline is the easiest way to avoid that mistake because the vectorizer, selector, and classifier are all refit inside cross-validation or train-test evaluation correctly.
That keeps the reported metrics honest.
Inspect Which Terms Were Selected
Feature selection should not be a black box. You should inspect the chosen terms and confirm they make sense.
This is where you often discover preprocessing problems such as noisy tokens, template fragments, or IDs that should have been cleaned earlier.
Choose k With Validation, Not Guesswork
There is no universally correct number of retained features. A small sweep is usually better than copying a value from another tutorial.
In practice, you often choose the smallest k that preserves the performance target well enough, because fewer features can reduce memory use and training cost.
Common Pitfalls
The first pitfall is applying chi-squared to features that can become negative. The second is doing vectorization and feature selection before the train-test split, which leaks class information.
Another issue is treating high-scoring terms as if they were causal explanations. They are discriminative features, not automatically meaningful business concepts.
Finally, do not pick k once and never revisit it. The best setting changes with preprocessing, n-gram choices, and corpus size.
Summary
- Chi-squared is a strong baseline selector for sparse text features.
- It works with both TF and TF-IDF as long as the feature values stay non-negative.
- Fit the vectorizer and selector inside a training pipeline to avoid leakage.
- Inspect selected terms so preprocessing errors do not hide inside the model.
- Choose the number of retained features through validation, not by copying a fixed number blindly.

