GQA

搜索文档
时空压缩!剑桥大学提出注意力机制MTLA:推理加速5倍,显存减至1/8
机器之心· 2025-06-11 08:24
Transformer架构的核心地位 - Transformer的自注意力机制在推理时能灵活建模长距离上下文,使其成为大语言模型不可替代的核心组件,尽管存在计算复杂度为二次方的问题 [1] - 线性复杂度的替代方案(如RNN、Linear Attention、SSM等)难以真正取代Transformer的地位,尤其在decoder-only架构广泛采用后,自注意力机制的重要性进一步凸显 [1] - 自注意力机制带来的KV缓存问题成为推理效率的关键瓶颈,其显存和带宽开销随模型参数维度扩大而显著上升,限制了推理长度与batch size [1] KV缓存的优化挑战与现有方案 - 当前主流大模型采用Grouped-Query Attention(GQA)机制,通过减少Key/Value头数量来减小KV缓存规模,在效率与效果间取得平衡 [5] - Multi-Query Attention(MQA)是GQA的极端形式,显存占用大幅减少但性能显著下降,而GQA成为大语言模型中最常见的注意力变体 [5] - DeepSeek团队提出的Multi-head Latent Attention(MLA)通过在隐空间压缩KV特征维度提升推理效率,但隐空间压缩幅度受限,KV缓存存储开销仍是瓶颈 [7] MTLA的创新与核心技术 - MTLA首次将时序压缩与隐空间压缩结合,在KV缓存的两个维度同时施加时空压缩策略,显著降低显存与计算成本 [2] - MTLA通过超网络动态融合相邻时间步信息,并设计步幅感知因果掩码确保训练与推理一致性,保持甚至略优于传统注意力机制的性能 [2][14] - 时间压缩率s=2时,MTLA对KV缓存的压缩程度与MQA相当但性能更优,且具备进一步压缩空间 [12][21] MTLA的训练与灵活性 - MTLA引入步幅感知因果掩码解决训练与推理行为不一致问题,保留所有中间状态KV表达,通过矩阵乘法实现高效并行计算 [14][16] - 解耦的旋转位置编码(decoupled RoPE)进一步提升了效率,MTLA在注意力机制与线性模型之间架起桥梁,提供效率与性能的灵活权衡 [17][18] - 当时间压缩率s足够大时,MTLA几乎只保留一个KV缓存,退化为线性序列建模方法 [17] MTLA的性能表现与应用潜力 - MTLA在语音翻译、文本摘要生成等任务中保持与标准MHA相当的质量,同时实现超过5倍推理速度提升和8倍显存占用降低 [20] - MTLA具备大规模部署潜力,尤其适合参数规模扩大和生成序列增长的大语言模型场景,有望成为自注意力模块的重要替代方案 [23] - MTLA实现代码已开源,但工程落地需社区持续推动,改动复杂度高于GQA和MQA [24][25]
Mamba核心作者新作:取代DeepSeek在用的注意力机制,专为推理打造
量子位· 2025-06-01 11:40
一水 发自 凹非寺 量子位 | 公众号 QbitAI 曾撼动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的实用替代品。 一言以蔽之,通过优化注意力机制的 内存使用 和 计算逻辑 ,在不牺牲模型生成质量的前提下,可显著提升大语言模型的推理效率和硬件资 源利用率,尤其在长上下文场景中优势更为突出。 相关论文公布后,一众研究者也赶来祝贺~ | Sara Hooker ...