Part9.第15章：Transformer(下)

整体架构

左半部分是编码器Encoder部分，右边部分是解码器Decoder部分。

编码器

全连接模块实现

class FeedForwardBlock(nn.Module):

    def __init__(self, d_model: int, d_ff: int, dropout: float) -> None:
        super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        # (batch, seq_len, d_model) --> (batch, seq_len, d_ff) --> (batch, seq_len, d_model)
        return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))

全连接模块里定义了两个层，第一个层将每个token embedding的维度从512扩展到2048，然后应用了ReLU激活，接着进行dropout。第二层将每个token embedding的维度从2048重新降维到512。
整个前馈网络可以表示为：
$FFN(x)=W_2\cdot Dropout(ReLU(W_1\cdot x+b_1))+b_2$

• 维度扩展-压缩模式
  d_model → d_ff → d_model
  先升维（增加表达能力）
  再降维（保持维度一致）

• 位置独立性
  对序列中的每个位置独立应用相同的前馈网络
  每个token的处理不依赖其他位置

Add & Norm 模块实现

class ResidualConnection(nn.Module):

    def __init__(self, features: int, dropout: float) -> None:
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(features)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

• 参数
  features: 特征维度，通常是 d_model
  dropout: dropout概率
• self.dropout = nn.Dropout(dropout)
  用于对子层输出进行正则化
• self.norm = LayerNormalization(features)
  归一化在残差连接中起到稳定训练的作用
• def forward(self, x, sublayer)
  x: 输入张量
  sublayer: 一个子层函数（如自注意力或前馈网络）

整个残差连接可以表示为：
$ResidualConnection(x)=x+Dropout(Sublayer(LayerNorm(x)))$

设计特点

1. 计算顺序：Pre-LN 架构

# Pre-LN: LayerNorm在子层之前
x → LayerNorm → Sublayer → Dropout → + x
# 对比 Post-LN: LayerNorm在残差连接之后
x → Sublayer → Dropout → + x → LayerNorm

Pre-LN（这里实现的）通常训练更稳定
Post-LN 是原始Transformer论文中的设计

2. 残差连接的重要性

残差连接的核心思想
output = input + transformation(input)

• 解决梯度消失: 梯度可以直接通过加法路径反向传播
• 恒等映射: 当子层学习效果不好时，网络可以退化为恒等映射
• 信息保留: 原始信息得到保留，新信息在残差上学习

3. Dropout的位置
   在残差连接之前应用dropout
   防止子层输出过拟合

编码器实现代码

class EncoderBlock(nn.Module):

    def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 feed_forward_block: FeedForwardBlock, dropout: float) -> None:
        super().__init__()
        # 定义多头自注意力模块
        self.self_attention_block = self_attention_block
        # 定义全连接模块
        self.feed_forward_block = feed_forward_block
        # 定义两个Add & Norm模块
        self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(2)])

    def forward(self, x, src_mask):
        # 第一个残差连接，跳过多头注意力模块
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
        # 第二个残差连接，跳过全连接模块
        x = self.residual_connections[1](x, self.feed_forward_block)
        return x


class Encoder(nn.Module):

    def __init__(self, features: int, layers: nn.ModuleList) -> None:
        super().__init__()
        # 传入的6个EncoderBlock
        self.layers = layers
        self.norm = LayerNormalization(features)

    def forward(self, x, mask):
        # 依次调用6个EncoderBlock
        for layer in self.layers:
            x = layer(x, mask)
        # 输出前进行Layer Norm
        return self.norm(x)

【代码解释】

• 【1】关键行: 创建包含2个残差连接的模块列表
  第一个用于自注意力层
  第二个用于前馈网络层

self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(2)])

• 【2】def forward(self, x, src_mask):
  x: 输入张量，形状为 [batch_size, seq_len, d_model]
  src_mask: 源序列掩码，用于防止注意力机制关注填充位置
• 【3】x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
  第一个残差连接：多头自注意力
  详细分解：
  self.residual_connections[0]: 第一个残差连接模块
  lambda x: self.self_attention_block(x, x, x, src_mask):
  这是一个匿名函数，将作为子层传递给残差连接
  self.self_attention_block(x, x, x, src_mask):
  三个 x 分别对应：query, key, value（自注意力）
  src_mask: 防止关注到填充位置
  计算流程：

原始输入x → LayerNorm → 自注意力 → Dropout → + 原始输入x

• 【4】x = self.residual_connections[1](x, self.feed_forward_block)
  第二个残差连接：前馈网络
  计算流程：

上一步输出 → LayerNorm → 前馈网络 → Dropout → + 上一步输出

• 【5】def forward(self, x, mask):
  参数:
  x: 输入张量（已加上位置编码）
  mask: 源序列掩码

一个编码器块的计算可以表示为：

其中 $l$ 表示第 $l$ 个编码器块。

编码器总结

编码器的作用是根据上下文，不断更新输入序列中每个token的Embedding。让最终输出序列token的embedding有最佳的语义。因为自注意力机制，可以一次性输入整个序列，序列内多个token可以并行处理，大大加快了处理速度。其中使用了位置编码增加了token的位置信息。通过多头注意力让token可以根据上下文信息修改自身embedding。通过全连接模块让每个token更新自身embedding。通过残差连接和LayerNorm让深度网络更容易训练。

解码器

以英文翻译中文任务为例，解码器的作用就是参考编码器的输出，以及自身已经翻译的内容，生成下一个中文token。

解码器的输入

编码器输出 编码器输出的是英文序列里每个token的embedding。它的维度为512。经过多层编码器的自注意力机制，每个token的embedding都已经根据上下文，计算出恰当的语义信息。它们将作为解码器输出的重要参考。

已经翻译出来的token序列 Transformer的编码器可以一次性输入完整的英文token序列，但在模型进行实际翻译时，需要解码器逐个生成对应的中文token序列。$<bos>$为解码器的初始输入，代表序列开始。在上图例子里，我们希望模型输入$<bos>$，能输出“我”这个token。然后将$“<bos>,我”$这两个token 序列作为输入，希望模型能输出“爱”这个token。依次类推，直到模型输出$<eos>$这个token。当然在这个过程中的每一步解码器也同时参考了编码器给出的英文token 序列的embedding。

对于Transformer模型在推理时，解码器确实如上述过程所述，是逐个生成中文token的。但是在训练时，因为我们已经知道英文对应的中文token序列，所以我们可以通过一种叫做带掩码的多头注意力机制（Masked Multi-Head Attention，MMHA）来实现并行化训练。
对于上图中中文token序列：$<bos>$ | 我 | 爱 | 吃 | 苹果，一共有5个token，我们可以构造一个下三角矩阵：

Mask矩阵每行表示当前token可以看到的token，1代表可以看到，0代表看不到。
比如第一行只有第一个位置为1，代表第一个token $<bos>$ 只能看到第一个token，也就是自己本身。

第二行可以看到前两个token。

Transformer里只有注意力计算时是跨token，根据上下文更新自身embedding的。在进行注意计算时，不光传入Q，K，V矩阵，还需要传入这个Mask矩阵。让每个token在注意里计算时，根据mask矩阵，只关注自己以及前边的token。这样在一次训练里就可以对整个中文序列一次性进行训练了。

在实际代码实现时，先利用Q、K矩阵进行注意力计算，在softmax()前，把Mask矩阵为0的位置的注意力结果填充成一个非常大的负值，这样经过softmax()后，这些位置的值就都为0。在接下来计算当前token的特征值时，序列中Mask为0的token的权重就为0，代码如下：

def attention(query, key, value, mask, dropout: nn.Dropout):
    d_k = query.shape[-1]
    # (batch, h, seq_len, d_k) --> (batch, h, seq_len, seq_len)
    attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        # 给mask为0的位置填入一个很大的负值
        attention_scores.masked_fill_(mask == 0, -1e9)
    # (batch, h, seq_len, seq_len)
    attention_scores = attention_scores.softmax(dim=-1)
    if dropout is not None:
        attention_scores = dropout(attention_scores)
    # (batch, h, seq_len, seq_len) --> (batch, h, seq_len, d_k)
    return (attention_scores @ value), attention_score

代码解释*

