TabFlex: Scaling Tabular Learning to Millions with Linear Attention
Authors: Yuchen Zeng
Tuan Dinh
Wonjun Kang
Andreas C Mueller
Venue: ICML 2025 spotlightposter (2025)
Keywords: Tabular Classification
Transformer
State-Space Models
Linear Attention
Scalability
Primary area: deep_learning->other_representation_learning
TL;DR: TabFlex: A linear attention-based model for scalable tabular classification, outperforming existing methods in speed and maintaining good performance on both small and large datasets.

Abstract: 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 2× speedup compared to TabPFN and a 1.5× 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.

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Yuchen Zeng 12
Tuan Dinh *3
Wonjun Kang45 45 Andreas C. Müeller
Abstract
2023), from vision (Bai et al., 2023b), audio (Chu al., 2023; 2024) to bio-signals (Wan et al., 2023) and protein sequences (Rives et al., 2019; Hayes et al., 2024). Tabu- lar data, as one of the most fundamental and critical data types in real-world applications - including recommenda- tion systems (Zhang et al., 2019), finance (Arun et al,, 2016), and medicine (Johnson et al., 2016) has attracted a great deal of attention and attempts to explore the potential of Transformer-basee models, particularly for tabular classifi- cation (Arik & Pfister, 2021: Hollmann et al., 2023; Huang et al., 2020; Dinh et al., 2022: Gorishniy et al., 2021). For example, the FT transformer (Gorishniy et al., 2021) con- verts each sample into a sequence of embeddings to use the transformer to make predictions. Tranfrnsformer (Huang et al., 2020) learns embeddings for categorical features, con- catenating them with continuous features. On the other hand, LIFT (Dinh et al., 2022) converts tabular data com- bined with feature names and task descriptions into textual sentences as input to LLMs. In particular, compared to tra- ditional methods for tabular data such as gradient- methods
Leveraging the in-context learning (ICL) capabil- ity of Large Language Models (LLMs) for tabu- lar classification has gained significant attention for its training-free free adaptability across diverse datasets. Recent advances, such as TABPFN, ex- cel in tabular small-scale 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 al- ternative complexity-quadie Our model, TABFLEX, efficiently handles tabu- lar datasets with thousands of features and hun- dreds of classes, seamlessly scaling to millions of samples. For instance. TABFLEX processes the poker-hand hand dataset with more than a million samples just 5 seconds. Our extensive evalua- tions demonstrate that TABFLEX can achieve over 2 speedup compared to TABPEN and 1.5x speedup over XGBoost, outperforming 25 tested baselines in terms of efficiency across a diverse range of datasets. Furthermore, TABFLEX re- mains highly effective large-scale datasets, de- livering strong performance with significantly re- duced computational costs, especially when com- bined with data-efficient techniques such as di- mensionality reduction and data sampling.
self-attention.
trees (Friedman, 2001), these transformer-bast often suffer from high latency overhead for training and inference, primarily due to their larger model sizes.
The recent TABPEN (Hollmann et al., 2023) addresses the latency limitations of caansffrmer-basee methods by utiliz- ing the in-context learning (ICL) capability (Brown et al., 2020) of LLMs for directly learning a new task from ex- amples without parameter updates, attaining superior effi- ciency and performance on small-scale datasets. In particu- lar, TABPFN incorporates all training and testing samples into single prompt and classifies the testing samples in one forward pass, making it highly efficient and effective on simple and small tabular datasets. However, TABPEN faces challenges with complex datasets that require large sample sizes for effective learning, primarily due to the seal- ability limitations imposed by the quadratic complexity of the attention mechanism, introducing difficulties for both the scalable pretraining and inference processes.
1. Introduction
Enhancing the applicability of the Transformer architec- ture (Vaswani et al., 2017) for diverse data modalities be- yond textual data and non-language tasks (Achiam et al., 2023; Brown et al., 2020; Bai et al., 2023a; Dubey et al., 2024) has achieved remarkable success (Gemini Team et al.,
'Equal contribution 'Work done during an internship at the Gray Systems Lab, Microsoft Uniieesiity of Musconsin-Maaiion P University of California San Francisco 'Furiosa AI SSouul Na- tional University 'Gray System Lab, Microsoft. Correspondence to: Andreas C. Müeller <amueller@ microsoft.com..
In this work, we address the limitations of the scalability of TABPFN and improve the effectiveness of Transformer- based methods for tabular classification. We first systemati- cally analyze scalable alternatives to attention mechanisms, focusing on state-space models (SSMs) within the Mamba model (Gu & Dao, 2024) and linear attention (Katharopou-
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
los et al., 2020). We find that (Finding 1) the inherent causal- ity of SSMs impedes ICL. performance compared to non- causal mechanisms. In contrast, (Finding 2) linear attention does not suffer from this limitation, maintaining comparable performance with improved computational efficiency. Thus, we develop TABFLEX leveraging linear attention as the attention mechanism, comprising three sub-models where each is optimized for different scenarios and selected based on dataset characteristies (e.g., sample size). TABFLEX efficiently handles tabular datasets with thousands of fea- tures and hundreds of classes, scaling to millions of samples. Via the comprehensive evaluation on a diverse range of datasets, we find that (Finding 3) TABFLEX consistently achieves competitively high performance with impressive computational efficiency compared to 25 baselines, includ- ing TABPEN and XGBoost. Notably, TABFLEX perform highly on poker-hand dataset with 1M+ samples in less than 5 seconds and attains high accuracies on image datasets such as MNIST (LeCun et al., 2010), Fashion-MNIST (Xiao et al., 2017), and CI (Krizhevsky et al., 2009) in less than one second. Furthermore. our ablation studies suggest that TABFLEX can seamlessly incorporate data-efficient techniques such as dimensionality reduction and data sam- pling for more computation reduction.
limited to small tabular classification datasets. To handle it, many concurrent variants are proposed. MixturePFN (Xu et al., 2025) improves scalability by routing new test sam- ples to a pool of scalable prompters using Sparse Mixture of In-Context Prompters, while LoCalPFN (Thomas et al., 2024) proposes retrieving a local subset of task-specific data for efficiently fine-tuning on. Ma et al. (2024) introduce in-context data distillation to optimize TabPFN's context and remove the data size constraint. Tune'Table (Feuer et al., 2024) scales TABPEN to large datasets by performing prefix tuning per dataset. TabPFNv2 (Hollmann et al., 2025) enhances TabPFN's accuracy in low-data regimes (fewer than 10,000 samples), complementing our focus on speed and scalability. Our method is also based on TABPEN, ex- tending its scalability to large datasets while maintaining and improving efficiency by simply replacing the softmax attention with linear attention.
A Attenionn mechanisms and scalable alternatives. As Transformers (Vaswani et al., 2017) face the scaling chal- lenge for long sequences due to the quadratic computational and memory complexity, scalable alternatives have been pro- posed (Orvieto et al., 2023; Sun et al., 2023). While RNNs provide efficient linear-time inference, they struggle with training efficiency and lack the parallelization capabilities of Transformer architectures. Linear attention (Katharopou- los et al., 2020) addresses both concerns by reformulating self-attention as a linear dot-product of kernel feature maps, reducing the computational complexity from quadratic to linear time. Furthermore, causal linear attention can be in- terpreted as a form of RNN, as it predicts based on a current token and a "hidden state," summarizing information from the previous tokens. State-space models (SSMs) address RNNs' drawbacks by considering linear RNNs with novel al algorithms for efficient training (Gu et al., 2021; 2022; Gu & Dao, 2024; Dao & Gu, 2024; Peng et al., 2023; Orvieto et al., 2023; Sun et al., 2023). Dao et al. (2022) identified an- other bottleneck in attention mechanisms' speed stemming from the relatively slow access to bighbhandwithh memory (HBM) in GPUs and proposed FlashAttention (Dao, 2024; Shah et al., 2024) to restructure attention computation to optimize the utilization of high-speed on-chip SRAM while minimizing access to slower HBM, enhancing the efficiency of GPU-based attention operations.
2. Related Works
Transformer-based approaches for tabular classification. The pioneering TabNet (Arik& & Pfister, 2021) applies un- supervised pre-training on masked tabular datasets to infer missing features, enhancing the model's understanding of data and features before supervised learning on feature selec- tion for the final decision boundary. TabTransformer (Huang et al., 2020) proposes handling categorical features by con- catenating their contextual embeddings into numerical fea- tures. FT-Trranformer (Gorishniy et al., 2021) converts samples to embedding sequences using a feature tokenizer for the transformer. LIFT (Dinh et al., 2022) converts each sample into sentence using a predefined template incor- porating the task description and feature names, as the nat- ural input to apply ICL in LLM. TabR (Gorishniy et al., 2024) proposes a carrivall-uugmente model with a custom kNN-like component to retrieve and extract signals from the nearest neighbors. BiSHop (Xu et al., 2024) establishes interconnected directional learning modules to process data column-wise and row-wise for tabular learning. XTab (Zhu et al., 2023) utilizes independent featurizers and federated learning to resolve inconsistent column types and quantities.
See Section A for an extended discussion of related works.
3. Preliminaries
We elucidate key concepts of TABPFN and two prominent scalable attention mechanisms (SSMs and linear attention).
The widely adopted transformer-base approaches for tab- ular CottiaattionTAAPPPI (Hollmann et al., 2023) is trained offline on synthetic datasets derived from previous distributions to perform ICI., allowing efficient inference in small-scale tabular classification tasks. However, it is
Implementation of ICL in TabPEN (Hollmann et al., 2023). Fig. 1 illustrates the design of TABPFN, where each sample is treated as a token, starting with training sam-
2

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
be reformulated as Convolutional Neural Network (CNN) as (3), enabling efficient and parallelizable training. SSMs ad- dress the quadratic time complexity problem w.r.t sequence length, as the output for each new token depends solely on the hidden states and the current token, in contrast to stan- dard attention mechanisms that attend to all previous tokens. Consequently, SSMs operate as a causal mechanism.
Input Feuture Projection
X 12
Label Prediction
MLP
MLP
(122, p22)
MLP
MLP
Attention
MLP
MI
1
MLP
MLP
MLP
Slwap1
h' h'(1) Ah(() + Br((), s(1) -Ch(l)
(1)
MLP
MLP
MLP
h = Ah + Bri y1 = Chi
(2)
1-1 B), K = (CB, CA B....... ...,. ",,
Plotam
MLP
MLP
MLP
Skoii.m
(yI,...;yo) Ve) (Z1....,,...) = K
(3)
Figure 1. Illustration of TABPEN's for classifying the entire dataset in one forward pass. In each layer, attention outputs for training sample positions attend to all other training samples, ensuring that predictions are invariant to the order of training samples. Conversely, attention outputs for test sample positions attend only to training samples, ensuring independent predictions for each test instance, unaffected by other test samples. The final classification for each test sample is derived by applying an MLP to the corresponding Transformer output at its respective position.
Linear attention. Assume a sequence with length n E N+ and embedding size E N We first focus on non-causal cases. For the i-th position, let qi RV k; E Rd, and v; E Rd denote the query, key, and value vectors, respec- tively, where
1...., n. In softmax attention, the sim- ilarity between qi and ky for any i * is computed as exp (q, kj). The attention output at the i-th position, de- noted as a; E R4, is obtained by averaging the values across all tokens weighted by their similarities. This pro- cess requires (n) complexity, as it necessitates comput- ing similarities with all n tokens. Linear attention reduces this complexity by replacing the similarity computation from exp(q, k;) with )(q:) ákk,),, where Φ Rdd + Rd is a feature conversion function. For linear attention out- puts (4) across all positions, we identify two common terms: Σj=1 ø (kz)) and Σj=1 & (kk)), which can be computed once. Consequently, for the linear output at position i, we only need to compute ((q.) and multiply it with these two statistics, resulting in O(1) complexity, thus significantly reducing computational demands.
ples and followed by testing samples. These samples are embedded (features x and labels y for training and only features x for testing samples) with MLPs before being concatenated. Outputs corresponding to training sample po- sitions are computed by attending to all other training sam- ples, while the outputs for test sample positions attend to the training samples- - enabling each test prediction to leverage the full training set without being influenced by other test samples. Test predictions are generated by projecting the Transformer outputs at test positions into probability distri- butions. This implementation is functionally equivalent to standard ICL but significantly more efficient. Standard ICL requires m (number of test samples) separate prompts, each containing all training samples and one test sample, necessi- tating m prediction passes. A notable feature of TABPFN the encoder with non-causal attention, allowing outputs within training sample positions to interact freely, rendering the order of training samples inconsequential.
(Softmax) a, Σj- exp (1, (T]k)) Ug
(4)
ΣØ=1 exp (g, k:))
Φ (q)) Σj=1 1 (ki,)
) U.
(Linear) a; $ (q:) Σ}= Φ (kij)
For causal cases, for position i, we replace Σj=1 with Ey= as each token attends only to previous tokens. The statistics then become Σ ø (kz) Uy and Σ (k;), which can be viewed as hidden states in RNNs. Thus, causal linear attention can be conceptualized as a linear RNN, which is also a variant of SSM.
State-Space Models (SSMs). The SSM framework is based on continuous system that transforms a one- dimensional signal x<t) E IR into y(t) E IR through an intermediate H H-dimensional latent state h(t) E R as shown in (1). Here. B E "x1 is the input transition vector and A E H*H is the state transition matrix. The latent state h(t) is then projected into the output yt) using the output mapping vector C' E RIXH For deep learning applications, discrete A and B replace continuous A and B through discretization methods, such as zero-order hold. This yields updated hidden state and output equations as shown in (2). While (2) is structured as linear RNN, it can
4. Architectural Exploration for Scalable Tabular Learning
We analyze State-Space Models and linear attention as at- tention architecture alternatives to enhance the scalability of TANBPFN focusing on tabular classification with ICI.
3

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Softmax Attention
Linear Attention
Mamba Transformer
0.82
Accuracy
0.80
Non-causal Masked
Accuracy
0.8
Causal Masked
0.78
0.6
0.76
0.4
0.2
0
1000
2000
3000
1.0
0.5
0.0
0.0
0.5
1.0
0.0
0.2
0.4
0.6
Number of Training Samples
Training Loss
Test Mean AUC
Runtime
(a) Effeet of causal masking. When more samples are provided, the non-causal model
(b) ICL performance comparison between Mamba and Transformer models. Re- sults show that Transformer-base models achieve lower training loss and higher AUC across 150 test datasets.
(c) Accuracy and runtime comparison of softmax and linear attention. Replacing softmax with linear attention preserves com- parable accuracy while significantly reduc- ing runtime.
shows better sample utilization and accu- racy, while the causal model's 's performance plateaus early and declines.
Figure Impact of model architecture on tabular classification performance. Please refer to Fig. 8 for a detailed breakdown.
4.1. Causal Model vs. Non-Causal Model
scenario. Our observation is supported by empirical stud- ies (Ding et al., 2024; Gong et al., 2023), which show that causal attention is suboptimal for ICL. Moreover, most the- oretical analyses of ICL assume non-causal attention (Ahn et al., 2023; Bai al., 2023c).
Ideally, the order of training samples (i.c., in-context demon- strations) in the prompt should not affect the final prediction. However, SSMs are inherently causal, computing outputs based on new inputs and hidden states derived from pre- vious inputs. This suggests potential drawback in this context. To validate our hypothesis regarding the subopti- mal performance of causal models in ICI., we conduct two experiments: (i) comparing the performance of TABPFN with a modified version of the same model that uses causal attention, and (ii) evaluating TABPFN against both its origi- nal version and model incorporating Mamba (specifically Mamba-II), a leading SSM-based architecture.
Mamba vs. Transformer. We further investigate whether Mamba, the most popular SSM- model, is suitable for ICL. We replicate TABPFN's training methodology, sub- stituting the transformer layer with an Mamba layer. To evaluate performance, we test the modified model on the same 150 validation datasets used in the TABPEN study (See Section F.3 for details). Fig. 2b visualizes the training loss and test mean AUC for both methods. The model with Mamba exhibits significantly higher training loss than the original TABPFN, with substantially lower test mean AUC. This experiment with a popular SSM model further demon- strates that SSMs underperform non-causal models in our specified tasks.
Causal Attention vs. Non-Causal Attention. Our first experiment compares the ICL capabilities of non-causal and causal attention mechanisms using the same experi- mental setup as TABPEN, shown in Fig. 2a. We replicate TABPFN's 's methodology for generating synthetic datasets from priors, training a modified version employing causal attention instead. For the inference, we generate 20 syn- thetic datasets that maintain a consistent 1000 test samples with varying numbers of training samples. We average the classification accuracy across 20 simulations.
4.2. Softmax Attention vs. Linear Attention
To address the quadratic complexity of standard attention mechanisms. linear attention has emerged as a popular alter- native (Katharopoulos al., 2020). To investigate its impact ICL in tabular classification, we replaced TABPFN's at- tention mechanism with linear attention and trained a model following the same strategy as TABPFN. We then evaluated both TABPFN and this linear attention model on 57 real datasets (used in Table 2 of McElfresh et al. (2023), where TABPFN achieved top performance among 19 methods for tabular classification). Fig. 2c visualizes the test accuracy and runtime. Our results demonstrate that linear attention does not decrease performance and significantly improves speed, making it a suitable method for scaling TABPFN to larger datasets. Finally, in Section 1. we demonstrate that linear attention significantly outperforms sliding window attention (Beltagy et al., 2020) in our setting.
We observe that non-causal attention generally outperforms causal attention. As more training samples are given, the accuracy of the non-causal model continues to improve. In contrast, the causal attention model shows accuracy im- provements only within very small range of training sam- ples, after which performance begins to decline with addi- tional samples. These findings indicate that TABPFN with non-causal attention functions as an effective ICL model, adeptly leveraging context from a large number of samples. Conversely, the same model equipped with causal attention fails to capitalize on the additional data, highlighting the su- periority of the non-causal approach in this tabular learning
4

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
5. TABFLEX: Scaling TabPFN for Large Datasets
Algorithm 1 Conditional Model Selection
Input: : A dataset D with a instances of d features
// Large dataset with few features if n 2 3K and I 100
Based on the empirical findings presented in Section 4, we identify non-causal linear attention as the optimal candidate to replace standard softmax attention in TABPEN. This sec- tion proceeds in two parts: first, we conduct a thorough anal- ysis the linear attention mechanism to ensure its efficient implementation.; subsequently, we leverage this efficient im- plementation train our proposed model, FABFLE Our approach aims to enhance the scalability and performance of tabular learning while maintaining computational efficiency.
then
return TABFLEX-LI00(D)
// High-dimensional datasets else if d > 100 or (d/n 2 0.2
and n 2 3K) then
if S 1000 then
return TABFLEX-HIK(D)
else
Apply random projection DD to reduce the number of features to 1000, yielding D' return TABFLEX- HIK(D')
Computation Analysis. Dao et al. (2022) demonstrates that significant wall-clock speedup for softmax attention can be achieved by optimizing the number of memory read- s/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. Based on this criterion, Yang et al. (2024) proposed PlashLinerAAttention for speeding up causal linear attention. This raises a natural question: can we further improve the speed of non-causal linear attention (we omit non-causal when it does not cause further confu- sion) by reducing the number of memory reads/writes? Our results in Theorem analyze the #HBM access and HBM memory usage of innidLLinaAAAtteeee concluding that further optimization is not necessary. In Section C.I, we first propose HBM-efficient linear atten- tion, and then show that the Py Torch implementation only incurs a marginal increase in terms of #HBM access and HBM memory usage, with FLOPS remaining unchanged. We provide more details, including the analysis of different attention mechanisms and actual memory usage and runtime visualization of these mechanisms in Section C.I. The re- sulting theorem below demonstrates that the straightforward PyTorch implementation of linear attention achieves linear HBM access, matching the performance of Ant tention after optimization. Consequently, our model adopts the straightforward implementation of linear attention.
// Small datasets else
return TABFLLEX-S1001D)
Mean Runtime
Mean AUC
20
0.75
0
0.70
TabPFN
TabFlex
TabPFN
TabFlex
and linear attention,
Figure 3. Runtime and AUC comparison of TABPFN and TABFLEX on validation datasets.
across a broader spectrum of datasets. Since models trained on large, nigh-dimensional datasets often struggle in small regions due to optimization issues, we introduce three spe- cialized models to address this limitation.
TABFLEX-S100 is trained on prompts with 1 1152 length (same as TABPFN), 100 features, and 10 classes. This is optimized for low-dimensional datasets. 'S' denotes standard configuration, '100' indicates feature capacity.
Theorem (High Bandwidth Memory Efficiency of Linear Attention). Let Q, K, V NxD represent the query, key, and value matrices for single attention head, where N is the sequence length and D is the embedding size. With any element-wise kernel feature mapping (e.g., elu(-) + 1), both causal Flash (Alg. 2) and non-cause linear attention (Listing 1) require O(ND) HBM accesses, O) HBM memory, and O(ND2) FLOPS to compute the attention output.
TABFLEX-L100 utilizes prompts of 50K length, 100 fea- tures, and 10 classes. This is designed for large, low- dimensional datasets. 'L' signifies larger sample size, and '100' represents feature count.
TABFLEX-HIK employs prompts of 50K length, 1K features, and 100 classes. This is suited for large, high- dimensional datasets. 'H' indicates high-dimensional ca- pabilities, and 'IK' denotes IK features.
Additional training details, including training loss, hy- perparameters, and other relevant information, are pro- vided in Section C.2. Our code is available at https: / //github. htw.:///co.mmmcccooo///tiii
TABFLEX. While TABPEN excels on small, simple datasets with fewer than 100 features and 10 classes, it struggles with more complex tasks, e.g., high-dimensional datasets or those with numerous classes. We aim to extend the use cases by training a model that maintains comparable speed to TABPEN while offering reasonable performance
We apply the conditional model selection strategy, shown in Alg. 1, select the model based on the target dataset's size
5

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
and dimensionality, ensuring optimal performance across di- verse data characteristics. The decision thresholds align with the training regimes of the models. TABFLEX-SI sharing TabPFN's training setup but with an updated architecture, is deployed similarly (n I 3K, d < 100). TABFLEX-L_100, trained on low-dimensional (d Σ 100) but larger datasets, is used for longer sequences (n 2 3K, d S 100). TABFLEX- HIK, trained on high-dimensional data, is assigned to han-
1.0
LightGBM
XGBoost
TabFlex
RandomForest TabPFN
Median AUC
0.8
LinearModel
DeclsionTree
KNN
0.6
Dataset Coverage Completed all datasets Completed subset of datasets
dle those cases accordingly. We note that performance is not highly sensitive to the chosen decision boundaries, sup- ported by our results in Section C.4.
0.4
0.0
0.1
0.2 Time per 1000 instances
0.3
0.4
0.5
In Fig. 3, we visualize the mean runtime and mean AUC comparison of TABPEN and TABFLEX on the validation datasets, comprising 40 datasets with varying sample sizes (up 100K), dimensions (up 3K), and number of classes (up 100). Detailed information about these datasets is pro- vided in Section C.3. Our analysis reveals that TABFLEEX not only exhibits superior performance but also demonstrates faster execution times compared to TABPEN.
Figure 4. Visualization of tested methods with processing times under 0.5 seconds per 1000 instances on the TabZilla hard benchmark. We report the median AUC across the completed datasets, several methods completed only subset of the datasets. Compared to other methods (XGBoost and TABPFN) that success- fully ran on all datasets, TABFLEX achieves a 2x speedup while maintaining relatively good performance.
6. Performance Evaluation of TABFLEX
tion: TuneTables (Feuer et al., 2024), a TF method, and HyperFast (Bonet et al., 2024), a Non-TF NN method.
We evaluate TABFLEX'S performance and speed across 1 115 OpenML tabular datasets (Vanschoren et al., 2013). See Section D.2 for the complete list of baselines, and detailed description of the models' implementation.
6.2. Evaluation on Simple Datasets
We evaluate models on two sets of data: 98 simple datasets from Table 1 and 57 small datasets from Table of McEI- fresh et al. (2023). The results are reported in Table 11 (Section D.3) and Table 1, respectively. For each dataset, we consider ten different train/test splits, computing the score mean and standard deviation, as well as the total run- time per 1000 instances. We then calculate the median and mean of these values across the entire set of datasets: 98 sim- ple datasets for Table 11 and 57 small datasets for Table I. Algorithms are ranked based on AUC and time. Our results demonstrate that TABFLEX achieves nearly identical perfor- mance to TABPEN on small, simple datasets while offering more than a 2x speedup. Compared to faster methods, such as Decision Tree and Linear Model in Table 1 11, and Deci- sion Tree, Linear Model, LightGBM, and KNN in Table 1, their performance significantly inferior to TABFLEX.
6.1. Experimental Setup
Unless otherwise stated, we follow the identical experiment setup of McElfresh et al. (2023) for baseline benchmarking.
Datasets. For classification tasks, we consider datasets of two difficulty levels. For simpler tasks, we use two collec- tions of datasets-98 and 57 total- oorignalyyy reported in Tables and of McElfresh et al. (2023), which are char- acterized by smaller sample sizes and lower feature dimen- sions. For more challenging tasks, we evaluate the methods on the TabZilla hard benchmark, which includes 36 difficult datasets-11 1 of which are high-dimennn -ighedimensional (100-2000 fea- tures) and large-scale (2 50K instances). Detailed dataset information, including names and characteristics, is pro- vided in Section D.I, with additional datasets and results available in Section D.4.
6.3. Evaluation on Hard Datasets
In this experiment, we compare TABFLEX to baselines on the TabZilla hard benchmark (McElfresh et al., 2023), which includes 36 datasets. However, due the challenging na- ture of the datasets in the TabZilla hard benchmark, many baselines fail to execute successfully. In Fig. 4, visu- alize the Median AUC and the runtime per 1000 instances across the 36 datasets, with methods that successfully ex- ecuted on all datasets marked as stars, and methods that failed to execute on some datasets marked as circles. This figure focuses on efficient methods, excluding those slower
Baselines. We evaluate our approach against a compre- hensive set of baselines, as considered by McElfresh al. (2023), including (i) four classical methods, e.g., Random Forest (Liaw et al., 2002), (ii) three Gradient Boosted De- cision Trees (GBDT) methods, e.g., XGBoost (Chen & Guestrin, 2016), (iii) ten Non- Transfornfrr Neural Network (Non-TP Mon-TFMM NN) methods, e.g., SAINT (Somepalli et al., 2021), (iv) four Transformer (TF) methods, e.g., TABPEN, with two recent methods designed for scaling tabular classifica-
6

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Mean AUC
Std. AUC
Time 1000 inst.
Algorithm
Class
median
mean
mean
median
median
mean
TabPFN (Hollmann al., 2023)
TF
0.97
0.90
0.21
0.15
0.82
1.04
TABFLEX (Ours)
TF
0.96
0.89
0.22
0.16
0.29
0.48
CatBoost (Prokhorenkova et al., 2018)
GBDT
0.95
0.89
0.23
0.16
2.59
19.51
ResNet (He et al., 2016)
Non-TF TF NN
0.93
0.84
0.24
0.16
13.90
23,40
SAINT (Somepalli et al., 2021)
TF
0.93
0.84
0.24
0.20
173.63
195.16 16
RandomForest (Liaw al., 2002)
Classical
0.92
0.86
0.24
0.17
0.45
0.61
XGBoost (Chen & Guestrin, 2016)
GBDT
0.91
0.86
0.24
0.18
0.49
0.95
HyperFast (Bonet al., 2024)
Non-TF TF NN
0.91
0.83
0.22
0.17
64.38
136.74
DANet (Chen al., 2022)
Non-TF NN
0.89
0.80
0.25
0.19
67.70
78.21
SVM (Cortes, 1995)
Classical
0.87
0.75
0.28
0.22
0.71
87.84
NODE (Popov al., 2019)
Non-TF NN
0.86
0.80
0.24
0.18
157.18
194.07
DeepFM (Guo et al., 2017)
Non-TF TF- NN
0,86
0.79
0.28
0.27
5.48
5,95
FTTransformer (Gorishniy et al., 2021)
TF
0.84
0.78
0.25
0.21
25.40
33.34
LightGBM (Ke al., 2017)
GBDT
0.83
0.76
0.28
0.21
0.25
0.67
MLP-rtdl (Gorishniy ct al., 2021)
Non-TP TF NN
0.83
0.74
0.26
0.20
12.65
22.97
LinearModel (Cox, 1958)
Classical
0.81
0.71
0.27
0.21
0.05
0.06
TuneTables (Feuer et ul., 2024)
TF
0.80
0.72
0.32
0.24
53.48
113.49
STG (Yamada et al., 2020)
Non-TP TF NN
0.79
0.67
0.29
0.23
18.46
21.26
TrbTrssffommrr (Huang et al., 2020)
TF
0.79
0.64
0.24
0.16
19.04
32.84
MLP (Rumelhart et al., 1986)
Non-TP TF NN
0.72
0.65
0.29
0.25
17.83
27.67
DecisionTree (Quinlan, 1986)
Classical
0.63
0.55
0.35
0.31
0.01
0.02
KNN (Cover& & Hart, 1967)
Classical
0.62
0.56
0.30
0.25
0.03
0.03
TabNet (Arik& & Pfister, 2021)
TF
0.56
0.50
0.42
0,40
34.66
42,09
VIME (Yoon al., 2020)
Non-TF NN
0.49
0.48
0.37
0.27
18.43
20.11
NAM (Agarwal et al., 2021)
Non-TF NN
0.33
0.38
0.38
0.31
147.30
341.58
Table Performance of algorithmms across datasets of size less than or equal 1250 (used in Table of McElfresh
(2023)).
The reported AUC values normalized. The "Time/1000 inst." column represents the combined training and test time for all datasets, divided by the total numher samples. Notably, TABF1,EX achieves top-2 performance, with significantly faster runtimes compared baselines of similar performance, and 2x speedup relative to TABPEN.
than 0.5 seconds per 1000 instances. We observe that only all
dataset, poker-hand, containing over one million samples, TABFLEX significantly outperforms other baselines, classi- fying all samples in just 4.88 seconds, while the fifth fastest method requires more than 500 seconds.
TABFLEX, TABPEN, and XGBoost successfully datasets. Notably, TABFLEX is faster and achieves better performance than TABPFN, and is faster than XGBoost while sacrificing only a small margin of performance.
6.4. Extension to Regression Tasks
Next, we focus on 11 nigh-dimensional and large datasets within the TabZilla hard benchmark. Since most baselines do not obtain complete results for all datasets, instead of comparing TABFLEX to specific baseline, report the 5th-best AUC and 5th-best runtime, using these values to summarize the general performance distribution of the base- lines. The results are presented in Table 2. We observe that, for these datasets, TABFLEX substantially outperforms TABPFN, While TABPEN follows McElfresh et al. (2023)'s 's strategy of using only 3000 training samples, TABFLEX utilizes all available training data, achieving superior perfor- mance with comparable or slightly higher processing times. TABFLEX exhibits competitive performance among base- lines while maintaining high efficiency. Notably, on large datasets with more than 50K instances, TABFLEX is signifi- cantly faster than the baselines. For instance, on the largest
So far, we have presented our evaluation on classification tasks. For the extension to regression, simple workaround is to convert the task into classification by discretizing the target range into bins. Shown in Table 3, we apply this technique for regression data with numerical features from Grinsztajn al. (2022), where targets are discretized into 10 and 100 uniform bins. For baselines, we linear regres- sion and the XGBoost Regressor (100 estimators, max depth 6), both with default parameters from the scikit- tetttleaat package. Although regression is not the primary goal of TABFLEX, demonstrates reasonable performance.
Furthermore, we compare TABFLEX with TaneTa- bles (Feuer et al., 2024) using their setup in Sec. D.5 and extend TABFLEX image datasets, shown in Section D.6.

