Transformer 注意力机制

总阅读量 加载中... 次

注意力机制（Attention）作为 Transformer 架构的核心组件，自提出以来一直是自然语言处理领域的研究热点。本文围绕 Transformer 中的四种注意力机制展开：Scaled Dot-Product Attention、Multi-Head Attention、Multi-Query Attention 以及 Grouped-Query Attention。

1. 缩放点积注意力

Scaled Dot-Product Attention 最早由 Vaswani 等人在论文《Attention Is All You Need》（2017）中提出，是 Transformer 架构的基础计算单元。该机制的核心思想源于早期的神经机器翻译工作，通过计算查询向量与键向量的点积相似度来确定注意力权重，进而对值向量进行加权求和。

Attention (Q, K, V) = Softmax (\frac{Q K^{T}}{\sqrt{d_{k}}}) V

矩阵	含义	矩阵维度 (单样本批量)
$Q$	查询矩阵 (Query)	$n \times d_{k}$
$K$	键矩阵 (Key)	$n \times d_{k}$
$V$	值矩阵 (Value)	$n \times d_{v}$
$Q K^{T}$	原始注意力分数	$n \times n$
$Softmax (\cdot)$	注意力权重分布	$n \times n$
Output	注意力机制最终输出	$n \times d_{v}$

1.1 核心要素定义

假设输入序列包含 $n$ 个向量，每个向量的特征维度为 $d_{model}$ 。在计算 Attention 之前，输入矩阵 $X \in R^{n \times d_{model}}$ 会通过三个不同的线性变换矩阵（权重矩阵），映射得到三个核心矩阵：

Query (查询矩阵) $Q$ ： $Q = X W_{Q}$ （其中 $W_{Q} \in R^{d_{model} \times d_{k}}$ ）
Key (键矩阵) $K$ ： $K = X W_{K}$ （其中 $W_{K} \in R^{d_{model} \times d_{k}}$ ）
Value (值矩阵) $V$ ： $V = X W_{V}$ （其中 $W_{V} \in R^{d_{model} \times d_{v}}$ ）

注：在标准的 Transformer 中，通常令 $d_{k} = d_{v} = d_{model} / {num}_{heads}$ 。

1.2 详细计算步骤

Attention 的计算主要分为以下 5 个步骤：

1.2.1 步骤 1：计算原始注意力分数

通过将 $Q$ 和 $K$ 的转置进行矩阵乘法，计算出 Query 和 Key 之间的相似度。点积结果越大，说明两个向量的相关性越高。

Scores = Q K^{T}

维度变化： $(n \times d_{k}) \times (d_{k} \times n) \to (n \times n)$
结果矩阵中的第 $(i, j)$ 个元素，代表第 $i$ 个单词对第 $j$ 个单词的原始注意力得分。

1.2.2 步骤 2：缩放操作

将原始分数除以 $\sqrt{d_{k}}$ （即 Key 向量维度的平方根）。

Scaled Scores = \frac{Q K^{T}}{\sqrt{d_{k}}}

目的：当 $d_{k}$ 很大时，点积的结果会非常大，导致后面经过 Softmax 函数时梯度变得极小（进入饱和区）。除以 $\sqrt{d_{k}}$ 可以起到方差缩放的作用，让训练更加稳定。

1.2.3 步骤 3：掩码操作

对缩放后的分数应用掩码（Mask），将需要屏蔽的位置设置为负无穷（ $- \infty$ ），使其在后续 Softmax 计算后的权重趋近于 0。

Masked Scores = Mask (\frac{Q K^{T}}{\sqrt{d_{k}}})

目的：
- Padding Mask：屏蔽输入序列中的 Padding 部分（通常是因为 batch 内句子长度不同，填充的无效 tokens）
- Causal Mask / Sequence Mask：在 Decoder 中，防止当前位置看到后续位置的信息（确保自回归特性）
实现方式：将需要屏蔽的位置乘以一个非常大的负数（如 $- 1 e 9$ 或 $- \infty$ ），Softmax 会对这些位置输出接近 0 的权重

