MHA & MQA & GQA & MLA

0. 引言

引入多头注意力机制是为了让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，从而增强模型的表达能力。

因为Transformer的encoder和decoder两个部分结构不同，在计算多头注意力时也会有些许不同。而后续的大模型大多数用的都是 decoder-only 结构，这种结构天然适合做KV-Cache，这是一种以空间换时间的策略，为了减少KV-Cache的内存占用就有了后面的MQA（Multi-Query Attention），GQA（Grouped-Query Attention）以及MLA（Multi-head Latent Attention）。

此篇博文介绍一下从MHA、MQA、GQA到MLA的进化历程。

简单来讲：

• 在 MHA（Multi Head Attention） 中，每个头有自己单独的 key-value 对。标准的多头注意力机制，h个Query、Key 和 Value 矩阵。
• 在 MQA（Multi Query Attention） 中只会有一组 key-value 对，多头注意力的一种变体，也是用于自回归解码的一种注意力机制。与MHA不同的是，MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。
• 在 GQA（Grouped Query Attention） 中，会对 attention 进行分组操作，query 被分为 N 组，每个组共享一