---

shown in Table To better understand this discrep- ancy, we conduct more fine-grained comparison between TABFLEX and XGBoost using synthetic datasets. XG- Boost is configured with tree depth of and 20 esti- mators to balance speed and accuracy. We evaluate the accuracy-runtime tradeoff across varying feature dimen- sions (100, 200, 400, 600, 800, 1000) and training sample sizes (1000, 2000, .... 12000), averaging results over 20 synthetic datasets with diverse distributions. TABFLEX consistently outperforms XGBoost in both accuracy and inference time when the feature dimensionality is below 800. As the number of features increases, the performance TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
AUC
Time (seconds)
Dataset
#Classes
#Features
#Instances
5th Best
TABPFN
TABFLEX
5th Best
TABPFN
TABFLEX
SpeedDating
120
8378
0.86
0.55
0.85
1.58
1.58
1.89
higgs
28
98050
0.79
0.72
0.76
3,46
2.82
4.92
enac-9
856
1080
1.00
0.48
0.96
0.51
0.51
3.80
albert
78
425240
0.71
0.70
33.98
13.46
audiology
24
69
226
0.92
0.82
0,81
0.13
0.23
0.26
jasmine
144
2984
0.86
0.70
0.86
0.68
1.27
0.99
nomao
118
34465
0.99
0.76
0.99
4.03
1.82
5,34
Bioresponse
1776
3751
0.85
0.50
0.75
2.49
1.29
12.38
MiniBooNE
50
130064
0.98
0.98
0.97
10.80
3.19
7.22
airlines
7
539383
0.70
0.63
0.64
6.53
9.73
4.20
poker-hand
10
10
10250009
0.54
0.72
0.84
504.52
15.36
4.88
Table Performance comparison TABFLEX, TABPEN, and other baselines on large, high-dimensional datasets from the TabZilla hard benchmark (McElfresh et al., 2023). Baseline results summarized by the 5th highest AUC and the 5th lowest runtime for each dataset. TABFLEX significantly outperforms TABPEN on these datasets, achieving comparable performance to other baselines while maintaining exceptional speed.
Dataset
TABFLEX
Linear Regression
XGBoost
epu_act
0,9622
0.7661
0.9872
pol
0.7770
0.4471
0.9876
0.7386
0.8984
wine_quality
0.1966
0.2842
0.4398
Ailerons
0.7284
0.8137
0.8272
houses
0.6803
0.6496
0.8469
house_166 16H
0.2519
0.1708
0.5276
diamonds
0.9085
0.9213
0.9477
Trazilian_housee
0.8943
0.3459
0.9828
Bike Biik_Sharing Demand
0.3796
0.3291
0.6995
2aaaaii-gree--de-221666
0.1547
0.3109
0,5732
house_sales
0.6656
0.7375
0.8732
sulfur medical. charges
0.4026
0,3068
0.7497
0.8173
0.8118
0.9790
MiamiHousiig20166
0.8112
0.7302
0.9306
gap narrows, and XGBoost eventually surpasses TABFLEX at 800 features. Nevertheless, TABFLeX achieves a stronger overall tradeoff across most settings.
superconduct
0.6867
0.7169
0.9086
yprop_4_1 abalone
0.0000
0.0449
0.0000
0.3689
0.4622
0.5125
Table Comparison of performance across regression tasks for TABFLEX, Linear Regression, and XGBoost. Regression datasets are from Grinsztajn et al. (2022). To extend TABFLEX to regression, we discretize the target variable into 10 and 100 uniform bins and use the better--erforming setting, converting the task to classification. Despite being designed for classification, TABFLEX delivers reasonable performance on regression tasks.
7.2. Incorporating Data-Efficient Techniques: Dimensionality Reduction and Data Sampling
Since TABFLEX utilizes the ICL for prediction, reducing the complexity of the data can further improve the inference efficiency. We combine TABFLEX with two data-efficient techniques (feature dimension reduction and training data sampling) and investigate how they affect the balance be- tween efficiency and classification performance. See Sec- tion E. 1 for detailed setup and datasets.
7. Ablation Studies
We conduct ablation studies, including a fine-grained com- parison with XGBoost and the integration of other data- efficient techniques. Performance and runtime trends with respect to training sample sizes, along with detailed experi- mental setups, are provided in Section E.
First, for dimensionality reduction, we apply three tech- niques: Principal components analysis (PCA) (Maćkiewicz & Ratajezak, 1993), Singular Value Decomposition (SVD), and random linear projection (Vempala, 2005). Fig. 6 presents the performance and latency of TABFLEX where feature dimensions are varied. Datasets are selected from Table where feature dimensions are greater than 100. On average, dimensions can be reduced to 10% reduce la- tency more than times while AUC scores remain the same
7.1. Fine-Grained Comparison with XGBoost
In Fig. 4, we observe a larger performance gap between FABFLEX and XGBoost compared to the simpler datasets
8

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
3.60
0.52
Seafee
0.4
0.85
0.50
0.775
K-Coaers
6.15
0.00
Madii
64
(a) 100 features
(b) 200 fcatures
(c) 400 features
21
94
50
190 12
23
42
33
40
70
ofasins Thileiy Data
Preimut Traning Dara
0.89
Figure 7. Varying training data with sampling methods. Results are measured on set datasets whose data size is greater than 1000 and whose dimension is lower than 100. For most datasets, only 20% of training data may required to preserve the perfor- mance (left), while significantly reducing the latency (right).
0.00
0.87
0.00
0.25
0,30
0.75
0.0
02
0.0
0.2
Time
Time
(d) 600 features
(e) 800 features
(I) 1000 features
Figure 5. Accuracy-runtime tradeoff of TABFLEX and XG- Boost across feature dimensions (100, 200, 400, 600, S00, 1000) and training sample sizes (1000, 2000, ..... 12000). Results averaged over 20 synthetic datasets with diverse data distrihutions. TABFLEX consistently achieves better performance and faster in- ference than XGBoost when the number of features is below 800. As dimensionality increases, the gap diminishes, with XGBoost overtaking TABFLEX at 800 features.
datasets with more than 2x speedup, while outperforming most other baselines with significantly reduced computa- tional time. Moreover, TABFLEX significantly outperforms TABPFN on larger and more complex datasets, becoming much faster than most other baselines on datasets larger than 100K samples, while maintaining performance on par with satt--of-theaan methods. We posit that TABFLEX further elevates the performance ceiling of neural network-based models tabular classification tasks.
10
5 0.
Limitations & Future Works. While our method achieves fast inference and competitive performance on datasets with around two thousand features, scaling to even larger feature spaces remains compelling research direc- tion. In this paper, we also explored adapting TABFLEX to image classification tasks on small-scale datasets such as MNIST and CIFAR-IO. Extending this approach to large- scale image classification could broaden its applicability, especially given its extremely fast inference and ability to generate predictions for all test samples simultaneously. Further extending TABFLEX to other modalities such as audio and video is another promising direction. This may require new strategies for synthetic data generation for pre- training, along with systematic analyses of the effects of architectural hyperparameters such as depth and embedding size. These efforts would help optimize the model for high- dimensional inputs and enhance its generalization across domains. Currently, our focus on tabular tasks is limited to classification. We attempted a naive extension to regression by discretizing the target range into bins and treating it as a classification problem. promising future direction is to ex- tend TA FL.EE regression tasks using more principled approach, such training on regression-specifi synthetic data. Lastly, TabPFNv2 is concurrent work that further improves the performance of TABPFN. Investigating how incorporating linear attention might impact TabPFNv2 is also an interesting question for future research.
06
Praection
10
20 13
90
100
20
Truttonnt Orgia Divarniaa
Figure 6. Varying feature dimension with dimensionality reduc- tion methods. Results are measured on set of datasets whose number of features is greater than 100. Dimension of features can reduced up to 90% to preserve the performance (left) with inference being 2x faster (right).
or even better than the original predictions (100%).
Similarly, we conduct ablation studies on training data size with different sampling methods (K-centers, K-medoids, un- certainty sampling, and random sampling), shown in Fig. 7. The tested datasets are selected from Table 9 where the data size is greater than 1000 instances and the feature dimen- sion is lower than 100. Fig. 7 illustrates that the original performance can be preserved with only 20% of training data while the latency can be significantly reduced.
8. Conclusion & Discussion
Conclusion. To extend TABPFN for ICL on larger and more challenging tabular classification tasks, in this paper, we conduct comprehensive exploration of scalable alterna- tives to attention, ultimately selecting non-causal linear at- tention. Through computational analysis for algorithmic op- timization of the implementation of linear attention, we de- velop our model, TABFLEX. We demonstrate that TABFLEX achieves comparable performance to TABPEN on small
9

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Impact Statement
in-context algorithm selection. Advances in neural infor- mation processing systems, 36:57125-57211 2023c.
This research tackles a key limitation in applying TABPEN to tabular classification tasks: scalability. By introducing TABFLEX, which leverages linear attention, enables ef- ficient processing of tabular datasets with millions of sam- ples and thousands of features. The primary impact lies in making in-context learning feasible for a broader range of large-scale tabular problems, which is particularly useful in domains such as finance and recommendation systems. Furthermore, the successful integration of linear attention in TABFLEX lays the groundwork for future studies on efficient attention mechanisms and model architectures for tabular and other non-NLP domains, as supported by promising preliminary results in image classification.
Beltagy, I., Peters, M. E., and Cohan, A. Long-
Former: The long-document transformer. arXiv preprint arXiv:200..051500 2020.
Bischl, B., Casalicchio, G., Feurer, M., Hutter, F., Lang, M., Mantovani, R. G., van Rijn, J. N., and Vanschoren, J. Openml benchmarking suites. arXiv:1708.0333112 [stat.MI.], 2019.
Bonet, D., Montserrat, D. M., Giró-i Nieto, X., and Toan- nidis, A. G. Hyperfast: Instant classification for tabular data. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 11114-11123, 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08874, 2023.
Chen, J., Liao, K., Wan, Y., Chen, D. Z., and Wu, Danets: Deep abstract networks for tabular data classification and regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 3930-3938, 2022.
Agarwal, R., Melnick, L., Frosst, N., Zhang, X., Lengerich, B., Caruana, R., and Hinton, G. E. Neural additive models: Interpretable machine learning with neural nets. Advances in neural information processing systems, 34: 4699-4711, 2021.
Chen, T. and Guestrin, C. Xgboost: scalable tree boosting system. In Proceedings of the 22nd acm sigkdd inter- national conference on knowledge discovery and data mining, pp. 785-794, 2016.
Ahamed, M. A. and Cheng, Q. Mambatab: A simple yet ef- fective approach for handling tabular data, arXiv preprint 2eXiv:2401.08867. 2024.
Chu, Y., Xu, J., Zhou, X.. Yang. Q., Zhang. S., Yan, Z., Zhou, C.,, and Zhou, J. Qwen-audio: Advancing univer- sal audio understanding via unified large-scale audio- language models. arXiv preprint arXiv:2311.07919, 2023.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36:45614-45650 2023.
Chu, Y., Xu, J., Yang, Q. Wei, H., Wei, X., Guo. Z., Leng, Y., Lv, Y., He, J., Lin, J., et al. Qwen2-audio technical report. arXiv preprint 2rXiv:2407.1077 2024.
Arik, S. S. and Pfister, T. Tabnet: Attentive interpretable tab- ular learning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 6679-6687, 2021.
Cortes, C. Support-vector networks. Machine Learning, 1995.
Arun, K., Ishan, G., and Sanmeet, K. Loan approval pre- diction based on machine learning approach. IOSR J. Comput. Eng, 18(3):18-22 11-11, 2016.
Cover, T. and Hart, P. Nearest neighbor pattern classification. IEEE transactions on information theory, 13(1):21 -7 1967.
Bai, J., Bai, S., Chu, Y., Cui, Z., Dang, K., Deng, X., Fan, Y., Ge, W., Han, Y., Huang, F, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023a.
Cox, D. R. Journal of the Royal Statistical Society Series B: Statisti- cal Methodology, 20(2):2152233 1958.
The regression analysis of binary sequences.
Bai, J., Bai, S., Yang, S., Wang. S., Tan, S., Wang, P., Lin, J., Zhou, C., and Zhou, J. Qwen-vl: A frontier large vision- language model with versatile abilities, arXiv preprint arXiv:2308. 12966, 2023b.
Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https: /openreview net/ forum? i id=m2n2Xyh
Bai, Y. Chen, F. Wang, H., Xiong, C., and Mei, S. Trans- formers as statisticians: Provable in-context learning with
10

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Dao, T. and Gu, A. Transformers are SSMs: General- ized models and efficient algorithms through structured state space duality. In Forty-first International Con- ference on Machine Learning, 2024. URL https: // //openrevien net / forum? wwtmmiddzztt@FFee...
Gorishniy, Y., Rubachev, I., Kartashev, N., Shlenskii, D., Kotelnikov, A., and Babenko, A. TabR: Tabular deep learning meets nearest neighbors. In International Con- ference on Learning Representations (ICLR), 2024.
Grinsztajn, L.,, Oyallon, E., and Varoquaux, G. Why do tree- based models still outperform deep learning on typical tabular data? Advances in neural information processing systems, 35:507-520, 2022.
Dao. T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. Flashattention: Fast and memory-efficeent exact atten- tion with 10-awareness. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho. K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //openrevien net forum? weeiiddgHHHgiiooo
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. First Conference on Lan- guage Modeling, 2024. URL https: openrevie...e.eeeeeeee net
Ding, N., Levinboim, T., Wu, J., Goodman, S., and Soricut, R. CausalLM is not optimal for in-context learning. In International Conference on Learning Representations (ICLR), 2024.
Gu, A., Johnson, I., Goel, K., Saab, K. K., Dao, T., Rudra, A., and Re, C. Combining recurrent, convolutional, and continuous-time models with linear state space layers. In Beygelzimer, A., Dauphin, Y., Liang. P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https: openrevieeeeeeeeeeeee net/ forum?id-y/Ade
Dinh, T., Zeng. Y., Zhang, R., Lin, Z., Gira, M., Rajput, S., yong Sohn, J., Papailiopoulos, D., and Lee, K. LIFT: Language-interfaeed fine-tuning for non-language ma- chine learning tasks. In Oh, A. H., Agarwal, A., Bel- grave, D., and Cho. K. (eds.), Advances in Neural In- formation Processing Systems, 2022. URL https: // //openreview.eoo www.:/iwww...teeeeeeeeeeeeeeeeee wwrumm/rummmmmo SttprrmmddddsSSMMMUSSaaS FJMEGIUfa.
Gu, A., Goel, K., and Re, C. Efficiently modeling long sequences with structured state spaces. In In- ternational Conference on Learning Representations, 2022. URL httpttttss///ppreeeeeeee id-uYLFozlv1AC.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A. et al. The Ilama herd of models. arXiv preprint arXiv:2407.21788, 2024.
net/ /
Guo, H., Tang, R., Ye, Y., Li, Z., and He, X. Deepfm: a factorization-machinn based neural network for ctr pre- diction. arXiv preprint arXiv:1703.0422 2017.
Feuer, B., Schirrmeister, R. T., Cherepanova, V., Hegde, C., Hutter, F., Goldblum, M., Cohen, N., and White, C. Tunetables: Context optimization for scalable prior-data fitted networks. arXiv preprint 44X 2024.
Hayes, T., Rao, R., Akin, H., Sofroniew, N. J., Oktay, D., Lin, Z. Verkuil, R., Tran, V. Q., Deaton, J., Wig- gert, M., Badkundri, R., Shafkat, I,, Gong, J., Derry, A., Molina, R. S., Thomas, N., Khan, Y. A., Mishra, C., Kim. C., Bartie, L. J., Nemeth, M., Hsu, P. D., Sercu, T., Candido. S., and Rives, A. Simulating 500 mil- lion years of evolution with language model. bioRxiv, 2024. doi: 10.1 01. 66005883. URL https: / / //doi. org/ 10. 1101/2024 07 01. 600583.
Feurer, M., Eggensperger, K., Falkner, S., Lindauer, M., and Hutter, F. Auto-sklearn 2.0: Hands-free automl via meta-learning. arXiv:2007.04074 [cs.LG], 2021.
Friedman, J. H. Greedy function approximation: a gradient boosting machine. Annals of statistics, pp. 1189-1232, 2001.
Gemini Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11800 2023.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Gong, Z., Liu, J., Wang, Q. Wang, J., Cai, X., Zhao, D., and Yan, R. Improving input-label mapping with demon- stration replay for in-context learning. arXiv preprint arXiv:2310.19572. 2023.
Hollmann, N., Müller, S., Eggensperger, K., and Hutter, F. TabPFN: A transformer that solves small tabular classification problems in second. In The Eleventh International Conference on Learning Representations, 2023. URL https:: http::///pppeeeeeeeeoeooe.mmoeep net. nwt.forutt? i dd=cp5PvcI6we__
Gorishniy, Y., Rubachev, I., Khrulkov, V., and Babenko, A. Revisiting deep learning models for tabular data. Ad- vances in Neural Information Processing Systems, 34: 18932-18943, 2021.
Hollmann, N., Müller, S., Purucker, L., Krishnakumar, A., Körfer, M., Hoo, S. B., Schirrmeister, R. T., and Hutter,

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
F. Accurate predictions on small data with a tabular foundation model. Nature, 637(8045):319-336, 2025.
Popov, S., Morozov, S., and Babenko, A. Neural oblivious decision ensembles for deep learning tabular data. arXiv preprint wwxXiv:999.003111 2019.
Huang, X., Khetan, A., Cvitkovic, M., and Karnin, Z. Tab- transformer: Tabular data modeling using contextual em- beddings. arXiv preprint wrXiv:2012.06677 2020.
Prokhorenkova, I.., Gusev, G., Vorobev, A., Dorogush, A. V., and Gulin, A. Catboost: unbiased boosting with categori- cal features. Advances in neural information processing systems, 31, 2018.
Johnson, A. E., Pollard, T. J., Shen, L.,, Lehman, L.-w. H., Feng, M.. Ghassemi, M.. Moody. B., Szolovits, P., An- thony Celi, L.,, and Mark, R. G. Mimic-iii, a freely ac- cessible critical care database. Scientific data, 3(1):1-9, 2016.
Qin, Z., Han, X., Sun, W., Li, D., Kong, L.,, Barnes, N., and Zhong, Y. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, 2022.
Katharopoulos, A., Vyas, A., Pappas. N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165, PMLR, 2020.
Quinlan, J. R. Induction of decision trees. Machine learning, 1:81-106, 1986.
Rives, A., Meier, J., Sercu, T., Goyal, S., Lin, Z., Liu, J., Guo, D., Ott, M., Zitmick, C.,, Ma, J., and Fergus, R. Bi- ological structure and function emerge from scaling unsu- pervised learning 250 million protein sequences. PNAS, 2019. doi: 10.1101/622803. URL https://wwi biorxiv. wtg/content..oooooooooooooomo monttent110.. 11111...11111111 101100/2220344....
Ke, G., Meng, Q.. Finley, T., Wang, T., Chen, W., Ma, W., Ye, Q., and Liu, T.-Y. Lightghm: A highly efficient gradient boosting decision tree. Advances in neural infor- mation processing systems, 30, 2017.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Rumelhart, D. E., Hinton, G. E., and Williams, R. J. Learn- ing representations by prokapropgggttng errors. nature, 323(6088):533-556 1986.
LeCun, Y., Cortes, C., and Burges, C. MNIST handwritten digit database. 2, 2010.
Shah, J., Bikshandi, G., Zhang. Y., Thakkar, V., Ramani, P., and Dao, Flashattention-3: Fast and accurate atten- tion with asynchrony and low-precision. arXiv preprint arXiv:2407.0866 2024.
Liaw, A., Wiener, M., et al. Classification and regression by randomforest. R news, 2(3): 1-)22,822, 2002.
Ma, J., Thomas, V., Yu, G., and Caterini, A. In- context data distillation with TabPFN. arXiv preprint arXiv:2402.06971, 2024.
Somepalli, G., Goldblum, M., Schwarzschild, A., Bruss, C. B., and Goldstein, T. Saint: Improved neural networks for tabular data via row attention and contrastive pre- training. arXiv preprint arXiv:2106.01 342, 2021.
Maćkiewicz, A. and Ratajezak, W. Principal components analysis (pea). Computers & Geosciences, 19(3):303- 342, 1993.
Sun, Y., Dong, L., Huang. S., Ma, S., Xia, Y., Xue, J., Wang. J., and Wei, F. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.086211 2023.
McElfresh, D. C., Khandagale, S., Valverde, J.. C, V. Ra- makrishnan, G., Goldblum, M., and White, C. When do neural nets outperform boosted trees on tabular data? In seveth-eeentth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, n forum?
Thielmann, A. F., Kumar, M., Weisser, C., Reuter, A., Säfken, B., and Samiee, S. Mambular: A sequen- tial model for tabular deep learning. arXiv preprint arXiv:2408.0622 2024.
2023. URL ptttps://operreeee id-GjVaXey42T.
Thomas, V., Ma, J., Hosseinzadeh, R., Golestanch, K., Yu, G., Volkovs, M., and Caterini, A. L. Retrieval & fine- tuning for in-context tabular models. In NeurIPS, 2024.
Orvieto. A., Smith, S. L., Gu, A., Fernando, A., Gulcehre, C., Pascanu, R., and De, S. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 86670-26699. PMI.R, 2023.
Vanschoren, J., van Rijn, J. N., Bischl, B., and Torgo, L. OpenML: Networked science in machine learning. SIGKDD Explorations, 15(2):49 60, 2013.
Peng, B., Alcaide, F., Anthony, Q. G., Albalak, A., Arcad- inho, S., Biderman, S., Cao, H., Cheng, X., Chung, M. N., Derczynski, L., et al. Rwkv: Reinventing rnns for the transformer era. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E.,, and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
12

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Vempala, S. The random projection method, volume 65. American Mathematical Soc., 2005.
Wan, Z., Li, M., Liu, S., Huang, J., Tan, H., and Duan, W. Eegformer: A transformer baseed brain activity classifica- tion method using eeg signal. Frontiers in Neuroscience, 17:1 148855, 2023.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIS, a novel image dataset for benchmarking machine learning algorithms. 2017.
Xu, C., Huang, Y.-C., Hu, J. Y.-C,, Li, W., Gilani, A., Goan, H.-S., and Liu, H. BiSHop: Bi-directional cellular learn- ing tabular data with generalized sparse modern hop- field model. In International Conference on Machine Learning (ICML), 2024.
Xu, D. Q. Cirik, F. O., Asadi, R., Sun, Y., and Wang, W. Mixture of in-context prompters for tabular PFNs. International Conference on Learning Representations (ICLR), 2025.
Yamada, Y., Lindenbaum, O., Negahban, S., and Kluger, Y. Feature selection using stochastic gates. Interna- tional conference on machine learning. pp. PMLR, 2020.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardwarc-efffcient train- ing. In Forty-first International Conference on Machine Learning, 2024.
Yoon, J., Zhang, Y., Jordon, J., and Van der Schaar, M. Vime: Extending the success of self-and sumi-supervised learn- ing to tabular domain. Advances Neural Information Processing Systems, 33:1 1013-11311443 2020.
You, H., Fu, Y., Wang, Z., Yazdanbakhsh, A., and Lin, Y. C. When linear attention meets autoregressive decod- ing: Towards more effective and efficient linearized large language models. In Forty-first International Conference on Machine Learning, 2024.
Zhang, M., Bhatia, K., Kumbong, H., and Re, C. The hedge- hog& & the poreupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Confer- ence on Learning Representations, 2024.
Zhang, S., Yao, L.,, Sun, A., and Tay, Y. Deep learning based recommender system: survey and new perspectives. ACM computing surveys (CSUR), 52(1):1-38, 2019.
Zhu, B., Shi, X., Erickson, N., Li, M., Karypis, G., and Shoaran, M. XTab: Cross-table pretraining for tabular transformers. In International Conference on Machine Learning (ICML), 2023.
13

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Appendix
Extended Section 2: Related Works
15
Supplement to Section 4: Architectural Exploration for Scalable Tabular Learning
16
B.1 Comparison with Other Attention Variant
16
B.2 Supplementary Visualizations
17
Supplement to Section 5: TABFLEX
17
C.I Computation Analysis of Various Attention Mechanism
17
C.I.1
HBM-Efficient Linear Attention
18
C.1.2 Simplified PyTorch Implementation of Linear Attention
20
C.2 Model Training
21
C,3 Validation Datasets
21
C.4 Sensitivity Analysis on Decision Boundaries
22
Supplement to Section 6: Performance Evaluation of TABFLEX
22
D.1 TabZilla Datasets
22
D.2 Extended Experiment Setup
24
D.3 Results on 98 Simple Datasets (Table 1, McElfresh et al. (2023))
24
D.4 Evaluation on Additional Datasets
24
D.5 Additional Comparison with Tune'Tables
24
D.6 Extending TABFLEX for Image Classification
24
Supplement to Section 7: Ablation Studies
27
E.1 Datasets for Section 7.2: Incorporating Data-Efficient Techniques
27
E.2 Performance and Runtime vs. Training Sample Size
29
14

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
A. Extended Section 2: Related Works
Classical Machine Learning Approaches for Tabular Classification. Classical machine learning algorithms have long been the foundation of tabular data classification. These methods include k-Nearest Neighbors (KNN) (Cover & Hart, 1967), Logistic Regression (Cox, 1958), Decision Trees (Quinlan, 1986), and Support Vector Machines (SVM) (Cortes, 1995). These classical models, while effective, often struggle to handle complex, high-dimensional tabular datasets, motivating the development of more sophisticated approaches.
Gradient-Boosting Decision Trees for Tabular Classification Gradient-boosting decision trees (GBDTs) (Friedman, 2001) have emerged as cornerstone in tabular classification, owing to their exceptional ability to capture intricate patterns in structured data. GBDTs refine their outputs to minimize errors by iteratively combining predictions from weak learners, resulting in high predictive accuracy. XGBoost (Chen & Guestrin, 2016) introduced weighted quantile sketching, advanced regularization techniques, and sparsity awareness, achieving mate-offthe-ane performance. LightGBM (Ke et al., 2017), a computationally efficient GBDT implementation. employs Gradient- based One-Side Sampling and a leaf-wise tree growth strategy, CatBoost (Prokhorenkova al., 2018) leverages symmetric trees and introduces ordered boosting, with a particular emphasis on effectively handling categorical features. These advancements have rendered GBDTss not only powerful but also versatile tools in the domain of tabular data, dominating tabular classification in terms of both speed and performance until the advent of TABPEN.
Transtormer-based Approaches for Tabular Classification. Recent years have witnessed numerous attempts to employ Transformers for tabular classification (Arik& & Pfister, 2021; Huang et al., 2020; Gorishniy et al., 2021; Dinh et al., 2022; Hollmann et al., 2023). These methods utilize Transformers in diverse ways to tackle tabular data. TabNet (Arik & Pfister, 2021), one of the pioneering efforts, applies unsupervised pre-training on masked tabular datasets to infer missing features, thereby enhancing the model's understanding of datasets and features. It then performs supervised learning on feature selection to obtain the final decision boundary, akin to decision trees. Huang et al. (2020) introduced Tab FabTransformer, which leverages Transformers to better handle categorical features by concatenating their contextual embeddings with numerical features. While TabTransformer processes categorical and continuous features separately, SAINT (Somepalli et al., 2021) projects both feature types into a shared embedding space before passing them through transformer blocks, thereby enhancing overall performance. FT-Transformer (Gorishniy et al., 2021) introduces a feature tokenizer to convert each example into a sequence of embeddings, enabling Transformers to process tabular datasets and make predictions. LIFT (Dinh et al., 2022) utilizes a pre-trained language model with parameter-efficiene each sample into complete sentence with feature names in the prediction prompt. TabR (Gorishniy et al., 2024) proposes a ettrival-auggented model with custom kNN-like component to retrieve and extract signals from the nearest neighbors. BiSHop (Xu et al., 2024) establishes interconnected directional learning modules to process data column- and row-wis wose for tabular learning. XTab (Zhu et al., 2023) utilizes independent featurizers and federated learning to resolve inconsistent column types and quantities.
fine-tuning, incorporating task descriptions and converting
TABPFN (Hollmann et al., 2023) is trained offline on synthetic datasets derived from prior distributions and performs ICL rather than additional parameter tuning for given dataset, enabling solve small tabular classification tasks within seconds. Prior to work, Tube'Table (Feuer et al., 2024) extended TABPFN scale large datasets by performing prefix-tuning for each dataset to achieve better performance. MixturePNN (Xu et al., 2025) improves scalability by routing new test samples to a pool of scalable prompters using Sparse Mixture of In- In-Context Prompters, while LoCalPEN (Thomas et al., 2024) proposes retrieving local subset of task-specific data for efficiently fine-tuning on. Ma et al. (2024) introduce in-context data distillation to optimize TabPFN's context and remove the data size constraint. Tune'Table (Feuer et al., 2024) scales TABPFN to large datasets by performing a prefix tuning per dataset. TabPFNv2 (Hollmann et al., 2025) enhances TabPFN's accuracy in low-data regimes (fewer than 10,000 samples), complementing our focus on speed and scalability, Our method is also based on TABPFN, extending its scalability large datasets while maintaining and improving efficiency by simply replacing the softmax attention with linear attention.
Attention Mechanisms and Scalable Alternatives. While attention in Transformers (Vaswani et al., 2017) is central to the strong performance of language models, it encounters scaling challenges for long sequences due to its quadratic computational and memory complexity. To overcome these limitations, several scalable alternatives have been proposed (Gu & Dao, 2024; Dao & Gu, 2024; Katharopoulos et al., 2020; Peng et al., 2023; Orvieto et al., 2023; Sun et al., 2023), all aiming to achieve subquadratic time complexity. In contrast, classical RNNs provide the advantage of efficient linear-time inference but suffer from limitations in training efficiency, lacking the parallelization capabilities of Transformer architectures. Linear
15

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
attention (Katharopoulos al., 2020) addresses both concerns by reformulating self-attention as a linear dot-product of kernel feature maps, reducing the computational complexity from quadratie linear time. Additionally, causal linear attention can be interpreted as a form of RNN, the model makes predictions based on a current token and "hidden state," which summarizes information from the previous tokens. State-space models (SSMs), another popular variant of RNNs, address the drawbacks of classical RNNs by considering linear RNNs and proposing novel algorithms for efficient training (Gu et al., 2021; 2022; Gu & Dao, 2024; Dao & Gu, 2024; Peng et al., 2023; Orvieto et al., 2023; Sun et al., 2023).
Dao et al. (2022) identified that another bottleneck in attention mechanisms' speed stems from the relatively slow access to high-bandwidth memory (HBM) in GPUs. To address this limitation, FlashAttention (Dao et al., 2022: Dao, 2024; Shah et al., 2024) restructures attention computation to optimize the utilization of high-speed on-chip SRAM while minimizing access to slower HBM, thereby enhancing the efficiency of GPU-based attention operations. FlashAttentioo strategically balances computational efficiency against memory bandwidth efficiency. Although the computational complexity in terms of sequence length remains quadratic, the optimizations introduced by FlashAttention significantly accelerate attention computation in wall-clock time.
Non-Transformer Neural Network-based Approaches for Tabular Classification. Non-Transformer neural networks, such as Multi-Layer Perceptrons (MLP) (Rumelhart et al., 1986), were explored for tabular classification long before Transformer-based methods, but their performance was limited. In recent years, several novel neural network techniques have been developed for this task, including ResNet (He et al., 2016), DANet (Chen et al., 2022), NODE (Popov et al., 2019), DeepFM (Guo et al., 2017). STG (Yamada et al., 2020), VIME (Yoon et al., 2020), and NAM (Agarwal et al., 2021). DeepFM (Guo et al., 2017) employs a factorization machine- neural network to learn from categorical data. Drawing inspiration from CatBoost, Popov et al. (2019) presents a novel neural network architecture designed specifically for tabular data, named Neural Oblivious Decision Ensembles (NODE). While self- and semi-supervised learning have demonstrated effectiveness in the domains of computer vision and natural language processing, Yoon et al. (2020) proposed Value Imputation and Mask Estimation (VIME), which represents the first attempt to address tabular tasks using a self- and semi-supervised learning framework. (Agarwal et al., 2021) proposed the Neural Additive Model (NAM), an interpretable neural network that maintains strong performance on tabular data. Yamada et al. (2020) proposed a feature selection method using stochastic gates (STG), which is a neural network-based and effective approach for tabular data. Chen et al. (2022) designed an abstract layer, specialized neural component for tabular data, and proposed Deep Abstract Networks (DANets) by stacking these layers.
Some approaches even replace Transformers with SSMs for tabular learning (Ahamed & Cheng, 2024; Thielmann et al., 2024). However, these methods require training on a per-dataset basis, leading to high computational costs, and they are generally slower than GBDTs for tabular classification tasks.
Linear Attention for In-Context Learning. Although linear attention has been reported to underperform in some language modeling tasks (You et al., 2024; Zhang et al., 2024; Qin et al., 2022), recent theoretical work demonstrates its effectiveness in in-context learning scenarios, where it can emulate gradient descent to achieve learning during inference (Ahn et al., 2023).
B. Supplement to Section 4: Architectural Exploration for Scalable Tabular Learning
B.1. Comparison with Other Attention Variant
In addition to the broad categories of all linear RNN variant models studied in this paper, we also consider another mechanism that enjoys linear complexity: sliding window attention (Beltagy et al., 2020). We show that TABFLEX achieves significantly better performance.
Method
#Class
#Features
#Instances
Sliding Window
Linear (Ours)
Poker-Hand
10
10
1,025,009
0.48
0.84
Airlines
2
7
539,383
0.48
0.64
Higgs
2
28
98,050
0.39
0.76
Table 4. Performance comparison of TABFLEX with Sliding Window attention.
16

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
0.8
Mamba Non-Causal Attn. Causal Attn.
Accuracy
0.6
Non-causal Masked Attention
Causal Masked Attention
0.4
Mamba
0.2
0.0
1000
2000
3000
Number of Training Samples
1.0
0.5
0.0
0.0
0.5
1.0
Training Loss
Test Mean AUC
(a) Effect of causal masking. When more samples are provided, the non-causal model shows better sample utilization and accuracy, while the causal model's performance plateaus early and declines. The Mamba variant performs poorly.
(b) ICL performance comparison between Mamba and Trans- former models. Results show that Transformer baseed models achieve lower training loss and higher AUC across 150 test datasets.
Mean: -0.003
Mean: -0.055
25
20
Density
15
Density
4
10
0
-0.100 00.075 -0 www.0.05000... 0.025 0.000 0.025 0.050 0.075
-0.4
-0.3
-0.2
00.1
0.0
0.1
Accuracy improvement (Linear Softmax)
Runtime Improvement (Linear Softmax)
(c) Accuracy and runtime comparison of softmax and linear attention. Replacing softmax with linear attention preserves comparable accuracy while significantly reducing runtime.
Figure 8. Impact of model architecture tabular classification performance. Detailed version of Fig. 2.
B.2. Supplementary Visualizations
Fig. 8 provides additional visualizations to complement the analysis in Fig. 2. It includes both TabPFN-Cause Causal-Masked and TabPFN-Mamba variants in the settings shown Fig. 2a and Fig. 2b, as the two experiments highlight different aspects of model behavior. To improve the interpretability of performance differences between softmax and linear attention across datasets, also include density plot illustrating the changes in runtime and accuracy when the softmax operator is replaced with linear attention.
C. Supplement to Section 5: TABFLEX
In this section, we provide detailed insights into TABFLEX, including training procedures, validation dataset selection, and analyses such as the sensitivity model selection thresholds on final performance.
C.1. Computation Analysis of Various Attention Mechanism
In this part, we provide a computational analysis of various attention mechanisms, comparing standard attention, FlashAtten- tion (specifically FlastAttention-I (Dao al., 2022)), causal PaashinneeAAttentiii (referred to Flash MlsshLinearAil intarAttention in
17

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Yang et al. (2024)), and non-causal linear attention. To clarify, FlashLinearAAtentiion is designed to reduce HBM access specifically for causal linear attention. For notational simplicity, we use the term "linear attention" to refer non-causal linear attention. For both linear attention and TlashhiinearAttentionn, we assume an element-wise kernel feature mapping. The assumption is reasonable, as popular kernel feature mappings like elu(-) + 1 (used here) and (eLU(-) are also element-wise. We denote this mapping as Φb
Algorithm/ 2 Causal FlashLineaurAttentioo Implementation (Yang et al., 2024)
Input: Matrices Q, K, V RND in HBM, on-chip SRAM of size M
Set block size B Initialize O (0) D E RNXD in HBM Divide Q into T T blocks QI,..., QT. of size x D each, and divide K, V into T = N
blocks Ki...., .T. and Vi... VT of size x D each Divide O into T blocks .,...., Or of size B x each On on-chip SRAM, construct causal mask, M E BxB On SRAM, initialize
S (0) Dx D RDDD for S j < T do
Load J, Vj Q2 O, from HBM to on-chip SRAM On chip, compute K2 +- [({K) On chip, compute Q, t o(2))
Write Oj f Q + ((Q, O M) V, to HBM On-chip, compute S + S KJ
3 end
Output: 0
We evaluate these mechanisms based on their High Bandwidth Memory (HBM) access, memory requirements, and floating- point operations (FLOP) when computing attention outputs given query, key, and value inputs. While Dao et al. (2022) have provided computations for standard attention and FlashAttention, we focus our analysis on causal FlashLinear Attention (detailed in Alg. 2) and HBM-efficient non-causal linear attention (developed by us and detailed in Alg. 3) in Section C.I.I. In practice, we employ a simplified PyTorch implementation of linear attention and demonstrate its efficiency, as it only causes marginal increases in HBM access and memory usage as we demonstrate in Section C.1.2. Furthermore, we present visualizations in Section C.1.2 that illustrate the time and CUDA memory consumption of these attention mechanisms across various sequence lengths and scenarios.
Algorithm: 3 HBBM-Ffffiiiett Implementation of Linear Attention
Input: Matrices Q, K, V E RND in HBM, on-chip SRAM of size M Set block size B Initialize O = (0)NxD D E NXD in HBM Divide Q into T = 1 blocks Qı....,QQ of size B x D each, and divide K, V into T = Tेr blocks K k....., Ky and V v.... Vr of size B D each Divide O into T blocks O Or of size B x D each On on-chip SRAM, initialize S = (0) DxD for 1 <is do Load K,, V. On chip, compute K; + (K)) On chip, compute S t S + K. V;
for 1 S j < T do Load O, On chip, compute Q; + (Q2)) Write Oj + Q, to HBM
Output: O
C.I.I. HBM DEFFIIEICENN LINEAR ATTENTION
Here, we analyze the number of HBM accesses, HBM memory, and FLOPS required by FlashL (Alg. 2) and linear attention (Alg. 3).
Let Q K, V E RD represent the query, key, and value matrices for single attention head, where N is the sequence length and is the embedding size. Both MaashlinearAteettin (Alg. 2) and linear attention (Alg. 3) require 5ND D HBM accesses to compute the attention output.
Lemma 2.
Proof of Lemma 2. For causal
(Alg. 2):
.Line& 8: Loading K, Vj, Q5 O, necessitates 4BD HBM accesses.
Line 11: Writing Oj requires BD HBM accesses.
18

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
These operations are executed 7' times, where 7 B
Thus, the total HBM accesses
5BD - 5BD = D.
B
For non-causal linear attention (Alg. 3):
Line 7: Loading K, V, requires 2BD HBM accesses.
Line II: Loading 2, O, requires 2BD HBM accesses.
Line 13: Writing Oj requires BD HBM accesses.
These operations are also repeated I' times, where T = NN Consequently, the total HBM accesses are:
N
5BD T 5BD
5ND.
Therefore, we conclude that both causal Flashl inear Attention and non-causal linear attention require 5ND HBM accesses to compute the attention output.
RND represent the query, key, and value matrices for single attention head, where N is the sequence length and D the embedding size. Both Flasttimearttenntion (Alg. 2) and linear attention (Alg. 3) require 4N D HBM memory compute the attention output.
Lemma 3.
Let Q K, V
Proof of Lemma 3. For both algorithms:
Storing Q, K, V requires 3ND memory.
Storing O requires N memory.
Total HBM memory usage: DND.
Lemma 4. Let Q.K.V E NRD represent the query, key, and value matrices for a single attention head, where N the sequence length and D is the embedding size. Both AlaattinnatAttentioon (Alg. 2) and linear attention (Alg. 3) require O(ND²) FLOPS to compute the attention output.
Proof of Lemma 4. For causal FlashLinearAAteentionn (Alg. 2):
Computing oi K2) and Φ(Q,) requires 2BD FLOPs.
Computing (Q, K M requires B2D 1) B² FLOPs.
The result of last step multiplied by V, requires B²(2D 1) BD(2B 1) FLOPs.
Computing Q requires B D(2D 1) FLOPs.
The addition of the last two steps requires BD FLOPs.
Computing K,T V, (line 10) requires (2B - 1) 22 FLOPs.
The addition of S and the last step requires 2 FLOPPs
19

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
The total number of FLOPs for iteration is:
2BD + (2 - 1) B (2 - +D - 1) B D(2D - 1) BD + (2B 1). D>2 n)
= + +DD 4BD²
These operations are repeated T times. The total number of FLOPs is:
(4B2 + 2BD 4BD*) - = O(ND)).
For non-causal linear attention (Alg. 3):
Computing ((K)) requires BD FLOPs.
Computing S + K. V; (line 9) requires 2BD² FLOPs.
Computing ((²)) (line 12) requires BD FLOPS.
Computing 2 (line 13) requires (2D D)DDDDDDDDD FLOPs.
These operations are repeated = times. The total number of FLOPs is:
(BD 4BD") T O
Thus, conclude that both algorithms require O(ND2) FLOPs compute the attention output.
C.1.2. SIMPLIFIED PYTORCH IMPLEMENTATION OF LINEAR ATTENTION
In our implementation, we adopt a straightforward PyTorch approach to linear attention rather than an HBM-efficient method. We employ the concise two-line implementation presented in Listing 1. the following lemma, we demonstrate that this straightforward implementation only incurs a marginal increase in HBM accesses and HBM memory usage.
def linear_ (q, k, v)
q: (batch, heads, seq_g, dim_qk) [k)
k: (batch, heads, kv, dim_qk) V: (batch, heads, kv, dim_v)
kv einsum ("bhnd, bhnm->bhdm", k, v) torch.ed naum ("bhld, bhōm->bh1m", q, kv) return c.conti igouu ()
Listing Straightforward PyTorch implementation of linear attention (Katharopoulos et al., 2020).
Theorem 1. Let Q, K, V E NxD represent the query, key, and value matrices for a single attention head, where N is the sequence length and D is the embedding size. Both causal FlashLinearAtteniun (Alg. 2) and non-causal linear attention (Listing 1) require ON HBM accesses, O(ND) HBM memory, and O(N FLOPS to compute the attention output.
Proof. Let consider the implementation in Listing and compare it to Alg. 3. PyTorch's optimized tensor computation ensures efficiency, with the primary distinction between Listing and Alg. 3 being the storage of kv in the former, which is equivalent to S E DDD in Alg. 3. This results in the following changes:
HBM Accesses: By Lemma Alg. requires 5ND HBM accesses. Due to the additional write and load operations for S E DDD Listing requires 5ND D - 2D2 HBM accesses.
HBM Memory Usage: By Lemma 3, Alg. 3 requires 4ND HBM memory usage. Due to the additional storage requirements
for S E DDD Listing requires 4ND 2 HBM memory usage.
20

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
The number of FLOPS remains unaffected. The analysis above, in conjunction with Lemmas 2, 3, and 4, yields the desired outcome.
In Table 5, we summarize the #HBM access, HBM memory, and FLOPSS required by standard attention (with naive PyTorch implementation), FlashAttention-l, FlashLincar (causal), and linear attention with both implementations.
FlashAttention
FlashLineerAAttention (Yang al., 2024)
Linear Attention
Standard Attention
(Dao et al., 2022)
Alg. 3
Listing
# HBM access
4N2 + 4ND
D2N
16N" D
2ND
5ND
5ND
5ND + 2D22
Memory
2N + 4ND
2N + 4ND
4ND
4ND
4ND + D
FLOPS
O(N2D)
O(N2D)
O(ND?)
O(ND?)
O(ND2)
Table Comparison of memory and computational costs across different attention mechanisms. FlashAttention improves the speed of standard attention by optimizing # HBM access. Flash causal linear attention takes a similar approach, achieving linear # HBM access. However, we show that non-causal linear attention already achieves linear # HBM access, matching the efficiency of flash causal linear attention without requiring any additional optimization on # HBM access.
Subsequently, we visualize the empirical execution time and CUDA memory utilization of FlashAttention-22, FlashLiin- carAttention, and linear attention in Fig. 9a and Fig. 9b, respectively. We vary the head dimension E {32, 64, 128, 256), the number of heads E {2,4, 8, 16), and the sequence length {2* 25 ....... 215). We focus on the self-attention case, randomly generating input (serving as key, query, and values) with a batch size of 10, and replicate the experiment times. The final values presented are aggregated across these 5 simulations. Notably, we were unable obtain results for FlashLinearAAtention in two configurations: (1) head dimension 256 with 8 heads, and (2) head dimension 256 with 16 heads, due to illegal memory access error incurred by the PyTorch package fla (Yang et al., 2024). Our observations from the figures indicate that both runtime and CUDA memory usage of -nesh Attention and linear attention exhibit linear growth with respect to sequence length, aligning with the predictions of Theorem 1.
C.2. Model Training
We implement linear attention with the feature function elu(-) + 1, adhering to the default implementation proposed by Katharopoulos et al. (2020). Unless otherwise specified, we adopt the training setup of TABPFN for TABFLEX-S100,, TABFLEX-L100, and TABFLEX-HIK Each model is trained on a single Nvidia A100 S0GB PCIe GPU.
Hyperparameters
Batch Size
Epoch
Learning Rate
#Steps/epoch
TABFLEX-S100
1210
8
3e-5
8192
TABFLEX-LI00
110
4
3e-5
8192
TABFLEX-HI
1410
3e-5
1024
Table Hyperparameters used for training TABFLeEX models. The number of steps per epoch indicates the quantity of synthetic datasets generated and used for training within each epach.
Table 6 summarizes the hyperparameters selected for training TABFI.EX-S100, TABFLEX-LI00, and TABFLEE--HK. For all three methods, we utilize the same embedding size of 512, consistent with TABPEN, We extend the feature capacity by modifying the first linear layer, which projects the features into embeddings - specifically, we increase the number of neurons responsible for receiving the features.
The training loss curves are illustrated in Fig. 10. We observe that as the number of features and the length of training dataset sequences increase, the training process becomes more cine-consuming. In fact, training robust FIBFL-X-HIKK model requires more than three weeks.
C.3. Validation Datasets
21

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
ElashAtttcatoon
Fasiillzeardttaationn (Causal)
FlashAttention Libear Attention (Noa Causal)
Flsaiiinaccttttiion (Causal)
Linear Attestion (Non-Causal)
2 Heads
4 Heads
Heads
16 Heads
1e10
Hoads
Hieads
8 Heads
16 Heads
Heed
5000
Head
0
Hoad 64
20000
0
20000
0
20000
0
20000
Head Dim. 256 Head Dim: 128 Head 64
0
20000
0
20000
0
20000
20000
1e10
5000
1-
20000
20000
20000
20000
0
20000
0
20000
0
20000
20000
128
1e10
5000
1-
20000
0
20000
20000
0
20000
0
20000
0
20000
0
20000
0
20000
Dim. 256
1e10
5000
0
20000
0
20000
0
20000
0
20000
0
20000
0
20000
0
20000
0
20000
Sequence Length
Sequence Length
Time
(b) CUDA Memory
Figure 9. Time and CUDA memory usage comparison of FlashAttention-2 (Dao, 2024), causal Flashl innarAAtention (Yang al., 2024), and linear attention (Katharopoulos et al., 2020) (implemented as Listing 1). Results for Flashl inearAttention two confgurratons: (1) head dimension 256 with 8 heads, and (2) head dimension 256 with 16 heads are missing. due to illegal memory access error incurred by the PoTorc package fla (Yang et al., 2024).
We select the validation datasets from the OpenML Au- toML Benchmark (Feurer et al., 2021) by choosing 10 datasets from each of the following sample size intervals: [0.1K, 1K), [IK, 10K), and [10K, 100K). To ensure di- versity the validation set, we also vary the number of classes and features within each interval. The details of all datasets used in validation summarized in Table 7.
1.0
1.0
TabFlex-S100
TabFlex-L100
0.8
0.8
TabFlex-H1K
Loss
0.6
0.6
0.4
0.4
C.4. Sensitivity Analysis on Decision Boundaries
0
50 Epoch
100
250000 500000 Wallclock Time
The decision thresholds align with the training regimes of the models. TABFLEX-S100, sharing TabPFN's 's training setup but with updated architecture, is deployed sim- ilarly (n < 3K, d < 100). TABFLEX-L.10), trained on low-dimensional (d < 100) but larger datasets, is used for longer sequences (n 2 3K, d S 100). TABFLEX-HIK, trained on high- nigh-dimensional data, is assigned to handle those cases accordingly.
Figure 10. Visualization of training loss for TABFLEX models function of epoch and wallclock time.
We note that performance is not highly sensitive the chosen decision boundaries. To demonstrate this, we conducted additional experiments on simple (low-dimension, small-size), dow-dineniional & large, and high- high-dimensional & large datasets two datasets per setting-an present the results for all three models in Table 8.
D. Supplement to Section 6: Performance Evaluation of TABFLEX
D.1. TabZilla Datasets
The results of our experiments on TabZilla-related datasets are reported Table 11, 1, and 2. (McElfresh et al., 2023) presents the details of the datasets used in their hard benchmark (Table 2) in Table of their paper. We provide the
22

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
OpenML did
Dataset
#Features
#Instances
#Classes
279
coota-streamiittevvab..u
75
45164
11
311
oil-spill
50
937
2
742
fri-c4-500-100
101
500
2
825
boston-corrected
21
506
2
833
bank32nh
33
8192
2
841
stock
10
950
2
920
fri-c2-500-00
51
500
2
940
water-treatment
37
527
981
kdd-intermet-usage
69
10108
2
1039
hiva-agnostic
1618
4229
2
1491
one-hundred lalll-pllanntgaeee
65
1600
1492
one-hundred- haapeddrred--laa---sappe
65
1600
100
1503
spoken-arabic-digit
15
263256
10
1515
micro-mass
1301
571
20
1536
volcamoes-b6
4
10130
5
1541
volcamoes-d4
4
8654
5
1549
autoUmiv-au6-750
41
750
8
40645
1001
1600
EDM-1-1
40672
fars
30
100968
8
40677
led24
25
3200
40693
xd6
10
973
2
40705
tokyol
45
959
2
40922
Run-or-walk-informatii
7
88588
2
40985
lamilnadu-clocttrcctt
4
45781
20
4 41082
USPS
257
9298
10
41 144
madeline
260
3140
2
41986
GTSRB-HOGO1
1569
51839
43
41988
GTSRB-HOG02
1569
51839
43
41989
GTSRB-HOG03
2917
51839
43
41990
GTSRB-HueHist
257
51839
43
41991
Kuzushiji-49
785
270912
49
42193
compas-two-years
14
5278
2
42206
porto-seguro
38
595212
2
42343
KDD98
478
82318
2
Table 7. Characteristics of datasets in our diverse validation set.
Dataset
Metric
TabPFN
TabFlex-S100
TabFlex-L110
TabFlex-HIK
Accuracy
0.79
0.82
0.79
0.75
credit-g
Time (s)
0.23
0.13
0.13
0.13
Simple
AUC.
0.78
0.78
0.77
0.78
diabet
Time (s)
0.15
0.09
AUC
0.89
0.89
0.90
0.89
bank-marketing
Time (s)
1.75
0.25
2.43
1.67
Now-Dimensional Large
AUC
0.94
0.94
0.95
0.94
elevators
Time
1.11
0.22
0.7
AUC
0.86
0.83
0.75
0.99
Time (s)
1.95
0.86
4.71
4.63
Hiihhdimennsooall Large
AUC
0.66
0.69
0.59
0.83
SpeedDating
Time
2.86
0.89
1.63
1.71
Table Performance of TARFLEX-SIM TABFLEX-L.100, and TABFLEX-HIK three types of datasets, using TabPFN baseline. We observe that all TABFLEX variants perform well on both simple and www-dimensional large datascts, demonstrating that performance fairly robust the choice of decision threshold.
23

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
specifications of the datasets used for our evaluation in Table 11 and Table 1 in Table 9 and Table 10, respectively.
D.2. Extended Experiment Setup
Baselines. We evaluate our approach against a comprehensive set of baselines, as considered by McElfresh et al. (2023). These include: (i) classical methods: Random Forest (Liaw al., 2002), SVM (Cortes, 1995), LinearModel (Cox, 1958), KNN (Cover& & Hart, 1967) and Decision Tree (Quinlan, 1986); (ii) Gradient Boosted Decision Trees (GBDT) methods: XGBoost (Chen & Guestrin, 2016), CatBoost (Prokhorenkova et al., 2018) 8), and LightGBM (Ke et al., 2017); (iii) Non- Transformer Neural Network (Non-TF NN) methods: SAINT (Somepalli et al., 2021), ResNet (He et al., 2016), DANet (Chen et al., 2022), NODE (Popov et al., 2019), MLP (Rumelhart et al., 1986), MLP-rtdl (Gorishniy et al., 2021), DeepFM (Guo et al., 2017), STG (Yamada et al., 2020), VIME (Yoon et al., 2020), and NAM (Agarwal et al., 2021); (iv) Transformer (TF) methods: TABPFN (Hollmann et al., 2023), FTTransformer (Gorishniy et al., 2021), TabNet (Arik & Pfister, 2021), I), and TabTransformer (Huang et al., 2020). The results for these methods, except TABPFN, are taken directly from McElfresh al. (2023), who conducted their experiments using V100 GPU, while our experiments are run may introduce slight variations in performance. Additionally, we incorporate two recent methods designed for scaling tabular classification: TuneTablees (Feuer et al., 2024), TF method, and HyperFast (Bonet et al., 2024), a Non- TF NN method.
A100 GPU, which
Note that not all baselines successfully ran on all datasets. Many methods face constraints and encounter issues, particularly with the TabZilla hard benchmark, often due to poor scalability. We explicitly indicate which methods failed to run smoothly across all datasets. Originally, TABPFN was limited to datasets with no more than 100 features and 10 classes. To facilitate fair comparison between TABFLEX and TABPEN, we implemented workarounds to prevent TABPFN from encountering errors. For datasets exceeding 100 features, we performed random feature selection. For those with more than 10 classes, we evaluated the accuracy of the nine most prevalent classes and marked all other classes as other, and incorrect. For TuneTTables, we directly import TuneTablesCli wwwsssaasssiiii siirsiiessrrr fier from their Python package t tunetables, Note that our results differ from those reported in their paper, as their study involved more extensive hyperparameter search, which significantly increased runtime. We also compare our methods with TuneTaables using the dataset split specified in their paper's setting, with results deferred Section D.5. Similarly, for HyperFast, we utilize HyperFast aaFFsttCaaa a i fier directly from their Python package hyperfast default parameters. Notably, HyperFast is meta-trained on many datasets use for evaluation.
D.3. Results on 98 Simple Datasets (Table 1, McElfresh et al. (2023))
The results, presented in Table 11, are consistent with the conclusions drawn in the main text.
D.4. Evaluation on Additional Datasets
We provide additional evaluation of TABFLEX on eight large datasets randomly selected from OpenML-CC18 Bench- marks (Bischl et al., 2019), after excluding the datasets contained in TabZilla's 's evaluation. As shown in Table 12, TABFLeX consistently outperforms TABPFN in terms of speed and achieves superior performance on the majority of the datasets.
D.5. Additional Comparison with TuneTabbes
As mentioned in Section 6, the results of TuneTabless presented in Table 13 of our main experiments use TuneTablesClassifierr However, we note that the original paper reported results after 30 iterations of hyper- parameter tuning. They also applied this process to TABPFN, using different subset of datasets as training samples at each iteration. In Table 13, we compare the performance of TABFLEX without any hyperparameter tuning to the results reported in their paper. TABFLEX remains competitive, particularly when the number of samples is limited. While Tune'ables tends to perform better with larger sample sizes due to its ability to update model parameters based on training data, TABFLEX maintains comparable performance while being significantly faster.
D.6. Extending TABFLEX for Image Classification
We explore the application of TABFLEX to image classification tasks, comparing it against MLP and ResNet architectures. Our evaluation uses straightforward configurations without extensive hyperparameter optimization to maintain reasonable computational costs. The MLP implementations include both two-layer and three-layer variants, each configured with 10
24

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Dataset
C
Dataset
C
Dataset
N
C
9
1473
sucmob
5
1156
adull-census
14
32561
kcl
21
2109
vehicle
18
846
4
breast-cancer
9
286
2
kc2
21
522
heart-h
13
294
mfeat-factors
216
2000
10
pe3
1563
jasmine
144
2984
mfeat-zernike
2000
10
po4
37
1458
phoneme
5
5404
dresses-sales
12
500
2
pel
21
1109
semeion
256
1593
10
mfeat-fourier
76
2000
10
cjs
2796
heart-c
303
balance-seaa scale
625
car
6
1728
kr-vs-kp
36
3196
bank-marketing
16
45211
tae
5
1,51
spambase
57
4601
car-ovaluation
21
1728
jml
10885
satimage
6430
cylinder-bands
540
dna
180
3186
mushroom
22
8124
mnreat-karhueen
64
2000
10
musk
167
6598
diabetes
8
768
credit-approval
15
690
2
wdbc
569
rabe rabe_266
2
120
ozone-level-Shr
2534
wilt
5
4839
breast-w
9
699
analcatdata dmft
4
797
ilpd
10
583
elevators
18
16599
monks monks-problems-
6
601
2
2
sick
28
3772
Satellite
36
5100
cardiotocography
35
2126
10
iris
4
150
fertility
9
100
PhishingWebsiie
11055
lymph
148
ionosphere
34
351
synthetic_come cenirtolccccoololllooooo
600
churn
20
5000
transplant
3
131
steel-plates-faus
27
1941
7
colic
22
368
eucalyptus
736
mfcat-
6
2000
10
morphological
ecoli
7
336
Australian
14
690
acute-
6
120
inflammations
autos
25
205
hayes-roth
4
160
analcatdata_ boxing
3
120
scene
299
2407
dermatology
366
analeatdata_chlamydiaa
100
profh
9
672
MiceProtein
77
1080
wall-robot- navigation
24
5456
4
colic
26
368
SpeedDating
120
8378
Bomalizing_llvesttor 2
130
labor
$7
tic-tac-tue
9
958
Click prediction small
39948
irish
500
hill-valley
100
1212
analcatdata authorshhp0
glass
9
214
page- blocks
10
5473
banknote-
4
1372
authentication
yeast
8
1269
lung-cancer
56
32
LED-display-
7
500
10
domain-7digg
sonar
60
208
qsar-biodeg
41
1055
visualizing-
3
111
environmental
splice
60
3190
fri fri_c3_100_5
5
100
postoperative-
8
88
patient-data
libras
104
360
10
ada_agnostt
48
4562
fooddtransfusion
4
748
service-center
anneal
38
898
fri fri_c0_100_S
100
Table Datascts utilized in the evaluation presented in Table respectively.
Here D, N, and denote the number features, instances, and classes,
25

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Dataset
#Features
#Instances
#Classes
Australian
14
690
2
LED-dixplay-domainn-2
7
500
10
MiceProtein
77
1080
8
ccute-inflammations
120
2
analcatdata, auttoosshhttttthippppp
70
841
4
analcatdata boxing!
3
2
analcatdata, callaatdat_chlaayydia
3
2
analcatdataa dmfft anneal
4
797
6
38
898
5
autos
25
205
6
balance-scale
4
625
3
blodd-trans
4
748
2
4
748
2
breast-cancer
2
breast-w
2
26
2
colic
22
2
credit-approval
15
690
2
cylinder-bands
37
540
2
dermatology
34
366
6
diabctes
8
768
2
dresscs-sales
12
500
2
ecoli
7
336
8
eucalyptus
fertility
9
2
ci__00110000 100_5
5
2
fri_c3_100_5
5
100
2
glass
9
214
6
hayes-roth
4
160
3
heart-c
13
303
2
heart-h
13
294
2
hill-valley
100
1212
2
ilpd
583
2
ionosphere
34
351
2
3
irish
5
2
kc2
21
2
labor
16
57
2
lung-cancer
56
32
3
lymph
18
148
4
2roble-rrobeems22
6
601
2
21
1109
2
cioopperatiee-patenn-dee
8
88
2
profb
2
giorebi
41
2
rabe
2
2
socmob
5
1156
2
sonar
60
208
2
synthetic_comie control
60
600
6
tae
5
151
3
tic-tac-toc
9
958
2
transplant
3
131
2
vehicle
18
4
caaualiznng_eviirnmmenn
3
2
sanallizing_lveesttoe
2
wdbe
30
2
yeast
8
1269
4
Table 10. Datasets utilized the evaluation presented in Table
26

---

TABFLEX:
Scaling Tabular Learning to Millions with Linear Attention
Mean AUC
Std. AUC
Time 1000 inst.
Algorithm
Class
median
mean
mean
median
median
mean
TabPFN (Hollmann et al., 2023)
TF
0.97
0.84
0.15
0.08
0.56
0.74
CatBoost (Prokhorenkova ct al., 2018)
GBDT
0.97
0.92
0.15
0.07
1.95
20.51
TABFLEX (Ours)
TF
0.96
0.90
0.15
0.08
0.22
0.37
XGBoost (Chen & Guestrin, 2016)
GRDT
0.96
0.91
0.16
0.09
0.38
0.85
Fores (aaaw et al., 2002)
Classical
0.95
0.90
0.16
0.09
0.32
0.47
SAINT (Somepalli et al., 2021)
TF
0.94
0.86
0.16
0.11
146.15
170.56
Hyperfust (Bonel al., 2024)
Non-TF NN
0.94
0.87
0.15
0.09
53,45
89.75
LightGBM (Ke ul., 2017)
GBDT
0.93
0.85
0.18
0.09
0.29
0.90
ResNet (He ul., 2016)
Non-TF NN
0,93
0.85
0.16
0.10
8.83
15.99
DANet (Chen et al., 2022)
Non-TF NN
0.92
0.85
0.16
0.08
57.18
64.29
NODE (Popov al.. 2019)
Non-TF NN
0.91
0.83
0.16
0.11
131.73
160.76
FTTransformer (Gorishniy et al., 2021)
TF
0.89
0.81
0.17
0.11
18.04
27.91
SVM (Cortes, 1995)
Classical
0.89
0.78
0.19
0.09
2.06
61.18
MLP-rtdl (Gorishniy ct al., 2021)
Non-TF NN
0.88
0.75
0.18
0.11
7.09
15.21
DeepFM (Guo 2017) 2017)
Non-TF NN
0.87
0.77
0.19
0.12
4.89
6.05
TabNet (Arik & Pfister. 2021)
TF
0.85
0.68
0.26
0.14
29.34
35.12
STG (Yamada al., 2020)
Non-TFNN
0.82
0.71
0.20
0.14
15.98
18.58
TuneTTables (Feuer ul., 2024)
TF
0.81
0.70
0.25
0.16
32,96
73.40
LineurModel (Cox, 1958)
Classical
0.78
0.67
0.19
0.14
0,03
0.04
MLP [Rumelhart al., 1986)
Non-TF NN
0.76
0.68
0.20
0.13
11.23
18.31
DecisionTree (Quinlan, 1986)
Classical
0.74
0.63
0.24
0.18
0,01
0.03
Tab Transfoorerr (Huang al., 2020)
0.72
0.17
0.13
13.45
22.05
KNN (Cover Hart, 1967)
Classical
0.70
0.21
0.14
0.03
0.05
VIME (Yoon al., 2020)
Non-TF NN
0.60
0.54
0.25
0.15
15.60
NAM (Agarwal al., 2021)
Non-TF NN
0.39
0.44
0.27
0.19
97.99
233.77
Table 11. Performance comparison algorithms across 98 simple datasets (as used in Table McElfresh al. (2023)). The reported AUC values are normalized. The "Time/1000 inst." column represents the combined training and test time for all datasets, divided by the total number samples. Notably, TABFLEX achieves top performance, with faster runtimes compared to baselines similar performance, and 2x speedup relative to TABPFN.
hidden neurons and trained for 70 epochs fixed learning rate of 0.001. The ResNet architecture employs residual blocks with main and hidden dimension sizes of 128 and 256, respectively. The experimental results demonstrate that TABFLEX achieves remarkable efficiency gains, operating 30x faster than the MLP and 400x faster than the ResNet while maintaining competitive performance. This represents a significant advancement in image classification efficiency, particularly noteworthy given that previous approaches like TABPFN were constrained to small, low-dimensional datasets.
We further evaluate TABFLEX on an 10-way classification CIFAR-10 image dataset (Krizhevsky al., 2009) of 60K samples for which each sample a color image of 32 32 size. We deploy two approaches to convert images to ID vectors. First, we flatten the RGB channels to obtain vector of 3072 dimensions. Second, utilize pretrained ResNet- 18 (He et al., 2016) 6) to obtain a semantically meaningful representation of the image as an 384-dim vector. For each approach, we then feed the ID vectors to TABFLEX as a tabular dataset. As shown in Table 14, with first approach (CIFARIO-ffattene)) achieve an AUC of 0.791 within seconds of inference. We find that reducing feature dimension to 30% with PCA leads to 4 times lower latency while preserving the AUC score. The second approach (CIFAR1O -meeeddinb) significantly increases AUC to 92.2% while reducing the inference latency by six times.
we
E. Supplement to Section 7: Ablation Studies
E.1. Datasets for Section 7.2: Incorporating Data-Efficient Techniques
For dimensionality reduction, we use the following datasets from OpenML: dna, musk, scene, jasmine, semeion, Speed- Dating, hill-valley, mfcat-factors, These datasets are selected from Table 9 where feature dimensions are greater than 100. For the random sampling experiment, we use the following datasets from OpenML: cmc, kcl, car, yeast, car-evaluation, morpholoooplll mfeat-zernike, the data size is greater than 1000 instances, and the feature dimension is lower than 100.
authenticationnn socmob. The tested datasets are selected from Table 9 where
27

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Mean AUC
Mean Time (seconds)
Dataset
#Features
#Instances
#Classes
TABPFN
TABFLEX
TABPFN
TABFLEX
kick
33
72983
2
0.663
0.684
13,330
3.096
Click-prediction- pre--iiooiiiinsallll 1220
10
39948
2
0.652
0.659
3.663
0.887
house-8L
9
22784
2
0.947
0.945
1.383
0.536
okcupid-stem
20
50789
3
0.825
0.828
6.152
1.511
volcanoes-bl
4
10176
5
0.660
0.663
0,349
0.202
volcanoes-b2
10668
5
0.651
0.652
0.375
0.217
kdd-internet-usage
69
10108
2
0.932
0.932
1,021
0.851
BNG(tic-tac-toe)
10
39366
2
0.836
0.835
3.626
1.111
Table 12, Performance comparison between TABPEN and TABFLEX on an additional large dataset. We observe that TABF1EEX consistently faster than TARPEN and outperforms it the majority of the datasets.
TABPFN
Tune Tables
TABFLEX
Dataset
Size
Acc.
Runtime (sec.)
Acc.
Runtime (sec.)
Acc.
Runtime (sec.)
breast-cancer
286
.765
29
.770
65
.793
heart-c
303
.888
40
.903
66
.903
ccoli
336
848
.843
66
.882
colic
368
.856
39
.892
66
.892
dresses-sales
500
578
41
.580
122
.580
cylinder-bands
540
800
41
.846
82
.796
climate
540
9559
59
.951
97
.963
balance-scale
625
.990
.955
1.000
blood- trodstraanfusion
748
.011
25
.782
56
.840
eme
1473
.544
91
.556
109
.605
2109
862
168
187
.867
bioresponse
3151
.797
638
.798
3012
.720
13
christine
5418
.742
666
.755
3920
.721
robert
10000
250
964
.414
2397
.333
dilbert
10000
922
761
.992
3749
.02
17
har
10299
.936
370
.981
2657
.918
ceg-eye-state
14980
.400
178
.986
1929
.837
elevators
16599
902
186
.902
1297
.907
riccardo
20000
.922
1395
.995
5247
.773
volkert
58310
.67
459
.693
6331
.61
12
higgs
67557
.671
931
.714
4084
.691
connect-4
98050
.668
931
.817
5395
.692
BNG (vote)
131072
.668
1976
.974
2493
.974
albert
425240
.642
2363
.658
17518
.637
airlines
539383
.6000
2602
.653
44434
BNG (labor)
1000000
937
5518
.677
7717
950
8
agrawall
1000000
.988
5158
.950
45504
.948
poker-hand
1025009
.511
2423
1.000
10471
.542
15
elick-prediction-smaall
1997410
.333
10421
.837
33148
.333
5
Table Accuracy comparison of TABPFN, TuneTables, and TABFLEX on test datasets from Feuer (2024). Results for TABPFN and Tablla are directly sourced from Feuer al. (2024), where hyperparameter tuning was performed 30 times for both methods. TABPEN. hyperparameters determine the subset dataset used in ICL. TABFLEX results are reported without hyperparameter tuning.
28

