Research map / Architectures and efficiency / Attention and transformers

Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Attention head and block redesign

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MiTA, an efficient attention that combines landmark-query compression (a narrow shared expert) with query-aware top-k routing to a tunable number of reusable deformable experts, framed by a fast-weight scaling taxonomy.

Problem

Full self-attention's fast-weight MLP has hidden width equal to sequence length N, so per-query cost grows with N and becomes unscalable for long sequences; prior MoE-style and top-k sparse attention either use rigid blocks or per-query private top-k sets.

Approach

Use a small set of average-pooled landmark queries to (1) compress the full KV cache into a shared expert via cross-attention (landmark values) and (2) gather each landmark's top-k keys/values as reusable, query-aware deformable routed experts; queries attend to the concatenation of both, with s=1 routed expert per query.

Key results

• On ImageNet-1K with DeiT-S trained from scratch, MiTA-DeiT-S with DWC and gating reaches 81.2% vs 79.9% for DeiT-S, and MiTA-DeiT-S without extras reaches 79.8%, above Agent-DeiT-S at 78.5%.
• On ADE20K with CBT head, MiTA-ViT-T reduces FLOPs by 42% (7G vs 12G) at a cost of 2.6 mIoU (36.5 vs 39.1).
• On LRA, MiTA reaches 58.91 average accuracy vs 58.19 for standard attention while cutting total training time by 77% (2.4 vs 10.6 hours).

Surprising findings

• Finetuning from standard-attention ViT pretraining transfers better to MiTA than to Agent Attention, and standard-attention-trained models generalize to MiTA at inference with over 95% of original performance.
• Performance is more sensitive to k (per-expert width) than to m (number of experts); routing is more critical than compression in the ablation.

Limitations

• Evaluation limited to vision tasks and small LRA setting; LLM decoding phase not evaluated.
• Segmentation accuracy drops (1.2-2.8 mIoU) because the backbone is not natively pretrained with MiTA; ImageNet top-1 is below some full-attention and ViT-5 baselines in the SOTA comparison.
• Gather operation is irregular memory access and may be memory-bound; s=1 used in most experiments, not fully exploiting expressiveness.

Experimental setup

Datasets: ImageNet-1K ImageNet-21K ADE20K Long Range Arena CIFAR-100

Benchmarks: ImageNet-1K ADE20K Long Range Arena

Metrics: top-1 accuracy mIoU FLOPs training throughput wall-clock time

Baselines: Agent Attention BRA MoBA Linear Attention MHLA PolaFormer FLatten Linformer Reformer Performer Nystromformer SViTE

Builds on: Mixture-of-Experts MoBA Deformable DETR Agent Attention DeiT FlashAttention

Closest papers by meaning

• kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers (ICLR 2025 · similarity 0.90)
• SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.89)
• Fast attention mechanisms: a tale of parallelism (NeurIPS 2025 · similarity 0.89)
• Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.89)
• Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
• ZeroS: Zero‑Sum Linear Attention for Efficient Transformers (NeurIPS 2025 · similarity 0.88)