LLM笔记（八）Transformer学习

核心思想: “Attention Is All You Need” (论文标题) - 完全抛弃了传统的循环神经网络 (RNN, LSTM, GRU) 和卷积神经网络 (CNN) 结构，仅依赖注意力机制来捕捉输入和输出之间的全局依赖关系。

来源: Google 2017年论文《Attention Is All You Need》

1. Transformer 整体架构

Transformer 模型通常采用 Encoder-Decoder 架构，特别适用于序列到序列 (Seq2Seq) 任务，如机器翻译。

• Encoder (编码器):
  • 接收输入序列 (例如，源语言句子)。
  • 将其转换为一系列连续的表示 (context vectors)。
  • 由 N 个相同的 Encoder Layer 堆叠而成。
• Decoder (解码器):
  • 接收 Encoder 的输出和目标序列的部分内容 (例如，已翻译的部分)。
  • 生成输出序列 (例如，目标语言句子)。
  • 由 N 个相同的 Decoder Layer 堆叠而成。

重要变体:

• Encoder-Only (例如 BERT, ViT): 适用于理解任务，如文本分类、命名实体识别、图像分类。
• Decoder-Only (例如 GPT系列): 适用于生成任务，如文本生成、语言建模。

整体架构代码-编码器解码器

class EncoderDecoder(nn.Module):  """  A standard Encoder-Decoder architecture. Base for this and many    other models.    """  def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):  super(EncoderDecoder, self).__init__()  self.encoder = encoder  self.decoder = decoder  self.src_embed = src_embed  self.tgt_embed = tgt_embed  self.generator = generator  def forward(self, src, tgt, src_mask, tgt_mask):  "Take in and process masked src and target sequences."  return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)  def encode(self, src, src_mask):  return self.encoder(self.src_embed(src), src_mask)  def decode(self, memory, src_mask, tgt, tgt_mask):  return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

class Generator(nn.Module):  "Define standard linear + softmax generation step."  def __init__(self, d_model, vocab):  super(Generator, self).__init__()  self.proj = nn.Linear(d_model, vocab)  def forward(self, x):  return log_softmax(self.proj(x), dim=-1)

---

2. 核心组件详解

2.1. 输入部分 (Input Embedding & Positional Encoding)

• Input Embedding (词嵌入):
  • 将输入的离散词元 (tokens) 转换为固定维度的向量表示。
  • 通常使用一个可学习的嵌入层 (Embedding Layer)。
• Positional Encoding (位置编码):
  • 为什么需要? Transformer 的自注意力机制本身不包含序列中词元的位置信息 (即它是排列不变的)。为了让模型知道词元的顺序，需要引入位置信息。
  • 如何实现?

    • 原始论文方法: 使用不同频率的正弦 (sin) 和余弦 (cos) 函数。对于位置 pos 和维度 i：

      PE(pos, 2i) = sin(pos / 10000^(2i/d_model))PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
      

      其中 d_model 是嵌入向量的维度。这种方式的好处是模型可以学习到相对位置信息，并且可以推广到比训练时更长的序列。

    • 可学习的位置编码: 将位置也视为一种可学习的嵌入向量。

  • 位置编码向量会与词嵌入向量相加。

class PositionalEncoding(nn.Module):  "Implement the PE function."  def __init__(self, d_model, dropout, max_len=5000):  super(PositionalEncoding, self).__init__()  self.dropout = nn.Dropout(p=dropout)  # Compute the positional encodings once in log space.  pe = torch.zeros(max_len, d_model)  position = torch.arange(0, max_len).unsqueeze(1)  div_term = torch.exp(  torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)  )  pe[:, 0::2] = torch.sin(position * div_term)  pe[:, 1::2] = torch.cos(position * div_term)  pe = pe.unsqueeze(0)  self.register_buffer("pe", pe)  def forward(self, x):  x = x + self.pe[:, : x.size(1)].requires_grad_(False)  return self.dropout(x)

---

2.2. 注意力机制 (Attention Mechanism)

相关博文：注意力机制说明

• Scaled Dot-Product Attention (缩放点积注意力):

  • 输入: Query (Q), Key (K), Value (V) 三个矩阵。
    • Q (查询): 当前词元的表示，它要去"查询"其他词元。
    • K (键): 其他词元的表示，用来和 Q 匹配，计算相似度。
    • V (值): 其他词元的表示，根据 Q 和 K 的相似度进行加权求和，得到当前词元的新表示。
  • 计算过程:
    1. 计算得分 (Score): Score = Q @ K^T (Q 乘以 K 的转置)
    2. 缩放 (Scale): Scaled_Score = Score / sqrt(d_k)，其中 d_k 是 K (或 Q) 的维度。缩放是为了防止点积结果过大，导致 softmax 函数梯度过小。
    3. (可选) Masking (掩码): 在计算 Softmax 之前应用。
    4. 计算权重 (Attention Weights): Attention_Weights = softmax(Scaled_Score_with_Mask)
    5. 加权求和: Output = Attention_Weights @ V

  代码实现 (attention 函数):

  def attention(query, key, value, mask=None, dropout=None):"Compute 'Scaled Dot Product Attention'"d_k = query.size(-1) # 获取 Q, K 的最后一个维度 (特征维度 d_k)# 1. 计算得分 (Score): Q @ K^T# query: [batch, h, seq_len_q, d_k]# key.transpose(-2, -1): [batch, h, d_k, seq_len_k]# scores: [batch, h, seq_len_q, seq_len_k]scores = torch.matmul(query, key.transpose(-2, -1))# 2. 缩放 (Scale)scores = scores / math.sqrt(d_k)# 3. (可选) Masking (掩码)if mask is not None:# mask 通常是一个布尔张量，True/1 表示有效位置，False/0 表示要屏蔽的位置# masked_fill 将 mask 中为 False/0 的位置在 scores 中对应的值替换为 -1e9 (一个很大的负数)scores = scores.masked_fill(mask == 0, -1e9)# 4. 计算注意力权重 (Attention Weights)# softmax 在最后一个维度 (seq_len_k，即 Key 的序列长度维度) 上进行p_attn = scores.softmax(dim=-1)# (可选) Dropout on attention weightsif dropout is not None:p_attn = dropout(p_attn)# 5. 加权求和: Attention_Weights @ V# p_attn: [batch, h, seq_len_q, seq_len_k]# value: [batch, h, seq_len_v, d_v] (通常 seq_len_k == seq_len_v)# output: [batch, h, seq_len_q, d_v]return torch.matmul(p_attn, value), p_attn # 返回加权和结果以及注意力权重本身
  

  • 这个函数是 Scaled Dot-Product Attention 的核心实现，它接收 Q, K, V 以及可选的掩码和 dropout。
  • 图示 完美对应了这个计算流程。

2.2.1. Padding Mask (填充掩码)

• 目的:

  • 在批量处理（batch processing）时，一个批次中的输入序列（无论是源语言句子还是目标语言句子）通常长度不同。
  • 为了能够将它们整合成一个固定大小的张量（tensor）进行高效的并行计算，我们会对较短的序列进行填充（padding），通常使用一个特殊的 <pad> 标记 (或其对应的ID，如0)。
  • 这些填充标记本身不包含任何有意义的语义信息。如果在计算注意力时将它们也考虑进去，模型可能会错误地将注意力分配给这些无意义的填充部分，从而影响模型的性能和学习效率。
  • Padding Mask 的作用就是告诉模型在计算注意力时忽略这些填充标记。

• 如何工作:

  1. 识别填充位置: 首先，需要确定输入序列中哪些位置是填充标记。这通常通过检查词元的ID是否等于预定义的填充ID来完成。

  2. 创建掩码张量: 基于填充位置，创建一个与注意力得分矩阵（Q * K^T 之后，softmax 之前）形状相同或可广播的掩码张量。

     • 对于非填充位置，掩码值为 0。
     • 对于填充位置，掩码值为一个非常大的负数 (例如 -1e9 或 -infinity)。

  3. 应用掩码: 在计算 Softmax 之前，将这个掩码张量==加==到缩放后的注意力得分上：

     masked_scores = scaled_scores + padding_mask
     

  4. Softmax 的效果:

     • 当一个得分加上一个非常大的负数后，其值会变得非常小。
     • 经过 Softmax 函数 exp(x) / sum(exp(x_i)) 计算后，这些位置的注意力权重会趋近于 0。
     • 这样，模型在后续计算加权和时，就几乎不会给填充位置分配任何权重，相当于忽略了它们。

• 使用位置:

  • Encoder 的 Self-Attention 层: 对源语言输入序列的填充部分进行掩码。
  • Decoder 的 Self-Attention 层: 对目标语言输入序列的填充部分进行掩码 (在训练时，目标序列也可能被填充)。
  • Decoder 的 Encoder-Decoder Attention (Cross-Attention) 层: 这里掩码的是 Encoder 输出中对应于源语言序列填充的部分。Decoder 在关注 Encoder 输出时，不应该关注源序列的填充部分。

代码实现 (在 Batch 类中创建 src_mask 和部分 tgt_mask):

class Batch:def __init__(self, src, tgt=None, pad=2): # pad 通常是 padding token 的 IDself.src = src# 创建 src_mask:# (src != pad) 会生成一个布尔张量，padding 位置为 False，非 padding 为 True# .unsqueeze(-2) 增加一个维度以匹配注意力分数的形状 (e.g., for broadcasting with multi-head)# src_mask 形状: [batch_size, 1, src_len]# 在 attention 函数中，mask == 0 的地方会被填充为 -1e9self.src_mask = (src != pad).unsqueeze(-2)if tgt is not None:self.tgt = tgt[:, :-1]self.tgt_y = tgt[:, 1:]# tgt_mask 的创建会涉及到 padding mask 和 sequence maskself.tgt_mask = self.make_std_mask(self.tgt, pad)self.ntokens = (self.tgt_y != pad).data.sum()@staticmethoddef make_std_mask(tgt, pad):"Create a mask to hide padding and future words."# 1. 创建目标序列的 Padding Mask# (tgt != pad) 形状: [batch_size, tgt_len]# .unsqueeze(-2) 形状: [batch_size, 1, tgt_len]tgt_pad_mask = (tgt != pad).unsqueeze(-2)# ... (Sequence Mask 的创建和合并见下文) ...return tgt_mask # 返回的是合并后的掩码

• self.src_mask = (src != pad).unsqueeze(-2):
  • 这一行直接根据源序列 src 和填充符号的 ID pad 创建了源序列的 Padding Mask。
  • src != pad 会产生一个布尔张量，其中 True 表示非填充位置，False 表示填充位置。
  • unsqueeze(-2) 是为了将掩码调整为 [batch_size, 1, seq_len] 的形状，使其能够与多头注意力中 [batch_size, num_heads, seq_len_q, seq_len_k] 形状的 scores 张量进行广播操作（当 mask == 0 时，scores 的相应位置被填充）。
• 在 make_std_mask 中，tgt_pad_mask = (tgt != pad).unsqueeze(-2) 同样为目标序列创建了 Padding Mask。

在 MultiHeadedAttention 中的应用:

