需要你对MHA、MQA、GQA有足够了解,相信本文能帮助你对MLA有新的认识。

本文内容都来自https://www.youtube.com/watch?v=0VLAoVGf_74,如果阅读本文出现问题,建议直接去看一遍。

        按照Deepseek设定一些参数值:输入token长度n=10,注意力头数目n_h=128,每个注意力头的隐含层维度d_h=128,transformer block层数 l =61,使用fp16存储参数。

        先来看MHA的kv-cache计算:

cachesize = 2*n*n_h*d_h*l*2 = 2*10*128*128*61*2 = 40MB

(第一个2是因为要保存K和V,第二个2是因为fp16占2bit)

        MQA和GQA的思路是通过不同注意力头之间共享参数,减少注意力头数目n_h来达到降低开销的目的。

        这样的问题是参数的共享会导致模型效果下降,毕竟原本有128个头,128份KV参数,每份KV参数都会计算出不一样的注意力分布,让模型能更好的根据所有的注意力分布去预测下一个词,而现在128份参数变成了1份,预测效果下降是必然的。

        如何解决这个问题?如何只保留1份参数,但又能计算出128个不同的注意力分布呢?

        MLA给出的答案是,只保留原本128分参数中共有的部分,而每份参数独有的部分则提取出来,不进行保存。

        这里就碰到了MLA第一个比较难理解的点,就是怎么找出128个W_K的共有部分和独有部分?(只以K为例,V也是一样的) 

        答案是不用去找,而是从一开始就用两个矩阵,分别去学习共有部分和独有部分。也就是下图中的W_DKV和W_UK,其中W_DKV学习共有部分,W_UK学习独有部分。也就是说128个注意力头,会共用W_DKV,但是每个注意力头的W_UK是独有的,这样保证了128个注意力头能计算出128个不同的注意力分布。

这里就会碰到MLA第二个比较难理解的点,为什么最后kv-cache只用保存L_KV,而不用保存K和V?

答案是根本就不存在K和V,MLA很巧妙的利用矩阵乘法,把W_UK与W_Q融合,把W_UV和W_O融合。至于为什么能这样做,可以从公式中找出答案。

说不存在W_UK和W_UV其实并不严谨,但是这样可以更方便去理解,其实这里所谓的把W_UK与W_Q融合是指输入先经过W_Q,紧跟着就经过W_UK,从结果上来看,跟先把W_UK与W_Q相乘得到W_QUK,然后输入经过W_QUK的效果是一样的。

        原本,加入W_DKV后,注意力的计算公式为:

 A = QK^T = (XW_Q)(XW_{DKV}W_{UK})^T

        按照矩阵运算,上述公式可以写成下述形式:

A = (XW_Q)(W_{UK}^TW_{DKV}^TX^T) = (XW_QW_{UK}^T)(W_{DKV}^TX^T) = (XW_QW_{UK}^T)(XW_{DKV})^T = (XW_{QUK})L_{KV}^T

        我们完全可以将W_QW_{UK}^T视作一个矩阵W_{QUK},它和W_Q并没有什么本质区别,只是维度需要调整(当然实际实现上还是两个矩阵,分开来学习)。从上式中,我们发现注意力计算公式中的K消失了

        然后是最终输出O的计算:

O = AVW_O = A(XW_{DKV}W_V)W_O = A(XW_{DKV})(W_VW_O) = AL_{KV}W_{VO}

        同理,这样就能把W_V融进W_O中,我们能够发现,最终输出的计算公式中,V也消失了

        最后的效果如下图,我们需要保存的只有L_KV,它是128个注意力头共用的,所以只需要保存一份,存储开销计算如下,整个计算公式中完全不需要考虑注意力头数目:

cachesize = n*d_h'*l*2 = 10*576*61*2 = 0.7MB

        开销降低40/0.7,约57倍,也就是deepseek技术报告中公布的压缩倍数。 

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