Transformer 模型结构详解及代码实现!

一、Transformer简要发展史

以下是Transformer模型发展历史中的关键节点：

Transformer架构于2017年6月推出。原本研究的重点是翻译任务。随后推出了几个有影响力的模型，包括：

这个列表并不全面，只是为了突出一些不同类型的 Transformer 模型。大体上，它们可以分为三类：

Transformer 默认都是大模型，除了一些特例（如 DistilBERT）外，实现更好性能的一般策略是增加模型的大小以及预训练的数据量。其中，GPT-2 是使用「transformer 解码器模块」构建的，而 BERT 则是通过「transformer 编码器」模块构建的。

二、Transformer 整体架构

论文中给出用于中英文翻译任务的 Transformer 整体架构如下图所示：

可以看出Transformer架构由Encoder和Decoder两个部分组成：其中Encoder和Decoder都是由N=6个相同的层堆叠而成。Multi-Head Attention 结构是 Transformer 架构的核心结构，其由多个 Self-Attention 组成的。其中，

Transformer 架构更详细的可视化图如下所示:

1. 输入模块

(1) Tokenizer预处理

在基于Transformer的大模型LLM中，输入通常为字符串文本。由于模型无法直接处理自然语言，因此需要借助Tokenizer对输入进行预处理。具体流程如下：

• 分词（Tokenization）：将输入文本按规则切分为一个个词元（token），如单词、子词或特殊符号。
• 词表映射（Vocabulary Mapping）：每个 token 被映射到一个唯一的整数 ID，该 ID 来自预训练模型所使用的词汇表。
• 生成 input_ids 向量（矩阵）：最终输出是一个由 token ID 构成的向量（或矩阵），作为模型输入。

以下是以 Hugging Face 的 transformers 库为例，展示如何使用 BertTokenizer 和 BertModel 完成输入文本的预处理和编码：

from transformers import BertTokenizer, BertModel
import torch

# 1. 加载预训练的 BERT tokenizer 和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 2. 输入文本
text = "A Titan RTX has 24GB of VRAM"

# 3. 分词并映射为 token ID 序列
inputs = tokenizer(text, return_tensors="pt", truncatinotallow=True, padding=True)

# 输出 token IDs
print("Token IDs:", inputs['input_ids'])

# 4. 传入模型，获取输出
outputs = model(**inputs)

# 5. 获取最后一层的隐藏状态表示
last_hidden_states = outputs.last_hidden_state
print("Last hidden states shape:", last_hidden_states.shape)

原始输入文本 "A Titan RTX has 24GB of VRAM" 通过 tokenizer 完成分词和词表映射工作，生成的输入 ID 列表：

[101, 138, 28318, 56898, 12674, 10393, 10233, 32469, 10108, 74727, 36535, 102]

其中，

• 101 表示 [CLS] 标记；
• 102 表示 [SEP] 标记；
• 其余为对应 token 在词表中的索引。

在所有基于 Transformer 的 LLM 中，唯一必须的输入是 input_ids，它是由 Tokenizer 映射后的 token 索引组成的整数向量，代表了输入文本在词表中的位置信息。

(2) Embedding 层

在基于 Transformer 的大型语言模型（LLM）中，嵌入层（Embedding Layer）是将输入 token ID 映射为向量表示的核心组件。其作用是将离散的整数索引转换为连续、稠密的向量空间表示，从而便于后续神经网络进行语义建模。

✅ 万物皆可 Embedding：虽然最常见的是词嵌入（Word Embedding），但图像、语音等也可以通过嵌入层映射为向量形式，实现统一建模。

例如，mnist 数据集中的图片，可以通过嵌入层来表示，如下图所示，每个点代表一个图片(10000*784)，通过嵌入层，将图片的像素点转化为稠密的向量，然后通过 t-SNE/pca 降维，可以看到图片的空间分布。