class MultiHeadedAttention(nn.Module):# ...def forward(self, query, key, value, mask=None): #这里的 mask 可以是 src_mask 或 tgt_maskif mask is not None:# Same mask applied to all h heads.mask = mask.unsqueeze(1) # -> [batch_size, 1, 1, seq_len_k] for src_mask (if seq_len_q=1 for src_attn)# or [batch_size, 1, tgt_len, tgt_len] for tgt_mask# ...x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout # mask 传递给 attention 函数)# ...

• 传递给 MultiHeadedAttention 的 mask 参数（可以是 src_mask 或 tgt_mask）会进一步通过 mask.unsqueeze(1) 调整维度，以适应多头注意力的计算，然后传递给 attention 函数内部进行 scores.masked_fill(mask == 0, -1e9) 操作。

---

2.2.2. Sequence Mask (Look-ahead Mask / Subsequent Mask / Causal Mask)

• 目的:

• 这种掩码仅用于 Decoder 的 Self-Attention 层。

  • Decoder 的任务是生成目标序列，通常是自回归 (auto-regressive) 的，即生成当前词元时，只能依赖于已经生成的之前词元，而不能"看到"未来的词元。
  • 如果在训练时（通常使用 “teacher forcing”，即直接将完整的目标序列喂给 Decoder），Decoder 的 Self-Attention 能够访问到当前位置之后的词元，那么模型就会"作弊"，直接从后面的词元中获取信息来预测当前词元，而不是真正学习预测。
  • Sequence Mask 的作用就是防止 Decoder 在预测第 i 个词元时，关注到第 i+1, i+2, … 等后续词元的信息，确保预测的单向性。

• 如何工作:

  1. 创建掩码张量: 针对目标序列，创建一个与注意力得分矩阵形状相同的上三角矩阵（或根据具体实现，是对角线以上的元素）。

     • 对于允许关注的位置 (即当前位置 i 及其之前的位置 j <= i)，掩码值为 0。
     • 对于不允许关注的位置 (即未来位置 j > i)，掩码值为一个非常大的负数 (例如 -1e9 或 -infinity)。
     • 例如，对于一个长度为4的序列，当计算第2个词元 (索引1) 的注意力时，它可以关注词元0和词元1，但不能关注词元2和词元3。

  2. 应用掩码: 同样，在计算 Softmax 之前，将这个掩码张量加到缩放后的注意力得分上：

     masked_scores = scaled_scores + sequence_mask
     

  3. Softmax 的效果:

     • 与 Padding Mask 类似，未来位置的注意力得分加上一个非常大的负数后，其 Softmax 权重会趋近于 0。
     • 这样，模型在计算当前位置的表示时，只会从之前和当前位置获取信息。

• 使用位置:

  • 仅在 Decoder 的第一个 Multi-Head Self-Attention 子层中使用。
  • Encoder 不需要这种掩码，因为 Encoder 的目的是编码整个输入序列，它可以同时看到所有词元。
  • Decoder 的 Encoder-Decoder Attention (Cross-Attention) 层也不需要这种掩码，因为它关注的是 Encoder 的完整输出，而不是目标序列的未来部分。

代码实现 (subsequent_mask 函数 和 Batch.make_std_mask 方法):

def subsequent_mask(size): # size 是目标序列的长度"Mask out subsequent positions."attn_shape = (1, size, size) # (1, tgt_len, tgt_len)# torch.triu(tensor, diagonal=k) 返回张量的上三角部分，对角线以下元素置0# diagonal=1 表示主对角线以上的元素保留，主对角线及以下为0 (如果我们想要屏蔽的是严格的未来)# 这里用 torch.ones 创建一个全1的矩阵，然后取其上三角部分(diagonal=1)# 这意味着 (pos_q < pos_k) 的位置为 1，其余为 0# .type(torch.uint8) 可能是为了旧版兼容性，现在可以直接用布尔subsequent_mask_values = torch.triu(torch.ones(attn_shape), diagonal=1)# 我们希望的是允许看过去和现在，屏蔽未来# 所以，(pos_q >= pos_k) 的位置应该是 True (允许)，(pos_q < pos_k) 的位置应该是 False (屏蔽)# 因此，原始的 subsequent_mask_values 中为 1 的地方（未来）应该对应 mask 中的 False/0# 原始的 subsequent_mask_values 中为 0 的地方（过去和现在）应该对应 mask 中的 True/1return subsequent_mask_values == 0 # True 表示允许，False 表示屏蔽

• subsequent_mask(size) 函数创建了一个形状为 [1, size, size] 的布尔掩码。
• torch.triu(torch.ones(attn_shape), diagonal=1) 会生成一个矩阵，其中 (row, col) 位置如果 col > row (即未来位置)，则为 1，否则为 0。
• subsequent_mask_values == 0 将其反转，使得如果 col <= row (当前或过去位置)，则为 True (允许)，否则为 False (屏蔽)。这个 True/False 的含义与 attention 函数中 mask == 0 表示屏蔽是一致的。

在 Batch.make_std_mask 中合并 Padding Mask 和 Sequence Mask:

    @staticmethoddef make_std_mask(tgt, pad): # tgt 是 Decoder 的输入 (tgt[:, :-1])"Create a mask to hide padding and future words."# 1. 创建目标序列的 Padding Mask (如上所述)# tgt_pad_mask: [batch_size, 1, tgt_len], True 表示非填充，False 表示填充tgt_pad_mask = (tgt != pad).unsqueeze(-2)# 2. 创建 Sequence Mask# subsequent_mask(tgt.size(-1)) 返回 [1, tgt_len, tgt_len]# True 表示允许看的位置 (当前和过去)，False 表示要屏蔽的未来位置seq_mask = subsequent_mask(tgt.size(-1)).type_as(tgt_pad_mask.data)# 3. 合并两个掩码# 只有当一个位置既不是填充 (tgt_pad_mask 为 True)# 并且也不是未来位置 (seq_mask 为 True) 时，# 该位置才被认为是有效的 (最终 tgt_mask 为 True)。# 逻辑与 (&) 操作确保了这一点。tgt_mask = tgt_pad_mask & seq_maskreturn tgt_mask

• tgt_pad_mask & seq_mask: 这里的 & 是逐元素的逻辑与操作。
  • tgt_pad_mask 确保了 padding 位置（值为 False）在最终的 tgt_mask 中也为 False (被屏蔽)。
  • seq_mask 确保了未来位置（值为 False）在最终的 tgt_mask 中也为 False (被屏蔽)。
  • 因此，tgt_mask 中为 True 的位置，必须同时满足“非填充”和“非未来”两个条件。这个 tgt_mask 会被传递给 Decoder 的 self_attn (Masked Multi-Head Self-Attention) 模块。

总结与对比:

特性	Padding Mask	Sequence Mask (Look-ahead Mask)
目的	忽略无意义的填充标记	防止 Decoder “看到” 未来词元
作用对象	序列中的 <pad> 标记	目标序列中当前位置之后的词元
应用层	Encoder Self-Attention	Decoder Self-Attention (仅此)
	Decoder Self-Attention
	Decoder Encoder-Decoder Attention
如何实现	将填充位置的得分变为极大负数	将未来位置的得分变为极大负数
代码对应	(input != pad_id)	torch.triu(...) == 0
	结合 scores.masked_fill(mask==0, -1e9)	结合 scores.masked_fill(mask==0, -1e9)
原因	提高效率，避免错误学习	保证自回归生成过程的正确性

在 Decoder 的 Self-Attention 层中，Batch.make_std_mask 生成的 tgt_mask 同时实现了 Padding Mask 和 Sequence Mask 的功能。这个组合掩码被传递给 attention 函数，通过 scores.masked_fill(mask == 0, -1e9) 应用，确保了在计算注意力权重时，既不关注填充词元，也不关注未来的词元。

---

2.3. Multi-Head Attention (多头注意力)

注意力机制说明

多头注意力是 Transformer 模型中的一个核心创新，它通过并行地执行多次 Scaled Dot-Product Attention 来增强模型的表达能力。

2.3.1. 为什么需要多头注意力? (Motivation)

• 捕捉多样化的相关性: 单一的注意力机制（即只有一个 “head”）在计算注意力权重时，可能只能学习到输入序列中一种特定类型的依赖关系或特征模式。例如，它可能擅长捕捉句法依赖，但忽略了语义上的相似性。
• 在不同表示子空间中学习: 多头注意力允许模型在不同的“表示子空间”（representation subspaces）中共同关注来自不同位置的信息。每个 “head” 可以学习关注输入的不同方面。
• 避免平均效应: “With a single attention head, averaging inhibits this.” 这句话指出，如果只有一个注意力头，其输出是基于单一加权平均得到的，这可能会“抑制”模型捕捉到多种不同类型的细微特征的能力，因为这些特征可能在平均过程中被模糊或抵消。多头机制通过让每个头专注于不同的子空间，然后将它们的结果结合起来，从而缓解了这个问题。

2.3.2. 如何实现? (Mechanism)

多头注意力的实现过程可以分解为以下几个步骤，对应于其数学表达式和代码实现：

1. 线性投影 (Linear Projections):

   • 原始的 Query (Q), Key (K), 和 Value (V) 向量（通常维度为 d_model）分别通过 h 组不同的、可学习的线性变换（权重矩阵 W_i^Q, W_i^K, W_i^V）投影到更低的维度。
   • 对于每个头 i (从 1 到 h)：
     • Q_i = Q @ W_i^Q
     • K_i = K @ W_i^K
     • V_i = V @ W_i^V
   • 通常，投影后的维度 d_k (Query 和 Key 的维度) 和 d_v (Value 的维度) 会被设置为 d_model / h。这样，所有头的计算总和与原始维度的单头注意力计算量相似。
   • 代码对应: self.linears 中的前三个线性层分别对应于所有头的 Q, K, V 投影。例如，self.linears[0] 用于投影 Q，self.linears[1] 用于投影 K，self.linears[2] 用于投影 V。这些线性层将 d_model 维的输入映射到 d_model 维的输出（因为 d_model = h * d_k），然后通过 .view() 和 .transpose() 操作将其分割并重排成 h 个头的形式。

2. 并行执行 Scaled Dot-Product Attention:

   • 对每个投影后得到的 (Q_i, K_i, V_i) 三元组并行地执行 Scaled Dot-Product Attention。
   • head_i = Attention(Q_i, K_i, V_i) = softmax( (Q_i @ K_i^T) / sqrt(d_k) ) @ V_i
   • (这里的 Attention 函数包含了可选的 Masking 操作，如代码中的 attention 函数所示)。
   • 这将产生 h 个不同的输出向量 head_i，每个 head_i 的维度是 d_v。
   • 代码对应: attention(query, key, value, mask=mask, dropout=self.dropout) 函数被调用，此时 query, key, value 已经是经过投影、分割和重排后的多头张量，形状为 [nbatches, self.h, seq_len, self.d_k]。attention 函数内部会并行计算所有头的注意力。

3. 拼接 (Concatenation):

   • 将这 h 个注意力头的输出 head_i 在最后一个维度上拼接起来。
   • Concat(head_1, ..., head_h) 得到一个维度为 h * d_v 的向量。
   • 由于通常 d_v = d_k = d_model / h，所以拼接后的维度是 h * (d_model / h) = d_model。
   • 代码对应:

     # x 是 attention 函数的输出，形状是 [nbatches, self.h, seq_len, self.d_k]
     x = (x.transpose(1, 2)  # 交换头和序列长度维度 -> [nbatches, seq_len, self.h, self.d_k].contiguous()      # 确保内存连续.view(nbatches, -1, self.h * self.d_k) # 重塑为 [nbatches, seq_len, d_model]
     )
     

     这一系列操作有效地将不同头的输出“并排”起来。

