在Transformer架构中，多头自注意力机制（MHA，Multi-Head Self-Attention）是其核心组件之一，主要由以下几个步骤实现：

1 输入序列的线性变换

多头自注意力机制的输入是一个序列，通常表示为一个矩阵$X$，其形状为$(N,L,D_{\text{model}})$，其中：

• $N$是批量大小（batch size）。
• $L$是序列长度（sequence length）。
• $D_{model}$是模型的维度（embedding size），指的是输入序列的特征维度，即每个单词或标记（token）的嵌入向量的维度。

在多头自注意力机制中，输入序列$X$会被分别映射到三个不同的矩阵：查询（Query）、键（Key）和值（Value）。这一过程通过线性变换实现，具体如下：

• 查询（Query）矩阵：$Q=X{W}^Q$
• 键（Key）矩阵：$K=X{W}^K$
• 值（Value）矩阵：$V=X{W}^V$

其中，$W^Q、W^K、W^V$是可学习的权重矩阵，它们的形状均为$(D_{\text{model}},D_{\text{model}})$，$D_{\text{model}}$为模型的纬度，因此乘积后所得的$Q,K,V$形状仍然为$(N,L,D_{\text{model}})$。这些权重矩阵将输入序列 $X$ 变换到不同的空间，分别用于计算查询、键和值矩阵。

2 多头机制的实现

多头自注意力机制的核心思想是将输入序列分割成多个子空间，也可以理解为多“头”（heads），每个“头”独立地计算注意力，从而能够捕捉到输入序列中不同位置的特征和关系。

上一步计算的查询矩阵、键矩阵和值矩阵 Q 、 K 和 V 分别被分割成多个“头”（heads）。假设我们有 h 个头，每个头的维度$d_k=\frac{D_{\text{model}}}{h}$ 。假如模型维度（嵌入向量的维度）为512，“头”的数量为4，则分割后每个头的维度 $d_k$ 为128。

分割过程如下：

• 将 $Q$ 分割成 h 个子矩阵 $Q_1,Q_2,Q_3\dots ,Q_h$，每个子矩阵（头）的形状为 $(N,L,d_k)$。分割后的矩阵$Q^{\prime }$形状为$(N,L,h,d_k)$。
• 将 $K$ 分割成 h 个子矩阵 $K_1,K_2,K_3\dots ,K_h$，每个子矩阵（头）的形状为 $(N,L,d_k)$。分割后的矩阵$K^{\prime }$形状为$(N,L,h,d_k)$。
• 将 $V$ 分割成 h 个子矩阵 $V_1,V_2,V_3\dots ,V_h$，每个子矩阵（头）的形状为 $(N,L,d_k)$。分割后的矩阵$V^{\prime }$形状为$(N,L,h,d_k)$。
  注意：虽然每个子矩阵的头形状相同，但是其中的数值不相同。

那分割过程是如何完成的？可以通过矩阵乘法和重塑操作实现。
例如我们将查询矩阵 $Q$ 进行分割：

• 首先将 $Q$ 的权重矩阵 $W^Q$ 分割为 $h$ 个子矩阵$W_1^Q,W_1^Q\dots W_h^Q$，每个子矩阵的形状为$(D_{\text{model}},d_k)$，然后打包为一个整体形状为$(D_{\text{model}},h,d_k)$的矩阵 $W_{\text{n}ew}^Q$ 。
• 然后将矩阵$W_{\text{n}ew}^Q$与 $Q$ 相乘，得到 $Q$ 的重塑矩阵（分割后的矩阵） $Q^{\prime }$，即 $Q^{\prime }=Q{W}_{\text{n}ew}^Q$，由于 $Q$ 的形状为$(N,L,D_{\text{model}})$，以 $Q$ 的最后维度 $D_{\text{model}}$ 匹配 $W_{\text{n}ew}^Q$的 $D_{\text{model}}$ 维度进行乘积，因此两者乘积的结果 $Q^{\prime }$ 形状为 $(N,L,h,d_k)$。

3 每个头的注意力计算

对于每个头 $i(i=1,2,\dots h)$，共 $h$ 个头，计算注意力分数（attention scores）和加权和（weighted sum）。

• 注意力分数：计算查询 $Q$ 分割后的一个子矩阵 $Q_i$ 和键 $K$ 分割后的一个子矩阵 $K_i$ 转置矩阵 $K_i^T$ 的乘积，然后除以一个缩放因子 $\sqrt{d_k}$ 用于缩放乘积结果，防止数值过大导致梯度消失或爆炸，$d_k$为“头”的维度，由之前计算所得。因此注意力分数计算公式为：
  $${\text{scores}}_i=\frac{Q_i{K}_i^T}{\sqrt{d_k}}$$

  其中，$Q_i$的形状为 $(N,L,d_k)$， $K_i$ 的转置矩阵形状为 $(N,d_k,L)$，以 $Q_i$的维度 $d_k$执行乘法，与 $K_i$ 的维度 $d_k$ 进行匹配相乘，则结果 ${\text{scores}}_i$ 的形状为 $(N,L,L)$。

• 注意力权重：对计算得到的注意力分数应用Softmax函数，得到注意力权重：
  $${\text{weights}}_i=\text{Softmax}({\text{scores}}_i)$$

• 加权和：将注意力权重与对应分割后的值子矩阵 $V_i$ 相乘，得到每个头的输出：
  $${\text{output}}_i={\text{weights}}_i{V}_i$$
  其中，${\text{weights}}_i$ 的形状为 $(N,L,L)$，$V_i$ 的形状为 $(N,L,d_k)$，同样以 ${\text{weights}}_i$ 的最后维度 $L$ 执行乘法，与 $V_i$ 的维度 $L$ 进行匹配相乘，因此 ${\text{output}}_i$ 的形状为 $(N,L,d_k)$。

根据此计算方法将每个“头”的 ${\text{output}}_i$ (第 $i=1,2,\dots h$ 个头的输出）计算出来。

4 多头输出的合并

将所有头的输出合并为一个矩阵，然后通过一个线性变换进行整合：

• 合并：将所有头的输出拼接在一起：
  $$\text{concat}=[{\text{output}}_1,{\text{output}}_2,\dots {\text{output}}_h]$$
  其中，${\text{output}}_i$ 的形状为 $(N,L,d_k)$，且$d_k\times h=D_{\text{model}}$，因此 $\text{concat}$ 的形状为$(N,L,D_{\text{model}})$。
  $\text{concat}$ 也被称为 $\text{Attention}$，它即是Transformer多头自注意力机制计算的最终结果。

  可将上述的所有计算过程合并成一个公式：
  $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$

• 线性变换：通过一个线性变换整合多头的输出，用于后续计算：
  $$\text{output}=\text{concat}{W}^O$$
  其中，$W^O$是一个可学习的权重矩阵，形状为$(D_{\text{model}},D_{\text{model}})$。

DeepSeek使用的Transformer架构中，使用的是优化后的多头自注意力机制，名为多头潜在注意力（MLA，Multi-Head Latent Attention），主要优化的点在于：能够显著降低内存占用和计算开销，同时保持模型性能。

如果对此感兴趣可关注后续文章。