• 【1】def attention(query, key, value, mask, dropout: nn.Dropout):
  query: 查询向量，形状为 (batch, num_heads, seq_len_q, d_k)
  key: 键向量，形状为 (batch, num_heads, seq_len_kv, d_k)
  value: 值向量，形状为 (batch, num_heads, seq_len_kv, d_v)
  mask: 注意力掩码，用于屏蔽某些位置
  dropout: Dropout层实例
• 【2】注意力计算核心公式

# (batch, h, seq_len, d_k) --> (batch, h, seq_len, seq_len)
attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)

【2.1】key.transpose(-2, -1): 转置键矩阵的最后两个维度
形状: (batch, num_heads, seq_len_kv, d_k) → (batch, num_heads, d_k, seq_len_kv)
【2.2】矩阵乘法 query @ key.transpose(…):

查询: [batch, num_heads, seq_len_q, d_k]
键转置: [batch, num_heads, d_k, seq_len_kv]
结果: [batch, num_heads, seq_len_q, seq_len_kv]

计算查询和键的点积
表示每个查询位置与每个键位置的相关性
【2.3】缩放因子 / math.sqrt(d_k):
数学公式: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$
为什么需要缩放？防止点积过大导致softmax梯度消失

• 【3】掩码处理

if mask is not None:
    # 给mask为0的位置填入一个很大的负值
    attention_scores.masked_fill_(mask == 0, -1e9)

将掩码为0的位置填充为负无穷（这里用-1e9近似）
在softmax后，这些位置的概率会趋近于0。

整个过程可以用公式表示：

$Q,K,V$: 查询、键、值矩阵
$d_k$: 键向量的维度
$M$: 掩码矩阵（负无穷表示需要屏蔽的位置）

完整流程图示如下：

查询向量Q        键向量K
    ↓             ↓
    ┌─────────────┐
    │ Q·K^T / √d_k │ ← 计算点积并缩放
    └─────────────┘
          ↓
    ┌─────────────┐
    │  掩码处理   │ ← 应用掩码（如需要）
    └─────────────┘
          ↓
    ┌─────────────┐
    │   softmax   │ ← 转换为概率分布
    └─────────────┘
          ↓
    ┌─────────────┐
    │   dropout   │ ← 随机丢弃（训练时）
    └─────────────┘
          ↓
    ┌─────────────┐
    │   · V       │ ← 加权求和值向量
    └─────────────┘
          ↓
       输出

Mask机制只能在训练时一次性对解码器所有位置的token进行训练。但是在模型训练好进行生成时，因为事先并不知道答案，只能一步生成一个新的token，然后拼接到解码器的输入，再生成下一个。

交叉注意力 Cross-Attention

观察Transformer模型架构图的Decoder Block部分，它也是6个，每个Decoder Block有2个注意力模块+1个全连接模块。
Corss-Attention是非常重要的一步，它是解码器从编码器获取信息的唯一途径。从模型架构图中可以看到，交叉注意力模块的K和V矩阵来自编码器的输出，Q矩阵来自解码器部分。所以解码器根据当前翻译的需要提出查询向量q，和所有编码器输出的k进行匹配，计算注意力。最终得到编码器输出的v的注意力加权值。
代码实现：

def forward(self, x, encoder_output, src_mask, tgt_mask):
    # 第一个自注意力模块，Q、K、V都来自自身。
    x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
    # 第二个交叉注意力模块的Q矩阵来自Decoder，K,V矩阵来自Encoder的输出。
    x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output,encoder_output, src_mask))
    # 全连接模块。
    x = self.residual_connections[2](x, self.feed_forward_block)
    return x

代码解释：

• 【1】def forward(self, x, encoder_output, src_mask, tgt_mask):
  x: 解码器输入（目标序列的嵌入+位置编码）
  encoder_output: 编码器的输出
  src_mask: 源序列掩码（用于编码器-解码器注意力）
  tgt_mask: 目标序列掩码（用于解码器自注意力）
• 【2】解码器自注意力（掩码自注意力）

x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
# 残差链接结构
# x → LayerNorm → 自注意力 → Dropout → + x

• • self_attention_block: 解码器的自注意力模块
    参数说明：
    第一个 x: query（来自解码器自身）
    第二个 x: key（来自解码器自身）
    第三个 x: value（来自解码器自身）
    tgt_mask: 目标序列掩码
