引言
很长一段时间里，我都想阅读并实现 Transformer 论文。在读过几篇配图清晰的文章之后，比如 Transformers from Scratch 和 The Illustrated Transformer，这些概念对我来说仍然有些抽象，因为我还没有亲自运行过代码。很多人提到，完整实现会相当繁琐。不过，在 ChatGPT 的帮助下，我搭出了一个最小可行产品，让这些想法变得更具体。当我看到训练 epoch 的输出和生成文本时，原本抽象的概念开始呈现出更清晰的形态。
在这篇文章中，我们将借助 OpenAI 的 ChatGPT，在 PyTorch 中创建一个小型 Transformer 风格的语言模型。我们会构建一个玩具数据集，对其进行分词，创建词表，编码输入，并训练模型来预测下一个 token。这个模型包含 token embedding、位置编码、多头自注意力、逐位置前馈网络、残差连接和层归一化。
这有意设计成一个小型教学示例。它并不是为了生成一个实用的语言模型，但已经足够让 Transformer 的主要组成部分更容易被观察和理解。
1. 准备数据集
首先，创建一个只有三句话的小型数据集：
sentences = [
"The quick brown fox jumped over the lazy dog.",
"Advancements in AI have transformed the way we interact with technology.",
"Yesterday, the stock market experienced a significant decline due to geopolitical tensions."
]2. 对数据集进行分词
在这个最小示例中，我们通过按空白字符拆分每个句子来进行分词：
tokenized_sentences = [sentence.split() for sentence in sentences]这种做法很简单，但也有局限。标点符号仍然会附着在单词上，所以 dog. 和 dog 会被视为不同的 token。在真实项目中，应使用适合模型和语料的 tokenizer。
3. 创建词表
接下来，根据分词后的句子创建词表。我们加入用于填充、句子开头和句子结尾的特殊 token。
word2idx = {"[PAD]": 0, "[CLS]": 1, "[SEP]": 2}
for sentence in tokenized_sentences:
for token in sentence:
if token not in word2idx:
word2idx[token] = len(word2idx)
idx2word = {idx: word for word, idx in word2idx.items()}4. 编码并填充数据集
现在将每个 token 转换为整数 ID，加入特殊 token，并把所有序列填充到相同长度。
max_seq_len = max(len(sentence) for sentence in tokenized_sentences) + 2
encoded_sentences = []
for sentence in tokenized_sentences:
encoded_sentence = [word2idx["[CLS]"]] + [word2idx[word] for word in sentence] + [word2idx["[SEP]"]]
encoded_sentence += [word2idx["[PAD]"]] * (max_seq_len - len(encoded_sentence))
encoded_sentences.append(encoded_sentence)这里的 + 2 是为 [CLS] 和 [SEP] 两个 token 预留空间。否则，在加入特殊 token 之后，最长序列就没有足够的位置了。
5. 定义 Transformer 模型
下面是主要组件的一个紧凑 PyTorch 实现。
import torch
import torch.nn as nn
import torch.optim as optim
class PositionalEncoding(nn.Module):
def __init__(self, d_model):
super(PositionalEncoding, self).__init__()
self.d_model = d_model
def forward(self, x):
seq_len = x.size(1)
pe = torch.zeros(seq_len, self.d_model)
position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, self.d_model, 2).float()
* (-torch.log(torch.tensor(10000.0)) / self.d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).to(x.device)
x = x + pe
return x
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, nhead):
super(MultiHeadAttention, self).__init__()
assert d_model % nhead == 0
self.nhead = nhead
self.head_dim = d_model // nhead
self.qkv_linear = nn.Linear(d_model, d_model * 3)
self.fc = nn.Linear(d_model, d_model)
self.scale = self.head_dim ** -0.5
def forward(self, x):
batch_size, seq_len, _ = x.size()
qkv = self.qkv_linear(x)
qkv = qkv.view(batch_size, seq_len, self.nhead, 3 * self.head_dim)
qkv = qkv.transpose(1, 2)
q, k, v = qkv.chunk(3, dim=-1)
attn_scores = torch.matmul(q, k.transpose(-1, -2)) * self.scale
attn_weights = torch.softmax(attn_scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_len, -1)
attn_output = self.fc(attn_output)
return attn_output
class FeedForwardNetwork(nn.Module):
def __init__(self, d_model, dim_feedforward):
super(FeedForwardNetwork, self).__init__()
self.fc1 = nn.Linear(d_model, dim_feedforward)
self.fc2 = nn.Linear(dim_feedforward, d_model)
self.relu = nn.ReLU()
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
class TransformerBlock(nn.Module):
def __init__(self, d_model, nhead, dim_feedforward):
super(TransformerBlock, self).__init__()
self.mha = MultiHeadAttention(d_model, nhead)
self.ffn = FeedForwardNetwork(d_model, dim_feedforward)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(0.1)
def forward(self, x):
attn_output = self.mha(x)
x = self.norm1(x + self.dropout(attn_output))
ffn_output = self.ffn(x)
x = self.norm2(x + self.dropout(ffn_output))
return x
class TransformerModel(nn.Module):
def __init__(self, vocab_size, d_model, nhead, num_layers, dim_feedforward):
super(TransformerModel, self).__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model)
self.transformer_blocks = nn.ModuleList(
[TransformerBlock(d_model, nhead, dim_feedforward) for _ in range(num_layers)]
)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = self.pos_encoding(x)
for block in self.transformer_blocks:
x = block(x)
x = self.fc(x)
return x注意力模块会把每个 token 的表示投影成 query、key 和 value 向量。query-key 点积用于估计每个 token 应该在多大程度上关注其他 token。随后，value 向量会根据这些注意力权重进行组合。
6. 训练模型
实例化模型，定义损失函数和优化器，然后训练若干个 epoch。
vocab_size = len(word2idx)
d_model = 8
nhead = 2
num_layers = 1
dim_feedforward = 16
model = TransformerModel(vocab_size, d_model, nhead, num_layers, dim_feedforward)
criterion = nn.CrossEntropyLoss(ignore_index=word2idx["[PAD]"])
optimizer = optim.Adam(model.parameters(), lr=0.001)
input_data = torch.tensor(encoded_sentences[:-1], dtype=torch.long)
target_data = torch.tensor(encoded_sentences[1:], dtype=torch.long)
num_epochs = 200
for epoch in range(num_epochs):
optimizer.zero_grad()
output = model(input_data)
loss = criterion(output.view(-1, vocab_size), target_data.view(-1))
loss.backward()
optimizer.step()
if (epoch + 1) % 20 == 0:
print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
这个玩具训练设置使用每个句子来预测下一个编码后的句子，所以它只是对张量形状、优化过程和模型结构的粗略演示。若要采用更常规的下一个 token 语言建模设置，可以通过将同一序列平移一个 token 来构造 input-target 对：
all_sequences = torch.tensor(encoded_sentences, dtype=torch.long)
input_data = all_sequences[:, :-1]
target_data = all_sequences[:, 1:]然后继续使用相同的模型输出和交叉熵损失模式进行训练。
结论
在这篇文章中，我们借助 OpenAI 的 ChatGPT，用 PyTorch 创建了一个简单的 Transformer 风格模型。这个练习涵盖了最小数据集、空白分词、词表构建、位置编码、多头注意力、前馈层、残差连接以及一个简短的训练循环。
这类 MVP 的关键价值并不在于性能，而在于可见性。运行一个小模型，可以让 Transformer 的抽象组件更容易被检查、调试，并与原始讲解中的图示对应起来。
若要查看我使用的提示词和生成的代码，请访问我的 GitHub 仓库：BuildChachaGPTWithChatGPT。