Text Feature Selection with TF-IDF and Chi-Square in Scikit-learn

Contents

Figure 1: Text Feature Selection with TF-IDF and Chi-Square — ten thousand columns, a few hundred that matter

A TF-IDF matrix for even a modest text corpus can have tens of thousands of columns, most of them words that appear in a handful of documents and tell the classifier nothing. Feature selection trims that vocabulary to the terms that actually relate to your labels.

Chi-square (χ²) is the classic filter for this job: it is fast, works directly on sparse matrices, and produces a ranked list you can read. This guide shows how it works, how to use it correctly inside a pipeline, how to avoid a leakage trap that inflates results, how to pull per-class top terms, and when other approaches serve you better.

Do You Need Feature Selection at All?

Not always. Linear models with regularization, such as logistic regression and linear SVMs, often handle huge sparse vocabularies well without any selection, and pruning can even discard useful rare terms. Selection earns its place when:

1. Memory or latency matters, and a smaller vocabulary means smaller models and faster inference.
2. The dataset is small relative to the vocabulary, so noise features invite overfitting.
3. You need interpretability: a ranked list of discriminative terms is a useful artifact in itself.
4. The classifier is sensitive to irrelevant features, as naive Bayes and k-NN can be.
5. You want a quick baseline before more elaborate modeling.

Before reaching for chi-square, try the cheap, unsupervised cuts in the vectorizer itself: min_df to drop very rare terms, max_df to drop near-ubiquitous ones, and max_features to cap the vocabulary. These often deliver most of the benefit with no risk of label leakage. The vocabulary-building side of this is covered in the TF-IDF guide.

What Chi-Square Measures

For each term, chi-square tests whether its occurrence is independent of the class label. The intuition: if "goal" appears in 90 of 100 hockey posts but only 5 of 400 other posts, its distribution across classes is far from what independence would predict, so it gets a high score. A term like "the" appears about equally everywhere, so its score is near zero.

The statistic sums (observed − expected)² / expected across classes, where "expected" is what you would see if the term were spread in proportion to class sizes.

Three Practical Points

Point	Consequence
Features must be non-negative	TF-IDF, counts, and booleans qualify; standardized features do not
Each feature is scored independently	Redundancy and interactions between terms are ignored
On TF-IDF, treat it as a ranking	The classical test assumes counts, so the p-values are approximate

With TF-IDF input, treat it as a ranking, not a formal test. The classical test assumes counts, and scikit-learn applies the same formula to whatever non-negative values you supply, so the p-values are approximate. Use the scores to rank features, and do not read the p-values as strict evidence, especially since you are running thousands of tests at once [1][2].

Basic Usage

We will classify four newsgroup categories, using scikit-learn's 20 newsgroups loader (which downloads the data on first use). Headers, footers, and quotes are stripped, because they give the model shortcuts like email domains that do not reflect real text understanding [5]:

import numpy as np
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_selection import SelectKBest, chi2
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import Pipeline

cats = ["sci.space", "rec.sport.hockey", "talk.politics.guns", "comp.graphics"]
data = fetch_20newsgroups(subset="train", categories=cats,
                          remove=("headers", "footers", "quotes"))
X_text, y = data.data, data.target

def make_pipe(k):
    return Pipeline([
        ("tfidf", TfidfVectorizer(sublinear_tf=True, min_df=2, max_df=0.5,
                                  ngram_range=(1, 2), stop_words="english")),
        ("select", SelectKBest(chi2, k=k)),           # k="all" keeps everything
        ("clf", LogisticRegression(max_iter=1000)),
    ])

full = make_pipe("all").fit(X_text, y)
print("vocabulary size:", len(full.named_steps["tfidf"].vocabulary_))

cv = StratifiedKFold(5, shuffle=True, random_state=0)
for k in ("all", 5000, 2000, 1000, 300, 100):
    scores = cross_val_score(make_pipe(k), X_text, y, cv=cv, scoring="accuracy")
    print(f"k={str(k):>5}  accuracy {scores.mean():.3f} ± {scores.std():.3f}")

Typically accuracy stays nearly flat as k shrinks from the full vocabulary down to a few thousand terms, then drops only gradually, until the selector starts cutting genuinely informative features. That flat region is the point: you can often shrink the model by an order of magnitude at little cost. Your exact numbers will vary.

Everything lives inside one Pipeline, which matters enormously, as the next section shows.

The Leakage Trap

Chi-square uses the labels, so it is supervised. If you select features on the whole dataset and then cross-validate, the selection step has already seen the validation labels, and your accuracy is inflated. The effect is dramatic with text, where the number of features dwarfs the number of documents. This experiment uses random labels, so honest accuracy should be about 50%:

rng = np.random.default_rng(0)
y_random = rng.integers(0, 2, size=len(X_text))

Xv = TfidfVectorizer(min_df=2, stop_words="english").fit_transform(X_text)  # unsupervised step

# WRONG: select on all the data first, then cross-validate
X_sel = SelectKBest(chi2, k=300).fit_transform(Xv, y_random)
wrong = cross_val_score(LogisticRegression(max_iter=1000), X_sel, y_random, cv=cv).mean()

# RIGHT: selection happens inside each training fold
pipe = Pipeline([("select", SelectKBest(chi2, k=300)),
                 ("clf", LogisticRegression(max_iter=1000))])
right = cross_val_score(pipe, Xv, y_random, cv=cv).mean()

print(f"selection outside CV: {wrong:.3f}   selection inside CV: {right:.3f}")

With random labels there is nothing to learn, yet selecting 300 features on all the data first produces accuracy well above chance, while the pipeline version hovers around 0.5.

Setup	Labels contain	Reported accuracy
Selection outside CV	No signal to learn	Well above chance
Selection inside the pipeline	No signal to learn	≈ 0.5, as it should be

The first number is the one that gets people into trouble: it looks like a great model and fails in production. (The TF-IDF step uses no labels, so fitting it outside the loop here is a minor leak. In real work, keep it inside the pipeline too, as above.)

The rule is the same as in the preprocessing guide in this series: anything that learns from labels belongs inside the cross-validation loop, and the final test set stays untouched until the end [4].

Choosing k

Treat k as a hyperparameter and tune it with cross-validation:

from sklearn.model_selection import GridSearchCV

grid = GridSearchCV(
    make_pipe(1000),
    {"select__k": [200, 500, 1000, 2000, 5000],
     "clf__C": [0.3, 1, 3, 10]},
    cv=cv, scoring="f1_macro", n_jobs=-1,
)
grid.fit(X_text, y)
print(grid.best_params_, round(grid.best_score_, 3))

Tune k together with the classifier's regularization, since the best C depends on how many features you give the model. When scores are close across several values of k, prefer the smallest k within one standard error of the best, because it gives a smaller, faster model with indistinguishable accuracy.

Per-Class Top Terms

SelectKBest(chi2) on a multiclass problem aggregates evidence across all classes into one score per term. To see which terms characterize each class, run chi-square one-vs-rest:

tfidf = TfidfVectorizer(sublinear_tf=True, min_df=2, max_df=0.5,
                        ngram_range=(1, 2), stop_words="english")
Xv = tfidf.fit_transform(X_text)          # for exploration only, not for evaluation
names = np.array(tfidf.get_feature_names_out())

for i, label in enumerate(data.target_names):
    in_class = (y == i)
    scores, _ = chi2(Xv, in_class.astype(int))
    class_mean = np.asarray(Xv[in_class].mean(axis=0)).ravel()
    rest_mean = np.asarray(Xv[~in_class].mean(axis=0)).ravel()
    scores = np.nan_to_num(scores)
    scores[class_mean <= rest_mean] = 0          # keep only terms over-represented in this class
    top = np.argsort(scores)[::-1][:10]
    print(f"{label:<20}", ", ".join(names[top]))

The direction filter matters. Chi-square measures dependence, not direction, so a term that is unusually rare in a class scores highly just like one that is unusually common. Without the filter, you would see terms that characterize what a class is not.

Reading these lists is also a quick way to catch artifacts: if the top terms include email fragments, signatures, or formatting debris, the model may be learning shortcuts you should clean out.

Alternatives and Complements

