[Paper Review] MambaTab: A Plug-and-Play Model for Learning Tabular Data

Md Atik Ahamed, Qiang Cheng|arXiv (Cornell University)|Jan 16, 2024

Data Quality and ManagementDecision Sciences3 citations

TL;DR

MambaTab proposes a plug-and-play, lightweight model for tabular data using Mamba, a structured state-space model (SSM) variant, enabling end-to-end supervised learning with minimal preprocessing. It achieves state-of-the-art performance across 8 benchmark datasets using less than 1% of the parameters of Transformer-based baselines, while supporting efficient feature incremental learning and linear parameter scaling.

ABSTRACT

Despite the prevalence of images and texts in machine learning, tabular data remains widely used across various domains. Existing deep learning models, such as convolutional neural networks and transformers, perform well however demand extensive preprocessing and tuning limiting accessibility and scalability. This work introduces an innovative approach based on a structured state-space model (SSM), MambaTab, for tabular data. SSMs have strong capabilities for efficiently extracting effective representations from data with long-range dependencies. MambaTab leverages Mamba, an emerging SSM variant, for end-to-end supervised learning on tables. Compared to state-of-the-art baselines, MambaTab delivers superior performance while requiring significantly fewer parameters, as empirically validated on diverse benchmark datasets. MambaTab's efficiency, scalability, generalizability, and predictive gains signify it as a lightweight, "plug-and-play" solution for diverse tabular data with promise for enabling wider practical applications.

Motivation & Objective

• To address the high computational and preprocessing demands of existing deep learning models for tabular data.
• To enable efficient, scalable, and generalizable learning on tabular data with minimal hyperparameter tuning.
• To support feature incremental learning, where new features are added sequentially without retraining on old data.
• To develop a lightweight, out-of-the-box solution that maintains high performance across diverse tabular datasets.
• To demonstrate that structured state-space models (SSMs), particularly Mamba, can effectively model tabular data with strong representation learning.

Proposed method

• MambaTab employs Mamba, an efficient SSM variant, as the core architecture for modeling tabular sequences.
• It uses residual connections to stack multiple Mamba blocks, enabling deeper representations with linear parameter growth.
• The model applies layer normalization to embedded representations to stabilize training and improve performance.
• It performs end-to-end training with minimal data preprocessing, including no need for one-hot encoding or feature scaling in most cases.
• The architecture supports both vanilla supervised learning and feature incremental learning, where new features are added without discarding prior data.
• The output head is flexible and can be adapted for both classification and future regression tasks.

Experimental results

Research questions

• RQ1 Can a structured state-space model (SSM) like Mamba achieve superior performance on tabular data compared to Transformers and CNNs?
• RQ2 Does MambaTab maintain high performance with significantly fewer parameters than state-of-the-art models?
• RQ3 Can MambaTab effectively handle feature incremental learning, where features are added sequentially?
• RQ4 How does the model perform across diverse tabular datasets with minimal data preprocessing?
• RQ5 What is the impact of architectural components like layer normalization and batch size on model stability and performance?

Key findings

• MambaTab outperforms state-of-the-art baselines, including Transformer-based models, on 8 public tabular datasets in both vanilla supervised and feature incremental learning settings.
• The model achieves superior AUROC scores—e.g., 0.771 on CG and 0.862 on CB datasets—when layer normalization is applied, demonstrating its effectiveness.
• MambaTab uses less than 1% of the parameters of comparable Transformer-based models while maintaining or exceeding their performance.
• The model shows minimal sensitivity to batch size, with stable performance across batch sizes from 60 to 140, indicating strong generalization.
• Scaling Mamba blocks linearly increases parameter count but maintains consistent performance, showing strong information retention and scalability.
• Ablation studies confirm that layer normalization improves performance by 0.012 AUROC on average, justifying its inclusion in the architecture.