Attention on 酒中仙

深入解析 Transformer 架构

Sat, 07 Feb 2026 20:34:25 -0800

深入解析 Transformer

注意力机制通过动态加权的方式，克服了传统 Seq2Seq 模型中的“信息瓶颈”问题。但是，这些模型依然依赖于 RNN 来处理序列信息，也就是说它们必须按顺序，一个词元接一个词元地进行计算，这在处理长序列时效率低下，并且存在长距离依赖信息丢失的问题。

2017年，Google 的研究团队发表了一篇名为《Attention Is All You Need》的论文，提出了一种全新的架构——Transformer ¹。这篇论文的标题很有冲击力，其思想也同样有颠覆性。它抛弃了传统的 RNN 和卷积网络，整个模型基于注意力机制来构建。Transformer 的提出在自然语言处理领域具有划时代的意义。它不仅凭借其出色的并行计算能力极大地提升了训练效率，还更有效地捕捉了文本中的长距离依赖关系，为后续的 BERT、GPT 等大规模预训练模型的诞生提供了架构基础。

一、自注意力机制

从根本上说，要让模型理解一段文本，就需要提取其“序列特征”，即将文本中所有词元的信息以某种方式整合起来。RNN 通过依次传递隐藏状态来顺序地整合信息，而 Transformer 则选择了一条截然不同的道路。其核心是 自注意力机制。它不再依赖于顺序计算，而是将提取序列特征的过程看作是输入序列“自己对自己进行注意力计算”。序列中的每个词元都会“审视”序列中的所有其他词元，来动态地计算出最能代表当前词元上下文含义的新表示。与上一节介绍的交叉注意力不同，在自注意力中，Query、Key、Value 均来源于同一个输入序列。

举个例子，在句子“苹果公司发布了新款手机，它采用了最新的芯片”中，要理解代词“它”指的是“新款手机”而不是“苹果公司”，模型就需要将“它”与句子中的其他词元进行关联。自注意力机制正是通过计算“它”对句中其他所有词的注意力权重来实现这一点的。

1.1 自注意力与交叉注意力的区别

从结构上看，自注意力与交叉注意力的区别在于信息的来源和流动方向。在交叉注意力机制中，信息在两个不同的序列之间流动。通常，Query 来自解码器（代表当前的目标序列状态），而 Key 和 Value 来自编码器的所有输出（代表完整的源序列信息）。其目的是在生成目标序列的每一步时，从源序列中寻找最相关的信息。

而在自注意力机制中，信息则是在同一个序列内部进行流动和重组。它的 Query, Key, 和 Value 都来自同一个输入序列。其目的是为了捕捉输入序列内部的依赖关系，重新计算序列中每个词元的表示，使其包含更丰富的上下文信息。

总结来说，尽管底层的加权求和计算方式相似，但两者在架构上的目标完全不同：

交叉注意力：用于对齐和整合两个不同序列之间的信息。
自注意力：用于理解和重构单个序列内部的依赖关系。

1.2 自注意力的计算过程

自注意力的计算过程与上一节介绍的 QKV 范式完全一致，关键区别在于 Q, K, V 的来源。

（1）生成 Q, K, V 向量：

对于输入序列中的每一个词元，首先获取其词嵌入向量 $x_i$。然后，将该向量分别与三个可学习的、在整个模型中共享的权重矩阵 $W^Q, W^K, W^V$ 相乘，生成该词元专属的 Query 向量 $q_i$、Key 向量 $k_i$ 和 Value 向量 $v_i$。

$$ q_i = x_i W^Q \\ k_i = x_i W^K \\ v_i = x_i W^V $$

这三个矩阵的作用是将原始的词嵌入向量投影到不同的、专门用于注意力计算的表示空间中，赋予了模型更大的灵活性。

（2）计算注意力分数：

为了计算第 $i$ 个词元的新表示，需要用它的 Query 向量 $q_i$ 去和所有词元（包括它自己）的 Key 向量 $k_j$ 计算点积，得到注意力分数。

$$ \text{score}(i, j) = q_i \cdot k_j $$

（3）缩放与归一化：

将得到的分数除以一个缩放因子 $\sqrt{d_k}$（$d_k$ 是 Key 向量的维度），然后通过 Softmax 函数进行归一化，得到最终的注意力权重 $\alpha_{ij}$。这个缩放步骤的目的与上一节中介绍的一致，都是为了在训练过程中保持梯度稳定。当向量维度 $d_k$ 较大时，点积结果的方差会增大，可能将 Softmax 函数推向其梯度极小的区域，从而导致梯度消失，影响模型学习。进行缩放可以有效缓解这个问题。

$$ \alpha_{ij} = \text{softmax}\left(\frac{q_i \cdot k_j}{\sqrt{d_k}}\right) $$

加权求和：

使用计算出的权重 $\alpha_{ij}$ 对所有词元的 Value 向量 $v_j$ 进行加权求和，得到第 $i$ 个词元经过自注意力计算后得到的新表示 $z_i$。

$$ z_i = \sum_j \alpha_{ij} v_j $$

通过这个过程，输出向量 $z_i$ 不再仅仅包含原始词元 $x_i$ 的信息，而是融合了整个序列中所有与之相关词元的信息，成为一个上下文感知的、更丰富的表示。其本质可以理解为：序列中的每个词元都同时扮演着“查询（Q）”、“键（K）”和“值（V）”三种角色。通过计算查询与其他所有词元的键之间的相关性，来决定如何加权融合所有词元的值，从而为每个词元生成一个全新的、深度融合了全局上下文信息的表示。

既然 Q, K, V 都来自同一个输入 X，为什么不直接用 X 计算，而要引入三个独立的权重矩阵 $W^Q, W^K, W^V$？甚至，为什么是三个，而不是两个或四个？

这可以类比在图书馆查资料的过程：

Query (Q) - 要问的问题：代表了我们主动想查询的意图。

Key (K) - 书的索引/标签：代表了书本内容的关键特征，用于被动地和你的问题进行匹配。

Value (V) - 书的具体内容：代表了书本实际包含的信息。

我们的“问题”和书本的“索引”可能都源于同一个知识领域（同一个输入 X），但它们在信息检索这个任务中扮演的角色是截然不同的。$W^Q, W^K, W^V$ 这三个矩阵的作用，就是让模型学会将原始输入 X 投影到三个功能不同的空间中，分别去扮演好“查询者”、“被查询的索引”和“信息提供者”这三种角色。Q-K 配对解决了“如何定位相关信息”的问题，而 V 提供了“应该提取什么信息”的答案。这个三元组结构在功能上是完备且高效的，所以成为了注意力机制的标准范式。

1.3 矩阵运算与并行化

上述步骤描述的是单个词元 $i$ 的计算过程。在实际应用中，如果采用循环的方式逐个计算每个词元的 $z_i$，效率会非常低下。自注意力的巨大优势在于其并行计算能力，这通过将整个过程表达为矩阵运算来实现。

假设整个输入序列的词嵌入矩阵为 $X$（维度为 [sequence_length, embedding_dim]），可以一次性计算出所有词元的 Q, K, V 矩阵：

$$ Q = X W^Q \\ K = X W^K \\ V = X W^V $$

然后，整个自注意力的输出矩阵 $Z$ 可以通过一个公式完成计算：

$$ Z = \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

这个公式与上一节中介绍的通用注意力公式完全相同。这里的主要区别不在于数学运算，而在于输入的来源：

在上一节的交叉注意力中，Q 来自一个序列（解码器），而 K 和 V 来自另一个序列（编码器）。
在当前的自注意力中，矩阵 Q、K 和 V 全部派生自同一个输入序列 X。

所以，同一个数学范式，根据输入来源的不同，被用于解决两个不同的问题，一个是两个序列之间的对齐，另一个是单个序列内部的依赖关系建模。在这个公式中， $QK^T$ 的计算结果是一个维度为 [sequence_length, sequence_length] 的注意力分数矩阵，其中第 $i$ 行第 $j$ 列的元素表示第 $i$ 个词元对第 $j$ 个词元的注意力分数（未归一化的 logits）。注意力权重来自对缩放分数应用 Softmax 后得到的归一化系数。

1.4 PyTorch 实现自注意力

本节完整代码

从概念上讲，自注意力的计算可以分解为对序列中每个词元进行循环操作，这种方式虽然直观但效率极低。因此，现代深度学习框架中的实现都采用了矩阵运算的方式。通过将整个序列的 Q, K, V 看作矩阵，利用一次大规模的矩阵乘法，就能并行地完成所有词元之间的相关性计算。下面是这种并行化版本的实现：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


class SelfAttention(nn.Module):
 """自注意力模块"""
 def __init__(self, hidden_size):
 super(SelfAttention, self).__init__()
 self.hidden_size = hidden_size
 self.q_linear = nn.Linear(hidden_size, hidden_size)
 self.k_linear = nn.Linear(hidden_size, hidden_size)
 self.v_linear = nn.Linear(hidden_size, hidden_size)

 def forward(self, x):
 q = self.q_linear(x)
 k = self.k_linear(x)
 v = self.v_linear(x)

 scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.hidden_size)
 attention_weights = torch.softmax(scores, dim=-1)
 context = torch.matmul(attention_weights, v)

 return context

__init__: 初始化了三个 nn.Linear 层，它们分别对应将输入映射到 Q, K, V 空间的权重矩阵 $W^Q, W^K, W^V$。
forward:
- q_linear(x), k_linear(x), v_linear(x)：将形状为 [batch_size, seq_len, hidden_size] 的输入张量 x 分别通过三个线性层，一次性地为序列中的所有词元计算出 Q, K, V 矩阵。
- torch.matmul(q, k.transpose(-2, -1)): 这是实现并行计算的核心。通过将 K 矩阵的最后两个维度转置（seq_len, hidden_size -> hidden_size, seq_len），再与 Q 矩阵相乘，直接得到了一个 [batch_size, seq_len, seq_len] 的分数矩阵。该矩阵中的 scores[b, i, j] 代表了批次 b 中第 i 个词元对第 j 个词元的注意力分数。
- / math.sqrt(self.hidden_size)：执行缩放操作，防止梯度消失。
- torch.softmax(scores, dim=-1)：对分数的最后一个维度（seq_len）进行 Softmax，得到归一化的注意力权重。
- torch.matmul(attention_weights, v)：将权重矩阵与 V 矩阵相乘，完成了对所有词元的 Value 向量的加权求和，得到最终的上下文感知表示。

二、多头注意力机制

仅仅用一组 $W^Q, W^K, W^V$ 矩阵进行一次自注意力计算，相当于只从一个“视角”来审视文本内在的关系。然而，文本中的关系是多层次的，例如，一组参数可能学会了关注代词（如 “它” 指向谁）的关系，但可能忽略了动作的执行者（主谓宾）等其他类型的关系。

为了让模型能够综合利用从不同维度和视角提取出的信息，Transformer 引入了多头注意力机制 (Multi-Head Attention)。其思想非常直接：并行地执行多次自注意力计算，每一次计算都是一个独立的“头 (Head)”。每个头都拥有一组自己专属的 $W^Q_i, W^K_i, W^V_i$ 权重矩阵，并且可以学习去关注一种特定类型的上下文关系。

那么，多头注意力与我们之前讨论的“增加 A, B, C 等新角色”有什么不同呢？

一个关键的区别：多头注意力不是通过增加 A, B, C 等新角色来深化单次注意力计算的复杂性，而是通过并行运行多个独立的 QKV 计算单元来拓宽其广度。

再次使用图书馆的类比：

增加 A, B, C：相当于给一个图书管理员一套更复杂的工具，让他一次性处理问题(Q)、索引(K)、内容(V)之外，还要考虑主题(A)、背景(B)等，这会使单次查询过程变得非常复杂。

多头注意力：相当于雇佣一个各有所长的专家团队（比如 8 个管理员，即 8 个“头”）。每个专家都只使用标准高效的 QKV 工具，但他们各自有独特的视角（独立的 $W^Q_i, W^K_i, W^V_i$ 矩阵）。一个专家可能专攻语法，另一个专攻语义。最后，将所有专家的报告汇总起来，得到一个更全面、更丰富的结论。

因此，多头注意力机制为模型提供了从不同子空间、不同视角审视信息的能力，而不是改变注意力计算本身的范式。

具体流程如下：

（1）并行计算：假设有 $h$ 个头，那么就初始化 $h$ 组不同的权重矩阵 $(W^Q_0, W^K_0, W^V_0), (W^Q_1, W^K_1, W^V_1), \dots, (W^Q_{h-1}, W^K_{h-1}, W^V_{h-1})$。

（2）独立注意力：对于输入序列，每个头都独立地执行一次完整的自注意力计算，产生一个输出矩阵 $Z_i$。

（3）拼接与投影：将所有 $h$ 个头的输出矩阵 $Z_0, Z_1, \dots, Z_{h-1}$ 在特征维度上进行拼接 (Concatenate)。

（4）最终输出：将拼接后的巨大矩阵乘以一个新的权重矩阵 $W^O$，将其投影回原始的输入维度，得到多头注意力机制的最终输出。

多头机制允许模型在不同的表示子空间中共同学习上下文信息。例如，一个头可能专注于捕捉长距离的语法依赖，而另一个头可能更关注局部的词义关联。这种设计极大地增强了模型的表达能力。

在实践中，为了保持计算总量不变，通常会将原始的词嵌入维度 embedding_dim 均分给 $h$ 个头。例如，如果 embedding_dim=512，有 h=8 个头，那么每个头产生的 Q, K, V 向量维度就是 d_k = d_v = 512 / 8 = 64。计算时，先将输入 $X$ 分别投影到 $h$ 组低维的 Q, K, V 向量，并行计算后，再将结果拼接并投影回 embedding_dim 维度。

2.1 PyTorch 实现多头注意力

多头注意力是通过并行运行多个独立的自注意力“头”，并融合它们的输出来增强模型的表达能力。一个低效的实现是简单地创建多个 SelfAttention 实例并拼接结果。而高效的实现则是将多个头的计算逻辑合并到一次矩阵运算中。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


class MultiHeadSelfAttention(nn.Module):
 """多头自注意力模块"""
 def __init__(self, hidden_size, num_heads):
 super(MultiHeadSelfAttention, self).__init__()
 assert hidden_size % num_heads == 0, "hidden_size 必须能被 num_heads 整除"

 self.hidden_size = hidden_size
 self.num_heads = num_heads
 self.head_dim = hidden_size // num_heads

 self.q_linear = nn.Linear(hidden_size, hidden_size)
 self.k_linear = nn.Linear(hidden_size, hidden_size)
 self.v_linear = nn.Linear(hidden_size, hidden_size)
 self.wo = nn.Linear(hidden_size, hidden_size)

 def forward(self, x):
 batch_size, seq_len, _ = x.shape

 q = self.q_linear(x)
 k = self.k_linear(x)
 v = self.v_linear(x)

 # 拆分多头
 q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
 k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
 v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

 # 并行计算注意力
 scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
 attention_weights = torch.softmax(scores, dim=-1)
 context = torch.matmul(attention_weights, v)

 # 合并多头结果
 context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)

 # 输出层
 output = self.wo(context)

 return output

__init__:
- head_dim：计算出每个头的维度，即 hidden_size / num_heads。
- q_linear, k_linear, v_linear：与单头类似，但这里的线性层输出维度仍然是 hidden_size。这是为了一次性计算出所有头所需的总特征。
- wo：对应于多头注意力机制中的输出权重矩阵 $W^O$，用于融合所有头的信息。
forward:
- 线性变换: 与单头版本相同，得到总的 Q, K, V 矩阵。
- 拆分多头:
  - .view(batch_size, seq_len, self.num_heads, self.head_dim): 首先，将 hidden_size 维度逻辑上拆分为 num_heads 和 head_dim 两个维度。此时张量形状变为 [batch, seq_len, num_heads, head_dim]。
  - .transpose(1, 2): 然后，交换 seq_len 和 num_heads 维度，得到 [batch, num_heads, seq_len, head_dim]。这一步是为了让 num_heads 成为一个类似批次 (batch) 的维度，使得后续的矩阵乘法可以在每个头内部独立、并行地进行。
- 并行计算注意力: torch.matmul(q, k.transpose(-2, -1)) 现在是一个四维张量的乘法。PyTorch 会自动地将其解释为在第 0 和第 1 维（batch 和 num_heads）上进行批处理，而对最后两个维度执行矩阵乘法。这样就实现了所有头的注意力分数计算的并行化。
- 合并多头: 这是拆分操作的逆过程。
  - .transpose(1, 2): 先将 num_heads 和 seq_len 维度换回来，形状变为 [batch, seq_len, num_heads, head_dim]。
  - .contiguous(): 由于 transpose 操作可能导致张量在内存中不是连续存储的，需要调用 .contiguous() 来确保内存连续，之后才能安全地使用 .view()。
  - .view(batch_size, seq_len, self.hidden_size): 最后，将 num_heads 和 head_dim 两个维度重新合并成 hidden_size 维度，完成了所有头输出的拼接。
- 输出投影: 将合并后的结果通过 wo 线性层，得到最终输出。

三、Transformer 整体结构

理解了自注意力和多头注意力之后，就可以从一个更高的视角来审视 Transformer 的整体结构了。通过图 4-3 可以看出它依然是一个 Encoder-Decoder 架构，但其内部是由几个标准化的“积木”堆叠而成的。

图 4-3 Transformer 架构

Transformer 的 Encoder 和 Decoder 都是由 N 个（原论文中 N=6）功能相同的层（Layer）堆叠而成。下面我们分别来看它们的内部构造。

3.1 编码器（Encoder）

编码器的作用是“理解”和“消化”输入的整个序列，为序列中的每个词元生成一个富含上下文信息的表示。一个标准的编码器层由两个主要的子层构成，分别是多头自注意力层（Multi-Head Self-Attention Layer）和位置前馈网络（Position-wise Feed-Forward Network）。每个子层的输出都经过了**残差连接（Add）与层归一化（Norm）**处理。所以，一个编码器层内部的数据流可以表示为 x -> Sublayer1(x) -> Add & Norm -> Sublayer2(...) -> Add & Norm。

关键特性：

注意力类型：编码器中的多头注意力层是双向的自注意力。这意味着在计算时，序列中的任何一个词元都可以“看到”序列中的所有其他词元（包括它自己、它前面的和它后面的）。
功能：由于其双向性，编码器非常擅长理解完整的输入文本，并为每个词元生成一个深度融合了上下文信息的表示。
应用：通过大量堆叠编码器层而构建的模型（Encoder-Only 架构），如 BERT，在文本分类、命名实体识别等自然语言理解（NLU）任务上取得了巨大成功。

3.2 解码器（Decoder）

解码器的作用是基于编码器对原始输入的理解，并结合已经生成的部分，来逐个生成下一个词元。为了完成这个更复杂的任务，一个标准的解码器层（Decoder Layer）比编码器层多了一个注意力子层，总共包含三个子层。分别是带掩码的多头自注意力层（Masked Multi-Head Self-Attention Layer）、交叉注意力层（Cross-Attention Layer）和位置前馈网络（Position-wise Feed-Forward Network）。同样，解码器的每个子层也都采用了残差连接和层归一化。

关键特性：

子层 1：带掩码的自注意力
- 与编码器层相比，这是解码器的第一个主要区别。解码器在生成序列时必须是自回归的，即在生成第 $t$ 个词元时，只能依赖于已经生成的前 $t-1$ 个词元，而不能“看到”未来的信息。
- 为了在并行的自注意力计算中实现这一点，需要引入掩码 (Masking)。在计算 Softmax 之前，一个“未来词元掩码”会被应用到注意力分数上，将所有未来位置的分数设置为一个极小的负数（如 -inf），这样在经过 Softmax 之后，这些位置的注意力权重就会变为 0，从而确保了模型的单向性。
子层 2：交叉注意力
- 这是连接编码器和解码器的桥梁。这一层的实现与上一节中描述的交叉注意力一致。
- 它的 Key 和 Value 来自于编码器的最终输出，而 Query 则来自于解码器前一个子层（即带掩码的自注意力层）的输出。
- 这一层允许解码器在生成每个词元时，能够“关注”到输入序列的所有部分，从而有针对性地提取所需信息。
子层 3：逐位置前馈网络 (FFN)
- 这部分与编码器中的 FFN 相同，为模型提供非线性变换能力。
大模型应用：通过大量堆叠解码器层而构建的模型（Decoder-Only 架构），如 GPT 系列，由于其天然的自回归生成能力，引领了当前大语言模型（LLM）的发展浪潮。

