4. 从零实现Transformer编码器：多头注意力层、前馈神经网络层、残差连接与层归一化

好，咱们今天来点硬核的。前面几章我们把Transformer的理论框架搭起来了，但说实话，光看公式和论文，你很难真正理解它的精髓。我个人习惯是——必须亲手把代码敲出来，哪怕是最基础的版本。

这一章，我们就从零开始，手撸一个Transformer编码器。别怕，我会带着你一步步走。核心就三个模块：多头注意力、前馈神经网络、还有残差连接+层归一化。嗯，这三个东西拼起来，就是编码器的全部秘密。

核心观点：Transformer编码器 = 多头自注意力 + 前馈网络 + 残差连接 + 层归一化。没有花哨的技巧，但每一步都值得深究。

Transformer 编码器结构图 输入序列 (X) 多头自注意力 (Multi-Head Attention) Q, K, V 拆分 → 缩放点积 → 拼接 残差连接 + 层归一化 (Add & Norm) X + Attention(X) → LayerNorm 前馈神经网络 (FFN) Linear → ReLU → Linear 残差连接 + 层归一化 (Add & Norm) X + FFN(X) → LayerNorm 编码器输出 第 1 个子层 第 2 个子层

4.1 多头注意力层：让模型学会「多角度观察」

先聊聊多头注意力。你想想看，如果只用一个注意力头，模型只能从一个角度去理解序列中词与词的关系。但现实世界很复杂——一个词可能同时和多个词有不同层面的关联。比如「苹果」这个词，在「苹果很好吃」和「苹果发布了新手机」里，它关注的对象完全不同。

多头注意力的思路很简单：把Q、K、V拆成多份，每份独立做注意力计算，最后再拼回去。这样每个头可以学到不同的注意力模式。

我的经验：头数一般设为8或16，但别盲目增加。我在一个项目里试过32个头，结果训练速度慢了一倍，效果反而没提升。8个头通常是个不错的起点。

来看代码实现。我会用PyTorch风格，但核心逻辑是框架无关的。

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model必须能被n_heads整除"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度

        # 线性变换层：生成Q、K、V
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)  # 输出投影

        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 1. 线性变换 + 拆分成多头
        # 形状: (batch, seq_len, d_model) -> (batch, seq_len, n_heads, d_k)
        Q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 2. 缩放点积注意力
        # 计算注意力分数: Q * K^T / sqrt(d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 加权求和得到输出
        attn_output = torch.matmul(attn_weights, V)

        # 3. 拼接所有头 + 输出投影
        # (batch, n_heads, seq_len, d_k) -> (batch, seq_len, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        output = self.W_o(attn_output)

        return output, attn_weights

这段代码里有个细节我想强调一下：缩放因子 sqrt(d_k)。为什么要除这个？因为当d_k比较大时，点积的结果会很大，softmax的梯度会变得非常小，模型就学不动了。这个缩放操作，说白了就是让注意力分布更「平滑」一些。

我曾经踩过的坑：在实现时忘记对mask做 -1e9 的填充，而是直接用了0。结果模型在解码阶段疯狂关注padding位置，训练损失死活降不下去。后来排查了半天才发现是mask没处理好。记住：被mask的位置，在softmax之前要变成负无穷。

4.2 前馈神经网络层：给模型加点「非线性」

注意力层本质上是一个加权求和，它是线性的。如果只有注意力层，那整个模型就退化成线性模型了，表达能力有限。所以我们需要一个前馈神经网络（FFN）来引入非线性变换。

标准的FFN结构很简单：两个线性层，中间夹一个ReLU激活函数。第一个线性层把维度从d_model映射到d_ff（通常是4倍），第二个线性层再映射回来。

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 线性变换 -> ReLU -> Dropout -> 线性变换
        return self.linear2(self.dropout(F.relu(self.linear1(x))))

就这么简单？对，就这么简单。但你别小看它。我见过很多初学者觉得FFN「太简单了，没啥好学的」，结果在调参时完全不知道d_ff该设多大。

参数 常见取值 说明
d_model 512 / 768 / 1024 模型整体维度，也是输入输出的维度
d_ff 2048 / 3072 / 4096 FFN中间层维度，通常是d_model的4倍
dropout 0.1 防止过拟合，训练时随机丢弃一部分神经元

我个人的习惯：d_ff设成d_model的4倍，这是一个经过大量实验验证的经验值。如果你资源有限，可以试试2倍或3倍，效果可能会差一些，但训练速度快不少。

4.3 残差连接与层归一化：让训练更稳定

好，现在两个核心模块都有了。但直接堆叠它们会出问题——梯度消失。层数一深，底层的梯度就传不上来了。怎么办？两个技巧：残差连接和层归一化。

残差连接：把输入直接加到输出上。这样梯度可以「抄近道」直接传到前面。公式很简单：output = LayerNorm(x + sublayer(x))。

层归一化：对每个样本的所有特征做归一化，让数据分布更稳定。和BatchNorm不同，LayerNorm不依赖batch大小，特别适合序列模型。

class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        # 计算均值和方差
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        # 归一化 + 缩放平移
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

嗯，这里要注意：eps是个很小的数，防止除以0。我一般设1e-6，太小了数值不稳定，太大了又会影响归一化效果。

4.4 组装编码器：把一切拼起来

现在我们把三个模块组装成一个完整的编码器层。一个编码器层 = 多头注意力 + 残差连接 + 层归一化 + FFN + 残差连接 + 层归一化。

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 子层1：多头自注意力 + 残差连接 + 层归一化
        attn_output, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 子层2：前馈网络 + 残差连接 + 层归一化
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))

        return x

注意看，我用了x + self.dropout(attn_output)这种写法。Dropout加在残差连接之前，这是原论文的做法。有些实现会把Dropout放在后面，效果差别不大，但保持一致比较好。

关键点总结：

  • 多头注意力：拆分QKV -> 独立计算 -> 拼接输出
  • FFN：升维 -> ReLU -> 降维，引入非线性
  • 残差连接：解决深层网络梯度消失问题
  • 层归一化：稳定训练过程，加速收敛

最后，如果你想堆叠多个编码器层，只需要循环调用这个TransformerEncoderLayer就行。我见过有人堆了6层，也有人堆了12层甚至24层。层数越多，模型容量越大，但训练难度也越大。我个人建议从6层开始试，效果不够再加。

好了，编码器的核心实现就这些。你可能会觉得代码量不大，但每一行都有它的道理。我当年第一次手撸Transformer时，光是调试多头注意力的维度匹配就花了一整天。嗯，如果你在实现过程中遇到问题，不妨回头看看我标注的那些注意事项——大概率能帮你省下不少时间。

公众号:蓝海资料掘金营，微信deep3321
← 上一章 📖 返回目录 下一章 →