• • tgt_mask 的作用：
    前瞻掩码（Look-Ahead Mask）: 防止当前位置关注未来的位置
    填充掩码: 防止关注到填充位置
• 【3】编码器-解码器注意力（交叉注意力）

x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask))
# 残差链接结构
上一步输出 → LayerNorm → 交叉注意力 → Dropout → + 上一步输出

• • lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask)
    cross_attention_block: 编码器-解码器注意力模块
    参数说明：
    第一个 x: query（来自解码器）
    第二个 encoder_output: key（来自编码器）
    第三个 encoder_output: value（来自编码器）
    src_mask: 源序列掩码
• • 为什么Q、K、V不同？
    Query: 来自解码器（目标语言）
    Key & Value: 来自编码器输出（源语言）
    这使得解码器可以"查询"编码器的信息
• • src_mask 的作用：
    防止解码器关注到编码器输出的填充位置
    确保只关注有效的源语言信息
• 【4】前馈神经网络

x = self.residual_connections[2](x, self.feed_forward_block)
残差连接结构：
上一步输出 → LayerNorm → 前馈网络 → Dropout → + 上一步输出

• • 前馈网络特点：
    与编码器中的前馈网络结构相同
    独立处理每个位置
    提供非线性变换能力

完整流程图示

解码器输入 (x)
       ↓
┌─────────────────────────────────┐
│ 第1个残差连接: 解码器自注意力         │
│ - Query: x                        │
│ - Key: x                          │
│ - Value: x                        │
│ - 掩码: tgt_mask (前瞻+填充掩码)    │
└─────────────────────────────────┘
       ↓
┌─────────────────────────────────┐
│ 第2个残差连接: 编码器-解码器注意力  │
│ - Query: 上一步输出               │
│ - Key: encoder_output           │
│ - Value: encoder_output         │
│ - 掩码: src_mask (填充掩码)       │
└─────────────────────────────────┘
       ↓
┌─────────────────────────────────┐
│ 第3个残差连接: 前馈神经网络        │
│ - 独立处理每个位置                │
│ - 增加非线性表达能力              │
└─────────────────────────────────┘
       ↓
解码器块输出

解码器最终会接一个线性分类头，输入是512，输出维度是字典的大小。线性层的输出再经过softmax()之后，就是字典里每个token作为输出下一个token的概率值。

解码器数学表示

自注意力: 捕捉目标序列的内部依赖关系；
交叉注意力: 建立源语言和目标语言之间的对应关系；
前馈网络: 增加非线性表达能力。

不同类型注意力的对比：

# 解码器自注意力（掩码自注意力）
self_attention_output = self_attention(
    query=x,     # 来自解码器
    key=x,       # 来自解码器
    value=x,     # 来自解码器
    mask=tgt_mask  # 前瞻掩码+填充掩码
)

# 编码器-解码器注意力（交叉注意力）
cross_attention_output = cross_attention(
    query=x,             # 来自解码器
    key=encoder_output,  # 来自编码器
    value=encoder_output,  # 来自编码器
    mask=src_mask        # 填充掩码
)

自注意力模块：为什么需要前瞻掩码

因果性约束

• 解码器生成目标序列时，当前位置只能看到当前和之前的token
• 防止模型在训练时"作弊"——看到未来的信息
• 这是序列生成任务的核心约束

假设生成句子：“我 喜欢 机器 学习”

步骤1: 输入: [<s>]          -> 输出: "我"
步骤2: 输入: [<s>, "我"]     -> 输出: "喜欢"
步骤3: 输入: [<s>, "我", "喜欢"] -> 输出: "机器"
步骤4: 输入: [<s>, "我", "喜欢", "机器"] -> 输出: "学习"

每个步骤只能看到已生成的部分
如果不使用前瞻掩码会怎样？

# 无掩码的训练（错误！）
输入: [<s>, "我", "喜欢", "机器", "学习"]
输出: "我 喜欢 机器 学习"

模型会在训练时看到整个目标句子!!!
推理时：模型只能看到部分句子 -> 性能严重下降!!!

交叉注意力：为什么只需要填充掩码

