大模型学习入门——Day3:注意力机制

本系列笔记的教材：快乐学习大模型-DataWhale团队

注意力机制

注意力机制最先源于计算机视觉领域，其核心思想为当我们关注一张图片，我们往往无需看清楚全部内容而仅将注意力集中在重点部分即可。而在自然语言处理领域，我们往往也可以通过将重点注意力集中在一个或几个 token，从而取得更高效高质的计算效果。
注意力机制有三个核心变量：Query（查询值）、Key（键值）和 Value（真值）。注意力机制就是根据一个查询（Query），来判断一组信息（Key 和 Value）中哪些部分更重要，然后将重要的部分信息加权组合起来，得到一个最终结果。
举个例子

• 我们的数据（键-值对 Key-Value）:
  • Keys = ["apple", "banana", "chair"] (信息的“标签”)
  • Values = [10, 5, 2] (信息本身)
• 我们的查询（Query）:
  • Query = "一种水果"

第一步：计算相似度分数 (Query · Key)
我们需要一个方法来衡量我们的查询（Query）和每个键（Key）的相似程度。在真实的神经网络中，这是通过高维向量的点积来完成的。在第一章中我们知道词由词向量表示，语义相近的词他们的词向量也非常接近。
$$x=Q{K}^T$$
此处的 K 即为将所有 Key 对应的词向量堆叠形成的矩阵。基于矩阵乘法的定义，x 即为 q 与每一个 k 值的点积。现在我们得到的 x 即反映了 Query 和每一个 Key 的相似程度。
第二步：缩放分数 (Scale)
在真实的 Transformer 模型中，为了防止分数过大导致梯度消失，会将其除以一个缩放因子 $\sqrt{d_k}$（$d_k$ 是Key向量的维度）。
第三步：计算注意力权重 (Softmax)
接下来，我们需要将这些原始分数转换成一个总和为 1 的、真正的“注意力权重”分布。这一步通过 Softmax 函数完成。

公式:
$$\text{Softmax}(x_i)=\frac{e^{x_i}}{{\sum }_j{e}^{x_j}}$$
权重解读:
当查询是“一种水果”时，模型会把 $softmax(x_1)$ 的注意力放在 “apple” 上，$$softmax(x_2)$ 的注意力放在 “banana” 上，而只有 $softmax(x_3)$ 的注意力放在 “chair” 上。
第四步：加权求和得到最终输出 (Multiply by Values)
最后一步，我们将上一步得到的注意力权重，与我们最初的 值 (Values) 进行加权求和。

公式:
$$\text{Output}=\sum _i({\text{weight}}_i\cdot V_i)$$

注意力机制核心计算公式

给定查询（Query）、键（Key）、值（Value）矩阵，注意力得分的计算公式为：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

其中：

• $Q\in {\mathbb{R}}^{n\times d_k}$ 表示查询矩阵
• $K\in {\mathbb{R}}^{m\times d_k}$ 表示键矩阵
• $V\in {\mathbb{R}}^{m\times d_v}$ 表示值矩阵
• $d_k$ 为键向量的维度
• $\mathrm{softmax}$ 用于归一化注意力权重

注意力权重矩阵（score）：

$$\mathrm{score}=\frac{Q{K}^{\top }}{\sqrt{d_k}}$$

最终加权求和：

$$\mathrm{output}=\mathrm{softmax}(\mathrm{score})V$$
pytorch代码实现

'''注意力计算函数'''
def attention(query, key, value, dropout=None):'''args:query: 查询值矩阵key: 键值矩阵value: 真值矩阵'''# 获取键向量的维度，键向量的维度和值向量的维度相同d_k = query.size(-1) # 计算Q与K的内积并除以根号dk# transpose——相当于转置scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)# Softmaxp_attn = scores.softmax(dim=-1)if dropout is not None:p_attn = dropout(p_attn)# 采样# 根据计算结果对value进行加权求和return torch.matmul(p_attn, value), p_attn

自注意力机制

在开始计算前，我们先回顾一下注意力机制的三个核心角色：查询(Query)、键(Key)、值(Value)。整个过程就是：用我的 Query (Q) 去和所有元素的 Key (K) 算一个“相关度分数”，然后根据这个分数给所有元素的 Value (V) 分配一个权重，最后把所有加权后的 Value 合并起来，就得到了一个融合了全局上下文的新表示。
所谓自注意力，即是计算本身序列中每个元素对其他元素的注意力分布，即在计算过程中，Q、K、V 都由同一个输入通过不同的参数矩阵计算得到。在 Encoder 中，Q、K、V 分别是输入对参数矩阵 Wq、Wk、Wv 做积得到，从而拟合输入语句中每一个 token 对其他所有 token 的关系。

