[Paper Review] E2E-FS: An End-to-End Feature Selection Method for Neural Networks

Brais Cancela, Verónica Bolón‐Canedo+1 more|arXiv (Cornell University)|Dec 14, 2020

Fault Detection and Control Systems4 citations

TL;DR

E2E-FS is a novel end-to-end feature selection method for neural networks that combines the accuracy of Lasso-like regularization with the interpretability of filter methods by enforcing a hard constraint on the maximum number of features selected. It uses gradient descent to jointly optimize model training and feature selection in a single step, achieving state-of-the-art performance across diverse datasets including image and microarray data.

ABSTRACT

Classic embedded feature selection algorithms are often divided in two large groups: tree-based algorithms and lasso variants. Both approaches are focused in different aspects: while the tree-based algorithms provide a clear explanation about which variables are being used to trigger a certain output, lasso-like approaches sacrifice a detailed explanation in favor of increasing its accuracy. In this paper, we present a novel embedded feature selection algorithm, called End-to-End Feature Selection (E2E-FS), that aims to provide both accuracy and explainability in a clever way. Despite having non-convex regularization terms, our algorithm, similar to the lasso approach, is solved with gradient descent techniques, introducing some restrictions that force the model to specifically select a maximum number of features that are going to be used subsequently by the classifier. Although these are hard restrictions, the experimental results obtained show that this algorithm can be used with any learning model that is trained using a gradient descent algorithm.

Motivation & Objective

• To address the trade-off between feature selection accuracy and model interpretability in embedded methods.
• To enable precise control over the number of selected features while maintaining high performance.
• To develop a method compatible with any model trained via gradient descent, including deep neural networks.
• To reduce computational cost compared to recursive methods like SVM-RFE or SFS by avoiding multiple model retraining.
• To provide a unified, efficient, and scalable solution for feature selection in high-dimensional machine learning tasks.

Proposed method

• The method formulates feature selection as a constrained optimization problem where a binary mask vector 𝛾 ∈ {0,1}^F selects features, with ||𝛾||₁ ≤ M enforcing a maximum of M features.
• It introduces a differentiable relaxation of the binary constraint using a softplus function, enabling gradient-based optimization via backpropagation.
• The loss function combines the model's prediction error with a non-convex regularization term that encourages sparsity in 𝛾.
• The algorithm is trained end-to-end using standard optimizers (e.g., Adam, SGD), with the final mask 𝛾 used to select the top M features.
• The approach is compatible with any differentiable model, including fully connected networks and convolutional neural networks (CNNs).
• The method uses a two-stage training process: first, a warm-up phase to stabilize the mask, followed by joint training with fixed hyperparameters (e.g., weight decay 5e-4, 70 epochs).

Experimental results

Research questions

• RQ1 Can a single-step, end-to-end method achieve both high accuracy and controlled feature count in neural network training?
• RQ2 How does the proposed method compare to recursive or iterative feature selection techniques in terms of computational efficiency and performance?
• RQ3 Does enforcing a hard upper bound on feature count improve generalization or model interpretability without sacrificing accuracy?
• RQ4 Can the method be effectively applied to diverse data types, including image and high-dimensional biological data?
• RQ5 How does the differentiable relaxation of the binary mask affect convergence and feature selection stability?

Key findings

• E2E-FS-Soft achieved state-of-the-art accuracy on Fashion-MNIST, CIFAR-10, and CIFAR-100, outperforming DFS, SFS, iSFS, and SFS+DFS across all feature counts.
• On CIFAR-10, E2E-FS-Soft achieved 93.68% accuracy with 392 features, a 1.3% improvement over the next best method (SFS+DFS at 92.59%).
• On CIFAR-100, E2E-FS-Soft reached 70.30% accuracy with 1536 features, surpassing DFS (67.42%) and SFS+DFS (64.94%) by over 5%.
• The method reduced training time to approximately 4180 seconds on CIFAR-100, significantly faster than iSFS (~39 hours) and SFS+DFS (7260 seconds).
• Performance gains were most pronounced with low feature counts: 29% improvement on CIFAR-100 when selecting only 153 features.
• The method showed limitations on binary datasets like MNIST, where it underperformed DFS, suggesting sensitivity to data type characteristics.