Mixture-of-Top-k注意力：通过可扩展快速权重实现高效注意力
机器学习
2026-05-12 v5 计算机视觉与模式识别
摘要
Transformer中的标准自注意力机制可视为一个两层快速权重MLP，其权重由输入动态诱导，其隐藏维度等于序列长度N。随着上下文的扩展，这种N宽度MLP的表达能力增加，但对于极长序列会变得不可扩展。最近，这种快速权重视角激发了Mixture-of-Experts(MoE)注意力机制，将序列划分为刚性块，将其作为快速权重专家，并稀疏地将token路由到这些专家。在本文中，我们将这一视角提升为统一的高效注意力机制框架，解释它们如何通过路由或压缩来实现快速权重的可扩展性，并将其组织为五维分类法。然后，我们提出了Mixture-of-Top-k注意力(MiTA)，其采用一小组标志查询来收集top-k注意力键值对作为查询感知且可变形的路由专家，同时将N宽度MLP压缩为更窄的共享专家。因此，MiTA提高了先前MoE注意力从刚性到可变形的快速权重专家灵活性，以及先前top-k注意力从特定查询集合到可重用top-k集合的可扩展性。我们在视觉任务上进行了大量实验，显示MiTA在有效性和效率方面具有优势，同时也发现了诸如出现令牌修剪效应以及易于从标准注意力推广等有趣的属性。代码已在https://github.com/QishuaiWen/MiTA上提供。
引用
@article{arxiv.2602.01219,
title = {Mixture-of-Top-k Attention: Efficient Attention via Scalable Fast Weights},
author = {Qishuai Wen and Zhiyuan Huang and Xianghan Meng and Wei He and Chun-Guang Li},
journal= {arXiv preprint arXiv:2602.01219},
year = {2026}
}
备注
Code is available at https://github.com/QishuaiWen/MiTA