A Performance-Driven Benchmark for Feature Selection in Tabular Deep Learning

Tabular data is ubiquitous across scientific and industrial applications of machine learning. Practitioners often curate tabular datasets by including exhaustive sets of available features or by hand-engineering additional features. Under such a procedure, real-world tabular datasets can quickly accumulate a large volume of features, many of which are not useful for downstream models. Training on such a large number of features, including noisy or uninformative ones, can cause overfitting. To avoid overfitting, practitioners filter out and remove features using automated feature selection methods (Guyon and Elisseeff, 2003; Liu and Yu, 2005).

The literature contains a wide body of work proposing or evaluating feature selection approaches that use classical machine learning algorithms in their selection criteria or which select features for training classical machine learning algorithms downstream (Tibshirani, 1996; Kohavi and John, 1997). Over the past few years, deep learning for tabular data has become competitive and is increasingly adopted by practitioners (Arık and Pfister, 2021; Gorishniy et al., 2021; Somepalli et al., 2021). While the community has identified that neural networks are especially prone to overfitting on noisy features (Grinsztajn et al., 2022), a thorough evaluation of feature selection methods for downstream tabular neural networks is lacking.

To address this absence, we benchmark feature selection methods in tabular deep learning setting by evaluating the selected features via the performance of neural networks trained on them downstream. In addition to popular existing methods, we select features using the attention maps of tabular transformer models (Gorishniy et al., 2021), and we further propose Deep Lasso – an input-gradient-based analogue of Lasso (Tibshirani, 1996) for deep tabular models. Whereas many previous works on feature selection use entirely synthetic datasets (Wah et al., 2018; Passemiers et al., 2023; Bolón-Canedo et al., 2013) or create extraneous features by concatenating random noise onto existing features (Dinh and Ho, 2020; Borisov et al., 2019), we conduct our benchmark on real datasets, and explore three different ways to construct extraneous features: random noise features, corrupted features, and second-order features that serve as a prototypical example of feature engineering.

In our experiments, we find that while many feature selection methods can differentiate between informative features and noise within reason, they may fail under more challenging settings. Notably, Deep Lasso selects features which achieve significantly better downstream neural network performance than previous methods when selecting from corrupted or second-order features. ¹¹1 Our code for benchmark and Deep Lasso is available at https://github.com/vcherepanova/tabular-feature-selection

Our primary contributions are summarized as follows:

• •

  We construct a challenging feature selection benchmark comprising real-world datasets with extraneous uninformative, corrupted, and redundant features.

• •

  We benchmark different feature selection algorithms for deep tabular models, including recent tabular transformer architectures, on the basis of downstream performance.

• •

  We propose a generalization of Lasso for deep neural networks, which leverages input gradients to train models robust to changes in uninformative features. We show that Deep Lasso outperforms other feature selection methods, including tree-based methods, in the most challenging benchmark setups.

In the following section, we provide a brief overview of recent developments in tabular deep learning, and of feature selection in machine learning more broadly.

We construct a challenging feature selection benchmark that uses real datasets and includes multiple approaches for the controlled construction of extraneous features. In all cases, we evaluate feature selection methods based on downstream neural network performance, i.e. the practically relevant metric of success for a feature selection procedure. We conduct experiments with both MLP architectures and the recent transformer-based deep tabular FT-Transformer architecture (Gorishniy et al., 2021) as downstream models. Our benchmark comprises 12 datasets with 3 types of additional features. These datasets are collected and adapted based on recent tabular benchmark papers (Gorishniy et al., 2021, 2022; Rubachev et al., 2022) and include ALOI (AL), California Housing (CA), Covertype (CO), Eye Movements (EY), Gesture (GE), Helena (HE), Higgs 98k (HI), House 16K (HO), Jannis (JA), Otto Group Product Classification (OT), Year (YE) and Microsoft (MI). Among these datasets, there are eight classification datasets and four regression datasets. We measure downstream model performance using accuracy for the classification tasks and RMSE for the regression tasks. Additional details concerning these datasets can be found in Appendix B.1.

For each benchmarking experiment, we perform extensive hyperparameter tuning for both feature selection algorithms and downstream models with respect to the downstream model performance using the Bayesian hyperparameter optimization engine Optuna (Akiba et al., 2019). We select the best hyperparameters based on validation metrics and report test metrics computed over 10 random model initializations (seeds). Details concerning final hyperparameters can also be found in the Appendix sections D, C.3.

In the following section, we present a motivating experiment. In Section 5 we discuss our benchmark design in detail. Section 6 presents experimental benchmark results using feature selection methods.