---

TABFLEX: Scaling Tabular Learning to Millions with Linear Attention
Two-Layer MLP
Taye MLP
ResNet
TABFLeX (Ours)
Dataset
AUC
Time (8)
AUC
Time (8)
AUC
Time (8)
AUC
Time (s)
MNIST
0.924
23.547 x30.5x)
0.959
23.060 (29.9x)
0.948
Fashion-MNIST
0.793
23.340 (28.8x)
0.853
23.604 (29.1x)
.90
398.45 (491.1))
0.979
0,810
CIFAR-I0 (flateend)
0.791
5.872
CIFAR-10 (embedding)
0.922
0.989
Table 14. Performance comparison of TABFLEX against baseline models image datasets. *Note: MLP and ResNet require significantly more time for training and inference, compared to TABFLEX. Missing evaluation will be provided in the supplementary.
0.700
0.65
0.0225
Accuracy
0.775
9,0215
0.0220
0.770
Time
9,0210
0,0210
0.755
9.0005
0.700
2,5000
0.00
25:0
5003
3500
10500
12500
2500
5000
7500
12500
2500
7500 10503 12520
2500
5900
7500
10000
12500
Samples
Samples
Samples
Samples
800 features
1000 features
Figure 11. Accuracy and runtime versus the number of samples. Two settings are considered: (a) 800 features and (b) 1000 features. Each curve averaged over 20 synthetic datasets with varying data distributions, generated the same algorithm as employed TabPFN (Hollmann et al., 2023).
E.2. Performance and Runtime Training Sample Size
We have demonstrated that TABFLEX performs well across diverse tasks. Here, provide more fine-grained analysis, examining how performance and runtime vary with the number of training samples. Following the setup of TABPFN (Holl- mann et al., 2023), we generate synthetic datasets with sample sizes ranging from 1,000 to 12,000 and feature dimensions of 800 and 1,000. Results are averaged over 20 synthetic datasets and presented in Fig. 11. We observe that accuracy consistently improves with more samples, while runtime increases linearly with the sample size, regardless feature count.
29


