LLM架构(2) - 从零开始训练一个AI

前言

在上一篇文章中，我们已经学习了《Attention Is All You Need》论文中提出的初代 Transformer 架构。它是 LLM 的理论基石，虽然严格意义上它不是一个大语言模型，因为它的设计目的主要是为了机器翻译，而不是自回归语言生成模型，但不影响它依然是如今所有 LLM 的爷爷。在这一系列文章中，我们正是以此作为起点，依次去学习初版 Transformer、GPT-1、GPT-2、LLaMA、Qwen、DeepSeek 等一系列开源的大语言模型，希望从这个演化的过程中一步一步地了解它们在架构上面的进化，特别是思想层面的演进。

首先，需要说明的是学习这些东西的目的是什么？是为了在 AI 时代，更深刻的学习并理解 AI 的本质是什么。

而学习一个东西最好的方式，就是将它拆解出来，理解每一个零件，了解它底层的原理，然后从零开始一步一步地再将它拼装起来。正如儿时的我最大的爱好，就是将买回家的每一个玩具都拆开然后再拼装起来，虽然最后总是会多出几个零件和螺丝，但在这个过程中，我享受到了比玩这个玩具更大的乐趣，因为人类的快乐其实只源于两点：破坏和建造。

言归正传，这篇文章将紧接着上一篇继续学习，如果你没有阅读过上一篇，阅读到此请停止继续阅读，因为在本篇文章中，我将基于上一篇文章里提到的理论基础和数学公式，以代码的形式去将它们一一实现，从零开始构建论文中描述的 Transformer Base 模型。

模型介绍

在本文中，我将尝试先用 PyTorch 代码搭建一个论文中描述的 Transformer Base 模型，然后用一个小数据集在 NVIDIA A100 单卡上训练出一个本地小模型，并以此模型为基础实现一个支持：英文-意大利语，英文-中文，中文-文言文的多语言 AI 翻译机器人。

我一共会训练出 3 个本地小模型，每个模型实现一种语言的翻译，以下是其中一个小模型的基本信息：

模型信息	实际值
数据集样本量	966,865
源词表大小	9,303
目标词表大小	12,058
序列长度	128
模型维度(`d_model`)	512
Transformer Block层数	6
Attention heads	8
Feed-forward维度	2,048
模型参数量	61 M
模型权重文件大小	735 MiB

备注：除了数据集不同，其他模型的参数和论文中的基本保持一致。之所以更换更小的数据集，主要是因为没有卡，训练大数据集太浪费时间了，我的重点是学习模型的架构设计，而不是复现论文的实验结果。

最终开发出来的 AI 翻译机器人截图如下：

模型架构

在上一篇文章中，我们主要根据论文学习了 Transformer 的模型架构，如下：

其中主要有几个部分组成：

1. Input Embedding
2. Positional Encoding
3. Encoder & Decoder

1. Multi-Head Attention
2. Add & Norm
3. Feed Forward Network

4. Linear + Softmax

接下来，我们将一一用代码来实现其中每个部分，并最后用它去训练我们的模型。

Input Embedding

在实现 Input Embedding 之前，我们需要先准备用于训练的数据集，这里我将使用和论文中差不多的 WMT 数据集，以及另外2个语言的数据集：

• opus_books (英文-意大利语)
• wmt19 (英文-中文)
• Chinese_modern_classical（中文-文言文）

在准备好数据集之后，我们需要先训练一个 Tokenizer （分词器），并根据这个数据集，使用我们的分词器构建我们自己的词典表（ Vocabulary ）。

然后我们将数据集随机的拆分为两个部分：90%为训练数据（ Training Dataset ），10%为验证数据（ Validation Dataset )。

因为国内访问 Huggingface 总是不稳定，所以这里我是将数据集下载到本地文件，然后加载到内存：

# 数据集dataset = load_dataset("parquet", data_files=local_dataset_path, split='train')# 分词器tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))tokenizer.pre_tokenizer = Whitespace()# 词典表trainer = WordLevelTrainer(special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"], min_frequency=2)tokenizer.train_from_iterator(get_all_sentences(ds, lang), trainer=trainer)

注：其中中文会使用 BPE 分词。

有了分词器和词典表以后，我们就可以开始编写 Input Embedding 了：

def init():    self.embedding = nn.Embedding(vocab_size, d_model)def forward(self, x):    return self.embedding(x) * math.sqrt(self.d_model)

这里的计算公式来自论文中的：

In the embedding layers, we multiply those weights by .

