State-Space Models for Tabular Prior-Data Fitted Networks

AI summary

Overview

Research area: Machine learning — tabular foundation models, specifically Prior-Data Fitted Networks (PFNs) and efficient sequence-model backbones (state space models).

Technical level: Intermediate. Familiarity with Transformers, attention complexity, and in-context learning helps, but the paper's core argument is accessible.

Scope: This paper asks whether bidirectional state space models, specifically Hydra, can replace the Transformer inside TabPFN for tabular classification while preserving accuracy and removing the quadratic scaling bottleneck.

What This Paper Is About

TabPFN performs tabular classification by feeding an entire labeled dataset into a pretrained Transformer and producing predictions in a single forward pass. Because the Transformer's attention has quadratic complexity in sequence length (which corresponds to the number of table rows), this approach is limited to small datasets. The authors test Mamba and Hydra, two linear-time state space models, as drop-in replacements for the Transformer backbone, and address the fact that SSMs are sensitive to input order even though row order in a table carries no meaning.

Key Contributions

1. Hydra as a TabPFN backbone. The paper replaces the Transformer encoder in TabPFN with a stack of Hydra layers, a bidirectional state space model built on quasiseparable matrix mixers, retaining the original embedding format (feature values concatenated with class labels) and requiring only retraining on the same synthetic prior.
2. Repeated Context Permutations (RCP). A simple inference-time procedure that predicts on r randomly shuffled orderings of the context rows and averages the resulting predicted probability distributions, reducing order dependence at a linear cost of r times the inference time.
3. An empirical comparison of Mamba, Hydra, and Transformer TabPFN. The authors evaluate all three on 30 filtered OpenML CC-18 multiclass datasets, each randomly split into training and test sets 16 times, reporting mean and standard error.
4. A quantitative treatment of order sensitivity. They use KL divergence between predictions on two shuffled contexts as a direct measure of how much the row ordering affects a model's output distribution.

Main Findings

• Linear complexity translates into much larger feasible inputs. On an H100 with 80GB VRAM, Transformer-based TabPFN failed at 2^16 rows because the quadratic self-attention matrix exceeded available memory, while Hydra only failed at 2^18 rows, and there it hit PyTorch's 32-bit indexing limit rather than a hardware limit. The paper notes the Transformer's scalability can be improved by a constant factor of roughly 20 with FlashAttention.
• Hydra outperforms Mamba on average. Mamba showed higher variance than Hydra in both AUC OvO and accuracy. Hydra achieved 3.6% higher accuracy than Mamba on average, which the authors attribute to the advantage of bidirectional over unidirectional processing.
• Hydra stays close to the Transformer. Hydra's average difference to the Transformer-based TabPFN was 1.1%, and Hydra attained the best performance on some datasets.
• RCP reduces order sensitivity. Increasing the number of repeated context permutations decreased the KL divergence between predictions from shuffled contexts, confirming the intended effect.
• RCP gives a small, variance-limited accuracy gain. Accuracy increased with the number of permutations up to a point, which the authors attribute to averaging out outliers caused by disadvantageous row orderings. They describe this effect as comparably small relative to the large variance, and they skip AUC for this analysis because no significant improvements were observed there.
• The evaluation is deliberately constrained. Datasets were filtered to at most 2000 rows, at most 100 features, and at most 10 classes to fit TabPFN's constraints, and context length was limited to 1000 rows.

Methodology in Plain English

The researchers took the existing TabPFN training pipeline and swapped out its Transformer encoder for two alternatives: Mamba, which reads the table's rows in one direction only, and Hydra, which reads them in both directions. Because the new backbones have different internal structure, they retrained each model on the same synthetic task distribution used for the original TabPFN — a Transformer run of 48 hours, Mamba 52 hours, and Hydra 134 hours, each on a single Nvidia A40 GPU with 48GB of memory.

To test scalability, they generated random tensors with 99 feature columns and row counts from 2^5 up to 2^18, ran each model 10 times on an NVIDIA H100 80GB node, and measured inference time until memory ran out. To test quality, they compared accuracy and AUC OvO across 30 OpenML CC-18 datasets, computing each SSM's difference from the Transformer per dataset and reporting averages over 16 random train/test splits. To test order sensitivity, they duplicated a run with shuffled context rows and measured the KL divergence between the two predicted class distributions, then repeated this under varying numbers of context permutations.

Why This Matters

The work targets the root cause of TabPFN's scalability limit — the quadratic attention mechanism — rather than working around it with post-hoc tricks. It shows that a bidirectional SSM is a viable backbone for tabular prior-data fitted networks, and it makes the case that bidirectionality is the key ingredient, since the unidirectional Mamba lost considerably more accuracy.

Real-world applications:

• Clinical and biomedical tabular datasets, where sample sizes often sit at the edge of TabPFN's feasible range and where the paper's own benchmark includes datasets such as breast-w, diabetes, and MiceProtein.
• Financial and credit risk scoring, matching datasets like credit-g, credit-approval, and banknote-authentication, where inference latency and memory footprint matter in production.
• Industrial and sensor-based tabular prediction, such as steel-plates-fault and climate-model-simulation-crashes, where feature counts and class counts stay modest but row counts can grow.
• Any setting requiring calibrated predictions in milliseconds without gradient-based adaptation, which is the central selling point of PFNs and is preserved here.

Industry relevance: the practical argument is cost and scale. Linear-time inference over rows means larger tables can be processed on the same hardware, and the memory ceiling moves far enough away that the remaining barrier becomes a software indexing limit rather than a hardware one.

Future Directions

• Test SSMs on much longer contexts. The authors limited context to 1000 rows to match prior TabPFN work and name longer contexts, for example more than 10k rows (the current limit for TabPFNv2), as the most promising next step.
• Further reduce row-order dependence. They call for additional work on mitigating the impact of row order on SSM-based tabular PFN predictions.
• Search for beneficial orderings. The paper supports the assumption that certain context row orderings may actively improve SSM performance, and suggests the optimal ordering may differ between unidirectional and bidirectional SSMs.
• Broader hyperparameter study. The authors state that thorough hyperparameter optimization was infeasible due to long training times, and that the proposed steps-per-epoch and batch size from prior TabPFN work caused catastrophic forgetting in both Mamba and Hydra, forcing manual changes.

Target Audience

Researchers and practitioners working on tabular machine learning, in-context learning, and efficient sequence architectures. It is especially relevant to those building or deploying tabular foundation models, to engineers concerned with inference memory and latency on large tables, and to readers interested in whether state space models can substitute for attention outside of language and long-sequence domains. Readers with a background in Transformers and attention complexity will get the most out of the experimental sections.

Authors’ abstract

Recent advancements in foundation models for tabular data, such as TabPFN, demonstrated that pretrained Transformer architectures can approximate Bayesian inference with high predictive performance. However, Transformers suffer from quadratic complexity with respect to sequence length, motivating the exploration of more efficient sequence models. In this work, we investigate the potential of using Hydra, a bidirectional linear-time structured state space model (SSM), as an alternative to Transformers in TabPFN. A key challenge lies in SSM's inherent sensitivity to the order of input tokens - an undesirable property for tabular datasets where the row order is semantically meaningless. We investigate to what extent a bidirectional approach can preserve efficiency and enable symmetric context aggregation. Our experiments show that this approach reduces the order-dependence, achieving predictive performance competitive to the original TabPFN model.