TabFlex: Scaling Tabular Learning to Millions with Linear Attention
This work addresses scalability issues for researchers and practitioners using tabular data with millions of samples, offering an incremental improvement over existing methods.
The paper tackles the problem of scaling tabular learning to large datasets by enhancing TabPFN with linear attention mechanisms, resulting in TabFlex, which processes a million-sample dataset in 5 seconds and achieves over 2x speedup compared to TabPFN and 1.5x over XGBoost.
Leveraging the in-context learning (ICL) capability of Large Language Models (LLMs) for tabular classification has gained significant attention for its training-free adaptability across diverse datasets. Recent advancements, like TabPFN, excel in small-scale tabular datasets but struggle to scale for large and complex datasets. Our work enhances the efficiency and scalability of TabPFN for larger datasets by incorporating linear attention mechanisms as a scalable alternative to complexity-quadratic self-attention. Our model, TabFlex, efficiently handles tabular datasets with thousands of features and hundreds of classes, scaling seamlessly to millions of samples. For instance, TabFlex processes the poker-hand dataset with over a million samples in just 5 seconds. Our extensive evaluations demonstrate that TabFlex can achieve over a 2x speedup compared to TabPFN and a 1.5x speedup over XGBoost, outperforming 25 tested baselines in terms of efficiency across a diverse range of datasets. Furthermore, TabFlex remains highly effective on large-scale datasets, delivering strong performance with significantly reduced computational costs, especially when combined with data-efficient techniques such as dimensionality reduction and data sampling.