Pretrained model boosts feature screening for complex data patterns

PACE: Plug-and-Play Contextual Embedding for Feature Screening with Pretrained Tabular Foundation Models

Machine Learning

Summary

Finding the important pieces of data can be tricky, especially when relationships are complicated or hidden. The authors created a way called PACE that uses a powerful pre-learned model to transform each piece of data before deciding if it matters. This approach helps computers better spot relevant information and improves predictions without much extra work. Their tests show PACE works better than some other methods and is efficient enough to use more widely.

What this means in practice

  • For data scientists: Improve data preprocessing by removing irrelevant features using pretrained embeddings to capture complex feature interactions before model training.
  • For machine learning engineers: Incorporate pretrained column encoders for efficient feature selection that speeds up training of predictive models on high-dimensional tabular datasets.

Authors

Qi Qin, Erbo Li, Ting Wei, Zizhou Huang, Zixuan Qin, Wu Wang, Yifan Sun

Abstract

In high-dimensional tabular learning, feature screening provides a lightweight, model-agnostic way to remove irrelevant features before model fitting. However, scoring raw values directly can miss nonlinear or distributional structure. We introduce PACE (Plug-and-Play Contextual Embedding), which inserts a frozen tabular foundation model (TFM) column encoder before an existing feature-scoring rule, expanding each feature into a higher-dimensional contextual representation. Across controlled studies, PACE improves raw-space screening of complex nonlinear dependence with only modest additional encoding cost. These gains translate to downstream prediction on TALENT datasets: PACE-DC improves binary AUC by 0.077 and multiclass macro-AUC by 0.064, with a median normalized RMSE improvement of 0.063 across ten learners. Matched random-weight and random-feature controls show that PACE gains from pretrained structure beyond generic dimensional expansion. PACE further achieves favorable performance--time trade-offs against task-fitted selectors and attribution-based methods, positioning pretrained column geometry as a reusable upstream primitive for tabular learning.