Source-linked AI summary
LassoNet: A Neural Network with Feature Sparsity
Ismael Lemhadri, Feng Ruan, Louis Abraham, Robert Tibshirani
TL;DR
Feature selection is well established for linear models but remains challenging for expressive neural networks and can miss interaction-driven relevance. LassoNet addresses this gap with hierarchical constraints that jointly learn sparse input features and nonlinear predictors, producing a regularization path; experiments report broad superiority across feature-selection and regression tasks.
Problem
Feature selection methods must identify relevant features while supporting nonlinear neural-network predictors, but filter methods can miss interactions and Lasso applies only to linear models.
Method
LassoNet uses a skip-layer hierarchy requiring hidden-unit participation to be accompanied by an active linear representative, with joint optimization of linear and nonlinear components.
Results
LassoNet generally outperforms state-of-the-art methods for feature selection and regression across tested datasets.
Takeaways & Limitations
LassoNet provides a regularization path with varying feature sparsity and can extend to other learning tasks with limited implementation changes.
Takeaways & Limitations
The ℓ1 penalty biases weights downward, so optimal predictive performance may require retraining the selected network to debias it.
Abstract
from arXiv · showhide
Much work has been done recently to make neural networks more interpretable, and one obvious approach is to arrange for the network to use only a subset of the available features. In linear models, Lasso (or $\ell_1$-regularized) regression assigns zero weights to the most irrelevant or redundant features, and is widely used in data science. However the Lasso only applies to linear models. Here we introduce LassoNet, a neural network framework with global feature selection. Our approach enforces a hierarchy: specifically a feature can participate in a hidden unit only if its linear representative is active. Unlike other approaches to feature selection for neural nets, our method uses a modified objective function with constraints, and so integrates feature selection with the parameter learning directly. As a result, it delivers an entire regularization path of solutions with a range of feature sparsity. On systematic experiments, LassoNet significantly outperforms state-of-the-art methods for feature selection and regression. The LassoNet method uses projected proximal gradient descent, and generalizes directly to deep networks. It can be implemented by adding just a few lines of code to a standard neural network.
1. Introduction
Feature selection can reduce cost and improve interpretability, but traditional approaches have limitations for nonlinear neural networks. LassoNet extends sparse feature selection to neural networks through hierarchical constraints and jointly optimized linear and nonlinear components.
- Feature selection removes redundant or noisy features, reducing downstream computational expense while improving insight into feature–outcome relationships.
- Experiments: On the MICE Protein Dataset, about 35 proteins suffice to obtain maximal classification accuracy.
- Related work: Filter methods may miss features that matter mainly through interactions, while Lasso is limited to linear models.
- Experiments: Figure 2 illustrates simultaneous pixel selection and digit classification on MNIST, with pixel importance defined by changes in predicted probability under maximum-intensity perturbation.
- Proposed method: LassoNet uses an input-to-output skip connection so hidden-layer weights for a feature can be nonzero only when its skip connection is active.
- Proposed method: The linear and nonlinear components are optimized jointly, allowing LassoNet to capture arbitrary nonlinearity and generally outperform state-of-the-art feature-selection and regression methods.
2. Problem Formulation
The formulation seeks a predictive function using a subset of features without restricting the function class to linear or quadratic forms. Its central challenge is balancing expressive nonlinear modeling with exponentially difficult subset selection.
- Global feature selection is formulated for supervised learning, with the goal of finding the best function f*(x) for predicting responses.
- The objective minimizes empirical reconstruction error over a feature subset S and a user-specified loss function.
- Choosing among feature subsets is combinatorial, growing exponentially with the number of features and becoming NP-hard even for simple linear models.
- The function class must remain expressive enough to approximate solutions ranging from linear regression to deep fully connected neural networks.
3. Our proposal: LassoNet
LassoNet extends feature sparsity to feed-forward neural networks by tying nonlinear feature use to a linear residual representative. A penalty path and hierarchical proximal updates provide controllable sparsity while retaining nonlinear modeling.
- Varying the ℓ1 penalty transforms combinatorial feature selection into a continuous regularization-path search, with larger λ encouraging sparser models.
- The hierarchy constraint bounds first-layer weights by the corresponding residual coefficient, making zero residual weights exclude features globally.
- LassoNet uses a residual connection so a feature can enter hidden units only when its residual weight is active.
- LassoNet recovers Lasso when M = 0 and approaches an unregularized feed-forward network as M tends to infinity.
- The formulation jointly fits linear and nonlinear components, allowing arbitrary nonlinearities while selecting relevant input features through one coefficient per feature.
- Training alternates ordinary gradient descent with a hierarchical proximal update on the residual and first hidden layers, requiring only minor changes to standard networks.
4. Optimization
LassoNet optimizes its regularization path with proximal gradient descent and a specialized hierarchical proximal solver. Dense-to-sparse warm starts improve generalization, while warm starts and pruning make the path computationally efficient.
- Dense-to-sparse warm starts outperform sparse-to-dense starts because sparse-to-dense optimization can become trapped in poorly generalizing local minima.
- Hier-Prox globally solves the non-convex proximal inner problem through feature-wise decomposable optimality conditions.
- Hier-Prox has O(p · log p) complexity, with negligible overhead relative to gradient computation and straightforward implementation in standard deep-learning frameworks.
- Warm starts and Hier-Prox make subsequent regularization-path training cheap by effectively pruning a previously trained dense model.
- The regularization path can be trained at a cost essentially equivalent to training one model, while avoiding advance knowledge of the optimal feature count.
- ℓ1 regularization slightly biases weights downward, so retraining the selected reduced network can improve predictive performance after model selection.
5. Experiments
LassoNet is evaluated on real-world protein, image, voice, object, and smartphone datasets against established feature-selection methods. Across experiments, it is generally the strongest performer, while ISOLET results show the highest classification accuracy across tested feature counts and both learners.
- Datasets: The experiments cover protein, image, voice, object, and smartphone datasets previously used to benchmark feature-selection methods.The datasets include Mice Protein, MNIST, MNIST-Fashion, ISOLET, COIL-20, and Smartphone Human Activity Recognition.
- Methodology: LassoNet is compared with HSIC-LASSO, Fisher Score, and principal feature analysis across varying numbers of selected features.Other methods typically select features first, after which separate downstream learners are applied.
- Methodology: Classification accuracy is evaluated using a one-hidden-layer feed-forward network and an Extremely Randomized Trees classifier.These learners are used to assess the value of the selected feature sets.
- Results: Across all tested k values and both learners on ISOLET, LassoNet achieves the highest classification accuracy.Figure 5 compares the one-hidden-layer neural network and Extremely Randomized Trees results.
- Results: LassoNet ties or outperforms the remaining methods in all reported instances and is the strongest performer in the large majority of cases.The comparison uses classification errors in the decoder-network experiments and corresponding tree-based results in the appendix.
- Results: Unsupervised LassoNet reconstructions on MNIST generally identify the digit correctly while preserving stylistic features such as orientation and thickness.The displayed reconstructions use an intermediate penalty corresponding to about 50 active features.
6. Extension to Unsupervised Feature Selection
The paper extends LassoNet to unsupervised feature selection by replacing classification with reconstruction. The approach selects a common subset of input features, and MNIST reconstructions generally preserve digit identity and selected stylistic details.
- Background: Unsupervised feature selection is motivated by settings where prediction tasks are unknown and missing features must be imputed with minimal distortion.The goal is to identify important features and redundant measurements without committing to a specific downstream task.
- Method: The unsupervised extension replaces the neural-network classifier with a decoder network trained using reconstruction loss.The decoder reconstructs X from the input through a feed-forward neural network.
- Method: A group-LASSO penalty enforces the same selected feature set across all reconstructed inputs.The resulting optimization uses the Group-Hier-Prox algorithm.
- Method: Training initializes a feed-forward network, follows a penalty path, computes gradients by backpropagation, and updates the active-feature count from nonzero linear coefficients.The procedure iterates while active features remain and applies the group hierarchical proximal operator.
7. Extension to Matrix Completion
The matrix-completion extension uses a neural network to model nonlinear low-dimensional structure without assuming that the data matrix is low rank. It also performs feature selection, and the reported results favor this approach when linear reconstruction is inappropriate.
- Background: Matrix completion predicts unobserved entries in a large sparse matrix, a setting where measurements may be costly or limited.The observed entries are indexed by Ω.
- Method: Unlike Soft-Impute and related methods, LassoNet’s matrix-completion extension makes no low-rank assumption about the underlying matrix.Existing approaches commonly use low-rank assumptions, whereas LassoNet targets arbitrary nonlinear low-dimensional structure.
- Method: The method iteratively trains a feed-forward network on reconstruction loss and updates the imputed input using the network’s reconstruction.Soft-Impute instead uses singular-value thresholding to find a linear low-dimensional structure.
- Method: After dense-model training, Group-Hier-Prox prunes the imputation model to retain only a small set of input features.Feature selection is performed along the λ-path.
- Experiments: The MICE Protein experiment hides 80% of entries for training, uses 10% for validation, and reserves 10% for testing.Early stopping with a patience of 10 epochs controls reconstruction-network training.
- Results: About 30% lower reconstruction error is reported for LassoNet than Soft-Impute before feature selection on MICE Protein data.The figure reports test mean-squared error over 25 trials with different initializations.
- Results: The results suggest LassoNet can outperform linear reconstruction when the data follow an arbitrary nonlinear low-dimensional structure.The paper also identifies feature selection as useful when measuring features is costly or unavailable.
8. Sparsity in Learned Features
The paper notes that sparsity over raw inputs is not always useful when features lack stable interpretability. It proposes extending sparsity to learned convolutional filters or latent units as future work.
- Scope boundary: Raw-feature sparsity is limited when the input features are not interpretable.The paper identifies this as a setting where inducing sparsity directly in those features is not helpful.
- Computer vision: For computer vision, selecting individual pixels may be inappropriate because corresponding pixels can lack consistent meaning across unregistered images.The proposed alternative is filter sparsity over learned convolutional filters.
- Computer vision: The paper suggests applying LassoNet to convolutional-layer outputs to select relevant visual filters and reduce the need for architecture search.This is presented as future work rather than an implemented extension in the paper.
- Auto-encoders: A related future direction is selecting the most important hidden units in auto-encoder bottleneck layers.This would further reduce auto-encoder complexity beyond latent-variable learning.
9. Discussion
LassoNet provides neural-network feature selection through hierarchy-constrained nonconvex optimization, producing a regularization path at roughly the cost of one model. It is general, easy to implement, and requires validation because it does not quantify statistical significance.
- Method: LassoNet uses hierarchy constraints and proximal gradient descent to produce models with varying numbers of selected features.The method alternates stochastic gradient descent with an analytically solved subproblem across a regularization path.
- Efficiency: The entire regularization path costs essentially as much as training a single model.The initial dense model explores and converges over models with different feature counts.
- Practical use: Unlike methods requiring a prespecified feature count, LassoNet provides a path spanning different numbers of selected features.The additional hierarchy coefficient has default value M = 10 in the reported experiments.
- Generality and ease of use: LassoNet extends to unsupervised reconstruction and matrix completion, while implementation requires only a few changes to a standard feed-forward network.Its runtime also remains similar to training one model over the full path, with hardware acceleration and parallelization improving performance.
- Limitation: LassoNet does not provide p-values or statistical significance quantification, so selected features require hypothesis testing or domain-specific analysis.This limitation is shared with the other feature-selection methods compared in the paper.
Appendix A. Additional Experiments
Across additional datasets, LassoNet generally maintains high classification accuracy while generating feature-count comparisons at essentially unchanged computational cost.
- Additional results: LassoNet continues to achieve high, though not always highest, accuracy across the additional datasets.The experiments report results for different numbers of selected features, complementing the fixed-k = 50 comparison in Table 1.
- Computational cost: Other methods require multiples of the table-generation cost to produce curves across feature counts, whereas LassoNet’s cost remains essentially unchanged.The difference arises because competing methods typically rerun from scratch for each requested k, while LassoNet produces a path.
- Dataset coverage: Figures 11 and 12 report additional results on the COIL and ACTIVITY datasets.The supplied captions identify the datasets but do not state a specific comparative outcome.
Appendix B. Proofs
The appendix proves correctness properties for hierarchy-proximal operators by reducing constrained optimization to feature-wise subproblems and characterizing the global optimum through ordered coordinates and a uniquely selected index.
- Optimization decomposition: Each iteration performs vanilla gradient descent followed by a constrained minimization that decomposes into d independent single-feature problems.This decomposition is the basis for the Hier-Prox optimization step.
- Operator guarantees: Hier-Prox and Hier-Prox-Group are stated to return global optima of their respective optimization problems.The proof analyzes the unconstrained and grouped operators separately.
- Global-optimum characterization: The proof orders coordinates by decreasing absolute value and defines candidate solutions b_s for each possible active-count index s.The global solution is identified among these candidates.
- Objective analysis: The objective is analyzed piecewise through F_s and a remainder term r_s independent of b, allowing the global minimum to be located over the candidate intervals.The appendix states that the global minimum b* equals b_s* and establishes the corresponding deferred claims.
- Proof structure: A unique s* is established through an increasing function h satisfying h(s*) ≤ 0 < h(s* + 1).The appendix then uses this index to identify the global minimizer b* = b_s*.
D.1 LassoNet Architecture
The reported implementation uses PyTorch feed-forward networks with ReLU activations, tunes hidden-layer width by validation accuracy, and draws datasets from official or UCI sources.
- Architecture: LassoNet is implemented in PyTorch as a one-hidden-layer feed-forward network with ReLU activation.A two-hidden-layer network is also used for the matrix completion problem.
- Model selection: The hidden-layer width is varied across [d/3, 2d/3, d, 4d/3], and the highest-validation-accuracy network is evaluated on the test set.Here d denotes the total number of features.
- Datasets: MNIST and MNIST-Fashion come from their official sources, while the remaining datasets come from the UCI Repository.