Mixture-of-Top-k Attention: Efficient Attention via Scalable Fast Weights
cs.LG, cs.CV
Submitted: 2026-02-01
Updated: 2026-09-27
Comments: NeurIPS 2026; Code is available at https://github.com/QishuaiWen/MiTA
Code: https://github.com/QishuaiWen/MiTA
License: http://creativecommons.org/licenses/by/4.0/
The gist: The vanilla self-attention mechanism in Transformers can be viewed as a two-layer fast-weight MLP, whose weights are dynamically induced by inputs and whose hidden dimension is equal to the sequence
Terminology
Abstract
The vanilla self-attention mechanism in Transformers can be viewed as a two-layer fast-weight MLP, whose weights are dynamically induced by inputs and whose hidden dimension is equal to the sequence length N. As the context extends, the expressive capacity of such an N-width MLP increases, but it becomes unscalable for extremely long sequences. Recently, this fast-weight perspective has motivated the Mixture-of-Experts (MoE) attention mechanism, which partitions the sequence into rigid blocks, treats them as fast-weight experts, and sparsely routes the tokens to them. In this paper, we elevate this perspective to a unifying framework for efficient attention mechanisms, interpreting them as making fast weights scalable through either routing or compression, and organizing them into a five-dimensional taxonomy. Then, we propose Mixture-of-Top- k Attention (MiTA), which employs a small set of landmark queries to gather top- k attended key-value pairs as query-aware and deformable routed experts, while compressing the N-width MLP into a narrower shared expert. Consequently, our MiTA improves the flexibility of prior MoE attention from rigid to deformable fast-weight experts, as well as the scalability of prior top- k attention from query-specific set to reusable top- k set. We conduct extensive experiments on vision tasks showing the superior effectiveness and efficiency of our MiTA, and also uncovering intriguing properties such as an emergent token-pruning effect and easy generalization from standard attention. Code is available at https://github.com/QishuaiWen/MiTA.
Sources
- ViT$^3$: Unlocking Test-Time Training in Vision
- Parallel Loop Transformer for Efficient Test-Time Computation Scaling
- MoBA: Mixture of Block Attention for Long-Context LLMs
- MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head
- Linformer: Self-Attention with Linear Complexity
- DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
- VMoBA: Mixture-of-Block Attention for Video Diffusion Models
- Mixture of Contexts for Long Video Generation
- MoGA: Mixture-of-Groups Attention for End-to-End Long Video Generation
- Distilling the Knowledge in a Neural Network
- Titans: Learning to Memorize at Test Time
- General-Purpose In-Context Learning by Meta-Learning Transformers
- Scaling Laws for Neural Language Models
- Longformer: The Long-Document Transformer
- Online normalizer calculation for softmax
- ViT-5: Vision Transformers for The Mid-2020s
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks