logo蛋烘糕.

不写博客的工程师不是好的产品经理

AI 技术演进与核心算法实战 | 第五篇:Transformer 架构深潜:Positional Encoding、Layer Norm 与前馈网络的手写实现(NanoGPT 复现)

Cover Image for AI 技术演进与核心算法实战 | 第五篇:Transformer 架构深潜:Positional Encoding、Layer Norm 与前馈网络的手写实现(NanoGPT 复现)
蛋烘糕
蛋烘糕

Self-Attention 是 Transformer 的灵魂,但只有灵魂还不够,它还需要一副强大的躯壳。今天,我们就来拼装这副躯壳。

上一篇中,我们解密了 Self-Attention 的核心机制,明白了词与词之间是如何通过 Q、K、V 进行动态交流的。然而,纯粹的 Attention 机制有两个致命的缺陷:

  1. 它是“色盲”加“脸盲”的(丢失位置信息):Self-Attention 将输入视为一个“无序的词袋”,无论“狗咬人”还是“人咬狗”,计算出的 Attention 分数是一样的,因为它完全不知道词的顺序。
  2. 它容易“营养不良”或“走火入魔”(梯度问题与特征表达能力不足):仅靠线性变换的 Attention,无法拟合复杂的非线性语言规律,同时深层网络极易出现梯度消失或爆炸。

为了解决这些问题,Transformer 引入了三大法宝:位置编码(Positional Encoding)层归一化(Layer Normalization)前馈网络(Feed-Forward Network, FFN)

本篇是 《AI 技术演进与核心算法实战》第一模块的第五篇。我们将结合 Andrej Karpathy 的 NanoGPT 项目,用手写代码和图解,带你彻底搞透这三大组件的底层逻辑,并最终拼装出一个完整的 Transformer Block!


1. 注入时间之魂:Positional Encoding

为什么 Transformer 那么快?因为它用“并行计算”取代了 RNN 的“串行计算”。 但代价是什么?代价是丢失了时序信息。

为了让模型知道每个词在句子中的位置,我们需要给每个词贴上一个“位置标签”。这就是位置编码(Positional Encoding)。

1.1 绝对位置编码的直觉

最简单的想法是:第一个词贴上标签 1,第二个词贴上 2…… 但这样做会引发两个致命问题:

  1. 特征在数值上被“淹没”(尺度失衡) 在深度学习中,Embedding 层输出的词向量数值通常非常小(分布在 [1,1][-1, 1] 附近)。如果位置索引是 1000,并且与词向量直接相加(1000 + 0.3 = 1000.3),原本代表单词核心语义的微小数值就会在巨大的位置数值面前显得微乎其微。这会导致网络过度关注“位置”,而忽略了“语义”,同时极大的输入值容易让后续的激活函数进入饱和区,造成梯度消失。

  2. 模型难以泛化到更长的句子(外推性崩溃) 如果模型在训练时最多只见过 512 个词长度的句子,它的权重参数只适应了 1512 这样的位置数值。当推理阶段遇到长度为 1000 的句子时,5131000 属于模型从未见过的分布外(OOD)巨大数值。这些异常输入会导致神经网络输出不可预测的巨大激活值,使模型完全崩溃。

为了解决这两个痛点,Google 在论文中给出了一个极其优雅的解法:使用不同频率的正弦和余弦函数

三角函数(sin\sincos\cos)的值域永远被限制在 [1,1][-1, 1] 之间,完美解决了特征被淹没的问题。同时,三角函数的周期性规律不仅能帮助模型计算相对距离,还具备平滑延伸的特性,极大缓解了长度泛化困难的问题。

想象一个二进制计数器:

  • 个位:0 1 0 1 0 1...(变化极快)
  • 十位:0 0 1 1 0 0...(变化中等)
  • 百位:0 0 0 0 1 1...(变化极慢)

Transformer 的位置编码也是类似原理,只是把离散的 0 和 1 换成了连续的 sincos 曲线。向量的第一维变化最快(高频),最后一维变化最慢(低频)。

1.2 图解位置编码

词 Embedding + 位置 Encoding = 带有位置的输入 高频 (维 0) 中频 (维 d/2) 低频 (维 d)