1.2.4 步骤 4：归一化

对掩码处理后的分数在行方向上应用 Softmax 函数，将其转化为概率分布（所有权重相加为 1）。

Attention Weights = Softmax (Mask (\frac{Q K^{T}}{\sqrt{d_{k}}}))

维度： $(n \times n)$
此时矩阵里的数值就是最终的注意力权重（矩阵中每行的和为 1）。

1.2.5 步骤 5：加权求和

用计算出的注意力权重矩阵去乘以 Value 矩阵 $V$ ，实现对 Value 的加权聚合。

Attention (Q, K, V) = Softmax (Mask (\frac{Q K^{T}}{\sqrt{d_{k}}})) V

维度变化： $(n \times n) \times (n \times d_{v}) \to (n \times d_{v})$
最终输出的矩阵中，每一个向量都融合了整个序列中与其相关的其他向量的信息。

1.3 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47import torch
from torch import nn, Tensor
import torch.nn.functional as F
from typing import Any

class ScaleDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention

    完整计算公式：
    $$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

    """
    def __init__(self) -> None:
        super(ScaleDotProductAttention, self).__init__()

    def forward(self, 
                Q: Tensor, 
                K: Tensor, 
                V: Tensor, 
                mask: Tensor = None, 
                return_attention_socre: bool = False
    ) -> tuple[Tensor, Any] | Tensor:
        """
        :param Q: Tensor, shape = (batch_size, num_heads, seq_len, d_k) - 查询矩阵
        :param K: Tensor, shape = (batch_size, num_heads, seq_len, d_k) - 键矩阵
        :param V: Tensor, shape = (batch_size, num_heads, seq_len, d_v) - 值矩阵 (通常 d_k == d_v)
        :param mask: Tensor, shape = (batch_size, num_heads, seq_len, seq_len] - 掩码矩阵（可选，用于Transformer中的Padding或Decoder的Causal Mask）
        :param return_attention_socre: 是否返回注意力分数
        :return:
            output: Tensor, shape = (batch_size, num_heads, seq_len, d_k)
        """
        d_k = Q.shape[-1] # 向量的维度,即d_k
        # STEP 1: 计算原始注意力分数 \frac{QK^T}{\sqrt(d_k)}
        scores = torch.matmul(Q, K.transpose(-1, -2))
        # STEP 2: 缩放
        scores = scores / torch.sqrt(torch.tensor(d_k, device=Q.device))
        # STEP 3: Mask
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        # STEP 4: 归一化softmax 得到注意力权重
        attention_scores = F.softmax(scores, dim=-1)
        # STEP 5: 加权求和
        output = torch.matmul(attention_scores, V)
        if return_attention_socre:
            return output, attention_scores
        return output

在 Vaswani 等人提出 Scaled Dot-Product Attention 之后，同一论文中进一步引入了多头注意力机制（Multi-Head Attention），以增强模型的表达能力。

2. 多头注意力

Multi-Head Attention（多头注意力机制） 是在 Scaled Dot-Product Attention（缩放点积注意力）的基础上发展而来的。

它的核心思想是：与其只计算一次复杂的注意力（单头），不如将输入拆分为多个低维的子空间（多个头），在每个子空间内独立地计算注意力，最后再将所有头的输出拼接并线性映射回去。这样做能让模型同时关注来自不同位置、不同表示子空间的信息。

MultiHead (Q, K, V) = Concat ({head}_{1}, \dots, {head}_{h}) W^{O}

where {head}_{i} = Softmax (\frac{Q_{i} K_{i}^{T}}{\sqrt{d_{k}}}) V_{i}

2.1 核心要素定义

假设输入序列矩阵为 $X \in R^{n \times d_{model}}$ （其中 $n$ 为序列长度， $d_{model}$ 为模型的隐藏层维度），我们要计算 $h$ 个头的注意力。

在多头注意力中，我们需要 $h$ 组不同的线性变换矩阵。对于第 $i$ 个头（ $i = 1, 2, \dots, h$ ），其权重矩阵分别定义为：

$W_{Q}^{(i)} \in R^{d_{model} \times d_{k}}$
$W_{K}^{(i)} \in R^{d_{model} \times d_{k}}$
$W_{V}^{(i)} \in R^{d_{model} \times d_{v}}$

通常，为了保持总参数量和计算量与单头一致，我们会令 $d_{k} = d_{v} = d_{model} / h$ 。

2.2 详细计算步骤

Multi-Head Attention 的计算主要分为以下 4 个步骤：

2.2.1 步骤 1：线性映射

将输入矩阵 $X$ 分别乘以 $h$ 组不同的权重矩阵，得到每个头专属的 $Q_{i}, K_{i}, V_{i}$ 矩阵：

Q_{i} = X W_{Q}^{(i)}, K_{i} = X W_{K}^{(i)}, V_{i} = X W_{V}^{(i)}

维度：每个头的 $Q_{i}, K_{i}$ 维度为 $(n \times d_{k})$ ， $V_{i}$ 维度为 $(n \times d_{v})$ 。

2.2.2 步骤 2：独立计算每个头的注意力（缩放点积注意力）

每个头独立运行标准的缩放点积注意力机制，得到该头的输出 ${head}_{i}$ ：

{head}_{i} = Attention (Q_{i}, K_{i}, V_{i}) = Softmax (\frac{Q_{i} K_{i}^{T}}{\sqrt{d_{k}}}) V_{i}

维度： ${head}_{i} \in R^{n \times d_{v}}$

2.2.3 步骤 3：拼接

将所有 $h$ 个头计算出来的输出矩阵在特征维度（列方向）上横向拼接在一起：

Concat ({head}_{1}, {head}_{2}, \dots, {head}_{h})

维度变化：由于有 $h$ 个维度为 $(n \times d_{v})$ 的矩阵拼接，拼接后的总维度变为 $(n \times (h \times d_{v}))$ 。
因为 $h \times d_{v} = d_{model}$ ，所以拼接后的矩阵维度重新恢复到了 $(n \times d_{model})$ 。

2.2.4 步骤 4：最后的线性变换

为了让多头聚集起来的信息能够充分融合，拼接后的结果会通过一个最终的输出权重矩阵 $W^{O} \in R^{d_{model} \times d_{model}}$ 进行线性映射：

MultiHead (Q, K, V) = Concat ({head}_{1}, {head}_{2}, \dots, {head}_{h}) W^{O}

维度： $(n \times d_{model}) \times (d_{model} \times d_{model}) \to (n \times d_{model})$

2.3 矩阵层面的高效实现（并行化）

在实际工程实现中，并不会真的用 for 循环去一个头一个头地计算。标准的并行化做法是：

一次性投影：直接用一个大矩阵将 $X$ 映射成总维度为 $(n \times d_{model})$ 的 $Q, K, V$ 。
维度重塑（Reshape）：将维度从 $(n \times d_{model})$ 改为 $(n \times h \times d_{k})$ 。
维度置换（Transpose）：转换为 $(batch_size, h, n, d_{k})$ 。
批量矩阵乘法（BMM）：利用高维矩阵乘法，一条指令同时计算所有样本、所有头的注意力，极大提升了 GPU 的运行效率。

2.4 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56import torch
from torch import nn, Tensor
import torch.nn.functional as F
from typing import Any

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Attention
    """
    def __init__(self, n_heads: int, d_model: int):
        super(MultiHeadAttention, self).__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.n_heads = n_heads # 注意力头的数量
        self.d_model = d_model # 模型的总维度
        # 每个 head 的维度
        self.d_k = d_model // n_heads
        # 线性映射：Q, K, V
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        # 输出映射
        self.W_O = nn.Linear(d_model, d_model, bias=False)
        self.attention = ScaleDotProductAttention()

    def forward(self, 
                Q: torch.Tensor, 
                K: torch.Tensor, 
                V: torch.Tensor, 
                mask: torch.Tensor = None
    ):
        """
        :param Q: Tensor, shape = (batch_size, seq_len, d_model)
        :param K: Tensor, shape = (batch_size, seq_len, d_model)
        :param V: Tensor, shape = (batch_size, seq_len, d_model)
        :param mask: Tensor, shape = (batch_size, seq_len, d_model)
        :return:
            output: Tensor, shape = (batch_size, seq_len, d_model)
        """
        batch_size = Q.shape[0]
        # STEP 1: 线性变换 将原始的 Q, K, V 映射到新的特征空间
        q_proj = self.W_q(Q) # (batch_size, seq_len, d_model)
        k_proj = self.W_k(K) # (batch_size, seq_len, d_model)
        v_proj = self.W_v(V) # (batch_size, seq_len, d_model)
        # STEP 2: reshape + transpose 拆成多头 
        # 形状从 (batch_size, seq_len, d_model) 变为 (batch_size, n_heads, seq_len, d_k)
        q_heads = q_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k_heads = k_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v_heads = v_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        # STEP 3: 调用attention (batch_size, n_heads, seq_len, d_k)
        context = self.attention(q_heads, k_heads, v_heads, mask=mask)  
        # STEP 4: 拼接多头
        context = context.transpose(1, 2).contiguous()
        output = context.view(batch_size, -1, self.d_model)
        # STEP 5: 线性映射
        output = self.W_O(output)
        return output

