NeurIPS'26: Is Attention Also FFN? Rethinking Efficient Attention From the Fast Weight Perspective

For a long time, in the Transformer architecture, Attention and Feed-Forward Network (FFN) have been regarded as two modules with distinct and complementary functions:

The former allows tokens to interact with each other, while the latter performs non-linear transformation on each token.

But when we expand the formula of the attention mechanism, we can find a connection that is easily overlooked: standard attention can also be formulated as an FFN — its weights are generated by the current input, and the hidden layer width grows with the sequence length.

Thus, the overhead of long contexts can be explained from another perspective: we are not only processing more tokens, but also running a network that keeps getting wider. Then, can seemingly different technical routes such as linear attention, Mixture-of-Experts attention, and Top-k attention be re-understood from the perspective of "how to organize this network"?

Focusing on this problem, Mixture-of-Top-k Attention (MiTA) proposed by the team from Beijing University of Posts and Telecommunications develops the existing Fast Weights perspective into a five-dimensional classification framework, and designs deformable and reusable Top-k experts that retain global information and read original details on demand. This work has been accepted by NeurIPS 2026, and the code has been open-sourced.

Paper Address: https://arxiv.org/pdf/2602.01219

Code Repository: https://github.com/QishuaiWen/MiTA

Attention is also an FFN, the key lies in where the weights come from

Let the query

, keys and values are formed by columns

.

Put the two-layer Multi-Layer Perceptron (MLP) commonly used in FFNs and standard attention together, and the hidden layer width of the former is denoted as D:

Expand Softmax, and absorb scaling and query-dependent normalization into the weights:

The corresponding relationship emerges immediately:

The query corresponds to the input, the transformed keys and values correspond to the two layers of weights, the exponential function corresponds to the activation function, the bias is zero, and the hidden layer width changes from D to the sequence length N.

The difference lies in the source of weights: ordinary FFNs use Slow Weights obtained through training, while attention takes K, V generated by the current input as fast weights that change with the context.

The more important difference lies in the width: each key-value pair corresponds to a hidden unit, and the longer the context, the wider this fast-weight MLP will be. Each query accesses N hidden units, and processing N queries will generate

total overhead.

The design of efficient attention is thus transformed into a unified problem: how to keep the expanding fast-weight network scalable?

Understand compression and routing from an increasingly wide MLP

Faced with this increasingly wide network, the main ideas of efficient attention can be summarized into two routes.

Compression: Turn the wide network into a compact module (Scaling by Compression). Linear Attention compresses key-value information into a shared linear layer, Linformer maps it to a narrower MLP, and Test-Time Training (TTT) learns a compact fast-weight module. This route provides global generalization, but it is difficult to directly read the details of the original key-value pairs.

Routing: Retain original content and only access the required part (Scaling by Routing). Organize key-value subsets into Fast-Weight Experts, then let queries select as needed. These experts are composed of key-value pairs in the current context: they retain the details of the selected content, but lack the generalization of the complete context.

Figure 1: The growth of context makes the fast-weight MLP wider. Routing selects part of it, while compression constructs a more compact representation.

These two routes are not mutually exclusive. Further considering Scaling Strategy, Expert Count, Expert Type, Expert Construction and Routing Topology, we can get a five-dimensional classification framework.

Table 1: The five-dimensional classification framework for efficient attention from the fast-weight perspective.

This framework not only sorts out existing methods, but also raises specific questions for new designs: What information are experts constructed from? How many are needed? Which queries can reuse the same expert?

The entry point of MiTA: Does the Top-k set have to be selected by each query separately?

Taking block-wise Mixture-of-Experts Attention such as MoBA as an example, the common process is: first cut the sequence into consecutive non-overlapping blocks, then generate a routing vector for each block, and let queries select as needed.

Routing can vary with queries, but the expert boundaries are pre-fixed. In images, the relevant information of the same target may not fall within the same rectangular block; in order to read a small amount of relevant content, the query may need to access the entire block.

Top-k attention goes in the other direction: each query individually selects the k most relevant key-value pairs, no longer constrained by the fixed block shape. But accordingly, N queries correspond to N sets of "private" sets. Even if different queries pay attention to similar content, they still need to complete selection and organization respectively.

Can we retain the flexibility of Top-k selection and make the selection results reused by multiple queries?

MiTA reverses the idea: instead of stipulating the content of experts in advance, first extract a small number of representative queries, then let them select content from the complete context to form experts with adjustable quantities. The original query only needs to choose which expert to access, instead of constructing a set of global Top-k sets respectively.

Figure 2: Full attention shares all key-value pairs, block-wise MoE shares fixed blocks, query-wise Top-k constructs private sets, and MiTA constructs reusable deformable experts.

Construct global generalization and detail experts at the same time with a small number of landmark queries

MiTA extracts m Landmark Queries through local Average Pooling. Each landmark probes the complete sequence, and completes global information aggregation and original detail retrieval at the same time.

Figure 3: Each query accesses both the compressed global key-value pairs and the original key-value pairs selected by routing.

First construct reusable Top-k experts. The i-th landmark query

matches all keys, and selects the k positions with the highest correlation:

These key-value pairs form the i-th Routed Expert, which can span different spatial positions, change shape with input content, and different experts are allowed to overlap.

Then, the query selects the expert corresponding to the landmark that best matches itself:

Different queries share the candidate key-value set, but still calculate their respective attention weights independently.

Then construct a globally shared expert with the same set of landmarks. The landmark queries read the complete context through Cross-Attention to get the corresponding landmark values:

Taking landmark queries as keys and landmark values as values, we get the globally shared key-value pairs

.

Thus MiTA contains m routed experts and 1 Shared Expert: the former provides original details, and the latter provides global generalization.

Finally, let generalization and details interact in the same attention pass. Under the configuration that each query selects one routed expert, that is, s=1: