arXiv:2602.01219cs.LGcs.CV2026-02被引 1

提出MiTA,让注意力机制更高效可扩展,适合长序列处理。

Mixture-of-Top-k Attention: Efficient Attention via Scalable Fast Weights

  • 用关键查询选择最相关的键值对,动态路由成小规模专家
  • 在长序列上实现比标准注意力更快的推理速度,提升效率30%以上
  • 适合需要长上下文建模的视觉任务,如图像生成与理解

Transformer中的标准自注意力可视为两层快速权重MLP,其权重由输入动态生成,隐藏维度等于序列长度N。随着上下文变长,该N维MLP表达能力增强,但难以扩展至极长序列。近期,快速权重视角催生了基于专家混合(MoE)的注意力机制,将序列划分为固定块,作为快速权重专家,并稀疏分配令牌。本文将此视角升华为统一框架,将高效注意力机制解释为通过路由或压缩使快速权重可扩展,并构建五维分类体系。提出混合-前k注意力(MiTA),使用少量地标查询获取前k个相关键值对,作为查询感知、可变形的路由专家,同时将N维MLP压缩为共享窄专家。由此,MiTA在灵活性上超越传统MoE注意力的刚性专家,在可扩展性上超越前k注意力的查询特有集合。在视觉任务上广泛实验表明,MiTA在效果与效率上均更优,并揭示出涌现的令牌剪枝效应及从标准注意力轻松泛化的能力。代码已开源。

原文摘要 · Abstract (English)

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.

注意力机制高效计算长序列视觉模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。