Recent contributions to the ongoing competition between tabular neural networks and gradient boosted decision trees (GBDT) have found that neural networks are more susceptible to noise than GBDT on small to medium datasets (up to 10,000 samples) (Grinsztajn et al., 2022). We scale this experiment to larger datasets and showcase it as a motivating example for feature selection methods and benchmarks specific to deep tabular models.

We explore the influence of uninformative features on tabular neural networks and assess the performance of MLP and FT-Transformer models on datasets containing varying numbers of uninformative Gaussian noise features. For reference, we also include a GBDT model into our comparison as implemented in the popular XGBoost package (Chen and Guestrin, 2016). Figure 1 illustrates the relationship between the performance of the three models and the proportion of uninformative features in these datasets. Similarly to Grinsztajn et al. (2022), we observe that the MLP architecture, on average, exhibits more overfitting to uninformative features compared to XGBoost, motivating the need for careful feature selection with tabular neural networks. Interestingly, as seen from the slope of the blue and green curves in Figure 1, the FT-Transformer model is roughly as robust to noisy features as the XGBoost model. The fact that the performance of the FT-Transformer model is not as severely affected by noisy features could be attributed to the ability of the transformer architecture to filter out uninformative features through its attention mechanism. Inspired by this observation, we further investigate the effectiveness of utilizing the attention map importance within FT-Transformer Gorishniy et al. (2021) as a feature selection method in our benchmark study.

It is not common for real-world datasets to contain completely random noise features with no predictive power whatsoever, although engineered features often exhibit varying degrees of redundancy and noise. Nonetheless, feature selection algorithms are often evaluated on datasets containing spurious features generated from Gaussian noise. Not only does this differ significantly from real-world feature selection scenarios, but it also presents a relatively straightforward task for many feature selection algorithms to eliminate these random features. In light of these considerations, we propose an alternative approach to establish a more challenging and realistic feature selection benchmark by introducing three distinct methods for crafting additional features:

Random Features. In the simplest scenario, we sample uninformative features from a Gaussian distribution and concatenate them with the original dataset features.

Corrupted Features. To simulate a scenario with noisy, yet still relevant features, we sample extraneous features from the original ones and corrupt them with Gaussian noise. In addition, we conduct experiments with Laplace noise corruption and report experimental results in Appendix F.1.

Second-Order Features. To simulate a feature engineering scenario with redundant information contained in engineered features, we add second-order features, i.e. products of randomly selected original features.

Note that higher-order features are not spurious and are often used by data scientists precisely because they can contain useful information. Therefore, selecting higher order features instead of original ones may not be a bad thing. Such feature selection algorithms must be evaluated in terms of downstream model performance as we do in the following section.

To gauge the difficulty of the proposed benchmark, we explore how often different feature selection algorithms rank extraneous features among the top-$k$ most important features, where $k$ represents the number of original features in the datasets. From Figure 2, we observe that all methods select fewer random features than they do corrupted or second-order features. Additionally, to quantify the overall agreement between different feature selection methods, we analyze the average pair-wise Spearman correlation between the rankings of features generated by different selection algorithms. Notably, the setup involving random extra features exhibits the highest correlation, indicating that filtering out random features is relatively straightforward and all feature selection algorithms behave similarly. In contrast, the setup with second-order extra features has the lowest rank correlation implying greater disparity in selection preferences among the algorithms.

We note, that different feature selection algorithms may select similar features, but rank features differently within "important" and "unimportant" groups, which would results in lower rank correlation even in the settings where feature selection is straightforward (i.e. selecting from random features).

In this section, we benchmark various feature selection methods. In particular, we consider the following feature selection approaches:

Univariate Statistical Test. This classical analysis checks the linear dependence between the predictors and the target variable. It selects features based on the ANOVA F-values for classification problems and univariate linear regression test F-values for regression problems.

Lasso uses $L_{1}$ regularization to encourage sparsity in a linear regression model (Tibshirani, 1996). After this sparse regression, features are ranked with respect to the magnitudes of their coefficients in the model.

First-Layer Lasso (1L Lasso) is an extension of Lasso for MLPs with multiple layers. It applies a Group Lasso penalty to the weights of the first layer parameters:

where $W^{(j)}$ is the j-th column of weight matrix of the first hidden layer corresponding to the $j$-th feature. Similarly to Lasso, First-Layer Lasso ranks features with respect to their grouped weights in the first layer.

Adaptive Group Lasso (AGL) is an extension of the Group Lasso regularization method (Dinh and Ho, 2020). Similarly to the First-Layer Lasso, it applies a Group Lasso penalty to the weights of the first layer parameters, however each group of coefficients is weighted with an adaptive weight parameter:

where $\hat{W}$ is the Group Lasso estimate of $W$. Adaptive Group Lasso then ranks features with respect to their grouped weights in the first layer.

