Chi-2 feature selection
TF vectors
`TF-IDF`
vectors
text analysis
machine learning

Perform Chi-2 feature selection on TF and TFIDF vectors

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

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.

python
1from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
2from sklearn.feature_selection import SelectKBest, chi2
3from sklearn.linear_model import LogisticRegression
4from sklearn.model_selection import train_test_split
5from sklearn.pipeline import Pipeline
6from sklearn.metrics import f1_score
7
8texts = [
9    "refund requested for delayed shipment",
10    "invoice mismatch and payment failure",
11    "the team won the final game",
12    "great defense and midfield passing",
13    "credit card charged twice refund needed",
14    "coach explains training strategy",
15    "billing issue unresolved",
16    "new striker signs with club",
17]
18labels = ["support", "support", "sports", "sports", "support", "sports", "support", "sports"]
19
20x_train, x_test, y_train, y_test = train_test_split(
21    texts, labels, test_size=0.25, random_state=0, stratify=labels
22)
23
24
25def build_pipeline(vectorizer):
26    return Pipeline([
27        ("vec", vectorizer),
28        ("sel", SelectKBest(chi2, k=8)),
29        ("clf", LogisticRegression(max_iter=1000)),
30    ])
31
32for name, vectorizer in {
33    "tf": CountVectorizer(ngram_range=(1, 2)),
34    "tfidf": TfidfVectorizer(ngram_range=(1, 2)),
35}.items():
36    model = build_pipeline(vectorizer)
37    model.fit(x_train, y_train)
38    pred = model.predict(x_test)
39    print(name, f1_score(y_test, pred, average="macro"))

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.

python
1import numpy as np
2from sklearn.feature_extraction.text import TfidfVectorizer
3from sklearn.feature_selection import SelectKBest, chi2
4
5vec = TfidfVectorizer(min_df=1, ngram_range=(1, 2))
6X = vec.fit_transform(texts)
7selector = SelectKBest(chi2, k=6)
8selector.fit(X, labels)
9
10feature_names = np.array(vec.get_feature_names_out())
11mask = selector.get_support()
12
13for term, score in sorted(
14    zip(feature_names[mask], selector.scores_[mask]),
15    key=lambda pair: pair[1],
16    reverse=True,
17):
18    print(term, round(score, 4))

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.

python
1from sklearn.model_selection import cross_val_score
2
3for k in [100, 500, 1000]:
4    pipeline = Pipeline([
5        ("vec", TfidfVectorizer(min_df=1)),
6        ("sel", SelectKBest(chi2, k=k)),
7        ("clf", LogisticRegression(max_iter=1000)),
8    ])
9    scores = cross_val_score(pipeline, texts * 20, labels * 20, cv=3, scoring="f1_macro")
10    print(k, scores.mean())

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.

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.