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

表形式データ用の状態空間モデルに基づく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

概要

TabPFNなどの表形式データ用基盤モデルの最近の進展により、事前学習されたTransformerアーキテクチャがベイズ推論を高い予測性能で近似できることが実証されました。しかし、Transformerはシーケンス長に対して二次の計算量を持つため、より効率的なシーケンスモデルの探索が求められています。本研究では、双方向線形時間構造化状態空間モデル(SSM)であるHydraをTabPFNのTransformerの代替案として使用する可能性を調査します。主な課題は、SSMが入力トークンの順序に対して本質的に敏感であることです。これは、行の順序が意味的に無意味である表形式データセットにとって望ましくない特性です。双方向アプローチが効率性を保持し、対称的なコンテキスト集約を可能にする程度を調査します。実験結果から、このアプローチが順序依存性を低減し、元のTabPFNモデルと競争力のある予測性能を達成することが示されました。

研究背景と動機

解決すべき問題: 本研究は、表形式データ用基盤モデルにおけるTransformerアーキテクチャの計算効率の問題に対処しています。特に、O(n²)の計算量複雑性は大規模データセット上でのスケーラビリティを制限しています。
問題の重要性: TabPFNは表形式データの基盤モデルとして優れた性能を示し、ミリ秒単位でベイズ推論の近似を完了できますが、Transformerベースのアーキテクチャは大規模データ処理時にメモリと計算のボトルネックに直面しています。
既存手法の限界:
- Transformerの自己注意機構は二次の計算量を持つ
- TransformerをMambaで直接置き換えると、入力シーケンスの順序に対する感度が導入される
- 表形式データでは行の順序は意味的に無意味であり、これはSSMの因果性設計と矛盾する
研究動機: 構造化状態空間モデル(SSM)をTransformerの代替案として探索し、線形計算量の効率性の利点を保持しながら、双方向処理機構を通じて入力順序への依存性を低減することです。

核心的貢献

HydraベースのTabPFNアーキテクチャの提案: 双方向構造化状態空間モデルHydraをTabPFNに統合し、表形式データの線形時間計算量処理を実現しました。
繰り返しコンテキスト順列(RCP)技術の導入: 入力を複数回ランダムに順列化し、予測結果を平均化することで、SSMのシーケンス順序に対する感度をさらに低減します。
顕著なスケーラビリティの向上: 元のTabPFNと比較して、新手法は2桁大きなデータセット(2¹⁵行から2¹⁷行に拡張)を処理できます。
競争力のある予測性能の維持: OpenML CC-18ベンチマークテストにおいて、Hydraベースの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の平均値を返す

技術的革新点

双方向性による順序感度の解決: 単方向のMambaと比較して、Hydraの双方向処理はコンテキスト情報を対称的に集約でき、入力順序への依存性を低減します。
線形計算量: 準可分離行列乗算器を通じてO(n)計算量を実現し、Transformerの O(n²)と比較して顕著な利点があります。
RCP戦略: 複数回のランダム順列と結果平均化を通じて順序感度をさらに低減する革新的な方法で、これは表形式データの特性に対する定制化設計です。