A Neural Network Alternative to Tree-based Models

TL;DR

This paper introduces sTAB-Net, a neural network model with built-in sparsity via attention mechanisms, outperforming tree-based models on biological tabular datasets and enabling better interpretability and insights.

Contribution

The paper presents a novel neural network architecture that incorporates sparsity through attention mechanisms, improving performance and interpretability on tabular data.

Findings

sTAB-Net outperforms traditional tree-based models on biological datasets.

The model allows for direct feature importance extraction.

sTAB-Net achieves state-of-the-art results, surpassing post-hoc interpretability methods.

Abstract

Tabular datasets are widely used in scientific disciplines such as biology. While these disciplines have already adopted AI methods to enhance their findings and analysis, they mainly use tree-based methods due to their interpretability. At the same time, artificial neural networks have been shown to offer superior flexibility and depth for rich and complex non-tabular problems, but they are falling behind tree-based models for tabular data in terms of performance and interpretability. Although sparsity has been shown to improve the interpretability and performance of ANN models for complex non-tabular datasets, enforcing sparsity structurally and formatively for tabular data before training the model, remains an open question. To address this question, we establish a method that infuses sparsity in neural networks by utilising attention mechanisms to capture the features' importance in tabular datasets. We show that our models, Sparse TABular NET or sTAB-Net with attention mechanisms, are more effective than tree-based models, reaching the state-of-the-art on biological datasets. They further permit the extraction of insights from these datasets and achieve better performance than post-hoc methods like SHAP.