哥伦比亚
38 min read Hero Image LLM / transformer

Happy-LLM 章节疑难点(1) Transformer 架构

通过 Happy-LLM 学习LLM的基础内容

第二章 Transformer

01 注意力机制

1. 最终value的意义

我在入门注意力机制,其中Q和K计算后是得到不同key的注意力分数,那使用注意力分数计算得到的最终的value,这个值是用来做什么的呢,能使用一段文本给我举例介绍吗

最终计算得到的 Value,是当前词在融合了整个句子的上下文语境后,得到的一个“全新版”的特征表示。 它不再是一个孤立的词,而是一个“懂周围环境”的词。为了更直观地理解,我们用一段文本来举例。

假设我们有两句话:

  1. “我 昨天 吃 了 一个 苹果”
  2. “我 昨天 买 了 一股 苹果”

在这个例子中,我们把目光聚焦在 “苹果” 这个词上。刚读取到“苹果”时,它的初始状态(原始向量)只包含了字典里的基础含义——可能是水果,也可能是科技公司。

现在,我们用注意力机制来更新“苹果”的特征: 第一步:计算注意力分数(Q 和 K 的作用)

  • 当模型处理第一句话时,“苹果”(Q)去环顾四周的词(K)。它发现自己和“吃”(K)的相关性极高,因此它们之间的注意力分数接近满分。
  • 在第二句话中,“苹果”(Q)环顾四周,发现和“吃”没关系了,但和“股”(K)以及“买”(K)的相关性极高,获得了很高的分数。

第二步:提取并混合内容(V 的作用,也就是你的问题所在) 注意力分数计算出来后,我们要用它来对所有词的 内容/实质含义(Value) 进行加权求和。

  • 在第一句话里: “苹果”的新表示 = 90%“吃”的 V + 10% 其他词的 V + 自身基础的 V。最终得到的 Value: “苹果”这个词吸收了“吃”的属性,它的特征向量发生偏移,牢牢地指向了“一种汁水丰富的物理水果”。

  • 在第二句话里: “苹果”的新表示 = 90%“股”的 V + 10% 其他词的 V + 自身基础的 V。最终得到的 Value: “苹果”这个词吸收了“金融、股票”的属性,它的特征向量发生偏移,牢牢地指向了“一家市值庞大的科技巨头”。

通过这套机制,大模型在阅读时就不会断章取义,而是拥有了真正的“全局视野”。

2. 使用向量进行计算

在Happy-LLM中,对于符号命名并不严谨,并且会有之前是简单概念后面就是公式的割裂感,导致理解会出现偏差。

深度学习引入了向量空间。这时候:在真实的注意力机制里,Query、Key 和 Value 全部都是高维向量(Vector),而不是单个的数字

可以这样来理解这个从标量到向量的升级: 在游戏开发或者构建 3D 场景时,一个物体(比如苹果)不可能只用一个数字来表示。

  • Key 向量 (K):相当于这个物体的“属性标签”。比如苹果的 KK 可能是 [甜度: 0.8, 颜色: 红色(RGB), 是否可食用: 1.0]
  • Query 向量 (Q):相当于玩家发出的“搜索条件”。比如玩家肚子饿了,系统发出一个 QQ[甜度需求: 0.9, 颜色无所谓, 是否可食用: 1.0]

1. 点积算的是什么? 我们将玩家的 QQ 向量和苹果的 KK 向量进行点积运算(Q⋅K)。点积计算的既包含了向量的模大小,也包含了它们在多维空间中的方向夹角(相似度)。点积结果越大,说明苹果的属性标签和玩家的需求越匹配(注意力分数高)。

2. 最后的 Value 是什么? 在算出了相似度(注意力分数)之后,我们要提取的最终 Value 也是一个高维向量 (V),是单独给出的,与 KK 向量的模值没有关系,可以把它想象成这个苹果在游戏里的完整数据。 所以,最终的计算过程是:

最终结果=(注意力分数apple×Vapple)+(注意力分数banana×Vbanana)+\text{最终结果} = (\text{注意力分数}_{\text{apple}} \times V_{\text{apple}}) + (\text{注意力分数}_{\text{banana}} \times V_{\text{banana}}) + \dots

如果 “fruit” 和 “apple” 的点积相似度是 0.9,和 “chair” 的相似度是 0.01。那么模型最终拿到手里的是 90%的苹果完整特征向量1%的椅子特征向量 的混合体。

3. 符号命名与具体计算

点积计算

这里的词向量矩阵k,为什么里面是V_apple,V_banana,V_chair,为什么用V开头,有什么联系吗;还有帮我展示线性代数计算的具体过程,

Happy-LLM 在符号命名上不够严谨 在标准的注意力机制(Attention)论文和绝大多数代码里:

  • K (Key) 矩阵是由各个词的 Key 向量组装而成的,通常应该写成 kapple,kbanana,kchairk_{apple}, k_{banana}, k_{chair}
  • V (Value) 矩阵是由各个词的 Value 向量组装而成的,通常写成 vapple,vbanana,vchairv_{apple}, v_{banana}, v_{chair}

在这里用字母 vv,纯粹是因为在数学和物理中,小写字母 vv 经常被用作 向量(Vector)的通用缩写。想表达的意思仅仅是“苹果的词向量”、“香蕉的词向量”,而不是注意力机制里的“Value”。

线性代数计算的具体过程 把公式 attention(Q,K,V)=softmax(qKT)Vattention(Q,K,V) = \text{softmax}(qK^T)V 拆解开,假设我们现在的词向量维度很小(比如只有 2 维),词库里只有 3 个词(苹果、香蕉、椅子)。

第一步:计算 qqKTK^T 的乘积(求相似度得分) 假设我们的 Query 向量 qq(代表 “fruit”):

q=[1.00.5]q = \begin{bmatrix} 1.0 & 0.5 \end{bmatrix}

这是一个 1×21 \times 2 的行向量。 我们的 Key 矩阵 KK 包含了 3 个词的 Key 向量:

K=[0.90.60.80.70.10.1]apple的Keybanana的Keychair的KeyK = \begin{bmatrix} 0.9 & 0.6 \\ 0.8 & 0.7 \\ 0.1 & 0.1 \end{bmatrix} \begin{array}{l} \leftarrow \text{apple的Key} \\ \leftarrow \text{banana的Key} \\ \leftarrow \text{chair的Key} \end{array}

这是一个 3×23 \times 2 的矩阵。 在矩阵乘法中,为了让 qq 能和 KK 里的每一行做点积,我们需要把 KK转置(Transpose)KTK^T(变成 2×32 \times 3 的矩阵):

KT=[0.90.80.10.60.70.1]K^T = \begin{bmatrix} 0.9 & 0.8 & 0.1 \\ 0.6 & 0.7 & 0.1 \end{bmatrix}

现在计算 x=qKTx = qK^T

x=[1.00.5][0.90.80.10.60.70.1]x = \begin{bmatrix} 1.0 & 0.5 \end{bmatrix} \begin{bmatrix} 0.9 & 0.8 & 0.1 \\ 0.6 & 0.7 & 0.1 \end{bmatrix}

计算结果:

  • 对 apple: (1.0×0.9)+(0.5×0.6)=1.2(1.0 \times 0.9) + (0.5 \times 0.6) = 1.2
  • 对 banana: (1.0×0.8)+(0.5×0.7)=1.15(1.0 \times 0.8) + (0.5 \times 0.7) = 1.15
  • 对 chair: (1.0×0.1)+(0.5×0.1)=0.15(1.0 \times 0.1) + (0.5 \times 0.1) = 0.15

得到原始得分向量 x=[1.21.150.15]x = \begin{bmatrix} 1.2 & 1.15 & 0.15 \end{bmatrix}。这其实就是 qq 向量在各个 kk 向量上投影长度总和。

第二步:Softmax 归一化(将得分转化为权重比例)

原始得分 xx 有大有小,不方便直接当比例用。Softmax 函数会把它们转换成相加等于 1 的概率分布:

softmax(xi)=exiexj\text{softmax}(x_i) = \frac{e^{x_i}}{\sum e^{x_j}}

x=[1.2,1.15,0.15]x = [1.2, 1.15, 0.15] 代入:

权重向量 w[0.450.430.12]\text{权重向量 } w \approx \begin{bmatrix} 0.45 & 0.43 & 0.12 \end{bmatrix}

“fruit” 把 45% 的注意力给了 apple,43% 给了 banana,只有 12% 给了毫无关系的 chair。

第三步:乘以 VV 矩阵(提取并混合特征) 现在,真正的 Value 矩阵 VV 登场了,假设这 3 个词对应的 Value 特征矩阵是:

V=[5.02.04.03.01.09.0]apple的Valuebanana的Valuechair的ValueV = \begin{bmatrix} 5.0 & 2.0 \\ 4.0 & 3.0 \\ 1.0 & 9.0 \end{bmatrix} \begin{array}{l} \leftarrow \text{apple的Value} \\ \leftarrow \text{banana的Value} \\ \leftarrow \text{chair的Value} \end{array}

这也是一个 3×23 \times 2 的矩阵。 最终的计算就是把权重向量 ww 乘上矩阵 VV

Final Result=[0.450.430.12][5.02.04.03.01.09.0]\text{Final Result} = \begin{bmatrix} 0.45 & 0.43 & 0.12 \end{bmatrix} \begin{bmatrix} 5.0 & 2.0 \\ 4.0 & 3.0 \\ 1.0 & 9.0 \end{bmatrix}

计算规则就是按比例将 VV 里的每一行加起来:

  • 最终向量的第一个维度:(0.45×5.0)+(0.43×4.0)+(0.12×1.0)=4.09(0.45 \times 5.0) + (0.43 \times 4.0) + (0.12 \times 1.0) = 4.09
  • 最终向量的第二个维度:(0.45×2.0)+(0.43×3.0)+(0.12×9.0)=3.27(0.45 \times 2.0) + (0.43 \times 3.0) + (0.12 \times 9.0) = 3.27

最终输出:[4.09​3.27​]

结果 [4.093.27]\begin{bmatrix} 4.09 & 3.27 \end{bmatrix} 就是融合当前语境后,“fruit” 这个词在向下一层传递的全新特征向量。

4. 缩放点积注意力

缩放点积注意力公式 这个公式被称为缩放点积注意力(Scaled Dot-Product Attention)。比基础公式,多了一个除以 dk\sqrt{d_k} 的操作。

1.高维度导致“方差爆炸”

假设我们的 Query 和 Key 向量里的数字,都是从一个均值为 0、方差为 1 的标准正态分布中随机抽取的。

点积的计算过程是按位相乘再相加:qkT=i=1dkqikiq \cdot k^T = \sum_{i=1}^{d_k} q_i k_i

  • 如果维度 dk=2d_k = 2,点积结果是由 2 个随机数相乘再相加,结果的波动范围(方差)大概是 2。
  • 如果维度 dk=512d_k = 512,点积结果是由 512 个随机数累加而成。根据概率论,累加的项越多,结果的方差就越大。此时点积结果的方差会变成 512。

直观理解: 维度越高,两个向量点积后产生的数值就越容易出现极端的“大数”或“极小数”。这就导致了不同词与词之间的初始注意力得分差距被拉得极其巨大。

2.Softmax 的“赢家通吃”效应

Softmax 的公式包含指数运算 exe^x ,指数函数的特点是:对数值的差异极其敏感。 假设我们有三个词的初始得分:

  • 低维度 (dk=2d_k=2) 时,得分差距较小,比如:[2, 0, -1] 经过 Softmax 后:[0.84, 0.11, 0.04] \leftarrow 注意力分配健康
  • 高维度 (dk=512d_k=512) 时,方差爆炸,得分差距被放大,比如:[50, 0, -20] e50e^{50} 是一个天文数字,而 e0=1e^0=1e20e^{-20} 趋近于 0。 经过 Softmax 后:[0.999999..., 0.000000..., 0.000000...] \leftarrow 赢家通吃

当维度很大时,Softmax 会把几乎 100% 的注意力全给了那个初始得分稍微高一点点的词,而其他词的注意力被压到了 0。

3.梯度消失

在反向传播更新神经网络时,我们需要计算 Softmax 的导数(梯度),Softmax 输出某个概率值 pp 时,它的导数大致与 p(1p)p(1 - p) 成正比。

  • p=0.5p = 0.5 时,导数最大,网络学得最快。
  • 当发生了“赢家通吃”,也就是 p1p \approx 1p0p \approx 0 时: 导数 1×(11)=0\approx 1 \times (1 - 1) = 0 或者 导数0×(10)=0导数 \approx 0 \times (1 - 0) = 0

结论: 如果不加干预,在高维度下,Softmax 的输出会牢牢卡在 1 和 0 边缘。此时梯度等于 0,无论怎么训练,参数都不会再更新了。

既然维度 dkd_k 让方差变成了 dkd_k,那么在数学上,只要我们把点积的结果除以 dk\sqrt{d_k},就能把方差重新拉回到 1。

5. attention 代码理解

'''注意力计算函数'''
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)
    # Softmax
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
        # 采样
     # 根据计算结果对value进行加权求和
    return torch.matmul(p_attn, value), p_attn

理解张量(Tensor)的维度(Shape)变化

我们假设当前的输入张量形状为 (Batch, Seq_Len, Dim),也就是 (批次大小, 句子长度, 词向量维度)。我来为你逐行拆解这段代码里的核心函数:

  1. 获取缩放因子:query.size(-1)
d_k = query.size(-1) 
  • 作用:获取张量的最后一个维度的大小。
  • 原理解析:在 PyTorch 中,-1 代表倒数第一个维度。由于 query 的形状是 (Batch, Seq_Len, Dim),所以 query.size(-1) 提取出来的就是词向量的维度(比如 512)。这就是 dkd_k
  1. 矩阵转置:key.transpose(-2, -1)
key.transpose(-2, -1)
  • 作用:将 key 张量的最后两个维度互换。
  • 原理解析:这是为了满足矩阵乘法的规则。
    • key 原本的形状是 (Batch, Seq_Len, Dim)
    • 我们只想要交换句子长度词向量维度,而不破坏前面的 Batch 结构。所以我们用 transpose(-2, -1),把形状变成了 (Batch, Dim, Seq_Len)。这对应了数学公式里的 KTK^T
  1. 批量矩阵乘法:torch.matmul(...)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
  • 作用:执行 Q×KTQ \times K^T,计算每个词对其他词的原始注意力得分,并除以 dk\sqrt{d_k} 完成缩放。
  • 原理解析torch.matmul 函数,它支持批量矩阵乘法(BMM)。它会忽略前面的 Batch 维度,只对后两个维度执行标准的矩阵乘法:
    • (Batch, Seq_Len, Dim) ×\times (Batch, Dim, Seq_Len) \rightarrow 结果的形状变成了 (Batch, Seq_Len, Seq_Len)
    • 这个 (Seq_Len, Seq_Len) 的方阵,就是注意力得分矩阵 比如句子有 10 个词,这就生成了一个 10×1010 \times 10 的网格,记录了每个词和其他 9 个词的相关性。
  1. 归一化分配:scores.softmax(dim=-1)
p_attn = scores.softmax(dim=-1)
  • 作用:在最后一个维度上执行 Softmax 操作,把原始得分变成概率权重。
  • 原理解析scores 的形状是 (Batch, Seq_Len, Seq_Len)dim=-1 意味着我们在每一行(也就是针对每一个具体的 Query 词)去横向比较它对所有 Key 词的得分,并把它们压缩成相加等于 1 的权重比例。
  1. 随机失活:dropout(p_attn)
if dropout is not None:
    p_attn = dropout(p_attn)
  • 作用:防止模型过拟合。
  • 原理解析:它会随机把注意力矩阵里的一些权重强制变成 0。这就像在学习时故意遮住一部分信息,强迫模型去关注更多的上下文,而不是只死死盯着某一个词。

02 自注意力机制

1. 理解自注意力

在真实的神经网络里,Q、K、V 到底是从哪里来的?

我们暂时抛开干巴巴的符号,用一个 “机器翻译” 的真实场景来拆解这段话。 假设我们正在用 Transformer 模型把英文 "Apple is delicious" 翻译成中文 "苹果很好吃"

1.什么是“两段序列”?

在翻译任务中,天然存在两段不同的文本序列:

  • 序列 A(历史信息/待生成信息):我们目前已经翻译出来的中文,比如 ["苹果", "很"]。它在 Transformer 中被称为 Decoder(解码器) 的输入。

  • 序列 B(编码信息/源信息):我们最初输入的英文原句,即 ["Apple", "is", "delicious"]。它在 Transformer 中被称为 Encoder(编码器) 的输出。

2.Q、K、V 是如何在这两段序列中分配的?

后半段话的核心:Q 来自于一个序列,K 与 V 来自于另一个序列。 当我们想预测 "很" 后面的下一个词时:

  • Query (Q) 来自序列 A(中文):模型当前的状态是 "很"。它发出了一个查询请求(Q):“我已经翻译了‘苹果很’,下一个词我该去英文原句里找哪个词来翻译?”
  • Key (K) 来自序列 B(英文):英文句子里的每个词 ["Apple", "is", "delicious"] 都举起自己的标签(K)。delicious 的标签可能写着“形容词、味道好”。
  • Value (V) 也来自序列 B(英文):英文句子里的每个词不仅有标签,还包含了它实际的语义内容(V)。

此时,中文的 "很"(Q)和英文的 "delicious"(K)一拍即合,注意力分数极高。于是,模型顺理成章地提取了 "delicious" 的内容(V),从而成功预测出下一个中文字是 "好吃"

3.怎么理解“很少存在额外的真值 Value”?

这句话是为了防止产生误解而加的补充说明,虽然表述得有些绕。 在数学公式 Attention(Q,K,V)Attention(Q, K, V) 里,Q、K、V 是三个独立的变量。很多初学者会误以为,在实际应用中我们需要找三段不同的文本来分别充当这三个角色。

作者是在强调:实际上我们只有两段物理文本(序列 A 和 序列 B)

  • 序列 A 单独提供 Q。
  • 序列 B 同时包揽了 K 和 V 的角色。

在底层代码中,序列 B 的同一个词向量(比如 "delicious"),分别乘以了两个不同的权重矩阵(WkW_kWvW_v),变身成了 K 向量和 V 向量。K 用来负责和 Q 算分数匹配,V 用来负责提供真实的特征内容。并没有凭空多出来一个“第三序列”去充当 Value。

4.为什么要用自注意力

为什么要在 Encoder(编码器)里不厌其烦地让句子里的每一个词都去和其他所有词算一遍相关性

简单用一句话概括:为了让每个原本孤立的词,进化成一个“完全懂这句话上下文”的词。

自然语言最复杂的地方在于,词是没有绝对固定含义的,它的意思完全由周围的词决定。让每一个 token 互相计算相关性,主要是为了解决以下三个核心痛点:

  • 1.解决致命的“代词指代”问题(指代消解) 语言中充满了简写和代词。我们来看深度学习中极其著名的一个测试用例:
  • 句子 A: “这只动物没有穿过那条街道,因为太累了。”
  • 句子 B: “这只动物没有穿过那条街道,因为太宽了。”

对于计算机来说,“它”只是一个毫无感情的代词 token。如果不看全局,模型根本不知道“它”是谁。

有了自注意力机制:

  • 句子 A 中,当“它”作为 QQ 去寻觅全句的 KK 时,由于句尾有“累了”这个语境,“它”会和“动物”产生极高的共鸣(注意力分数极高)。混合 VV 之后,“它”的向量特征就变成了“疲惫的动物”。
  • 句子 B 中,由于句尾有“宽了”,“它”的 QQ 会和“街道”的 KK 产生极高分数。混合后,“它”的向量特征就变成了“宽阔的街道”。

通过让全句所有的词互相“握手”打分,模型完美破解了语言中的歧义。

  • 2.重塑多义词的灵魂(上下文融合)

假设句子是:“我今天去苹果店买了一部手机,顺便吃了一个苹果。” 在输入给 Encoder 之前,这两个“苹果”在底层词典里的初始词向量是一模一样的。这显然不合理。 通过自注意力机制:

  • 第一个“苹果”会敏锐地和“店”、“买”、“手机”产生高昂的注意力分数,它最终输出的 VV 会变成一个充满科技属性的实体向量。
  • 第二个“苹果”会和“吃”产生极高的关联,它的特征会偏移向水果属性。

编码器的目的,就是要把原始文本“翻译”成机器能懂的高级特征矩阵。自注意力让输出的每一个向量都不再是死板的字典词,而是动态的、定制化的

  • 3.全局视野 在 Transformer 诞生之前,大家用的都是 RNN(循环神经网络)或 LSTM。这类模型像人类读书一样,从左到右,一个词一个词地读。 这带来一个致命弱点:如果句子很长,当它读到第 100 个词时,早就把第 1 个词的细节忘光了(长距离依赖衰减)。

在自注意力的矩阵运算中,不管词汇在句首还是句尾,它们之间的“物理距离”永远是 O(1)O(1)。第 1 个词和第 1000 个词可以直接进行点积 QKTQ \cdot K^T 对话,没有任何信息损耗。这使得大语言模型拥有了极其恐怖的上下文统筹能力。

最后一句话提到了“给 Q、K、V 的输入传入同一个参数”,这正是 Self-Attention 的精髓所在。 在之前讲的交叉注意力中,QQ 来自中文句子,KKVV 来自英文句子。

但在编码器中,没有外人,只有当前这一句话。 假设输入文本经过最初步的转化后,变成了一个张量矩阵 XX。模型会使用这同一个张量 XX,去分别乘以三个不同的线性权重矩阵(WQ,WK,WVW^Q, W^K, W^V),从而自己分裂出自己的 Q、K、V: Q=XWQQ = X \cdot W^Q K=XWKK = X \cdot W^K V=XWVV = X \cdot W^V

这就像是一个人(输入序列 XX)站在三面不同角度的镜子前,照出了自己的三个分身。然后这三个分身互相打分、互相融合,最终完成了一次对自我的深度审视和升华。

2. 掩码自注意力

第一,掩码注意力(Masked Attention)是用在解码器里面吗? 是的。截图里提到的这种“上三角掩码(Causal Mask)”是解码器的标志性特征。

第二,它是用于计算编码器送来的 K、V 与解码器自身的 Q 吗?

掩码代码,是用在解码器的掩码自注意力层(Masked Self-Attention),此时的 Q、K、V 全部来自解码器自身

而“Q来自解码器,K、V来自编码器”的层叫做交叉注意力层(Cross-Attention),在这个层里,解码器是可以纵览编码器所有内容的,通常不需要加这种遮蔽未来的上三角掩码。

假设我们在做机器翻译,要把 “I love you” 翻译成 “我 爱 你”。 在训练阶段,我们其实是把整句正确答案 ["我", "爱", "你"] 一次性喂给解码器的。

当我们训练模型预测 “我” 的下一个词是 “爱” 时,模型面前其实已经摆着完整的 ["我", "爱", "你"] 序列了。如果不加限制,第 1 个词 “我” 的 QQ 去和全句的 KK 算相似度时,它会偷偷“偷看”到后面的答案,这就变成了作弊

为了防止作弊,我们人为引入了一个掩码矩阵(Mask)。规则很简单:每个词只能看到它自己和它之前的词,绝对不能看到它后面的词。

Mask 矩阵维度一般为 (1, seq_len, seq_len)

scores = scores + mask[:, :seqlen, :seqlen]
  • scores:这是原始的注意力打分矩阵(形状是 seqlen x seqlen),也就是 QKTQ \cdot K^T 的结果。
  • mask[:, :seqlen, :seqlen]:这是一个预先生成好的矩阵。在这个矩阵中,下三角区域(包括对角线)全为 0,而上三角区域全是负无穷(-inf,在代码里通常用一个极小的负数比如 -1e9 表示)
  • 核心逻辑:当 scoresmask 相加时:
    • 对于过去和现在的词(下三角):原始分数 +0=+ 0 = 分数不变
    • 对于未来的词(上三角):原始分数 +()=+ (-\infty) = -\infty
scores = F.softmax(scores.float(), dim=-1).type_as(xq)

这短短一行包含了三个非常专业的工程操作:

  1. scores.float() (精度提升):在训练大模型时,为了省显存,输入 xq 通常是半精度浮点数(FP16 或 BF16)。但在做 Softmax 这种包含指数运算的操作时,半精度极其容易溢出(变成 NaN)。所以这里先强制把数据转成单精度(FP32)以保证计算的绝对稳定。
  2. F.softmax(..., dim=-1) (核心魔法): 由于我们在上一步把“未来词”的分数都变成了 -\infty。根据 Softmax 的公式包含 exe^xe=0e^{-\infty} = 0 这就是掩码最精妙的地方!经过 Softmax 后,所有本来试图偷看“未来词”的注意力权重,被物理意义上硬生生地变成了 0。 模型把 100% 的注意力被迫全部分配给了当前词和过去的词。
  3. .type_as(xq) (精度还原):Softmax 算完(且安全了)之后,把 FP32 的数据重新降级回半精度(比如 FP16),和最原始的输入 xq 保持一致,继续往下传给后续的神经网络,深藏功与名。

在 Decoder(解码器)中,数据的真实流动是这样的两段式

  1. 第一关:掩码自注意力(Masked Self-Attention)
    • Q、K、V 全是解码器目前生成的中文序列。
    • 必须使用 -inf 掩码代码,严格防止偷看后面的中文。
  2. 第二关:交叉注意力(Cross-Attention)
    • Q 是刚才第一关输出的中文特征,K 和 V 是编码器送来的英文全句特征。
    • 不用上三角掩码。因为英文原句早就写死了,不管当前翻译到哪个中文字,它都有权力去看整句英文上下文。

掩码自注意力的K,V算出来是做什么的

在掩码自注意力中,KKVV 被计算出来,并且当场就被“消耗”掉了

在真实的解码器块(Decoder Block)里,数据是像接力赛一样流转的。我们把解码器里的这两个注意力层,比作两次会议:

掩码自注意力

  • 参与者: 只有解码器自己生成的中文词(比如 ["<BOS>", "我", "爱"])。
  • 生成 Q、K、V: 这里的中文词向量 XX,克隆出 QselfQ_{self}KselfK_{self}VselfV_{self}

在这个层里,它们当场就进行了完整的注意力计算

“爱” 的 QselfQ_{self},去和过去词的 KselfK_{self}(“<BOS>”, “我”, “爱”)算分数,然后混合它们对应的 VselfV_{self}

  • 第一场会议的输出: 得到一个全新升级的 XX'。这个 XX' 里的“爱”,已经完美融合了“我”的语境。

  • K 和 V 的结局: 第一场会议结束,这些 KselfK_{self}VselfV_{self} 已经完成了它们的使命(被合成了 XX'

交叉注意力

现在,带着升级版的 XX',解码器来到了第二关,要去看编码器翻译的英文原句了。

  • Q 的来源: 刚才第一场会议输出的升级版中文特征 XX',化身成了这里的 QcrossQ_{cross}。(代表:我现在的中文语境是这样的,我该找哪个英文词?)
  • K、V 的来源: 编码器送来的英文原句(比如 ["I", "love"]),分裂成了这里的 KcrossK_{cross}VcrossV_{cross}
  • 计算: 中文的 QcrossQ_{cross} 和英文的 KcrossK_{cross} 算分数,提取英文的 VcrossV_{cross}

工程现实中的 KV Cache

虽然在数学推导上,第一场会议的 KselfK_{self}VselfV_{self} 算完就被丢弃了,但在真实的代码推理(Inference)阶段,系统会把它们悄悄存进显存里(这就是大名鼎鼎的 KV Cache)。

为什么?

因为当模型下一轮要预测“你”这个词时,前置序列变成了 ["<BOS>", "我", "爱", "你"]

如果没有 Cache,模型为了计算“你”的分数,又要让“我”和“爱”重新生成一遍它们的 KselfK_{self}VselfV_{self}。这太浪费算力了!

有了 KV Cache,模型只需要计算最新词“你”的 QselfQ_{self},然后直接去显存里把之前算好存起来的 KselfK_{self}VselfV_{self} 调出来,拼在一起算点积就可以了。

为了让你彻底看清 Q、K、V 在解码器这两层中是如何“接力”的,我为你做了一个解码器双层注意力数据流模拟器。你可以一步步点击,看看张量是如何变形并被消耗掉的:

总结一下:

掩码自注意力里的 K 和 V,不仅算了,而且是大算特算!它们的唯一作用,就是在那一层里把孤立的词汇揉捏成带有历史语境的全新特征(XX',然后功成身退,把接力棒(XX')交给下一层。

3. 解码器结构

整个解码器(Decoder)数据流的物理起点是哪里?

在数据进入那套复杂的 Q、K、V 矩阵乘法之前,必须有一个最原始的输入张量 XX。这个 XX 并不是凭空产生的,它是从我们人类的文本经过一系列物理转换后得到的。

依然用把 “I love you” 翻译成 “我 爱 你” 为例。

这个原始的 XX 的诞生,主要经历了三个极其关键的预处理步骤

第一步:获取当前的文本序列

解码器的输入 XX,本质上就是“目前已经翻译出来的结果”。 但是这里有一个问题:当模型刚开始翻译,连第一个字都还没吐出来的时候,输入是什么?

为了启动这个过程,深度学习里引入了一个特殊的占位符,通常叫 <BOS> (Begin Of Sentence,句子开头) 或者 [START]

  • 第一轮生成: 此时解码器手里只有 ["<BOS>"]
  • 第二轮生成: 假设上一轮模型成功预测了“我”,此时解码器手里的序列变成了 ["<BOS>", "我"]
  • 第三轮生成: 变成了 ["<BOS>", "我", "爱"]。 (每一轮都通过掩码自注意力+交叉注意力) 这个不断变长的中文序列,就是 XX 的数据源头。

第二步:查字典(Token Embedding 词向量嵌入)

计算机不认识 <BOS> 也不认识“我”,它只认识数字。所以,我们要把文本序列变成向量。

模型内部自带一本极其庞大的“字典”(Embedding Table)。字典里有几万个词,每个词都对应着一个长度为 Dim(比如 512 维)的浮点数向量。

  • 序列 ["<BOS>", "我", "爱"] 被送进字典。
  • 查表后,它变成了一个形状为 (3, 512) 的矩阵。
    • 第 1 行是 <BOS> 的专属 512 维向量。
    • 第 2 行是“我”的专属 512 维向量。
    • 第 3 行是“爱”的专属 512 维向量。

这就是 词嵌入(Word Embedding)

第三步:贴上时间戳(Positional Encoding 位置编码)

如果只做第二步,会有个致命缺陷:自注意力机制是同时处理所有词的,它根本没有“先后顺序”的概念。在它眼里,“我爱你”和“你爱我”计算出来的结果是一模一样的。

为了让模型知道词的先后顺序,研究人员发明了位置编码(Positional Encoding)

这也是一个形状为 (3, 512) 的矩阵,但它里面的数字和词意无关,纯粹是代表物理位置的数学公式(通常是正弦和余弦函数)。

  • 第 1 行代表“绝对位置 1”的信号。
  • 第 2 行代表“绝对位置 2”的信号。
  • 第 3 行代表“绝对位置 3”的信号。

第四步:张量 XX 诞生

最后一步:把词嵌入矩阵和位置编码矩阵,直接加起来 X=Word Embedding+Positional EncodingX = \text{Word Embedding} + \text{Positional Encoding} 现在,我们得到了一个形状依然是 (3, 512) 的矩阵 XX。 这个 XX 里的每一行(每一个词向量),既包含了词汇本身的含义(我是谁),又包含了它在句子里的顺序(我在哪)。

它会被再次正式送入解码器的第一层(掩码自注意力层)。然后,这个新的 XX 会克隆出三份,分别去乘以 WQW^QWKW^KWVW^V 矩阵,裂变成属于它自己的 Q、K、V,开始执行我们之前自注意力打分机制(后面可能生成终结符 <EOS>)

4. 多头注意力机制

1). 为什么使用多头注意力

你的前提完全正确:单次自注意力机制,确实已经让一个词和全句所有的词都进行了计算。 既然每个词都看遍了全句,为什么还要搞个“多头(Multi-Head)”呢?

