8.6 Transformer Encoder：把 Self-Attention 堆起来

Author

jshn9515

Published

2026-05-03

Modified

2026-05-03

前面几节里，我们已经把 Transformer 里最核心的几个组件拆开看过了。

Self-attention 让序列中的每个 token 都可以直接和其他 token 交换信息；multi-head attention 让模型可以从多个表示空间并行地观察序列；positional encoding 则补上了 self-attention 本身缺少的顺序信息。

但到目前为止，我们看到的还只是一些单独的组件。真正的 Transformer Encoder 并不是只做一次 self-attention，而是把这些组件组织成一个可以反复堆叠的模块。这个模块通常叫做 Transformer Encoder Block，或者简称 Encoder Block 或 Encoder Layer。

从整体上看，一个 Transformer Encoder 就是多个 encoder block 堆起来：

\[ X \rightarrow \text{Embedding} + \text{Positional Encoding} \rightarrow \text{Encoder Block}_1 \rightarrow \cdots \rightarrow \text{Encoder Block}_N \]

这一节我们就来看看，Transformer encoder 是怎么把 self-attention 变成一个完整网络结构的。

from collections.abc import Callable

import dnnl.nn as dnn
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

print('PyTorch version:', torch.__version__)
PyTorch version: 2.12.0+xpu

8.6.1 为什么要堆叠多层 Self-Attention？

先看一个直觉问题：既然一层 self-attention 已经可以让每个 token 看到所有 token，为什么还要堆很多层？

原因是：看到所有 token，并不等于已经完成了足够复杂的理解。

一层 self-attention 做的事情，主要是让每个位置根据当前表示，从其他位置聚合信息。它确实建立了全局连接，但这种连接只是一次信息混合。对于复杂语言现象，模型通常需要经过多轮加工，才能形成更高层的表示。

比如下面这个句子：

The animal didn’t cross the street because it was too tired.

模型要理解 it 指的是 animal，可能不只是看 it 和 animal 的相似度这么简单。它还需要综合语义、句法结构、常识关系以及上下文中的描述。第一层可能先捕捉局部词义和短语关系；更高层可能进一步组合这些信息，形成更抽象的句子级理解。

所以，多层 Transformer Encoder 的作用不是简单重复同一个操作，而是让表示逐层变得更抽象。

我们可以粗略理解为：

  • 底层更偏向局部和词级信息；
  • 中间层逐渐组合短语和上下文关系；
  • 高层更偏向任务相关的语义表示。

当然，这只是一个直观解释，不代表每一层都有固定的人类可以理解和解释的功能。但从网络结构上看，堆叠多层确实给了模型反复加工信息的能力。这其实和堆叠 CNN 或 RNN 是同一个道理。

8.6.2 Encoder Block 的整体结构

一个标准的 Transformer Encoder block 主要包含两个子模块：

  1. Multi-Head Self-Attention
  2. Position-wise Feed-Forward Network

除此之外，每个子模块外面还会配上两个非常重要的结构：

  1. Residual Connection
  2. Layer Normalization

所以，一个 encoder block 可以写成下面这个形式：

\[ H = \operatorname{LayerNorm}(X + \operatorname{MultiheadAttention}(X)) \]

\[ Y = \operatorname{LayerNorm}(H + \operatorname{FFN}(H)) \]

这里的 \(X\) 是当前层的输入表示，\(Y\) 是当前层的输出表示，FFN 是 Position-wise Feed-Forward Network。

从数据流上看，就是：

图 1：Transformer Encoder Block 结构图
图 1：Transformer Encoder Block 结构图

这就是最基本的 encoder block 结构。

需要注意的是，上面的写法是比较容易理解的 Post-LN 形式，也就是先做子模块，再做残差连接和 LayerNorm。原始 Transformer 使用的就是这种结构。后来很多现代 Transformer 会使用 Pre-LN 结构，也就是先对输入做 LayerNorm，再送入子模块。我们后面写代码时会采用更常见也更稳定的 Pre-LN 写法。

8.6.2.1 第一部分：Multi-Head Attention

Encoder block 的第一部分是 multi-head attention。需要注意的是，这里的 attention 是 self-attention，\(Q\)、\(K\)、\(V\) 都来自同一个输入 \(X\)：

\[ Q = XW_Q,\quad K = XW_K,\quad V = XW_V \]

然后计算：

\[ \operatorname{MultiheadAttention}(Q,K,V) \]

