Transformer模型解析与实例：搭建一个自己的预测语言模型

1. 前言

2. Transformer 的核心结构

2.1 编码器（Encoder）

2.2 解码器（Decoder）

2.3 位置编码（Positional Encoding）

3. 使用 PyTorch 构建 Transformer

3.1 导入所需的模块：

3.2 定义位置编码

3.3 构建 Transformer 模型

3.4 训练模型

3.5 完整代码

4. 总结

1. 前言

Transformer 是一种革命性的深度学习架构，最初由 Vaswani 等人在 2017 年的论文《Attention Is All You Need》中提出。它通过引入自注意力机制（Self-Attention），解决了传统 RNN 和 LSTM 在处理长序列时的效率和性能问题。Transformer 的核心思想是让模型能够同时关注序列中所有位置的信息，而不是像 RNN 那样逐个处理序列元素。

本文将深入剖析 Transformer 的核心结构，并通过 PyTorch 实现一个完整的 Transformer 模型，帮助读者全面理解这一架构的原理和应用，同时也作为LLM学习的第一章。

2. Transformer 的核心结构

首先的首先，文章中每个单词都会被映射为一个高维的向量，称其为embedding层，该映射后的向量则称为嵌入向量。

2.1 编码器（Encoder）

Transformer 的编码器由多个相同的层堆叠而成，每一层包含两个主要子模块：

多头自注意力机制（Multi-Head Self-Attention）
前馈神经网络（Feed-Forward Network）

多头自注意力机制

多头自注意力机制是 Transformer 的核心组件之一。它的作用是让模型能够同时关注序列中不同位置的元素，并捕捉它们之间的关系。

自注意力机制（Self-Attention）

自注意力机制通过计算查询向量（Query）、键向量（Key）和值向量（Value）之间的点积，得到一个注意力分数，用于加权求和得到输出向量。

输入表示：输入是一个序列，每个元素是一个向量 xi。
线性变换：将输入向量分别映射到查询向量 Q、键向量 K 和值向量 V：

其中 WQ,WK,WV 是可学习的权重矩阵。
计算注意力分数：通过点积计算查询向量和键向量之间的相似度：

其中 dk 是键向量的维度，用于缩放点积以稳定训练。

多头注意力（Multi-Head Attention）

多头注意力通过将输入向量映射到多个不同的注意力头，每个头独立计算注意力分数，最后将所有头的输出拼接在一起。这种方法可以捕捉不同粒度的特征。

通过 WO 矩阵将拼接后的输出投影到模型所需的维度。

其中：

前馈神经网络

每个编码器层的第二个子模块是一个前馈神经网络，用于进一步处理注意力机制的输出。这个网络对每个位置的元素独立应用相同的变换：

前馈神经网络对每个位置的词向量独立进行处理，这意味着它可以在所有位置上并行计算。这种设计充分利用了现代硬件（如 GPU）的并行计算能力，显著提高了模型的效率。

残差连接与层归一化

为了稳定训练，编码器的每个子模块都使用了残差连接（Residual Connection）和层归一化（Layer Normalization）：

2.2 解码器（Decoder）

解码器的结构与编码器类似，但多了一个额外的子模块：

掩码多头自注意力机制（Masked Multi-Head Self-Attention）
编码器-解码器注意力（Encoder-Decoder Attention）
前馈神经网络

掩码多头自注意力机制

掩码多头自注意力机制的作用是防止解码器在生成目标序列时看到未来的位置信息。通过在注意力分数中应用掩码（将未来位置的分数设为负无穷），确保每个位置只能关注当前位置及之前的位置。

编码器的多头自注意力：查询（Q）、键（K）和值（V）都来自同一个输入序列。
解码器的多头自注意力：查询（Q）、键（K）和值（V）都来自解码器的上一层输出，但受到掩码的限制。

编码器-解码器注意力

编码器-解码器注意力用于将解码器的输出与编码器的输出结合起来。查询向量来自解码器，而键和值向量来自编码器，这样解码器可以利用编码器生成的上下文信息。

2.3 位置编码（Positional Encoding）

由于 Transformer 没有像 RNN 那样的隐状态来捕捉序列顺序，因此需要引入位置编码。位置编码通过将序列的位置信息嵌入到词向量中，使模型能够感知序列的顺序。

词嵌入与位置编码相加，而不是拼接，他们的效率差不多，但是拼接的话维度会变大，所以不考虑。

位置编码通常通过正弦和余弦函数实现：

其中 pos 是位置，i 是维度（某个位置向量中的第 i 个维度）。

3. 使用 PyTorch 构建 Transformer

3.1 导入所需的模块：

import torch
import torch.nn as nn
import torch.optim as optim
import math

3.2 定义位置编码

位置编码通过正弦和余弦函数实现：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, 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, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-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).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

torch.arange：生成一个从 0 到 max_len-1 的一维张量，步长为 1。
unsqueeze(1)：在张量的第 1 个维度（索引从 0 开始）增加一个维度。这将一维张量 (max_len,) 转换为二维张量 (max_len, 1)。