LLM 中，单词 token 需要经过 Embedding 层，Embedding 层的作用是将输入的离散化表示（例如 token ids）转换为连续的低维向量表示，其由单词 Embedding 和位置 Embedding （Positional Encoding）相加得到，通常定义为 TransformerEmbedding 层。

① 单词嵌入（Token Embedding）」

自然语言处理中，输入文本通常是以符号形式存在的词汇，而这些离散符号无法直接被神经网络处理。因此需要一个可学习的嵌入矩阵将每个 token 转换为固定维度的向量。

工作原理：

a. 输入是一个形状为 [batch_size, seq_len] 的整数张量，表示每个 token 在词表中的索引；

b. 输出是一个形状为 [batch_size, seq_len, d_model] 的三维张量，其中：

• d_model 是嵌入维度（如 512 或 768）；
• 每个 token 对应一个 d_model 维的向量；

c. 嵌入层权重矩阵大小为 [vocab_size, d_model]，参数量为：

在 PyTorch 中，词嵌入层通常使用 torch.nn.Embedding 模块实现，其作用是将 token 的索引转换为低维语义向量表示。

✅ 输入与输出说明

• embedding_dim 是嵌入向量的维度，也称为词向量维度；
• 它通常被设置为 d_model 或 h，即后续 Transformer 层使用的隐藏层维度（如 512 或 768）.

📐 示例代码：构建 Token Embedding 层

from transformers import BertTokenizer
import torch.nn as nn

## 1, 使用 BERT tokenizer 将批量输入的字符串文本序列转化为 input_ids
tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased") 
batch_text = ["A Titan RTX has 24GB of VRAM", "I have a dog and cat"]
inputs = tokenizer(batch_text, return_tensors="pt", truncation=True, padding=True)
input_ids = inputs["input_ids"]

# 2. 创建一个 nn.Embedding 层
vocab_size = tokenizer.vocab_size  # 词表大小取决于你加载的具体 tokenizer 模型
embedding_dim = 512  # 嵌入向量的维度，参考 transformer 论文的大小
embedding_layer = nn.Embedding(vocab_size, embedding_dim)

# 3. 通过 nn.Embedding 层，将输入的 IDs 映射到嵌入向量
embedded_output = embedding_layer(input_ids)

# 4. 输出嵌入向量的形状
print("嵌入向量的形状：", embedded_output.shape)  # (batch_size, sequence_length, embedding_dim), torch.Size([2, 12, 512])

# 5. 打印嵌入向量
print(embedded_output)

程序运行后输出结果如下所示:

② 位置嵌入（Positional Encoding）」

由于 Transformer 不依赖于 RNN 的顺序性建模方式，它必须显式地引入位置信息，以保留 token 在序列中的位置特征。

为此，Transformer 使用了 Sinusoidal Positional Encoding（正弦/余弦位置编码）：

其中：

• pos：token 在序列中的位置；
• i：维度索引；
• d_model：嵌入维度。

③ TransformerEmbedding 层集成

transformer 输入模块有三个组成部分：文本/提示词、分词器（Tokenizer）和嵌入层（Embeddings）。输入模块的工作流程和代码实现如下所示:

矩阵的每一列表示一个 token 的嵌入向量。

class PositionalEncoding(nn.Module):
    """
    compute sinusoid encoding.
    """
    def __init__(self, d_model, max_len, device):
        """
        constructor of sinusoid encoding class

        :param d_model: dimension of model
        :param max_len: max sequence length
        :param device: hardware device setting
        """
        super(PositionalEncoding, self).__init__()

        # same size with input matrix (for adding with input matrix)
        self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False  # we don't need to compute gradient

        pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        # 1D => 2D unsqueeze to represent word's position

        _2i = torch.arange(0, d_model, step=2, device=device).float()
        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])
        # "step=2" means 'i' multiplied with two (same with 2 * i)

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        # compute positional encoding to consider positional information of words

    def forward(self, x):
        # self.encoding
        # [max_len = 512, d_model = 512]

        batch_size, seq_len = x.size()
        # [batch_size = 128, seq_len = 30]

        return self.encoding[:seq_len, :]
        # [seq_len = 30, d_model = 512]
        # it will add with tok_emb : [128, 30, 512]         