3.3 组件解析

下面来详细解析一下构成上述“层”的几个重要组件。

3.3.1 位置前馈网络 (FFN)

这是一个由两次线性变换和一个激活函数组成的全连接网络，它独立地应用于序列中的每一个位置。

作用：特征变换。注意力子层内部主要包含 Softmax 归一化；逐位置的非线性主要由 FFN 提供（常见 ReLU/GELU²）。
内部结构：其常见的结构是“升维-激活-降维”。
$$ \text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2 $$
- 第一次线性变换（ $W_1$ ）通常会将输入维度 embedding_dim 放大到 4 倍（4 * embedding_dim）。这种升维操作的作用是将特征投影到一个更高维的空间，以便提取更丰富、更复杂的模式。
- 使用一个激活函数（如 ReLU）进行非线性处理。
- 第二次线性变换（ $W_2$ ）再将维度从 4 * embedding_dim 压缩回原始的 embedding_dim。这种降维操作可以看作是对高维特征的筛选和压缩，保留最重要的信息。

将中间层维度设为 4 倍的做法，主要是继承自原始论文的经验设定，并因其良好效果而被后续模型广泛采用，并非有严格的理论证明。

3.3.2 残差连接与层归一化 (Add & Norm)

为了让这些层能够成功地“堆叠”起来，每个子层的后面都连接了这个组合。

Add (残差连接)：解决了深度网络的“模型退化”问题。从反向传播的角度看，子层的输出可以写成 $y = x + \text{Sublayer}(x)$。在计算梯度时， $\frac{\partial y}{\partial x} = 1 + \frac{\partial \text{Sublayer}(x)}{\partial x}$。其中“1”的存在为梯度创建了一条“高速公路”，确保无论网络有多深，梯度都能至少以大小为 1 的程度回传到最浅层，极大地稳定了训练过程。同时，这也要求模型中所有子层的输入和输出维度必须保持一致，以便进行元素相加。
Norm (层归一化) ³：用于稳定训练过程。它独立地对每个样本的每个词元的特征向量（即 hidden_size 维度）进行标准化，使其均值变为 0，方差变为 1（但这并不假设其原始分布为正态分布）。更重要的是，它引入了两个可学习的参数 $\gamma$（缩放）和 $\beta$（偏移），让模型可以自主学习最佳的数据分布，其完整公式为：
$$ y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta $$
这使得模型既能享受到归一化带来的稳定性，又具备了根据任务需要恢复或调整原始分布的能力。与主要用于计算机视觉的 Batch Normalization（对一个批次中所有样本的同一特征进行归一化）相比，Layer Normalization 不受批次大小的影响，更适合处理长度可变的自然语言序列。

这个设计使得模型可以通过简单地增加“层”的数量（即深度）和特征维度（即宽度）来进行扩展，为后来参数量巨大的语言模型奠定了基础。

原论文采用 Post-LN（Sublayer → Add → LayerNorm）。许多现代实现（如 GPT 系列）采用 Pre-LN（LayerNorm → Sublayer → Add），训练更稳定、更易加深，但功能等价。

3.4 位置编码

自注意力机制的主要缺陷在于其 位置无关性。由于计算是完全并行的，模型无法感知词元的顺序。例如，“猫追狗”和“狗追猫”这两个句子，在自注意力看来，它们的词元集合完全相同，因此会为“猫”和“狗”生成相同的上下文表示，这显然是错误的。为了解决这个问题，Transformer 在将词嵌入向量输入模型之前，为它们加入了一个 位置编码 (Positional Encoding) 向量。其工作方式非常直接：

$$ input_\text{embedding} = token_\text{embedding} + positional_\text{encoding} $$

这个额外注入的向量为每个词元提供了其在序列中的位置信息。这是一种 绝对位置编码，即每个位置（如第 0、1、2 个位置）都有一个固定的编码向量。在实践中，主要有两种实现方式：

（1）可学习的位置编码 (Learned Positional Encoding) - 在 Encoder-only 模型（如 BERT）中常见；而近年的大型解码器式模型多采用相对/旋转类位置编码（如 RoPE⁴）。 - 其实现非常简单：创建一个 nn.Embedding 层，大小为 [max_sequence_length, hidden_size]。max_sequence_length 是模型能处理的最大序列长度，这是一个重要的超参数（在很多模型配置文件中被称为 max_position_embeddings）。在训练时，模型会像学习词嵌入一样，自动学习出每个位置（0, 1, 2, …）最合适的向量表示。

（2）基于三角函数的固定编码 (Sinusoidal Positional Encoding) - 这是原版 Transformer 论文中使用的方法，它不需要学习。 - 使用不同频率的正弦和余弦函数来为每个位置生成一个独特的、固定的编码向量：

 $$
PE_{(\text{pos}, 2i)} = \sin\left(\frac{\text{pos}}{10000^{\frac{2i}{d_{\text{embedding}}}}}\right) \\
PE_{(\text{pos}, 2i+1)} = \cos\left(\frac{\text{pos}}{10000^{\frac{2i}{d_{\text{embedding}}}}}\right)
$$
其中：
- $pos$ 是词元在序列中的绝对位置（如第 0 个、第 1 个词...）。
- $i$ 是编码向量中的维度索引（从 0 到 $d_{embedding}$/2）。公式通过 $i$ 来同时计算偶数维度 $2i$ 和奇数维度 $2i+1$ 的值，因此 $i$ 的取值范围只需达到维度总数的一半。
- $d_{embedding}$ 是词嵌入的维度。
公式利用不同频率的正弦和余弦函数，为 $d_{embedding}$ 维编码向量的每一个维度（$2i$ 对应偶数位，$2i+1$ 对应奇数位）计算一个特定的值。由于每个位置 $pos$ 和每个维度 $i$ 的组合都是独一无二的，所以这种方法能为序列中的每个位置生成一个完全独特的编码向量。这种方法的优势是不同位置的编码向量之间存在固定的线性关系，这可能有助于模型推断出词元间的相对位置。其主要优点是不需要训练，并且理论上可以外推到比训练时遇到的更长的序列。

绝对 vs. 相对位置编码

上述两种方法都属于绝对位置编码，因为它们为每个绝对位置（第 1 个、第 10 个等）分配一个特定的编码。然而，这种方式在处理超长文本时可能存在泛化性问题。因此，许多现代的大语言模型（如 Transformer-XL, Llama）转而采用相对位置编码。这种方法不再关注词元的绝对位置，而是直接在注意力计算中建模词元之间的相对距离（例如，“当前词”与“前 2 个词”之间的关系），这被证明在处理长序列时更有效、更灵活。

3.5 注意力掩码

掩码是 Transformer 模型中一个重要的机制。其主要目的是确保解码器在生成序列时的自回归特性，即不能“看到”未来的信息。此外，作为一个通用的工程实践，掩码也被用来处理批量训练中因句子长度不同而引入的填充（Padding）问题。Transformer 主要使用以下两种掩码：

（1）因果掩码：因果掩码专用于解码器的带掩码的自注意力（Masked Self-Attention）子层，是为了确保解码过程遵循自回归（Auto-regressive）特性，即生成第 $i$ 个词元时只能依赖前 $i-1$ 个词元的信息，而绝不能“偷看”到 $i$ 及之后位置的内容。它的实现核心是确保注意力权重矩阵呈现下三角矩阵的形态。对于长度为 $T$ 的序列，在 $[T, T]$ 的矩阵中，主对角线及以下的位置被标记为可关注（如 True 或 0），而主对角线以上的位置则被标记为屏蔽（如 False 或 1）。在计算 Softmax 之前，所有被屏蔽位置的注意力分数会被加上一个极大的负数（如 -inf），迫使其注意力权重归零，从而物理上切断了信息的向后传播路径。

（2）填充掩码：填充掩码广泛应用于编码器和解码器的所有注意力层，目的是解决变长序列批量处理时的**填充（Padding）**问题。由于填充词元（如 <pad>）本身不携带语义信息，若模型对其分配注意力，不仅浪费计算资源，还会引入噪声干扰。填充掩码的作用就是在计算注意力分数后，将所有涉及填充词元的位置（无论是作为查询 Query 还是作为键 Key）的对应分数强制设为极大的负数（如 -1e9 或负无穷）。假设有一个注意力分数矩阵，维度为 [batch_size, num_heads, seq_len, seq_len]。填充掩码会是一个 [batch_size, 1, 1, seq_len] 的矩阵（或可广播的形状），标记了哪些位置是填充。在进行 Softmax 之前，这个掩码会被加到分数矩阵上。经过 Softmax 运算后，这些负无穷位置的注意力权重会趋近于 0，从而在后续的加权求和中被完全忽略。

在解码器的自注意力层中，这两种掩码通常会结合使用，确保模型既不会关注到未来的信息，也不会关注到填充位。

3.6 解码器推理与 KV 缓存

解码器在训练和推理时的行为有很大不同。训练时，模型可以看到完整的“正确答案”序列，并通过注意力掩码来并行计算所有位置的损失。然而，在推理时，模型必须逐个生成词元，这是一个自回归的过程：

（1）输入 [BOS]（开始符），生成第一个词 token_1。

（2）输入 [BOS], token_1，生成第二个词 token_2。

（3）输入 [BOS], token_1, token_2，生成第三个词 token_3。

（4）… 直到生成 [EOS]（结束符）或达到最大长度。