举个例子
准备工作：我们的迷你 NLP 任务

• 输入句子: “The cat sat” (为了计算简单，我们用3个单词)
• 简化假设:
  • 每个单词的词嵌入 (Embedding) 维度 d_model = 4。
  • 注意力头的数量 h = 1 (后文多头注意力部分的时候h=2)。
  • 每个头的维度 d_k = d_v = d_model / h = 4 / 1 = 4。

1. 初始输入和权重矩阵

我们假设已经通过词嵌入层得到了每个单词的初始向量，并随机初始化了将它们转换为 Q, K, V 的权重矩阵。

• 输入词嵌入矩阵 X (3个单词, 每个4维):

  X = [[1, 0, 1, 0],  # The[0, 1, 1, 0],  # cat[1, 1, 0, 1]]  # sat
  

  需要学习的权重矩阵 (我们手动设定一些简单的整数值):
  • W_Q (用于生成 Query): [[1,0,1,0], [0,1,1,0], [1,1,0,1], [1,0,0,1]]
  • W_K (用于生成 Key): [[0,1,0,1], [1,0,1,1], [0,0,1,0], [1,1,0,0]]
  • W_V (用于生成 Value): [[0,0,1,1], [1,1,0,0], [0,1,0,1], [1,0,1,0]]

2. 计算 Q, K, V 矩阵

这是自注意力的第一步：Q = X @ W_Q, K = X @ W_K, V = X @ W_V
现在我们有了计算所需的 Q, K, V 矩阵。

Q = X @ W_Q  # [[2, 1, 1, 1], [1, 2, 1, 1], [2, 1, 2, 2]]
K = X @ W_K  # [[0, 1, 1, 1], [1, 0, 2, 1], [2, 2, 1, 2]]
V = X @ W_V  # [[1, 0, 2, 1], [1, 2, 0, 1], [2, 2, 2, 2]]

3. 计算注意力分数 scores = QK^T

Q 矩阵的每一行（代表一个词的Query）都要和 K 矩阵的每一行（代表每个词的Key）做点积。
$K^T=\left(\begin{array}{ccc}0& 1& 2\\ 1& 0& 2\\ 1& 2& 1\\ 1& 1& 2\end{array}\right)$

$Q{K}^T=\left(\begin{array}{cccc}2& 1& 1& 1\\ 1& 2& 1& 1\\ 2& 1& 2& 2\end{array}\right)\cdot \left(\begin{array}{ccc}0& 1& 2\\ 1& 0& 2\\ 1& 2& 1\\ 1& 1& 2\end{array}\right)=\left(\begin{array}{ccc}2& 5& 9\\ 4& 4& 9\\ 6& 8& 12\end{array}\right)$
4. 缩放 (Scale)

将分数除以 $\sqrt{d_k}=\sqrt{4}=2$。
scaled_scores = scores / 2
$=\left(\begin{array}{ccc}1.0& 2.5& 4.5\\ 2.0& 2.0& 4.5\\ 3.0& 4.0& 6.0\end{array}\right)$

5. Softmax 归一化

对 scaled_scores 的每一行独立应用 Softmax。我们以第一行为例：

• $e^{1.0}\approx 2.718$, $e^{2.5}\approx 12.182$, $e^{4.5}\approx 90.017$
• Sum = 2.718 + 12.182 + 90.017 = 104.917
• weights_row1 = [2.718/104.9, 12.182/104.9, 90.017/104.9] = [0.026, 0.116, 0.858]

计算所有行后，得到注意力权重矩阵 A:
A = [[0.026, 0.116, 0.858], [0.035, 0.035, 0.930], [0.024, 0.064, 0.912]]

6. 加权求和 Output = A @ V

$$Output=\left(\begin{array}{ccc}0.026& 0.116& 0.858\\ ...& & \end{array}\right)\cdot \left(\begin{array}{cccc}1& 0& 2& 1\\ 1& 2& 0& 1\\ 2& 2& 2& 2\end{array}\right)$$

计算第一行（“The” 的新表示）：
output_row1 = 0.026 * [1,0,2,1] + 0.116 * [1,2,0,1] + 0.858 * [2,2,2,2]
≈ [1.858, 1.948, 1.768, 1.858]

这个 [1.858, 1.948, 1.768, 1.858] 就是 “The” 在融合了整个句子上下文之后的新表示。

多头注意力机制

现在我们使用 2个头 (h=2)，每个头的维度 d_k=2。
步骤 1: 分割 Q, K, V 到不同的“头”

在真实实现中，我们会使用不同的 W 矩阵为每个头生成独立的 Q, K, V。为简化，我们直接将之前计算的 Q, K, V 矩阵按列劈开。

• Q_head1 = [[2,1], [1,2], [2,1]] (Q的前2列)
• Q_head2 = [[1,1], [1,1], [2,2]] (Q的后2列)
• K_head1 = [[0,1], [1,0], [2,2]] (K的前2列)
• K_head2 = [[1,1], [2,1], [1,2]] (K的后2列，此处为简化调整值)
• V_head1 = [[1,0], [1,2], [2,2]] (V的前2列)
• V_head2 = [[2,1], [0,1], [2,2]] (V的后2列)

步骤 2: 并行计算每个头的注意力

现在我们对每个头独立地重复上面的自注意力计算过程。$\sqrt{d_k}=\sqrt{2}\approx 1.414$。

Head 1:

• scores_h1 = (Q_head1 @ K_head1.T) / 1.414
• weights_h1 = softmax(scores_h1)
• Z1 = weights_h1 @ V_head1 (假设计算结果为 [[1.8, 1.5], [1.7, 1.6], [1.9, 1.8]])

Head 2:

• scores_h2 = (Q_head2 @ K_head2.T) / 1.414
• weights_h2 = softmax(scores_h2)
• Z2 = weights_h2 @ V_head2 (假设计算结果为 [[1.5, 1.9], [1.4, 1.8], [1.6, 2.0]])

步骤 3: 拼接与最终线性变换

1. 拼接 (Concatenate) 所有头的输出:
   Z = Concat(Z1, Z2)
= [[1.8, 1.5, 1.5, 1.9], # "The" 的新多头表示
[1.7, 1.6, 1.4, 1.8], # "cat" 的新多头表示
[1.9, 1.8, 1.6, 2.0]] # "sat" 的新多头表示

2. 最终线性变换:
   将拼接后的矩阵 Z 乘以一个最终的输出权重矩阵 W_O（也是需要学习的），得到多头注意力的最终输出。
   MultiHead_Output = Z @ W_O

掩码注意力机制

这个机制实现起来非常巧妙。它不是真的删除或隐藏未来的词，而是在计算注意力权重的过程中，通过一个“掩码 (Mask)”来让未来位置的注意力权重在数学上变为零。这个过程发生在计算完注意力分数、但在进行 Softmax 之前。

具体步骤：

1. 和标准自注意力一样，我们先计算出注意力分数矩阵 Scores = QK^T / sqrt(d_k)。
   假设这是我们对句子 “The cat sat” 计算出的（已缩放的）注意力分数：

原始分数矩阵 (Scaled Scores):

      The   cat   sat   (Keys)
The [[1.0,  2.5,  4.5],
cat  [2.0,  2.0,  4.5],
sat  [3.0,  4.0,  6.0]]
(Queries)

2. 我们创建一个“掩码矩阵 (Mask Matrix)”。这是一个上三角矩阵，其中我们希望保留的位置（当前和过去的位置）值为 0，希望屏蔽掉的位置（未来的位置）为一个非常大的负数（比如 -1e9 或负无穷）。
   我们需要一个上三角部分为 $-\infty$ 的掩码（0 代表保留，1 代表屏蔽，这里我们用 $-\infty$ 来表示屏蔽的效果）。

[[  0, -∞, -∞],[  0,   0, -∞],[  0,   0,   0]]

3. 我们将这个掩码矩阵加到注意力分数矩阵上。
   Masked Scores = Scaled Scores + Mask
   $=\left(\begin{array}{ccc}1.0& 2.5& 4.5\\ 2.0& 2.0& 4.5\\ 3.0& 4.0& 6.0\end{array}\right)+\left(\begin{array}{ccc}0& -\infty & -\infty \\ 0& 0& -\infty \\ 0& 0& 0\end{array}\right)=\left(\begin{array}{ccc}1.0& -\infty & -\infty \\ 2.0& 2.0& -\infty \\ 3.0& 4.0& 6.0\end{array}\right)$
4. 最后，对加了掩码的分数矩阵进行 Softmax 计算。

• 第一行 (处理 “The”): softmax([1.0, -∞, -∞])

  • $e^{1.0}\approx 2.718$
  • $e^{-\infty }=0$
  • $e^{-\infty }=0$
  • Sum = 2.718 + 0 + 0 = 2.718
  • 权重: [2.718/2.718, 0/2.718, 0/2.718] = [1.0, 0.0, 0.0]
  • 解读: “The” 只能关注它自己。

• 第二行 (处理 “cat”): softmax([2.0, 2.0, -∞])

  • $e^{2.0}\approx 7.389$
  • $e^{2.0}\approx 7.389$
  • $e^{-\infty }=0$
  • Sum = 7.389 + 7.389 + 0 = 14.778
  • 权重: [7.389/14.778, 7.389/14.778, 0/14.778] = [0.5, 0.5, 0.0]
  • 解读: “cat” 可以关注 “The” 和它自己，但不能关注 “sat”。

• 第三行 (处理 “sat”): softmax([3.0, 4.0, 6.0])

  • 因为没有未来位置，所以这一行的计算与标准自注意力相同。
  • 权重结果约为: [0.024, 0.064, 0.912]
  • 解读: “sat” 可以关注所有词。

最终的注意力权重矩阵 (Masked Attention Weights):
A_masked = [[1.0, 0.0, 0.0 ],
[0.5, 0.5, 0.0 ],
[0.024, 0.064, 0.912]]

这个权重矩阵接下来会像标准自注意力一样，与 V (值矩阵) 相乘，得到最终的输出。
公式

$$\text{Attention}(Q,K,V)=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}}+M)V$$