class TokenEmbedding(nn.Embedding):
    """
    Token Embedding using torch.nn
    they will dense representation of word using weighted matrix
    """

    def __init__(self, vocab_size, d_model):
        """
        class for token embedding that included positional information
        :param vocab_size: size of vocabulary
        :param d_model: dimensions of model
        """
        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)

class TransformerEmbedding(nn.Module):
    """
    token embedding + positional encoding (sinusoid)
    positional encoding can give positional information to network
    """

    def __init__(self, vocab_size, max_len, d_model, drop_prob, device):
        """
        class for word embedding that included positional information
        :param vocab_size: size of vocabulary
        :param d_model: dimensions of model
        """
        super(TransformerEmbedding, self).__init__()
        self.tok_emb = TokenEmbedding(vocab_size, d_model)
        # self.position_embedding = nn.Embedding(max_len, embed_size)
        self.pos_emb = PositionalEncoding(d_model, max_len, device)
        self.drop_out = nn.Dropout(p=drop_prob)

    def forward(self, x):
        tok_emb = self.tok_emb(x)
        pos_emb = self.pos_emb(x)
        return self.drop_out(tok_emb + pos_emb)

2. Multi-Head Attention 结构

Encoder 和 Decoder 结构中公共的 layer 之一是 Multi-Head Attention，其是由多个 Self-Attention 并行组成的。Encoder block 只包含一个 Multi-Head Attention，而 Decoder block 包含两个 Multi-Head Attention (其中有一个用到 Masked)。

(1) Self-Attention 结构

Self-Attention 中文翻译为自注意力机制，论文中叫作 Scale Dot Product Attention，它是 Transformer 架构的核心，使得每个 token 能够关注整个序列中的其他 token，从而建立全局依赖关系。其结构如下图所示：

✅ 在本文中，Self-Attention 层与论文中的 ScaleDotProductAttention 层意义一致，实现方式完全相同。

🧮 数学定义

Self-Attention 的计算过程可以表示为：

其中：

•  ：Query 向量；
• ：Key 向量；
• ：Value 向量；
• ：Query 和 Key 的维度；
• Softmax 对注意力分数按最后一个维度归一化；
• ：用于缩放点积，防止 softmax 梯度消失；

输入来源：

• 输入词向量经过 Embedding 层后，进入位置编码层；
• 再通过线性变换（Linear 层），分别生成 Query、Key 和 Value 向量；
• 这三个向量的形状通常为 [batch_size, seq_len, d_k] 或 [seq_len, d_k]。

计算步骤如下：

① 计算注意力分数矩阵

• 其中 是 Key 张量的转置；
• 点积结果是一个 [seq_len, seq_len] 的注意力得分矩阵；
• 使用 softmax 归一化，得到注意力权重。

② 应用掩码（可选）

• 在 Decoder 中使用 Masked Self-Attention，防止未来信息泄露；
• 若传入 mask，将对应位置设为极小值（如 -1e9）以抑制其影响。

③ 加权聚合 Value 向量

• 将 softmax 后的注意力权重与 Value 相乘，得到上下文感知的输出张量；
• 输出维度保持与输入一致：[batch_size, seq_len, d_v]。

🧱 代码实现：

import torch
import math
import torch.nn as nn

