Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention
Authors: Giorgio Roffo, Hazem Abdelkawy, Nilli Lavie, Luke Palmer
Organizations: Equixly API Security, Italy · Toyota Motor Europe, Belgium · MindVisionLabs, UK · Institute of Cognitive Neuroscience, University College London, UK · GlimpseML, UK
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).
Figures & tables
Figure 1: Comparison of attention graphs. Visualization of ViT-L/16 attention maps on ImageNet. Softmax attention distributes focus across background regions, while InfSA variants produce sharper, object-aligned activations.
Figure 2: (a) InfSA in a Pre-LN ViT block. Two InfSA variants within standard Transformer scaffolding: (1) Pure InfSA uses full attention with ReLU and Frobenius normalization, accumulating discounted outputs across layers; (2) Linear InfSA computes soft token scores, pools values per head, and broadcasts context with per-layer scaling. Both are drop-in compatible with Transformer blocks. (b) Efficiency by complexity tier (4L, 10242 ). Inference throughput vs. energy per image for nine attention mechanisms, colored by asymptotic complexity. InfViT Linear ( O(N) , red star) achieves the highest throughput at the lowest energy cost.
Figure 3: Softmax attention (1-hop) vs. InfSA (Neumann series). Left: Frobenius-normalized A^ ( ∥A^∥F=1 ); row sums vary, unlike softmax. Middle: Absorbing Markov chain M=γA^ ; dashed red arrows show absorption Ri=1−∑jMij into a . Right: Starting from the all-ones input e (Eq. 8 ): softmax attention (row-stochastic, 1-hop) ranks token 0 first via column sums, since many tokens directly attend to it. InfSA iterates M further; at n=2 the chain 0→3→4 redirects mass to token 4, and the Katz centrality cin correctly identifies token 4 as globally most important—the multi-hop outcome standard self-attention misses.
Figure 4: Efficiency dashboard (4L-64H, inference at 10242 ). Speed-up over Standard ViT, absolute throughput, and energy per image. InfViT Linear ( O(N) ) achieves 13.4× speed-up at 0.87 J/img.
Table 1: Scalability benchmark on A100 40 GB (batch = 1). Inference at 10242 (4 096 tokens, patch 16); training at 5122 . Energy: E=Pˉ⋅Δt ; Pˉtr=300 W, Pˉinf=200 W. Type: Q = quadratic, L = linear/sub-quadratic, I = InfSA. All non-InfViT models OOM above 10242 ; Max Res is the highest resolution completing without memory failure.
Figure 5: Attention quality summary. MoRF-AOC, ROC-AUC, and PR-AUC (%). Both InfSA variants outperform Standard ViT by 20–34 pp. Full curves in the Supplement (Figs. 7 , 8 ).
Figure 6: Accuracy vs. parameters. ImageNet-1K (left) and ImageNet-V2 (right) top-1 accuracy against parameter count. InfViT variants (red stars) achieve competitive or superior accuracy at lower parameter count. On IN-V2, all four InfViT models exceed every baseline.
Metric
Mean
Std. dev.
Cosine similarity
0.985
0.036
Spearman correlation
0.937
0.132
Table 2: Alignment between the Perron eigenvector v of A^=ϕ(QQ⊤) and the Linear-InfSA weight vector a (Eq. 27 ), across 512 per-head samples on a frozen trained checkpoint.
Figure 7: MoRF degradation and ROC curves. Left: MoRF degradation (1 000 images, mean ±1σ ); steeper drop = more focused attention. Right: patch-level ROC against ImageNet bounding boxes (2 088 images). InfSA variants consistently outperform Standard ViT.
Figure 8: LeRF retention and Precision–Recall curves. Left: LeRF retention (Least Relevant First removal; AUC ↑ ; 1 000 images, mean ±1σ ). Right: Precision–Recall curves against ImageNet bounding boxes (2 088 images, mean ±1σ ). Pure InfSA achieves the highest LeRF AUC (65.3%), and Linear InfSA the highest PR-AUC (76.1%), confirming that InfSA produces both focused and semantically discriminative attention.
Method
Year
Params (M)
FLOPs (G)
IN-1K
IN-V2
Softmax / standard attention
ViT-B/16 [ 13 ]
’21
86
17.6
81.5
74.1
DeiT-B [ 74 ]
’21
86
17.5
81.8
71.5
HorNet-B [ 43 ]
’22
88
15.5
84.3
—
SG-Former-B [ 45 ]
’23
78
15.6
84.7
—
SMT-L [ 35 ]
’23
81
17.7
84.6
—
Table 3: Top-1 accuracy (%) on ImageNet-1K (IN-1K) and ImageNet-V2 (IN-V2), grouped by attention type. Bold = best in column. — = not reported. All models trained on IN-1K only (no external data).
Path-decay γ
Activation
Metric
0.3
0.5
0.7
0.9
ReLU
GELU
∣x∣
Top-1 (%)
84.2
84.3
84.7
83.5
84.7
83.5
82.8
MoRF AOC (%)
72.5
74.7
76.0
75.8
76.0
73.2
75.0
LeRF AUC (%)
62.1
63.4
64.9
65.3
64.0
62.8
65.5
ROC-AUC (%)
75.1
75.4
76.0
76.2
76.6
74.8
77.2
Latency (ms)
37/11
35 /12
36/ 11
37/12
35 /11
36/12
36/12
Table 4: Full ablation results for Linear InfViT-4L-64H. Left: path-decay γ sweep (ReLU fixed). Right: activation sweep ( γ=0.7 fixed). Bold = selected default.
Figure 9: 4L vs. 24L depth comparison (inference at 10242 ). Throughput (left) and energy (right) for the 4L-64H and 24L-16H configurations. InfViT Linear leads in both regimes; the gap narrows at 24L due to the higher fixed overhead of deeper networks, but the ranking is preserved.
Figure 10: Attention maps on ImageNet-1K validation samples. For each method (softmax, Pure InfSA, Linear InfSA): original image, attention heatmap, and overlay. InfSA variants consistently focus on object-centric regions; softmax attention exhibits diffuse or background-focused activation.
Figure 11: Latency vs. input resolution (RTX 5090, 32 GB). Dotted lines: training; solid: inference. Linear InfViT scales near-linearly, with differences becoming clearer at high resolutions. Note: this figure uses an RTX 5090 (32 GB); all other benchmarks in the paper use an A100 40 GB. Relative scaling trends are consistent across GPUs.