如果按照这个流程直接计算，效率会非常低下。例如，在生成 token_3 时，模型需要为 [BOS] 和 token_1 重新计算它们的 Q, K, V 向量并参与注意力计算。但事实上，[BOS] 和 token_1 的 Key 和 Value 向量在之前的步骤中已经被计算过了。

为解决这种冗余计算，推理时会采用一项关键的优化技术：KV 缓存。

基本原理：对于解码器的每一层，都缓存下截至当前时刻已经计算出的所有词元的 Key 和 Value 向量。
工作流程：在生成第 $t$ 个词元时，模型只需要为当前输入的第 $t-1$ 个词元计算出它自己的 $q_{t-1}, k_{t-1}, v_{t-1}$。然后，它从缓存中取出历史的 $K_{cache} = [k_0, k_1, …, k_{t-2}]$ 和 $V_{cache} = [v_0, v_1, …, v_{t-2}]$。最后，将新的 $k_{t-1}, v_{t-1}$ 追加到缓存中，并用 $q_{t-1}$ 与更新后的完整 $K_{cache}, V_{cache}$ 进行注意力计算。

通过 KV 缓存，每次解码步骤的计算量从与整个已生成序列长度的平方（$O(T^2)$）相关，降低到只与序列长度（$O(T)$）线性相关，极大地加速了文本生成的速度，是实现高效大模型推理的常用技术之一。需要注意，KV 缓存占用会随步数线性增长（$O(T)$），在多层多头设置下需关注显存开销。

四、Transformer 代码实践

本节完整代码

4.1 项目结构设计

为了更好地理解 Transformer 的内部工作机制，接下来尝试从零实现一个完整的 Transformer 模型。我们会采用**“先整体框架，后组件实现”的思路，拆分多个文件来构建项目。在前面我们详细分析了 Transformer 的几大核心组件，分别是位置编码**、多头注意力、前馈网络以及归一化。为了体现这些组件的独立性和复用性，我们将遵循模块化的设计原则，将它们拆分到 src/ 目录下的独立文件中，而将模型的组装和运行逻辑放在根目录的 main.py 中。目录设计如下：

1
2
3
4
5
6
7
8


code/C4/transformer/
├── src/
│ ├── transformer.py # 核心框架：定义 Transformer、EncoderLayer 和 DecoderLayer
│ ├── attention.py # 核心组件：多头注意力机制 (MultiHeadAttention)
│ ├── ffn.py # 核心组件：前馈神经网络 (FeedForward)
│ ├── norm.py # 辅助组件：层归一化 (LayerNorm)
│ └── pos.py # 辅助组件：位置编码 (PositionalEncoding)
└── main.py # 入口脚本：组装模型并演示前向传播

4.2 搭建整体框架

在开始编写具体的注意力机制或前馈网络之前，我们可以先在 src/transformer.py 中勾勒出模型的高层架构。这种**“自顶向下”**的编程方式有助于我们理清数据流向。通过前面的学习我们知道，Transformer 宏观上是一个 Encoder-Decoder 架构，所以首先要实现的主要是以下几个部分：

Embedding 层：将输入的 token ID 转换为连续的向量表示，并加上位置编码以保留序列顺序信息。
Encoder 堆叠：由 $N$ 个 EncoderLayer 串联而成，负责深度提取和理解输入序列的特征。
Decoder 堆叠：由 $N$ 个 DecoderLayer 串联而成，负责基于 Encoder 的输出逐步生成目标序列。
Output 层：一个线性层，将解码器的最终输出映射回词表大小，用于计算下一个词的概率分布。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43


# src/transformer.py
import torch.nn as nn
from .pos import PositionalEncoding # 稍后实现
# ... 导入其他组件

class Transformer(nn.Module):
 def __init__(self, src_vocab_size, tgt_vocab_size, dim=512, n_heads=8, n_layers=6, ...):
 super().__init__()

 self.dim = dim
 # 1. 嵌入层与位置编码
 # src_embedding: 将源语言序列映射为向量 (Encoder输入)
 self.src_embedding = nn.Embedding(src_vocab_size, dim)
 # tgt_embedding: 将目标语言序列映射为向量 (Decoder输入)
 self.tgt_embedding = nn.Embedding(tgt_vocab_size, dim)
 self.pos_encoder = PositionalEncoding(dim, max_seq_len)
 self.dropout = nn.Dropout(dropout)

 # 2. 编码器与解码器堆叠
 # 使用 ModuleList 来存储层列表，支持按索引访问和自动注册参数
 self.encoder_layers = nn.ModuleList([
 EncoderLayer(dim, n_heads, hidden_dim, dropout) for _ in range(n_layers)
 ])
 self.decoder_layers = nn.ModuleList([
 DecoderLayer(dim, n_heads, hidden_dim, dropout) for _ in range(n_layers)
 ])

 # 3. 输出头
 self.output = nn.Linear(dim, tgt_vocab_size)

 def forward(self, src, tgt):
 # 1. 生成掩码 (Padding Mask & Causal Mask)
 src_mask, tgt_mask = self.generate_mask(src, tgt)

 # 2. 编码器前向传播
 enc_output = self.encode(src, src_mask)

 # 3. 解码器前向传播
 dec_output = self.decode(tgt, enc_output, src_mask, tgt_mask)

 # 4. 输出 Logits
 return self.output(dec_output)
 return logits

有了这个骨架，接下来的任务就是填充 EncoderLayer 和 DecoderLayer，而它们又依赖于更底层的组件。

4.3 实现核心组件

（1）位置编码 (src/pos.py)

在 src/transformer.py 中我们引入了 PositionalEncoding，它是 Transformer 处理序列顺序的关键。这里我们实现论文中的正弦位置编码。位置编码的核心在于初始化阶段，我们会预先计算好一个足够长的编码矩阵。它的计算公式使用了不同频率的正弦和余弦函数：

$$ PE(pos, 2i) = \sin(pos / 10000^{2i/d_{model}}) \\ PE(pos, 2i+1) = \cos(pos / 10000^{2i/d_{model}}) $$

在 __init__ 方法中，我们一次性生成这个矩阵，并将其注册为 buffer。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35


import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
 """
 正弦位置编码
 Transformer 论文中使用固定公式计算位置编码，不涉及可学习参数。
 """
 def __init__(self, dim, max_seq_len=5000):
 super().__init__()

 # 创建一个足够长的 PE 矩阵 [max_seq_len, dim]
 pe = torch.zeros(max_seq_len, dim)

 # 生成位置索引 [0, 1, ..., max_seq_len-1] -> [max_seq_len, 1]
 position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)

 # 计算分母中的 div_term: 10000^(2i/dim) = exp(2i * -log(10000)/dim)
 # 这种对数变换的计算方式在数值上更稳定
 div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))

 # 填充 PE 矩阵
 # 偶数维度用 sin，奇数维度用 cos
 pe[:, 0::2] = torch.sin(position * div_term)
 pe[:, 1::2] = torch.cos(position * div_term)

 # 增加 batch 维度: [1, max_seq_len, dim] 以便广播
 pe = pe.unsqueeze(0)

 # 注册为 buffer
 # register_buffer 的作用是告诉 PyTorch：
 # 1. 'pe' 是模型状态的一部分，会随模型保存和加载 (state_dict)。
 # 2. 'pe' 不是模型参数 (Parameter)，优化器更新时不会更新它。
 self.register_buffer('pe', pe)

在前向传播中，我们的任务就是将位置编码加到输入的词嵌入上。由于我们预先生成的 pe 矩阵可能比当前的输入序列 x 要长，所以需要根据 x 的实际长度对 pe 进行切片。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


...
class PositionalEncoding(nn.Module):
 def __init__(self, dim, max_seq_len=5000):
 ...

 def forward(self, x):
 """
 Args:
 x: 输入的词嵌入序列 [batch_size, seq_len, dim]
 Returns:
 加上位置编码后的序列 [batch_size, seq_len, dim]
 """
 # 截取与输入序列长度对应的位置编码并相加
 # x.size(1) 是 seq_len
 # self.pe 的形状是 [1, max_seq_len, dim]，切片后会自动广播到 batch_size
 x = x + self.pe[:, :x.size(1), :]
 return x

最后，我们可以编写一段简单的测试代码来验证维度是否正确。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


if __name__ == "__main__":
 # 准备参数
 batch_size, seq_len, dim = 2, 10, 512
 max_seq_len = 100

 # 初始化模块
 pe = PositionalEncoding(dim, max_seq_len)

 # 准备输入
 x = torch.zeros(batch_size, seq_len, dim) # 输入为0，直接观察PE值

 # 前向传播
 output = pe(x)

 # 验证输出
 print("--- PositionalEncoding Test ---")
 print(f"Input shape: {x.shape}")
 print(f"Output shape: {output.shape}")

输出如下：

1
2
3


--- PositionalEncoding Test ---
Input shape: torch.Size([2, 10, 512])
Output shape: torch.Size([2, 10, 512])

（2）多头注意力 (src/attention.py)

这是 Transformer 中最复杂的组件，用于从不同的“表示子空间”中提取信息。在初始化阶段，我们需要定义四个主要的线性层：wq, wk, wv 用于将输入投影到 Q, K, V 空间，wo 用于将多头注意力的输出投影回原始维度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


# src/attention.py
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
 def __init__(self, dim, n_heads, dropout=0.1):
 super().__init__()
 self.dim = dim
 self.n_heads = n_heads
 self.head_dim = dim // n_heads

 # 定义 Wq, Wk, Wv 矩阵
 # 这里我们使用一个大的线性层一次性计算所有头的 Q, K, V
 self.wq = nn.Linear(dim, dim)
 self.wk = nn.Linear(dim, dim)
 self.wv = nn.Linear(dim, dim)

 # 最终输出的线性层 Wo
 self.wo = nn.Linear(dim, dim)

 self.dropout = nn.Dropout(dropout)

这部分前向传播的重点是“分头”操作。我们不直接对 [batch, seq_len, dim] 进行计算，而是将其 reshape 为 [batch, n_heads, seq_len, head_dim]，这样就可以利用矩阵运算并行地处理所有头。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68