图解说明: 从上图中我们可以直观地看到 Transformer 注入位置信息的过程:

  1. 直接相加(Add):位置编码(Positional Encoding)的维度与词向量(Token Embedding)完全相同,两者通过元素级相加(Element-wise Addition)融合在一起,形成最终“带有位置的输入”。
  2. 频率递减的波形:图中下方的三条波浪线代表了位置编码向量中不同维度的数值变化规律。
    • 高频(红色波浪):对应向量的较前维度(如维度 0)。波峰波谷变化非常密集,类似于二进制中的“个位”,对相邻位置极其敏感,用于区分近距离的词。
    • 中频(蓝色波浪):对应向量的中间维度。波长变长,变化适中。
    • 低频(绿色波浪):对应向量的较后维度(如维度 d)。波形非常平缓,类似于二进制中的“百位”,在一个较长的句子范围内才完成一次周期变化,帮助模型感知词在句子中的宏观/全局相对位置。

1.3 手写实现(PyTorch)

在目前的实践中(如 GPT 系列),通常采用更简单的可学习的位置嵌入(Learned Positional Embedding),而不是固定的正余弦。NanoGPT 采用的正是这种方案:

import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    def __init__(self, block_size, n_embd):
        super().__init__()
        # block_size 是最大序列长度,n_embd 是词向量维度
        # 直接创建一个可学习的 Embedding 矩阵 (block_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        
    def forward(self, x):
        # x 的 shape: (batch_size, seq_len, n_embd)
        B, T, C = x.shape
        # 生成 0 到 T-1 的位置索引
        pos = torch.arange(0, T, dtype=torch.long, device=x.device) # (T)
        # 获取位置向量并与输入相加
        pos_emb = self.pos_emb(pos) # (T, C)
        return x + pos_emb # 利用广播机制 (B, T, C) + (T, C) -> (B, T, C)

物理意义: 词的最终表达 = “它是什么意思”(Token Embedding) + “它在哪个位置”(Positional Embedding)。这两者在向量空间中完美融合。


2. 驯服狂暴的梯度:Layer Normalization

当我们在堆叠多层 Self-Attention 时,数值会变得非常大,导致梯度爆炸,网络根本无法收敛。 为了让每一层的输出“冷静”下来,我们需要 Layer Normalization(层归一化)

2.1 为什么是 Layer Norm 而不是 Batch Norm?

在图像领域(CNN),大家习惯用 Batch Norm(批归一化),它是跨越 Batch 对同一个通道进行归一化。 但在自然语言中,句子的长度是不一样的(有长有短),如果跨句子求平均,长句后面的词只能和填充(Padding)的 0 一起算,这毫无意义。

因此,NLP 中采用 Layer Norm在每一个词的词向量内部进行归一化。不管句子多长,每个词管好自己的向量就行。

2.2 图解 Layer Norm

Batch Norm 跨样本同维度计算 Layer Norm 单样本同词汇内部计算 NLP首选

Batch Norm (批归一化) 的执行过程: 想象你有一个班级的学生成绩单,包含“数学”、“语文”、“英语”三门课。Batch Norm 的做法是:把全班所有人的“数学”成绩拿出来,计算这门课的平均分和方差,然后让每个人的数学成绩减去平均分再除以标准差。接着对“语文”、“英语”做同样的操作。 在深度学习中,它是在同一个特征维度上,跨越所有样本(Batch) 计算均值和方差。这种方式在计算机视觉(CV)中很常用,因为图片的特征(如 RGB 通道)在不同图片间有明确的对应关系。但在 NLP 中,句子的长度不一样,某个词的位置在不同句子中可能代表完全不同的含义,强行跨样本算均值就不合理了。

Layer Norm (层归一化) 的执行过程: 还是用学生成绩单举例。Layer Norm 的做法是:单独看“张三”这名同学,把他自己的“数学”、“语文”、“英语”成绩加起来算一个平均分和方差,然后把他的各科成绩按照他自己的平均分做归一化。李四也独立计算自己的。 在 NLP 中,它是针对每一个单独的样本(也就是一个句子或一个词),在这个样本内部的所有特征维度上计算均值和方差。因为文本数据长度可变,Layer Norm 不依赖于 Batch 的大小,只关注当前词自身的特征分布,因此在 Transformer 等 NLP 模型中表现更好,成为了首选。