尽管 Multi-Head Attention 在建模能力上表现优异，但随着模型规模增大和序列长度增加，其推理效率问题日益突出。针对这一问题，Shazeer 等人在论文《Fast Transformer Decoding: One Write-Head is All You Need》（2019）中提出了 Multi-Query Attention（MQA），通过让多个 Query 头共享同一组 Key 和 Value 来显著降低推理时的 KV Cache 开销。

3. 多查询注意力

Multi-Query Attention (MQA) 是由 Noam Shazeer 在 2019 年提出的一种 Attention 变体。

它的核心动机是为了解决大模型推理（Generation）时的性能瓶颈。在传统的 Multi-Head Attention (MHA) 中，每个独立头都有自己的 $Q, K, V$ 矩阵。在推理时，为了避免重复计算，模型会把历史 Token 的 $K$ 和 $V$ 缓存起来（即 KV Cache）。随着序列变长，KV Cache 会吞噬极大的显存带宽和空间。

MQA 的核心思想极其精简：让所有的 Query 头共享同一组 Key 和 Value 头。

3.1 核心要素定义

假设输入序列矩阵为 $X \in R^{n \times d_{model}}$ ，我们要计算 $h$ 个 Query 头。

与 MHA 为每个头都准备独立的 $W_{Q}, W_{K}, W_{V}$ 不同，MQA 只准备：