其中，multi-head attention 里的每个 head 都是在不同的投影空间里做 scaled dot-product attention。

这一部分的作用是让不同 token 之间交换信息。经过 self-attention 之后，每个 token 的表示不再只包含它自己的词义和位置信息，而是融合了整句上下文的信息。所以，通过不断堆叠 encoder block，就可以把一串原本相对孤立的 token embedding，逐步变成一串上下文化的 token representation。

8.6.2.2 第二部分：Residual Connection

在 multi-head attention 之后，Transformer 不会直接把 attention 的输出送到下一步，而是会加上原来的输入：

\[ X + \operatorname{MultiheadAttention}(Q,K,V) \]

这就是残差连接（Residual Connection），也是我们前面在 ResNet 里学过的一个重要结构。

这样做有两个好处。第一，它可以保留原来的信息。如果某些信息暂时不需要被修改，残差路径可以让它更容易传递到后面的层。第二，它可以缓解深层网络训练困难的问题。如果没有残差连接，每一层的输出都必须经过子模块本身，原来的输入信息必须由这个子模块学习并保留下来。随着 encoder block 越堆越深，这会让信息和梯度的传递都变得更困难。残差连接提供了一条更直接的信息通路，让原表示可以继续向后传递，而子模块只需要学习在原表示基础上的增量修改，让深层 Transformer 更容易训练。

所以，Transformer encoder 不是简单地一层覆盖一层，而是每一层都在原表示上做逐步修正。

8.6.2.3 第三部分：Layer Normalization

残差连接之后，Transformer 通常会接一个 Layer Normalization。LayerNorm 的作用是对每个 token 的特征维度做归一化。假设某个 token 的表示是：

\[ X \in \mathbb{R}^{d_\mathrm{model}} \]

LayerNorm 会在这个向量内部计算均值和方差，然后做标准化：

\[ \operatorname{LayerNorm}(X) = \gamma \frac{X - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta \]

其中，\(\gamma\) 和 \(\beta\) 是可学习参数。

直观来说，LayerNorm 可以让每一层的表示尺度更加稳定。否则，经过多层 attention、FFN 和残差连接之后，激活值的均值和方差可能不断变化，表示的尺度也可能逐层放大或缩小，这会让深层网络更难训练。LayerNorm 通过对每个 token 的特征维度进行归一化，帮助模型维持相对稳定的激活分布，从而缓解训练不稳定的问题。

至于为什么 Transformer 里更常用 LayerNorm，而不是 BatchNorm，主要是因为 LayerNorm 不依赖 batch 统计量。对于变长序列来说，batch 里的句子长度、padding 数量和 token 位置都可能不同，如果用 BatchNorm，归一化结果会受到当前 batch 组成的影响。LayerNorm 则是对每个 token 自己的特征维度做归一化，因此更适合 Transformer 这类序列模型。

8.6.2.4 第四部分：Position-wise Feed-Forward Network

Encoder block 的第四部分是 feed-forward network，通常简称 FFN。它的形式很简单：

\[ \operatorname{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2 \]

其中，\(\sigma\) 是非线性激活函数。原始 Transformer 使用的是 ReLU，现代模型中也经常使用 GELU 或 SiLU。

如果输入维度是 \(d_\mathrm{model}\)，FFN 通常会先把维度扩展到更大的 \(d_\mathrm{ff}\)，再投影回 \(d_\mathrm{model}\)：

\[ d_\mathrm{model} \rightarrow d_\mathrm{ff} \rightarrow d_\mathrm{model} \]

比如常见设置是：

\[ d_\mathrm{ff} = 4d_\mathrm{model} \]

也就是说，FFN 先把每个 token 的表示映射到一个更高维的空间里做非线性变换，再压回原来的维度。

这里我们要注意一个关键词：position-wise。Position-wise 的意思是，这个 FFN 会独立地作用在每一个位置上。它不会直接让不同 token 之间交换信息。也就是说，对于输入：

\[ H = [h_1, h_2, \dots, h_n] \]

FFN 做的是：

\[ \operatorname{FFN}(H) = [\operatorname{FFN}(h_1), \operatorname{FFN}(h_2), \dots, \operatorname{FFN}(h_n)] \]

不同位置使用同一套 FFN 参数，但计算时彼此独立。

到此，我们可以把 encoder block 中两大模块的分工理解成：

  • Self-attention 负责 token 之间的信息交换；
  • FFN 负责对每个 token 的表示做进一步非线性加工。

一个负责混合上下文，一个负责加工表示。两者组合起来，才形成完整的 Transformer Encoder block。

8.6.3 Pre-LN 和 Post-LN

前面为了理解方便，我们把 encoder block 写成了：

\[ H = \operatorname{LayerNorm}(X + \operatorname{MultiheadAttention}(X)) \]

\[ Y = \operatorname{LayerNorm}(H + \operatorname{FFN}(H)) \]

这种写法叫 Post-LN，因为 LayerNorm 放在子模块和残差连接之后。

但在很多现代 Transformer 实现中，更常见的是 Pre-LN：

\[ H = X + \operatorname{MultiheadAttention}(\operatorname{LayerNorm}(X)) \]

\[ Y = H + \operatorname{FFN}(\operatorname{LayerNorm}(H)) \]

也就是先对输入做 LayerNorm，再送入 attention 或 FFN，最后再加回残差。

从概念上讲，Pre-LN 和 Post-LN 的核心结构是一致的：都有 attention、FFN、残差连接和 LayerNorm。区别只是 LayerNorm 放在子模块之前还是之后。研究发现，Pre-LN 相比 Post-LN 的一个好处是训练更稳定，尤其是在模型很深的时候。因为残差路径更加直接，梯度可以更顺畅地沿着残差分支传播。

为了使代码更接近现代实现，下面我们采用 Pre-LN 写法。

8.6.4 Transformer Encoder Block 的 PyTorch 实现

下面我们实现一个简化版 Transformer Encoder block。这里为了避免重复实现 multi-head attention，我们直接使用 8.4 节里实现的 MultiheadAttention。

class TransformerEncoderLayer(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        dim_feedforward: int = 2048,
        activation: Callable[[Tensor], Tensor] = F.relu,
        bias: bool = True,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.self_attn = dnn.MultiheadAttention(
            d_model,
            num_heads,
            bias=bias,
            dropout=dropout,
        )

        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias)
        self.activation = activation

        self.norm1 = nn.LayerNorm(d_model, bias=bias)
        self.norm2 = nn.LayerNorm(d_model, bias=bias)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(
        self,
        src: Tensor,
        src_mask: Tensor | None = None,
        src_key_padding_mask: Tensor | None = None,
    ) -> Tensor:
        x = src + self._sa_block(
            self.norm1(src),
            attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask,
        )
        x = x + self._ff_block(self.norm2(x))
        return x

    def _sa_block(
        self,
        x: Tensor,
        attn_mask: Tensor | None,
        key_padding_mask: Tensor | None,
    ) -> Tensor:
        x, _ = self.self_attn(
            x, x, x,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=False,
        )  # fmt: skip
        x = self.dropout1(x)
        return x

    def _ff_block(self, x: Tensor) -> Tensor:
        x = self.linear1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear2(x)
        x = self.dropout2(x)
        return x


x = torch.randn(2, 32, 512)  # (batch_size, seq_len, d_model)
encoder_layer = TransformerEncoderLayer(d_model=512, num_heads=8)

with torch.inference_mode():
    output = encoder_layer(x)

print('Encoder block output shape:', output.shape)
Encoder block output shape: torch.Size([2, 32, 512])

这段代码里最关键的是两条残差路径：

x = src + self._sa_block(...)

以及：

x = x + self._ff_block(...)

第一条残差对应 multi-head self-attention，第二条残差对应 FFN。

8.6.5 堆叠多个 Encoder Block

有了单个 encoder block 以后，完整的 Transformer Encoder 就很容易写了。我们只需要把它重复堆叠多次。

class TransformerEncoder(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        num_layers: int,
        dim_feedforward: int = 2048,
        activation: Callable[[Tensor], Tensor] = F.relu,
        bias: bool = True,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.layers = nn.ModuleList(
            [
                TransformerEncoderLayer(
                    d_model=d_model,
                    num_heads=num_heads,
                    dim_feedforward=dim_feedforward,
                    activation=activation,
                    bias=bias,
                    dropout=dropout,
                )
                for _ in range(num_layers)
            ]
        )
        self.norm = nn.LayerNorm(d_model)

    def forward(
        self,
        src: Tensor,
        mask: Tensor | None = None,
        src_key_padding_mask: Tensor | None = None,
    ) -> Tensor:
        output = src
        for layer in self.layers:
            output = layer(
                output,
                src_mask=mask,
                src_key_padding_mask=src_key_padding_mask,
            )

        output = self.norm(output)
        return output


x = torch.randn(2, 32, 512)  # (batch_size, seq_len, d_model)
encoder = TransformerEncoder(d_model=512, num_heads=8, num_layers=6)

with torch.inference_mode():
    output = encoder(x)

print('Encoder output shape:', output.shape)
Encoder output shape: torch.Size([2, 32, 512])

需要注意的是，这个 encoder 的输入 x 是已经经过 token embedding 和 position embedding 之后的表示。也就是说，在送入 encoder 之前，我们需要先把 token id 转换成 embedding 向量，并加上位置编码。

原始的序列输入是 token id，形状通常是：

(batch_size, seq_len)

经过 token embedding 和 position embedding 之后，变成：

(batch_size, seq_len, d_model)

然后依次通过多个 encoder block，最后输出仍然是：

(batch_size, seq_len, d_model)

这意味着 encoder 会为输入序列中的每个 token 输出一个上下文化表示。比如，我们输入的一句话有 10 个 token，那么 encoder 输出也会有 10 个向量。每个向量都对应一个位置，但已经融合了整句话的上下文信息。

8.6.6 Padding Mask：忽略补齐出来的位置

在实际训练中，一个 batch 里的句子长度通常不一样。为了把它们拼成一个张量，我们会把短句子补齐到相同长度。

比如：

I love deep learning [PAD] [PAD]
Transformers are very powerful models

这里的 [PAD] 只是为了对齐长度，并不是真正的输入内容。模型在做 self-attention 时，不应该关注这些 padding token。但是，我们要怎么告诉模型哪些位置是 padding 呢？

这就需要 key_padding_mask。

在 PyTorch 的 nn.MultiheadAttention 里，key_padding_mask 的形状通常是：

(batch_size, seq_len)

其中，True 表示这个位置应该被 mask 掉，也就是不让其他 token 注意到它。

例如：

key_padding_mask = input_ids == pad_token_id

然后传入 Encoder：

output = encoder(input_ids, key_padding_mask=key_padding_mask)

这样 self-attention 在计算时，就会忽略 padding 位置。

这一点在 NLP 任务里非常重要。否则模型可能会把 [PAD] 也当成有意义的 token，从而影响表示学习。

8.6.7 Encoder 的输出可以用来做什么

Transformer Encoder 的输出是一组上下文化 token 表示：

\[ H = [h_1, h_2, \dots, h_n] \]

不同任务会用这些输出做不同事情。

如果是序列标注任务，比如命名实体识别，我们通常会使用每个位置的输出：

\[ h_1, h_2, \dots, h_n \]

因为每个 token 都需要一个预测结果。

如果是句子分类任务，比如情感分类，我们通常需要把整句话表示成一个向量。常见做法有两种。

第一种是加入一个特殊的 [CLS] token，然后使用 [CLS] 位置的输出作为整句表示：

\[ h_{\mathrm{cls}} \]

第二种是对所有 token 的输出做 pooling，比如 average pooling：

\[ h_{\mathrm{avg}} = \frac{1}{n}\sum_{i=1}^{n}h_i \]

这也是为什么 Transformer Encoder 可以用于很多不同任务。它本身负责把输入序列编码成上下文化表示；至于后面接分类头、标注头还是检索头，则取决于具体任务。

8.6.8 本章小结

这一节里，我们把前面学过的 self-attention、multi-head attention 和 positional encoding 组合了起来，形成了完整的 Transformer encoder。

一个 encoder block 主要由两部分组成：multi-head self-attention 和 position-wise feed-forward network。前者负责让不同 token 之间交换上下文信息，后者负责对每个 token 的表示做非线性加工。每个子模块外面还会配上残差连接和 LayerNorm，让深层网络更容易训练。

把多个 encoder block 堆叠起来以后，模型就可以对输入序列进行多轮信息交互和表示加工。最终输出的每个 token 向量，都不再只是它自己的 embedding，而是融合了整个上下文后的表示。

到这里，我们已经看完了 Transformer Encoder 的核心结构。下一节我们继续往前走，看 Transformer decoder。它和 encoder 很像，但多了一个关键限制：在生成时，当前位置不能偷看未来的 token，否则模型就不是在预测下一个词，而是在抄答案。这就引出了 masked self-attention。