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

This paper introduces Infinite Self-Attention (InfSA) and its linear-time variant, Linear-InfSA, a spectral reformulation of self-attention as a diffusion process on token graphs that achieves state-of-the-art ImageNet accuracy and enables efficient, memory-free inference at ultra-high resolutions (up to 9216×9216) by replacing the quadratic softmax cost with a Neumann series approximation.

The Gist

Imagine you are trying to understand a massive, crowded room full of people (the "tokens" in an AI model). Your goal is to figure out who is important and what the main topic of conversation is.

The Problem: The "Quadratic" Bottleneck

In standard AI models (called Transformers), every person in the room has to look at every other person to decide who to listen to.

• The Analogy: If there are 10 people, everyone makes 100 connections. If there are 1,000 people, that's 1,000,000 connections. If you have a high-resolution photo (like a 4K or 8K image), the "room" has hundreds of thousands of people.
• The Result: The computer gets overwhelmed. It runs out of memory, takes forever to think, and uses a massive amount of electricity. This is the "quadratic cost" mentioned in the paper.

The Solution: "Infinite Self-Attention" (InfSA)

The authors, Giorgio Roffo and Luke Palmer, propose a new way to listen to the room. Instead of everyone shouting at everyone else at once, they treat the room like a game of "telephone" or a "rumor mill" that spreads information through the crowd.

They call this Infinite Self-Attention. Here is how it works in simple terms:

1. The "Rumor Mill" (Graph Diffusion)

Imagine you drop a pebble in a pond. The ripples spread out.

• Standard AI: Tries to calculate the exact shape of every single ripple instantly. Too much math!
• InfSA: Lets the ripples spread naturally. It asks, "If I start a rumor at Person A, how many people will hear it after 1 hop? After 2 hops? After 10 hops?"
• The Magic: It doesn't just look at who is standing next to you (1 hop); it looks at who is connected to your friends, and their friends, and so on. This is called multi-hop interaction. It finds the "influencers" of the room—people who are central to the conversation, even if they aren't talking to you directly.

2. The "Absorbing State" (Stopping the Rumor)

In math, there's a concept called an "absorbing Markov chain."

• The Analogy: Imagine the rumor spreads, but there's a tiny chance at every step that the rumor "dies out" (someone forgets it).
• Why it helps: This prevents the AI from getting confused by noise. It ensures that the most important people (the ones the rumor keeps reaching) stand out clearly, while background noise fades away. This makes the AI's "attention map" much sharper and more focused on the actual object in a photo, rather than the background.

3. The "Super-Speed" Version (Linear-InfSA)

Calculating all those ripples for a huge crowd is still hard. So, the authors created a shortcut called Linear-InfSA.

• The Analogy: Instead of tracking every single person, the AI finds the "Main Vibe" of the room.
• How it works: It uses a mathematical trick (finding the "principal eigenvector") to instantly guess who the most important people are without doing the heavy lifting of checking every pair.
• The Result: It's like having a superpower where you can instantly know the most important person in a stadium of 100,000 people, without needing to talk to everyone.

Why This Matters (The Real-World Wins)

1. It's Super Fast and Cheap:

   • The paper tested this on a massive 8K resolution image (like a huge, detailed painting). Standard AI models crashed (ran out of memory).
   • Linear-InfSA handled it easily. It was 13 times faster and used 13 times less energy than the standard model. It's like switching from a gas-guzzling truck to a sleek electric bike.

2. It Sees Better:

   • Standard AI often gets distracted by the background (like focusing on the grass when looking at a dog).
   • InfSA focuses sharply on the "dog." In tests, it was much better at identifying exactly where an object is in a picture.

3. It's Smarter:

   • Even with fewer parameters (less "brain" size), the new model scored higher on standard tests (ImageNet) than the older, bigger models. It's a more efficient way of thinking.

The Bottom Line

The authors took a complex idea from graph theory (how things connect in a network) and applied it to AI vision. They turned the AI's "attention" from a chaotic, expensive shouting match into a structured, efficient flow of information.

In short: They taught the AI to listen to the whole room by following the flow of conversation, rather than trying to hear every single voice at once. This makes it faster, cheaper, and much better at seeing what actually matters.

Technical Summary

Here is a detailed technical summary of the paper "Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention".

1. Problem Statement

Standard Transformer architectures rely on Softmax Self-Attention, which incurs a quadratic computational and memory cost ($O(N^2)$) relative to the sequence length $N$. This bottleneck severely limits scalability in high-resolution computer vision tasks (e.g., 4K–9K images) and long-context language modeling.

• Scalability: Existing efficient attention mechanisms (e.g., Linformer, Performer, FlashAttention) often approximate or sparsify the attention matrix without a principled model of multi-hop token interactions.
• Interpretability: Standard attention maps often distribute focus diffusely across background regions rather than concentrating on semantically relevant objects.
• Energy Efficiency: The quadratic cost dominates the energy budget of data centers, contributing to the projected doubling of data center energy consumption by 2030.

2. Methodology: Infinite Self-Attention (InfSA)

The authors propose Infinite Self-Attention (InfSA), a spectral reformulation of self-attention that treats token interactions as a diffusion process on a content-adaptive graph.

A. Theoretical Foundation: Graph Diffusion & Neumann Series

Instead of a single-step aggregation, InfSA models attention as the accumulation of multi-hop paths over an infinite horizon, truncated by a discount factor $\gamma$.

