TabSwift: An Efficient Tabular Foundation Model with Row-Wise Attention

Technical Summary: TABSWIFT

Problem Statement
Tabular foundation models (TFMs), such as TabPFN, have demonstrated competitive performance on small-to-medium datasets by leveraging in-context learning (ICL) without test-time parameter updates. However, recent advancements in TFMs often rely on increasingly complex architectures (e.g., alternating row/column attention) to capture feature-wise heterogeneity. These complex designs significantly increase inference costs and memory usage, limiting their practical deployment in latency-sensitive scenarios. While lightweight row-wise attention-only backbones (like TabPFN v1) are efficient, they risk hitting a performance ceiling on highly heterogeneous tasks. The central question addressed is whether a lightweight, row-wise attention-only backbone can achieve competitive foundation-model performance when enhanced with modern stabilization and pretraining techniques, thereby offering a superior accuracy-efficiency trade-off.

Methodology
The authors propose TABSWIFT, a lightweight tabular foundation model that retains the efficient row-wise attention-only structure while introducing three key components:

1. Stabilized Row-Wise Backbone:

   • Architecture: The model treats each table row as a token and applies self-attention only across rows (instances), avoiding the expensive column-wise attention passes found in models like TabPFN v2 or TabICL.
   • Gated Attention: To improve optimization stability during pretraining, the authors incorporate an element-wise gated attention mechanism. This applies a multiplicative gate to the scaled dot-product attention (SDPA) output before the projection, introducing input-dependent sparsity and non-linearity with negligible overhead.
   • Learnable Register Tokens: A small set of learnable register tokens is prepended to the input sequence. These tokens act as latent slots to aggregate global context and task-level information, helping the model maintain and refine dataset-level representations across layers.

2. Unified Pretraining Formulation:

   • TABSWIFT is trained on a unified objective that supports both classification and regression within a single pretrained model.
   • The pretraining pipeline utilizes a large pool of synthetic tabular tasks (20,000 steps, 512 tasks per step). For each synthetic task, targets are processed to provide both a discretized classification target and a normalized regression target, allowing simultaneous optimization of separate lightweight MLP heads for each task type.

3. Adaptive Layer-Wise Early-Exit Mechanism:

   • To further reduce inference latency, the model employs a per-query adaptive early-exit strategy.
   • Mechanism: Lightweight prediction heads and "exit heads" are attached to the last $E$ transformer layers. The exit head takes the current query representation concatenated with a pooled summary of the register tokens to predict a reliability score.
   • Policy: During inference, the model processes layers sequentially. If the predicted reliability score for a specific query exceeds a threshold $\tau$, the model exits early and outputs the prediction. This allows "easy" samples to be processed with fewer layers while "hard" samples traverse the full depth, enabling "anytime" inference.

Key Contributions

• TABSWIFT Model: Demonstrates that a row-wise attention-only backbone, when stabilized by gated attention and augmented with register tokens, achieves performance competitive with heavier, more complex TFMs (e.g., TabPFN v2, TabICL) while maintaining significantly lower inference costs.
• Unified Architecture: Introduces a single pretrained model capable of handling both classification and regression tasks, simplifying deployment.
• Adaptive Inference: Proposes a reliability-predictive early-exit strategy that dynamically adjusts computation depth per sample, substantially reducing average inference time with minimal loss in predictive performance.
• Open Source: The code is made available to the community.

Experimental Results
The authors evaluated TABSWIFT on the TALENT benchmark, a large-scale collection of 300 binary, multi-class, and regression datasets.

• Performance: TABSWIFT achieves an average rank comparable to the strongest baselines (TabPFN v2, TabICL, LimiX) on both classification (AUC) and regression (RMSE) tasks. Statistical testing (Wilcoxon-Holm) indicates it is often not significantly different from top-performing methods.
• Efficiency: Inference time measurements show TABSWIFT operates at a substantially lower latency than complex TFMs. Figure 1 illustrates that TABSWIFT attains a markedly better accuracy-efficiency trade-off, positioned closer to the ideal point of high performance and low latency.
• Ablation Studies: Experiments confirm that both gated attention and register tokens contribute to performance gains over a baseline row-wise attention model.
• Early-Exit Analysis: Varying the exit threshold $\tau$ reveals a favorable Pareto frontier. The learned exit heads, conditioned on register tokens, outperform entropy-based stopping baselines. Visualizations of the embedding space show that the model correctly identifies ambiguous or overlapping samples as requiring deeper computation, while confident samples exit early.

Significance and Claims
The paper claims that TABSWIFT enables efficient and anytime tabular in-context learning suitable for practical deployments where inference budgets are strict. By revisiting the minimalist row-wise design and applying modern stabilization techniques, the authors argue that the community does not necessarily need to adopt computationally expensive alternating attention patterns to achieve state-of-the-art results. The work suggests that TABSWIFT serves as a strong, efficient component for model selection or ensembling, offering a robust solution for latency-sensitive serving in real-world tabular applications.