信息访问权限

• 解码器访问编码器输出时，可以访问整个源序列
• 因为源序列是已知的完整输入（如源语言句子）
• 没有"因果性"约束，只有"有效性"约束

为什么不需要前瞻掩码？

• 1.源序列是静态的：训练和推理时，编码器已经处理了完整的源序列
• 2.全序列信息可用：翻译时，目标词的生成可以依赖于源序列的任何位置
• 3.对齐需求：目标词可能需要与源序列的任意位置对齐

举例：英译中
源英文: "I love machine learning" (4个词)
目标中文: "我 喜欢 机器 学习" (4个词)

交叉注意力允许：
- "我" 可以关注 "I"
- "喜欢" 可以关注 "love"  
- "机器" 可以关注 "machine"
- "学习" 可以关注 "learning"
- 每个目标词都可以关注整个源序列

两者对比：

解码器实现代码

class DecoderBlock(nn.Module):
    def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 cross_attention_block: MultiHeadAttentionBlock, feed_forward_block: FeedForwardBlock,
                 dropout: float) -> None:
        super().__init__()
        self.self_attention_block = self_attention_block
        self.cross_attention_block = cross_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(3)])

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
        # 交叉注意力模块的Q矩阵来自Decoder，K,V矩阵来自Encoder的输出
        x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output,encoder_output, src_mask))
        x = self.residual_connections[2](x, self.feed_forward_block)
        return x

class Decoder(nn.Module):
    def __init__(self, features: int, layers: nn.ModuleList) -> None:
        super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, encoder_output, src_mask, tgt_mask)
        return self.norm(x)

完整架构图示

解码器输入 (目标序列嵌入 + 位置编码)
       ↓
┌─────────────────────────────────┐
│  Decoder Block 1                │
│  ┌─────────────────────────┐    │
│  │ 1. 掩码自注意力           │    │
│  │    - Q,K,V: 解码器自身    │    │
│  │    - 掩码: tgt_mask      │    │
│  └─────────────────────────┘    │
│  ┌─────────────────────────┐    │
│  │ 2. 交叉注意力            │    │
│  │    - Q: 解码器           │    │
│  │    - K,V: 编码器输出     │    │
│  │    - 掩码: src_mask      │    │
│  └─────────────────────────┘    │
│  ┌─────────────────────────┐    │
│  │ 3. 前馈神经网络          │    │
│  └─────────────────────────┘    │
└─────────────────────────────────┘
       ↓
┌─────────────────────────────────┐
│  Decoder Block 2                │
│  ...                            │
└─────────────────────────────────┘
       ↓
       ...
       ↓
┌─────────────────────────────────┐
│  Decoder Block N                │
│  ...                            │
└─────────────────────────────────┘
       ↓
     层归一化
       ↓
   解码器输出

训练与推理的区别
训练模式：

# 训练时：使用完整目标序列进行教师强制
def train_step(self, src, tgt, src_mask, tgt_mask):
    # 前向传播
    logits = self.forward(src, tgt, src_mask, tgt_mask)
    
    # 计算损失（只计算非填充位置）
    loss = self.compute_loss(logits, tgt)
    return loss

推理模式：

def generate(self, src, src_mask, max_len=50, start_token=1):
    """自回归序列生成"""
    batch_size = src.size(0)
    device = src.device
    
    # 编码源序列
    encoder_output = self.encode(src, src_mask)
    
    # 初始化目标序列（只有开始标记）
    tgt = torch.ones(batch_size, 1, dtype=torch.long, device=device) * start_token
    
    for step in range(max_len):
        # 创建当前步的目标序列掩码
        tgt_mask = create_look_ahead_mask(tgt.size(1)).to(device)
        
        # 解码
        decoder_output = self.decode(tgt, encoder_output, src_mask, tgt_mask)
        
        # 预测下一个词
        logits = self.project(decoder_output[:, -1:, :])  # 只取最后一个位置
        next_token = torch.argmax(logits, dim=-1)
        
        # 添加到序列中
        tgt = torch.cat([tgt, next_token], dim=1)
        
        # 检查是否所有序列都生成了结束标记
        if (next_token == END_TOKEN).all():
            break
    
    return tgt