Tydra：用于表格数据的高效混合模型

# Tydra：用于表格数据的高效混合模型  
来源：https://arxiv.org/html/2608.21199  
Saurabh Mathur、Maurice Kraus、Sriraam Natarajan、Kristian Kersting  

###### 摘要  
基于Transformer的表格基础模型（如TabPFN）在预测性能上表现强劲，但其计算成本随上下文长度呈二次增长。另一方面，基于次二次状态空间模型（SSM）的替代方案（如Hydra）以精度换取效率。为平衡两者，我们提出Tydra——一种用于表格上下文学习的混合Transformer–状态空间模型（SSM）架构，交替使用注意力层和SSM层。在30个OpenML数据集上，Tydra在保留TabPFN大部分预测性能的同时，将推理时间减少了30%。Tydra还优于规模约大十倍的Hydra模型，同时提供更快的推理速度。结果表明，混合架构是表格基础模型的一个有前景的发展方向。

## 引言  
表格基础模型（如TabPFN）在多个表格预测任务中取得了强劲的预测性能，无需任务特定训练（10 [https://arxiv.org/html/2608.21199#bib.bib1]）。然而，TabPFN的Transformer主干随上下文长度呈二次扩展，使得在大规模或长上下文表格任务中的推理成本高得令人望而却步。这种成本对无法使用服务器托管推理的机构尤其具有限制性。医院和其他受监管的机构通常因数据保护原因禁止将患者记录发送到外部服务器，必须在计算资源有限的本地硬件上运行推理。虽然Hydra架构通过基于次二次状态空间模型（SSM）的架构降低了这一成本，但其预测性能不及TabPFN（12 [https://arxiv.org/html/2608.21199#bib.bib4]；13 [https://arxiv.org/html/2608.21199#bib.bib3]）。即使扩大Hydra的规模，这种差距依然存在；将TabPFN缩小以匹配Hydra的效率会降低其精度。两种纯架构都无法同时实现高精度和低推理成本。  

参见标题图1：Tydra在表格数据上表现迅速且不牺牲预测性能。Tydra与TabPFN和Hydra在OpenML数据集上的推理速度和预测性能对比。颜色编码了组合的精度-速度分数（F2，偏向速度）。阴影椭圆表示每个模型平均值的一个标准差范围。星形标记了理想角落：最快且最准确。Hydra 16M速度快但预测性能明显落后，而TabPFN和Hydra 160M准确但慢；Tydra实现了接近TabPFN和Hydra 160M的预测性能，同时推理速度快两倍以上（2×）。混合架构（交替使用注意力层和SSM层）最近在语言模型中改善了这种精度-效率权衡（16 [https://arxiv.org/html/2608.21199#bib.bib5]），尽管这种架构在表格基础模型中尚未被探索。表格数据在行和列上是置换不变的，不同于混合架构所设计的固定顺序语言序列。为此，我们提出Tydra，它将TabPFN的注意力层与Hydra的SSM层交错。我们表明，这种混合化在中小型规模数据上以显著更低的推理成本匹配了TabPFN的精度，使得强大的表格上下文学习在服务器托管推理不可用的场景中变得可行。  

总结来说，我们做出以下贡献：  
1. 我们提出了Tydra——首个用于表格上下文学习的混合Transformer–SSM架构。  
2. 我们表明，Tydra在30个OpenML数据集上以高达30%的更低推理时间匹配了TabPFN的精度，同时显著优于Hydra。  
3. 我们对Tydra的架构族进行了广泛研究，展示了Hydra和TabPFN层不同组合与比例的优势。  

本文结构如下：首先讨论相关工作，然后介绍Tydra架构，最后展示实证结果。

## 相关工作  
Tydra与多个研究方向相关，即TabPFN、状态空间模型和混合模型。

### TabPFN  
TabPFN（表格先验数据拟合网络）是一种表格基础模型，通过上下文学习执行表格分类。给定整个训练集和测试查询，它通过单次前向传递进行预测。其架构基于Transformer，自注意力机制被调整为跨行和列的置换不变性。这个调整后的Transformer在数百万个合成分类任务上进行离线元训练，无需任何数据集特定训练即可在广泛表格数据集上获得强大的预测精度。然而，由于自注意力随表格数据集行数呈二次扩展，推理可能在计算上昂贵。自原始发布以来，已有多个扩展针对TabPFN的范围和扩展。TabPFN v2（10 [https://arxiv.org/html/2608.21199#bib.bib1]）将原始仅分类模型（9 [https://arxiv.org/html/2608.21199#bib.bib16]）扩展到回归并增加了模型容量，而后续版本TabPFN v2.5和v3（7 [https://arxiv.org/html/2608.21199#bib.bib17]）大致将Transformer深度加倍并提高了支持的类别数量。这些连续版本提高了精度和任务覆盖范围，但保留了相同的注意力机制，因此继承了其随行数二次增长的成本。认识到这一瓶颈，另一系列工作提出了*事后*策略来扩展TabPFN的实际应用范围，无需重新训练或架构更改：子采样和分而治之方案将大型数据集分割为TabPFN大小的块并聚合预测（19 [https://arxiv.org/html/2608.21199#bib.bib18]），以及特定领域的适应（如用于时间序列预测的TabPFN-TS（11 [https://arxiv.org/html/2608.21199#bib.bib19]））。这些努力表明扩展限制已被广泛认识，但它们是绕过Transformer的行扩展行为，而非消除它。相比之下，我们的方法在架构层面解决了这一瓶颈。

### 状态空间模型  
状态空间模型（SSMs）已成为克服RNN和Transformer局限性的有效方法。源自连续时间动态系统，它们实现了序列长度的近线性复杂度（17 [https://arxiv.org/html/2608.21199#bib.bib7]）。Mamba（8 [https://arxiv.org/html/2608.21199#bib.bib8]）是一个突出例子，引入了选择性状态空间机制，自适应地过滤序列中的信息；这种选择性使Mamba能够实现线性时间推理，同时在多个领域匹配或超越Transformer性能（8 [https://arxiv.org/html/2608.21199#bib.bib8]）。Hydra（12 [https://arxiv.org/html/2608.21199#bib.bib4]）通过拟可分离矩阵混合器将Mamba扩展到非因果设置，实现双向上下文聚合同时保留Mamba的效率——这一特性与表格数据尤其相关，因为行没有自然序列顺序。Hydra先前已被应用于表格设置（13 [https://arxiv.org/html/2608.21199#bib.bib3]）。虽然TabPFN在标准中小型数据集上的推理速度和预测精度均优于Hydra，但在更大数据集规模上，Hydra成为更好的选择，因为TabPFN的二次成本变得过高，而Hydra的次二次扩展使其保持实用。这表明SSMs相对于基于注意力的表格基础模型提供了真正的效率优势，但以预测性能为代价，特别是在标准中小型数据集上。

### 混合模型  
结合Transformer层与更高效次二次层的混合架构已成为效率-精度权衡的有效解决方案。混合语言模型在常见的推理、数学和编码任务中取得了有前景的结果（16 [https://arxiv.org/html/2608.21199#bib.bib5]；14 [https://arxiv.org/html/2608.21199#bib.bib6]）。已提出多种混合策略，包括交替Transformer和状态空间块（15 [https://arxiv.org/html/2608.21199#bib.bib10]）、在递归主干上复用少量注意力块（6 [https://arxiv.org/html/2608.21199#bib.bib9]），以及在相同输入上并行计算递归和注意力分支（5 [https://arxiv.org/html/2608.21199#bib.bib11]）。然而，混合化迄今几乎只针对语言建模探索，仅有限扩展到其他模态（20 [https://arxiv.org/html/2608.21199#bib.bib12]），并且据我们所知，先前没有工作为表格基础模型引入混合架构。Tydra正好填补了这一空白，采用交错策略结合TabPFN的注意力层和Hydra的状态空间层。

## Tydra架构族  
特征编码器 \( E_x \)  
标签编码器 \( E_y \)  
特征向量 \((X_{\mathrm{Train}}, X_{\mathrm{Test}})\)  
训练标签 \( Y_{\mathrm{Train}} \)  
行级表示 \(\text{Concatenate}(E_x(X_{\mathrm{Train}}) + E_y(Y_{\mathrm{Train}}), E_x(X_{\mathrm{Test}}))\)  
Hydra层 双向状态空间混合  
Transformer编码器层 多头注意力和前馈网络 \(\times N\)  
预测头 具有GELU的MLP解码器  
测试行logits → softmax → 类概率 \( P(Y_{\mathrm{Test}} \mid X_{\mathrm{Train}}, Y_{\mathrm{Train}}, X_{\mathrm{Test}}) \)

图2：Tydra的混合架构。特征向量被投影为行级嵌入，训练行添加了标签嵌入。得到的序列通过四个Hydra–Transformer编码器对（共八层），每对顺序为Hydra后接注意力。MLP解码器为测试行生成类logits，通过softmax转换为类概率。

| 模型 | 架构 |
|------|------|
| Hydra {8H} | HHHHHHHT |
| Tydra {4HT} | THTHTHTH |
| Tydra {4TH} | HTHTHTHT |
| Tydra {6T} | HHTTTTTTHT |
| Tydra {2H}{4T}{2H} | HHTTTTHH |
| Tydra {T}{6H}{T} | THHHHHH |
| Tydra {2T}{4H}{2T} | TTHHHH |

纯架构与Tydra混合架构族  
图3：Tydra架构族。评估的Hydra和Tydra模型的架构。块从底部输入层到顶部输出层排序。T表示Transformer注意力层，H表示Hydra层。

我们现在介绍Tydra——一种用于表格上下文学习的混合架构族，它将注意力层与Hydra的状态空间（SSM）层交错。本节描述Tydra的架构、模型初始化和训练。

### 混合架构  
图2（https://arxiv.org/html/2608.21199#Sx3.F2）展示了我们的混合先验拟合架构Tydra<sup>[1]</sup>。为不失一般性，我们将完全交错的Tydra {4HT}模型简称为Tydra。其他Tydra变体可通过修改主干部分表示。给定训练特征 \( X_{\mathrm{train}} \)、训练标签 \( Y_{\mathrm{train}} \) 和测试特征 \( X_{\mathrm{test}} \)，Tydra估计 \( p(Y_{\mathrm{test}} \mid X_{\mathrm{train}}, Y_{\mathrm{train}}, X_{\mathrm{test}}) \)。Tydra将Hydra高效的双向序列混合与自注意力提供的内容相关交互相结合。每个表格行由一个token表示。特征和标签使用线性编码器 \( E_x: \mathbb{R}^n \rightarrow \mathbb{R}^m, E_y: \mathbb{R} \rightarrow \mathbb{R}^m \) 分别嵌入，在本例中 \( n=10 \)，\( m=512 \)。对于有标签的训练行 \( (x_i, y_i) \)，初始表示由两个嵌入相加形成： \( z_i^{(0)} = E_x(x_i) + E_y(y_i) \)，而无标签的测试行仅由其特征嵌入表示： \( z_i^{(0)} = E_x(x_i) \)。因此，训练行和测试行共享相同的 \( m \) 维token空间，而标签仅对训练行可见。得到的序列由 \( K=4 \) 个有序的Hydra–TabPFN对处理，共八层：
\[
\underbrace{[\mathrm{Hydra} \rightarrow \mathrm{TabPFN}] \times K}_{2K \text{层}}.
\]
每一层保持 \( m \) 维表示，允许Hydra和注意力层无需额外投影层即可交错。Hydra层使用12 [https://arxiv.org/html/2608.21199#bib.bib4]的双向状态空间混合器，遵循13 [https://arxiv.org/html/2608.21199#bib.bib3]对其适应表格先验拟合网络的调整。每个Transformer层包含多头自注意力和具有隐藏维度 \( 2m \) 的两层前馈网络。两个子层都使用残差连接和层归一化。最后，与测试行对应的表示通过两层MLP： \( \mathbb{R}^m \rightarrow \mathbb{R}^{2m} \rightarrow \mathbb{R}^C \)，使用GELU激活，其中 \( C \) 表示类别数。得到的类logits通过softmax归一化以获得预测类概率。

### 双向状态空间混合  
表格数据没有规范的行顺序，即置换 \( X_{\mathrm{Train}} \) 的行不应改变模型的预测。这种置换不变性区分了语言和表格上下文学习。标准SSM从左到右处理序列。因此，天真地用TabPFN的集合Transformer层替换SSM会破坏置换不变性。Hydra层通过用双向状态空间混合器替换标准SSM来解决此问题，该混合器允许每一行在任一方向上关注其他每一行，同时仍允许线性时间评估。设 \( L \) 表示序列长度， \( D \) 表示通道数。形式上，应用于 \( \mathbf{X} \in \mathbb{R}^{L \times D} \) 的因果选择性SSM（如单向Mamba中使用）充当矩阵混合器 \( \mathbf{Y} = \mathbf{M} \mathbf{X} \)，其中混合器矩阵的条目 \( m_{ts} \) 为：
\[
m_{ts} = \begin{cases}
\mathbf{c}_t^\top \Big( \prod_{k=s+1}^{t} \mathbf{A}_k \Big) \mathbf{b}_s, & t > s, \\
\delta_t, & t = s, \\
\overleftarrow{\mathbf{c}}_t^\top \Big( \prod_{k=t}^{s-1} \overleftarrow{\mathbf{A}}_k \Big) \overleftarrow{\mathbf{b}}_s, & t < s,
\end{cases}
\]
其中 \( \mathbf{A}_k \in \mathbb{R}^{N \times N} \) 和 \( \mathbf{b}_k, \mathbf{c}_k \in \mathbb{R}^N \) 是位置 \( k \) 处的离散化状态、输入和输出矩阵， \( N \) 是SSM状态维度。这样的 \( \mathbf{M} \) 是 \( N \)-*半可分离*的：从下三角形取的每个子矩阵的秩最多为 \( N \)，这恰好允许 \( O(L) \) 递归评估，但也强制 \( \mathbf{M} \) 严格因果（对角线上方为零）。Hydra（12 [https://arxiv.org/html/2608.21199#bib.bib4]）通过将混合器参数化为*拟可分离*矩阵来解除这一限制：