4. 最终线性变换 (Final Linear Projection):

   • 将拼接后的向量再通过一个可学习的线性变换 (权重矩阵 W^O) 得到最终的多头注意力输出。这个 W^O 的维度是 [h * d_v, d_model] (在代码中是 [d_model, d_model])。
   • MultiHead(Q, K, V) = Concat(head_1, ..., head_h) @ W^O
   • 代码对应: self.linears[-1](x)。self.linears 中的第四个线性层 (self.linears[3]) 就是这里的 W^O。

2.3.3. 数学表达式回顾

• Scaled Dot-Product Attention (单头核心):
  $$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
• Multi-Head Attention:
  $$\begin{gathered} \mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{hea}{\mathrm{d}}_1,...,\mathrm{hea}{\mathrm{d}}_{\mathrm{h}})W^O \\ \text{where}\mathrm{hea}{\mathrm{d}}_{\mathrm{i}}=\mathrm{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V) \end{gathered}$$
• 参数矩阵维度:
  • W^Q_i ∈ R^(d_model × d_k)
  • W^K_i ∈ R^(d_model × d_k)
  • W^V_i ∈ R^(d_model × d_v)
  • W^O ∈ R^(h*d_v × d_model)
• 论文中的典型设置:
  • h = 8 (平行注意力头数)
  • d_k = d_v = d_model / h = 64 (如果 d_model = 512)
  • 这种设置使得多头注意力的总计算成本与维度为 d_model 的单头注意力相似，因为每个头的维度降低了。

2.3.4. PyTorch 代码解读 (MultiHeadedAttention 类)

def attention(query, key, value, mask=None, dropout=None):  "Compute 'Scaled Dot Product Attention'"  d_k = query.size(-1)  scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)  if mask is not None:  scores = scores.masked_fill(mask == 0, -1e9)  p_attn = scores.softmax(dim=-1)  if dropout is not None:  p_attn = dropout(p_attn)  return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):def __init__(self, h, d_model, dropout=0.1):super(MultiHeadedAttention, self).__init__()assert d_model % h == 0  # 确保 d_model 可以被头数 h 整除self.d_k = d_model // h  # 每个头的维度self.h = h               # 头数# 创建 4 个线性层: 3 个用于 Q, K, V 的初始投影 (W_i^Q, W_i^K, W_i^V 的集合)# 第 4 个用于最终的输出投影 (W^O)# 注意: 这里每个线性层输出维度是 d_model，之后会通过 view 和 transpose 分割成 h 个头self.linears = clones(nn.Linear(d_model, d_model), 4)self.attn = None         # 用于存储注意力权重，方便可视化或分析self.dropout = nn.Dropout(p=dropout)def forward(self, query, key, value, mask=None):if mask is not None:# 对所有头应用相同的 mask，在 batch 维度后增加一个头维度mask = mask.unsqueeze(1) # mask: [nbatches, 1, seq_len_q, seq_len_k] or [nbatches, 1, 1, seq_len_k]nbatches = query.size(0)# 1) 线性投影 + 重塑为多头形式# query, key, value 初始形状: [nbatches, seq_len, d_model]query, key, value = [lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)for lin, x in zip(self.linears, (query, key, value))]# 投影后 lin(x) 形状: [nbatches, seq_len, d_model]# .view(...) 形状: [nbatches, seq_len, self.h, self.d_k]# .transpose(1, 2) 形状: [nbatches, self.h, seq_len, self.d_k] (这是 attention 函数期望的输入格式)# 2) 应用 Scaled Dot-Product Attention (并行处理所有头)x, self.attn = attention( # self.attn 存储的是 p_attnquery, key, value, mask=mask, dropout=self.dropout)# x (输出) 形状: [nbatches, self.h, seq_len_q, self.d_k]# self.attn (注意力权重) 形状: [nbatches, self.h, seq_len_q, seq_len_k]# 3) 拼接 + 最终线性变换x = (x.transpose(1, 2) # -> [nbatches, seq_len_q, self.h, self.d_k].contiguous()     # 确保内存是连续的，以便 view 操作.view(nbatches, -1, self.h * self.d_k) # -> [nbatches, seq_len_q, d_model])# del query, key, value # 可选，用于及时释放显存return self.linears[-1](x) # -> [nbatches, seq_len_q, d_model]

2.4. Feed-Forward Network (FFN, 前馈网络)

在 Transformer 的每个 Encoder Layer 和 Decoder Layer 中，紧随多头注意力子层 (及其 Add & Norm) 之后，都有一个按位置前馈网络 (Position-wise Feed-Forward Network, FFN)。这个 FFN 为模型提供了非线性和进一步的特征变换能力。

2.4.1. FFN 的结构与功能

• 结构: FFN 是一个相对简单的两层全连接神经网络。

  1. 第一个线性变换: 将输入 x (维度为 d_model) 投影到一个更高的中间维度 d_ff (通常 d_ff 是 d_model 的几倍，例如 d_ff = 4 * d_model，在原始论文中 d_model=512, d_ff=2048)。
     output_1 = x @ W_1 + b_1
  2. 激活函数: 对第一个线性变换的输出应用一个非线性激活函数，通常是 ReLU (Rectified Linear Unit)。
     activated_output = ReLU(output_1) = max(0, output_1)
     (一些现代 Transformer 变体可能使用 GELU 或 SwiGLU 等其他激活函数。)
  3. (可选) Dropout: 在激活函数之后、第二个线性变换之前，可以应用 Dropout 以进行正则化。
  4. 第二个线性变换: 将激活后的中间表示 (维度为 d_ff) 投影回原始维度 d_model。
     FFN_output = activated_output @ W_2 + b_2

• 数学表达式:
  $$\mathrm{FFN}(x)=\max (0,x{W}_1+b_1)W_2+b_2$$
  (如果包含 dropout，则是在 max(0, ...) 之后，W_2 之前应用。)

• 功能与目的:

  • 非线性变换: 注意力机制本身主要是线性的（除了 softmax），FFN 引入了必要的非线性，增强了模型的表达能力。
  • 特征交互: 允许在每个词元位置的表示内部，不同维度之间进行更复杂的交互和信息转换。
  • 维度扩展与压缩: 通过先扩展到 d_ff 再压缩回 d_model，模型可以在一个更高维的空间中学习更丰富的特征表示，然后再将这些信息提炼回原始维度。
  • 独立处理: FFN 的 “Position-wise” 特性意味着它独立地应用于序列中的每个词元（或位置）的表示上。

2.4.2. “Position-wise” 的含义

• 独立应用: 对于输入序列中的每个词元 x_i (其中 i 是词元在序列中的位置)，FFN 都使用完全相同的权重矩阵 W_1, b_1, W_2, b_2 来进行计算。
• 参数共享: 所有位置共享同一组 FFN 参数。这大大减少了模型的参数量，相比于为每个位置都学习一套独立的 FFN 参数。
• 并行计算: 由于每个位置的 FFN 计算是独立的，因此可以高效地并行处理序列中的所有位置。在实践中，这通常通过对整个序列张量 (形状如 [batch_size, sequence_length, d_model]) 应用 1D 卷积 (核大小为1) 或直接使用 nn.Linear (它会自动在除最后一个维度外的其他维度上广播) 来实现。

2.4.3. PyTorch 代码解读 (PositionwiseFeedForward 类)

class PositionwiseFeedForward(nn.Module):"Implements FFN equation."def __init__(self, d_model, d_ff, dropout=0.1):super(PositionwiseFeedForward, self).__init__()# 第一个线性层: d_model -> d_ffself.w_1 = nn.Linear(d_model, d_ff)# 第二个线性层: d_ff -> d_modelself.w_2 = nn.Linear(d_ff, d_model)# Dropout 层self.dropout = nn.Dropout(dropout)def forward(self, x):# x 的形状通常是 [batch_size, sequence_length, d_model]# 1. 应用第一个线性变换 (self.w_1)# 2. 应用 ReLU 激活函数 (.relu())# 3. 应用 Dropout (self.dropout(...))# 4. 应用第二个线性变换 (self.w_2)return self.w_2(self.dropout(self.w_1(x).relu()))

• d_model: 输入和输出的维度。
• d_ff: FFN 内部隐藏层的维度。
• self.w_1: 对应数学公式中的 W_1 (和 b_1)。
• self.w_2: 对应数学公式中的 W_2 (和 b_2)。
• forward 方法清晰地实现了 FFN 的计算流程：Linear -> ReLU -> Dropout -> Linear。
  • self.w_1(x): PyTorch 的 nn.Linear 会自动处理 x 的前导维度 (batch 和 sequence length)，只在最后一个维度 (d_model) 上进行线性变换。
  • .relu(): 逐元素应用 ReLU。
  • self.dropout(...): 应用 dropout。
  • self.w_2(...): 再次应用线性变换。

2.4.4. 与 Embeddings 类的关系 (补充)

class Embeddings(nn.Module):def __init__(self, d_model, vocab):super(Embeddings, self).__init__()self.lut = nn.Embedding(vocab, d_model) # Look-Up Tableself.d_model = d_modeldef forward(self, x):# x 是词元 ID 序列# self.lut(x) 将 ID 转换为 d_model 维的词嵌入向量# 乘以 sqrt(d_model) 是原始 Transformer 论文中的一个细节，# 目的是调整嵌入的尺度，可能与后续的位置编码结合方式有关return self.lut(x) * math.sqrt(self.d_model)

这个 Embeddings 类是 Transformer 输入管道的一部分，它负责将输入的离散词元 ID 转换为 d_model 维的密集向量表示。这些词嵌入向量（通常与位置编码相加后）会成为 Encoder (或 Decoder 的 Masked Self-Attention) 的初始输入。

FFN 处理的是这些经过了（多层）自注意力机制和Add & Norm操作后的 d_model 维向量。 FFN 进一步对这些已经包含了上下文信息的词元表示进行变换。

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

在每个编码器（Encoder）和解码器（Decoder）层中，每一个子层（如多头注意力机制或前馈网络）之后都会跟随一个"Add & Norm"步骤。这包含两个操作：一个残差连接 (Add)，然后是一个层归一化 (Norm)。

• Residual Connection (残差连接 - Add):

  • 工作原理: 子层（例如，多头注意力）的输出会与其自身的输入相加。

    子模块的输出 = 子模块的输入 + 子模块(子模块的输入)
    

    所以，如果 x 是子层（例如多头注意力）的输入，Sublayer(x) 是该子层学习到的函数（即子层的输出），那么残差连接的输出就是 x + Sublayer(x)。

  • 目的与优势:

    1. 缓解梯度消失 (Combating Vanishing Gradients): 这是主要的好处。在非常深的网络中，当梯度通过许多层反向传播时，它们可能会变得非常小，使得早期层难以有效学习。残差连接为梯度提供了一条"捷径"或更直接的路径回传，使得梯度更容易传播到较早的层，从而有助于训练更深的模型。
    2. 实现恒等映射 (Enabling Identity Mapping): 如果某个特定的层或变换 Sublayer(x) 没有益处（即对当前任务来说是多余的），网络可以学习使 Sublayer(x) 的输出接近于零。在这种情况下，输出 x + Sublayer(x) 将近似于 x，这意味着该层有效地学习了一个恒等映射并"跳过"了该变换。这使得层在不需要复杂函数时更容易学习简单的函数，防止它们降低性能。
    3. 更容易优化 (Easier Optimization): 通过允许层学习对恒等映射的修改（即学习残差 Sublayer(x)），优化过程可能会变得更平滑且更容易。
    4. 特征复用 (Feature Re-use): 原始输入 x 被直接传递，确保了早期阶段的信息在经过转换时不会丢失。