$h$ 组不同的 Query 权重矩阵： $W_{Q}^{(1)}, W_{Q}^{(2)}, \dots, W_{Q}^{(h)} \in R^{d_{model} \times d_{k}}$
仅仅 1 组 Key 权重矩阵： $W_{K} \in R^{d_{model} \times d_{k}}$
仅仅 1 组 Value 权重矩阵： $W_{V} \in R^{d_{model} \times d_{v}}$

通常，每个头的维度仍满足 $d_{k} = d_{v} = d_{model} / h$ 。

3.2 详细计算步骤

MQA 的计算步骤与标准注意力机制相似，但由于 $K$ 和 $V$ 的共享特性，矩阵的维度和广播（Broadcasting）逻辑发生了变化：

3.2.1 步骤 1：线性映射

投影时，Query 依然有 $h$ 个不同的结果，而 $K$ 和 $V$ 只有单份：

对于第 $i$ 个 Query 头： $Q_{i} = X W_{Q}^{(i)} \in R^{n \times d_{k}}$
公共的 Key 矩阵： $K = X W_{K} \in R^{n \times d_{k}}$
公共的 Value 矩阵： $V = X W_{V} \in R^{n \times d_{v}}$

3.2.2 步骤 2：独立计算每个查询头的注意力

每个独立的 Query 头 $Q_{i}$ 都去和同一个 $K$ 计算注意力分数，并和同一个 $V$ 进行加权求和。

