通过嵌入门控机制,大幅压缩大模型推理缓存,同时保持高精度。
EG-MLA: Embedding-Gated Multi-head Latent Attention for Scalable and Efficient LLMs
- 在潜空间引入令牌特异性嵌入门控,精细调节压缩后的键值向量。
- 相比MHA减少91.6%缓存占用,比MLA额外节省59.9%内存且提升准确率。
- 适用于超大规模语言模型部署,支持百亿参数级扩展,适合追求高效推理的场景。
降低键值(KV)缓存大小是实现大语言模型(LLMs)高效推理的关键步骤,尤其在延迟和内存受限条件下。尽管多头注意力(MHA)具有强表征能力,但带来显著内存开销。近期提出的多头潜空间注意力(MLA)通过将键值表示压缩至共享潜空间,实现了性能与缓存效率之间的更好权衡。尽管如此,现有方法在不损失性能的前提下仍存在进一步压缩的空间。本文提出嵌入门控多头潜空间注意力(EG-MLA),一种对MLA的新型扩展,在保持高性能的同时进一步减少缓存。EG-MLA在潜空间引入令牌特异性嵌入门控机制,实现对压缩后键值向量的细粒度调制,仅需极少额外计算。相较于MHA,EG-MLA实现超过91.6%的缓存缩减,性能损失可忽略;相较MLA,其在多种推理基准上持续提升任务准确率,并实现高达59.9%的额外内存节省。理论分析表明,嵌入门控诱导隐式高阶交互,实证评估显示其在不同模型规模和压缩比例下均具鲁棒泛化能力。特别地,我们成功将EG-MLA扩展至超过10亿参数规模,验证其在大规模语言模型部署中的实际可行性。这些结果确立了EG-MLA作为内存与计算高效的注意力机制,支持现代大模型的可扩展高性能推理。
原文摘要 · Abstract (English)
Reducing the key-value (KV) cache size is a crucial step toward enabling efficient inference in large language models (LLMs), especially under latency and memory constraints. While Multi-Head Attention (MHA) offers strong representational power, it incurs significant memory overhead. Recent work on Multi-head Latent Attention (MLA) mitigates this by compressing KV representations into a shared latent space, achieving a better trade-off between performance and cache efficiency. While MLA already achieves significant KV cache reduction, the scope for further compression remains limited without performance loss. In this paper, we propose \textbf{Embedding-Gated Multi-head Latent Attention (EG-MLA)}, a novel extension of MLA that further reduces KV cache size while enhancing representational expressiveness. EG-MLA introduces a token-specific embedding gating mechanism applied in the latent space, enabling fine-grained modulation of compressed KV vectors with minimal additional computation. Compared to MHA, EG-MLA achieves over 91.6\% reduction in KV cache size with negligible performance degradation. Relative to MLA, EG-MLA consistently improves task accuracy across diverse reasoning benchmarks while achieving up to 59.9\% additional memory savings. Our theoretical analysis highlights how embedding gating induces implicit high-order interactions, and empirical evaluations demonstrate robust generalization across model scales and compression regimes. Notably, we successfully scale EG-MLA to over 1 billion parameters, demonstrating its practical viability for large-scale LLM deployment. These results establish EG-MLA as a memory- and compute-efficient attention mechanism that enables scalable, high-performance inference in modern LLMs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。