🤖 machine learning

Neural Additive and Basis Models with Feature Selection and Interactions

This paper proposes integrating a trainable feature selection layer into Neural Additive and Basis Models to overcome computational bottlenecks in high-dimensional settings, thereby enabling efficient modeling of feature interactions while maintaining interpretability and competitive performance.

Original authors: Yasutoshi Kishimoto, Kota Yamanishi, Takuya Matsuda, Shinichi Shirakawa

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yasutoshi Kishimoto, Kota Yamanishi, Takuya Matsuda, Shinichi Shirakawa

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the outcome of a complex event, like whether a patient will recover or if a loan will be approved. You have a massive list of clues (features) to help you decide.

The Problem: The "Black Box" vs. The "Too-Heavy" Box
In the world of Artificial Intelligence, there are two main ways to solve these puzzles:

The Black Box (Deep Neural Networks): These are incredibly smart and accurate, but they work like a magic trick. You put data in, and an answer comes out, but no one knows how the machine made that decision. In fields like medicine or law, this is dangerous because you need to trust the reasoning.
The Transparent Box (Generalized Additive Models): These are like a team of specialists. Each specialist looks at just one clue (e.g., "age" or "income") and gives a score. The final answer is just the sum of all these scores. This is easy to understand because you can see exactly how each clue contributed. However, traditional transparent boxes are often not smart enough to handle complex relationships between clues.

The Existing Solution: NAM and NBM
Researchers created "Neural Additive Models" (NAM) and "Neural Basis Models" (NBM). Think of these as upgrading the "Transparent Box." Instead of simple specialists, they use tiny, smart AI brains (Neural Networks) to look at each clue. This makes them both smart and transparent. You can still see how each clue affects the result.

The Bottleneck: The "Too Many Clues" Crisis
Here is where the old method breaks down.

The Interaction Problem: Sometimes, clues work better together. For example, "Age" and "Income" might predict something better when looked at together rather than separately. To do this, the old models tried to create a tiny AI brain for every possible pair of clues.
The Math Nightmare: If you have 1,000 clues, there are nearly 500,000 possible pairs. Trying to train 500,000 tiny AI brains at once is like trying to hire 500,000 employees to work in a tiny office. It crashes the computer, takes forever, and becomes impossible to manage.
The High-Dimensional Wall: When datasets get huge (thousands of clues), the old models simply stop working. They run out of memory and time.

The New Solution: The "Smart Selector"
The authors of this paper propose a simple but powerful fix: Feature Selection.

Imagine you are a detective with a massive evidence board. Instead of trying to analyze every single piece of evidence at once (which is impossible), you hire a Smart Selector.

How it works: The model starts with a "menu" of all possible clues. During training, it learns to assign a "weight" to each clue.
The Filter: It uses a special mathematical filter (called entmax) that acts like a sieve. It keeps the most important clues and effectively "turns off" the useless ones.
The Result: Instead of trying to build 500,000 AI brains for every pair of clues, the model only builds brains for the top 50 or 500 most important clues and pairs.

The Benefits

Speed: Because the model ignores the junk data, it runs much faster. The paper shows that while the old models crash on datasets with 1,000+ clues, the new models (called NAM-FS and NBM-FS) handle them smoothly.
Interactions: It allows the model to look at pairs of clues working together (like "Age + Income") without getting overwhelmed by the math.
Transparency: Even though it's selecting the best clues, it remains transparent. You can still see exactly which clues were chosen and how they influenced the final decision.

The Results
The authors tested this on six different real-world datasets with thousands of clues.

Performance: Their new models were just as accurate, or even better, than the best existing "transparent" models.
Efficiency: They were significantly faster and could handle data sizes that the old models couldn't even touch.
Comparison: They also proved that letting the model learn which clues to pick during training is better than trying to pick the clues manually before starting.

In Summary
This paper introduces a way to make smart, transparent AI models that can handle massive amounts of data. It does this by teaching the AI to ignore the noise and focus only on the most important clues and their relationships, solving the problem of models that are either too dumb to be useful or too heavy to run.