Approach	What it buys	What it costs
min_df / max_df / max_features	Cheap, leak-free vocabulary control	No use of labels at all
mutual_info_classif	Non-linear dependence, sparse input	Slower, noisier on sparse text
f_classif (ANOVA)	Another univariate score	Assumptions less suited to sparse term data
SelectFromModel + L1	Accounts for redundancy	Needs a model fit, slower
TruncatedSVD (LSA)	Dense components, smaller input	Loses term interpretability
HashingVectorizer	Fixed dimensionality, no vocabulary	Collisions, no term names
Pretrained embeddings	Semantic coverage	Interpretability, sometimes cost

Unsupervised pruning means min_df, max_df, and max_features in the vectorizer. Cheap and leak-free.

Mutual information (mutual_info_classif) captures non-linear dependence and handles sparse input, but is slower and noisier on sparse text [1].

ANOVA F-test (f_classif) is another univariate score, with assumptions less suited to sparse term data.

Embedded selection with SelectFromModel trains a sparse linear model and keeps the terms with the largest weights. Unlike chi-square, it accounts for redundancy, because correlated terms compete for weight:

from sklearn.feature_selection import SelectFromModel
from sklearn.svm import LinearSVC

l1_select = SelectFromModel(
    LinearSVC(C=0.3, penalty="l1", dual=False),
    max_features=1000, threshold=-np.inf,         # keep the 1,000 largest |weights|
)
pipe_l1 = Pipeline([
    ("tfidf", TfidfVectorizer(sublinear_tf=True, min_df=2, max_df=0.5,
                              ngram_range=(1, 2), stop_words="english")),
    ("select", l1_select),
    ("clf", LogisticRegression(max_iter=1000)),
])
print(cross_val_score(pipe_l1, X_text, y, cv=cv).mean())

Compare this against chi-square at the same feature count in the same cross-validation setup, since which wins depends on your data.

Recursive feature elimination is thorough but expensive on large vocabularies.

Dimensionality reduction (TruncatedSVD, also called latent semantic analysis) replaces terms with dense components, which can help but removes the interpretability that makes term selection attractive.

Hashing (HashingVectorizer) caps dimensionality with no vocabulary at all, at the price of collisions and no term names.

Pretrained embeddings sidestep vocabulary selection entirely, trading interpretability and sometimes cost for semantic coverage, as discussed in the TF-IDF guide. The filter, wrapper and embedded comparison covers the families in more depth.

Limitations of Chi-Square

1. Univariate: it ranks terms one at a time, so near-duplicates (a term and its bigrams) all score high together and crowd out diversity.
2. Unreliable for rare terms: with tiny expected counts, the statistic becomes unstable and can inflate scores for terms seen only a couple of times, which is one reason min_df is a good companion.
3. Not a proper test on TF-IDF: p-values are approximate, and with thousands of terms you face a multiple-testing problem, so use it for ranking.
4. Non-negative inputs only.
5. Class-imbalance effects: large classes dominate the aggregated multiclass score, so consider per-class selection and combining the lists if small classes matter.
6. Can discard useful rare discriminators, such as a rare term that perfectly marks a small class.
7. It cannot see interactions, which a linear model cannot use anyway, but a tree or neural model could.

A Practical Workflow

1. Build a baseline with no selection: TF-IDF, sensible min_df and max_df, and a regularized linear model.
2. Add chi-square selection inside the pipeline only if you need a smaller model, speed, or a cleaner vocabulary.
3. Tune k with cross-validation, jointly with regularization, and favor the smallest k within one standard error of the best.
4. Inspect the selected terms, both globally and per class, for leakage artifacts and boilerplate.
5. Compare against an embedded selector like L1 and, if accuracy is the goal, against a pretrained-embedding baseline.
6. Save the whole fitted pipeline with joblib, so serving applies exactly the same vocabulary and selection.
7. Monitor vocabulary drift in production, since new terms the selector never saw will not be used.

Common Pitfalls

1. Selecting features on the full dataset before cross-validation or the train-test split.
2. Reading chi-square p-values from TF-IDF as rigorous evidence.
3. Interpreting high chi-square as "this term indicates the class", when it may indicate the term's absence from the class.
4. Ignoring rare-term instability by skipping min_df.
5. Setting k larger than the vocabulary, which raises an error.
6. Selecting on imbalanced data without considering minority-class terms.
7. Choosing k by looking at the test set.
8. Leaving metadata and signatures in the text, so the selector latches onto shortcuts.
9. Forgetting that the selector and vectorizer must ship together in production.
10. Refitting the vectorizer at serving time on live traffic. The vocabulary must be frozen at fit time, or the column indices mean something different from what the selector was trained on.

