TabFlex: Scaling Tabular Learning to Millions with Linear Attention

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.

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.

Thanks for raising this point.

While 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).

Experiment results: https://shorturl.at/tH02z

TabFlex 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:

• (i) Training limitation: TabFlex is pretrained on datasets with up to 1K features and struggles to generalize beyond that range.
• (ii) Distribution shift: PFN-generated synthetic datasets may not reflect real-world distributions, leading to slightly degraded performance at real datasets.

Nevertheless, 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.

Suggestion: Analyze performance and speedup as a function of the number of training samples.

Great suggestion. We evaluate TabFlex on a set of synthetic 1000-feature datasets generated from prior distributions developed by PFN. We vary the number of training samples and visualize the change in accuracy and inference latency.

Experiment results: https://shorturl.at/dWap2

As 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.

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?

Great 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.

Q: How would incorporating linear attention impact TabPFNv2?

TabPFNv2 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.

Comment: The performances on image classification should be evaluated with the SoTA (i.e., vision architectures) to highlight accuracy-latency trade-off.

We 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.

Suggested references.

Thanks for the suggestion. MixturePFN (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. Ma et al 2024 introduce in-context data distillation to optimize TabPFN’s context and remove the data size constraint. We integrated this discussion into our revision.

Minor comments:

• Typos: We revised the typos
• Impact statement: We revised to clarify the confusion.
• Fig.2: thanks, we updated Fig. 2.
• benchmarks: we added the explanation.
• 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.

---

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.

We thank Reviewer K5HQ for the detailed review and constructive suggestions. We 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. Please find our answers to your comments and questions as follows.

---

Comment: The theoretical claim seems good, but the kernel feature mapping lacks a good explanation.

Thank 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.

See Extended version of Theorem 1 here: https://shorturl.at/MFFwF

Comment: The decision thresholds (of conditional model selection) seem somewhat arbitrary.

The 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.

We 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. Our new results for all three models, presented in this table https://shorturl.at/HFaOD demonstrate our claim.

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.

Based 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.

Suggested References.

Based on your suggestion, we added the following discussion. TabR proposes a retrieval-augmented model with a custom kNN-like component to retrieve and extract signals from the nearest neighbors. BiSHop establishes interconnected directional learning modules to process data column-wise and row-wise for tabular learning. XTab utilizes independent featurizers and federated learning to resolve inconsistent column types and quantities.

Q: Is there any theory or existing work to support that non-causal attention generally outperforms causal attention?

A 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.

Comment: The evaluation lacks regression tasks, though the claim of this paper is on classification.

Thanks 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."

Comment: Mamba vs. Transformer and Softmax Attention vs. Linear Attention: More analysis into the main text?

We 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.

---

References:

• [1] Ding, Nan, et al. "CausalLM is not optimal for in-context learning." ICLR (2024).
• [2] Gong, Zhuocheng, et al. "Improving Input-label Mapping with Demonstration Replay for In-context Learning." arXiv (2023).
• [3] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." NeurIPS (2023).
• [4] Bai, Yu, et al. "Transformers as statisticians: Provable in-context learning with in-context algorithm selection." NeurIPS (2023).

Final notes: Thanks again for your positive feedback and thoughtful suggestions.

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.

---

Comment: The performance margin between TabFlex and XGBoost seems substantial and worth discussing potential reasons.

Thank 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].

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

TabFlex 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:

• (i) Training limitation: TabFlex is pretrained on datasets with up to 1K features and struggles to generalize beyond that range.
• (ii) Distribution shift: PFN-generated synthetic datasets may not reflect real-world distributions, leading to slightly degraded performance at real datasets.

Nevertheless, 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.

Suggested references: TabPFNv2 should be mentioned as concurrent work somewhere.

Thank 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.

Comment: No theoretical claim.

Thanks 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.

Comment: Subsampling of 3000 for TabPFN seems a bit arbitrary.

Thanks 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.

---

References:

[1] McElfresh, Duncan, et al. "When do neural nets outperform boosted trees on tabular data?" NeurIPS (2023).

[2] Feuer, Benjamin, et al. "Tunetables: Context optimization for scalable prior-data fitted networks." NeurIPS (2024).

[3] Müller, Samuel, et al. "Transformers can do bayesian inference." ICLR (2022).

[4] Hollmann, Noah, et al. "Accurate predictions on small data with a tabular foundation model." Nature (2025).

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.