arXiv:2603.00175cs.CV2026-03

提出线性自注意力机制,突破图像模型长序列计算瓶颈。

Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention

  • 用谱方法重构建模自注意力,将注意力视为图上随机游走的累积过程。
  • 线性版本仅需固定大小状态,支持4096×4096输入,推理速度提升13倍。
  • 在ImageNet上比传统Transformer高3.2分,且显存占用最低,适合大图任务。

softmax自注意力的二次计算复杂度限制了视觉Transformer在高分辨率场景下的扩展。本文提出无限自注意力(InfSA),通过谱重构将每层注意力视为内容自适应的令牌图上的扩散步骤,利用折扣的Neumann级数累积多跳交互。该方法将自注意力与经典图中心性(Katz、PageRank、特征向量中心性)关联,实现可解释的令牌权重分配。进一步发现Neumann核等价于吸收马尔可夫链的基本矩阵,即令牌中心性等于其被访问的期望次数。随后提出线性-InfSA,无需构建完整注意力矩阵即可近似主特征向量,仅需与每头维度dh成比例的固定大小辅助状态,兼容视觉变压器,支持4096×4096训练和9216×9216推理(约33.2万令牌)。在4层ViT(5350万参数,224×224下59 GFLOPs)上,线性-InfSA在ImageNet-1K上达到84.7%的准确率,较同等深度的softmax ViT提升3.2个百分点。在ImageNet-V2上,InfViT变体优于所有对比基线(最高79.8%对76.8%),显示分布偏移下的鲁棒性。在A100 40GB GPU上,线性-InfViT每秒处理231张图像,能耗0.87焦/图像,吞吐量与能效分别为同规模ViT的13倍和13倍,并是唯一能在9216×9216下完成推理而不出现显存不足的模型。线性近似与二次算子的主导特征向量高度一致(余弦相似度0.985)。

原文摘要 · Abstract (English)

The quadratic cost of softmax attention limits Transformer scalability in high-resolution vision. We introduce Infinite Self-Attention (InfSA), a spectral reformulation that treats each attention layer as a diffusion step on a content-adaptive token graph, accumulating multi-hop interactions through a discounted Neumann series over attention matrices. This links self-attention to classical graph centrality (Katz, PageRank, eigenvector centrality) for interpretable token weighting. We also show the Neumann kernel equals the fundamental matrix of an absorbing Markov chain, so a token's centrality is its expected number of random-walk visits before absorption. We then propose Linear-InfSA, a linear-time variant that approximates the principal eigenvector of the implicit attention operator without forming the full attention matrix. It keeps an auxiliary state of fixed size proportional to per-head dimension dh (independent of sequence length N), is drop-in compatible with Vision Transformers, and supports stable training at 4096 by 4096 and inference at 9216 by 9216 (about 332k tokens). In a 4-layer ViT (53.5M parameters, 59 GFLOPs at 224 by 224), Linear-InfSA reaches 84.7% top-1 on ImageNet-1K, a +3.2 point architectural gain over an equal-depth softmax ViT trained with the same recipe. On ImageNet-V2, InfViT variants outperform all compared baselines (up to 79.8% vs 76.8%), indicating robustness under distribution shift. On an A100 40GB GPU, Linear-InfViT runs at 231 images/s and 0.87 J/image (13x better throughput and energy than equal-depth ViT) and is the only tested model to complete 9216 by 9216 inference without out-of-memory. The linear approximation closely matches the dominant eigenvector of the quadratic operator (cosine 0.985).

自注意力线性模型视觉变换器高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。