分组潜在注意力机制GLA

搜索文档
Mamba核心作者新作:取代DeepSeek在用的注意力机制,专为推理打造
猿大侠· 2025-06-02 12:22
核心观点 - 提出两种专为推理优化的注意力机制GTA和GLA,在保持模型性能不变的情况下,将解码速度和吞吐量最高提升2倍,显著优化长上下文推理能力[1][5] - GTA是GQA的有效替代品,KV缓存用量减少约50%;GLA是MLA的实用替代品,解码速度更快,某些情况下比FlashMLA快2倍[2][3][11] - 通过优化注意力机制的内存使用和计算逻辑,提升大语言模型的推理效率和硬件资源利用率,尤其在长上下文场景中优势突出[5][8] 技术贡献 GTA(分组绑定注意力机制) - 将不同查询头的键和值状态进行组合与重用,减少内存传输次数[15] - 多头注意力的头分为若干组,每组内的头共享相同的Key和Value参数,仅查询参数独立[15] - 相比GQA,通过参数绑定实现更彻底的KV重复利用,KV缓存减少约50%[16][28] GLA(分组潜在注意力机制) - 通过共享联合潜在表示减少每个设备需要加载的KV缓存量,减少内存访问量[19] - 增加每字节内存加载的计算量,减少对内存带宽的依赖,保持并行可扩展性[18] - 解码速度比FlashMLA快2倍,随着序列长度从1K增加到64K,优势更明显[30] 实验验证 模型性能 - 在四种规模模型(183M/433M/876M/1471M)上测试,GTA在中大型模型上优于GQA,GLA与MLA相当[22][23] - 下游任务(Winogrande/SciQ等7个基准)整体表现差距不大,但GTA和GLA可保持或提高从中型到XL尺寸的性能[24][25] 效率指标 - GTA相比GQA减少约50%的KV缓存,验证参数绑定+分组重用的有效性[28] - GLA在64个并发请求的输出吞吐量上均优于MLA,处理长上下文时吞吐量更高[31][34] - MLA计算瓶颈达610 TFLOPS/s,GLA尚未饱和(360 TFLOPS/s),资源利用率更高[29] 行业影响 - 研究团队来自普林斯顿大学,核心作者Tri Dao曾提出Mamba架构和FlashAttention系列工作,在优化Transformer领域具有权威性[46][48][49] - 成果已应用于DeepSeek Coder V2 Base(236B)模型,验证了GLA在处理长上下文时的吞吐量优势[34] - 该研究是迈向推理"理想"架构的第一步,未来可能进一步推动大模型推理效率的提升[50]