Research map / Architectures and efficiency / Attention and transformers
TabFlex: Scaling Tabular Learning to Millions with Linear Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Replaces TabPFN's softmax attention with non-causal linear attention to scale in-context tabular classification to millions of samples and thousands of features, via three specialized models selected by dataset shape.
Problem
TabPFN's quadratic softmax attention limits its use on large, high-dimensional, or many-class tabular datasets, hurting both pretraining and inference scalability.
Approach
Empirically shows that causal models (SSMs/Mamba) hurt in-context learning for tabular data while non-causal linear attention keeps accuracy, then builds TabFlex with linear attention and three checkpoints (S100, L100, H1K) chosen by a conditional model-selection rule.
Key results
- On 57 small datasets, TabFlex matches TabPFN's median AUC (0.96 vs 0.97) at 0.29 vs 0.82 s per 1000 instances, a >2x speedup.
- On poker-hand (1M+ samples), TabFlex reaches 0.84 AUC in 4.88 s, versus 0.72 AUC for TabPFN at 15.36 s.
- On Bioresponse, TabFlex attains 0.75 AUC versus 0.50 for TabPFN, while being slower (12.38 s vs 1.29 s).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace softmax attention with non-causal linear attention (ablation) | TabPFN · 57 real tabular datasets | Accuracy and runtime | Helped |
| Use causal attention instead of non-causal attention (ablation) | TabPFN · Synthetic tabular ICL datasets | Classification accuracy as training samples increase | Hurt |
| Replace Transformer layers with Mamba layers (ablation) | TabPFN · 150 validation datasets | Mean AUC | Hurt |
| Replace linear attention with sliding-window attention (ablation) | TabFlex · Poker-Hand | Performance: 0.48 -> 0.84 | Hurt |
| Replace linear attention with sliding-window attention (ablation) | TabFlex · Airlines | Performance: 0.48 -> 0.64 | Hurt |
| Replace linear attention with sliding-window attention (ablation) | TabFlex · Higgs | Performance: 0.39 -> 0.76 | Hurt |
| Reduce feature dimensions using PCA, SVD, or random projection (ablation) | TabFlex · Datasets with more than 100 features | AUC: 10% to reduce latency more than 2 times | No clear effect |
| Use 20% of training data via sampling (ablation) | TabFlex · Datasets with more than 1000 instances and fewer than 100 features | Classification performance: 20% | No clear effect |
| Use non-causal linear attention rather than softmax attention (ablation) | TabPFN · Tabular classification | Runtime | Helped |
Surprising findings
- Causal (Mamba-style) attention underperforms non-causal attention for tabular ICL, with accuracy declining as more training samples are added in the causal setting.
- Sliding-window attention performs much worse than linear attention on large datasets (e.g., poker-hand 0.48 vs 0.84 AUC).
Limitations
- Scaling to feature spaces beyond about two thousand features is not addressed.
- Regression is handled only by naive discretization into bins, and performance is weak on several regression datasets.
- Only classification is the primary focus; image experiments are limited to small-scale datasets such as MNIST and CIFAR-10.
Experimental setup
Datasets: OpenML tabular datasets TabZilla hard benchmark MNIST Fashion-MNIST CIFAR-10
Benchmarks: OpenML 98 simple datasets 57 small datasets (TabPFN Table 2 set) TabZilla hard benchmark (36 datasets) poker-hand regression datasets
Metrics: AUC runtime per 1000 instances latency
Baselines: TabPFN XGBoost CatBoost LightGBM RandomForest TuneTables HyperFast SAINT FTTransformer ResNet
Builds on: TabPFN linear attention Mamba FlashAttention TuneTables
Closest papers by meaning
- TabICL: A Tabular Foundation Model for In-Context Learning on Large Data (ICML 2025 · similarity 0.91)
- TabDPT: Scaling Tabular Foundation Models on Real Data (NeurIPS 2025 · similarity 0.91)
- Large Scale Transfer Learning for Tabular Data via Language Modeling (NeurIPS 2024 · similarity 0.91)
- Retrieval & Fine-Tuning for In-Context Tabular Models (NeurIPS 2024 · similarity 0.91)
- TabPFN Unleashed: A Scalable and Effective Solution to Tabular Classification Problems (ICML 2025 · similarity 0.90)
- ConTextTab: A Semantics-Aware Tabular In-Context Learner (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.