Reviews:
[{"review_id": "GnsYooV4cy", "reviewer": "Reviewer_qdGi", "summary": "This paper evaluates scaling the  tabular in-context model TabPFN to larger dataset sizes by using linear attention to circumvent the quadratic memory complexity of regular attention. They first compare to state-space-models (SSMs) like MAMBA, finding SSMs to underperform, attributed to their causal nature (instead of the permutation-equivariant nature of regular or linear attention). They find linear-attention TabPFN to retain most if the predictive accuracy of TabPFN on a large benchmark, while substantially (> 2 times) speeding up the prediction time and allowing to apply TabPFN to substantially larger datasets up to a million datapoints.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer qdGi for the constructive feedback and suggestions. We are encouraged that Reviewer qdGi found that our main claim (on linear attention and its performance) is well-supported by a sensible experiment design on the well-established benchmark. Please find our answers to your comments as follows.\n\n--- \n\n> **Comment**: The performance margin between TabFlex and XGBoost seems substantial and worth discussing potential reasons.\n\nThank you for the suggestion. First, to better understand the regimes where TabFlex and XGBoost outperform each other, we conducted synthetic experiments with feature counts of [100, 200, 400, 600, 800, 1000], generating 20 datasets per count using PFN priors [3]. \n\n**`Experiment results`**:  Accuracy-runtime tradeoff curves are visualized in the following figure: https://anonymous.4open.science/r/icml25_rebuttal_tabflex-7455/Figure1_tablfex_vs_xgboost.png\n\nTabFlex outperforms XGBoost in both accuracy and inference speed when the feature count is below 1K. However, for >1K features, while TabFlex remains faster, it lags in accuracy, especially on the real-world datasets used in the paper. We believe this is due to:\n\n* (i) Training limitation: TabFlex is pretrained on datasets with up to 1K features and struggles to generalize beyond that range.\n* (ii) Distribution shift: PFN-generated synthetic datasets may not reflect real-world distributions, leading to slightly degraded performance at real datasets.\n\nNevertheless, we note that TabFlex offers a great balance between computational cost and performance in many practical scenarios. We now explicitly include this discussion in the revised paper, emphasizing TabFlex's practical efficiency-performance trade-off.\n\n\n> **Suggested references**: TabPFNv2 should be mentioned as concurrent work somewhere. \n\nThank you for the suggestion. TabPFNv2 [4] improves TabPFN’s performance in the low-data regime (fewer than 10,000 samples), which is complementary to our focus on speed and scalability. We’ve added a mention of TabPFNv2 in the updated version as a promising direction for future work that could be combined with our approach.\n\n> **Comment**: No theoretical claim.\n\nThanks for the comment. We do include Theorem 1 (page 5) that supports the high bandwidth memory efficiency of linear attention deployed in our proposed method, demonstrating that the straightforward implementation of linear attention achieves linear HBM access, matching the performance of FlashLinearAttention after optimization.\n\n> **Comment**: Subsampling of 3000 for TabPFN seems a bit arbitrary. \n\nThanks for the comment. We follow the TabZilla framework [1], which recommends subsampling to 3000 due to the quadratic scaling of TabPFN's runtime and memory usage. This setting is commonly adopted [2] for benchmarking TabPFN across diverse datasets. \n\n---\n\n*References:*\n\n[1] McElfresh, Duncan, et al. \"When do neural nets outperform boosted trees on tabular data?\" NeurIPS (2023).\n\n[2] Feuer, Benjamin, et al. \"Tunetables: Context optimization for scalable prior-data fitted networks.\" NeurIPS (2024).\n\n[3] Müller, Samuel, et al. \"Transformers can do bayesian inference.\" ICLR (2022).\n\n[4] Hollmann, Noah, et al. \"Accurate predictions on small data with a tabular foundation model.\" Nature (2025).\n\n**Final Note**: We hope that these answers will allay any concerns about our work and convince the reviewer that it will be a welcome contribution to the ICML community. If there are additional questions that we can address to further support our case, please let us know.", "claims_and_evidence": "In general the claims regarding linear attention and its performance for TabPFN seeem well-supported, although the random subsampling of 3000 for TabPFN seems a bit arbitrary, a subanalysis looking at a range of subset sizes may be interesting here.\n[Update: seems 3000 is common practice]\n\n\"Notably, TABFLEX is faster and achieves better performance than TABPFN, and is faster than XGBoost while sacrificing only a small margin of performance.\" -> To me, this does not look like a small margin, this seems like a very substantial margin, and worth discussing potential reasons for it as well.\n[Update: authors provided some analysis]", "methods_and_evaluation_criteria": "Yes the benchmark datasets make sense, as written above 3000 seems a bit arbitrary. [Update: seems 3000 is common practice]", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "As written,  evaluate on well-established benchmarks which make sense.", "relation_to_broader_scientific_literature": "Insights into how well in-context tabular methods can be scaled, in terms of providing results for one particular way to do so.", "essential_references_not_discussed": "TabPFNv2  (https://www.nature.com/articles/s41586-024-08328-6) should be mentioned as concurrent work somewhere.\n[Update: promised to be done]", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gYztc1gIJk", "reviewer": "Reviewer_K5HQ", "summary": "The paper proposes an in-context learning architecture for tabular learning. The overall method is incremental to TabPFN.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer K5HQ for the detailed review and constructive suggestions. \nWe greatly appreciate that Reviewer K5HQ found that (i) our main claim (on improving scalability over TABPFN by linear attention) is well-supported by both theoretical analysis and corresponding empirical evaluation, (ii) our method is sensible with (iii) very good evaluation on large benchmarks and good theoretical claim. We are also encouraged that the Reviewer finds our supplementary material good and well-designed.\nPlease find our answers to your comments and questions as follows.\n\n---\n\n> **Comment**: The theoretical claim seems good, but the kernel feature mapping lacks a good explanation.\n\nThank you for the comment. Theorem 1 originally does not assume any kernel feature mapping. We update it to incorporate elementwise kernel mappings (commonly used, e.g., $\\text{elu}(\\cdot) + 1$), while the original statement still holds. Since the mapping is applied elementwise after loading queries and keys, it adds no extra HBM access or memory, and the additional $2DN$ FLOPS is negligible compared to the total $O(ND^2)$ FLOPS.\n\nSee **`Extended version of Theorem 1`** here: https://shorturl.at/MFFwF\n\n> **Comment**: The decision thresholds (of conditional model selection) seem somewhat arbitrary.\n\nThe decision thresholds for model selection are **not arbitrary but aligned with the training regimes of the models**. TabFlex-S100, sharing TabPFN’s training setup but with an updated architecture, is deployed similarly $(n ≤ 3K, d ≤ 100)$. TabFlex-L100, trained on low-dimensional $(d ≤ 100)$ but larger datasets, is used for longer sequences $(n \\geq 3K, d ≤ 100)$. TabFlex-H1K, trained on high-dimensional data, is assigned to handle those cases accordingly.\n\nWe note that **performance is not highly sensitive to the chosen decision boundaries**. To demonstrate this, we conducted *additional experiments* on simple (low-dimension, small-size), low-dimensional & large, and high-dimensional & large datasets—two datasets per setting.\nOur new results for all three models, presented in this table **https://shorturl.at/HFaOD** demonstrate our claim.\n\n> **Comment**: The claim *TabFlex generalizes image classification tasks well* is not well supported since it only contains MLP and ResNet and only two classes of very simple image datasets. Remove this section or move it to the appendix.\n\nBased on your comment, we have moved the results to the Appendix and clarified that image classification is not a core claim. Our goal was to explore TabFlex’s adaptability beyond tabular data; the 10-class MNIST and CIFAR10 results offer preliminary support for this, now framed as a side observation.\n\n> **Suggested References**. \n\nBased on your suggestion, we added the following discussion.\nTabR proposes a retrieval-augmented model with a custom kNN-like component to retrieve and extract signals from the nearest neighbors. \nBiSHop establishes interconnected directional learning modules to process data column-wise and row-wise for\ntabular learning.\nXTab utilizes independent featurizers and federated learning to resolve inconsistent column types and quantities.\n\n\n> **Q**: Is there any theory or existing work to support that non-causal attention generally outperforms causal attention?\n\nA few empirical works [1,2] support our observation that causal attention is suboptimal in the ICL setting. For theoretical work, there is no direct comparison, but most of the theoretical work on ICL is based on non-causal attention [3,4]. We will add them into related works. \n\n\n> **Comment**: The evaluation lacks regression tasks, though the claim of this paper is on classification.\n\nThanks for the comment. We update our limitation discussion to acknowledge this: \"For tabular tasks, our current focus is limited to classification. A simple workaround for regression is to discretize the target range into bins and treat it as a classification problem. An interesting future work is extending TabFlex to regression tasks with a more principled approach involving training on regression-specific synthetic data.\"\n \n\n> **Comment**: Mamba vs. Transformer and Softmax Attention vs. Linear Attention: More analysis into the main text?\n\nWe do not include additional theoretical analysis due to the analytical complexity. However, we recognize this as an interesting future direction and will add it to the discussion section.\n\n--- \n*References*:\n* [1] Ding, Nan, et al. \"CausalLM is not optimal for in-context learning.\" ICLR (2024).\n* [2] Gong, Zhuocheng, et al. \"Improving Input-label Mapping with Demonstration Replay for In-context Learning.\" arXiv (2023).\n* [3] Ahn, Kwangjun, et al. \"Transformers learn to implement preconditioned gradient descent for in-context learning.\" NeurIPS (2023). \n* [4] Bai, Yu, et al. \"Transformers as statisticians: Provable in-context learning with in-context algorithm selection.\" NeurIPS  (2023).\n\n**Final notes**: Thanks again for your positive feedback and thoughtful suggestions.", "claims_and_evidence": "- **TabFlex improves scalability over TABPFN by using linear attention instead of quadratic attention.**: This claim is well supported by their theoretical analysis aligns with empirical results. \n- **TabFlex generalizes image classification tasks well.**: I don't really think this claim is well supported since it only contains MLP and ResNet and only two classes of very simple image datasets. I suggest the authors remove this section from the main text or move it to the appendix. \n- **Non-causal attention generally outperforms causal attention**: the author just empirically analyzes this idea. Is there any theory or existing work to support that claim? \n- **Mamba vs. Transformer** and **Softmax Attention vs. Linear Attention**: there is also lack of mathematically analysis but that should be fine. If more analysis put on the main text it will be better.", "methods_and_evaluation_criteria": "- The proposed methods make sense in the tabular learning area. The method use linear attention to reduce the cost, conditional model selection for different dataset sizes and feature dimensions (but the decision thresholds seem somewhat arbitrary).\n- The evaluation is very good and includes a lot of datasets and benchmarks. The main issue is lacking of regression tasks even the claim of this paper is on classification.", "theoretical_claims": "The theoretical claim of **Linear Attention Reduces Complexity** seems good to me. \nHowever, the kernel feature mapping lacks of good explanation.", "experimental_designs_or_analyses": "The experimental design and analysis do not seems to have any issues since it follows the most popular tabular benchmarks.", "relation_to_broader_scientific_literature": "The paper's method is an incremental to the previous method: TabPFN and shows the potential goodness of to the tabular learning field.", "essential_references_not_discussed": "Some most recent tabular deep learning methods are not mentioned and discussed, such as [1], [2], [3]\n\n[1] Gorishniy, Yury, et al. \"Tabr: Tabular deep learning meets nearest neighbors in 2023.\" arXiv preprint arXiv:2307.14338 (2023).  \n[2] Xu, Chenwei, et al. \"Bishop: Bi-directional cellular learning for tabular data with generalized sparse modern hopfield model.\" arXiv preprint arXiv:2404.03830 (2024).  \n[3] Zhu, Bingzhao, et al. \"Xtab: Cross-table pretraining for tabular transformers.\" arXiv preprint arXiv:2305.06090 (2023).", "other_strengths_and_weaknesses": "I don't see other strengths and weaknesses.", "other_comments_or_suggestions": "Just remove or move the image classification task from main text or to the appendix.", "questions_for_authors": "I don't have questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KTEn5D3nHY", "reviewer": "Reviewer_1peU", "summary": "This paper explores the use of linear attention for TabPFN, to overcome its limitations in terms of scalability. Indeed, TabPFN is limited by the quadratic complexity of self-attention, making it inefficient for datasets with more than a few 1,000s samples. The authors:\n- Demonstrate experimentally that linear attention is preferable to causal linear attention as in State-Space models.\n- Introduce TabFlex, a TabPFN variant with linear attention, comprising three models trained on datasets of varying sizes (up to 1,152 vs. 50K samples, 100 vs. 1K features, and 10 vs. 100 classes). TabFlex selects the appropriate model based on dataset size and class count.\n- Evaluate TabFlex on a benchmark of small datasets (98 + 57 datasets), showing it matches TabPFN's performance—already the best among many baselines including boosted trees—while achieving a 2× speedup.\n- Evaluate TabFlex on the Tabzilla hard benchmark (36 datasets) where it remains on par with TabPFN, offering 2× speedup. It is however inferior to other models including XGBoost, albeit slightly faster.\n- Evaluate TabFlex on vectorized CIFAR-10 images.\n- Show that feature and sample downsampling can further improve TabFLEX efficiency without affecting performances.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 1peU for the insightful feedback and constructive suggestions. We are very excited that you truly enjoyed reading our work and finds that (i) our claims are well supported by clear and convincing evidence on well-chosen evaluation criteria, providing interesting finding on performance perservation, (ii) our method is 'natural' and sound with rigorous analysis, (iii) our paper is very well-written. Please find our answers to your comments and questions as follows.\n\n\n> **Comment**: TabFlex does not provide SoTA performance in harder cases (e.g., larger datasets), suggesting that linear attention scaling alone may be insufficient, especially compared to XGBoost.\n\nThanks for raising this point.\n\nWhile our goal is not to surpass SoTA but to improve the efficiency-performance trade-off over TabPFN, we ran additional experiments comparing TabFlex with XGBoost on 20 synthetic datasets (from PFN priors). \n\n\n**`Experiment results`**: https://shorturl.at/tH02z\n\nTabFlex outperforms XGBoost in both accuracy and inference speed when the feature count is below 1K. However, for >1K features, while TabFlex remains faster, it lags in accuracy, especially on the real-world datasets used in the paper. We believe this is due to:\n* (i) Training limitation: TabFlex is pretrained on datasets with up to 1K features and struggles to generalize beyond that range.\n* (ii) Distribution shift: PFN-generated synthetic datasets may not reflect real-world distributions, leading to slightly degraded performance at real datasets.\n\nNevertheless, TabFlex offers a great balance between computational cost and performance in many practical scenarios. We now explicitly include this discussion in the revised paper, emphasizing TabFlex's practical efficiency-performance trade-off.\n\n> **Suggestion**: Analyze performance and speedup as a function of the number of training samples. \n\nGreat suggestion. \nWe evaluate TabFlex on a set of synthetic 1000-feature datasets generated from prior distributions developed by PFN. \nWe vary the number of training samples and visualize the change in accuracy and inference latency.\n\n**`Experiment results`**:  https://shorturl.at/dWap2\n\nAs the number of training samples increases, the accuracy indeed improves, and the inference latency (or speed) linearly scales to the number of training samples, as expected.\n\n\n\n> **Q**: Will pre-training across all dataset sizes simultaneously allow a single model to perform well in all cases instead of using three separate models?\n\nGreat question. For TabFlex-H1K, we do pre-train across all dataset sizes simultaneously to leverage the full sample set. However, this significantly slows down training (see Fig. 8 in the Appendix), and even after convergence, the model slightly underperforms compared to TabFlex-S100 on small datasets (see  https://shorturl.at/HFaOD). This trade-off is why we opted for an ensemble approach with thresholds based on each model’s specialized training regime.\n\n> **Q**: How would incorporating linear attention impact TabPFNv2?\n\nTabPFNv2 improves TabPFN's performance in the low-sample regime (<10K examples). Since our work targets scalability—an orthogonal aspect—extending our approach to TabPFNv2 could potentially yield benefits in both scalability and performance. We agree this is a promising and interesting direction, and we will add this point to the discussion as suggested.\n\n\n> **Comment**: The performances on image classification should be evaluated with the SoTA (i.e., vision architectures) to highlight accuracy-latency trade-off.\n\nWe clarify that image classification is not the focus of our work. The evaluation was intended to explore TabFlex’s adaptability beyond tabular data. We agree that comparing with SoTA vision models would better highlight the limitations of tabular approaches and will note this as a future direction.\n\n\n> **Suggested references**.\n\nThanks for the suggestion.\nMixturePFN (Xu et al 2024) improves scalability by routing new test samples to a pool of scalable prompters using Sparse Mixture of In-Context Prompters, while LoCalPFN (Thomas et al 2024) proposes retrieving a local subset of task-specific data for efficiently fine-tuning on. \nMa et al 2024 introduce in-context data distillation to optimize TabPFN’s context and remove the data size constraint.\nWe integrated this discussion into our revision.\n\n> **Minor comments**:\n- Typos: We revised the typos\n- Impact statement: We revised to clarify the confusion.\n- Fig.2: thanks, we updated Fig. 2.\n- benchmarks: we added the explanation.\n- Fig. 6: We vary the fraction as a common choice for the datasets of different sizes. We added visualization over the sample size to provide a clearer signal. \n\n---\n**Final Note:** Thanks again for your insight and thoughtful suggestions. We hope that our responses and new experiment results help you better appreciate our work as well as support accepting our paper.", "claims_and_evidence": "Yes, claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are well chosen.", "theoretical_claims": "There is one theoretical claim comparing the computational and memory complexity of causal FlashLinearAttention and non-causal linear attention.  I did not check the correctness of this claim.", "experimental_designs_or_analyses": "The experimental designs are sound.", "relation_to_broader_scientific_literature": "This paper investigates whether TabPFN’s scalability limitations can be addressed with linear attention while preserving performance.", "essential_references_not_discussed": "The following recent papers address TabPFN’s scalability challenges using approaches different from linear attention. Positioning this work in relation to these references would provide useful context.\n\n[1] Xu et al 2024, Mixture of In-Context Prompters for Tabular PFNs. \n\n[2] Thomas et al 2024, Retrieval & Fine-Tuning for In-Context Tabular Models. \n\n[3] Ma et al 2024, In-Context Data Distillation with TabPFN.", "other_strengths_and_weaknesses": "**Strengths**\n* I truly enjoyed reading this paper: it is very well written, very pedagogical, and its claims are well supported by the experiments.\n\n* Given the limitations of TabPFN, exploring whether linear attention could enhance scalability without significant performance loss is a natural direction. This paper provides a rigorous answer.\n\n* The finding that linear attention tends to preserve performance compared to the original self-attention in TabPFN is interesting.\n\n**Weaknesses**\n* A limitation of TabFlex is that in harder cases (may be for larger datasets in particular), it does not provide state-of-the-art performances. In Fig. 4, TabFlex seems significantly outperformed by XGBoost, and overall does not compare favourably to the 5th best method in Table 2. It indicates that scaling the training to larger datasets thanks to linear attention is not enough to obtain sota performance. \n\n* The paper does not analyze performance and speedup as a function of the number of training samples. Since TabFlex replaces the quadratic complexity of TabPFN with linear complexity via linear attention, one would expect the speedup to scale with $n$, the number of samples. It would be helpful to highlight this experimentally.\n\nMoreover, linear attention enables TabFlex to train on larger datasets, suggesting that its performance could improve relative to TabPFN as the dataset size increases. However, the current hard benchmark combines large datasets with high-dimensional and other challenging datasets, making it difficult to isolate this effect. A more targeted analysis of how TabPFN, TabFlex, and other baselines such as XGBoost perform as a function of $n$ would be valuable.\n\n* Since TabFlex consists of three models, a natural question is whether it could be pre-trained across all dataset sizes simultaneously, allowing a single model to perform well in all cases instead of using three separate models. Have the authors explored this possibility?\n\n* I am not convinced of the significance of the results on image classification experiments. The performances should be evaluated in relation to the state-of-the-art (i.e. vision architectures) as fast inference time alone may not be useful if it comes at the cost of significantly lower accuracy.\n\n* Since TabPFN has now been significantly outperformed by TabPFN2 [4] (released in January 2025), a natural question is how incorporating linear attention would impact TabPFN2.\n\n[4] Hollmann et al, Accurate predictions on small data with a tabular foundation model", "other_comments_or_suggestions": "**Typos**\n\nL.173 With more samples are provided => When\n\nL. 261: reported in respectively reported in Tables 1 and 2\n\nL. 420: relatively well performance\n\n**Minor comments**\n\n* In the impact statement: “Our approach enables an efficient solution for utilizing LLMs” => as TabPFN and TabFlex are not LLMs, I think that the vocabulary should be updated to avoid confusions.\n\n* In Figure 2, it would be nice to have both TabPFN-causal-masked and TABPFN-Mamba on both Figure2a and Figure2b, as both experiments (in Fig. 2a and Fig. 2b) are interesting and do not reveal the same information. \n\n* It is difficult to see from Fig 2c the overall accuracy difference between Softmax attention and Linear attention across all datasets. May be a scatterplot delta_runtime vs delta_accuracy could improve the readability.\n\n* It would be helpful to clarify the specifics of the benchmarks chosen in this paper, such as explaining roughly what simple and hard benchmarks mean.\n\n* On fig. 6 about data subsampling: I am a bit surprised that there is to little improvement after using 20% of the training data. Maybe looking at fraction of training data rather than absolute number of samples hides the signal, i.e, may be having 3,000 samples helps compared to 1,000 samples, but 10,000 instead of 5,000 not so much.", "questions_for_authors": "See weaknesses.\n\nIn particular the second point on analyzing the effects of the number of samples on performance and speed-up.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission investigates making the TabPFN-class of models more scalable. Itshows that a linear attention is preferrable to state space models and scales well the framework.  The submission led to interest and discussion with the reviewers, which appreciated the gain in scalability and the empricial study. The work should be situated with regards to the TabPFNv2 paper which came out a bit before the deadline and markedly improves upon the original TabPFN.
Decision: Accept (spotlight poster)