{head}_{i} = Attention (Q_{i}, K, V) = Softmax (\frac{Q_{i} K^{T}}{\sqrt{d_{k}}}) V

维度： ${head}_{i} \in R^{n \times d_{v}}$
注意：这里的 $K^{T}$ 和 $V$ 对所有的 ${head}_{i}$ 来说是完全一模一样的。

3.2.3 步骤 3：拼接

将这 $h$ 个头得到的输出矩阵在特征维度上横向拼接：

Concat ({head}_{1}, {head}_{2}, \dots, {head}_{h})

维度： $h \times (n \times d_{v}) \to (n \times d_{model})$

3.2.4 步骤 4：最后的线性变换

与 MHA 一样，通过一个输出权重矩阵 $W^{O} \in R^{d_{model} \times d_{model}}$ 融合多头信息：

MultiQuery (Q, K, V) = Concat ({head}_{1}, \dots, {head}_{h}) W^{O}

维度： $(n \times d_{model})$

3.3 完整公式总结

MultiQuery (Q, K, V) = Concat ({head}_{1}, \dots, {head}_{h}) W^{O}

where {head}_{i} = Softmax (\frac{Q_{i} K^{T}}{\sqrt{d_{k}}}) V

虽然数学公式看起来和 MHA 极其相似，但请注意其本质区别：在 ${head}_{i}$ 的计算中，所有的 $i$ 共享相同的 $K$ 和 $V$ 。

3.4 并行化矩阵维度的变化对照

在现代深度学习框架中，通常会利用张量的广播机制（Broadcasting）来实现 MQA 的并行化计算。可以通过下表清晰地看出 MHA 与 MQA 在张量维度上的断层式差异：

矩阵 / 张量	MHA (多头注意力) 维度	MQA (多查询注意力) 维度
Query ( $Q$ )	`[batch_size, num_heads, seq_len, d_k]`	`[batch_size, num_heads, seq_len, d_k]`
Key ( $K$ )	`[batch_size, num_heads, seq_len, d_k]`	`[batch_size, 1, seq_len, d_k]`
Value ( $V$ )	`[batch_size, num_heads, seq_len, d_v]`	`[batch_size, 1, seq_len, d_v]`
Score Matrix	`[batch_size, num_heads, seq_len, seq_len]`	`[batch_size, num_heads, seq_len, seq_len]` (通过广播计算)

3.5 为什么 MQA 能极大加速推理？

减少显存占用：在推理时，KV Cache 的大小缩减到了原来的 $1 / h$ （如果模型有 32 个头，KV 显存直接缩减 32 倍）。
打破带宽瓶颈：大模型自回归生成时，瓶颈不在于 GPU 计算力（Math-bound），而在于 GPU 内存带宽（Memory-bound）。由于每次读取的 $K, V$ 变少，GPU 可以花更少的时间去搬运内存，从而大幅提升 Token 的生成速度。

3.6 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53import torch
from torch import nn, Tensor
import torch.nn.functional as F
from typing import Any