• Layer Normalization (层归一化 - Norm):

  • 工作原理: 层归一化在残差连接_之后_应用。

    模块最终输出 = LayerNorm(子模块的输入 + 子模块(子模块的输入))
    

    它独立地对批次中每个样本/词元的特征（embedding vector中的所有元素）进行归一化。

    1. 对于每个词元的嵌入向量，它计算其所有特征（向量中的元素）的均值和方差。

    2. 然后使用这个均值和方差对特征进行归一化：

       归一化特征 = (特征 - 均值) / sqrt(方差 + epsilon)
       

       （epsilon 是一个用于数值稳定性的很小的常数，防止除以零）。

    3. 最后，它应用一个可学习的仿射变换：

       输出 = gamma * 归一化特征 + beta
       

       gamma（缩放因子）和 beta（平移因子）是每个特征维度可学习的参数，允许网络在需要时缩放和平移归一化后的输出，如果这是最优的，甚至可能恢复原始表示。

  • 目的与优势:

    1. 稳定训练 (Stabilizing Training): 它通过将激活值保持在更一致的范围内，有助于稳定训练过程中的隐藏状态动态。这减少了"内部协变量偏移"（Internal Covariate Shift，即随着前一层参数的改变，每层输入的分布在训练过程中发生变化）的问题，使得训练更加鲁棒。
    2. 加速收敛 (Accelerating Convergence): 通过稳定激活值，层归一化可以使用更高的学习率并加速收敛。
    3. 降低对初始化的敏感性 (Reducing Sensitivity to Initialization): 归一化过程使得模型对权重的初始值不那么敏感。
    4. 与批量大小无关 (Independence of Batch Size): 与批量归一化（Batch Normalization，它在批量维度上进行归一化）不同，层归一化的统计数据是按样本计算的。这使得它即使在小批量大小或对于批量统计可能不那么有意义的序列数据（如NLP任务）也有效。

原始 Transformer 中的顺序 (“Post-LN”): 在原始的《Attention Is All You Need》论文中，层归一化是在残差连接之后应用的：

Output = LayerNorm(x + Sublayer(x))

这通常被称为 Post-LN (后置层归一化)。

替代方案 (“Pre-LN”): 后来的研究发现，在子层和残差连接_之前_应用层归一化有时可以带来更稳定的训练，特别是对于非常深的 Transformer：

Output = x + Sublayer(LayerNorm(x))

这被称为 Pre-LN (前置层归一化)。虽然原始论文使用的是 Post-LN，但由于其改进的训练稳定性和在没有仔细的学​​习率预热（warmup）情况下训练更深模型的能力，许多现代实现和变体（如GPT-2及之后的一些模型）可能会选择 Pre-LN。

3. Encoder 结构

Transformer 的 Encoder 负责处理输入序列（例如，源语言句子），并将其转换为一系列上下文相关的向量表示 (context vectors)。这些向量随后可以被 Decoder 用来生成目标序列，或者在 Encoder-Only 模型 (如 BERT) 中直接用于下游任务。

整个 Encoder 由 N 个完全相同的 Encoder Layer 堆叠而成。每个 Encoder Layer 都执行相同的转换操作，但拥有独立的权重。这种堆叠结构允许模型逐步构建更复杂、更抽象的输入序列表示。

(图示：单个 Encoder Layer 的内部结构，展示了两个主要的子层：多头自注意力和前馈网络，以及它们之间的残差连接和层归一化。)

3.1. 单个 Encoder Layer 的构成

如图所示，一个标准的 Encoder Layer 主要由两个子层 (sub-layers) 构成：

1. Multi-Head Self-Attention Layer (多头自注意力层):

   • 功能: 这是 Encoder 的核心组件。它允许输入序列中的每个词元关注序列中的所有其他词元（包括其自身），从而计算出每个词元在当前上下文中的加权表示。这使得模型能够捕捉序列内部的长距离依赖关系。
   • 输入: 来自前一层 (或输入嵌入层) 的词元向量序列。
   • 输出: 经过自注意力加权和组合后的词元向量序列，与输入序列等长。
   • Masking: 在自注意力计算中，通常会应用 Padding Mask，以确保模型忽略输入序列中由于批处理而添加的填充 (padding) 词元。

2. Position-wise Feed-Forward Network (FFN, 按位置前馈网络):

   • 功能: 这是一个简单的全连接前馈网络，它独立地应用于序列中的每个词元位置。也就是说，所有位置共享相同的 FFN 权重，但每个位置的计算是分开的。
   • 结构: 通常由两个线性变换和一个 ReLU (或 GELU 等) 激活函数组成： FFN(x) = max(0, xW_1 + b_1)W_2 + b_2。
   • 目的: 在自注意力捕捉了全局上下文信息后，FFN 提供了额外的非线性变换和特征提取能力，进一步处理每个位置的表示。它也可以被看作是在不同维度间进行信息交互。

3.2. 子层连接：Add & Norm (残差连接与层归一化)

在每个 Encoder Layer 中，上述两个子层 (多头自注意力和 FFN) 的输出并不会直接作为下一部分的输入。它们都分别包裹在一个 “Add & Norm” 操作中：

• Residual Connection (Add - 残差连接):

  • 操作: 子层的输出与其自身的输入相加： output = x + Sublayer(x)。
  • 目的:
    • 缓解梯度消失问题，使梯度更容易在深层网络中传播。
    • 允许模型学习恒等映射，如果子层没有学到有用的信息，可以直接“跳过”该层。
    • 加速训练过程。

• Layer Normalization (Norm - 层归一化):

  • 操作: 对残差连接后的输出进行归一化。它独立地对每个样本（序列中的每个词元向量）的特征进行归一化。

class LayerNorm(nn.Module):
“Construct a layernorm module (See citation for details).”

def __init__(self, features, eps=1e-6):  super(LayerNorm, self).__init__()  self.a_2 = nn.Parameter(torch.ones(features))  self.b_2 = nn.Parameter(torch.zeros(features))  self.eps = eps  def forward(self, x):  mean = x.mean(-1, keepdim=True)  std = x.std(-1, keepdim=True)  return self.a_2 * (x - mean) / (std + self.eps) + self.b_2# LayerNorm 实现 (参考提供代码)# mean = x.mean(-1, keepdim=True)# std = x.std(-1, keepdim=True)# normalized_x = (x - mean) / (std + eps)# output = gamma * normalized_x + beta# (其中 gamma 是 self.a_2, beta 是 self.b_2)```
*   **目的**:*   稳定训练过程中的隐藏层激活值，减少内部协变量偏移。*   加速模型收敛，并降低对权重初始化的敏感性。*   与批量大小无关，适用于序列数据。

连接方式的变体:
原始 Transformer 论文采用的是 Post-LN 结构，即 LayerNorm(x + Sublayer(x))。
提供的 SublayerConnection 代码实现了一种 Pre-LN 结构：x + Sublayer(LayerNorm(x))。

class SublayerConnection(nn.Module):  """  A residual connection followed by a layer norm.    Note for code simplicity the norm is first as opposed to last.    """  def __init__(self, size, dropout):  super(SublayerConnection, self).__init__()  self.norm = LayerNorm(size)  self.dropout = nn.Dropout(dropout)  def forward(self, x, sublayer):  "Apply residual connection to any sublayer with the same size."  return x + self.dropout(sublayer(self.norm(x)))
# SublayerConnection 实现 (Pre-LN)
# class SublayerConnection(nn.Module):
#     ...
#     def forward(self, x, sublayer):
#         return x + self.dropout(sublayer(self.norm(x))) # norm(x) 先于 sublayer

Pre-LN 在一些后续研究中被发现可以带来更稳定的训练，尤其对于非常深的网络。

3.3. Encoder Layer 的 PyTorch 实现思路

构建逻辑：

# EncoderLayer 定义
class EncoderLayer(nn.Module):def __init__(self, size, self_attn, feed_forward, dropout):super(EncoderLayer, self).__init__()self.self_attn = self_attn  # 传入实例化的多头自注意力模块self.feed_forward = feed_forward # 传入实例化的前馈网络模块# 创建两个 SublayerConnection 实例，分别用于自注意力和FFNself.sublayer = clones(SublayerConnection(size, dropout), 2)self.size = size # 通常是 d_modeldef forward(self, x, mask): # x 是输入序列，mask 是 padding mask# 第一个子层: 多头自注意力 + Add & Norm# 使用 lambda 函数包装 self_attn 调用，以匹配 SublayerConnection 的接口x = self.sublayer[0](x, lambda x_input: self.self_attn(x_input, x_input, x_input, mask))# 第二个子层: 前馈网络 + Add & Normx = self.sublayer[1](x, self.feed_forward)return x

• size: 代表模型维度 d_model。
• self_attn: 一个多头自注意力模块实例。
• feed_forward: 一个前馈网络模块实例。
• dropout: 应用于子层输出的 dropout 概率。
• sublayer: 使用 clones 函数（一个简单的模块列表深拷贝工具）创建了两个 SublayerConnection 实例。SublayerConnection 封装了 “Add & Norm” 的逻辑（在提供的代码中是 Pre-LN）。
• 在前向传播中，输入 x 首先通过第一个 SublayerConnection 和自注意力模块，然后其输出再通过第二个 SublayerConnection 和前馈网络模块。

3.4. 整体 Encoder 结构

整个 Encoder 由 N 个这样的 EncoderLayer 实例堆叠而成。

# Encoder 定义
class Encoder(nn.Module):def __init__(self, layer_prototype, N): # layer_prototype 是一个 EncoderLayer 实例super(Encoder, self).__init__()self.layers = clones(layer_prototype, N) # 堆叠 N 个 EncoderLayer# 在所有 Encoder Layer 之后，再额外应用一个 LayerNormself.norm = LayerNorm(layer_prototype.size)def forward(self, x, mask): # x 是初始的输入嵌入+位置编码，mask 是 padding maskfor layer in self.layers:x = layer(x, mask) # 数据依次流过每个 EncoderLayerreturn self.norm(x) # 最终输出前再进行一次层归一化

• 输入 x (通常是词嵌入与位置编码之和) 和 mask (padding mask) 依次通过堆叠的 N 个 EncoderLayer。
• 在所有 N 个层处理完毕后，通常还会再进行一次层归一化 (self.norm)，然后才输出最终的编码器表示。这个最后的层归一化步骤在原始论文的图示中并不总是明确画出，但一些实现（如提供的代码）会包含它。

4. Decoder 结构

Transformer 的 Decoder 负责根据 Encoder 提供的输入序列的上下文表示 (context vectors) 和已经生成的部分目标序列，来生成下一个目标词元。这个过程是自回归的，即逐个词元生成，直到产生一个特殊的结束符 <EOS> 或达到预设的最大长度。