...
class MultiHeadAttention(nn.Module):
 def __init__(self, dim, n_heads, dropout=0.1):
 ...

 def forward(self, q, k, v, mask=None):
 batch_size = q.size(0)

 # 1. 线性投影
 # [batch, seq_len, dim] -> [batch, seq_len, dim]
 q = self.wq(q)
 k = self.wk(k)
 v = self.wv(v)

 # 2. 分头 (Split Heads)
 # 变换形状: [batch, seq_len, n_heads, head_dim] 
 # 然后转置: [batch, n_heads, seq_len, head_dim] 以便并行计算
 q = q.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)
 k = k.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)
 v = v.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)

 # 3. 计算缩放点积注意力 (Scaled Dot-Product Attention)
 # scores: [batch, n_heads, seq_len, seq_len]
 scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

 # 4. 应用掩码 (Masking)
 if mask is not None:
 # mask == 0 的位置被填充为负无穷，Softmax 后变为 0
 scores = scores.masked_fill(mask == 0, float('-inf'))

 # 5. Softmax 与加权求和
 attn_weights = torch.softmax(scores, dim=-1)

 if self.dropout is not None:
 attn_weights = self.dropout(attn_weights)

 # context: [batch, n_heads, seq_len, head_dim]
 context = torch.matmul(attn_weights, v)

 # 6. 合并多头 (Concat Heads)
 # [batch, n_heads, seq_len, head_dim] -> [batch, seq_len, n_heads, head_dim]
 # -> [batch, seq_len, dim]
 context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.dim)

 # 7. 输出层投影
 output = self.wo(context)

 return output

# 单元测试
if __name__ == "__main__":
 # 准备参数
 batch_size, seq_len, dim = 2, 10, 512
 n_heads = 8

 # 初始化模块
 mha = MultiHeadAttention(dim, n_heads)

 # 准备输入 (Query, Key, Value 相同)
 x = torch.randn(batch_size, seq_len, dim)

 # 前向传播
 output = mha(x, x, x)

 # 验证输出
 print("--- MultiHeadAttention Test ---")
 print(f"Input shape: {x.shape}")
 print(f"Output shape: {output.shape}")

输出如下：

1
2
3


--- MultiHeadAttention Test ---
Input shape: torch.Size([2, 10, 512])
Output shape: torch.Size([2, 10, 512])

（3）前馈神经网络 (src/ffn.py)

标准的 Transformer FFN 是一个简单的两层全连接网络，中间包含激活函数。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


# src/ffn.py
import torch.nn as nn

class FeedForward(nn.Module):
 def __init__(self, dim, hidden_dim, dropout=0.1):
 super().__init__()
 self.w1 = nn.Linear(dim, hidden_dim) # 升维
 self.w2 = nn.Linear(hidden_dim, dim) # 降维
 self.dropout = nn.Dropout(dropout)

 def forward(self, x):
 # 线性变换 -> ReLU -> Dropout -> 线性变换
 return self.w2(self.dropout(torch.relu(self.w1(x))))

if __name__ == "__main__":
 # 准备参数
 batch_size, seq_len, dim = 2, 10, 512
 hidden_dim = 2048

 # 初始化模块
 ffn = FeedForward(dim, hidden_dim)

 # 准备输入
 x = torch.randn(batch_size, seq_len, dim)

 # 前向传播
 output = ffn(x)

 # 验证输出
 print("--- FeedForward Test ---")
 print(f"Input shape: {x.shape}")
 print(f"Output shape: {output.shape}")

输出如下：

1
2
3


--- FeedForward Test ---
Input shape: torch.Size([2, 10, 512])
Output shape: torch.Size([2, 10, 512])

（4）层归一化 (src/norm.py)

层归一化 (Layer Normalization) 是 Transformer 中用来稳定训练的组件。与 Batch Normalization 不同，它是在最后一个维度（即特征维度 dim）上进行归一化的。公式如下：

$$ y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta $$

其中 $\gamma$ 和 $\beta$ 是可学习的缩放和平移参数。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48


import torch
import torch.nn as nn

class LayerNorm(nn.Module):
 """
 层归一化 (Layer Normalization)
 公式: y = (x - mean) / sqrt(var + eps) * gamma + beta
 """
 def __init__(self, dim, eps=1e-6):
 super().__init__()
 self.eps = eps
 # 可学习参数 gamma (缩放) 和 beta (偏移)
 # nn.Parameter 会被自动注册为模型参数
 self.gamma = nn.Parameter(torch.ones(dim))
 self.beta = nn.Parameter(torch.zeros(dim))

 def forward(self, x):
 # x: [batch_size, seq_len, dim]
 # 在最后一个维度 (dim) 上计算均值和方差
 # keepdim=True 保持维度以便进行广播计算
 mean = x.mean(-1, keepdim=True)
 # unbiased=False 使用有偏估计 (分母为 N)，与 PyTorch 默认行为一致
 var = x.var(-1, keepdim=True, unbiased=False)

 # 归一化
 x_norm = (x - mean) / torch.sqrt(var + self.eps)

 # 缩放和平移
 return self.gamma * x_norm + self.beta

# 单元测试
if __name__ == "__main__":
 # 准备参数
 batch_size, seq_len, dim = 2, 10, 512

 # 初始化模块
 ln = LayerNorm(dim)

 # 准备输入
 x = torch.randn(batch_size, seq_len, dim)

 # 前向传播
 output = ln(x)

 # 验证输出
 print("--- LayerNorm Test ---")
 print(f"Input shape: {x.shape}")
 print(f"Output shape: {output.shape}")

输出如下：

1
2
3


--- LayerNorm Test ---
Input shape: torch.Size([2, 10, 512])
Output shape: torch.Size([2, 10, 512])

4.4 组装与运行

（1）完善核心框架 (src/transformer.py)

之前我们只搭建了 Transformer 类的骨架，现在我们利用已经实现好的组件，按“编码器层 → 解码器层 → 辅助方法”的顺序来补全 src/transformer.py。编码器层，这部分包含一个多头自注意力子层和一个前馈网络子层，每个子层后面都接残差连接和层归一化（Post-LN 结构），代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


import torch
import torch.nn as nn
import math
# 导入组件
from .attention import MultiHeadAttention
from .ffn import FeedForward
from .norm import LayerNorm
from .pos import PositionalEncoding

class EncoderLayer(nn.Module):
 def __init__(self, dim, n_heads, hidden_dim, dropout=0.1):
 super().__init__()
 # 多头自注意力子层
 self.attention = MultiHeadAttention(dim, n_heads, dropout)
 self.attention_norm = LayerNorm(dim)
 # 前馈网络子层
 self.feed_forward = FeedForward(dim, hidden_dim, dropout)
 self.ffn_norm = LayerNorm(dim)

 self.dropout = nn.Dropout(dropout)

 def forward(self, x, mask=None):
 # 子层 1：自注意力
 _x = x
 x = self.attention(x, x, x, mask) # Q=K=V=x
 x = self.attention_norm(_x + self.dropout(x))

 # 子层 2：前馈网络
 _x = x
 x = self.feed_forward(x)
 x = self.ffn_norm(_x + self.dropout(x))

 return x

接下来是解码器层，这部分比编码器层多了一个“交叉注意力”子层，先是带掩码的自注意力，再是对编码器输出的交叉注意力，最后是前馈网络。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


class DecoderLayer(nn.Module):
 def __init__(self, dim, n_heads, hidden_dim, dropout=0.1):
 super().__init__()
 # 1. 带掩码的自注意力
 self.self_attention = MultiHeadAttention(dim, n_heads, dropout)
 self.self_attn_norm = LayerNorm(dim)
 # 2. 交叉注意力
 self.cross_attention = MultiHeadAttention(dim, n_heads, dropout)
 self.cross_attn_norm = LayerNorm(dim)
 # 3. 前馈网络
 self.feed_forward = FeedForward(dim, hidden_dim, dropout)
 self.ffn_norm = LayerNorm(dim)

 self.dropout = nn.Dropout(dropout)

 def forward(self, x, enc_output, src_mask, tgt_mask):
 # 子层 1：带掩码的自注意力
 _x = x
 x = self.self_attention(x, x, x, tgt_mask)
 x = self.self_attn_norm(_x + self.dropout(x))

 # 子层 2：交叉注意力（Q 来自解码器，K/V 来自编码器输出）
 _x = x
 x = self.cross_attention(x, enc_output, enc_output, src_mask)
 x = self.cross_attn_norm(_x + self.dropout(x))

 # 子层 3：前馈网络
 _x = x
 x = self.feed_forward(x)
 x = self.ffn_norm(_x + self.dropout(x))

 return x

最后在 Transformer 主类中，我们需要补全相关的辅助方法。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48


class Transformer(nn.Module):
 def __init__(self,
 src_vocab_size,
 tgt_vocab_size,
 dim=512,
 n_heads=8,
 n_layers=6,
 hidden_dim=2048,
 max_seq_len=5000,
 dropout=0.1):
 # ... 初始化嵌入层、位置编码、编码器/解码器堆叠以及输出层等 ...
 self._init_parameters()

 def _init_parameters(self):
 for p in self.parameters():
 if p.dim() > 1:
 nn.init.xavier_uniform_(p)

 def generate_mask(self, src, tgt):
 # src_mask: [batch, 1, 1, src_len]，pad token 假设为 0
 src_mask = (src != 0).unsqueeze(1).unsqueeze(2)

 # tgt_mask: [batch, 1, tgt_len, tgt_len]，结合 pad mask 和 causal mask
 tgt_len = tgt.size(1)
 tgt_pad_mask = (tgt != 0).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, tgt_len]
 tgt_subsequent_mask = torch.tril(torch.ones((tgt_len, tgt_len), device=tgt.device)).bool()
 tgt_mask = tgt_pad_mask & tgt_subsequent_mask.unsqueeze(0)
 return src_mask, tgt_mask

 def encode(self, src, src_mask):
 x = self.src_embedding(src) * math.sqrt(self.dim)
 x = self.pos_encoder(x)
 x = self.dropout(x)
 for layer in self.encoder_layers:
 x = layer(x, src_mask)
 return x

 def decode(self, tgt, enc_output, src_mask, tgt_mask):
 x = self.tgt_embedding(tgt) * math.sqrt(self.dim)
 x = self.pos_encoder(x)
 x = self.dropout(x)
 for layer in self.decoder_layers:
 x = layer(x, enc_output, src_mask, tgt_mask)
 return x

 # 前向传播
 def forward(self, src, tgt):
 ...

