补全llm知识体系的地基：KV Cache与注意力机制的改进

MHA：多头注意力机制

• 在MHA中，我们假设最初输入X乘上QKV的权重矩阵之后，得到了尺寸为[s, d]的QKV矩阵（seq_len, dimension）
• 假设多头的数目为head，则QKV都被切片为head个[s, d/head]子矩阵
• 这里，如果每一轮都全部计算的话，QKV内存占用非常大

KV Cache

• 由于在生成时，存在mask机制，导致即使在生成后续字符时，计算已生成字符的结果仍保持一模一样
• 因此，不需要每生成一个字符，都重新计算所有的QKV，只需保留之前的结果，并计算最后一行传入Q的QKV结果
• 之前的KV缓存结果就是KV Cache
• 注意，缓存的是之前序列输入时，计算得到的K和V矩阵，新一轮输入时，只计算新的token，再做拼接即可
• 理论上来说，Attention模块的输出的seq_len维中，之前所有行的结果都是完全一致的，只有最后一行需要用新的输入计算。但如果不缓存KV的话，还是要计算出全量的KV矩阵（使用了之前的输出），才能计算出最后一行的Attention结果（因为前面的全都没有被mask）。这里应该主要是访存瓶颈和内存占用瓶颈，而非节省是否用前面序列的Q参与计算的计算量

Multi-Query Attention：MQA公用KV

• MQA中，保有head个Q的同时，只同时存储一个KV。
• 这样，需要保存的KV矩阵的数目就减小到了1 / head
• 但这样也同时导致，KV的参数量变小了。原来可以用K、V学习的语言键值信息，现在全部要靠Q来学习。虽然在其他地方补回参数量可以改善一部分，但还是影响表现的

Group-Query Attention：成组公用KV

• 介于MHA和MQA中间的折中
• 设置group个组，每个组保留一个KV，即每head/group个query公用一个KV Cache
• 主流用法，如Llama，ChatGLM等

Multi-Latent Attention：低秩分解，RoPE和Deepseek

• LoRA的低秩分解方法在减小KV Cache上也同样有效。实际上，从输入X到QKV矩阵的过程本身也是在做低秩分解。此处特别地，我们引入一个线性变换矩阵C，在X与W_k,W_v相乘之前，先乘一个C，这样就降低了KV权重矩阵的尺寸（原来需要的宽度是X的宽度d，现在只需要C的宽度d_c）
• 但这样的话，GQA就退化成了MHA（或者说，单纯做个低秩分解对GQA并没有什么提升，还会使得每个Head的KV Cache在被C投影之后不一样了）
• 计算注意力时，我们用的\(Q*K^T = x * W_q * (x * W_k)^T\)。但在低秩分解之后，变成了\(Q*K^T = x * W_q * (x * C * W_k)^T = x * (W_q * W_k^T) * (x * C)^T\)。我们把中间的\((W_q * W_k^T)\)作为Q的投影矩阵（新的W_q，只需计算一次），就实现了用投影结果\((x * C)\)来替代K。对于V同样也有类似的操作。
• 至此，只要取c的宽度为g(d_k+d_v)，就得到了一个表达能力更强且完全等价的GQA
• 如果c的宽度取的更小，那么得到了缓存矩阵\((x * C)\)也会更小，这样虽然会因为低秩损失掉一点信息，但投影矩阵的表达能力又弥补了这一点。所以，这样就能降低KV Cache了。（计算效率和模型能力的trade-off）
• 更进一步，如果在输入之前加上RoPE位置编码，则\((W_q * W_k^T)\)项中间会多出一项由在序列中的位置差决定的旋转矩阵。这个旋转矩阵就在迭代中需要多次计算了，所以简单的MLA兼容不了RoPE
• 最终，可以在输入时加宽维度来解决位置编码，新宽度取d_c的四分之一（from DeepseekV2方案）
• 另外，对Q也做低秩分解，能够减少训练占用和梯度大小
  

参考：
https://kexue.fm/archives/10091
https://zhuanlan.zhihu.com/p/21151178690