【注意力MHA,MQA,GQA,MLA】

注意力机制优化简明图解

1. 多头注意力（MHA）

图示：

Input --> [Attention Head 1]
       --> [Attention Head 2]
       --> [Attention Head 3]
       --> ...
       --> [Attention Head N]
       --> [Concatenate] --> Output

公式：

$\text{Output}=\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_N)$
${\text{head}}_i=\text{Attention}(Q,K,V)$

2. 多查询注意力（MQA）

图示：

Input --> [Shared Keys & Values]
       --> [Attention Head 1]
       --> [Attention Head 2]
       --> [Attention Head 3]
       --> ...
       --> [Concatenate] --> Output

公式：
$\text{Output}=\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_N)$
${\text{head}}_i=\text{Attention}(Q,K_{\text{shared}},V_{\text{shared}})$

3. 分组查询注意力（GQA）

图示：

Input --> [Attention Group 1]
       --> [Attention Group 2]
       --> ...
       --> [Concatenate] --> Output

公式：
$\text{Output}=\text{Concat}({\text{group}}_1,{\text{group}}_2,\dots ,{\text{group}}_M)$
${\text{group}}_j=\text{Attention}(Q_{{\text{group}}_j},K_{{\text{group}}_j},V_{{\text{group}}_j})$

4. 多头潜在注意力（MLA）

图示：

Input --> [Compressed Keys & Values]
       --> [Attention Head 1]
       --> [Attention Head 2]
       --> [Attention Head 3]
       --> ...
       --> [Concatenate] --> Output

公式：
$\text{Output}=\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_N)$
${\text{head}}_i=\text{Attention}(Q,K_{\text{compressed}},V_{\text{compressed}})$

低秩键值联合压缩公式：

$$K_{\text{compressed}}=U_K\cdot S_K\cdot V_K^T$$
$$V_{\text{compressed}}=U_V\cdot S_V\cdot V_V^T$$

图示概述

1. MHA： 每个头独立操作，最终结果拼接。
2. MQA： 多个头共享键和值，只计算一次查询，减少计算量。
3. GQA： 查询分组，每组共享键和值，进一步减少计算量。
4. MLA： 键和值进行压缩，减少内存和计算需求。

这些方法通过不同的策略优化注意力机制，提高了计算效率，降低了内存消耗，使Transformer模型在实际应用中更加高效。