（2）运行主程序 (main.py)

现在所有的零件都准备好了，我们可以在 main.py 中将它们组装起来，并运行一个简单的前向传播测试。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50


import torch
from src.transformer import Transformer

def main():
 # 超参数
 src_vocab_size = 100
 tgt_vocab_size = 100
 dim = 512
 n_heads = 8
 n_layers = 6
 hidden_dim = 2048
 max_seq_len = 50
 dropout = 0.1

 # 实例化模型
 model = Transformer(
 src_vocab_size,
 tgt_vocab_size,
 dim,
 n_heads,
 n_layers,
 hidden_dim,
 max_seq_len,
 dropout
 )

 # 模拟输入数据
 batch_size = 2
 src_len = 10
 tgt_len = 12

 # 随机生成 src 和 tgt 序列 (假设 pad_token_id=0)
 # 确保没有 pad token 影响简单测试，或者手动插入
 src = torch.randint(1, src_vocab_size, (batch_size, src_len))
 tgt = torch.randint(1, tgt_vocab_size, (batch_size, tgt_len))

 # 前向传播
 output = model(src, tgt)

 print("Model Architecture:")
 # print(model)
 print("\nTest Input:")
 print(f"Source Shape: {src.shape}")
 print(f"Target Shape: {tgt.shape}")

 print("\nModel Output:")
 print(f"Output Shape: {output.shape}") # 预期 [batch_size, tgt_len, tgt_vocab_size]

if __name__ == "__main__":
 main()

输出如下：

1
2
3
4
5
6
7
8


Model Architecture:

Test Input:
Source Shape: torch.Size([2, 10])
Target Shape: torch.Size([2, 12])

Model Output:
Output Shape: torch.Size([2, 12, 100])

参考文献

注意力机制

Sat, 07 Feb 2026 19:34:25 -0800

注意力机制

在上一节的结尾，讨论了标准 Seq2Seq 架构存在的一个核心缺陷：信息瓶颈。编码器需要将源序列的所有信息，不论长短，全部压缩成一个固定长度的上下文向量 $C$。这种机制在处理长序列时，很容易丢失序列开头的关键信息，同时也无法让解码器在生成不同词元时，有选择性地关注输入的不同部分。

用上一节提到的对联任务举例，当上联是“两个黄鹂鸣翠柳”时，期望模型在生成下联时：

生成第一个词“一行”时，主要关注上联的“两个”。
生成第二个词“白鹭”时，主要关注上联的“黄鹂”。
…

但是标准的 Seq2Seq 架构的模型在生成“一行”、“白鹭”、“上青天”的每一个词时，所依赖的都是同一个、包含了整个上联概要的上下文向量 $C$。模型缺乏一种动态的、有倾向性的“关注”能力。为了解决这个问题，注意力机制 (Attention Mechanism) ¹被提出。

一、注意力机制的设计原理

注意力机制的原理，可以通俗地理解为从“一言以蔽之”到“择其要者而观之”的转变。人类在进行阅读理解或翻译时，并不会将整个句子或段落的信息平均地记在脑海里。当回答特定问题或翻译特定词组时，我们的注意力会自然地聚焦到原文中的相关部分。

注意力机制就是对这种认知行为的模拟。它的原理是在解码器生成每一个词元时，不再依赖一个固定的上下文向量，而是允许它“回头看”一遍完整的输入序列，并根据当前解码的需求，自主地为输入序列的每个部分分配不同的注意力权重，然后基于这些权重将输入信息加权求和，生成一个动态的、专属当前时间步的上下文向量。通过这种方式，模型便获得了“择其要者而观之”的能力：

在生成“一行”时，模型可以学会将最大的权重分配给“两个”所对应的编码器状态。
在生成“白鹭”时，则将最大的权重分配给“黄鹂”所对应的状态。

这个动态计算的权重，就是注意力权重；而整个动态计算上下文向量的过程，就是注意力机制。

二、注意力机制的动机与推导

为了更直观地理解注意力机制的必要性，可以跟随一个逐步深入的思路。

2.1 问题的根源与固定对齐策略的局限

标准的 Seq2Seq 模型之所以表现不佳，根源是它试图将源序列的所有信息无差别地压缩进一个向量。但在对联这类任务中，输入和输出之间存在着明显的局部对应关系。一个直观的想法是能不能建立一种固定的对齐策略？例如，在生成下联第一个词时，就只使用上联第一个词的编码信息；生成第二个词时，就只用第二个词的信息，以此类推。

这个想法可以表示为：

$C_1 = h_1$ (生成第一个词的上下文是第一个编码状态)
$C_2 = h_2$ (生成第二个词的上下文是第二个编码状态)
…

这种方法在处理像对联这样长度相等、词序对应的“特例”时似乎是可行的。但它的局限性非常明显：

（1）要求序列等长：对于不等长的序列（如中英文翻译），这种一对一的映射关系立刻失效。

（2）对齐关系僵化：它假设了输入和输出的对齐关系是固定不变的，但实际任务中的对应关系可能非常复杂（如一对多、多对一）。

这种固定对齐策略过于理想化，缺乏通用性。我们需要一种更灵活、更具普适性的方法。

2.2 注意力机制的动态加权原理

既然只取一个输入信息过于绝对，那么退一步，是否可以把所有输入信息都利用起来，但给它们分配不同的“重要性”呢？这就是通过动态加权进行对齐的思想，即加权求和。我们可以为解码的第 $t$ 步，动态地计算一个上下文向量 $C_t$，它由编码器所有的隐藏状态 $(h_1, h_2, \dots, h_{T_x})$ 加权求和得到：

$$ C_t = \sum_{j=1}^{T_x} \alpha_{tj} h_j $$

其中， $\alpha_{tj}$ 就是在解码第 $t$ 个词时，分配给输入第 $j$ 个词的注意力权重。

在这个思路下，前面提到的“固定对齐策略”可以看作是它的一个特例。例如，当 $\alpha_{11}=1$ 且其他所有 $\alpha_{1j}=0$ 时，就实现了 $C_1 = h_1$ 的效果。

2.3 如何确定权重？

加权求和的思路虽然灵活，但它引入了一个新的问题：权重 $\alpha_{tj}$ 从何而来？

这个权重显然不能是固定的。它必须是动态的，应该根据当前的解码需求来决定。例如，当解码器正要生成与“黄鹂”对应的词时，权重 $\alpha_{t, \text{黄鹂}}$ 就应该最大。所以我们需要一个额外的模块或机制，它能够：

（1）审视当前解码器的状态（例如，解码器上一时刻的隐藏状态 $h^{\prime}_{t-1}$）。

（2）将这个状态与编码器的每一个隐藏状态 $h_j$ 进行比较。

（3）根据比较结果，生成一组相应的权重 $(\alpha_{t1}, \alpha_{t2}, \dots, \alpha_{t,T_x})$。

让模型自行学习如何根据当前上下文来计算这组权重，正是注意力机制的关键。

三、注意力机制详解

带有注意力机制的 Encoder-Decoder 模型，其整体结构与标准 Seq2Seq 类似，主要区别在于解码器部分。编码器的工作保持不变，但是需要向解码器提供所有时间步的隐藏状态序列 $(h_1, h_2, \dots, h_{T_x})$，而不仅仅是最后一个时间步的状态。解码器在生成第 $t$ 个目标词元 $y_t$ 时，会通过三步进行“注意力计算”，来动态生成该时刻的上下文向量 $C_t$。这个过程通常以上一时刻的解码器隐藏状态 $h^{\prime}_{t-1}$ 为起点。

3.1 注意力计算三部曲

（1）计算相似度

使用解码器上一时刻的隐藏状态 $h^{\prime}_{t-1}$ 与编码器的每一个隐藏状态 $h_j$ 计算一个分数，这个分数衡量了在当前解码时刻，应当对第 $j$ 个输入词元投入多少“关注”。

$$ e_{tj} = \text{score}(h^{\prime}_{t-1}, h_j) $$

这个分数越高，代表关联性越强。计算这个分数的方式有很多种，例如简单的点积、或者引入一个可学习的神经网络层。

（2）计算注意力权重

得到输入序列所有位置的注意力分数 $(e_{t1}, e_{t2}, \dots, e_{t,T_x})$ 后，为了将它们转换成一种“权重”的表示，可使用 Softmax 函数对其进行归一化。这样，就能得到一组总和为 1、且均为正数的注意力权重 $(\alpha_{t1}, \alpha_{t2}, \dots, \alpha_{t,T_x})$。

$$ \alpha_{tj} = \text{softmax}(e_{tj}) = \frac{\exp(e_{tj})}{\sum_{i=1}^{T_x} \exp(e_{ti})} $$

这组权重 $\alpha_t$ 构成了一个概率分布，清晰地表明了在当前解码步骤 $t$，注意力应该如何分配在输入序列的各个位置上。

（3）加权求和，生成上下文向量

最后，使用上一步得到的注意力权重 $\alpha_{tj}$，对编码器的所有隐藏状态 $h_j$ 进行加权求和，从而得到当前解码时刻 $t$ 专属的上下文向量 $C_t$。

$$ C_t = \sum_{j=1}^{T_x} \alpha_{tj} h_j $$

这个 $C_t$ 向量，由于是根据当前解码需求动态生成的，它比原始 Seq2Seq 的那个固定向量 $C$ 包含了更具针对性的信息。

3.2 结合上下文进行预测

得到动态上下文向量 $C_t$ 后，模型会将其与当前解码器自身的输入词元 $y_{t-1}$ 的词嵌入结合起来（最常见的方式是将两者拼接），形成一个新的、信息更丰富的向量。

最后，将这个拼接后的向量连同上一时刻的状态 $h^\prime_{t-1}$ 一起送入解码器的 RNN 单元，计算出当前时刻的状态 $h^\prime_{t}$，并基于 $h^\prime_{t}$ 预测出最有可能的输出词元 $y_t$。整个过程可以通过图 4-2 来概括：

图 4-2 Attention 工作流程

3.3 一种高效的注意力打分函数

计算相关性分数的函数有多种设计，其中一种非常高效的方法，是直接计算查询向量（ $h^\prime_{t-1}$ ）和键向量（ $h_j$ ）的点积，并对其进行缩放。这种思想也是后续通用注意力框架的核心。其计算方式非常简洁：

$$ \text{score}(h^\prime_{t-1}, h_j) = \frac{{h^\prime_{t-1}}^T \cdot h_j}{\sqrt{d_k}} $$

其中：

${h^\prime_{t-1}}^T \cdot h_j$ 就是两个向量的点积。点积是衡量向量相似度的一种有效方式。
$d_k$ 是键向量（在这里是编码器隐藏状态）的维度。
除以 $\sqrt{d_k}$ 是一个关键的缩放步骤。当向量维度 $d_k$ 很大时，点积的结果的方差也会很大，这可能导致一些维度的值非常大，从而将 Softmax 函数推向其梯度极小的区域（即概率值极端地趋近于 0 或 1），造成梯度消失，使模型难以训练。通过除以 $\sqrt{d_k}$ 进行缩放，可以有效缓解这个问题，使训练过程更加稳定。

3.4 注意力机制的价值

引入注意力机制，不仅仅是对 Seq2Seq 架构的一个小修补，它还带来了一个全新的视角。

（1）克服信息瓶颈，提升性能

最直接的好处是，注意力机制彻底打破了信息必须被压缩成一个固定长度向量的限制。解码器在每一步都可以直接访问到源序列的全部信息，并根据需要动态聚焦。这使得模型在处理尤其是长序列时，性能得到了巨大的提升。

（2）提供可解释性，实现“词对齐”

注意力机制的另一个巨大价值在于它提供了很好的可解释性。注意力权重矩阵 $\alpha$ 本身就蕴含了丰富的信息。可以将这个矩阵可视化，来观察当模型生成某个输出词时，它的“注意力”主要集中在输入的哪些词上。

四、查询-键-值 (QKV) 范式

为了将刚刚描述的注意力计算过程抽象出来，形成一个更通用的思想，可以引入一个概念框架，即 查询-键-值（Query-Key-Value, QKV） 。这个范式将注意力的计算过程类比为一次信息检索：

查询（Query）：代表了当前的需求或意图。在 Seq2Seq 中，这就是解码器在生成下一个词元前的状态 $h^\prime_{t-1}$，可以理解为它在“查询”：“根据我现在的情况，我最需要输入序列的哪部分信息？”
键（Key）：可以看作是输入序列中各个信息片段的“标签”或“索引”，用于和查询进行匹配。在 Seq2Seq 中，输入序列的每个词元的隐藏状态 $h_j$ 都对应一个“键”。
值（Value）：是与“键”对应的实际信息内容。在基础的注意力机制中，“键”和“值”通常是相同的，都来自于编码器的隐藏状态 $h_j$。

无论形式如何变化，注意力机制的本质都可以概括为：通过查询（Q） 和一系列键（K） 计算相关性（权重），然后利用这个权重，对与各个键对应的值（V） 进行加权求和，得到最终的输出。

具体计算过程，可以用一个凝练的数学公式来统一表达，这就是 缩放点积注意力（Scaled Dot-Product Attention） ，它也是 Transformer 模型的核心组件之一²：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

这个公式准确概括了注意力的计算步骤：

（1）$QK^T$：计算查询矩阵 $Q$ 和键矩阵 $K$ 的转置的点积，得到原始的注意力分数。

（2）$\sqrt{d_k}$：对分数进行缩放，以维持训练稳定性，其中 $d_k$ 是键向量的维度。

（3）softmax(…)：通过 Softmax 函数将分数归一化，得到注意力权重。

（4）…V：将得到的权重矩阵与值矩阵 $V$ 相乘，进行加权求和，得到最终的输出。

这个通用的范式很实用，是理解后续 Transformer 等更先进模型的基础。在不同的任务中，只需要思考如何定义场景中的 Q、K、V 即可应用注意力机制。

在刚刚讨论的 Seq2Seq 中：Q 是解码器状态，K 和 V 都是编码器状态序列。
在后续的自注意力 (Self-Attention) 机制中，Q, K, V 将全部来源于同一个序列自身。

此外，为了增加模型的表达能力，还可以在计算注意力之前，对原始的 Q, K, V 向量各自通过一个独立的全连接层进行线性变换，得到新的 Q’, K’, V’，再用它们进行注意力的计算。这种做法可以让模型学习到在不同的“子空间”中进行信息匹配和聚合。

五、PyTorch 实现与代码解析

本节完整代码

5.1 整体思路

要在 PyTorch 中实现一个带注意力的 Seq2Seq 模型，需要对上节的代码进行一些关键的调整：

（1）编码器 Encoder：

forward 函数的返回值需要改变。除了最后一个时间步的隐藏状态 (hidden, cell)，还需要返回所有时间步的输出 outputs，这正是注意力机制计算所需要的 Key 和 Value。
如果编码器是双向 (Bidirectional) 的，其输出维度会是 hidden_size * 2。那么就需要增加一个线性层对其进行降维，或将其状态进行合并，以便与单向的解码器状态维度相匹配。

（2）新增 Attention 模块：

创建一个独立的 nn.Module 类来实现注意力的计算逻辑。其 forward 方法接收解码器状态 (Query) 和编码器所有输出 (Keys/Values)，返回计算得到的上下文向量。

（3）解码器 Decoder：

解码器的结构变化是最大的。它需要实例化一个 Attention 模块。
其 forward 函数通常以循环的方式逐个时间步解码，因为在第 $t$ 步计算注意力时需要依赖第 $t-1$ 步的解码器状态。需要强调的是，RNN 解码本身就是按时间步顺序计算，不能在时间维度并行；Attention 并未改变这一点。相较于“整序列一次性送入 RNN”的写法，逐步解码更便于在每步显式计算注意力并灵活插入教师强制等策略。在循环的每一步，它都会调用 Attention 模块计算上下文向量，并将其与当前词元的词嵌入融合后，再送入 RNN 单元。

5.2 编码器

为支持 Attention，编码器通常使用双向 RNN 以捕获更丰富的上下文，并需要返回所有时间步的输出序列作为注意力计算的 Key 和 Value。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


class Encoder(nn.Module):
 def __init__(self, vocab_size, hidden_size, num_layers):
 super(Encoder, self).__init__()
 self.embedding = nn.Embedding(
 num_embeddings=vocab_size,
 embedding_dim=hidden_size
 )
 self.rnn = nn.LSTM(
 input_size=hidden_size,
 hidden_size=hidden_size,
 num_layers=num_layers,
 batch_first=True,
 bidirectional=True # 使用双向LSTM
 )
 self.fc = nn.Linear(hidden_size * 2, hidden_size)

 def forward(self, x):
 embedded = self.embedding(x)
 outputs, (hidden, cell) = self.rnn(embedded)

 # 将双向RNN的输出通过线性层降维，使其与解码器维度匹配
 outputs = torch.tanh(self.fc(outputs))

 return outputs, hidden, cell

bidirectional=True：启用双向 LSTM，使原始 RNN outputs 维度变为 (batch, src_len, hidden_size * 2)。
self.fc：定义一个线性层，将拼接后的双向输出映射回 hidden_size 维度；经过 self.fc 和 tanh 后，outputs 维度回到 (batch, src_len, hidden_size)，方便后续计算。
return outputs, ...：返回降维后的所有时间步输出 outputs (作为后续的 K 和 V)，以及原始的最终状态 hidden 和 cell。

5.3 注意力模块的两种实现

这是模型的核心，我们通过两个版本的实现来体现具体的演进思路。

5.3.1 无参数的注意力

这个版本直接使用缩放点积来计算注意力，不引入额外的可学习参数，对应了注意力机制最基础的数学思想。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


class AttentionSimple(nn.Module):
 """1: 无参数的注意力模块"""
 def __init__(self, hidden_size):
 super(AttentionSimple, self).__init__()
 # 确保缩放因子是一个 non-learnable buffer
 self.register_buffer("scale_factor", torch.sqrt(torch.FloatTensor([hidden_size])))

 def forward(self, hidden, encoder_outputs):
 # hidden shape: (num_layers, batch_size, hidden_size)
 # encoder_outputs shape: (batch_size, src_len, hidden_size)

 # Q: 解码器最后一层的隐藏状态
 query = hidden[-1].unsqueeze(1) # -> (batch, 1, hidden)
 # K/V: 编码器的所有输出
 keys = encoder_outputs # -> (batch, src_len, hidden)

 # energy shape: (batch, 1, src_len)
 energy = torch.bmm(query, keys.transpose(1, 2)) / self.scale_factor

 # attention_weights shape: (batch, src_len)
 return torch.softmax(energy, dim=2).squeeze(1)

forward:
- 此方法的思路是：相似度越高的向量，其点积越大。直接利用这一数学特性来衡量 Query 与各个 Key 的关联程度。
- query = hidden[-1].unsqueeze(1): 提取解码器上一时间步的最终隐藏状态，作为当前解码需求的查询 (Query)。
- energy = torch.bmm(...): 通过矩阵乘法，一次性计算出 Query 向量与所有 Key 向量（即编码器各时刻的输出）的点积。这个点积结果 energy 直接反映了它们之间的原始相似度分数。除以一个缩放因子是为了让训练过程更稳定。
- return torch.softmax(...): 使用 Softmax 函数将原始分数转换成一个标准的概率分布，即最终的注意力权重。分数越高的位置，获得的权重也越大。

5.3.2 带参数的注意力

这个版本引入了可学习的参数（一个线性层和一个向量 v），让模型可以自主学习如何更好地对齐 Query 和 Keys。从 QKV 的来源来看，由于查询（Query）来自解码器，而键（Key）和值（Value）来自编码器，因此这种机制也可以称为交叉注意力 (Cross-Attention)。