class ScaleDotProductAttention(nn.Module):
    def __init__(self):
        """
        初始化 Self-Attention 层，仅包含一个 softmax 操作。
        """
        super(ScaleDotProductAttention, self).__init__()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor = None):
        """
        Self-Attention 前向传播函数

        :param Q: Query 向量，形状为 [batch_size, seq_len, d_k]
        :param K: Key 向量，形状为 [batch_size, seq_len, d_k]
        :param V: Value 向量，形状为 [batch_size, seq_len, d_v]
        :param mask: 掩码张量，形状为 [batch_size, seq_len, seq_len]
        :return: 
            output: 加权后的 Value 向量，形状为 [batch_size, seq_len, d_v]
            attn_weights: 注意力权重矩阵，形状为 [batch_size, seq_len, seq_len]
        """
        # 1. 计算 QK^T 得到注意力分数
        K_T = K.transpose(-1, -2)  # [batch_size, d_k, seq_len]
        scores = torch.matmul(Q, K_T) / math.sqrt(Q.size(-1))  # [batch_size, seq_len, seq_len]

        # 2. 如果有 mask，应用掩码（例如 Decoder 中防止看到未来词）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 3. 应用 softmax 得到注意力权重
        attn_weights = self.softmax(scores)  # [batch_size, seq_len, seq_len]

        # 4. 权重 × Value 得到最终输出
        output = torch.matmul(attn_weights, V)  # [batch_size, seq_len, d_v]

        return output, attn_weights

📐 示例调用与输出解析：

# 创建 Q、K、V 张量
Q = torch.randn(5, 10, 64)  # [batch_size=5, seq_len=10, d_k=64]
K = torch.randn(5, 10, 64)
V = torch.randn(5, 10, 64)

# 创建 Self-Attention 层
attention = ScaleDotProductAttention()

# 前向传播
output, attn_weights = attention(Q, K, V)

# 打印输出形状
print(f"ScaleDotProductAttention output shape: {output.shape}")      # [5, 10, 64]
print(f"attn_weights shape: {attn_weights.shape}")                # [5, 10, 10]

(3) Multi-Head Attention

Multi-Head Attention(MHA)是在Self-Attention基础上引入的一种增强机制。其核心理念是：将输入向量空间划分为多个子空间，在每个子空间中独立计算Self-Attention，最后将多个子空间的输出拼接在一起并进行线性变换，从而得到最终的输出。

对于 MHA，之所以需要对 Q、K、V 进行多头（head）划分，其目的是为了增强模型对不同信息的关注。具体来说，多组 Q、K、V 分别计算 Self-Attention，每个头自然就会有独立的 Q、K、V 参数，从而让模型同时关注多个不同的信息，这有些类似 CNN 架构模型的多通道机制。

下图是论文中 Multi-Head Attention 的结构图。

从图中可以看出， MHA 结构的计算过程可总结为下述步骤:

• 将输入 Q、K、V 张量进行线性变换（Linear 层），输出张量尺寸为 [batch_size, seq_len, d_model]；
• 将前面步骤输出的张量，按照头的数量（n_head）拆分为 n_head 子张量，其尺寸为 [batch_size, n_head, seq_len, d_model//n_head]；
• 每个子张量并行计算注意力分数，即执行 dot-product attention 层，输出张量尺寸为 [batch_size, n_head, seq_len, d_model//n_head]；
• 将这些子张量进行拼接 concat ，并经过线性变换得到最终的输出张量，尺寸为 [batch_size, seq_len, d_model]。

📐 数学表达式

其中：

• ：第 i 个 head 的可学习参数；
• ：最终输出的线性变换矩阵；
• Concat表示将各个 head 的输出拼接在一起。

(4) Multi-Head Attention 实现

import torch
import math
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    """Multi-Head Attention Layer"""
    
    def __init__(self, d_model, n_head):
        """
        Args:
            d_model: 模型嵌入维度（通常为 512 或 768）；
            n_head: 注意力头的数量（如 8）；
        """
        super(MultiHeadAttention, self).__init__()
        
        # 初始化参数
        self.n_head = n_head
        self.attention = ScaleDotProductAttention()  # 使用前面定义的 Self-Attention
        
        # 线性变换层
        self.w_q = nn.Linear(d_model, d_model)       # Query 变换
        self.w_k = nn.Linear(d_model, d_model)       # Key 变换
        self.w_v = nn.Linear(d_model, d_model)       # Value 变换
        self.fc = nn.Linear(d_model, d_model)         # 输出投影层

    def forward(self, q, k, v, mask=None):
        """
        Args:
            q: Query 张量，[batch_size, seq_len, d_model]
            k: Key 张量，[batch_size, seq_len, d_model]
            v: Value 张量，[batch_size, seq_len, d_model]
            mask: 掩码张量，[batch_size, seq_len, seq_len]
        """
        # Step 1: 线性变换
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)

        # Step 2: 拆分到多个 head
        q = self.split(q)   # [batch_size, n_head, seq_len, d_tensor]
        k = self.split(k)   # [batch_size, n_head, seq_len, d_tensor]
        v = self.split(v)   # [batch_size, n_head, seq_len, d_tensor]

        # Step 3: 计算每个 head 的 attention
        sa_output, attn_weights = self.attention(q, k, v, mask)

        # Step 4: 拼接所有 head 的输出
        mha_output = self.concat(sa_output)  # [batch_size, seq_len, d_model]

        # Step 5: 最终线性变换
        mha_output = self.fc(mha_output)

        return mha_output, attn_weights

    def split(self, tensor):
        """
        拆分输入张量为多个 head

        Args:
            tensor: [batch_size, seq_len, d_model]
        Returns:
            [batch_size, n_head, seq_len, d_tensor]
        """
        batch_size, seq_len, d_model = tensor.size()
        d_tensor = d_model // self.n_head  # 每个 head 的维度
        
        # reshape + transpose 实现拆分
        tensor = tensor.view(batch_size, seq_len, self.n_head, d_tensor)
        tensor = tensor.transpose(1, 2)  # [batch_size, n_head, seq_len, d_tensor]
        
        return tensor

    def concat(self, sa_output):
        """
        拼接多个 head 的输出

        Args:
            sa_output: [batch_size, n_head, seq_len, d_tensor]
        Returns:
            [batch_size, seq_len, d_model]
        """
        batch_size, n_head, seq_len, d_tensor = sa_output.size()
        d_model = n_head * d_tensor
        
        # transpose + reshape 实现合并
        sa_output = sa_output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        
        return sa_output

📊 示例调用与输出解析：

# 定义参数
d_model = 512
n_head = 8
seq_len = 10
batch_size = 32

# 创建 Q、K、V 张量
Q = torch.randn(batch_size, seq_len, d_model)
K = torch.randn(batch_size, seq_len, d_model)
V = torch.randn(batch_size, seq_len, d_model)

# 构建 MHA 层
mha_layer = MultiHeadAttention(d_model=d_model, n_head=n_head)

# 前向传播
output, weights = mha_layer(Q, K, V)

# 打印输出形状
print("MHA Output Shape:", output.shape)      # [32, 10, 512]
print("Attn Weights Shape:", weights.shape)   # [32, 8, 10, 10]

3. Encoder结构

Transformer 中的 Encoder 是整个模型中用于编码输入序列的部分。它由 N=6 个相同的 encoder block 堆叠而成。每个 encoder block 主要包含两个子层（sub-layers）：多头自注意力机制（Multi-Head Self-Attention）和位置全连接前馈网络（Position-wise Feed Forward Network）。

这两个子层之间都使用了 残差连接（Residual Connection） 和 层归一化（Layer Normalization），以增强训练稳定性和模型表达能力。

下图中红色框选部分表示一个标准的 Encoder Block，其内部结构如下：

由以下四个关键部分构成。

(1) 每一层的计算流程（以单个 encoder block 为例）

① 多头自注意力机制（Multi-Head Self-Attention）

• 输入：嵌入后的张量 
• 输出：通过自注意力加权后的新张量 sa_output

sa_output, attn_weights = MultiHeadAttention(q=x, k=x, v=x, mask=src_mask)

其中，

• Query、Key 和 Value 来自同一个输入 ；
• 可选 mask 通常用于屏蔽 padding token 或控制位置感知范围。

② 残差连接 + 层归一化（Sublayer 1）

x = x + dropout(sa_output)
x = layer_norm(x)

• 应用残差映射，缓解梯度消失问题；
• 使用 LayerNorm 对每个 token 的向量进行标准化处理；
• 整体目标：提升模型表达能力与训练稳定性。

③ 位置全连接前馈网络（Position-wise FeedForward）

定义为： 

nn.Sequential(
    nn.Linear(d_model, d_ff),
    nn.ReLU(),
    nn.Linear(d_ff, d_model),
    nn.Dropout(drop_prob)
)

• d_model：模型隐层维度（如 512）；
• d_ff：FeedForward 网络中间维度（如 2048）；
• ReLU 导致非线性更强的语义表达。

④ 再次残差连接 + 层归一化（Sublayer 2）

x = x + dropout(ffn_output)
x = layer_norm(x)

• 保证模型在经过复杂变换后仍能保留原始信息；
• 达成对上下文感知表示的稳定学习。

(2) 维度变化说明（输入输出保持一致）

无论经过多少层 Encoder block，每个 block 的输入与输出形状始终一致：

(3) PyTorch 模块封装示例

import torch
import torch.nn as nn

class EncoderBlock(nn.Module):
    def __init__(self, d_model, n_head, d_ff, drop_prob=0.1):
        """
        Args:
            d_model: 嵌入维度（例如 512）
            n_head: 多头数量（通常设为 8）
            d_ff: Feed Forward 网络中间维度（通常为 2048）
            drop_prob: Dropout 概率
        """
        super(EncoderBlock, self).__init__()
        self.attention = MultiHeadAttention(d_model, n_head)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(drop_prob)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(drop_prob)

    def forward(self, x, src_mask=None):
        # Step 1: Multi-Head Self-Attention
        sa_output, _ = self.attention(x, x, x, src_mask)
        x = self.norm1(x + self.dropout(sa_output))

        # Step 2: Position-wise FeedForward
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))

        return x