这是 Transformer 论文里的一个小但重要的缩放，原因主要是：让 token embedding 的数值尺度和 positional encoding 的数值尺度更匹配。因为在 Transformer 里，输入是：input_embedding + positional_encoding，如果 embedding 初始化后数值比较小，那么直接相加时，里面 positional encoding 可能会占比较大，导致模型一开始更“看重位置”，而不是 token 本身。所以论文把 embedding 乘上让 embedding 的尺度变大一些，这里的乘法是为了让“这个词是什么”的信息不要在一开始被位置信息淹没。

Positional Encoding

在得到 Input Embedding 之后，还需要将位置信息编码进去。论文中的数学公式是：

这里用代码的实现方式，如下：

# PE(pos, 2i)positional_encoding[:, 0::2] = torch.sin(position * div_term)# PE(pos, 2i + 1)positional_encoding[:, 1::2] = torch.cos(position * div_term)

然后按照论文里面说的将 Input embedding 和 positional_encoding 两个矩阵相加即可：

The positional encodings have the same dimension the embeddings, so that the two can be summed.

不同维度使用不同频率的正余弦函数，本质上是在让模型能够通过向量之间的线性关系，推导出 Token 之间的相对距离关系。这也是为什么 Transformer 即使没有循环结构，也依然能够感知序列顺序。

Multi-Head Attention

接下来就是 Transformer 中最核心的部分：Self-Attention。在论文中，多头注意力（Multi-Head Attention）的定义如下：

根据这个公式，我们需要四个线性层：

# Wqself.w_q = nn.Linear(d_model, d_model, bias=False)# Wkself.w_k = nn.Linear(d_model, d_model, bias=False)# Wvself.w_v = nn.Linear(d_model, d_model, bias=False)# Woself.w_o = nn.Linear(d_model, d_model, bias=False)

随后会将特征拆分为多个 attention head, 每个 head 只关注其中一部分维度，例如 d_model 为 512，分为 8 个 head, 那么每个 head 只需要关注 512 / 8 = 64 个维度：

self.d_k = d_model // h # Dimension of vector seen by each head

这样不同的 head 就能够在不同的表示子空间中学习不同的 attention 模式。

首先使用 Wq , Wk 和 Wv 将输入映射为 Query、Key 和 Value：

query = self.w_q(q)key = self.w_k(k)value = self.w_v(v)

然后将其拆分为多个 head：

query = query.view(batch_size, seq_len, self.h, self.d_k).transpose(1, 2)key = key.view(batch_size, seq_len, self.h, self.d_k).transpose(1, 2)value = value.view(batch_size, seq_len, self.h, self.d_k).transpose(1, 2)

然后计算每个 head 的 attention scores ，公式如下：

对应的代码实现为：

# 计算 Query 与 Key 的相似度attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)# Mask，确保不能看未来attention_scores.masked_fill_(mask == 0, -1e9)# 归一化,转换为注意力权重attention_scores = attention_scores.softmax(dim=-1)# 将注意力作用到 Value 上x = (attention_scores @ value)

注：其中的Mask ，在 Decoder 的 Attention 中会使用 causal mask，防止当前位置看到未来 token。在 Encoder 的 Attention 中，mask 通常用于屏蔽 padding token。

你可以把 Query 理解成“我正在寻找什么”，把 Key 理解成“我能提供什么信息”，而点积本质上是在计算两者的匹配程度。匹配越高，说明当前 Token 越应该关注另一个 Token。

最后，将多个 head 的结果重新拼接回原始维度，并通过 W^O 融合不同 head 学习到的信息：

x = x.transpose(1, 2).contiguous().view(    x.shape[0], -1, self.h * self.d_k)self.w_o(x)

Add & Norm

然后是残差连接（Residual Connection），这是深度学习中的一个重要发明。在残差连接提出之前，神经网络一旦变得很深，就容易出现梯度消失、梯度爆炸以及训练退化等问题，导致模型难以继续优化。而引入残差连接后，模型即使拥有很深的网络结构，也依然能够保持良好的训练效果。

在 Transformer 中，通常会将子层的输出与输入直接相加：

x = x + sublayer(x)

为了进一步稳定训练过程，在进入下一层之前，还会使用 Layer Normalization 对特征进行归一化，从而缓解梯度问题并加速模型收敛，其公式如下：

代码实现如下：

# 求平均值mean = x.mean(dim = -1, keepdim = True)# 求标准差std = x.std(dim = -1, keepdim = True) # (batch, seq_len, 1)# 归一化并进行缩放与平移return self.alpha * (x - mean) / std + self.eps + self.bias

其中 eps 是一个极小值，用于避免除零错误。这虽然没有在论文中重点强调，但属于工程实现中的常见技巧。

