Mixture-of-Top-k Attention, or MiTA, starts from a simple but easy-to-miss claim: in a Transformer, attention and the feed-forward network are usually treated as separate and complementary modules, yet standard attention can also be written as an FFN once its formula is expanded. In that form, the weights are not fixed training-time parameters. They are generated by the current input, and the hidden width grows with sequence length.
That shift in view changes how the cost of long-context modeling is described. The paper argues that the overhead does not come only from processing more tokens. It also comes from running a fast-weight network that keeps getting wider as the context grows. From there, a team at Beijing University of Posts and Telecommunications develops a broader framework for efficient attention and proposes MiTA, a method built around reusable Top-k experts. The work has been accepted to NeurIPS 2026, and both the paper and the code are available publicly.
Attention as an FFN, with different weights and a width tied to sequence length
The paper places a standard two-layer multilayer perceptron next to standard attention and makes the correspondence explicit. After expanding Softmax and absorbing scaling and query-dependent normalization into the weights, the mapping becomes clear: the query acts as the input, transformed keys and values play the role of the two weight matrices, the exponential function acts like the activation, the bias is zero, and the hidden width changes from D in a conventional FFN to N, the sequence length.
Under that reading, a regular FFN uses slow weights learned during training, while attention uses keys and values produced from the current context as fast weights. Width is the other major difference. Each key-value pair can be seen as one hidden unit. As the context gets longer, the fast-weight MLP gets wider. Each query visits N hidden units, and processing N queries creates the full quadratic burden.
The paper then turns efficient attention into a single design question: how can a fast-weight network that keeps expanding remain scalable?
Two broad routes: compression and routing
The authors group efficient attention methods into two main lines.
Compression tries to turn a wide network into a compact module. Linear Attention compresses key-value information into shared linear layers, Linformer maps it into a narrower MLP, and Test-Time Training learns compact fast-weight modules. This path gives a global summary, but it is less direct when raw key-value details need to be read out.
Routing goes the other way. It keeps the original content and only accesses the needed subset. A method can organize key-value subsets as fast-weight experts and let each query choose among them. That preserves the fine detail of selected content, but it does not summarize the full context as naturally.
Figure 1 in the paper illustrates that split: as context grows, the fast-weight MLP gets wider; routing selects part of it, while compression builds a compact representation.
From that point, the authors define a five-dimensional taxonomy consisting of scaling strategy, expert count, expert type, expert construction, and routing topology. The stated goal is not only to organize prior work, but also to force concrete design questions: what information should build an expert, how many experts are needed, and which queries can reuse the same expert.
MiTA’s central question: does every query need its own Top-k set?
To motivate MiTA, the paper contrasts block-based mixture-of-experts attention with per-query Top-k attention.
Methods such as MoBA split a sequence into continuous, non-overlapping blocks, build routing vectors for those blocks, and let queries choose among them. Routing can vary from query to query, but expert boundaries are fixed in advance. In image tasks, relevant content from the same object may not fall inside the same rectangular block, so a query may end up reading an entire block to get a small amount of useful information.
Top-k attention removes that fixed shape by letting each query choose its own most relevant k key-value pairs. The tradeoff is that N queries produce N private sets. Even if several queries care about similar content, each still has to perform its own selection and organization.
MiTA asks whether the flexibility of Top-k selection can be preserved while making the selected content reusable across multiple queries.
The method reverses the usual order. Instead of defining expert content first, it extracts a small number of representative queries and lets those queries select content from the full context to form experts. The original queries only need to decide which expert to visit. They no longer have to build a global Top-k set one by one.
Landmark queries build both routed experts and a shared global expert
MiTA begins by using local average pooling to extract m landmark queries. Each landmark probes the full sequence and serves two roles at once: it summarizes global information and retrieves raw detail from the original context.
For routed experts, the i-th landmark query matches against all keys and selects the k most relevant positions. Those key-value pairs form the i-th routed expert. The paper stresses that these experts can span different spatial locations, change shape with the input, and overlap with one another.
After that, each original query chooses the expert associated with the landmark that matches it best. Multiple queries can therefore share the same candidate key-value set, while still computing their own attention weights independently.
The same landmarks are also used to build a shared expert. Landmark queries read the full context through cross-attention and produce landmark values. Using the landmark queries as keys and the landmark values as values gives a globally shared key-value set.
MiTA therefore contains m routed experts and 1 shared expert. The routed experts supply raw detail; the shared expert supplies a global summary.
In the s=1 setting described in the paper, each query selects one routed expert. The model then concatenates that expert’s key-value pairs with the shared expert’s pairs and applies one unified Softmax normalization. This lets Top-k retrieval and global summarization share landmark scores, and the score from a query to a landmark can be reused for routing. Compression and routing are not only complementary in function; part of the computation can be shared as well.
On ImageNet-1K, with m=k=25, each query visits 25 compressed key-value pairs and 25 raw key-value pairs, 50 in total. The scan over the full sequence and the construction of Top-k sets are concentrated in the small landmark set instead.
The paper’s Algorithm 1 first constructs experts, then groups queries by routing result, and finally merges outputs through Online Softmax. The result is lower query-expert attention cost, with additional acceleration from FlashAttention and grouped computation by expert.
ImageNet-1K: baseline MiTA stays close to full attention
On ImageNet-1K image classification, the authors first compare plain attention mechanisms under the same DeiT training setup, without adding extra components.
On DeiT-S, baseline MiTA reaches 79.8% Top-1 accuracy, compared with 79.9% for full attention, a gap of 0.1 percentage point. Parameter count stays unchanged, while FLOPs drop from 4.6G to 4.4G. Under the same setup, Agent Attention posts 78.5%.
When extra components commonly used in efficient visual attention are added, MiTA improves further. With Depth-Wise Convolution, it reaches 80.6%. With gating added on top, it reaches 81.2%. The paper explicitly notes that these are standard add-on gains and reports them as such.
With the ViT-5 architecture design plus DWC, MiTA-ViT-5-S reaches 81.7%, versus 82.2% for ViT-5-S, while FLOPs fall from 4.7G to 4.5G.
ADE20K: fewer accessed tokens on longer visual sequences
For ADE20K semantic segmentation, the paper replaces the attention modules in a pretrained ViT with MiTA and then trains the segmentation model. With m=k=49, each query accesses only 98 key-value pairs, while full attention accesses 1,024 or 1,600 pairs under the two input resolutions studied.
Using ViT-S+CBT as one example, FLOPs fall from 33G to 25G, a reduction of about 24%, and mean Intersection over Union changes from 45.8% to 44.5%. In the ViT-T setting, FLOPs are reduced by about 42%, while mIoU drops by 2.6 percentage points.
LRA: average accuracy edges up while training time falls sharply
Across the five tasks in Long Range Arena, MiTA records 58.91% average accuracy, versus 58.19% for full attention. Total training time drops from 10.6 hours to 2.4 hours, a reduction of about 77%.
These tests were run on a single 24GB RTX 4090, and the full-attention baseline used FlashAttention.
The compression branch also lowers the gather cost of raw key-value pairs. In an ablation where each query accesses the same number of entries, adding the compression branch shortens training time from 5.5 hours to 2.4 hours and improves average accuracy from 58.81% to 58.91%, compared with the routing-only version.
For inference throughput, the paper evaluates a three-layer Transformer with embedding dimension 128 on an RTX 4090, setting m*k≈N/4 and tuning batch size for each configuration. As sequence length grows, the speedup widens. At lengths of 64K, 256K, and 1024K, MiTA is about 19x, 56x, and 159x faster than full attention, respectively.
Long-document classification: gains at 16K, 32K, and 64K
The arXiv long-document classification task requires a model to identify the category of a long paper. The paper presents it as a useful test because it demands both overall understanding and retrieval of fine-grained clues scattered through the document.
At 16K, 32K, and 64K lengths, MiTA reaches 88.0%, 92.9%, and 98.7% accuracy, respectively, leading all other methods listed in the table. Compared with standard attention, that is an increase of 18.2, 3.2, and 1.7 percentage points.
Two additional observations from the paper
The first is an emergent token-pruning effect. Even without an explicit pruning module, the selected key-value positions in routed experts gradually concentrate on class-related regions as the network gets deeper, while background positions are chosen less often. Figure 5 shows the input image in the leftmost column and the selected positions across layers in the remaining columns, aggregated across attention heads. Later layers focus more tightly on class-relevant regions. Background tokens not selected into routed experts can still contribute through the compression branch.
The second is weight reuse and adjustable inference budgets. After freezing model parameters and changing only the attention mechanism at inference time, the paper reports good compatibility between MiTA and full attention.
When a ViT pretrained with full attention is switched to MiTA and then fine-tuned on ImageNet-1K, ViT-S and ViT-L reach 80.9% and 85.5%, only 0.3 and 0.4 percentage point behind full attention. The paper also says that, with model parameters unchanged, adjusting the number of experts m and expert width k at inference can produce even higher accuracy under some settings. The authors describe this as algorithmic generalization, opening the door to lower training cost and flexible inference budgets.
From asking what attention is to asking how it can be designed
The paper closes by returning to its opening premise. Writing attention as an FFN does more than line up two formulas. It creates a new set of design dimensions for efficient attention: compression determines how global information is retained, expert construction determines which raw content should be organized together, and routing determines which queries can reuse that content. In the authors’ view, ideas once studied inside FFNs can feed back into attention design and may eventually weaken the hard boundary between the two.
That makes MiTA not just another efficient attention variant, but also a proposal about how to think about architecture design at a lower level.
References:
- Paper: https://arxiv.org/pdf/2602.01219
- Code: https://github.com/QishuaiWen/MiTA
This article is based on a post from the WeChat public account "新智元," written by 新智元.