• 这里的 M 就是我们的掩码矩阵，其中需要屏蔽的位置是 $-\infty$，其他位置是 0。

代码实现

import torch.nn as nn
import torch'''多头自注意力计算模块'''
class MultiHeadAttention(nn.Module):def __init__(self, args: ModelArgs, is_causal=False):# 构造函数# args: 配置对象super().__init__()# 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵assert args.n_embd % args.n_heads == 0# 模型并行处理大小，默认为1。model_parallel_size = 1# 本地计算头数，等于总头数除以模型并行处理大小。self.n_local_heads = args.n_heads // model_parallel_size# 每个头的维度，等于模型维度除以头的总数。self.head_dim = args.dim // args.n_heads# Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd# 这里通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，# 不理解的读者可以自行模拟一下，每一个线性层其实相当于n个参数矩阵的拼接self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)self.wk = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)self.wv = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)# 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)# 注意力的 dropoutself.attn_dropout = nn.Dropout(args.dropout)# 残差连接的 dropoutself.resid_dropout = nn.Dropout(args.dropout)# 创建一个上三角矩阵，用于遮蔽未来信息# 注意，因为是多头注意力，Mask 矩阵比之前我们定义的多一个维度if is_causal:mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))mask = torch.triu(mask, diagonal=1)# 注册为模型的缓冲区self.register_buffer("mask", mask)def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):# 获取批次大小和序列长度，[batch_size, seq_len, dim]bsz, seqlen, _ = q.shape# 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)# 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)# 因为在注意力计算中我们是取了后两个维度参与计算# 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度而不是直接按注意力输入展开，是因为view的展开方式是直接把输入全部排开，# 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的目标xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)xq = xq.transpose(1, 2)xk = xk.transpose(1, 2)xv = xv.transpose(1, 2)# 注意力计算# 计算 QK^T / sqrt(d_k)，维度为 (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)# 掩码自注意力必须有注意力掩码if self.is_causal:assert hasattr(self, 'mask')# 这里截取到序列长度，因为有些序列可能比 max_seq_len 短scores = scores + self.mask[:, :, :seqlen, :seqlen]# 计算 softmax，维度为 (B, nh, T, T)scores = F.softmax(scores.float(), dim=-1).type_as(xq)# 做 Dropoutscores = self.attn_dropout(scores)# V * Score，维度为(B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)output = torch.matmul(scores, xv)# 恢复时间维度并合并头。# 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)# contiguous 函数用于重新开辟一块新内存存储，因为Pytorch设置先transpose再view会报错，# 因为view直接基于底层存储得到，然而transpose并不会改变底层存储，因此需要额外存储output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)# 最终投影回残差流。output = self.wo(output)output = self.resid_dropout(output)return output