🤖 AI

SEMA: a Scalable and Efficient Mamba like Attention via Token Localization and Averaging

This paper introduces SEMA, a scalable and efficient attention mechanism that combines token localization to prevent weight dispersion with arithmetic averaging to capture global context, thereby outperforming existing linear attention and vision Mamba models in image classification tasks.

Original authors: Nhat Thanh Tran, Fanghui Xue, Shuai Zhang, Jiancheng Lyu, Yunling Zheng, Yingyong Qi, Jack Xin

Published 2026-07-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nhat Thanh Tran, Fanghui Xue, Shuai Zhang, Jiancheng Lyu, Yunling Zheng, Yingyong Qi, Jack Xin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to recognize a cat in a photo. To do this, the robot uses a special tool called "attention," which acts like a spotlight. This spotlight scans the entire image, deciding which pixels are important (like the cat's ears) and which are just background noise (like a blurry tree). The problem is that as the photo gets bigger, the spotlight has to check every single pixel against every other pixel. If the photo is small, this is fast. But if the photo is huge, the robot gets overwhelmed, spending so much time checking connections that it runs out of battery before it can even say "cat."

Scientists have tried to fix this by making the spotlight "linear," meaning it checks pixels in a simpler, faster line. But there's a catch: these fast spotlights often lose their focus. As the image gets bigger, they start treating every single pixel as equally important, like a flashlight that has been turned down so low it just glows evenly over the whole room. The robot stops seeing the cat and starts seeing a gray blur. This paper, written by researchers from the University of California, Irvine, and Qualcomm AI Research, tackles this exact problem. They want to build a spotlight that is both fast enough to handle giant images and sharp enough to still see the details.

The researchers, led by Nhat Thanh Tran and colleagues, discovered a mathematical truth that explains why these fast spotlights fail: as the number of pixels (or "tokens") grows toward infinity, the attention mechanism naturally spreads out and becomes useless. They call this the "dispersion" phenomenon. It's like trying to hear a whisper in a stadium; if you try to listen to everyone at once, you end up hearing nothing but noise. The paper proves that no matter how you tweak the math of these fast attention tools, they will eventually lose their focus if the image gets too big.

To solve this, the team invented a new method called SEMA (Scalable and Efficient Mamba like Attention). Instead of fighting the math, they decided to work with it. They realized that while the spotlight should focus on specific details, it also needs a way to understand the "big picture" without getting lost. So, they designed a two-part system. First, they use Token Localization, which is like giving the spotlight a small window to look through. It only focuses on a tiny neighborhood of pixels at a time, ensuring it never gets overwhelmed or distracted. This keeps the focus sharp.

But looking through a small window has a downside: the robot might miss the whole cat if it's only looking at one paw. To fix this, SEMA adds a second step: Averaging. This is like taking a quick, blurry snapshot of the entire room and adding it to the detailed view. The researchers proved mathematically that when the image gets huge, the best way to capture the "global" feel of the image is simply to average everything out. By combining the sharp, local window view with this simple, global average, SEMA gets the best of both worlds.

The paper shows that this approach works incredibly well. When they tested SEMA on standard image datasets like ImageNet-1K, it outperformed other popular models, including the recent "Mamba" models, especially when the images were very large. For example, when the image size was increased to 1024x1024 pixels, other models struggled significantly, dropping in accuracy, while SEMA stayed strong and even improved. The researchers also found that SEMA was faster than its competitors, taking less time to process large images.

One of the most interesting parts of the paper is how they handled the "dispersion" problem. Instead of trying to force the attention to stay sharp in a way that breaks the math, they accepted that for huge images, the attention should spread out. They used this spreading nature to their advantage by using a simple average to represent the global information. It's a bit like realizing that if you have a million friends, you can't remember every detail about every single one, but you can remember the general vibe of the group. SEMA remembers the details of the immediate neighborhood and the general vibe of the whole group, allowing it to recognize the cat perfectly, no matter how big the photo is.

The authors tested their idea on various tasks, from recognizing objects to finding specific parts of an image (segmentation). In every case, SEMA showed it could scale up to larger images without losing its mind. They even showed that as the images got bigger, the "averaging" part of their system became more important, proving their theory that this simple step is crucial for handling massive amounts of data. While the paper focuses on computer vision, the team suggests this idea could help with other problems involving long sequences of data, like analyzing huge medical scans.

In short, SEMA is a clever trick that admits when a spotlight gets too wide to be useful and switches to a different strategy: look closely at the details nearby, and take a quick, simple average of the rest. It's a scalable, efficient, and mathematically sound way to help computers see the world clearly, even when the world gets very big.

Technical Summary: SEMA – A Scalable and Efficient Mamba-like Attention via Token Localization and Averaging

1. Problem Statement

The paper addresses two primary challenges in applying attention mechanisms to computer vision tasks:

Computational Complexity: Vanilla full attention (softmax-based) exhibits quadratic computational complexity ( $O(n^2)$ ) relative to the input size, which is prohibitive for high-resolution images.
The Dispersion Phenomenon: Linear attention variants, designed to achieve linear complexity, often suffer from an inability to focus on relevant information. As the number of tokens ( $n$ ) increases, these mechanisms tend to "disperse," assigning nearly equal weights to all keys, thereby losing the discriminative power inherent in softmax attention.

While recent works like Mamba (state space models) and MILA (Mamba-inspired linear attention) have attempted to resolve these issues, the authors observe that linear attention mechanisms still lack expressive focusing power, and existing solutions for dispersion (e.g., adaptive temperature or differential attention) are often heuristic or still theoretically prone to dispersion in the limit.

2. Methodology

2.1 Theoretical Foundation: Generalized Attention and Dispersion

The authors first establish a Generalized Attention framework that encompasses both vanilla softmax attention and linear attention. They define a $\Phi$ -normalized attention mechanism where the attention coefficients are determined by a continuous normalization function $\phi$ .

Key Theoretical Result (Dispersion Theorem):
The paper proves that for any continuous normalization function $\phi$ , as the number of tokens $n$ tends to infinity, the attention coefficients inevitably disperse. That is, the attention weight assigned to any specific key approaches zero ( $\Theta(1/n)$ ), effectively assigning equal weight to all tokens.

This holds for vanilla softmax ( $\phi(x) = \exp(x)$ ) and linear attention ( $\phi(x) = x$ ).
The proof relies on the compactness of the input feature space and the continuity of the functions involved.
The authors note that while differential attention attempts to mitigate noise, it theoretically remains subject to dispersion ( $\Theta(1/n)$ ).

2.2 Proposed Solution: SEMA

Motivated by the proven dispersion property, the authors design SEMA (Scalable and Efficient Mamba-like Attention). The core philosophy is to embrace the asymptotic behavior of global attention (dispersion) while preserving local focusing capabilities.

SEMA combines two components:

Token Localization (Window Attention): To maintain the "focusing" property of vanilla attention on high-frequency features, SEMA employs a localized window attention mechanism. This restricts each query to interact only with a fixed number of keys within a pre-defined window, preventing dispersion within that local context.
Theoretical Averaging (Global Approximation): To capture the global aspect of attention without incurring quadratic complexity or suffering from the dispersion of full linear attention, SEMA incorporates an arithmetic averaging term.
- Based on the asymptotic analysis, the authors argue that in the limit of large $n$ , the global component of full attention can be approximated by a simple average of all value vectors ( $\frac{1}{n}\sum v_j$ ).
- This acts as a "homogeneous mixing" or low-pass filter, aggregating global information efficiently.

Formulation:
The SEMA attention for a query $q$ , keys $K$ , and values $V$ with window size $w$ is defined as:
$\text{SEMA}(Q, K, V) = A_{\Phi}(Q, K, V, w) + \frac{1}{n} \sum_{j=1}^{n} v_j$
Where $A_{\Phi}$ is the windowed attention (using softmax or linear variants) and the second term is the broadcasted global average.

2.3 Architecture

SEMA is integrated into a Mamba-like macro-architecture (similar to MILA). The authors replace the linear attention blocks in MILA with SEMA blocks. Unlike MILA, SEMA includes a linear projection for the values $V$ , which the authors found necessary for performance, slightly increasing the parameter count but maintaining efficiency.

3. Key Contributions

Theoretical Proof of Dispersion: The paper provides a rigorous mathematical proof that generalized attention (including softmax and linear variants) disperses as the token count approaches infinity, regardless of the specific continuous normalization function used.
SEMA Design: A novel attention mechanism that utilizes token localization to avoid local dispersion and arithmetic averaging to theoretically approximate the global aspect of attention, offering a scalable alternative to linear attention.
Empirical Validation: Demonstrating that SEMA outperforms recent vision Mamba models (e.g., VMamba) and linear attention baselines (MILA) on standard benchmarks, particularly as image resolutions increase.

4. Experimental Results

4.1 ImageNet-1K Classification

Standard Resolution (224²): SEMA-T achieves 83.7% Top-1 accuracy, outperforming MILA-T (83.3%) and DefMamba (83.5%) with similar parameter counts (~26M).
Scalability (Larger Resolutions): The most significant gains appear at higher resolutions.
- At 1024², VMamba-T drops significantly (74.7% without fine-tuning, 83.4% with), whereas SEMA-T maintains robustness (75.4% without, 84.2% with fine-tuning).
- The gap between SEMA and VMamba widens to 7.3% at 1024² without fine-tuning, indicating SEMA's superior resilience to token count increases.

4.2 Downstream Tasks

COCO (Instance Segmentation): Using Mask R-CNN, SEMA-T achieves state-of-the-art or top-2 performance across all metrics (AP, AP50, AP75) for both 1x and 3x training schedules, surpassing MILA-T and VMamba-T.
ADE20K (Semantic Segmentation): SEMA-T achieves 48.2 mIOU (single-scale), outperforming VMamba-T (47.9), and remains competitive in multi-scale settings.

4.3 Efficiency

Inference Time: SEMA demonstrates superior scalability in runtime. At 1024² resolution, SEMA is approximately 33% faster than MILA. At 2048², the speedup is even more pronounced, confirming the linear complexity advantage of the averaging component.

4.4 Ablation Studies

Value Projection: Removing the value projection layer reduces accuracy by ~0.8%, highlighting its importance.
Averaging Component: At 224², the averaging term contributes minimally (~0.2%). However, as resolution increases to 672², its contribution grows significantly (+2.2% without fine-tuning), validating the theoretical claim that averaging becomes crucial as token counts rise and dispersion effects intensify.

5. Significance and Claims

The paper claims that SEMA offers a theoretically grounded solution to the dispersion problem in linear attention. By proving that dispersion is an inherent property of global attention in the limit, the authors justify the use of arithmetic averaging as a valid global approximation.

The significance of SEMA lies in its ability to:

Scale Effectively: It maintains performance and efficiency as image resolutions (and token counts) increase, a known weakness of current Mamba-based vision models.
Balance Local and Global: It successfully combines the focusing power of window attention with the global context of averaging, avoiding the quadratic cost of full attention and the "blurry" focus of pure linear attention.
Provide a Robust Alternative: It serves as a scalable and efficient alternative for computer vision tasks, particularly in scenarios involving large inputs or long sequences (e.g., medical imaging), where token counts can reach orders of magnitude higher than standard image classification.

The authors conclude that while the averaging component is marginal at standard resolutions, it provides a critical "safety net" and performance boost in the large-token regime, making SEMA a robust choice for scalable vision architectures.