Problem
Deep neural networks (DNNs) offer high representation capabilities but often lack interpretability, limiting their use in high-reliability fields like medicine and law. While Generalized Additive Models (GAMs) and their neural variants, the Neural Additive Model (NAM) and Neural Basis Model (NBM), provide a balance of interpretability and performance, they face significant scalability issues. Standard NAM and NBM architectures assign a separate neural network (NN) to every input feature. When extended to capture pairwise feature interactions (creating GA2Ms), the naive approach requires preparing shape functions for all possible feature pairs. This results in a computational complexity of $O(D^2)$ , where $D$ is the number of features. Consequently, training these models on high-dimensional datasets (e.g., $D > 1,000$ ) becomes intractable due to excessive memory consumption and computational costs. Furthermore, managing and visualizing thousands of shape functions compromises the interpretability that GAMs are designed to provide.

Methodology
The authors propose incorporating a differentiable feature selection mechanism directly into the training of NAM and NBM to resolve these bottlenecks. The proposed models, termed NAM-FS and NBM-FS, introduce a feature selection layer that dynamically selects a predefined number of single features ( $K_1$ ) and feature pairs ( $K_2$ ) during the training process.

Key technical components include:

Differentiable Selection: Instead of discrete feature selection, which is non-differentiable, the authors utilize entmax (a generalized version of softmax and sparsemax) to relax the selection process. This allows the model to optimize feature selection weights via stochastic gradient descent (SGD) in an end-to-end manner.
Temperature Annealing: A temperature parameter $\tau$ is introduced and annealed during training. As $\tau$ decreases, the entmax output approaches a one-hot vector, effectively fixing the selected features for the shape functions.
Architecture: The models replace the full set of $D$ (or $D(D-1)/2$ ) shape functions with a fixed, smaller set of $K_1$ univariate and $K_2$ bivariate NNs. The input to these NNs is a weighted combination of the original features determined by the learned selection weights.
Implementation Optimization: The authors optimized the implementation of the NAM layers using torch.bmm (batch matrix-matrix product) in PyTorch, achieving a 2–6x speedup compared to previous 1-D convolution implementations.

Key Contributions

Feature Selection Integration: The paper introduces a simple, end-to-end feature selection mechanism for NAM and NBM, enabling their application to high-dimensional datasets where vanilla versions fail.
Scalability with Interactions: The method allows for the efficient incorporation of pairwise feature interactions (GA2M style) even in high-dimensional settings by limiting the number of interaction terms to a manageable $K_2$ .
Computational Efficiency: The proposed models significantly reduce the number of trainable parameters and multiply-accumulate operations (MACs), resulting in higher throughput compared to vanilla NAM, NBM, and their interaction-based variants (NA2M, NB2M).
Performance Validation: The authors demonstrate that NAM-FS and NBM-FS achieve better or comparable performance to state-of-the-art GAMs and GA2Ms on high-dimensional classification tasks, while outperforming vanilla NAM/NBM that rely on pre-selected features.

Results
Experiments were conducted on six high-dimensional classification datasets (ranging from 561 to 5,000 features).

Throughput: Vanilla NA2M and NB2M could not run on datasets with more than 100 features due to memory constraints. In contrast, NAM-FS and NBM-FS maintained stable throughput even with up to 50,000 features.
Accuracy: On high-dimensional datasets (e.g., guillermo with 4,296 features and Gisette with 5,000 features), NAM-FS(500) and NBM-FS(500) achieved the highest accuracy among all GAMs tested.
Interactions: Models incorporating pairwise interactions (NA2M-FS and NB2M-FS) with $K_2=500$ outperformed state-of-the-art GA2Ms (EB2M and NODE-GA2M) on several datasets.
Feature Selection Efficacy: When compared to NAM/NBM trained on features pre-selected via mutual information, the proposed end-to-end feature selection methods generally yielded higher accuracy, demonstrating the effectiveness of learning feature selection jointly with the model parameters.

Significance and Claims
The paper claims that by integrating feature selection into the training loop, NAM and NBM can be scaled to high-dimensional datasets without sacrificing their core advantage: interpretability. The authors emphasize that their method is simple yet effective, allowing for the visualization of a manageable number of shape functions while capturing complex feature interactions. They note that while the models are highly effective for high-dimensional data, their advantages diminish in low-dimensional settings (fewer than 100 features). The work positions these models as a practical solution for applications requiring both high predictive accuracy and global interpretability in complex, high-dimensional domains.

More like this