与 Encoder 类似，整个 Decoder 也是由 N 个完全相同的 Decoder Layer 堆叠而成。

(图示：单个 Decoder Layer 的内部结构，展示了三个主要的子层：掩码多头自注意力、多头编码器-解码器注意力（交叉注意力）和前馈网络，以及它们之间的残差连接和层归一化。)

4.1. 单个 Decoder Layer 的构成

如图所示，一个标准的 Decoder Layer 比 Encoder Layer 多一个子层，总共包含三个主要的子层：

1. Masked Multi-Head Self-Attention Layer (掩码多头自注意力层):

   • 功能: 与 Encoder 中的自注意力类似，它允许目标序列中的每个词元关注目标序列中先前已生成的词元（包括其自身）。
   • 关键区别 (Masking):
     • Sequence Mask (Look-ahead Mask / Causal Mask): 这是 Decoder 自注意力的核心特性。它确保在预测第 i 个词元时，模型只能访问到位置 i 及其之前 (0 到 i) 的词元信息，而不能 “看到” 未来位置 (i+1 及之后) 的词元。这是为了保证自回归生成的特性，即预测当前词只能依赖于过去。
     • Padding Mask: 与 Encoder 一样，如果目标序列被填充，也需要应用 Padding Mask 来忽略填充词元。
   • 输入: 来自前一层 (或目标序列的输入嵌入层) 的目标词元向量序列。
   • 输出: 经过掩码自注意力加权和组合后的目标词元向量序列，与输入序列等长。

2. Multi-Head Encoder-Decoder Attention Layer (Cross-Attention / 源注意力层):

   • 功能: 这是连接 Encoder 和 Decoder 的桥梁。它允许 Decoder 的每个位置关注整个输入序列 (来自 Encoder 的输出) 的所有位置。这使得 Decoder 能够从源序列中提取与当前生成任务最相关的信息。
   • 输入:
     • Query (Q): 来自 Decoder 前一个子层（即掩码自注意力层）的输出。
     • Key (K) 和 Value (V): 来自 Encoder 最终层的输出 (memory in code)。
   • Masking: 这里的 Masking 主要是指对 Encoder 输出中的 Padding Mask (src_mask in code)。如果源序列有填充，Decoder 在关注 Encoder 输出时应忽略这些填充部分。
   • 输出: 一个结合了源序列信息的向量序列，代表了 Decoder 在当前步骤对源序列的理解。

3. Position-wise Feed-Forward Network (FFN, 按位置前馈网络):

   • 功能与结构: 与 Encoder 中的 FFN 完全相同，独立地应用于序列中的每个词元位置，提供额外的非线性变换。
   • 输入: 来自前一个子层（即编码器-解码器注意力层）的输出。
   • 输出: 经过 FFN 处理后的词元向量序列。

4.2. 子层连接：Add & Norm

与 Encoder 一样，Decoder Layer 中的这三个子层 (掩码自注意力、编码器-解码器注意力、FFN) 的输出也都分别包裹在一个 “Add & Norm” 操作中，即先进行残差连接，然后进行层归一化。这同样有助于缓解梯度问题、稳定训练并加速收景。

• SublayerConnection: 代码中同样使用 SublayerConnection 来封装这种 “Add & Norm” 逻辑 (在提供的代码中是 Pre-LN 变体: x + Sublayer(LayerNorm(x)))。

4.3. Decoder Layer 的 PyTorch 实现思路

# DecoderLayer 定义
class DecoderLayer(nn.Module):def __init__(self, size, self_attn, src_attn, feed_forward, dropout):super(DecoderLayer, self).__init__()self.size = size  # 通常是 d_modelself.self_attn = self_attn      # 传入实例化的 Masked Multi-Head Self-Attention 模块self.src_attn = src_attn        # 传入实例化的 Multi-Head Encoder-Decoder Attention 模块self.feed_forward = feed_forward # 传入实例化的前馈网络模块# 创建三个 SublayerConnection 实例，分别用于上述三个子层self.sublayer = clones(SublayerConnection(size, dropout), 3)def forward(self, x, memory, src_mask, tgt_mask):# x: 当前目标序列的嵌入 (或上一 Decoder Layer 的输出)# memory: Encoder 的最终输出# src_mask: 源序列的 Padding Mask (用于 Encoder-Decoder Attention)# tgt_mask: 目标序列的 Mask (结合了 Sequence Mask 和 Padding Mask, 用于 Masked Self-Attention)m = memory # 别名，可读性# 第一个子层: Masked Multi-Head Self-Attention + Add & Normx = self.sublayer[0](x, lambda x_input: self.self_attn(x_input, x_input, x_input, tgt_mask))# 第二个子层: Multi-Head Encoder-Decoder Attention + Add & Norm# Query 来自 x (上一子层输出), Key 和 Value 来自 m (Encoder 输出)x = self.sublayer[1](x, lambda x_input: self.src_attn(x_input, m, m, src_mask))# 第三个子层: Position-wise Feed-Forward Network + Add & Normx = self.sublayer[2](x, self.feed_forward)return x

• self_attn: 对应图中的 “Masked Multi-Head Attention”。
• src_attn: 对应图中的 “Multi-Head Attention”，即 Encoder-Decoder Cross-Attention。
• feed_forward: 对应图中的 “Feed Forward”。
• sublayer: 创建了三个 SublayerConnection 实例。
• 在前向传播中：
  1. 目标序列 x 首先通过自身的掩码自注意力 (self_attn)，并应用 tgt_mask。
  2. 然后，其输出作为 Query，与 Encoder 的输出 memory (作为 Key 和 Value) 进行编码器-解码器注意力 (src_attn) 计算，并应用 src_mask。
  3. 最后，结果通过前馈网络 (feed_forward)。
     所有步骤都由 SublayerConnection 包裹。

4.4. 整体 Decoder 结构

整个 Decoder 由 N 个这样的 DecoderLayer 实例堆叠而成。

# Decoder 定义
class Decoder(nn.Module):def __init__(self, layer_prototype, N): # layer_prototype 是一个 DecoderLayer 实例super(Decoder, self).__init__()self.layers = clones(layer_prototype, N) # 堆叠 N 个 DecoderLayer# 在所有 Decoder Layer 之后，再额外应用一个 LayerNormself.norm = LayerNorm(layer_prototype.size)def forward(self, x, memory, src_mask, tgt_mask):# x: 目标序列的嵌入 + 位置编码 (在第一层时)# memory: Encoder 的最终输出# src_mask: 源序列的 Padding Mask# tgt_mask: 目标序列的 Maskfor layer in self.layers:x = layer(x, memory, src_mask, tgt_mask) # 数据依次流过每个 DecoderLayerreturn self.norm(x) # 最终输出前再进行一次层归一化

• Decoder 的输入包括目标序列的部分嵌入 x，Encoder 的输出 memory，以及两种掩码 src_mask 和 tgt_mask。
• 数据依次通过堆叠的 N 个 DecoderLayer。
• 与 Encoder 类似，在所有 N 个层处理完毕后，通常也会再进行一次层归一化 (self.norm)。这个输出随后会送入一个线性层和 Softmax 层来预测下一个词元的概率。

---

5. 输出部分 (Output Generation)

当 Transformer 的 Decoder 完成了对目标序列当前位置的编码（即 N 个 Decoder Layer 都处理完毕后），其最终的输出向量（通常维度为 d_model）需要被转换为对词汇表中下一个词元的概率预测。这个过程通常包含以下两个步骤：

1. 线性变换 (Linear Transformation):

   • Decoder 最后一层（或整体 Decoder 模块，在所有层和最后的 LayerNorm 之后）的输出向量序列，其每个位置的向量维度是 d_model。
   • 这个 d_model 维的向量首先会经过一个线性层 (Linear Layer)。
   • 这个线性层的目的是将 d_model 维的向量映射到词汇表大小 (vocab_size) 的维度。
   • 如果词汇表中有 V 个词元，那么这个线性层的输出维度就是 V。输出的每个元素可以看作是对应词汇表中词元的 “logit” 或 “得分”。
   • 数学表示: Logits = DecoderOutput @ W_proj + b_proj
     其中 W_proj 的维度是 [d_model, vocab_size]，b_proj 的维度是 [vocab_size]。

2. Softmax 函数:

   • 线性层输出的 vocab_size 维的 logits 向量随后会通过一个 Softmax 函数。
   • Softmax 函数将这些 logits 转换为一个概率分布，使得所有词元的概率总和为1，并且每个词元的概率值都在 0 到 1 之间。
   • 这个概率分布 P(token_i) 表示模型预测词汇表中第 i 个词元是下一个正确词元的概率。
   • 数学表示:

     P(token_i | context) = softmax(Logits)_i = exp(logit_i) / Σ_j exp(logit_j)
     

     其中 context 代表了到目前为止的输入和已生成的目标序列。

总结公式:

Probabilities = softmax(Linear(Decoder_Final_Output))

其中 Linear 是指上述的线性变换层，Decoder_Final_Output 是指 Decoder 堆栈最后一层（并经过最终 LayerNorm）的输出。

5.1. PyTorch 代码实现 (Generator 类)

在提供的代码中，这个输出生成步骤被封装在 Generator 类中：

class Generator(nn.Module):"Define standard linear + softmax generation step."def __init__(self, d_model, vocab_size): # 注意，这里参数名是 vocab，但其含义是词汇表大小super(Generator, self).__init__()# 线性层，将 d_model 维的输入映射到 vocab_size 维的 logitsself.proj = nn.Linear(d_model, vocab_size)def forward(self, x):# x 是 Decoder 的最终输出，形状通常是 [batch_size, sequence_length, d_model]# self.proj(x) 进行线性变换，输出形状是 [batch_size, sequence_length, vocab_size]# log_softmax 在最后一个维度（词汇表维度）上计算 log(softmax(logits))# 返回的是对数概率，这在计算 NLLLoss (Negative Log Likelihood Loss) 时很方便return log_softmax(self.proj(x), dim=-1)

• __init__(self, d_model, vocab_size): 初始化一个线性层 self.proj，其输入维度是 d_model，输出维度是 vocab_size。
• forward(self, x):
  • 接收 Decoder 的最终输出 x。
  • 通过 self.proj(x) 进行线性变换。
  • 使用 log_softmax(..., dim=-1) 计算对数 Softmax 概率。
    • 为什么是 log_softmax 而不是 softmax? 在训练时，如果使用负对数似然损失 (NLLLoss)，log_softmax 的输出可以直接作为其输入。NLLLoss(log_softmax(input), target) 等价于 CrossEntropyLoss(input, target)。直接输出 log_softmax 可以提高数值稳定性，并简化损失函数的计算。
    • 如果需要原始概率，可以对 log_softmax 的输出再取指数 exp()。

5.2. 在整体模型 (make_model) 中的集成

在 make_model 函数中，Generator 类的实例被创建并传递给 EncoderDecoder 模型：

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):# ... (其他组件的创建) ...model = EncoderDecoder(# ... (Encoder, Decoder, Embeddings 的创建) ...Generator(d_model, tgt_vocab), # <--- Generator 实例作为 EncoderDecoder 的一部分)# ... (参数初始化) ...return model

