# TabFlex: Scaling Tabular Learning to Millions with Linear Attention

arXiv: https://arxiv.org/abs/2506.05584 · HTML: https://mlsanity.com/paper/2506.05584/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 79 references that are in the map.
- First posted on arXiv 5 Jun 2025.

## Builds on (most cited first, 30 of 51)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [DeepFM: A Factorization-Machine based Neural Network for CTR Prediction](https://mlsanity.com/paper/1703.04247/) (2017)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [TabNet: Attentive Interpretable Tabular Learning](https://mlsanity.com/paper/1908.07442/) (2019)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models](https://mlsanity.com/paper/2311.07919/) (2023)
- [Qwen2-Audio Technical Report](https://mlsanity.com/paper/2407.10759/) (2024)
- [OpenML: networked science in machine learning](https://mlsanity.com/paper/1407.7722/) (2014)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [TabTransformer: Tabular Data Modeling Using Contextual Embeddings](https://mlsanity.com/paper/2012.06678/) (2020)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [Neural Additive Models: Interpretable Machine Learning with Neural Nets](https://mlsanity.com/paper/2004.13912/) (2020)

## Built on by (most cited first, 15 of 15)

- [TabPFN-Wide: Continued Pre-Training for Extreme Feature Counts](https://mlsanity.com/paper/2510.06162/) (2025)
- [ICR-RL: Deep Reinforcement Learning via In-Context Regression](https://mlsanity.com/paper/2509.11259/) (2025)
- [End-to-End Compression for Tabular Foundation Models](https://mlsanity.com/paper/2602.05649/) (2026)
- [One Pipeline, Many Transformers: Pattern-Specific Imputation Specialists for Tabular Missing Data](https://mlsanity.com/paper/2510.02625/) (2025)
- [State-Space Models for Tabular Prior-Data Fitted Networks](https://mlsanity.com/paper/2510.14573/) (2025)
- [Mitigating Label Shift in Tabular In-Context Learning via Test-Time Posterior Adjustment](https://mlsanity.com/paper/2605.04363/) (2026)
- [CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching](https://mlsanity.com/paper/2606.11473/) (2026)
- [When Tabular Foundation Models Meet Strategic Tabular Data: A Prior Alignment Approach](https://mlsanity.com/paper/2605.19662/) (2026)
- [Privacy Vulnerabilities of Attention Layers in Tabular Foundation Models and Protection of High-Risk Queries](https://mlsanity.com/paper/2606.26021/) (2026)
- [Memory Efficient Tabular Foundation Models](https://mlsanity.com/paper/2607.27546/) (2026)
- [Are Tabular Foundation Models Robust to Realistic Query Distribution Shifts in Microbiome Data?](https://mlsanity.com/paper/2606.24995/) (2026)
- [Lost in Aggregation: How Benchmarks Overlook Irreplaceable Model Strengths](https://mlsanity.com/paper/2608.18919/) (2026)
- [Xiaomi-TabLDM: A Tabular Foundation Model Technical Report](https://mlsanity.com/paper/2609.03880/) (2026)
- [Adapting Linear-Time Architectures for Tabular In-Context Learning](https://mlsanity.com/paper/2609.36337/) (2026)
- [GIST: Generator-state-Informed Supervision for Tabular Foundation Models](https://mlsanity.com/paper/2610.12904/) (2026)

## Abstract (arXiv)

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.