class MultiQueryAttention(nn.Module):
    """
    Multi Query Attention 
    """
    def __init__(self, n_heads: int, d_model: int):
        super(MultiQueryAttention, self).__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.n_heads = n_heads
        self.d_model = d_model
        self.d_k = self.d_model // self.n_heads
        self.W_q = nn.Linear(self.d_model, self.d_model, bias=False)
        self.W_k = nn.Linear(self.d_model, self.d_k, bias=False)
        self.W_v = nn.Linear(self.d_model, self.d_k, bias=False)
        self.attention = ScaleDotProductAttention()
        self.W_O = nn.Linear(self.d_model, self.d_model, bias=False)

    def forward(self, 
                Q: torch.Tensor, 
                K: torch.Tensor, 
                V: torch.Tensor, 
                mask: torch.Tensor = None):
        """
        :param Q: Tensor, shape = (batch_size, seq_len, d_model)
        :param K: Tensor, shape = (batch_size, seq_len, d_model)
        :param V: Tensor, shape = (batch_size, seq_len, d_model)
        :param mask: Tensor, shape = (batch_size, 1, seq_len, seq_len)
        :return:
        """
        batch_size = Q.shape[0]
        # 1. 线性映射 (Linear Projection)
        q_proj = self.W_q(Q)
        k_proj = self.W_k(K)
        v_proj = self.W_v(V)
        # 2. 维度重塑与变换 (Reshape & Transpose)
        # Q: [batch_size, len_q, d_model] -> [batch_size, num_heads, len_q, d_k]
        q_heads = q_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        # K 和 V 增加一个 Head 维度（设为 1 即可），以便后续进行广播计算
        # [batch_size, len_k, d_k] -> [batch_size, 1, len_k, d_k]
        k_heads = k_proj.unsqueeze(1)
        v_heads = v_proj.unsqueeze(1)
        # STEP 3: 调用attention (batch_size, n_heads, seq_len, d_k)
        context = self.attention(q_heads, k_heads, v_heads, mask=mask)  
        # STEP 4: 拼接多头
        context = context.transpose(1, 2).contiguous()
        output = context.view(batch_size, -1, self.d_model)
        # STEP 5: 线性映射
        output = self.W_O(output)
        return output

MQA 虽然有效降低了推理开销，但过度共享 KV 导致模型表达能力下降。为解决这一问题，Ainslie 等人在论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》（2023）中提出了 Grouped-Query Attention（GQA），通过引入分组机制在 MHA 和 MQA 之间取得平衡。

4. 分组查询注意力

Grouped-Query Attention (GQA) 是由 Ainslie 等人在 2023 年提出的一种 Attention 变体。

它是 Multi-Head Attention (MHA) 和 Multi-Query Attention (MQA) 的折中方案。MHA 虽然表达能力强，但推理时 KV Cache 显存开销巨大；MQA 虽然极大地压缩了 KV Cache，但由于所有头强行共享一组 KV，会导致模型表达能力下降。

GQA 的核心思想是：将 Query 头分成若干个组（Groups），每一个组内的所有 Query 头共享同一组 Key 和 Value 头。

4.1 核心要素定义

假设输入序列矩阵为 $X \in R^{n \times d_{model}}$ 。我们设定：

Query 的总头数为 $h$ 。
Key 和 Value 的总头数为 $g$ （即分组数）。
每个组内包含的 Query 头数为 $m = h / g$ 。

在 GQA 中，权重矩阵的定义如下：

$h$ 组不同的 Query 权重矩阵： $W_{Q}^{(1)}, W_{Q}^{(2)}, \dots, W_{Q}^{(h)} \in R^{d_{model} \times d_{k}}$
$g$ 组不同的 Key 权重矩阵： $W_{K}^{(1)}, W_{K}^{(2)}, \dots, W_{K}^{(g)} \in R^{d_{model} \times d_{k}}$
$g$ 组不同的 Value 权重矩阵： $W_{V}^{(1)}, W_{V}^{(2)}, \dots, W_{V}^{(g)} \in R^{d_{model} \times d_{v}}$