(4) 封装整个 Encoder 模块

有了 EncoderBlock 后，我们可以将它 重复 N 次 构建完整的 Encoder：

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, n_head, d_ff, drop_prob=0.1):
        """
        Args:
            num_layers: encoder block 堆叠层数（原论文为 6）
            d_model: 模型维度（如 512）
            n_head: 注意力头数（如 8）
            d_ff: FeedForward 网络维度（如 2048）
        """
        super(TransformerEncoder, self).__init__()
        self.blocks = nn.ModuleList([
            EncoderBlock(d_model, n_head, d_ff, drop_prob)
            for _ in range(num_layers)
        ])

    def forward(self, x, mask=None):
        for block in self.blocks:
            x = block(x, mask)
        return x

📈 示例调用

# 创建输入张量
x = torch.randn(batch_size=32, seq_len=20, d_model=512)  # [32, 20, 512]

# 构建 Encoder
encoder = TransformerEncoder(num_layers=6, d_model=512, n_head=8, d_ff=2048)
output = encoder(x)

print("Encoder 输出形状:", output.shape)  # [32, 20, 512]

4. Decoder结构

Decoder是Transformer架构中用于生成输出序列的部分。与Encoder类似，它由N=6个相同的Decoder block堆叠而成，但结构更为复杂。

(1) Decoder Block 的核心组件

一个标准的Decoder block包含三个主要子层：

• Masked Multi-Head Self-Attention
• Encoder-Decoder Attention
• Position-wise Feed Forward Network

每个子层后面都跟随 残差连接（Residual Connection） 和 层归一化（Layer Normalization）。

