MiTA Attention: Efficient Fast-Weight Scaling via a Mixture of Top-$k$ Activations

TL;DR

This work addresses the quadratic bottleneck of full attention in Transformers for long sequences by reframing efficient attention as fast-weight scaling. It introduces MiTA, a hybrid routing-compression mechanism that uses landmark queries to compress the key-value cache and to form deformable, top-$k$ activated experts, yielding an attention operator with complexity $O(N(m+ks))$. A five-dimensional taxonomy for fast-weight scaling is proposed, and MiTA is shown to fit within this framework by combining a fixed set of base experts ($m$) with query-dependent routing ($s$) to achieve flexible, scalable attention. Preliminary experiments on Vision tasks (ImageNet-1K, ADE20K) and long-sequence benchmarks (Long Range Arena) demonstrate competitive accuracy and improved throughput, validating MiTA as a promising direction for efficient attention with potential for broader applications.

Abstract

The attention operator in Transformers can be viewed as a two-layer fast-weight MLP, whose weights are dynamically instantiated from input tokens and whose width equals sequence length $N$. As the context extends, the expressive capacity of such an $N$-width MLP increases, but scaling its fast weights becomes prohibitively expensive for extremely long sequences. Recently, this fast-weight scaling perspective has motivated the Mixture-of-Experts (MoE) attention, which partitions the sequence into fast-weight experts and sparsely routes the tokens to them. In this paper, we elevate this perspective to a unifying framework for a wide range of efficient attention methods by interpreting them as scaling fast weights through routing and/or compression. Then we propose a compress-and-route strategy, which compresses the $N$-width MLP into a narrower one using a small set of landmark queries and constructs deformable experts by gathering top-$k$ activated key-value pairs for each landmark query. We call this strategy a Mixture of Top-$k$ Activations (MiTA), and refer to the resulting efficient mechanism as MiTA attention. Preliminary experiments on vision tasks demonstrate the promise of our MiTA attention and motivate further investigation on its optimization and broader applications in more challenging settings.

Figures (12)

• Figure 1: Fast-weight scaling and its two scaling strategies. As the context extends, the width of the two-layer fast-weight MLP induced by full attention increases accordingly. We categorize efficient fast-weight scaling approaches into two strategies: a) scaling by routing and b) scaling by compression, and illustrate each of them with a representative method: Mixture of Block Attention (MoBA) Lu:2025-MoBA and Test-time Training (TTT) sun2024learning. A pioneering method that combines both strategies is Native Sparse Attention (NSA) Yuan:ACL25-NSA.
• Figure 2: Illustration for our MiTA attention. In full attention, each query attends to all key-value pairs. In our MiTA attention, it attends to the concatenation of a small number of the compressed key-value pairs and a routed subset of the full key-value pairs.
• Figure 3: Visualization of experts' gathered key-value pairs, and routed queries. The red box marks the local window from which the landmark query is obtained via average pooling. The attention heatmap (averaged over heads) indicates key-value pairs within each expert (top row) and the queries routed to it (bottom row). Notably, neither the expert’s key–value pairs nor the routed queries are confined to the local window.
• Figure 4: The token pruning effect of our MiTA. Each row visualizes, for each layer, the positions of key-value pairs (aggregated over heads) selected as experts; the leftmost image shows the original input. In later layers, most tokens are effectively "pruned" (i.e., not selected as experts), and attention concentrates on class-relevant regions. The examples are sampled from the ImageNet-1K training set.
• Figure 5: The layer-wise positional overlap (between the key–value pairs gathered by an expert and the queries routed to it) is quantified by mIoU, averaged over experts and heads.