“交叉”体现在哪里？

“交叉”一词形象说明了注意力计算的信息来源是两个不同的序列。

查询：来自解码器序列的当前状态（例如 $h^\prime_{t-1}$），代表了“我现在需要什么信息”。

键和值：来自编码器处理完整个源序列后产生的所有状态（例如 $h_1, h_2, \dots, h_{T_x}$），代表了“这里有全部的原始信息可供查询”。

信息从编码器序列流向解码器序列，两者通过注意力机制进行互动和对齐，因此被称为“交叉注意力”。与之相对的是自注意力 (Self-Attention)，其查询、键、值均来自同一个序列。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


class AttentionParams(nn.Module):
 """2: 带参数的注意力模块"""
 def __init__(self, hidden_size):
 super(AttentionParams, self).__init__()
 self.attn = nn.Linear(hidden_size * 2, hidden_size)
 self.v = nn.Parameter(torch.rand(hidden_size))

 def forward(self, hidden, encoder_outputs):
 src_len = encoder_outputs.shape[1]
 hidden_last_layer = hidden[-1].unsqueeze(1).repeat(1, src_len, 1)

 energy = torch.tanh(self.attn(torch.cat((hidden_last_layer, encoder_outputs), dim=2)))
 attention = torch.sum(self.v * energy, dim=2)

 return torch.softmax(attention, dim=1)

__init__:
- 此方法的核心是创建一个小型的神经网络（self.attn 和 self.v），让模型自主学习如何判断 Query 和 Key 之间的相关性，而不是使用固定的点积运算。
- self.attn: 一个线性层，它将解码器状态（Query）和编码器状态（Key）拼接后的信息进行变换，学习它们之间复杂的对齐关系。
- self.v: 一个可学习的向量，它的作用是将 self.attn 计算出的多维对齐信息，最终转化为一个单一的注意力分数。
forward:
- hidden_last_layer = ...: 将代表当前 Query 的解码器状态复制，使其能与每一个编码器状态进行配对。
- energy = torch.tanh(...): 将配对好的 Query 和 Key 拼接起来，一同送入 self.attn 线性层。这一步会计算出一个“能量”向量，tanh 激活函数则为其增加了非线性表达能力。
- attention = torch.sum(...): 利用可学习的 v 向量与这个向量进行点积，将多维的能量信息“压缩”成一个最终的、未经归一化的注意力分数。
- return torch.softmax(attention, dim=1): 和之前一样，使用 softmax 将分数转换为标准的注意力权重。

5.4 通用解码器

为了能同时适配上述两种 Attention 模块，我们设计一个通用的解码器。其核心改动是在每个时间步都调用 Attention 模块，并将计算出的上下文向量融入到当前步的输入中。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36


class DecoderWithAttention(nn.Module):
 def __init__(self, vocab_size, hidden_size, num_layers, attention_module):
 super(DecoderWithAttention, self).__init__()
 self.attention = attention_module
 self.embedding = nn.Embedding(
 num_embeddings=vocab_size,
 embedding_dim=hidden_size
 )
 self.rnn = nn.LSTM(
 input_size=hidden_size * 2, # 输入维度是 词嵌入(hidden_size) + 上下文向量(hidden_size)
 hidden_size=hidden_size,
 num_layers=num_layers,
 batch_first=True
 )
 self.fc = nn.Linear(hidden_size, vocab_size)

 def forward(self, x, hidden, cell, encoder_outputs):
 embedded = self.embedding(x.unsqueeze(1))

 # 1. 计算注意力权重
 # a shape: [batch, src_len]
 a = self.attention(hidden, encoder_outputs).unsqueeze(1)

 # 2. 计算上下文向量
 context = torch.bmm(a, encoder_outputs)

 # 3. 将上下文向量与当前输入拼接
 rnn_input = torch.cat((embedded, context), dim=2)

 # 4. 传入RNN解码
 outputs, (hidden, cell) = self.rnn(rnn_input, (hidden, cell))

 # 5. 预测输出
 predictions = self.fc(outputs.squeeze(1))

 return predictions, hidden, cell

__init__:
- self.attention: 持有传入的 Attention 实例（可以是 AttentionSimple 或 AttentionParams）。
- self.rnn: input_size 变为 hidden_size * 2，因为它接收的是词嵌入向量和上下文向量拼接后的结果。
forward: 完整地演示了注意力的应用流程。
- a = self.attention(...): 调用 attention 模块计算权重 a。
- context = torch.bmm(a, encoder_outputs): 对应注意力计算的第三步。使用矩阵乘法，通过权重 a 对 encoder_outputs (Values) 进行加权求和，得到上下文向量 context。
- rnn_input = torch.cat(...): 将动态生成的上下文向量和当前词的嵌入向量拼接起来，形成一个信息更丰富的输入，再送入 RNN 进行解码。

5.5 Seq2Seq 包装模块

这个模块负责将上述组件串联起来，并处理好双向编码器与单向解码器之间的状态传递问题。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


class Seq2Seq(nn.Module):
 """带注意力的Seq2Seq"""
 def __init__(self, encoder, decoder, device):
 super(Seq2Seq, self).__init__()
 self.encoder = encoder
 self.decoder = decoder
 self.device = device

 def forward(self, src, trg, teacher_forcing_ratio=0.5):
 batch_size = src.shape[0]
 trg_len = trg.shape[1]
 trg_vocab_size = self.decoder.fc.out_features
 outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)

 encoder_outputs, hidden, cell = self.encoder(src)

 # 适配Encoder(双向)和Decoder(单向)的状态维度
 hidden = hidden.view(self.encoder.rnn.num_layers, 2, batch_size, -1).sum(dim=1)
 cell = cell.view(self.encoder.rnn.num_layers, 2, batch_size, -1).sum(dim=1)

 input = trg[:, 0]
 for t in range(1, trg_len):
 # 在循环的每一步，都将 encoder_outputs 传递给解码器
 # 这是 Attention 机制能够"回顾"整个输入序列的关键
 output, hidden, cell = self.decoder(input, hidden, cell, encoder_outputs)
 outputs[:, t, :] = output
 teacher_force = random.random() < teacher_forcing_ratio
 top1 = output.argmax(1)
 input = trg[:, t] if teacher_force else top1

 return outputs

状态适配: 编码器是双向的，其 hidden 状态形状为 (num_layers * 2, ...)。解码器是单向的，需要 (num_layers, ...) 的初始状态。这里的 hidden.view(...).sum(dim=1) 通过 view 操作将状态拆分为 (层数, 方向, ...)，然后在方向维度上求和，巧妙地将双向状态合并为单向状态。
循环解码: 正如之前所强调的，Attention 机制下的解码必须是串行的。在 for 循环的每一步，都将 encoder_outputs 完整地传递给解码器，确保解码器在每个时间步都能基于上一时刻的状态，动态计算出当前最需要的上下文信息。

这个实现完整地展示了 Attention 机制如何克服信息瓶颈问题：解码器不再只依赖于一个固定的上下文向量，而是在生成的每一步，都通过 Attention 模块动态地计算出一个与当前解码状态最相关的上下文向量，极大地提升了模型性能。

完整Python代码

六、注意力机制的类型

在注意力机制发展的早期，受限于当时的硬件计算能力，研究者们为了降低计算开销，提出了一些不同类型的注意力机制。

6.1 Soft Attention vs. Hard Attention

Soft Attention：这就是前文一直在详细讨论的机制。它为输入序列的所有位置都计算一个注意力权重，这些权重是 0 到 1 之间的浮点数（经 Softmax 归一化），然后进行加权求和。这种方式的优点是模型是端到端可微的，可以使用标准的梯度下降法进行训练。其缺点是在处理非常长的序列时，计算开销会很大。因为解码的每一步，都需要计算当前状态与所有输入状态的相似度。
Hard Attention³：与 Soft Attention 对所有输入进行加权不同，Hard Attention 在每一步只选择一个最相关的输入位置。可以看作是一种“非 0 即 1”的注意力分配，即选中的位置权重为 1，其他所有位置的权重均为 0。这样做的好处是计算量大大减少，因为不再需要进行全面的加权求和。但它的缺点也很突出：选择过程是离散的、不可微的，因此无法使用常规的反向传播算法进行训练，通常需要借助强化学习等更复杂的技巧。

6.2 Global Attention vs. Local Attention

这是另一组从计算范围角度区分的概念，出自于另一篇开创性的论文⁴。

Global Attention (全局注意力)：其思想与 Soft Attention 基本一致，即在计算注意力时，会考虑编码器的所有隐藏状态。
Local Attention (局部注意力)：这是一种介于 Soft Attention 和 Hard Attention 之间的折中方案。能够减少计算量，但又不像 Hard Attention 那样极端。其核心思想是，在每个解码时间步，只关注输入序列的一个局部窗口。它的工作流程通常是：

（1）预测对齐位置：首先，模型需要预测一个当前解码步最关注的源序列位置 $p_t$。这个位置可以通过一个小型神经网络，仅依赖于当前解码器状态 $h^\prime_t$ 来预测，从而避免了与所有编码器状态进行比较，降低了计算成本。预测公式可以设计为： $p_t = T_x \cdot \text{sigmoid}(W_p h’_t + b_p)$，其中 $T_x$ 是源序列长度， $W_p$ 和 $b_p$ 是可学习的参数。

（2）定义窗口：以预测出的 $p_t$ 为中心，定义一个大小为 $2D+1$ 的窗口，其中 $D$ 是一个超参数。

（3）局部计算：最后，模型只在这个窗口内的编码器状态上应用 Soft Attention 机制，计算权重并生成上下文向量。

早期 Local Attention 通过局部窗口显著降低复杂度并保持良好性能。尽管硬件与内核优化推动了全局注意力在常规长度任务中的普及，但其 $O(N^2)$ 成本在长序列、低延迟或资源受限场景仍是瓶颈，因此局部/稀疏/窗口化/混合注意力在这些场景依然常用。