特殊情况：当 $g = h$ 时，GQA 退化为标准的 MHA；当 $g = 1$ 时，GQA 退化为 MQA。通常大模型中会选择 $g = 8$ 。

4.2 详细计算步骤

GQA 的计算通过“分组共享”的逻辑进行，以下是具体的 4 个步骤：

4.2.1 步骤 1：线性映射

将输入矩阵 $X$ 通过各自的线性层进行投影。

得到 $h$ 个 Query 矩阵： $Q_{i} = X W_{Q}^{(i)} (i = 1, \dots, h)$
得到 $g$ 个 Key 矩阵： $K_{j} = X W_{K}^{(j)} (j = 1, \dots, g)$
得到 $g$ 个 Value 矩阵： $V_{j} = X W_{V}^{(j)} (j = 1, \dots, g)$

4.2.2 步骤 2：将查询映射到对应的键值组

每一个 Query 头 $Q_{i}$ 都有一个专属的 KV 组索引 $j$ 。对应关系为：

j = ⌊ (i - 1) / m ⌋ + 1

也就是说，前 $m$ 个 Query 头共享 $K_{1}, V_{1}$ ，接下来的 $m$ 个 Query 头共享 $K_{2}, V_{2}$ ，以此类推。

4.2.3 步骤 3：独立计算每个头的注意力

每个 Query 头 $Q_{i}$ 与它所属组的 $K_{j}$ 和 $V_{j}$ 进行标准的缩放点积注意力计算：

{head}_{i} = Attention (Q_{i}, K_{j}, V_{j}) = Softmax (\frac{Q_{i} K_{j}^{T}}{\sqrt{d_{k}}}) V_{j}

维度： ${head}_{i} \in R^{n \times d_{v}}$

4.2.4 步骤 4：拼接与最终映射

将所有 $h$ 个头的输出进行拼接，并通过最终的输出矩阵 $W^{O} \in R^{d_{model} \times d_{model}}$ 融合信息：

GQA (Q, K, V) = Concat ({head}_{1}, {head}_{2}, \dots, {head}_{h}) W^{O}

维度： $(n \times d_{model})$

4.3 并行化矩阵维度的变化对照

在工程实现中，我们不会用循环来处理分组，而是利用张量重塑（Reshape）和 PyTorch 的广播（Broadcasting）来并行计算。

以下是三种 Attention 机制在并行化计算时的张量维度对比（假设 batch_size 维度已省略）：

机制	Query ( $Q$ ) 维度	Key ( $K$ ) & Value ( $V$ ) 维度	广播与对齐方式
MHA	`[h, n, d_k]`	`[h, n, d_k]`	头数完全一致，1对1计算。
MQA	`[h, n, d_k]`	`[1, n, d_k]`	KV 头部维度为 1，自动广播给所有 $h$ 个 Q 头。
GQA	`[g, m, n, d_k]`	`[g, 1, n, d_k]`	将 Query 拆出组维度 $g$ 和组内头维度 $m$ ；KV 拆出组维度 $g$ 且组内头维度为 1。在组内进行广播。

4.4 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63import torch
from torch import nn, Tensor
import torch.nn.functional as F
from typing import Any

