表格基础模型注意力机制的基准测试

arXiv:2609.31306v3 公告类型：替换 摘要：TabPFN、Mitra 或 ConTextTab 等表格上下文学习器依赖于在潜在嵌入的二维序列上交替进行行注意力和列注意力计算。这些注意力模式与语言模型中的一维情形显著不同：行注意力涉及较长的序列，而列注意力处理的序列短得多，且表格数据的跨步内存布局使得生成连续张量的成本很高。此外，当前模型使用的隐藏维度相较于近期的语言模型较小。然而，高效注意力的研究主要集中于一维序列，二维表格场景尚未得到探索。为此，我们构建了一套可复现的基准测试环境，并在多个后端上研究表格注意力的独有特征——包括 Torch SDPA（efficient 和 cuDNN）、FlashAttention-2/3/4，以及仅用于推理的后端 vLLM 和 SageAttention——在三代 GPU（A100、H100、B200）上，针对符合实际的表格张量形状测量前向和反向计算的吞吐量。我们发现，列注意力与行注意力的最优后端选择不同，并且会随硬件和模型具体配置而变化：尽管针对各代 GPU 定制的 FlashAttention 实现总体表现最佳，但在序列较长的列注意力场景中，CuDNN 有时表现更优，性能交叉点取决于注意力头维度。在仅用于推理的后端中，SageAttention 在行注意力以及超过 16,000 行的大规模序列上表现良好。我们可复现的基准测试为未来改进表格原生注意力机制奠定了基础。自包含的基准测试与评估代码已公开发布于：https://github.com/SAP-samples/tabular-attention-benchmark