Tri Dao大神新作重磅来袭：取代DeepSeek的MLA，专为推理“量身定制”的注意力机制 来源：量子位曾撼动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。 分组绑定注意力机制GTAGTA的核心设计思路是：将不同查询头的键（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 推荐阅读 MIT新晋副教授何恺明《计算机视觉进展》课程，附Slides与资料下载 西电IEEE Fellow团队出品！最新《Transformer视觉表征学习全面综述》如何做好科研？这份《科研阅读、写作与报告》PPT，手把手教你做科研奖金675万！3位科学家，斩获“中国诺贝尔奖”！最新 2022「深度学习视觉注意力 」研究概述，包括50种注意力机制和方法！【重磅】斯坦福李飞飞《注意力与Transformer》总结，84页ppt开放下载！2021李宏毅老师最新40节机器学习课程！附课件+视频资料 欢迎大家加入DLer-大模型技术交流群！ ??长按识别，邀请您进群！