这里的核心盲区在于:“看遍了所有词” 不等于 “看懂了所有维度的关系”。

用一个“专家会诊”的比喻,来彻底拆解它。

1 直观比喻:为什么一个“头”不够用?

假设我们有一句话:“小明 笑着 咬了 一口 甜甜的 苹果”

单头注意力(Single-Head Attention 中,就只有一个专家在审视这句话。这个专家虽然看遍了全句所有的词,但他可能是一个 “语法专家”

  • 他的眼中只有主谓宾关系。当他看到“咬了”(Query)时,他的注意力绝大部分都分配给了“小明”(主语)和“苹果”(宾语)。
  • 他完美拟合了动作关系,但他完全忽略了情感和状态

而在多头注意力(Multi-Head Attention) 中,Transformer 雇佣了 8 个甚至 96 个不同领域的专家(Head)同时来看这句话:

  • 头 1(语法专家): 依然把“咬了”的注意力分配给“小明”和“苹果”。
  • 头 2(情感专家): 当看到“小明”时,把极高的注意力分配给了“笑着”。
  • 头 3(属性专家): 当看到“苹果”时,死死地盯住了“甜甜的”。
  • 头 4(指代专家): 专门寻找代词和实体的联系(虽然这句没有代词)。

结论: 每一个“头”确实都在做全局扫描(和所有词算分数),但因为它们各自掌管的参数权重矩阵(WQ,WK,WVW^Q, W^K, W^V)不同,它们**“眼中在乎的东西”完全不一样**。最后把这 8 个专家的意见拼在一起,大模型才能全方位无死角地理解这句话。

2). 如何实现多头注意力

在代码和张量(Tensor)层面,多头到底是怎么实现的?

假设我们的 hidden_size = 512,并且我们设置了 num_heads = 8(8 个头)。

如果不使用多头: 模型用完整的 512 维特征去算一次庞大的点积,得出一张 (seq_len, seq_len) 的注意力分数网格。由于 512 维揉在一起算总分,很多微妙的、局部的特征关系(比如“甜甜的”)很容易被占主导地位的宏观特征(比如主谓宾)给平均掉、掩盖掉

使用多头的精妙物理操作: Transformer 在拿到 512 维的 Q,K,VQ, K, V 张量后,做了一个极其聪明的动作——切片(Split) 它在 hidden_size 这个维度上,把 512 平均切成了 8 份!

  • 512÷8=64512 \div 8 = 64
  • 现在,每一个头(Head)只分到了 64 维的特征。

张量的维度在这个瞬间发生了华丽的变形: 原本的 (batch_size, seq_len, 512)reshape 成了 \rightarrow (batch_size, 8个头, seq_len, 64) 接下来,这 8 个头拿着自己专属的 64 维特征,互不干扰地、并行地去算自己的注意力分数。

  • 头 1 用 0~63 维算出一张 (seq_len, seq_len) 的关系网。
  • 头 2 用 64~127 维算出另一张完全不同的 (seq_len, seq_len) 关系网。 …
  • 一共算出了 8 张不同的关系网格!

最后一步:拼接(Concatenate) 等这 8 个头各自算完自己那 64 维特征融合后的全新 VV 后,Transformer 再把这 8 个 64 维的输出原封不动地拼接(Concat)回去64×8=51264 \times 8 = 512

当年提出 Transformer 的论文《Attention Is All You Need》中,关于多头注意力的原理解释是这样的:

为了让 8 个头关注不同的东西,我们为这 8 个头各自准备了独立且不同的小型权重矩阵(W1Q,W2QW8QW_1^Q, W_2^Q \dots W_8^Q),每个小矩阵的形状是 (512, 64)。初始词向量 XX 分别去乘这 8 个小矩阵,得到 8 个不同的、64 维的小 QQ

但这样写因为这就意味着你要写一个 for 循环,让 GPU 算 8 次小矩阵乘法。GPU 最讨厌零碎的小计算,它最擅长的是一次性算一个无比巨大的矩阵

所以要使用了一个极其优雅的数学等价替换:

  1. 拼装大矩阵: 把那 8 个 (512, 64) 的小权重矩阵,在内存里横向拼成一个巨大无比的权重矩阵 WQW^Q,它的形状变成了 (512, 512)
  2. 让初始词向量 XX 直接和这个巨大的 WQW^Q 做一次矩阵乘法。 Qbig=XWQQ_{\text{big}} = X \cdot W^Q 此时,我们得到了一整个 512 维的 QQ
  3. 切片分配(Split/Reshape): 然后,我们在得到这 512 维的 QQ 之后,拿刀把它均分成 8 份,每份 64 维。

在真正的代码里,这个“切片”操作只需要改变张量的“看法(View)”和“维度转置(Transpose)”。

数据流是这样的:

  1. 生成大 Q: 形状为 (batch_size, seq_len, 512)
  2. View变形(切片): 告诉系统,不要把它当成 512 维,把它当成 8 个 64 维。 形状变成了 (batch_size, seq_len, 8, 64)
  3. Transpose转置(为多头并行做准备): 把“头”的维度换到前面去。 形状最终变成 (batch_size, 8, seq_len, 64)

现在,GPU 会把这个张量看作是 8 个互不干扰的独立宇宙。在接下来的注意力分数计算 QKTQ \cdot K^T 中,这 8 个头会在底层的 CUDA 核心里同时进行点积计算。

在标准的 Transformer 架构(包括现在的 GPT、Gemini 等所有大语言模型)中,只要提到注意力机制,默认全部都是多头(Multi-Head)的

解码器(Decoder)里面其实包含了两种注意力层,它们全都是多头体制:

  1. 多头掩码自注意力(Masked Multi-Head Self-Attention): 也就是你问的这个。用于审视已经生成的文本,严格防止偷看未来。
  2. 多头交叉注意力(Multi-Head Cross-Attention): 用于在生成当前词时,去原文本(编码器输出)里寻找线索。

为什么掩码自注意力也必须是多头? 因为解码器在预测下一个词时,同样面临极其复杂的语言环境。

假设已经生成的中文是:“我 笑着 咬了 一口 甜甜的”,现在要预测下一个词(苹果)。 解码器同样需要请出所有的“专家(Heads)”来审视这句半成品:

  • 头 1(动作专家): 盯住“咬了”,说明下一个词大概率是食物。
  • 头 2(属性专家): 盯住“甜甜的”,进一步缩小范围。
  • 头 3(情感专家): 盯住“笑着”,奠定正向的情感基调。

所以,多头机制对于解码器来说同样是不可或缺的灵魂。

跑一遍这段底层代码的维度变化。这会彻底打通这两者的关系 假设 batch_size = 32seq_len = 5num_heads = 8,总维度 hidden_size = 512(切分后每个头 64 维)。

第一步:计算多头分数(Scores) 经过我们上一轮聊过的切片和转置后,QQKK 的维度变成了 (32, 8, 5, 64)。 它们做点积 QKTQ \cdot K^T(在最后两个维度上运算)后,64 维被消耗掉了,分数矩阵的维度变成了 Scores(32,8,5,5)Scores \rightarrow (32, 8, 5, 5)

(翻译:32 句话,8 个独立的专家头,每个头都算出了一张 5x5 的关系网格)

第二步:掩码(Mask)

现在高潮来了,我们的 Mask 矩阵应该是什么维度?需要给 8 个头分别准备不同的 Mask 吗?

不需要。

无论你是“语法专家”还是“情感专家”,都绝对不能偷看未来的词。防作弊规则对所有的头一视同仁。

所以,我们依然只生成那个最基础的、上三角全是 -inf5×55 \times 5 网格。为了利用 PyTorch 的广播(Broadcasting)机制,我们会把它的维度定义为: Mask(1,1,5,5)Mask \rightarrow (1, 1, 5, 5) 第三步:完美叠加(相加)

scores = scores + mask

在执行这行代码的瞬间,PyTorch 极其聪明地把这个 (1, 1, 5, 5) 的 Mask:

  1. batch 维度上复制了 32 份。
  2. head 维度上复制了 8 份。 瞬间覆盖到了所有批次、所有专家的每一张关系网格上!随后大家一起平平安安地去做 Softmax。

3). 理解公式

MultiHead(Q,K,V)=Concat(head1,...,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

这个公式用极其精炼的数学语言,把整个多头注意力的流水线严丝合缝地描述了出来。

为了彻底吃透它,我们应该从下往上看,也就是先看下面那行局部的 head 是怎么算的,再看上面那行整体是怎么拼接的。

第一步:局部计算 headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) 这行公式解释了第 ii 个“专家头”是如何独立工作的。

这里的 Q,K,VQ, K, V 其实指的是初始的输入张量(比如编码器输出的 512 维特征序列)。

  • WiQ,WiK,WiVW_i^Q, W_i^K, W_i^V 是什么? 它们就是我们上文提到的,属于第 ii 个专家的专属小型投影矩阵(在理论层面,它们的形状是 512 x 64)。
  • QWiQQW_i^Q 在干什么? 这是让初始的 512 维输入,去乘以这个专家的专属矩阵。这就好比让输入数据戴上了这位专家特制的“有色眼镜”。经过乘法后,原本 512 维的庞杂特征,被提炼、压缩成了 64 维的专属特征。
  • Attention(...)\text{Attention}(...) 在干什么? 当这三个 64 维的专属小 Q、小 K、小 V 准备好之后,它们就去执行我们最熟悉的那套“点积 \rightarrow 缩放 \rightarrow Softmax \rightarrow 乘以 V”的标准单头注意力流程。

结论: 经过这一行公式,第 ii 个头算出了一个融合了上下文的 64 维输出特征。如果有 8 个头,我们就得到了 8 份不同的 64 维特征。

第二步:全局组装

MultiHead(Q,K,V)=Concat(head1,...,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O

这行公式是整个多头机制的“收尾点睛之笔”,它分为两个至关重要的动作:

动作一:Concat\text{Concat}(无缝拼接)

把刚才那 8 个独立算出来的 64 维特征,像拼积木一样横向拼接起来。

64+64+64+(拼8次)=51264 + 64 + 64 + \dots \text{(拼8次)} = 512

此时,张量的维度重新胖回了原来的 512 维。但注意,这时的 512 维其实是 8 个互不相干的“知识区块”硬凑在一起的。

动作二:WOW^O(终极融合矩阵)

这是很多初学者最容易忽略,但工程上绝对不可缺少的一个矩阵!

WOW^O (Output Weight Matrix)是一个巨大的、形状为 512 x 512 的权重矩阵。

为什么必须要有 WOW^O

如果不乘 WOW^O,那拼接出来的 512 维特征就只是一个“大杂烩”。头 1(语法专家)和头 2(情感专家)的特征虽然靠在了一起,但它们彼此之间并没有交流。

将拼接后的特征再乘以 WOW^O,就相当于召开了一次总结大会WOW^O 这个“大老板”会将这 8 个专家的意见进行跨维度的加权混合、重新洗牌。经过这次终极运算后,输出的 512 维张量才是一个真正融会贯通、浑然一体的全新特征表示。

那K,V需要分成八个吗?

“既然我已经派出了 8 个不同的专家(8 个 QQ)去提问了,那让他们去查阅同一个公共的资料库(1 个 KK 和 1 个 VV)不就行了吗?为什么连资料库也要切成 8 份?”

直接回答你的结论:不行。KKVV 必须、且绝对要跟着分成 8 份(或者说通过 8 组不同的权重矩阵计算出来)。

如果不把 KKVV 分开,多头注意力机制就会彻底失效。我们可以从“业务逻辑”和“数学规则”两个极其硬核的维度来彻底解开这个谜团。

业务逻辑层:不同的专家,需要不同的“索引”和“内容”

还记得我们说的“专家会诊”比喻吗?多头注意力的核心,是让不同的头在不同的特征子空间(Sub-spaces) 工作。

假设输入词是“苹果”。

  • QQ (提问):专家在找什么。
  • KK (键/标签):词汇挂在身上的特征标签(用来和 Q 匹配)。
  • VV (值/实质内容):词汇实际包含的语义信息。

如果只分 8 个 QQ,但不分 KKVV(让大家共享 512 维的 KKVV):

  • 语法专家(头1的 QQ 来找主语。他看着那个庞杂的 512 维公共 KK(里面既有语法信息,又有颜色、味道、金融信息),他会被干扰,很难精准匹配。
  • 当他好不容易匹配上了,他提取出的公共 VV 是一个大杂烩。他明明只想提取“名词属性”,却被迫把“红彤彤、甜甜的、乔布斯”这些多余的信息全带走了。

这就是为什么要切分出 8 个 KK 和 8 个 VV

  • 输入 XX 乘以 W1KW_1^K,生成了头 1 专属的 K1K_1。这相当于给“苹果”贴上了纯粹的**“语法标签”**。
  • 输入 XX 乘以 W1VW_1^V,生成了头 1 专属的 V1V_1。这相当于把“苹果”的内容提纯成了纯粹的**“名词实体信息”**。
  • 现在,语法专家的 Q1Q_1 只和语法标签 K1K_1 匹配,提取纯粹的语法内容 V1V_1。整个链路干净、纯粹、互不干扰!

同理,**情感专家(头2)**拥有自己专属的 K2K_2(情感标签)和 V2V_2(情感内容)。这样,同一个词“苹果”,在 8 个头眼里,展现出了 8 种完全不同的侧面。这才是“全面深入拟合语言信息”的真谛。

如果你尝试在代码里写出“8 个 QQ,共享 1 个 KK”的代码,PyTorch 会在运行的第一秒就给你报红报错。

为什么?因为注意力机制的核心运算是点积(Dot Product)

Scores=QKT\text{Scores} = Q \cdot K^T

  • 如果你把 QQ 分成了 8 份,那么每一个小 QiQ_i 的维度就变成了 64 维

  • 如果你不分 KK,保留公共的 KK,那它的维度依然是 512 维

  • **数学法则规定:一个 64 维的向量,是无法和一个 512 维的向量做点积的

为了能和 64 维的 QiQ_i 算分数,KK 必须也被降维(切分)到 64 维,变成 KiK_i。 算完分数后,这个分数要乘以 VV 来混合特征。为了保证最后 8 个头拼接起来(64×864 \times 8)能刚好还原成 512 维,每一个小 ViV_i 也必须是 64 维。

工程代码是怎么写的?

在真实的 Transformer 代码底层,Q,K,VQ, K, V 的生成待遇是完全平等、完全一致的。

  1. 代码会同时初始化三个巨大的权重矩阵:WbigQ,WbigK,WbigVW_{big}^Q, W_{big}^K, W_{big}^V(形状都是 512 x 512)。
  2. 输入 XX 分别和这三个大矩阵相乘:
    • Qbig=XWbigQQ_{big} = X \cdot W_{big}^Q
    • Kbig=XWbigKK_{big} = X \cdot W_{big}^K
    • Vbig=XWbigVV_{big} = X \cdot W_{big}^V
  3. 然后,代码会极其公平地,把 Qbig,Kbig,VbigQ_{big}, K_{big}, V_{big} 同时切成 8 份。
    • 第 1 份的 Q1,K1,V1Q_1, K_1, V_1 组成第一套头 1。
    • 第 2 份的 Q2,K2,V2Q_2, K_2, V_2 组成第二套头 2。
    • 以此类推……

总结一下疑惑: 不能只让“提问者(Q)”分身,必须让“资料的索引(K)”和“资料的内容(V)”也跟着分身。只有这样,模型才能在不同的特征子空间里,进行互不干扰、极其纯粹的专业计算。

4). 代码理解

import torch.nn as nn
import torch

'''多头自注意力计算模块'''
class MultiHeadAttention(nn.Module):

    def __init__(self, args: ModelArgs, is_causal=False):
        # 构造函数
        # args: 配置对象
        super().__init__()
        # 隐藏层维度必须是头数的整数倍,因为后面我们会将输入拆成头数个矩阵
        assert args.dim % args.n_heads == 0
        # 每个头的维度,等于模型维度除以头的总数。
        self.head_dim = args.dim // args.n_heads
        self.n_heads = args.n_heads

        # Wq, Wk, Wv 参数矩阵,每个参数矩阵为 n_embd x dim
        # 这里通过三个组合矩阵来代替了n个参数矩阵的组合,其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积,
        # 不理解的读者可以自行模拟一下,每一个线性层其实相当于n个参数矩阵的拼接
        self.wq = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
        # 输出权重矩阵,维度为 dim x dim(head_dim = dim / n_heads)
        self.wo = nn.Linear(self.n_heads * self.head_dim, args.dim, bias=False)
        # 注意力的 dropout
        self.attn_dropout = nn.Dropout(args.dropout)
        # 残差连接的 dropout
        self.resid_dropout = nn.Dropout(args.dropout)
        self.is_causal = is_causal

        # 创建一个上三角矩阵,用于遮蔽未来信息
        # 注意,因为是多头注意力,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, dim) -> (B, T, dim)
        xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)

        # 将 Q、K、V 拆分成多头,维度为 (B, T, n_head, dim // n_head),然后交换维度,变成 (B, n_head, T, dim // n_head)
        # 因为在注意力计算中我们是取了后两个维度参与计算
        # 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度而不是直接按注意力输入展开,是因为view的展开方式是直接把输入全部排开,
        # 然后按要求构造,可以发现只有上述操作能够实现我们将每个头对应部分取出来的目标
        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_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)
        # 做 Dropout
        scores = 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, dim // n_head),再拼接成 (B, T, n_head * dim // 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

为了方便理解维度的变化,我们在脑海里预设一组真实大模型的典型参数:

  • batch_size (批次大小,bsz) = 32
  • seq_len (序列长度,T) = 10
  • n_embd (输入词向量维度) = 512
  • dim (注意力内部总维度) = 512
  • n_heads (头数,nh) = 8
  • head_dim (每个头的维度,hs) = 512 / 8 = 64

这部分在模型刚被创建时执行一次,用于在显卡里分配好所有的权重矩阵。

def __init__(self, args: ModelArgs, is_causal=False):
  • args: 这是一个配置对象(通常是一个数据类 Dataclass),里面打包了模型的所有超参数(比如维度、头数)。
  • is_causal: 这是一个布尔值开关。如果为 True,代表这是一个解码器(Decoder),需要加掩码防止偷看未来;如果为 False,代表是编码器,无需掩码。
super().__init__()
assert args.dim % args.n_heads == 0
  • super().__init__(): 继承 PyTorch nn.Module 的优良传统,初始化父类。
  • assert ...: 断言:总维度 dim 必须能被头数 n_heads 完美整除。除不尽的话程序直接在这里报错拦截。
self.head_dim = args.dim // args.n_heads
self.n_heads = args.n_heads
  • 计算并保存每个头专属的维度(head_dim = 64),并保存总头数(8)。
self.wq = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
self.wv = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)

nn.Linear(in_features, out_features, bias=False) 时,PyTorch 随机初始化并在显存里分配了一个形状为 (out_features, in_features)参数矩阵(权重矩阵 Tensor)

  • 输入维度 args.n_embd(512): 这是来自上一层的、最原始的词向量 XX。既然 XX 是 512 维的,要想让它发生矩阵乘法,权重矩阵的输入端必须也是 512 宽。这就保证了特征的完整流入。

  • 输出维度 self.n_heads * self.head_dim(8 × 64 = 512): 在上一行代码 assert args.dim % args.n_heads == 0 中,其实已经规定了总维度必须能被头数整除。所以 n_heads * head_dim 在数值上其实就等于 args.n_embd(也是 512)。

  • 核心动作:创建 Q, K, V 的三个权重矩阵。

  • 输入 args.n_embd:上游传过来的词向量维度(512)。

  • 输出 self.n_heads * self.head_dim:其实就是 512。这三个矩阵的形状都是 (512, 512)

  • bias=False:不要偏置项(不要那个 +b+ b)。在现代大模型(如 LLaMA)中,为了配合 RMSNorm 等归一化技术,通常会舍弃偏置项以加速训练并减少显存占用。

self.wo = nn.Linear(self.n_heads * self.head_dim, args.dim, bias=False)
  • 融合矩阵 WOW^O:负责把多头拼接后的 512 维特征,再次映射混淆,输出最终的 512 维特征。
self.attn_dropout = nn.Dropout(args.dropout)
self.resid_dropout = nn.Dropout(args.dropout)
self.is_causal = is_causal
  • Dropout (随机失活):训练时随机把一些神经元变成 0,防止死记硬背。attn_dropout 用在注意力分数上,resid_dropout 用在最终输出上。
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)
  • 构建掩码:只有当 is_causal=True 时才执行。即只有是解码器的时候
  • torch.full(...):生成一个形状为 (1, 1, 最大长度, 最大长度) 的全矩阵,里面全填满负无穷 -inf
  • torch.triu(..., diagonal=1):把这个矩阵的上三角区域(不含主对角线)保留下来(依然是 -inf),其他地方全变成 0
  • self.register_buffer (技巧):在 PyTorch 里,如果直接写 self.mask = mask,系统不会把它当成模型的一部分。用 register_buffer 是告诉 PyTorch:“这是一个固定法则(不参与梯度求导更新),但在保存模型参数(state_dict)或者把模型推入 GPU (.cuda()) 时,请把它一起带走。”

这是模型每一次推理或训练时,数据真正流过的路径。

def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
  • 接收三个张量输入。在自注意力中,传入的 q, k, v 其实是同一个张量 XX;在交叉注意力中,q 是解码器数据,k, v 是编码器数据。
bsz, seqlen, _ = q.shape
  • 解包形状:提取当前的批次大小(bsz = 32)和句子长度(seqlen = 10)。后面的 _ 代表我们忽略第三个维度(因为我们知道它是 n_embd)。
xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
  • 生成 QKV:让输入过一遍刚才定义的 Linear 矩阵。此时 xq, xk, xv 的形状都是 (32, 10, 512)
xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_heads, self.head_dim)
  • 逻辑切片:把最后的 512 维,看作 8 个头,每个头 64
  • 形状变为:(32, 10, 8, 64)
xq = xq.transpose(1, 2)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
  • 物理转置:交换第 1 维(seqlen)和第 2 维(n_heads)。
  • 形状变为:(32, 8, 10, 64)
  • 目的:把 n_heads 提到前面,让 GPU 把 32 * 8 = 256 个区块当成互相独立的批次,在底层的 CUDA 核心里并行地执行矩阵乘法。
scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
  • 核心数学运算:QKT/dkQ \cdot K^T / \sqrt{d_k}
  • xk.transpose(2, 3):针对 (32, 8, 10, 64) 的后两维转置,变成 (32, 8, 64, 10)
  • 点积(..., 10, 64) * (..., 64, 10) 结果为 (..., 10, 10)
  • 此时 scores 形状为:(32, 8, 10, 10)。这代表 32 个句子,8 个专家,每个人都画出了一张 10x10 的关系打分图。
  • / math.sqrt(self.head_dim):除以 64=8\sqrt{64} = 8。这就是我们之前讨论过的,防止方差爆炸、拯救 Softmax 梯度的关键缩放步骤。
if self.is_causal:
    assert hasattr(self, 'mask')
    scores = scores + self.mask[:, :, :seqlen, :seqlen]
  • 施加掩码:如果是解码器。
  • **为什么切片 :seqlen? 因为在 __init__ 里造了一个 max_seq_len(比如 4096)那么大的终极掩码。但当前这句话可能只有 10 个词长,所以只从左上角切出一块 10x10 的掩码来用。
  • scores (32, 8, 10, 10) 加上切好的 mask (1, 1, 10, 10),触发 PyTorch 广播机制,完美覆盖所有批次和头,把“未来”分数变成 -inf
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
  • 归一化与防溢出保护
  • .float():强制把半精度(FP16/BF16)的 scores 转成单精度(FP32)。因为 Softmax 里的 exe^x 指数爆炸极易导致 FP16 溢出报错 (NaN)。必须升精度。
  • F.softmax(..., dim=-1):在最后一个维度(行)上做 Softmax。把 -inf 变成 0,其他的变成加和为 1 的百分比权重。
  • .type_as(xq):算完 Softmax,安全了,赶紧把精度降回 xq 原本的半精度,省点显存留给后面的兄弟。
scores = self.attn_dropout(scores)
  • 加点噪音:随机把刚刚算好的一些注意力权重抹成 0。强迫模型不要过分依赖某一个特定的词。
output = torch.matmul(scores, xv)
  • 提取价值:拿权重 (32, 8, 10, 10) 去乘以 VV (32, 8, 10, 64)
  • 结果 output 形状为:(32, 8, 10, 64)
output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
  • 收网、拼接 (Concat)

  • transpose(1, 2):把头换回中间,形状变回 (32, 10, 8, 64)

  • contiguous():向系统申请一块连续的新内存,把数据搬过去规整好。

  • view(..., -1):强行压扁后两维。8 * 64 再次合体变成 512

  • 结果 output 形状为:(32, 10, 512)

output = self.wo(output)
output = self.resid_dropout(output)
return output
  • 让拼接好的特征穿过 WOW^O 矩阵,完成不同头特征的信息交流与融合。
  • 进行最后一次 Dropout。
  • 输出全新 (32, 10, 512) 特征张量交给下一层的“前馈神经网络(FFN)”继续处理。

03 Encoder-Decoder

Transformer 中的 Encoder,就是用于上述的编码过程;Decoder 则用于上述的解码过程。Transformer 结构,如图所示: transformer架构|450

1. 层归一化 (LayerNorm)

为什么不使用Batch Norm(批归一化)?

1). 批归一化缺点

批归一化存在一些缺陷,例如:

  • 当显存有限,mini-batch 较小时,Batch Norm 取的样本的均值和方差不能反映全局的统计分布信息,从而导致效果变差;
  • 对于在时间维度展开的 RNN,不同句子的同一分布大概率不同,所以 Batch Norm 的归一化会失去意义;
  • 在训练时,Batch Norm 需要保存每个 step 的统计信息(均值和方差)。在测试时,由于变长句子的特性,测试集可能出现比训练集更长的句子,所以对于后面位置的 step,是没有训练的统计量使用的;
  • 应用 Batch Norm,每个 step 都需要去保存和计算 batch 统计量,耗时又耗力

把神经网络里的数据想象成一次期末考试的成绩单

  • Batch(批次):一个班级里的学生人数(比如 32 个学生)。
  • Sequence Length(序列长度):每个人考了多少门课(句子长度,比如 10 个词)。
  • Hidden Size(特征维度):每门课的考卷上有多少道题(比如 512 维特征)。

Batch Norm 的计算逻辑是: 把这 32 个学生、在第 1 门课的第 1 道题上的分数抽出来,算一个平均分,然后归一化。

这就直接导致四个致命缺陷:

  1. “样本太少,统计失真”(受限于 Batch Size): 如果显存很小,一个 Batch 只能塞下 2 个句子(2 个学生)。用 2 个人的分数去代表平均水平,均值和方差会剧烈抖动,导致模型崩溃。

  2. “长短句参差不齐”(序列长度问题): 句子长度是不一样的。有的学生考了 10 门课,有的学生只考了 3 门课。当 Batch Norm 算到第 10 门课的平均分时,发现 32 个学生里只有 5 个人考了这门课,剩下的人全是 0(Padding 填充符)。这算出来的均值完全失去了统计学意义。

  3. “没见过的超长句子”(测试集的灾难): 在训练时,Batch Norm 需要记录每一道题的“历史平均分”。如果平时测试大家最多只考 20 门课(最长序列 20),但期末考试突然来了一个考了 50 门课的学生。对于第 21~50 门课,Batch Norm 根本没有历史记录可以用来归一化,当场宕机。

  4. “耗时又耗力”: 在时间维度展开算这么多乱七八糟的全局变量,内存访问极其不连续,拖慢了显卡的计算速度。

Layer Norm 的逻辑: 只针对某一个句子里的某一个词,把它自己身上的 512 维特征加起来,算出一个均值和方差,进行归一化。

Layer Norm 的优点:

  1. 不受 Batch Size 影响: 就算只有一个句子、一个词,Layer Norm 依然能完美算出它的均值和方差。
  2. 适应变长序列: 句子长短无所谓,因为归一化是针对“单个词内部”独立进行的,不用管别人长短。
  3. 测试阶段不需要历史统计量: 它在推理(生成)文本时,是实时计算当前词的统计量的,根本不需要像 Batch Norm 那样去保存训练时的滑动平均值。

2). 解读层归一化代码

class LayerNorm(nn.Module):
    ''' Layer Norm 层'''
    def __init__(self, features, eps=1e-6):
        super().__init__()
        # 线性矩阵做映射
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
    
    def forward(self, x):
        # 在统计每个样本所有维度的值,求均值和方差
        mean = x.mean(-1, keepdim=True) # mean: [bsz, max_len, 1]
        std = x.std(-1, keepdim=True) # std: [bsz, max_len, 1]
        # 注意这里也在最后一个维度发生了广播
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

注意,在我们上文实现的 Layer Norm 层中,有两个线性矩阵进行映射。

看看 Layer Norm 是如何在底层实现的。

class LayerNorm(nn.Module):
    ''' Layer Norm 层'''
    def __init__(self, features, eps=1e-6):
        super().__init__()
        # 线性矩阵做映射
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
  1. features 与 缩放平移参数 (self.a_2, self.b_2)Zj​​=σ2+ϵ​Zj​−μj​​
  • features 就是 hidden_size(比如 512)。

  • 为什么要有 a_2b_2 归一化操作(减去均值除以方差)会暴力地把原本的数据全部压缩到均值为 0、方差为 1 的正态分布里。这其实会破坏模型学到的特征(比如某个特征本来就该特别大,代表了某种强烈的情感)。 为了弥补这个损失,给模型发了两个可学习的参数

    • a_2 (缩放参数 Scale / Gain,即公式里的 γ\gamma):全初始化为 1。
    • b_2 (平移参数 Bias / Shift,即公式里的 β\beta):全初始化为 0。 模型在后续的训练中,如果发现“归一化破坏了特征”,它可以通过反向传播自动调节 a_2b_2 的值,把归一化后的数据重新拉伸和平移到最合适的分布。
  • nn.Parameter(...) 是告诉 PyTorch:这两个变量不是普通张量,而是需要跟着模型一起训练更新的参数。nn.Parameter 默认情况下 requires_grad=True,这意味着它会参与梯度计算和反向传播。

  1. eps=1e-6 ε 这是一个极小的数字(10610^{-6})。在后面的除法运算中,如果一个词的方差极其微小(接近 0),除以 0 会导致计算机直接报 NaN(Not a Number)错误。

    def forward(self, x):
        # 统计每个样本所有维度的值,求均值和方差
        mean = x.mean(-1, keepdim=True) # mean: [bsz, max_len, 1]
        std = x.std(-1, keepdim=True)   # std: [bsz, max_len, 1]
  1. 计算均值和方差
  • 输入 x 形状(bsz, max_len, features),例如 (32, 10, 512)

  • x.mean(-1):这里的 -1 指的是最后一个维度(features/512)。 这句代码在物理意义上执行的操作是:把每一个词包含的 512 个数字加起来求平均。

  • keepdim=True: 如果不加它,算完均值后,512 这个维度就彻底消失了,形状会变成 (32, 10)。 加上它之后,虽然 512 维被压扁成了一个单一的均值,但这个维度的“空壳”还在,形状变成了 (32, 10, 1)

        # 注意这里也在最后一个维度发生了广播
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
  • (x - mean)x 形状 (32, 10, 512)mean 形状 (32, 10, 1)。 PyTorch 触发广播机制,把 mean 里的那个 1 自动复制扩展成 512 份。这相当于让一个词身上的 512 个特征,每一个都减去了它们共同的平均值。
  • / (std + self.eps):同理,每一个减去均值后的特征,都除以了标准差,完成了最纯粹的标准化(变成均值 0,方差 1)。
  • self.a_2 * (...) + self.b_2: 这里算完归一化后的数据形状依然是 (32, 10, 512)。 此时去乘上 a_2(形状为 512 维的一维张量)和加上 b_2(512 维)。 PyTorch 再次利用广播机制,把这 512 维的缩放和平移参数,完美地套用到所有的 32 个批次和 10 个词的对应维度上

2. 残差连接

残差连接

# 注意力计算
h = x + self.attention.forward(self.attention_norm(x))
# 经过前馈神经网络
out = h + self.feed_forward.forward(self.fnn_norm(h))
  1. 公式逐项拆解:xnew=xold+Attention(LayerNorm(xold))x_{new} = x_{old} + \text{Attention}(\text{LayerNorm}(x_{old}))

假设 xoldx_{old} 是我们刚才讨论过的、包含“我 爱 学习”这句话的词向量矩阵。

  • 等号左边的 xx (即 xnewx_{new}):这是经过这一个子层(Sub-layer)处理后,输出的全新特征
  • 等号右边单独的 xx (即 xoldx_{old}):这是进入这个子层之前的原始特征。让它绕过中间所有的复杂计算,直接送到终点。
  • LayerNorm(x)\text{LayerNorm}(x):在进行复杂的注意力计算前,先把原始数据做一次层归一化,让数据变得“规整”。
  • MultiHeadSelfAttention(...)\text{MultiHeadSelfAttention}(...):在这堆规整的数据上,让 8 个专家去寻找词与词之间的上下文关系,提取出增量特征(我们叫它残差/Residual)。
  • 最核心的加号 +:把“直达快车道”送来的原始特征,和专家们刚刚提取出来的增量特征,直接相加

  1. 为什么非要加这个 xx

为什么不直接输出注意力算出来的结果,非要把原来的 xx 加回来呢?

痛点一:信息的遗忘

如果只保留注意力的结果,模型在经过几十层网络后,词向量在不断地融合上下文,它可能会忘了“自己到底是谁”。

比如“苹果”,在融合了“吃”之后,可能会变得太像一个动作的受体。

加上原来的 xx 后,公式的物理意义变成了:“保持我作为‘苹果’的本意不变,同时加上我吸收到的周围的语境信息。” 原来的特征不会丢失。

痛点二:“梯度消失”

在训练模型时,误差(梯度)需要从最后一层一层层往回传(连乘)。

如果网络有 100 层,很多小于 1 的数字连乘 100 次,梯度就会变成 0.0000001,导致最前面的层根本得不到更新。

但是有了这个 + x,根据微积分的求导法则,f(x)+xf(x) + x 的导数是 f(x)+1f'(x) + 1。 无论注意力的梯度 f(x)f'(x) 变得多小,甚至变成了 0,梯度都可以沿着 + 1 ,无损地直接传回上一层。这使得训练极深的大模型成为了可能。 残差连接效果对比

  1. 第二个公式的意义 output=x+FNN(LayerNorm(x))output = x + \text{FNN}(\text{LayerNorm}(x)) 这个公式的逻辑和上面完全一模一样,只是把中间从多头注意力换成前馈神经网络。
  • 注意力机制(Attention) 的作用是:让词和词之间交换信息
  • 前馈神经网络(FFN) 的作用是:在每个词自己的特征内部进行升维再降维的非线性变换,让特征更丰富。

在一个 Transformer 的编码器(Encoder)块里,数据就是这样连续:

先走一次 “注意力+残差” 提取全局关系, 再走一次 “FFN+残差” 丰富自身特征。 这样连续堆叠 32 层,就构成了一个庞大聪明的 LLM。

3. 代码

class Encoder(nn.Module):
    '''Encoder 块'''
    def __init__(self, args):
        super(Encoder, self).__init__() 
        # 一个 Encoder 由 N 个 Encoder Layer 组成
        self.layers = nn.ModuleList([EncoderLayer(args) for _ in range(args.n_layer)])
        self.norm = LayerNorm(args.n_embd)

    def forward(self, x):
        "分别通过 N 层 Encoder Layer"
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)

第一部分:初始化 __init__

def __init__(self, args):
    super(Encoder, self).__init__() 
    # 一个 Encoder 由 N 个 Encoder Layer 组成
    self.layers = nn.ModuleList([EncoderLayer(args) for _ in range(args.n_layer)])
    self.norm = LayerNorm(args.n_embd)
  1. [EncoderLayer(args) for _ in range(args.n_layer)]

    • 这是 Python 的列表推导式。假设配置文件里写了 args.n_layer = 12,这就意味着模型要造 12 层。
    • 这行代码会连续实例化 12 个 EncoderLayer 对象。
    • 核心重点: 这 12 个层的内部结构(多头注意力 + FFN)是一模一样的,但是它们内部的参数权重矩阵(WQ,WK,WVW^Q, W^K, W^V 等)是完全独立、互不干扰的!第一层负责提取浅层特征(比如词性),第十二层负责提取极其抽象的深层语义。
  2. nn.ModuleList(...)

    • 为什么不直接用 Python 自带的普通的 list,而非要套一层 nn.ModuleList
    • 这是 PyTorch 的硬性规定。 如果你用普通的 list 把这 12 层装起来,PyTorch 的底层引擎会“看不见”它们,在反向传播更新梯度时,这 12 层里的参数全都会被忽略,模型就废了。
    • 套上 nn.ModuleList,就像是在向 PyTorch 注册:“请把这个列表里的所有网络层都纳入你的计算图,好好计算梯度。”
  3. self.norm = LayerNorm(args.n_embd)

    • 在数据跑完所有的 N 层之后,最后再做一次层归一化(Layer Norm),稳定特征分布,然后送出编码器。

第二部分:前向传播 forward

def forward(self, x):
    "分别通过 N 层 Encoder Layer"
    for layer in self.layers:
        x = layer(x)
    return self.norm(x)