class GroupedQueryAttention(nn.Module):
    def __init__(self, n_heads: int, d_model: int, n_groups: int) -> None:
        super(GroupedQueryAttention, self).__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        assert n_heads % n_groups == 0, "n_heads must be divisible by n_groups"
        self.n_heads = n_heads
        self.d_model = d_model
        self.n_groups = n_groups
        self.group_size = self.n_heads // n_groups
        self.d_k = self.d_model // self.n_heads
        self.W_q = nn.Linear(self.d_model, self.d_model, bias=False)
        # KV 的总维度缩小为: 分组数 * 单头维度
        self.W_k = nn.Linear(self.d_model, self.n_groups * self.d_k, bias=False)
        self.W_v = nn.Linear(self.d_model, self.n_groups * self.d_k, bias=False)
        self.attention = ScaleDotProductAttention()
        self.W_O = nn.Linear(self.d_model, self.d_model, bias=False)

    def forward(self, 
                Q: torch.Tensor, 
                K: torch.Tensor, 
                V: torch.Tensor, 
                mask: torch.Tensor = None) -> Tensor:
        batch_size = Q.shape[0]
        # 步骤 1：映射成投影矩阵
        q_proj = self.W_q(Q)  # [B, L_q, d_model]
        k_proj = self.W_k(K)  # [B, L_k, num_groups * d_k]
        v_proj = self.W_v(V)  # [B, L_v, num_groups * d_k]
        # 步骤 2：调整维度，引入组(Group)概念
        # Q 拆分为 5D 张量: [B, L_q, num_groups, group_size, d_k] -> 
        # 转置为 [B, num_groups, group_size, L_q, d_k]
        q_heads = q_proj.view(batch_size, -1, self.n_groups, self.group_size, self.d_k).permute(0, 2, 3, 1, 4)
        # K 和 V 同样拆为 5D，但组内头数设为 1: [B, L_k, num_groups, 1, d_k] -> 
        # 转置为 [B, num_groups, 1, L_k, d_k]
        k_heads = k_proj.view(batch_size, -1, self.n_groups, 1, self.d_k).permute(0, 2, 3, 1, 4)
        v_heads = v_proj.view(batch_size, -1, self.n_groups, 1, self.d_k).permute(0, 2, 3, 1, 4)
        # 在内部自动转换 Mask 的维度
        if mask is not None:
            # 情况 A：如果传入的是纯 2D 掩码 [len_q, len_k] (例如全局共享的因果掩码)
            if mask.dim() == 2:
                # [len_q, len_k] -> [1, 1, 1, len_q, len_k]
                mask = mask.unsqueeze(0).unsqueeze(1).unsqueeze(2)
            # 情况 B：如果传入的是标准 3D 批处理掩码 [B, len_q, len_k] (例如带 Padding 的掩码)
            elif mask.dim() == 3:
                # [B, len_q, len_k] -> [B, 1, 1, len_q, len_k]
                mask = mask.unsqueeze(1).unsqueeze(2)
        # 步骤 3：调用 ScaledDotProductAttention
        # q_heads: [B, G, M, L_q, d_k]
        # k_heads: [B, G, 1, L_k, d_k] -> matmul 在第 2 维(M与1)自动广播
        context = self.attention(q_heads, k_heads, v_heads, mask=mask)
        # context 维度: [B, num_groups, group_size, L_q, d_k]
        # 步骤 4：恢复并还原维度 (将组和组内头数合并回原先的多头)
        # 先换回序列长度在前: [B, L_q, num_groups, group_size, d_k]
        context = context.permute(0, 3, 1, 2, 4).contiguous()
        # 展平回 3D 张量: [B, L_q, d_model]
        output = context.view(batch_size, -1, self.d_model)
        # 步骤 5：输出线性映射
        output = self.W_O(output)
        return output

5. 总结

本文综述了 Transformer 注意力机制的演进历程。从最初的 Scaled Dot-Product Attention 到 Multi-Head Attention，再到针对推理优化的 Multi-Query Attention 和折中方案 Grouped-Query Attention，研究者们始终在模型表达能力与推理效率之间寻求平衡。

机制	论文来源	KV 头数	推理效率	表达能力
MHA	Vaswani et al., 2017	$h$ 个	较低	强
MQA	Shazeer et al., 2019	1 个	高	弱
GQA	Ainslie et al., 2023	$g$ 个	中等	中等

综上所述，GQA 作为 MHA 与 MQA 的折中方案，在保持较好表达能力的同时显著降低了 KV Cache 开销，已成为当代大语言模型的主流选择。