• 这里的 tgt_vocab 是目标语言词汇表的大小。

• EncoderDecoder 类在其 forward 方法中，可能会在调用 self.decode(...) 之后，再将 Decoder 的输出传递给 self.generator 来得到最终的词元概率（或对数概率）。（虽然在提供的 EncoderDecoder.forward 中直接返回了 self.decode 的结果，但通常 Generator 会在最顶层被调用，或者 decode 方法的最终输出在某个地方会经过 generator）。在实际应用中，EncoderDecoder 的 forward 方法通常会返回 Generator 的输出，因为这才是模型对下一个词元的预测。

  更常见的做法是在 EncoderDecoder 的 forward 中：

  # class EncoderDecoder(nn.Module):
  #     ...
  #     def forward(self, src, tgt, src_mask, tgt_mask):
  #         decoder_output = self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
  #         return self.generator(decoder_output) # <--- 调用 generator
  

  或者在推理（生成）循环中显式调用。

5.3. 生成阶段的选择策略

一旦获得了每个词元的（对数）概率分布，模型在生成 (inference/decoding) 阶段需要从中选择下一个词元：

• Greedy Search (贪心搜索):

  • 在每一步，简单地选择当前概率最高的词元作为下一个输出词元。
  • next_token_id = argmax(Probabilities)
  • 这种方法简单快速，但可能不是全局最优的，因为它只考虑了当前一步的最佳选择。

• Beam Search (集束搜索):

  • 在每一步，保留 k (beam width) 个概率最高的候选序列。
  • 在下一步，从这 k 个候选序列分别扩展所有可能的下一个词元，并计算所有新生成序列的累积概率（通常是累积对数概率）。
  • 再次选择新的 k 个概率最高的序列。
  • 重复此过程，直到所有 k 个序列都生成了 <EOS> 符号或达到最大长度。
  • 最终选择累积概率最高的完整序列作为输出。
  • Beam search 通常能产生比贪心搜索质量更高的序列，但计算成本更高。

• Sampling (采样):

  • 从概率分布中随机采样下一个词元。
  • Temperature Sampling: 可以引入一个温度参数 T 来调整 Softmax 的平滑度。softmax(logits / T)。
    • T > 1: 概率分布更平滑，增加多样性，但可能降低连贯性。
    • T < 1: 概率分布更尖锐，更倾向于高概率词元，减少多样性。
    • T = 1: 标准 Softmax。
  • Top-k Sampling: 只从概率最高的 k 个词元中进行采样。
  • Top-p (Nucleus) Sampling: 只从累积概率超过阈值 p 的最小词元集合中进行采样。
  • 采样方法可以生成更多样化和创造性的文本，但有时可能牺牲一定的准确性和连贯性。

总结:
Transformer 的输出部分通过一个线性层将 Decoder 的高维表示映射到词汇表大小的 logits，然后通过 Softmax (或 log_softmax) 得到下一个词元的概率分布。在推理阶段，会采用如贪心搜索、集束搜索或各种采样策略从这个分布中选择实际输出的词元。

---

6. 训练 (Training)

训练 Transformer 模型的核心目标是调整其内部参数，使其能够根据输入序列准确地生成期望的输出序列。这通常在监督学习的框架下进行，利用大量的输入-输出配对数据（例如，在机器翻译任务中是源语言句子和对应的目标语言句子）。本节将详细介绍 Transformer 模型的训练机制，包括数据准备、损失计算、优化策略以及正则化技巧。

(图示：展示了训练时 Teacher Forcing 的概念以及如何逐个词元计算损失。Decoder 的输入是目标序列的前缀，目标是预测序列中的下一个词元。)

6.1. 数据准备：批处理与掩码 (Batches and Masking)

在将数据喂给模型进行训练之前，需要进行适当的预处理，包括将句子对组织成批次 (batches) 并创建必要的掩码 (masks)。

• 批处理 (Batching):
  • 为了提高计算效率，训练数据通常以批次的形式输入模型。
  • 论文中提到：“Sentence pairs were batched together by approximate sequence length. Each training batch contained a set of sentence pairs containing approximately 25000 source tokens and 25000 target tokens.” 这是一种动态批处理策略，目标是使每个批次包含相似数量的词元，而不是固定数量的句子，有助于更均匀地利用计算资源。
• 掩码 (Masking):
  • 源序列掩码 (src_mask): 用于在 Encoder 的自注意力和 Decoder 的编码器-解码器注意力中忽略源序列中的填充 (padding) 词元。如果一个词元是填充符，那么它不应该对其他词元的表示或注意力权重产生贡献。
  • 目标序列掩码 (tgt_mask): 用于 Decoder 的掩码自注意力层。它包含两种功能：
    1. 填充掩码: 忽略目标序列中的填充词元。
    2. 后续词元掩码 (Sequence Mask / Look-ahead Mask): 防止 Decoder 在预测当前位置 i 的词元时，关注到位置 i 之后的未来词元。这是保持 Decoder 自回归特性的关键。

提供的代码中 Batch 类负责封装这些逻辑：

class Batch:"""Object for holding a batch of data with mask during training."""def __init__(self, src, tgt=None, pad=2):  # pad=2 假设是 <blank> 或 padding 符号的 IDself.src = src# src_mask: (batch_size, 1, src_len)，在填充位置为 False/0self.src_mask = (src != pad).unsqueeze(-2)if tgt is not None:# Decoder的输入是目标序列的前缀 (去掉最后一个词元)self.tgt = tgt[:, :-1]# 真实的目标是目标序列的后缀 (去掉第一个词元，通常是<BOS>)self.tgt_y = tgt[:, 1:]# 创建标准的目标序列掩码 (包含填充和后续词元掩码)self.tgt_mask = self.make_std_mask(self.tgt, pad)# 计算批次中非填充目标词元的数量，用于损失归一化self.ntokens = (self.tgt_y != pad).data.sum()@staticmethoddef make_std_mask(tgt, pad):"Create a mask to hide padding and future words."# 1. 创建填充掩码tgt_pad_mask = (tgt != pad).unsqueeze(-2) # (batch_size, 1, tgt_len)# 2. 创建后续词元掩码# subsequent_mask(size) 返回一个上三角矩阵，对角线以上为 False/0seq_mask = subsequent_mask(tgt.size(-1)).type_as(tgt_pad_mask.data)# 3. 合并两种掩码 (通过逻辑与操作)tgt_mask = tgt_pad_mask & seq_maskreturn tgt_mask

• tgt 和 tgt_y: tgt 是 Decoder 的输入 (e.g., <BOS> w1 w2)，tgt_y 是对应的真实标签 (e.g., w1 w2 <EOS>)。这种错位是序列到序列模型训练的标准做法。

6.2. 训练循环 (Training Loop)

训练过程通常在一个循环中进行，每个循环称为一个周期 (epoch)，在每个周期中模型会遍历整个训练数据集。

提供的代码中的 run_epoch 函数展示了一个典型的训练周期流程：

class TrainState: # 用于跟踪训练状态step: int = 0accum_step: int = 0samples: int = 0tokens: int = 0def run_epoch(data_iter, model, loss_compute, optimizer, scheduler,mode="train", accum_iter=1, train_state=TrainState()
):start = time.time()total_tokens = 0total_loss = 0# ... (其他初始化)for i, batch in enumerate(data_iter):# 1. 前向传播out = model.forward(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)# 2. 计算损失loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)if mode == "train" or mode == "train+log":# 3. 反向传播loss_node.backward()# ... (更新训练状态)# 4. 梯度累积与参数更新if i % accum_iter == 0: # 每 accum_iter 步进行一次参数更新optimizer.step()optimizer.zero_grad(set_to_none=True)# 5. 学习率调度scheduler.step()# ... (累积损失和词元数，打印日志)return total_loss / total_tokens, train_state

• 梯度累积 (accum_iter): 当 GPU 显存不足以容纳大的批次时，可以使用梯度累积。模型处理 accum_iter 个小批次，累积它们的梯度，然后进行一次参数更新。这等效于使用一个 accum_iter 倍大的批次。

6.3. 损失函数 (Loss Function)

• 目标: 使模型输出的概率分布与真实目标词元的分布尽可能接近。
• 交叉熵损失 (Cross-Entropy Loss): 如前所述，这是标准选择。
  Loss = -sum(y_true * log(y_pred))
• 标签平滑 (Label Smoothing): Transformer 论文中采用的一项重要正则化技术。
  • 动机: 防止模型对预测过于自信（即为正确答案分配接近1的概率，其他答案接近0），这可能导致过拟合和泛化能力下降。
  • 实现: 在计算损失时，不使用严格的独热编码 (one-hot) 作为目标分布 y_true，而是创建一个平滑后的目标分布。如果真实类别是 c，平滑后的目标概率为：
    • P(y=c) = 1.0 - ε_ls (其中 ε_ls 是平滑因子，论文中为 0.1)
    • P(y≠c) = ε_ls / (V - 1) (其中 V 是词汇表大小，剩余的概率质量 ε_ls 均匀分配给其他 V-1 个错误类别)
    • (代码实现中略有不同：smoothing / (self.size - 2)，可能是因为排除了 padding token 和一个其他特殊 token)。
  • 损失计算: 通常使用 KL 散度损失 (KLDivLoss) 来计算模型输出的对数概率与这个平滑后的目标分布之间的差异。
  • 论文指出：“This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.”

提供的 LabelSmoothing 类实现了这个功能：

class LabelSmoothing(nn.Module):def __init__(self, size, padding_idx, smoothing=0.0):super(LabelSmoothing, self).__init__()self.criterion = nn.KLDivLoss(reduction="sum") # KL散度，求和而不是平均self.padding_idx = padding_idxself.confidence = 1.0 - smoothingself.smoothing = smoothingself.size = size # 词汇表大小self.true_dist = None # 用于存储平滑后的目标分布def forward(self, x, target): # x 是模型的 log_softmax 输出, target 是真实标签assert x.size(1) == self.sizetrue_dist = x.data.clone()# 将所有位置填充为 ε_ls / (V - K_special_tokens)true_dist.fill_(self.smoothing / (self.size - 2)) # -2 可能是排除了 padding 和另一个特殊符号# 将真实目标标签位置的概率设置为 1 - ε_lstrue_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)# padding 位置的概率设置为 0true_dist[:, self.padding_idx] = 0mask = torch.nonzero(target.data == self.padding_idx)if mask.dim() > 0:true_dist.index_fill_(0, mask.squeeze(), 0.0) # 确保 padding 行全为0self.true_dist = true_dist # 存储以便可视化return self.criterion(x, true_dist.clone().detach())

SimpleLossCompute 类则将 Generator (输出对数概率) 和损失标准 (如 LabelSmoothing) 结合起来：

class SimpleLossCompute:def __init__(self, generator, criterion):self.generator = generatorself.criterion = criteriondef __call__(self, x, y, norm): # x: Decoder输出, y: 真实目标, norm: ntokensx = self.generator(x) # x 变为对数概率 [batch, seq_len, vocab_size]# 计算损失，并按非填充词元数归一化sloss = (self.criterion(x.contiguous().view(-1, x.size(-1)), # [batch*seq_len, vocab_size]y.contiguous().view(-1)              # [batch*seq_len]) / norm # norm 是 batch.ntokens)return sloss.data * norm, sloss # 返回总损失和平均损失节点

