从零开始学大模型:Transformer核心架构
Table of Contents
本系列文章是学习自视频教程的学习笔记
1.2 从第一性原理理解自注意力
现在我们的视野里假设只有看见 token,怎么理解 token ?将其理解为是大模型的最小单位即可。
我们就把它当成一个向量(实际上向量只是 token 的一种表现形式),向量的维度是 d_model。每个 token 都有一个对应的向量表示,这些向量可以通过位置嵌入或者其他方式获得。
例如 我们有三个 token:A、B、C。每个 token 都被表示为一个 d_model 维的向量。我们可以将这些向量堆叠成一个矩阵 X,形状为 (3, d_model)。这就是3个 token 的表示。
在真实训练过程中,我们不单单只输入一个token,而是一个序列的 token。比如我们输入一个句子 “I love NLP”,这个句子被分成三个 token:[“I”, “love”, “NLP”]。每个 token 都有一个对应的向量表示,我们将这些向量堆叠成一个矩阵 X,形状为 (3, d_model)。
进一步的,甚至不只送入一个序列,而是一个 batch 的序列。比如我们有两个句子 “I love NLP” 和 “Transformers are great”,这两个句子被分成六个 token:[“I”, “love”, “NLP”, “Transformers”, “are”, “great”]。每个 token 都有一个对应的向量表示,形状为 (2, 3, d_model)。 这里我们假设每个句子有3个 token,batch size 是2。
同一个序列中的 token 之间一定是存在着某种关系的,我们希望捕捉到这种关系。我们总结成以下三个维度:
- 每个 token 自身有什么信息?
- 每个 token 希望和谁有关系?
- 每个 token 能够提供给别人什么信息?
我们可以通过三个矩阵来捕捉这三个方面: Wq:查询矩阵,捕捉每个 token 希望和谁有关系。 Wk:键矩阵,捕捉每个 token 自身有什么信息。 Wv:值矩阵,捕捉每个 token 能够提供给别人什么信息。
这些矩阵能够表达每个 token 的查询、键和值的信息。
使用 numpy 来手动计算自注意力
我们准备一个序列的玩具数据
batch=1, seq=3, d_model=4
import numpy as np
# 设置NumPy打印选项:保留4位小数,不使用科学计数法
np.set_printoptions(precision=4, suppress=True)
# 这里模拟3个token的嵌入向量,每个token用4维向量表示
X = np.array([[[0.1, 0.2, 0.3, 0.4], # token 0 的嵌入
[0.5, 0.4, 0.3, 0.2], # token 1 的嵌入
[0.0, 0.1, 0.0, 0.1]]], dtype=np.float32) # token 2 的嵌入
# 查询(Query)投影矩阵:用于生成"我在寻找什么"的表示
Wq = np.array([[ 0.2, -0.1],
[ 0.0, 0.1],
[ 0.1, 0.2],
[-0.1, 0.0]], dtype=np.float32)
# 键(Key)投影矩阵:用于生成"我提供什么信息"的表示
Wk = np.array([[ 0.1, 0.1],
[ 0.0, -0.1],
[ 0.2, 0.0],
[ 0.0, 0.2]], dtype=np.float32)
# 值(Value)投影矩阵:用于生成"我的实际内容"的表示
Wv = np.array([[ 0.1, 0.0],
[-0.1, 0.1],
[ 0.2, -0.1],
[ 0.0, 0.2]], dtype=np.float32)
接着将 输入 投影到查询、键和值空间
# 将输入投影到查询、键和值空间
Q = X @ Wq # (1,3,4) @ (4,2) = (1,3,2)
K = X @ Wk # (1,3,4) @ (4,2) = (1,3,2)
V = X @ Wv # (1,3,4) @ (4,2) = (1,3,2)
这里 只需要关注 d_model 的维度,batch 和 seq 维度不变。我们得到了查询矩阵 Q、键矩阵 K 和值矩阵 V。
我们看一下 QKV 形状
Q shape: (1, 3, 2)
Q (查询矩阵) =
[[ 0.01 0.07]
[ 0.11 0.05]
[-0.01 0.01]]
K shape: (1, 3, 2)
K (键矩阵) =
[[0.07 0.07]
[0.11 0.05]
[0. 0.01]]
V shape: (1, 3, 2)
V (值矩阵) =
[[ 0.05 0.07]
[ 0.07 0.05]
[-0.01 0.03]]
计算注意力分数(缩放点积)
# 计算注意力分数(缩放点积)
scale = 1.0 / np.sqrt(Q.shape[-1]) # scale = 1/sqrt(2) ≈ 0.707
attn_scores = (Q @ K.transpose(0,2,1)) * scale # (1,3,2) @ (1,2,3) = (1,3,3)
在这个例子中,d_model 是 4,经过投影后变成了 2,所以我们在计算注意力分数的时候需要除以 sqrt(2) 来进行缩放。
缩放因子 scale = 1/sqrt(2) = 0.7071 注意力分数矩阵 (Q @ K^T * scale):
- 行:查询位置(当前token)
- 列:键位置(要关注的token)
- 值:相似度分数
scale 是为了防止点积结果过大导致梯度消失或者爆炸。我们将查询矩阵 Q 和键矩阵 K 的转置相乘,得到注意力分数矩阵 attn_scores,形状为 (1,3,3),表示每个 token 对其他 token 的注意力分数。
注意力分数本质上是一个相似度矩阵,表示每个 token 对其他 token 的关注程度。我们可以通过 softmax 将这些分数转换为概率分布,不过在那之前,有一个叫做因果掩码的东西需要说一下。
本质上现在大语言模型是借助了 transformer 架构的 decoder 来实现的,而 decoder 的一个重要特性就是它是单向的,也就是说在生成下一个 token 的时候只能看到之前的 token,不能看到未来的 token。为了实现这个特性,我们需要在计算注意力分数的时候引入一个因果掩码(causal mask),这个掩码会将未来 token 的注意力分数设置为负无穷,这样在 softmax 之后这些分数就会变成0,保证了模型只能关注到之前的 token。
这不同于 BERT 这种双向模型,在 BERT 中是没有这个因果掩码的,因为 BERT 是一个 encoder-only 的模型,它在训练的时候可以同时看到左右两边的 token。
将注意力分数 应用 因果掩码
确保每个 token 只能关注之前的 token
mask = np.triu(np.ones((1,3,3), dtype=bool), k=1) # k=1表示主对角线上方为True
attn_scores = np.where(mask, -1e9, attn_scores) # 将被掩盖的位置设为极小值
因果掩码长这个样子:
因果掩码(上三角矩阵,True表示被掩盖):
[[0 1 1]
[0 0 1]
[0 0 0]]
掩码后的分数:
应用掩码后的注意力分数(被掩盖位置设为-1e9):
[[ 3.9598e-03 -1.0000e+09 -1.0000e+09]
[ 7.9196e-03 1.0324e-02 -1.0000e+09]
[ 1.3632e-11 -4.2426e-04 7.0711e-05]]
在这个例子中,我们创建了一个上三角矩阵 mask,主对角线以上的元素为 True,表示这些位置需要被掩盖。我们使用 np.where 将这些位置的注意力分数设置为一个非常小的值(-1e9),这样在后续的 softmax 计算中,这些位置的权重就会接近于0。
归一化注意力分数,得到注意力权重
公式如下
weights = np.exp(attn_scores - attn_scores.max(axis=-1, keepdims=True))
weights = weights / weights.sum(axis=-1, keepdims=True)
到此为止我们的注意力权重已经计算出来了,形状为 (1,3,3),表示每个 token 对其他 token 的关注程度。我们可以看一下这个权重矩阵:
Weights shape: (1, 3, 3)
注意力权重矩阵(因果掩码后)=
[[1. 0. 0. ]
[0.4994 0.5006 0. ]
[0.3334 0.3332 0.3334]]
解读:
- 第0行: token 0 只能看到自己,权重为 [1.0, 0.0, 0.0]
- 第1行: token 1 可以看到 token 0和1,权重和为1
- 第2行: token 2 可以看到所有token,权重和为1
将注意力权重应用到值矩阵 V 上
得到最终的输出
output = weights @ V # (1,3,3) @ (1,3,2) = (1,3,2)
大功告成
Output shape: (1, 3, 2)
Output (注意力输出) =
[[0.05 0.07 ]
[0.06 0.06 ]
[0.0367 0.05 ]]
说明:
- 每个输出向量是对应位置可见的所有值向量的加权平均
- 权重由注意力机制自动学习,反映了token之间的相关性
完整代码
"""
import numpy as np
# 设置NumPy打印选项:保留4位小数,不使用科学计数法
np.set_printoptions(precision=4, suppress=True)
# ============================================================================
# 1. 准备输入数据
# ============================================================================
# 玩具输入数据:batch=1, seq=3, d_model=4
# 这里模拟3个token的嵌入向量,每个token用4维向量表示
X = np.array([[[0.1, 0.2, 0.3, 0.4], # token 0 的嵌入
[0.5, 0.4, 0.3, 0.2], # token 1 的嵌入
[0.0, 0.1, 0.0, 0.1]]], dtype=np.float32) # token 2 的嵌入
# ============================================================================
# 2. 定义权重矩阵(在真实模型中这些是学习得到的参数)
# ============================================================================
# 为了演示的确定性,我们手动设置这些权重值
# 每个权重矩阵的形状都是 (d_model=4, d_k=2),将4维输入映射到2维
# 查询(Query)投影矩阵:用于生成"我在寻找什么"的表示
Wq = np.array([[ 0.2, -0.1],
[ 0.0, 0.1],
[ 0.1, 0.2],
[-0.1, 0.0]], dtype=np.float32)
# 键(Key)投影矩阵:用于生成"我提供什么信息"的表示
Wk = np.array([[ 0.1, 0.1],
[ 0.0, -0.1],
[ 0.2, 0.0],
[ 0.0, 0.2]], dtype=np.float32)
# 值(Value)投影矩阵:用于生成"我的实际内容"的表示
Wv = np.array([[ 0.1, 0.0],
[-0.1, 0.1],
[ 0.2, -0.1],
[ 0.0, 0.2]], dtype=np.float32)
# ============================================================================
# 3. 计算 Q, K, V(线性投影)
# ============================================================================
# 通过矩阵乘法将输入投影到查询、键、值空间
Q = X @ Wq # (1,3,4) @ (4,2) = (1,3,2)
K = X @ Wk # (1,3,4) @ (4,2) = (1,3,2)
V = X @ Wv # (1,3,4) @ (4,2) = (1,3,2)
print("=" * 70)
print("步骤1: 线性投影 - 计算 Q, K, V")
print("=" * 70)
print("Q shape:", Q.shape, "\nQ (查询矩阵) =\n", Q[0])
print("\nK shape:", K.shape, "\nK (键矩阵) =\n", K[0])
print("\nV shape:", V.shape, "\nV (值矩阵) =\n", V[0])
# ============================================================================
# 4. 计算注意力分数(缩放点积)
# ============================================================================
# 计算每个查询与所有键的相似度
# scale因子用于稳定梯度,防止softmax进入饱和区
scale = 1.0 / np.sqrt(Q.shape[-1]) # scale = 1/sqrt(2) ≈ 0.707
attn_scores = (Q @ K.transpose(0,2,1)) * scale # (1,3,2) @ (1,2,3) = (1,3,3)
print("\n" + "=" * 70)
print("步骤2: 计算注意力分数(缩放点积)")
print("=" * 70)
print(f"缩放因子 scale = 1/sqrt({Q.shape[-1]}) = {scale:.4f}")
print("注意力分数矩阵 (Q @ K^T * scale):")
print(" - 行:查询位置(当前token)")
print(" - 列:键位置(要关注的token)")
print(" - 值:相似度分数")
# ============================================================================
# 5. 应用因果掩码(Causal Mask)
# ============================================================================
# 创建上三角掩码矩阵,确保每个位置只能看到它之前的位置
# 这对于自回归语言模型至关重要,防止信息泄露
mask = np.triu(np.ones((1,3,3), dtype=bool), k=1) # k=1表示主对角线上方为True
attn_scores = np.where(mask, -1e9, attn_scores) # 将被掩盖的位置设为极小值
print("\n因果掩码(上三角矩阵,True表示被掩盖):")
print(mask[0].astype(int))
print("\n应用掩码后的注意力分数(被掩盖位置设为-1e9):")
print(attn_scores[0])
# ============================================================================
# 6. Softmax归一化得到注意力权重
# ============================================================================
# 对每一行进行softmax,使得每个查询对所有键的注意力权重和为1
# 使用数值稳定的softmax实现(减去最大值)
weights = np.exp(attn_scores - attn_scores.max(axis=-1, keepdims=True))
weights = weights / weights.sum(axis=-1, keepdims=True)
print("\n" + "=" * 70)
print("步骤3: Softmax归一化得到注意力权重")
print("=" * 70)
print("Weights shape:", weights.shape)
print("注意力权重矩阵(因果掩码后)=")
print(weights[0])
print("\n解读:")
print(" - 第0行: token 0 只能看到自己,权重为 [1.0, 0.0, 0.0]")
print(" - 第1行: token 1 可以看到 token 0和1,权重和为1")
print(" - 第2行: token 2 可以看到所有token,权重和为1")
# ============================================================================
# 7. 加权求和得到输出
# ============================================================================
# 使用注意力权重对值向量进行加权平均
out = weights @ V # (1,3,3) @ (1,3,2) = (1,3,2)
print("\n" + "=" * 70)
print("步骤4: 加权求和得到最终输出")
print("=" * 70)
print("Output shape:", out.shape)
print("Output (注意力输出) =\n", out[0])
print("\n说明:")
print(" - 每个输出向量是对应位置可见的所有值向量的加权平均")
print(" - 权重由注意力机制自动学习,反映了token之间的相关性")