LassoNet is a neural network architecture that incorporates feature selection (Lemhadri et al., 2021). LassoNet achieves feature sparsity by adding a skip (residual) layer and allowing the features to participate only if their skip-layer connection is active.

Random Forest (RF) is a bagging ensemble of decision trees, and it ranks features with respect to their contribution to the ensemble (Breiman, 2001). In particular, importance is calculated by measuring the decrease in impurity when that feature is used for splitting at each node of the decision trees in the forest.

XGBoost is a popular implementation for gradient boosted decision tree (Chen and Guestrin, 2016). XGBoost computes feature importance as the average gain across all splits in trees where a feature was used.

Attention Map Importance (AM) is computed for FT-Transformer model from one forward pass on the validation set. We follow Gorishniy et al. (2021) and calculate feature importance as the average attention map for the [CLS] token across all layers, heads, and validation samples.

Deep Lasso is our generalization of Lasso to deep tabular models (and, in fact, any differentiable model). Deep Lasso encourages feature gradient sparsity for deep tabular models by applying a Group Lasso penalty to gradients of the loss with respect to input features during training. Intuitively, this makes the model robust to changes in unimportant features. For train data $(X,Y)$ with $n$ samples and $m$ features, the Deep Lasso penalty is given by

where, for $j$-th feature (column of $X$), $\frac{{\partial\mathcal{L}_{\theta}(X,Y)}}{{\partial X^{(j)}}}=\left(\frac{{\partial\mathcal{L}_{\theta}(X,Y)}}{{\partial x_{1j}}},\frac{{\partial\mathcal{L}_{\theta}(X,Y)}}{{\partial x_{2j}}},\ldots,\frac{{\partial\mathcal{L}_{\theta}(X,Y)}}{{\partial x_{nj}}}\right).$ Once the model is trained with the above regularizer, the corresponding feature importance is provided by

Note that in the linear regression case, the classical Lasso is equivalent to the proposed input-gradient sparsity regularizer applied to model output since in the linear case input gradients are the weights of the linear model. We provide a formal proof of equivalence between Deep Lasso, classical Lasso, and First-Layer Lasso in the linear regression case in Appendix Section E. We also note that Deep Lasso is related to methods used to promote network explainability by leveraging input gradients (Sundararajan et al., 2017; Shrikumar et al., 2017; Smilkov et al., 2017; Levin et al., 2022). In addition, (Liu, 2021) examines input gradients for feature selection from a Bayesian perspective. This work uses Bayesian Neural Networks (BNNs) along with associated credible intervals to estimate uncertainty surrounding input gradients, choosing features based on the plausibility that their corresponding gradients are zero. On the one hand, this principled work comes with theoretical guarantees, but on the other hand it requires Hamiltonian Monte Carlo, an expensive sampler which does not scale to large datasets and models.

In this section we analyze which feature selection algorithms are more similar to each other in terms of feature ranking. In Figure 3, we present a heatmap of pair-wise Spearman correlation between feature rankings of different FS methods averaged across datasets. We focus on the setup involving second-order features, as we discovered the lowest level of agreement between methods within this configuration. It is worth noting that we provide results for methods fine-tuned for both MLP and FT-Transformer models, ensuring a comprehensive comparison.

We find related similarity patterns for both downstream models. In particular, we observe that Random Forest and XGBoost rankings exhibit high correlation, indicating a strong agreement between these both tree-based methods. We find that First-Layer Lasso is highly correlated with Adaptive Group Lasso, which is an extension of First-Layer Lasso. Additionally, Random Forest results are also correlated with results from the Univariate Test. Deep Lasso is most correlated with Lasso-based methods, such as classical Lasso, First-Layer Lasso and Adaptive Group Lasso and with Attention Map Importance.

In this paper, we present a challenging feature selection benchmark for deep tabular models, addressing the need for more realistic scenarios in feature selection evaluation. Our benchmark encompasses real-world datasets with extraneous uninformative, corrupted, and redundant features. Through comprehensive experimentation, we compare various feature selection methods on our proposed benchmark. Our findings reveal that while classical feature selection methods, including tree-based algorithms perform competitively in the random and corrupted setups, specialized deep tabular feature selection methods, like the proposed Deep Lasso outperform other methods in selecting from second-order features. This indicates the benefits of neural network inductive biases in feature selection algorithms. Overall, our study contributes a systematic new benchmark with analysis, a new feature selection method, and insights into improving the performance and robustness of deep tabular models. The benchmark code and Deep Lasso implementation are made available to facilitate reproducibility and practical usage.

This work was made possible by the ONR MURI program and the AFOSR MURI program. Commercial support was provided by Capital One Bank, the Amazon Research Award program, and Open Philanthropy. Further support was provided by the National Science Foundation (IIS-2212182), and by the NSF TRAILS Institute (2229885).