• Path Integration: The attention score between tokens is viewed as the sum of weights of all possible paths between them.
• Neumann Series: The infinite sum of discounted path powers $\sum_{t=1}^{\infty} \gamma^t A^t$ is closed-form solved using the Neumann kernel:
  $$\check{C} = (I - \gamma A)^{-1} - I$$
  where $A$ is the attention matrix and $\gamma$ is a decay factor.
• Absorbing Markov Chains: The authors establish a rigorous link between InfSA and Absorbing Markov Chains.
  • Tokens are treated as transient states; the diffusion process terminates at an "absorbing state."
  • The matrix $(I - \gamma A)^{-1}$ is the Fundamental Matrix of this chain.
  • Interpretation: The centrality score of a token corresponds to the expected number of random-walk visits to that token before absorption. This connects InfSA to classical centrality measures like Katz Centrality and PageRank.

B. Pure InfSA

• Mechanism: Replaces the standard Softmax operation with a Frobenius-normalized attention matrix ($\hat{A}$) followed by a ReLU activation to ensure non-negativity.
• Accumulation: The output is an accumulation of layer-wise contributions with geometric decay: $S_L = \sum_{l=1}^{L} \gamma^l Z^{(l)}$.
• Properties: Frobenius normalization ensures the operator is sub-stochastic ($\rho(\hat{A}) < 1$), guaranteeing convergence of the Neumann series and preventing the "oversmoothing" common in standard graph diffusion.

C. Linear-InfSA (The Scalable Variant)

To achieve linear complexity $O(N)$, the authors propose Linear-InfSA, which approximates the dominant eigenvector of the implicit attention operator without constructing the $N \times N$ matrix.

• Eigenvector Approximation: Based on the Perron-Frobenius theorem, the dominant eigenvector of a positive operator represents the limiting distribution of the diffusion process.
• Algorithm:
  1. Compute token energies via $\ell_2$ norms of queries ($e_i = \|Q_i\|_2$).
  2. Normalize to get soft weights $\alpha$ (replacing Softmax).
  3. Compute a "central query" $\bar{q}$ as a weighted average of keys.
  4. Score keys against $\bar{q}$ using a ReLU kernel and normalize again to get final attention weights $a$.
  5. Broadcast the resulting context vector to all tokens.
• Complexity: Reduces complexity from $O(N^2)$ to $O(N)$ with a fixed auxiliary state size of $O(d_h)$ (independent of sequence length $N$).
• Accuracy: Empirically recovers the principal eigenvector of the full quadratic operator with a cosine similarity of 0.985.

3. Key Contributions

1. Spectral Reformulation: Connects self-attention to graph diffusion, Katz centrality, and absorbing Markov chains, providing a mathematically grounded interpretation of token weighting.
2. Infinite Self-Attention (InfSA): Introduces a new attention mechanism that aggregates multi-hop dependencies via a truncated Neumann series, offering better structural interpretability than standard softmax.
3. Linear-InfSA: Proposes an $O(N)$ approximation that avoids the $N \times N$ matrix construction while maintaining the benefits of global context, making it compatible with standard Vision Transformer (ViT) blocks.
4. Scalability & Efficiency: Demonstrates the ability to process 9216×9216 resolution images (~332k tokens) on a single GPU, a task where standard and other efficient attention models fail due to Out-Of-Memory (OOM) errors.

4. Experimental Results

The models were evaluated on ImageNet-1K and ImageNet-V2 using a 4-layer ViT backbone (53.5M parameters).

• Classification Accuracy:
  • Linear-InfSA (4L): Achieved 84.7% Top-1 on ImageNet-1K, a +3.2 percentage point improvement over a standard 4-layer ViT baseline (81.5%) trained with the same recipe.
  • Generalization: On ImageNet-V2, InfSA variants outperformed all baselines (up to 79.8% vs. 76.8% for the best prior method), indicating superior robustness to distribution shifts.
• Attention Quality:
  • Localization: InfSA produces sharper, object-aligned attention maps.
  • Metrics: MoRF-AOC (Most Relevant First) reached 76.0% (vs. 42.6% for Softmax ViT), and Bounding-Box PR-AUC reached 76.1% (vs. 56.2%), confirming semantically grounded attention.
• Scalability & Efficiency (A100 40GB GPU):
  • Throughput: Linear-InfSA achieved 231 images/second at 1024² resolution, 13× faster than Standard ViT.
  • Energy: Consumed 0.87 J/image, a 13× improvement in energy efficiency over Standard ViT.
  • Extreme Resolution: It was the only model to successfully complete inference at 9216² resolution (332k tokens) without running out of memory.

5. Significance

This work bridges the gap between graph theory, stochastic processes, and deep learning architectures.

• Theoretical Impact: It provides a principled justification for linear attention mechanisms by grounding them in spectral graph theory and Markov chain limits, moving beyond heuristic approximations.
• Practical Impact: Linear-InfSA enables the deployment of Transformer models on ultra-high-resolution inputs (e.g., satellite imagery, medical scans, 8K video) that were previously computationally intractable.
• Sustainability: By reducing the computational complexity from quadratic to linear, the method significantly lowers the energy footprint of training and inference, addressing critical environmental concerns in AI development.

The code is available at huggingface.co/groffo/infinite-self-attention.