6.4. 优化器 (Optimizer)

• Adam: 论文中使用的优化器，参数设置为 β1=0.9, β2=0.98 和 ε=10^-9。
  • torch.optim.Adam(model.parameters(), lr=..., betas=(0.9, 0.98), eps=1e-9)

6.5. 学习率调度 (Learning Rate Scheduling)

Transformer 训练成功的关键之一是其独特的学习率调度策略。

• 公式:
  $$lrate=d_{\text{model}}^{-0.5}\cdot \min ({step\_num}^{-0.5},step\_num\cdot {warmup\_steps}^{-1.5})$$
• 行为: 学习率先从0线性增加到某个峰值 (在 warmup_steps 步达到)，然后按 step_num 的平方根倒数比例衰减。
• warmup_steps: 预热步数，论文中为 4000。
• 实现: 通常使用 torch.optim.lr_scheduler.LambdaLR 结合一个自定义的lambda函数来实现。

def rate(step, model_size, factor, warmup):if step == 0: step = 1 # 避免 0^(-0.5)return factor * (model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5)))
# ...
# lr_scheduler = LambdaLR(optimizer=optimizer, lr_lambda=lambda step: rate(step, *example))

6.6. 正则化 (Regularization)

除了上述的标签平滑，Transformer 还使用了：

• Dropout: 应用于每个子层（自注意力、FFN）的输出（在残差连接和层归一化之前，或在代码实现中是 Pre-LN，则在子层模块内部或之后，Add之前），以及词嵌入和位置编码相加之后。论文中基础模型使用 P_drop=0.1。
• 权重初始化: 论文中提到 “Initialize parameters with Glorot / fan_avg.” 代码中 make_model 函数包含了 nn.init.xavier_uniform_(p)。

6.7. 硬件与训练时间 (Hardware and Schedule - from paper)

• 论文中模型在 8 卡 NVIDIA P100 GPUs 上训练。
• 基础模型 (N=6, d_model=512, d_ff=2048, h=8): 每步约 0.4 秒，训练 100,000 步 (约12小时)。
• 大模型: 每步约 1.0 秒，训练 300,000 步 (约3.5天)。

7. KV 缓存 (Key-Value Cache)

在 Transformer 模型（尤其是 Decoder-Only 或 Encoder-Decoder 架构的 Decoder 部分）进行自回归 (auto-regressive) 推理生成序列时，KV 缓存是一种至关重要的优化技术，旨在显著提高生成速度并减少冗余计算。

7.1. 为什么需要 KV 缓存？

在自回归生成模式下，模型一次生成一个词元 (token)。例如，要生成句子 “A B C D”，过程如下：

1. 输入 $\text{<start>}$，模型预测 “A”。
2. 输入 $\text{<start> A}$，模型预测 “B”。
3. 输入 $\text{<start> A B}$，模型预测 “C”。
4. 输入 $\text{<start> A B C}$，模型预测 “D”。

注意到，在第3步计算 “C” 时，Transformer 的自注意力机制需要计算当前词元 “B” (作为 Query) 与之前所有词元 $\text{<start>}$, “A”, “B” (作为 Key 和 Value) 之间的注意力。其中，$\text{<start>}$ 和 “A” 的 Key ($K$) 和 Value ($V$) 向量在第2步生成 “B” 时已经计算过了。如果没有 KV 缓存，这些 $K$ 和 $V$ 向量会在每一步都被重新计算，造成巨大的计算浪费，尤其是当序列变长时。

KV 缓存的核心思想就是：对于已经处理过的词元，将其在自注意力层中计算出的 Key 和 Value 向量存储起来，在后续生成步骤中直接复用，避免重复计算。

7.2. KV 缓存的工作流程

以 Decoder 的自注意力层为例（Encoder-Decoder 架构中的 Cross-Attention 也可以应用类似思想，但 Encoder 的 $K$, $V$ 通常只需计算一次）：

1. 初始状态: 缓存为空。
2. 生成第 1 个词元 ($t=1$):
   • 输入是起始符 (e.g., $\text{<bos>}$)。
   • 模型为这个输入词元计算其 Query ($Q_1$)，Key ($K_1$)，和 Value ($V_1$) 向量。
   • 注意力计算：$\text{Attention}(Q_1,K_1,V_1)$。
   • 缓存: 将 $K_1$ 和 $V_1$ 存储到 KV 缓存中。
3. 生成第 2 个词元 ($t=2$):
   • 输入是第 1 个生成的词元。
   • 模型为这个新词元计算其 Query ($Q_2$)，Key ($K_2$)，和 Value ($V_2$) 向量。
   • 从缓存中检索: 取出之前存储的 $K_1,V_1$。
   • 注意力计算：$Q_2$ 需要与所有历史词元的 Key 进行交互。此时，有效的 Keys 是 $[K_1\text{ (cached)},K_2\text{ (new)}]$，有效的 Values 是 $[V_1\text{ (cached)},V_2\text{ (new)}]$。
     • $\text{AttnScores}=Q_2\cdot [K_1,K_2{]}^T$
     • $\text{Output}=\text{softmax}(\text{AttnScores})\cdot [V_1,V_2]$
   • 缓存: 将新计算的 $K_2$ 和 $V_2$ 追加到 KV 缓存中。缓存现在包含 $[K_1,K_2]$ 和 $[V_1,V_2]$。
4. 生成第 $i$-个词元 ($t=i$):
   • 输入是第 $i-1$ 个生成的词元。
   • 模型为这个新词元计算其 Query ($Q_i$)，Key ($K_i$)，和 Value ($V_i$) 向量。
   • 从缓存中检索: 取出之前存储的 $[K_1,...,K_{i-1}]$ 和 $[V_1,...,V_{i-1}]$。
   • 注意力计算：$Q_i$ 与 $[K_1,...,K_{i-1},K_i]$ 和 $[V_1,...,V_{i-1},V_i]$ 进行交互。
     • ${\text{AttnScores}}_i=Q_i\cdot [K_1,K_2,...,K_i{]}^T$
     • ${\text{Output}}_i=\text{softmax}({\text{AttnScores}}_i)\cdot [V_1,V_2,...,V_i]$
   • 缓存: 将新计算的 $K_i$ 和 $V_i$ 追加到 KV 缓存中。

这个过程在每个 Decoder 层的自注意力模块中都会发生。

7.3. KV 缓存出现的位置

KV 缓存主要应用于 Transformer 推理阶段的以下注意力模块：

1. Decoder 的 Masked Multi-Head Self-Attention 层:

   • 这是 KV 缓存最典型的应用场景。
   • 在自回归生成的每一步，当前生成词元的 $Q$ 会与**所有先前已生成词元（其 $K$, $V$ 已被缓存）以及当前词元自身（其 $K$, $V$ 刚被计算）**进行注意力计算。
   • 缓存的是已处理过的目标序列词元的 $K$ 和 $V$ 向量。

2. Decoder 的 Multi-Head Encoder-Decoder Attention (Cross-Attention) 层 (对于 Encoder-Decoder 模型如原始 Transformer, BART, T5):

   • 在这种注意力中，Query ($Q$) 来自 Decoder 的前一子层输出。
   • Key ($K$) 和 Value ($V$) 来自 Encoder 的最终输出。
   • 由于 Encoder 的输出在整个 Decoder 生成过程中是固定不变的，因此 Encoder 输出的 $K$ 和 $V$ 向量只需要在 Decoder 开始生成第一个词元之前计算一次，然后就可以被视为一个静态的 KV 缓存，供 Decoder 的每一层、每一步生成时重复使用。
   • 这种情况下，“缓存” 更像是 “预计算和复用”，因为这些 $K$, $V$ 值不会在 Decoder 生成过程中动态增长。

Encoder-Only 模型 (如 BERT) 在进行非自回归任务 (如分类、NER) 时，通常一次性处理整个输入序列，不需要逐步生成，因此不直接使用 KV 缓存这种动态优化。Decoder-Only 模型 (如 GPT 系列) 则严重依赖 KV 缓存进行高效的文本生成。

7.4. KV 缓存的优势与计算复杂度分析

KV 缓存带来的最显著优势是计算效率的提升。下面我们详细分析其计算复杂度：

7.4.1 不使用 KV 缓存的计算复杂度

假设我们要生成长度为 $L$ 的序列，每个词元的隐藏维度为 $d$。在传统的自回归生成中：

• 在第 $t$ 步（$1\le t\le L$）：

  • 需要计算 $t$ 个词元的 $Q$、$K$、$V$ 向量，复杂度为 $O(t\cdot d^2)$
  • 执行注意力计算：$Q\cdot K^T$，复杂度为 $O(t^2\cdot d)$
  • 应用 softmax 和与 $V$ 相乘，复杂度为 $O(t^2\cdot d)$
  • 因此第 $t$ 步的总复杂度为 $O(t\cdot d^2+t^2\cdot d)$

• 整个序列的总复杂度为：
  $$\sum _{t=1}^{L}O(t\cdot d^2+t^2\cdot d)=O(L^2\cdot d^2+L^3\cdot d)$$

• 对于较长序列，$L^3$ 项成为主导因素，总复杂度近似为 $O(L^3\cdot d)$

7.4.2 使用 KV 缓存的计算复杂度

• 在第 $t$ 步（$1\le t\le L$）：

  • 只需为当前新词元计算 $Q$、$K$、$V$ 向量，复杂度为 $O(d^2)$
  • $K$ 和 $V$ 从缓存中获取，无需重复计算
  • 执行注意力计算：$Q$ 与 $t$ 个缓存的 $K$ 向量相乘，复杂度为 $O(t\cdot d)$
  • 应用 softmax 和与 $V$ 相乘，复杂度为 $O(t\cdot d)$
  • 因此第 $t$ 步的总复杂度为 $O(d^2+t\cdot d)$

• 整个序列的总复杂度为：
  $$\sum _{t=1}^{L}O(d^2+t\cdot d)=O(L\cdot d^2+L^2\cdot d)$$

• 对于较长序列，这比不使用缓存的 $O(L^3\cdot d)$ 要高效得多

7.4.3 KV 缓存的其他优势

• 显著减少计算量: 如上所述，对于长度为 $L$ 的序列，KV 缓存将总复杂度从 $O(L^3\cdot d)$ 降低到 $O(L^2\cdot d+L\cdot d^2)$，这使得长序列生成成为可能。
• 大幅提升推理速度: 计算量的减少直接转化为推理时间的缩短。
• 内存权衡: 需要额外的内存空间来存储这些 $K$ 和 $V$ 向量。KV 缓存的空间复杂度约为 $O(L\cdot d\cdot h\cdot b\cdot n)$，其中：
• $L$ 是序列长度
• $d$ 是每个头的维度
• $h$ 是头的数量
• $b$ 是批次大小
• $n$ 是模型层数

对于非常长的序列，KV 缓存本身也会变得很大，催生了如 Multi-Query Attention (MQA) 和 Grouped-Query Attention (GQA) 等技术来减小 KV 缓存的体积。

7.5. 实现细节