这段代码看似极其简单,但它描绘了数据在神经网络中穿梭的最核心画面。

  • 初始的输入 xx(带着词嵌入和位置编码的张量)进入循环。
  • x = layer(x)
    • xx 进入第 1 层,经历多头注意力和 FFN,变成了一个更高级的 xnewx_{new}
    • 然后把这个 xnewx_{new} 重新赋值给 xx
    • 接着,这个升级版的 xx 进入第 2 层,继续融合、升级……
  • 就这样不断循环 NN 次。每一次循环,xx 都被当前层打磨得越来越“懂”这句话的宏观语义。
  • 最后,打磨了 NN 次的 xx 走出循环,穿过最后的 `self.norm(x)

04 搭建 Transformer

1. 位置编码

1). 对照解释

import numpy as np
import matplotlib.pyplot as plt
def PositionEncoding(seq_len, d_model, n=10000):
    P = np.zeros((seq_len, d_model))
    for k in range(seq_len):
        for i in np.arange(int(d_model/2)):
            denominator = np.power(n, 2*i/d_model)
            P[k, 2*i] = np.sin(k/denominator)
            P[k, 2*i+1] = np.cos(k/denominator)
    return P

P = PositionEncoding(seq_len=4, d_model=4, n=100)
print(P)

正弦/余弦位置编码(Sinusoidal Positional Encoding) 公式。 由于 Transformer 的自注意力机制是同时处理所有词的(并行计算),分不清“我爱你”和“你爱我”的区别。这段代码的作用,就是为句子里的每一个物理位置,生成一个独一无二的向量,让模型能够感知到词汇的先后顺序。

  1. 参数与矩阵初始化
def PositionEncoding(seq_len, d_model, n=10000):
    P = np.zeros((seq_len, d_model))
  • seq_len: 序列长度,也就是一句话里有几个词。
  • d_model: 词向量的维度(例如大模型常用的 512)。
  • n: 频率缩放的底数(论文中默认是 10000)。
  • P: 生成一个形状为 (seq_len, d_model) 的全 0 矩阵。这相当于准备了一张巨大的表格,每一行代表一个词的位置,每一列代表位置编码的一个维度。
  1. 双重循环(定位与计算)
for k in range(seq_len):
    for i in np.arange(int(d_model/2)):
  • 外层循环 k: 遍历句子中的每一个位置(第 0 个词、第 1 个词…)。在论文公式中,它对应变量 pospos
  • 内层循环 i: 遍历维度的“对”(pair)。因为我们会交替使用正弦和余弦,所以每次处理 2 个维度,循环次数自然是 d_model / 2
  1. 核心数学公式:计算频率分母
        denominator = np.power(n, 2*i/d_model)

这行代码在计算论文公式里的分母部分:100002i/dmodel10000^{2i / d_{model}}

  • 随着维度索引 i 的增大,这个分母 denominator 会呈指数级变大。
  • 分母越大,后续计算出来的波形频率就越低(波长越长)
  1. 正余弦交替赋值

        P[k, 2*i] = np.sin(k/denominator)
        P[k, 2*i+1] = np.cos(k/denominator)
  • 偶数列(2*i:填入正弦值 sin(pos/denominator)\sin(pos / denominator)
  • 奇数列(2*i+1:填入余弦值 cos(pos/denominator)\cos(pos / denominator)

为什么非要用不同频率的 Sin 和 Cos? 可以把这个位置编码矩阵想象成机械水表的几个齿轮:

  • 前面的维度(比如第 0、1 维):频率极高,像水表最右边走得飞快的小齿轮。位置 k 稍微变一点,它的值就剧烈变化。模型通过它来分辨相邻的两个词谁前谁后。
  • 后面的维度(比如第 510、511 维):频率极低,像水表最左边走得极慢的大齿轮。位置 k 走过几十个词,它的值才发生微小的变化。模型通过它来感知句子的宏观长度和远距离依赖位置编码
P = PositionEncoding(seq_len=4, d_model=4, n=100)

生成了一个 4 个词、每个词 4 维的模型。

  • 对于第 0 个词 (k=0):不管分母是多少,sin(0)=0\sin(0)=0cos(0)=1\cos(0)=1。所以矩阵的第 0 行永远是 [0, 1, 0, 1]
  • 对于后面的词:数值开始随着波浪形的频率产生变化,每一行的 4 个数字组合绝对不会重复。

2). 距离计算

举实例实现PE(pos+k) 可以用 PE(pos)计算

把频率分母(100002i/dmodel10000^{2i/d_{model}})简化成一个常数符号 ω\omega

  1. 明确已知条件 在特定的第 2i2i(偶数)和 2i+12i+1(奇数)维度上,绝对位置 pospos 的编码公式为: PE(pos,2i)=sin(ωpos)PE(pos, 2i) = \sin(\omega \cdot pos) PE(pos,2i+1)=cos(ωpos)PE(pos, 2i+1) = \cos(\omega \cdot pos)

  2. 开始推导 PE(pos+k)PE(pos+k)

现在,我们想知道距离当前位置 kk 个步长的那个词,它的位置编码 PE(pos+k)PE(pos+k) 是多少。 先看偶数维度 (2i2i):

PE(pos+k,2i)=sin(ω(pos+k))PE(pos+k, 2i) = \sin(\omega \cdot (pos + k)) 将其展开为 sin(ωpos+ωk)\sin(\omega \cdot pos + \omega \cdot k)。 套用三角函数公式 sin(A+B)=sin(A)cos(B)+cos(A)sin(B)\sin(A+B) = \sin(A)\cos(B) + \cos(A)\sin(B)PE(pos+k,2i)=sin(ωpos)cos(ωk)+cos(ωpos)sin(ωk)PE(pos+k, 2i) = \sin(\omega \cdot pos)\cos(\omega \cdot k) + \cos(\omega \cdot pos)\sin(\omega \cdot k) 这个时候,仔细看等号右边的式子:

  • sin(ωpos)\sin(\omega \cdot pos) 不就是 PE(pos,2i)PE(pos, 2i) 吗?
  • cos(ωpos)\cos(\omega \cdot pos) 不就是 PE(pos,2i+1)PE(pos, 2i+1) 吗?

把它们替换进去: 公式一: PE(pos+k,2i)=PE(pos,2i)cos(ωk)+PE(pos,2i+1)sin(ωk)PE(pos+k, 2i) = PE(pos, 2i)\cos(\omega \cdot k) + PE(pos, 2i+1)\sin(\omega \cdot k)

再看奇数维度 (2i+12i+1): PE(pos+k,2i+1)=cos(ω(pos+k))PE(pos+k, 2i+1) = \cos(\omega \cdot (pos + k)) 展开为 cos(ωpos+ωk)\cos(\omega \cdot pos + \omega \cdot k)。 套用公式 cos(A+B)=cos(A)cos(B)sin(A)sin(B)\cos(A+B) = \cos(A)\cos(B) - \sin(A)\sin(B)PE(pos+k,2i+1)=cos(ωpos)cos(ωk)sin(ωpos)sin(ωk)PE(pos+k, 2i+1) = \cos(\omega \cdot pos)\cos(\omega \cdot k) - \sin(\omega \cdot pos)\sin(\omega \cdot k) 同样进行替换: 公式二: PE(pos+k,2i+1)=PE(pos,2i+1)cos(ωk)PE(pos,2i)sin(ωk)PE(pos+k, 2i+1) = PE(pos, 2i+1)\cos(\omega \cdot k) - PE(pos, 2i)\sin(\omega \cdot k)

在深度学习的代码里,模型只认识矩阵乘法(线性变换)

把刚才推导出的公式一公式二写成线性代数里的矩阵乘法形式: [PE(pos+k,2i)PE(pos+k,2i+1)]=[cos(ωk)sin(ωk)sin(ωk)cos(ωk)][PE(pos,2i)PE(pos,2i+1)]\begin{bmatrix} PE(pos+k, 2i) \\ PE(pos+k, 2i+1) \end{bmatrix} = \begin{bmatrix} \cos(\omega \cdot k) & \sin(\omega \cdot k) \\ -\sin(\omega \cdot k) & \cos(\omega \cdot k) \end{bmatrix} \begin{bmatrix} PE(pos, 2i) \\ PE(pos, 2i+1) \end{bmatrix}

看中间那个 2×22 \times 2 的变换矩阵 [cos(ωk)sin(ωk)sin(ωk)cos(ωk)]\begin{bmatrix} \cos(\omega \cdot k) & \sin(\omega \cdot k) \\ -\sin(\omega \cdot k) & \cos(\omega \cdot k) \end{bmatrix}里面只有距离参数 kk,根本没有绝对位置 pospos

无论你是在句子的开头(pos=1pos=1),还是在句子的末尾(pos=100pos=100),只要两个词之间的距离是 kk(比如 k=2k=2,也就是“隔一个词”),那么从词 A 的位置向量变到词 B 的位置向量,所经历的线性变换矩阵是绝对固定、一模一样的。

3). 实际代码

class PositionalEncoding(nn.Module):
    '''位置编码模块'''

    def __init__(self, args):
        super(PositionalEncoding, self).__init__()
        # Dropout 层
        # self.dropout = nn.Dropout(p=args.dropout)

        # block size 是序列的最大长度
        pe = torch.zeros(args.block_size, args.n_embd)
        position = torch.arange(0, args.block_size).unsqueeze(1)
        # 计算 theta
        div_term = torch.exp(
            torch.arange(0, args.n_embd, 2) * -(math.log(10000.0) / args.n_embd)
        )
        # 分别计算 sin、cos 结果
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        # 将位置编码加到 Embedding 结果上
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return x

在大模型开发中,写 for 循环去算矩阵是极其低效的。这段代码利用了 PyTorch 张量广播(Broadcasting)数学等价代换,把原本需要双重循环的代码,压扁成了瞬间完成的矩阵并行运算。

第一部分:初始化 (__init__)

这段代码在模型被创建的那一刻运行,它的目标是一次性把最大长度的绝对位置编码全部算好,存起来备用

pe = torch.zeros(args.block_size, args.n_embd)
position = torch.arange(0, args.block_size).unsqueeze(1)
  • block_size:就是模型允许的最大序列长度(比如 2048)。
  • pe:造一张 2048 行、512 列的全零空白表格。
  • position:生成一个从 0 到 2047 的列向量。形状是 (2048, 1)。这就是那个 pospos 变量。
div_term = torch.exp(
    torch.arange(0, args.n_embd, 2) * -(math.log(10000.0) / args.n_embd)
)

Numpy 里的分母是怎么写的?100002i/dmodel10000^{2i / d_{model}}

但是在底层,直接算“大底数的指数次幂”容易导致数值溢出(Overflow) 或失去精度。

所以使用对数指数恒等式:ab=ebln(a)a^b = e^{b \cdot \ln(a)},并将其倒数为负指数。

我们推导一下代码里的这串式子: e2idmodelln(10000)=eln(100002idmodel)=100002idmodel=1100002idmodele^{-\frac{2i}{d_{model}} \cdot \ln(10000)} = e^{\ln(10000^{-\frac{2i}{d_{model}}})} = 10000^{-\frac{2i}{d_{model}}} = \frac{1}{10000^{\frac{2i}{d_{model}}}} 经过极其精妙的数学代换,它完美等价于原公式里的 1分母\frac{1}{\text{分母}} 而且 torch.expmath.log 在 GPU 上的计算速度极快且绝对稳定。 这里 div_term 算出来是一个行向量,形状为 (256,) (假设维度是 512)。

张量广播

pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
  • position * div_term:这里发生了一次的广播机制。一个形状为 (2048, 1) 的列向量,去乘一个形状为 (256,) 的行向量。PyTorch 会瞬间自动把它们扩展,生成一个 (2048, 256) 的庞大矩阵!
  • 0::21::2:这是 Python 的切片语法。0::2 代表从索引 0 开始,步长为 2(即 0, 2, 4… 偶数列)。1::2 代表奇数列。 所有的正弦和余弦值就被填入了这 2048 行矩阵的对应列中。

打包备用

pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
  • unsqueeze(0):在最前面加一个 Batch 维度。形状从 (2048, 512) 变成了 (1, 2048, 512)
  • register_buffer:在讲多头注意力掩码时遇到过,位置编码是固定的死规则,不需要被梯度更新。所以把它注册为 Buffer,让它跟着模型走,但别让优化器去动它。

第二部分:前向传播中的动态应用 (forward)

def forward(self, x):
    # 将位置编码加到 Embedding 结果上
    x = x + self.pe[:, : x.size(1)].requires_grad_(False)
    return x
  • x 的形状:输入的词向量,通常是 (batch_size, seq_len, 512)。注意,这里的 seq_len 可能只有 10,远远小于我们刚才造好的 2048。

  • : x.size(1)按需切片。既然你这句话只有 10 个词长,那我就只从那个 2048 行的庞大表格里,切出前 10 行给你用。

  • + (相加): 输入 xx 的形状:(batch_size, 10, 512) 切出来的 pepe 形状:(1, 10, 512)

    再次触发广播机制!这 1 个批次的位置编码,会被复制应用到 batch_size 包含的每一句话上。词的语义(xx)和词的物理位置(pepe)在这一刻融合。

2. 完整架构搭建

Transformer 模型结构LayerNorm 层是放在 Attention 层前面的,也就是“Pre Norm”结构。考虑到目前 LLM 一般采用“Pre-Norm”结构(可以使 loss 更稳定),实现时采用“Pre-Norm”结构。

class Transformer(nn.Module):
   '''整体模型'''
    def __init__(self, args):
        super().__init__()
        # 必须输入词表大小和 block size
        assert args.vocab_size is not None
        assert args.block_size is not None
        self.args = args
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(args.vocab_size, args.n_embd),
            wpe = PositionalEncoding(args),
            drop = nn.Dropout(args.dropout),
            encoder = Encoder(args),
            decoder = Decoder(args),
        ))
        # 最后的线性层,输入是 n_embd,输出是词表大小
        self.lm_head = nn.Linear(args.n_embd, args.vocab_size, bias=False)

        # 初始化所有的权重
        self.apply(self._init_weights)

        # 查看所有参数的数量
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    '''统计所有参数的数量'''
    def get_num_params(self, non_embedding=False):
        # non_embedding: 是否统计 embedding 的参数
        n_params = sum(p.numel() for p in self.parameters())
        # 如果不统计 embedding 的参数,就减去
        if non_embedding:
            n_params -= self.transformer.wte.weight.numel()
        return n_params

    '''初始化权重'''
    def _init_weights(self, module):
        # 线性层和 Embedding 层初始化为正则分布
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
    '''前向计算函数'''
    def forward(self, idx, targets=None):
        # 输入为 idx,维度为 (batch size, sequence length, 1);targets 为目标序列,用于计算 loss
        device = idx.device
        b, t = idx.size()
        assert t <= self.args.block_size, f"不能计算该序列,该序列长度为 {t}, 最大序列长度只有 {self.args.block_size}"

        # 通过 self.transformer
        # 首先将输入 idx 通过 Embedding 层,得到维度为 (batch size, sequence length, n_embd)
        print("idx",idx.size())
        # 通过 Embedding 层
        tok_emb = self.transformer.wte(idx)
        print("tok_emb",tok_emb.size())
        # 然后通过位置编码
        pos_emb = self.transformer.wpe(tok_emb) 
        # 再进行 Dropout
        x = self.transformer.drop(pos_emb)
        # 然后通过 Encoder
        print("x after wpe:",x.size())
        enc_out = self.transformer.encoder(x)
        print("enc_out:",enc_out.size())
        # 再通过 Decoder
        x = self.transformer.decoder(x, enc_out)
        print("x after decoder:",x.size())

        if targets is not None:
            # 训练阶段,如果我们给了 targets,就计算 loss
            # 先通过最后的 Linear 层,得到维度为 (batch size, sequence length, vocab size)
            logits = self.lm_head(x)
            # 再跟 targets 计算交叉熵
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # 推理阶段,我们只需要 logits,loss 为 None
            # 取 -1 是只取序列中的最后一个作为输出
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

最后,可以读一读 为什么你总觉得看不懂 Transformer?一次从0到1的架构大串讲 - 知乎