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

TL;DR

This paper introduces Infinite Self-Attention (InfSA), a spectral reformulation of self-attention that enables linear-time transformers with high scalability and robustness, demonstrated on vision tasks with significant efficiency gains.

Contribution

The authors propose a novel spectral reformulation called Infinite Self-Attention and a linear-time variant, Linear-InfSA, improving scalability and interpretability of transformers for high-resolution vision tasks.

Findings

Linear-InfSA achieves 84.7% top-1 accuracy on ImageNet-1K.

Linear-InfSA supports inference on 9216 by 9216 tokens, outperforming previous models.

Linear-InfSA runs 13x faster and more energy-efficient than equal-depth ViT.

Abstract

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).