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.
While linear-complexity attention mechanisms offer a promising alternative to Softmax attention for overcoming the quadratic bottleneck, training such models from scratch remains prohibitively expensive. Inheriting weights from pretrained Transformers provides an appealing shortcut, yet the fundamental representational gap between Softmax and linear attention prevents effective weight transfer. In this work, we address this conversion challenge from two perspectives: architectural alignment and representational alignment. We identify Test-Time Training (TTT) as a linear-complexity architecture whose two-layer dynamic formulation is structurally aligned with Softmax attention, enabling direct inheritance of pretrained attention weights. To further align representational properties, including key shift-invariance and locality, we introduce key instance normalization and a lightweight locality enhancement module. We validate our approach by linearizing Stable Diffusion 3.5 and introduce SD3.5-T5 (Transformer To Test Time Training). With only 1 hour of fine-tuning on 4×H20 GPUs, SD3.5-T5 achieves comparable text-to-image quality to the fine-tuned Softmax model, while accelerating inference by 1.32× and 1.47× at 1K and 2K resolutions. Code is available at https://github.com/LeapLabTHU/Transformer-to-TTT.
This paper introduces Exact Linear Attention (ELA), a mechanism that achieves linear computational complexity for Transformer attention by exploiting the exact decomposition property of kernel functions, thereby eliminating approximation error. We identify and address two key limitations of prior linear attention -- gradient explosion and token attention dilution -- by imposing kernel constraints that ensure non-negativity, discriminability, and geometric interpretability. Several kernel functions are proposed, including the Hadamard Exp Kernel, Summation Squared Euclidean Distance Kernel, and Subtraction Squared Euclidean Distance Kernel, each tailored for specific attention behaviors. Beyond the core attention formulation, the paper presents three engineering innovations: (1) a Hyper-Link structure that replaces traditional residual connections to mitigate gradient degradation; (2) a Memory Lobe module based on bidirectional linear attention, which captures "transformation flow" across layers to implement qualitative memory and an implicit reinforcement learning paradigm; and (3) a routing-score-based bias mechanism for Mixture-of-Experts (MoE) to improve interpretability and semantic alignment. Experimental results demonstrate that ELA achieves up to 6x faster decoding speed and 75% reduction in KV cache memory usage compared to full attention, while maintaining comparable or superior training performance. The proposed memory module accelerates convergence and enhances generalization. Furthermore, we extend the linear attention principle to vision models, yielding YOLO-LAT, which attains up to 4.3x GPU inference speedup and 7.9x parameter reduction with competitive detection accuracy. These results underline the broad applicability of exact linear attention for scaling Transformer models to ultra-long sequences and efficient visual tasks.
Vision Transformers (ViTs) based vision foundation models (VFMs) have achieved remarkable performance across diverse vision tasks, but suffer from quadratic complexity that limits scalability to long sequences. Existing linear attention approaches for ViTs are typically trained from scratch, requiring substantial computational resources, while linearization-based methods developed for large language model decoders do not transfer well to ViTs. To address these challenges, we propose ViT-AdaLA, a novel framework for effectively adapting and transferring prior knowledge from VFMs to linear attention ViTs. ViT-AdaLA consists of three stages: attention alignment, feature alignment, and supervised fine-tuning. In the attention alignment stage, we align vanilla linear attention with the original softmax-based attention in each block to approximate the behavior of softmax attention. However, residual approximation errors inevitably accumulate across layers. We mitigate this by fine-tuning the linearized ViT to align its final-layer features with a frozen softmax VFM teacher. Finally, the adapted prior knowledge is transferred to downstream tasks through supervised fine-tuning. Extensive experiments on classification and segmentation tasks demonstrate the effectiveness and generality of ViT-AdaLA over various state-of-the-art linear attention counterpart.
Yifan Li, Seunghyun Yoon, Viet Dac Lai +4
Michigan State University, East Lansing, MI, USA. · Adobe Research.