position：表示位置信息，形状为 (max_len, 1)。
div_term：用于计算 100002i/dmodel，形状为 (d_model//2,)。

(-math.log(10000.0) / d_model)：计算一个负的对数值，除以 d_model。这个值用于调整正弦和余弦波的频率。
0::2：表示从索引 0 开始，每隔两列取一个元素（即偶数列）。

pe.unsqueeze(0)：在 pe 的第 0 维增加一个维度，将形状从 (max_len, d_model) 转换为 (1, max_len, d_model)。
pe.transpose(0, 1)：交换第 0 维和第 1 维，得到形状为 (max_len, 1, d_model)。
pe = pe.unsqueeze(0).transpose(0, 1) 用于调整位置编码的形状，使其能够与输入词向量相加。

缓冲区用于存储模型中需要保存但不需要被优化器更新的张量。例如：

位置编码（Positional Encoding）：位置编码是模型的一部分，但不需要在训练过程中更新。
词向量（Embedding）：如果词向量是预训练的且不需要更新，也可以将其注册为缓冲区。

通过 self.register_buffer('pe', pe)，将 pe 注册为模型的缓冲区，这样：

pe 会成为模型的一部分，并在保存和加载模型时被自动保存和加载。
pe 不会被优化器更新，因为它不是可学习的参数。

3.3 构建 Transformer 模型

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, nlayers, nlayers, nhid, dropout)
        self.fc_out = nn.Linear(d_model, vocab_size)

        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        src = self.embedding(src) * math.sqrt(self.embedding.embedding_dim)
        src = self.pos_encoder(src)
        tgt = self.embedding(tgt) * math.sqrt(self.embedding.embedding_dim)
        tgt = self.pos_encoder(tgt)
        output = self.transformer(src, tgt, src_mask, tgt_mask)
        output = self.fc_out(output)
        return output

vocab_size：词汇表的大小，即词典中不同单词的数量。
d_model：词向量的维度，也是 Transformer 中每个层的输入和输出维度。
nhead：多头注意力机制中的头数。
nhid：前馈神经网络的隐藏层维度。
nlayers：编码器和解码器的层数。
dropout：Dropout 的概率，默认为 0.5。

设置模型类型为 'Transformer'，用于标识模型的架构。

self.fc_out = nn.Linear(d_model, vocab_size)

初始化输出层，将 Transformer 的输出映射回词汇表的大小，用于预测下一个单词。

initrange = 0.1

作用：定义了一个初始化范围，用于设置权重的初始值范围。
值：0.1 表示权重将被初始化为在 -0.1 到 0.1 之间的均匀分布。

self.embedding.weight.data.uniform_(-initrange, initrange)

作用：将词嵌入层的权重初始化为均匀分布。
uniform_：PyTorch 中的函数，用于将张量的值填充为指定范围内的均匀分布。
self.embedding.weight.data：访问词嵌入层的权重张量。
-initrange, initrange：指定均匀分布的范围。

self.fc_out.bias.data.zero_()

作用：将输出层的偏置初始化为零。

对于forward函数：

self.embedding：词嵌入层，将单词索引映射到词向量。
src 和 tgt：分别是输入序列和目标序列。
math.sqrt(self.embedding.embedding_dim)：对嵌入向量进行缩放，确保其具有合适的尺度，防止梯度消失或爆炸。
src_mask 和 tgt_mask：分别是输入序列和目标序列的掩码，用于处理填充部分或防止信息泄露。

其中

src = self.embedding(src) * math.sqrt(self.embedding.embedding_dim)

输入：src 的形状为 (seq_len, batch_size)，其中 seq_len 是序列长度，batch_size 是批量大小。
输出：src 的形状为 (seq_len, batch_size, d_model)，其中 d_model 是词向量的维度。
嵌入向量的方差：嵌入向量的初始化通常具有方差 1/d_model，乘以 sqrt(d_model) 后方差变为 1。

nn.Linear 是 PyTorch 中的全连接层，它接受任意维度的输入，只要最后一个维度是输入特征的维度即可。在这种情况下，输入张量的形状是 (tgt_seq_len, batch_size, d_model)，其中 d_model 是输入特征的维度。

3.4 训练模型

# 定义超参数
vocab_size = 10000  # 假设词汇表大小为 10000
d_model = 512       # 词向量维度
nhead = 8           # 多头注意力的头数
nhid = 2048         # 前馈网络的隐藏层维度
nlayers = 6         # 编码器和解码器的层数
dropout = 0.5       # Dropout 概率

# 创建模型
model = TransformerModel(vocab_size, d_model, nhead, nhid, nlayers, dropout)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in train_loader:
        src, tgt = batch  # 假设 src 和 tgt 是训练数据
        optimizer.zero_grad()
        output = model(src, tgt[:, :-1])  # 输入解码器时去掉最后一个词
        loss = criterion(output.view(-1, vocab_size), tgt[:, 1:].reshape(-1))  # 预测下一个词
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader)}")

3.5 完整代码

完整代码（数据集需要自己定义）如下方便调试：

import torch
import torch.nn as nn
import torch.optim as optim
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, 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, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-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).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, nlayers, nlayers, nhid, dropout)
        self.fc_out = nn.Linear(d_model, vocab_size)

        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        src = self.embedding(src) * math.sqrt(self.embedding.embedding_dim)
        src = self.pos_encoder(src)
        tgt = self.embedding(tgt) * math.sqrt(self.embedding.embedding_dim)
        tgt = self.pos_encoder(tgt)
        output = self.transformer(src, tgt, src_mask, tgt_mask)
        output = self.fc_out(output)
        return output

vocab_size = 10000  # 假设词汇表大小为 10000
d_model = 512       # 词向量维度
nhead = 8           # 多头注意力的头数
nhid = 2048         # 前馈网络的隐藏层维度
nlayers = 6         # 编码器和解码器的层数
dropout = 0.5       # Dropout 概率

# 创建模型
model = TransformerModel(vocab_size, d_model, nhead, nhid, nlayers, dropout)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in train_loader:
        src, tgt = batch  # 假设 src 和 tgt 是训练数据
        optimizer.zero_grad()
        output = model(src, tgt[:, :-1])  # 输入解码器时去掉最后一个词
        loss = criterion(output.view(-1, vocab_size), tgt[:, 1:].reshape(-1))  # 预测下一个词
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader)}")