• 数据结构: KV 缓存通常为每个 Decoder 层维护一组张量。这些张量的形状通常是 $[b,h,s,d_h]$，其中：
  • $b$ (batch_size): 处理的序列数量。
  • $h$ (num_heads): 多头注意力的头数。
  • $s$ (current_sequence_length): 当前已处理/生成的词元数量，这个维度会随着生成的进行而动态增长。
  • $d_h$ (head_dim): 每个注意力头的维度 ($\frac{d_{model}}{h}$)。
  • 通常会为 Key 和 Value 分别维护这样的张量，或者将它们堆叠在一个额外的维度上，例如 $[2,b,h,s,d_h]$，其中 2 代表 $K$ 和 $V$。或者在实现层面，每个 Transformer 层对象会持有 past_key_values 这样的属性。
• 动态扩展与预分配:
  • 在实践中，为了效率，可能会预先分配一个最大序列长度 $L_{max}$ 的缓存空间，然后用一个指针或索引跟踪当前实际填充到的长度。
  • 当 $s$ 增长时，新的 $K$, $V$ 向量会被拼接到已缓存张量的相应维度上，可表示为：
    $$K_{cached}^{new}=\text{concat}(K_{cached}^{old},K_{new}),\text{dim}=2$$
    $$V_{cached}^{new}=\text{concat}(V_{cached}^{old},V_{new}),\text{dim}=2$$
• 跨层传递: 在多层 Decoder 中，每一层的 KV 缓存是独立的。当模型生成下一个词元时，会将上一时间步所有层的 KV 缓存状态传递给当前时间步对应的层。

7.6. KV 缓存与掩码 (Masking) 的关系

• KV 缓存与 Sequence Mask (Look-ahead Mask) 紧密相关且协同工作。
• 在训练时，Sequence Mask 确保模型在预测第 $i$ 个词元时，只能关注到 $0$ 到 $i$ 位置的词元。这是通过添加一个掩码矩阵 $M$ 实现的：

$M_{ij}=\{\begin{array}{ll}0& \text{if }i\ge j\\ -\infty & \text{if }i<j\end{array}$

• 然后在计算注意力分数时应用该掩码：

$\text{Attention\_Scores}=Q\cdot K^T+M$

$\text{Attention\_Weights}=\text{softmax}(\text{Attention\_Scores})$

• 在带 KV 缓存的推理时，这种 “只能看过去” 的特性是通过缓存的构建方式自然实现的。因为在生成第 $i$ 个词元时，KV 缓存中只包含了 $0$ 到 $i-1$ 位置词元的 $K$ 和 $V$。当前词元 $i$ 的 $Q$ 会与这些缓存的 $K$, $V$ 以及自身新计算的 $K$, $V$ 进行交互。模型物理上无法访问未来词元的 $K$, $V$，因为它们尚未被计算和缓存。

• 因此，在推理时，显式的上三角掩码矩阵可能不再需要以同样的方式构建（或其作用由缓存机制隐式完成），注意力计算简化为：

${\text{Attention\_Scores}}_i=Q_i\cdot [K_1,K_2,...,K_i{]}^T$

${\text{Attention\_Weights}}_i=\text{softmax}({\text{Attention\_Scores}}_i)$

${\text{Output}}_i={\text{Attention\_Weights}}_i\cdot [V_1,V_2,...,V_i]$

8. 参考

• 原始论文: “Attention Is All You Need” - Ashish Vaswani, et al.
• 博客/教程:
  • Jay Alammar 的 “The Illustrated Transformer” (图解Transformer，非常经典)
  • Harvard NLP 的 “The Annotated Transformer” (带代码实现的详细解读)
  • 李沐老师的 Transformer 视频和文章
  • transformer-explainer
  • 交叉注意力说明
  • 动画展示3b1b
  • 三种注意力机制

9. 附录完整代码流程

9.1 核心注意力机制

基础注意力函数

def attention(query, key, value, mask=None, dropout=None):d_k = query.size(-1)scores = torch.matmul(query, key.transpose(-2,-1)) / math.sqrt(d_k)if mask is not None:scores = scores.masked_fill(mask==0, -1e9)p_attn = scores.softmax(dim=-1)if dropout is not None:p_attn = dropout(p_attn)return torch.matmul(p_attn, value), p_attn

这个函数实现了缩放点积注意力，计算步骤：

1. 计算query和key的点积
2. 除以缩放因子(√d_k)
3. 如果有mask则应用
4. 对结果应用softmax
5. 用attention权重对value加权求和

多头注意力

class MultiHeadedAttention(nn.Module):def __init__(self, h, d_model, dropout=0.1):super(MultiHeadedAttention, self).__init__()assert d_model % h == 0self.d_k = d_model // hself.h = hself.linears = clones(nn.Linear(d_model, d_model), 4)self.attn = Noneself.dropout = nn.Dropout(p=dropout)def forward(self, query, key, value, mask=None):if mask is not None:mask = mask.unsqueeze(1)nbatches = query.size(0)# 1) 线性映射并分割多头query, key, value = [lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)for lin, x in zip(self.linears, (query, key, value))]# 2) 应用注意力机制x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)# 3) 合并多头结果x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)# 4) 最后的线性层return self.linears[-1](x)

多头注意力将输入分成多个子空间并行处理：

1. 将输入线性变换并分割成h个头
2. 对每个头计算注意力
3. 合并多头结果
4. 通过最后一个线性层输出

9.2 前馈网络

class PositionwiseFeedForward(nn.Module):def __init__(self, d_model, d_ff, dropout=0.1):super(PositionwiseFeedForward, self).__init__()self.w_1 = nn.Linear(d_model, d_ff)self.w_2 = nn.Linear(d_ff, d_model)self.dropout = nn.Dropout(dropout)def forward(self, x):return self.w_2(self.dropout(self.w_1(x).relu()))

这是每个Transformer层中的前馈神经网络，由两个线性变换组成，中间有一个ReLU激活函数。

9.3 模型构建辅助函数

def clones(module, N):return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

这个函数复制模块N次，用于创建多层结构。

9.4 层归一化

class LayerNorm(nn.Module):def __init__(self, features, eps=1e-6):super(LayerNorm, self).__init__()self.a_2 = nn.Parameter(torch.ones(features))self.b_2 = nn.Parameter(torch.zeros(features))self.eps = epsdef forward(self, x):mean = x.mean(-1, keepdim=True)std = x.std(-1, keepdim=True)return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

层归一化对每一层的输出进行标准化，加速网络收敛。

9.5 残差连接

class SublayerConnection(nn.Module):def __init__(self, size, dropout):super(SublayerConnection, self).__init__()self.norm = LayerNorm(size)self.dropout = nn.Dropout(dropout)def forward(self, x, sublayer):return x + self.dropout(sublayer(self.norm(x)))

实现了残差连接，注意这里是先归一化再应用子层，最后加上输入。

9.6 编码器架构

编码器层

class EncoderLayer(nn.Module):def __init__(self, size, self_attn, feed_forward, dropout):super(EncoderLayer, self).__init__()self.self_attn = self_attnself.feed_forward = feed_forwardself.sublayer = clones(SublayerConnection(size, dropout), 2)self.size = sizedef forward(self, x, mask):x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))return self.sublayer[1](x, self.feed_forward)

每个编码器层包含一个自注意力层和一个前馈网络，通过残差连接和层归一化相互连接。

完整编码器

class Encoder(nn.Module):def __init__(self, layer, N):super(Encoder, self).__init__()self.layers = clones(layer, N)self.norm = LayerNorm(layer.size)def forward(self, x, mask):for layer in self.layers:x = layer(x, mask)return self.norm(x)

编码器由N个相同的层堆叠而成，最后有一个额外的层归一化。

9.7 解码器架构

解码器层

class DecoderLayer(nn.Module):def __init__(self, size, self_attn, src_attn, feed_forward, dropout):super(DecoderLayer, self).__init__()self.size = sizeself.self_attn = self_attnself.src_attn = src_attnself.feed_forward = feed_forwardself.sublayer = clones(SublayerConnection(size, dropout), 3)def forward(self, x, memory, src_mask, tgt_mask):m = memoryx = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))return self.sublayer[2](x, self.feed_forward)

解码器层比编码器层多一个注意力层，包含：

1. 掩码自注意力层（只关注已生成的输出）
2. 编码器-解码器注意力层（关注输入序列）
3. 前馈网络

完整解码器

class Decoder(nn.Module):def __init__(self, layer, N):super(Decoder, self).__init__()self.layers = clones(layer, N)self.norm = LayerNorm(layer.size)def forward(self, x, memory, src_mask, tgt_mask):for layer in self.layers:x = layer(x, memory, src_mask, tgt_mask)return self.norm(x)

解码器也是由N个相同的层堆叠而成，最后有一个额外的层归一化。

9.8 位置编码

class PositionalEncoding(nn.Module):def __init__(self, d_model, dropout, max_len=5000):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)# 预计算位置编码pe = torch.zeros(max_len, d_model)position = torch.arange(0, max_len).unsqueeze(1)div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))pe[:, 0::2] = torch.sin(position * div_term)pe[:, 1::2] = torch.cos(position * div_term)pe = pe.unsqueeze(0)self.register_buffer("pe", pe)def forward(self, x):x = x + self.pe[:, : x.size(1)].requires_grad_(False)return self.dropout(x)

位置编码为模型注入位置信息，使用正弦和余弦函数生成不同频率的波形。

9.9 序列掩码生成

def subsequent_mask(size):attn_shape = (1, size, size)subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)return subsequent_mask == 0

这个函数生成一个掩码，防止解码器看到未来的位置。

9.10 输出生成器

class Generator(nn.Module):def __init__(self, d_model, vocab):super(Generator, self).__init__()self.proj = nn.Linear(d_model, vocab)def forward(self, x):return log_softmax(self.proj(x), dim=-1)

将解码器的输出映射到词汇表大小的向量，并应用log_softmax。

9.11 完整Transformer模型

class EncoderDecoder(nn.Module):def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):super(EncoderDecoder, self).__init__()self.encoder = encoderself.decoder = decoderself.src_embed = src_embedself.tgt_embed = tgt_embedself.generator = generatordef forward(self, src, tgt, src_mask, tgt_mask):return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)def encode(self, src, src_mask):return self.encoder(self.src_embed(src), src_mask)def decode(self, memory, src_mask, tgt, tgt_mask):return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

9.12 模型构建函数

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):c = copy.deepcopyattn = MultiHeadedAttention(h, d_model)ff = PositionwiseFeedForward(d_model, d_ff, dropout)position = PositionalEncoding(d_model, dropout)model = EncoderDecoder(Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),nn.Sequential(Embeddings(d_model, src_vocab), c(position)),nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),Generator(d_model, tgt_vocab),)# 参数初始化for p in model.parameters():if p.dim() > 1:nn.init.xavier_uniform_(p)return model

这个函数将所有组件组装成一个完整的Transformer模型，默认参数设置和原始论文一致。

关键PyTorch技术细节

masked_fill 方法

用于在注意力计算中屏蔽某些位置，实现：

scores = scores.masked_fill(mask==0, -1e9)

masked_fill 对掩码值为True的位置用指定值（这里是-1e9）替换，在softmax后这些位置的权重接近0。

维度处理：

多头注意力中的核心维度变换：

lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)

1. 线性投影：(batch, seq_len, d_model)
2. 重塑：(batch, seq_len, h, d_k)
3. 转置：(batch, h, seq_len, d_k)

这让每个注意力头能独立处理序列，充分利用并行计算。