Feed Forward Network

接下来是前馈网络（Feed Forward Network, FFN）。在 Transformer 中，Multi-Head Attention 负责学习 Token 之间的依赖关系，而 FFN 则会对每个 Token 的特征进行独立的非线性变换，从而进一步增强模型的表达能力。

FFN 本质上是一个两层的全连接神经网络。第一层会先将特征从 512 维映射到更高维的 2048 维特征空间，以提升模型的非线性表达能力；随后第二层再将特征映射回原始的 512 维。论文里面的公式如下：

其中：

• 第一层线性变换负责升维
• ReLU 用于引入非线性
• 第二层线性变换负责降维

代码实现如下：

# 第一层网络（升维），w1 和 b1self.linear_1 = nn.Linear(d_model, d_ff)# 第二层网络（降维），w2 和 b2self.linear_2 = nn.Linear(d_ff, d_model)x = self.linear_1(x)x = torch.relu(x)x = self.linear_2(x)

Linear + Softmax

最后就是输出层，首先会经过一个 Linear 线性层，将 Transformer 输出的特征映射到词表（Vocabulary）大小的维度上，例如 d_model = 512，词表大小 vocab_size = 30000，那么 Linear 层就会将每个 Token 的 512 维特征映射成一个长度为 30000 的向量，这个向量中的每一个值都表示当前 Token 对应词表中某个单词的预测分数：

self.proj = nn.Linear(d_model, vocab_size)logits = self.proj(x)

这里输出的 logits 还不是概率，而是模型对词表中每个 token 给出的未归一化分数，因此接下来还需要经过 Softmax 将这些分数转换成概率分布，公式如下：

经过 Softmax 后，所有值都会被压缩到 [0, 1] 之间，并且所有概率之和为 1，这样模型就可以根据概率大小选择最有可能的下一个 Token：

probs = torch.softmax(logits, dim=-1)next_token = torch.argmax(probs, dim=-1)

在真正的大语言模型推理中，通常不会直接使用 argmax，而是会通过 sampling 策略增加生成文本的多样性。

Transformer

到此，我们就基本上 Transformer 模型架构中的每一个模块都实现了一遍。现在，只需要像搭积木一样，将这些模块按照论文中的结构组织起来，就能够真正搭建出完整的 Transformer 模型了，数据流如下：

Input Embedding    ↓Positional Encoding    ↓Encoder × N    ↓Decoder × N    ↓Projection Layer    ↓Vocabulary Logits

当你真正开始把这些模块一个一个拼装起来的时候，你会突然发现一件非常有意思的事情：Transformer 的伟大之处，其实并不在于某一个单独模块有多么复杂，相反，它里面绝大部分东西单独拿出来看，都并不算特别“惊艳”的发明。但 Transformer 真正厉害的地方在于：它第一次用一种极其优雅的方式将它们合理的组装了起来，而正是这种组合方式几乎彻底改变了整个人工智能的发展方向。

结语

亲自动手实现 2017 年论文中提出的初代 Transformer 架构，并从零开始训练出一个 AI 模型，在现如今其实已经非常简单了，但在这个过程中，我们依然可以学习到很多非常厉害的思想，特别是它在解决那些困扰了人类很多年的一些问题上用的那些很巧妙的方法，这点是最值得学习的地方。

最后，还是回到开头的话题。

在人人都在谈论 AI Agent、直接调用各种顶尖模型 API 就能完成大量工作的今天，学习这些近十年前的“老东西”看起来似乎有些过时。但我的看法恰恰相反：在 AI 时代，技术本身正在快速平权。你可以用 AI 提效，别人同样也可以；大家都在使用相同的模型和工具，真正拉开差距的，已经不再只是生产力本身，而是认知力、判断力和创新力。

而对于 AI 的理解，其实也存在层次：从“使用 AI Agent”，到“构建 AI Agent”，再到底层的“训练 AI 模型”。只有沿着这条链路不断深入，才可能真正理解 AI 的本质，并形成对未来更强的判断能力。

作为一个写了十几年代码的老程序员，我认为 AI 会彻底改变“程序员”这个职业，就像工业革命改变了纺织工人一样。未来真正重要的，可能不再只是编码能力，而是理解问题、构建系统以及驾驭 AI 的能力。

尤其是在生产力已经剧烈变化，而生产关系尚未完成重构的这个短暂过渡期里，谁能够更早理解这种变化，谁就更有机会抓住这场工业革命带来的红利。

下一篇文章中，我将正式进入 GPT 时代，看看 Transformer 是如何从一个用于机器翻译的模型，一步一步进化成今天真正的大语言模型的。