如下图右侧红框表示一个标准的 Decoder Block。

①「Masked Multi-Head Self-Attention」

这是Decoder的第一个注意力机制，用于处理目标语言的输入序列（即解码器自身的输入）。

• 使用 masking 技术 防止在预测当前词时看到未来的词，保持因果关系。
• 实现方式：通过 trg_mask 屏蔽未来信息。

x = self.mha1(q=dec_out, k=dec_out, v=dec_out, mask=trg_mask)

②「Encoder-Decoder Attention」

这是 Decoder 的第二个注意力机制，用于将 Encoder 的输出信息融合到 Decoder 中。

• Query (Q) 来自上一层 Decoder 的输出；
• Key (K) 和 Value (V) 来自 Encoder 的输出；
• 这样 Decoder 在生成每个词时都能关注到整个输入句子的信息。

x = self.mha2(q=x, k=enc_out, v=enc_out, mask=src_mask)

③ 「Position-wise Feed Forward Network」

这是一个简单的两层全连接网络，对每个位置的向量进行非线性变换。

x = self.ffn(x)

④「残差连接 + 层归一化（Add & Norm）」

每个子层都应用：

x = self.ln(x_residual + dropout(sublayer_output))

• 提升训练稳定性；
• 缓解梯度消失问题。

⑤「Decoder的完整实现」

DecoderLayer类：

class DecoderLayer(nn.Module):
    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):
        super(DecoderLayer, self).__init__()
        # 第一个 Multi-Head Attention: Masked Self-Attention
        self.mha1 = MultiHeadAttention(d_model, n_head)
        self.ln1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(p=drop_prob)

        # 第二个 Multi-Head Attention: Encoder-Decoder Attention
        self.mha2 = MultiHeadAttention(d_model, n_head)
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(p=drop_prob)

        # 前馈网络
        self.ffn = PositionwiseFeedForward(d_model, ffn_hidden)
        self.ln3 = nn.LayerNorm(d_model)
        self.dropout3 = nn.Dropout(p=drop_prob)

    def forward(self, dec_out, enc_out, trg_mask, src_mask):
        x_residual1 = dec_out
        # Step 1: Masked Self-Attention
        x = self.mha1(q=dec_out, k=dec_out, v=dec_out, mask=trg_mask)
        x = self.ln1(x_residual1 + self.dropout1(x))

        if enc_out is not None:
            # Step 2: Encoder-Decoder Attention
            x_residual2 = x
            x = self.mha2(q=x, k=enc_out, v=enc_out, mask=src_mask)
            x = self.ln2(x_residual2 + self.dropout2(x))

        # Step 3: Position-wise Feed Forward
        x_residual3 = x
        x = self.ffn(x)
        x = self.ln3(x_residual3 + self.dropout3(x))

        return x

Decoder 类：

class Decoder(nn.Module):
    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob, device):
        super().__init__()
        # 输入嵌入 + 位置编码
        self.emb = TransformerEmbedding(
            d_model=d_model,
            drop_prob=drop_prob,
            max_len=max_len,
            vocab_size=dec_voc_size,
            device=device
        )

        # 堆叠多个 Decoder Layer
        self.layers = nn.ModuleList([
            DecoderLayer(
                d_model=d_model,
                ffn_hidden=ffn_hidden,
                n_head=n_head,
                drop_prob=drop_prob
            ) for _ in range(n_layers)
        ])

        # 最终输出层：映射到目标词汇表大小
        self.linear = nn.Linear(d_model, dec_voc_size)

    def forward(self, trg, src, trg_mask, src_mask):
        # trg: 目标序列 [batch_size, trg_seq_len]
        # src: Encoder 输出 [batch_size, src_seq_len, d_model]
        trg = self.emb(trg)

        for layer in self.layers:
            trg = layer(trg, src, trg_mask, src_mask)

        # 输出：[batch_size, trg_seq_len, dec_voc_size]
        output = self.linear(trg)
        return output

