无限自注意力及其以外：面向线性 Transformer 的无限自注意力
摘要
Softmax 注意力的二次计算成本限制了 Transformer 在高分辨率视觉中的可扩展性。我们引入了无限自注意力（InfSA），一种谱方法，将每个注意力层视为内容自适应 token 图上的扩散步骤，通过注意力矩阵上的折扣 Neumann 级数累积多跳相互作用，将自注意力与经典图中心性（Katz、PageRank、特征向量中心性）关联，以实现可解释的 token 加权。我们还表明，Neumann 核等于吸收马尔可夫链的基本矩阵，因此 token 的中心性是其在吸收之前的预期随机漫步访问次数。随后，我们提出了 Linear-InfSA，一种线性时间变体，用于在不形成完整注意力矩阵的情况下近似隐式注意力算子的主特征向量。它保留了一个大小与每头维度 成正比的辅助状态（独立于序列长度 ），可无缝集成到 Vision Transformer 中，并支持在 4096×4096 和 9216×9216（约 33 万个 tokens）上的稳定训练和推理。在一个 4 层 ViT（5350 万参数，224×224 时 59 GFLOPs）中，Linear-InfSA 在 ImageNet-1K 上达到 84.7% 的 Top-1 准确率，相较于相同深度的 Softmax ViT 训练的相同配方提升了 3.2 个百分点。在 ImageNet-V2 上，InfViT 变体在所有比较基线中均表现优于（最高达 79.8% vs 76.8%），表明其在分布迁移下的鲁棒性。在 A100 40GB GPU 上，Linear-InfViT 的吞吐量达 231 张/秒，能耗为 0.87 J/张（相较于相同深度的 ViT 吞吐量和能耗分别提升约 13 倍），且是唯一一个在 9216×9216 上完成推理而不出现内存溢出的模型。该线性近似与二次算子的主特征向量之间夹角余弦为 0.985。
引用
@article{arxiv.2603.00175,
title = {Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention},
author = {Giorgio Roffo and Hazem Abdelkawy and Nilli Lavie and Luke Palmer},
journal= {arXiv preprint arXiv:2603.00175},
year = {2026}
}
备注
This work was initiated and primarily carried out while working at MindVisionLabs. We gratefully acknowledge the support of Toyota Motor Europe (TME) and Equixly API Security for this work