Wrapping Up

Chi-square feature selection gives text models a fast, interpretable way to cut a huge TF-IDF vocabulary down to the terms most dependent on the labels. Use SelectKBest(chi2, k=...) inside a pipeline so the selection is learned only from training folds, tune k by cross-validation, run it one-vs-rest with a direction filter when you want per-class term lists, and remember the leakage experiment: selection outside cross-validation can make random labels look learnable.

Should you use it? Start with min_df and max_df and a regularized linear model, and add chi-square when size, speed, or interpretability justifies it. Then verify that the model you ship scores the same inside a clean pipeline, because that is the only number that predicts what happens on new documents.

Frequently Asked Questions

Does chi-square feature selection need scaled or normalised input?

It needs non-negative input, not scaled input. Raw counts, TF-IDF values and booleans all qualify because every entry is zero or positive. Standardised features do not, because z-scoring produces negative values and scikit-learn raises an error. That is one reason chi-square suits text so well: a sparse TF-IDF matrix is non-negative by construction.

Why does selecting features before cross-validation inflate accuracy?

Because chi-square is supervised. If you score terms on the whole dataset and then cross-validate, the selector has already read the validation labels, so those labels leak into the feature set. With random labels the honest accuracy is about 50%, yet selecting 300 features on all the data first produces accuracy well above chance. The fix is to put SelectKBest inside the pipeline so it is refitted on each training fold.

Can I use chi-square p-values as a statistical test on text?

Not really. The classical test assumes counts, and scikit-learn applies the same formula to whatever non-negative values you supply, so with TF-IDF input the p-values are approximate. You are also running thousands of tests at once, which is a multiple-testing problem. Treat the scores as a ranking and use them to order or cut features, not as evidence of significance.

How do I choose k for SelectKBest?

Treat k as a hyperparameter and tune it with cross-validation, jointly with the classifier's regularisation, since the best C depends on how many features the model receives. When several values score within one standard error of the best, prefer the smallest, because it gives a smaller and faster model at indistinguishable accuracy. Also remember that k larger than the vocabulary raises an error.

Why do my top chi-square terms include words that are rare in the class?

Because chi-square measures dependence, not direction. A term that is unusually rare inside a class scores just as highly as one that is unusually common, so a raw ranking mixes both. Filter by direction: compute the mean TF-IDF inside the class and outside it, and zero out terms that are not over-represented inside the class before taking the top terms.

Should I use chi-square or SelectFromModel with an L1 model?

Chi-square is a univariate filter: it scores each term alone, so near-duplicate terms all rank highly and crowd out variety. SelectFromModel with an L1-penalised linear model is embedded selection and accounts for redundancy, because correlated terms compete for weight. Chi-square is faster and gives a readable global ranking; L1 often gives a better model. Compare both at the same feature count in the same cross-validation setup.

Sources

• [1] scikit-learn, "Feature selection" user guide — univariate selection, SelectKBest, the note that chi-square requires non-negative features, and the comparison with mutual_info_classif and f_classif: scikit-learn.org/modules/feature_selection
• [2] scikit-learn, chi2 — the exact statistic computed per feature, the non-negative requirement that raises on negative input, and the return of scores alongside p-values: scikit-learn.org/.../chi2
• [3] scikit-learn, TfidfVectorizer — min_df, max_df, max_features, sublinear_tf, ngram_range, and get_feature_names_out() used to map scores back to terms: scikit-learn.org/.../TfidfVectorizer
• [4] scikit-learn, "Common pitfalls and recommended practices" — the leakage patterns behind the random-label experiment, and why fitted transformers belong inside the pipeline: scikit-learn.org/common_pitfalls
• [5] scikit-learn, "The 20 newsgroups text dataset" — the loader, the remove=("headers", "footers", "quotes") option and the reason those fields act as shortcuts: scikit-learn.org/datasets/real_world

Related Articles

• Feature Engineering for Text Data: Bag of Words and TF-IDF
• Feature Selection: Filter, Wrapper and Embedded Methods
• Naive Bayes Classifier for Text and Tabular Data in Python