三、Transformer实现

基于前面实现的 Encoder 和 Decoder 组件，我们就可以实现 Transformer 模型的完整代码，如下所示:

import torch
from torch import nn

class Transformer(nn.Module):
    def __init__(self, src_pad_idx, trg_pad_idx, trg_sos_idx, enc_voc_size, dec_voc_size, 
                 d_model, n_head, max_len, ffn_hidden, n_layers, drop_prob, device):
        super().__init__()
        
        # 保存特殊标记的索引
        self.src_pad_idx = src_pad_idx  # 源语言填充索引
        self.trg_pad_idx = trg_pad_idx  # 目标语言填充索引
        self.trg_sos_idx = trg_sos_idx  # 目标语言起始符号索引
        self.device = device  # 设备信息
        
        # 构建 Encoder
        self.encoder = Encoder(
            d_model=d_model,  # 模型维度
            n_head=n_head,  # 注意力头数量
            max_len=max_len,  # 最大序列长度
            ffn_hidden=ffn_hidden,  # 前馈网络隐藏层维度
            enc_voc_size=enc_voc_size,  # 源语言词汇表大小
            drop_prob=drop_prob,  # Dropout 概率
            n_layers=n_layers,  # Encoder 层数
            device=device  # 设备信息
        )
        
        # 构建 Decoder
        self.decoder = Decoder(
            d_model=d_model,  # 模型维度
            n_head=n_head,  # 注意力头数量
            max_len=max_len,  # 最大序列长度
            ffn_hidden=ffn_hidden,  # 前馈网络隐藏层维度
            dec_voc_size=dec_voc_size,  # 目标语言词汇表大小
            drop_prob=drop_prob,  # Dropout 概率
            n_layers=n_layers,  # Decoder 层数
            device=device  # 设备信息
        )

    def forward(self, src, trg):
        # 创建源序列的 padding mask
        src_mask = self.make_pad_mask(src, src, self.src_pad_idx, self.src_pad_idx)
        
        # 创建目标序列到源序列的 mask
        src_trg_mask = self.make_pad_mask(trg, src, self.trg_pad_idx, self.src_pad_idx)
        
        # 创建目标序列的 padding mask 和因果mask的组合
        trg_mask = self.make_pad_mask(trg, trg, self.trg_pad_idx, self.trg_pad_idx) * \
                   self.make_no_peak_mask(trg, trg)

        # 编码器前向传播
        enc_src = self.encoder(src, src_mask)
        
        # 解码器前向传播
        output = self.decoder(trg, enc_src, trg_mask, src_trg_mask)
        
        return output

    def make_pad_mask(self, q, k, q_pad_idx, k_pad_idx):
        # 获取输入序列长度
        len_q, len_k = q.size(1), k.size(1)

        # 创建针对 key 的 mask
        # batch_size x 1 x 1 x len_k
        k_mask = k.ne(k_pad_idx).unsqueeze(1).unsqueeze(2)
        # batch_size x 1 x len_q x len_k
        k_mask = k_mask.repeat(1, 1, len_q, 1)

        # 创建针对 query 的 mask
        # batch_size x 1 x len_q x 1
        q_mask = q.ne(q_pad_idx).unsqueeze(1).unsqueeze(3)
        # batch_size x 1 x len_q x len_k
        q_mask = q_mask.repeat(1, 1, 1, len_k)
        
        # 组合两个 mask
        mask = k_mask & q_mask
        
        return mask

    def make_no_peak_mask(self, q, k):
        # 创建因果mask，防止解码器看到未来信息
        len_q, len_k = q.size(1), k.size(1)
        
        # 创建下三角矩阵，保证解码器只能关注当前词及之前的词
        mask = torch.tril(torch.ones(len_q, len_k)).type(torch.BoolTensor).to(self.device)
        
        return mask