Mamba核心作者新作：取代DeepSeek在用的注意力机制，专为推理打造

曾撼动 Transformer 统治地位的 Mamba 作者之一Tri Dao，刚刚带来新作——

提出两种专为推理 " 量身定制 " 的注意力机制。

在保持模型性能不变的情况下，将解码速度和吞吐量最高提升 2 倍，大大优化了模型的长上下文推理能力。

这项研究的三位作者均来自普林斯顿大学，论文主要有两大贡献：

其一，提出 Grouped-Tied Attention（GTA），与已集成到 LLaMA 3 的注意力机制 GQA 质量相当，但 KV 缓存用量减少约 50%。

其二，提出 Grouped Latent Attention（GLA），与 DeepSeek 所使用的注意力机制 MLA 质量匹配，但解码速度更快，某些情况下比 FlashMLA 快 2 倍。

按照作者之一 Ted Zadouri 的总结：

GTA 是 GQA 的有效替代品，而 GLA 是 MLA 的实用替代品。

一言以蔽之，通过优化注意力机制的内存使用和计算逻辑，在不牺牲模型生成质量的前提下，可显著提升大语言模型的推理效率和硬件资源利用率，尤其在长上下文场景中优势更为突出。

相关论文公布后，一众研究者也赶来祝贺 ~

那么，这项研究具体讲了些啥？

引入推理感知注意力机制

概括而言，论文核心引入了推理感知注意力机制，即针对模型推理阶段的内存冗余、计算低效、长上下文瓶颈等问题，重新设计注意力机制。

据 Tri Dao 介绍，这项研究的起点始于一个想法：

在推理驱动 AI 发展的时代，" 理想 " 架构应该是什么样子？

尤其在涉及长上下文推理时，当前的大语言模型（LLM）面临内存访问瓶颈和并行性限制两大难题。

就是说，模型生成文字时，每次都要从内存里调取大量 " 历史记录 "，不仅导致每个字生成变慢，而且只能按顺序生成、没法让多个芯片同时干活。

对此，团队打算从两个方向重新设计注意力机制：

更高的硬件效率：通过增加 " 每字节内存加载的计算量 "（算术强度），减少对内存带宽的依赖；

保持并行可扩展性：在不牺牲模型并行训练 / 推理能力的前提下优化解码速度。

而最终提出的 GTA 和 GLA，在减少 KV 缓存用量的同时，模型质量保持与现有方案相当，且解码速度显著提升。

这里提到的 " 现有方案 "，主要指早已闻名学术界的两种方法：

一是分组查询注意力（GQA）机制，它通过分组共享 KV 缓存减少内存占用，在视觉 Transformer（ViT）等任务中表现良好，适用于大规模数据处理，目前已应用于 Llama 3 等开源模型。

二是多头潜在注意力（MLA）机制，最早可追溯到《Attention Is All You Need》这篇论文，后被 DeepSeek 再次带火。它关注的是在不同层之间如何融合注意力信息，能减少每一层的冗余计算。

不过，由于 GQA 仍需为每组查询头存储独立 KV、MLA 并行优化不足，故仍需进一步改进。

下面分别展开团队提出的新方法 GTA 和 GLA。

分组绑定注意力机制 GTA

GTA 的核心设计思路是：将不同查询头的键（Key）和值（Value）状态进行组合与重用，减少内存传输次数。

具体而言（右图），它将多头注意力的头分为若干组（Group），每组内的头共享相同的 Key 和 Value 参数。计算时，同一组内的头使用相同的 KV 缓存，仅查询（Query）参数独立。

相比之下，中间传统的多头注意力机制（MHA）每个查询头都有独立的键和值，由于没有共享，导致它需要更多的内存来存储所有的键和值。

再对比 GQA 来看（左图），GQA 分组共享 KV 但每组仍独立存储，而 GTA 通过参数绑定实现了更彻底的 KV 重复利用。

分组潜在注意力机制 GLA

而 GLA 的设计则采用了双层结构：

潜在层（Latent Layer）：引入固定数量的潜在 Tokens，作为全局上下文的压缩表示，替代部分原始 Token 的 KV 缓存；

分组头机制：将查询头分组，每组头共享潜在 Token 的 KV，同时保留与原始 Token 的交互。

在解码过程中，对比 MLA（左图），GLA 通过共享联合潜在表示减少了每个设备需要加载的 KV 缓存量，从而减少了内存访问量。

并且由于每个设备上的 KV 缓存量减少了，更多的请求也可以同时处理。

「GQA 和 MLA」的有效替代品

那么，GTA 和 GLA 的效果究竟如何呢？