2.3 手写实现(核心逻辑)

Layer Norm 的公式很简单:先减去均值,再除以标准差,最后乘上缩放参数 gamma,加上平移参数 beta

class LayerNorm(nn.Module):
    def __init__(self, ndim, bias=True):
        super().__init__()
        # 可学习的缩放参数 gamma 和平移参数 beta
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        # input shape: (Batch, SeqLen, EmbeddingDim)
        # 在 Embedding 维度(最后一步)上计算均值和方差
        mean = input.mean(dim=-1, keepdim=True)
        var = input.var(dim=-1, keepdim=True, unbiased=False)
        
        # 归一化:减均值,除以标准差 (加 1e-5 防止除零)
        out = (input - mean) / torch.sqrt(var + 1e-5)
        
        # 应用可学习的仿射变换
        out = out * self.weight
        if self.bias is not None:
            out = out + self.bias
        return out

3. 赋予思考的深度:前馈网络(FFN)

Self-Attention 虽然强大,但它只做了一件事:让词与词之间交换信息。这本质上只是一个复杂的加权求和(线性变换)。 如果全是线性变换,多层网络就会退化成一层。我们需要非线性激活函数,这就是前馈网络(Feed-Forward Network, FFN)的作用。

3.1 放大与压缩的艺术

在 Transformer 中,FFN 的结构非常经典:

  1. 升维:先将维度放大 4 倍(n_embd -> 4 * n_embd)。
  2. 非线性:通过激活函数(如 GELU 或 ReLU)进行非线性映射。
  3. 降维:再将维度压缩回原来的大小(4 * n_embd -> n_embd)。

物理意义与通俗解释:

  • Self-Attention 负责“收集资料”(看上下文):就像你在写论文时,去图书馆查阅各种相关的书籍和文献,把它们汇总到一起。这个过程是在做“信息交流”,但还没有形成你自己的观点。
  • FFN 负责“闭门思考”(内部消化与推理):收集完资料后,你把自己关在房间里,开始思考、提炼和总结。
    • 升维(扩大 4 倍):相当于你在草稿纸上发散思维,把简单的几个关键词展开成各种可能的详细推论。维度越大,模型能“记住”和处理的特征组合就越丰富。
    • 非线性激活(GELU):相当于你的逻辑判断过程,“如果 A 成立,那么就 B,否则就抛弃 C”。如果没有这一步,再多的思考也只是一堆简单的线性叠加。
    • 降维(还原维度):思考结束后,你需要把长篇大论的推导过程,浓缩成论文里精炼的一句话,输出给下一个阶段。这就是将维度压缩回原来的大小,保证网络结构的统一性。

3.2 手写实现

class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            # 1. 升维:放大 4 倍
            nn.Linear(n_embd, 4 * n_embd),
            # 2. 非线性激活:GPT 系列通常使用 GELU
            nn.GELU(),
            # 3. 降维:还原回 n_embd
            nn.Linear(4 * n_embd, n_embd),
            # Dropout 防止过拟合
            nn.Dropout(0.1),
        )

    def forward(self, x):
        return self.net(x)

4. 拼装神迹:构建完整的 Transformer Block

现在,我们有了 Self-Attention(上篇)、Layer Norm 和 FFN。是时候把它们组装成一个真正的 Transformer Block 了!

在早期的 Transformer(如原始论文)中,Layer Norm 放在 Attention 之后(Post-LN)。但实践证明,将 Layer Norm 放在 Attention 之前(Pre-LN)能让训练极其稳定,GPT 家族全部采用了 Pre-LN 架构。

4.1 架构图解(Pre-LN)

Input Layer Norm Self-Attention + Layer Norm Feed Forward + 残差连接 残差连接

架构图含义与执行过程通俗解释:

  1. 主干道与小路(残差连接): 图中实线箭头是数据处理的“主干道”(经过各种复杂计算),而左侧的虚线箭头是一条“小路”(残差连接,Skip Connection)。不管主干道上的计算多么复杂,原始输入都可以通过小路直接走到终点参与汇合(图中的 + 号)。这保证了哪怕主干道学“废”了,网络至少不会比原来更差。
  2. 第一阶段:交流信息(Attention)
    • Input:词向量首先进入 Block。
    • Layer Norm:在进行复杂交流前,先对自己做一次“归一化”整理,把特征拉平(这就是 Pre-LN 的体现,提前整理能让后续训练极其稳定)。
    • Self-Attention:整理好后,词与词之间开始互相看、互相交流,提取上下文信息。
    • 加法汇合(+):交流完的结果,与一开始没处理过的原始输入(从小路走过来的)加在一起,形成第一阶段的输出。
  3. 第二阶段:深度思考(FFN)
    • Layer Norm:拿着第一阶段汇总好的结果,再次进行“归一化”整理。
    • Feed Forward:每个词关起门来,利用刚才交流得到的信息进行内部的深度逻辑推理(升维、激活、降维)。
    • 加法汇合(+):思考完的结果,再次与第二阶段刚开始时的输入加在一起,最终输出给下一个完整的 Transformer Block。

4.2 手写实现

结合图解,我们用代码将它们串联起来。注意残差连接(x = x + ...)的运用,它是深层网络能够训练成功的另一个关键。

class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        # 1. 归一化层 1
        self.ln_1 = LayerNorm(n_embd)
        # 2. 多头自注意力机制 (我们在上一篇详细讲过)
        # 这里假设 CausalSelfAttention 已经实现好
        self.attn = CausalSelfAttention(n_embd, n_head)
        # 3. 归一化层 2
        self.ln_2 = LayerNorm(n_embd)
        # 4. 前馈网络
        self.mlp = FeedForward(n_embd)

    def forward(self, x):
        # 注意这里的结构是 Pre-LN:先做 LayerNorm,再输入到模块,最后加上残差
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

4.3 代码运行结果

我们在本地测试了上面拼装好的完整的 Transformer Block。在配置为 Batch=2, SeqLen=8, EmbedDim=32, Heads=4 的情况下,运行结果如下:

正在初始化配置: Batch=2, SeqLen=8, EmbedDim=32, Heads=4

[1] 原始输入特征 x 的形状: torch.Size([2, 8, 32])
[2] 经过位置编码后的特征形状: torch.Size([2, 8, 32]) (未改变维度,数值被注入了时序特征)
[3] 经过一个完整的 Transformer Block (Pre-LN) 后的输出形状: torch.Size([2, 8, 32])

✅ 运行成功!模型的所有组件能够完美衔接工作。

从输出结果可以看出,无论内部的自注意力机制和前馈网络进行了多么复杂的空间映射和维度缩放,由于残差连接的巧妙设计,整个 Transformer Block 最终输出的张量形状与输入完全一致。这就是为什么我们可以像搭积木一样,将几十个甚至上百个这样的 Block 串联堆叠在一起,构建出拥有千亿参数的 GPT 大模型!


5. 总结

在这篇文章中,我们为裸奔的 Self-Attention 穿上了神装:

  1. Positional Encoding 赋予了模型时间的概念,让模型不再脸盲。
  2. Layer Norm 稳定了数值波动,降伏了狂暴的梯度。
  3. FFN 提供了非线性的推理空间,让模型拥有了真正的“思考”能力。
  4. 残差连接 (Residual Connection) 搭建了信息的高速公路,使得几十上百层的网络依然能够顺畅反向传播。

在下一篇中,我们将进入实战环节:预训练与微调。我们将揭秘大模型是如何通过 Masked LM 或 Causal LM 目标函数“读遍天下书”的,并深入探讨 LoRA/P-Tuning 等让普通玩家也能玩转大模型的高效微调黑科技。敬请期待!


参考文献与延伸阅读

  1. Attention Is All You Need (Vaswani et al., 2017):Transformer 架构的开山之作。
  2. NanoGPT by Andrej Karpathy:用极其简洁优雅的代码复现了 GPT,是学习大模型底层实现最好的开源项目。
  3. Layer Normalization (Ba et al., 2016):详细阐述了为什么在 RNN 和 NLP 任务中 Layer Norm 优于 Batch Norm。
  4. On Layer Normalization in the Transformer Architecture (Xiong et al., 2020):论证了 Pre-LN 相对于 Post-LN 的稳定性优势。
博客日历
2026年04月
SuMoTuWeThFrSa
29
30
31
01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
01
02
03
04
05
06
07
08
09
更多
--
--
--
--