2025-11-12T14:19:10.228100

State-Space Models for Tabular Prior-Data Fitted Networks

Koch, Wever, Raisch et al.

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.

academic

State-Space Models for Tabular Prior-Data Fitted Networks

基本信息

论文ID: 2510.14573
标题: State-Space Models for Tabular Prior-Data Fitted Networks
作者: Felix Koch, Marcel Wever, Fabian Raisch, Benjamin Tischler
分类: cs.LG
发表时间/会议: Proceedings of the 1st ICML Workshop on Foundation Models for Structured Data, Vancouver, Canada. 2025
论文链接: https://arxiv.org/abs/2510.14573

摘要

研究背景与动机

要解决的问题: 本研究针对表格数据基础模型中Transformer架构的计算效率问题，特别是其O(n²)的复杂度限制了在大型数据集上的可扩展性。
问题的重要性: TabPFN作为表格数据的基础模型展现了出色的性能，能够在毫秒级别内完成贝叶斯推理近似，但其基于Transformer的架构在处理大规模数据时面临内存和计算瓶颈。
现有方法的局限性:
- Transformer的自注意力机制具有二次复杂度
- 直接用Mamba替换Transformer会引入对输入序列顺序的敏感性
- 表格数据中行的顺序在语义上是无意义的，这与SSM的因果性设计相冲突
研究动机: 探索结构化状态空间模型(SSM)作为Transformer的替代方案，既保持线性复杂度的效率优势，又通过双向处理机制减少对输入顺序的依赖。

核心贡献

提出了基于Hydra的TabPFN架构: 将双向结构化状态空间模型Hydra集成到TabPFN中，实现线性时间复杂度的表格数据处理。
引入重复上下文排列(RCP)技术: 通过多次随机排列输入并平均预测结果来进一步减少SSM对序列顺序的敏感性。
实现了显著的可扩展性提升: 相比原始TabPFN，新方法能处理两个数量级更大的数据集(从2¹⁵行扩展到2¹⁷行)。
保持了竞争性的预测性能: 在OpenML CC-18基准测试中，Hydra-based TabPFN的准确率仅比原始模型低1.1%。

方法详解

任务定义

本文研究表格分类任务，其中：

输入: 包含训练和测试样本的完整表格数据集
输出: 对测试样本的类别概率预测
约束: 需要在单次前向传播中完成推理，无需梯度更新或微调

模型架构

1. Hydra架构替换

核心设计: 用Hydra层堆叠替换Transformer编码器
双向处理: 利用准可分离矩阵混合器实现双向状态空间建模
层结构: 每个Hydra层包含双向状态空间混合，后跟前馈变换

2. 嵌入策略保持

保留原始TabPFN的数据嵌入方法
每个输入表示为特征值和类标签的拼接
推理时通过边缘化所有可能的标签分配处理未标记数据

3. 重复上下文排列(RCP)

算法流程如下：

输入: 排列次数r, 上下文D, 测试样本xtest
输出: 预测的类别值
初始化空列表: outputs ← []
for i = 1 to r do
    打乱D的行: Dp ← shuffle(D)
    拼接xtest到Dp: Din ← Dp ∪ xtest
    预测: outputs[i] ← PFN.predict(Din)
end for
返回 outputs的平均值