团队在四种规模的模型上进行了实验，包括小型（183M）、中型（433M）、大型（876M）和 XL（1471M）。这些模型基于 FineWeb-Edu-100B 数据集训练，采用 GPT-3 架构和 Llama 3 分词器。

测试的指标主要分为两大类：

质量指标：困惑度（Perplexity）、下游任务准确率（Winogrande、SciQ 等 7 个基准）；

效率指标：每 Token 解码延迟、吞吐量、KV 缓存占用量。

实验对比了 GQA、MLA、FlashMLA、传统 MHA 等多种注意力机制。

困惑度实验显示，GTA 在中大型模型上优于 GQA，说明GTA 可能更适合模型的进一步扩展；而 GLA 在多数场景下与 MLA 相当，说明GLA 的设计是合理的，它能在并行计算和模型质量之间找到一个较好的平衡点。

几种方案在下游任务中（涵盖典型常识推理、逻辑推理和知识问答等场景）的整体表现差距不大。

但从变化趋势来看（下图为从中型到大型），GTA 和 GLA可以保持或提高从中型到 XL 尺寸的下游任务性能。

KV 缓存方面，不牺牲模型质量的前提下，GTA 相比 GQA 减少约 50% 的 KV 缓存，验证了 " 参数绑定 + 分组重用 " 的有效性。

同时，针对查询长度为 1 的情况，MLA 已接近计算瓶颈（达到 610 TFLOPS/s ），而 GLA 尚未使计算资源饱和（360 TFLOPS/s ）。

且随着序列长度从 1K 增加到 64K ，GLA 的解码速度比 FlashMLA 快 2 倍。

此外，在实时服务器性能测试中，对于 64 个并发请求的输出吞吐量（越高越好），相同并行方案下 GLA 的表现均优于 MLA。

接下来，团队还在 DeepSeek Coder V2 Base （236B）模型上，当使用 FP8 精度时，对比了二者在不同预填充长度和解码长度下的输出吞吐量。

结果显示，在预填充长度为 32K 和 64K 时，GLA-8 的输出吞吐量明显高于 MLA。这表明在处理长上下文时，GLA 在吞吐量上优于 MLA。

在处理不均衡负载时，GLA-8 同样展现出更高的输出吞吐量。这表明 GLA 在处理不同长度的请求时，能够更有效地利用资源，提高整体性能。

以上实验均验证了论文作者的说法，「GTA 和 GLA」是「GQA 和 MLA」的有效替代品。

论文作者均来自普林斯顿大学

论文作者包括 Tri Dao 在内一共三位，均来自普林斯顿大学。

Ted Zadouri，目前是普林斯顿大学博士生，研究方向为机器学习。

之前曾在英特尔有过两段实习经历（研究深度学习），还短暂在 AI 创企 Cohere 担任研究员。

Hubert Strauss，普林斯顿大学研究工程师，研究方向为机器学习和模型深度学习。

本科毕业于法国知名工程学校 Arts et M é tiers，之后在佐治亚理工学院取得运筹学硕士学位。

毕业后曾有多段实习和工作经历，成为普林斯顿大学工程师之前曾在一家公司担任机器学习工程师，负责模型训练和 Transformer 优化。

Tri Dao，目前是普林斯顿大学计算机科学助理教授，还是生成式 AI 初创公司 Together AI 的首席科学家。

他因提出一系列优化 Transformer 模型注意力机制的工作而闻名学界。

其中最有影响力的，是其作为作者之一提出了Mamba 架构，这一架构在语言、音频和基因组学等多种模态中都达到了 SOTA 性能。

尤其在语言建模方面，无论是预训练还是下游评估，Mamba-3B 模型都优于同等规模的 Transformer 模型，并能与两倍于其规模的 Transformer 模型相媲美。

另外他还参与发表了FlashAttention1-3 版本，FlashAttention 被广泛用于加速 Transformers，已经使注意力速度提高了 4-8 倍。

Anyway，回到这项研究，论文作者 Ted Zadouri 直言：

这只是迈向 test-time 推理 " 理想 " 架构的第一步！

论文：

https://arxiv.org/abs/2505.21487

代码：

https://github.com/Dao-AILab/grouped-latent-attention

参考链接：

[ 1 ] https://x.com/tri_dao/status/1928170648863473892

[ 2 ] https://x.com/gm8xx8/status/1927572103806554262

[ 3 ] https://x.com/tedzadouri/status/1928167296821854363

—  完  —

 量子位 AI 主题策划正在征集中！欢迎参与专题365 行 AI 落地方案，一千零一个 AI 应用，或与我们分享你在寻找的 AI 产品，或发现的AI 新动向。

也欢迎你加入量子位每日 AI 交流群，一起来畅聊 AI 吧～

一键关注 点亮星标

科技前沿进展每日见

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法！