Neural Additive and Basis Models with Feature Selection and Interactions

Abstract: Deep neural networks (DNNs) exhibit attractive performance in various fields but often suffer from low interpretability. The neural additive model (NAM) and its variant called the neural basis model (NBM) use neural networks (NNs) as nonlinear shape functions in generalized additive models (GAMs). Both models are highly interpretable and exhibit good performance and flexibility for NN training. NAM and NBM can provide and visualize the contribution of each feature to the prediction owing to GAM-based architectures. However, when using two-input NNs to consider feature interactions or when applying them to high-dimensional datasets, training NAM and NBM becomes intractable due to the increase in the computational resources required. This paper proposes incorporating the feature selection mechanism into NAM and NBM to resolve computational bottlenecks. We introduce the feature selection layer in both models and update the selection weights during training. Our method is simple and can reduce computational costs and model sizes compared to vanilla NAM and NBM. In addition, it enables us to use two-input NNs even in high-dimensional datasets and capture feature interactions. We demonstrate that the proposed models are computationally efficient compared to vanilla NAM and NBM, and they exhibit better or comparable performance with state-of-the-art GAMs.

• The paper introduces NAM-FS and NBM-FS, which embed differentiable feature selection into neural additive and basis models to enhance scalability and interpretability.
• Experiments on high-dimensional datasets show that these models achieve comparable or improved accuracy with significantly reduced computational demands.
• The study demonstrates that joint optimization of feature selection and interaction modeling enables efficient capture of univariate and bivariate relationships in critical domains.

Neural Additive and Basis Models with Feature Selection and Interactions: Technical Summary

Introduction and Motivation

This work addresses the interpretability bottleneck inherent in deep neural networks (DNNs), particularly for high-stakes domains such as healthcare and law. The neural additive model (NAM) and neural basis model (NBM) are established transparent variants of generalized additive models (GAMs), leveraging neural networks as nonlinear univariate shape functions, thereby combining interpretability with competitive predictive performance. However, as feature dimensionality increases or when pairwise interactions are incorporated (e.g., in GA$^2$Ms), both computational resource requirements and model management complexity (number of parameters, forward MACs, memory) become intractable for vanilla implementations. This paper proposes embedding differentiable feature selection into NAM and NBM architectures, enabling scalable modeling with explicit feature selection for both univariate and bivariate interactions.

Model Architecture

The proposed extensions, NAM-FS and NBM-FS, introduce trainable feature selection layers based on entmax parameterization, facilitating end-to-end differentiability. Input features are selected via sparse relaxations (entmax outputs), and the temperature parameter is annealed during training to achieve near one-hot selection, at which point feature selection becomes deterministic. Univariate ($K_1$) and bivariate ($K_2$) shape functions are parameterized by separate selection layers. During training, both shape function and selection weights are jointly optimized. Architecture is naturally extensible to multiclass outputs by linear combination.

Figure 1: NAM-FS architecture for single output, with sparse feature selection integrated before each shape function.

Figure 2: NBM-FS architecture for single output, employing shared basis functions with feature selection applied at input.

Computational Complexity and Throughput

The addition of feature selection substantially reduces the number of neural networks and computational operations required. For NAM, the parameter count reduces from $D P_{\mathrm{NAM}}$ to $K_1 (D + P_{\mathrm{NAM}})$. Pairwise interactions (NA$^2$M/NA$^2$M-FS) see a drastic drop from $D + D(D-1)/2$ NNs to $K_1 + K_2$. NBM-FS and NB$^2$M-FS benefit similarly in MACs and memory. Empirical throughput measurements on NVIDIA A100 GPU demonstrate scalability of NAM-FS and NBM-FS: throughput remains efficient even with tens of thousands of features, whereas vanilla NA$K_1$0M and NB$K_1$1M fail to execute beyond hundreds of features due to prohibitive costs.

Figure 3: Throughput (inputs processed per second) for NAM/NAM-FS/NA$K_1$2M/NA$K_1$3M-FS across increasing feature dimensionality D.

Experimental Results

Experiments were conducted on six high-dimensional classification datasets (up to $K_1$4). NAM-FS and NBM-FS, with $K_1$5, consistently matched or surpassed vanilla NAM/NBM and other state-of-the-art interpretable models (EBM, NODE-GAM) in test accuracy while achieving superior computational efficiency. The feature selection mechanism was evaluated against pre-selection via mutual information; the integrated selection yielded markedly superior results, especially in low $K_1$6, underscoring the efficacy of joint optimization. Models with learned feature interactions (NA$K_1$7M-FS, NB$K_1$8M-FS, up to $K_1$9) exhibited top performance among all interpretable methods on multiple benchmarks and enabled explicit interaction modeling in high-dimensional spaces unattainable by vanilla architectures.

Practical and Theoretical Implications

The proposed feature selection mechanism provides a scalable pathway for interpretable neural models in high-dimensional settings, critical for applications in genomics, sensor networks, and domains with dense tabular features. The approach maintains global interpretability (via visualizable shape functions and interaction heatmaps) and enables efficient learning of relevant features and interactions, dynamically balancing model complexity and computational cost through user-parametric $K_2$0, $K_2$1. Theoretically, differentiation via entmax relaxations advances the landscape of sparse neural model optimization and paves the way for efficient learning of higher-order interactions, albeit at a cost of interpretability. The empirical results demonstrate that the integrated feature selection outperforms conventional approaches (pre-selection), supporting end-to-end feature identification as a superior strategy.

Future Research Directions

Potential improvements include the extension to higher-order interactions beyond pairwise, though this will compromise interpretability. Incorporating sparse feature support—especially for applications involving dense versus sparse features—may further enhance scalability. Deployment to real-world biomedical, legal, or industrial datasets will validate robustness and practical relevance. Further optimization of selection mechanisms and integration with other interpretability tools may drive adoption in high-reliability domains.

Conclusion

This paper advances interpretable neural models by embedding differentiable feature selection into NAM and NBM, enabling scalable and efficient modeling with feature interactions in high-dimensional spaces. The approach yields strong empirical performance and preserves interpretability, providing a critical tool for applications requiring transparency and computational scalability (2606.19850).