DeepSeek中的多头潜在注意力(MLA)浅尝 原创

发布于 2025-2-11 11:03
浏览
0收藏

MLA是MHA的变体,因此先来看看MHA。

MHA(多头注意力)

MHA通过将输入向量分割成多个并行的注意力“头”,每个头独立地计算注意力权重并产生输出,然后将这些输出通过拼接和线性变换进行合并以生成最终的注意力表示。

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

Transformer 编码器块内的缩放点积注意力机制和多头注意力机制

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

MHA计算过程

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

MHA 能够理解输入不同部分之间的关系。然而,这种复杂性是有代价的——对内存带宽的需求很大,尤其是在解码器推理期间。主要问题的关键在于内存开销。在自回归模型中,每个解码步骤都需要加载解码器权重以及所有注意键和值。这个过程不仅计算量大,而且内存带宽也大。随着模型规模的扩大,这种开销也会增加,使得扩展变得越来越艰巨。

MLA(多头潜在注意力)

概念:

  • 多头注意力机制:Transformer 的核心模块,能够通过多个注意力头并行捕捉输入序列中的多样化特征。
  • 潜在表示学习:通过将高维输入映射到低维潜在空间,可以提取更抽象的语义特征,同时有效减少计算复杂度。

问题:

1.效率问题:传统多头注意力的计算复杂度为O(n2d),即随着序列长度的增长,键值(Key-Value,KV)缓存的大小也会线性增加,这给模型带来了显著的内存负担。) 2.表达能力瓶颈:难以充分捕捉复杂全局依赖。MLA 通过在潜在空间中执行注意力计算,有效降低复杂度,同时提升建模能力。

MLA 的提出:MLA 将多头注意力机制 与 潜在表示学习 相结合,解决MHA在高计算成本和KV缓存方面的局限性。

MLA的具体做法(创新点): 采用低秩联合压缩键值技术,优化了键值(KV)矩阵,显著减少了内存消耗并提高了推理效率。

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

如上图,在MHA、GQA中大量存在于keys values中的KV缓存——带阴影表示,到了MLA中时,只有一小部分的被压缩Compressed的Latent KV了。

并且,在推理阶段,MHA需要缓存独立的键(Key)和值(Value)矩阵,这会增加内存和计算开销。而MLA通过低秩矩阵分解技术,显著减小了存储的KV(Key-Value)的维度,从而降低了内存占用。

MLA的核心步骤:

  1. 输入映射->潜在空间

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

  1. 潜在空间中的多头注意力计算

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

  1. 映射回原始空间
    将多头注意力结果从潜在空间映射回原始空间:

DeepSeek中的多头潜在注意力(MLA)浅尝-AI.x社区

参考文献:https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf


本文转载自公众号大模型自然语言处理  作者:余俊晖

原文链接:​​https://mp.weixin.qq.com/s/CSK1t9o82un5QDphhOP6fQ​


©著作权归作者所有,如需转载,请注明出处,否则将追究法律责任
收藏
回复
举报
回复
相关推荐