第33天：从零构建迷你Transformer（代码详解）

Transformer 是 2017 年重塑自然语言处理（NLP）的架构。在 Transformer 出现之前，表现最好的模型要么使用循环神经网络（RNN），要么使用卷积神经网络（CNN）。Transformer 在机器翻译、文本摘要和问答等任务中带来了显著的性能提升。其核心优势在于注意力机制，该机制允许模型关注句子中任何位置的重要词汇，而不仅仅是邻近的词。如今你所看到的所有大型语言模型（LLM）——如 GPT、BERT、Llama——都是基于 Transformer 或其近亲构建的。

在 Transformer 出现之前，RNN 和 CNN 面临着巨大的局限。RNN 逐步处理输入，难以连接相距较远的词。CNN 能够识别局部模式，但要“看清”长文本，需要堆叠许多层。这两种类型都难以高效利用现代 GPU，因为它们的计算无法完全并行化。必须处理完输入的一部分才能继续处理下一部分。

相比之下，Transformer 一次性处理整个输入序列。其关键发明——自注意力（self-attention）——让每个词都能“查看”输入中的其他所有词。这使得模型能够生成丰富且具备上下文感知的表示，并支持快速、并行的计算。

让我们概述 Transformer 的主要组成部分：

• 输入嵌入（Input Embeddings）：模型将单词或标记转换为向量（数字列表），以捕捉基本含义和结构。
• 自注意力（Self-Attention）：根据彼此之间的关系，每个嵌入可以从同一序列中的其他所有嵌入中收集信息。这为每个标记从整个序列中获取上下文。
• 前馈网络（Feedforward Network）：每个具备上下文感知能力的嵌入通过一个简单的神经网络，进一步对其进行优化。
• 输出：通过堆叠层可以重复此过程。最终输出是一系列向量，可用于文本分类或下一个词生成等预测任务。

重点关注作为核心机制的自注意力。以句子“The cat sat on the mat.”为例。自注意力使得“mat”这个标记能够访问来自“the”、“cat”、“sat”和“on”的信息，从而在上下文中构建更丰富的含义。

以下是自注意力的核心步骤：

1. 对于每个输入标记（模型对单词的内部表示），创建三个不同的版本：查询（Query）、键（Key）和值（Value）。每个版本都是一个向量，通过将标记嵌入乘以小型矩阵得到。
2. 对于每一对标记，通过取一个标记的 Query 和另一个标记的 Key 来计算一个“分数”。这是通过点积完成的：将它们对应的元素相乘并求和。
3. 使用 softmax 函数对这些分数进行缩放和归一化，使它们针对每个 Query 相加等于 1。这就得到了注意力权重——即应该“关注”其他每个标记的程度。
4. 每个标记的新嵌入是所有 Value 向量的加权和，使用权重即为上述注意力权重。

例如，如果你有标记 A、B 和 C，A 的自注意力计算如下：

• A 应该多大程度关注自身？（A 的 Query 点积 A 的 Key）
• 多大程度关注 B？（A 的 Query 点积 B 的 Key）
• 多大程度关注 C？（A 的 Query 点积 C 的 Key）A 的新表示是 A、B 和 C 的 Value 向量的加权平均值。对每个标记重复此过程。

一个 Transformer 层包含两个主要部分：

• 多头自注意力模块（Multi-Head Self-Attention block）（通常并行运行多个注意力操作；为简化起见，此处我们仅使用单个头）
• 前馈模块（Feedforward block）（每个位置一个小神经网络）

Transformer 将这些层按顺序堆叠多层，从而创建深层且灵活的表示。

下面是一个可运行的、极简的 Transformer，使用 NumPy 逐步构建。此版本去除了非必要功能：仅包含单个自注意力头、基础的前馈网络，没有其他额外特性。

import numpy as np

def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)  # 为了数值稳定性
    e_x = np.exp(x)
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

class MiniTransformer:
    def __init__(self, d_model, seq_len, vocab_size):
        self.d_model = d_model
        self.seq_len = seq_len
        self.vocab_size = vocab_size

        # 随机词元嵌入 (vocab_size x d_model)
        self.token_embeddings = np.random.randn(vocab_size, d_model) / np.sqrt(d_model)

        # 自注意力投影矩阵 (每个都是 d_model x d_model)
        self.W_Q = np.random.randn(d_model, d_model) / np.sqrt(d_model)
        self.W_K = np.random.randn(d_model, d_model) / np.sqrt(d_model)
        self.W_V = np.random.randn(d_model, d_model) / np.sqrt(d_model)

        # 前馈网络：(d_model x d_ff) 然后 (d_ff x d_model)
        d_ff = d_model * 2  # 前馈网络中的隐藏层大小
        self.W1 = np.random.randn(d_model, d_ff) / np.sqrt(d_model)
        self.b1 = np.zeros((d_ff,))
        self.W2 = np.random.randn(d_ff, d_model) / np.sqrt(d_ff)
        self.b2 = np.zeros((d_model,))

    def forward(self, token_ids):
        # token_ids: (seq_len,) 整数序列（词元索引）
        x = self.token_embeddings[token_ids]  # (seq_len, d_model)

        # -- 自注意力机制 --
        Q = x @ self.W_Q      # (seq_len, d_model)
        K = x @ self.W_K      # (seq_len, d_model)
        V = x @ self.W_V      # (seq_len, d_model)

        # 注意力分数 (seq_len, seq_len)
        attn_scores = Q @ K.T / np.sqrt(self.d_model)
        attn_weights = softmax(attn_scores, axis=1) # (seq_len, seq_len)

        # 加权求和: (seq_len, seq_len) @ (seq_len, d_model) -> (seq_len, d_model)
        attn_output = attn_weights @ V

        # -- 前馈模块 --
        ff_hidden = np.maximum(0, attn_output @ self.W1 + self.b1)  # ReLU 激活函数
        ff_output = ff_hidden @ self.W2 + self.b2

        return ff_output  # (seq_len, d_model)

# 示例用法：
np.random.seed(42)
vocab_size = 10
d_model = 8  # 嵌入维度
seq_len = 5

# 一个示例“句子”：一组整数的词元 ID 列表
token_ids = np.array([2, 5, 3, 7, 1])

model = MiniTransformer(d_model=d_model, seq_len=seq_len, vocab_size=vocab_size)
output_vectors = model.forward(token_ids)
print(output_vectors.shape)  # (5, 8)
print(output_vectors)

这段代码做了什么？

• 接收一个假想的句子：一组词元 ID。
• 查找它们的嵌入向量。
• 应用自注意力机制，让每个词元都能混合来自其他所有词元的信息。
• 将更新后的向量通过一个小神经网络层进行细化。

这不是一个生产级系统。没有训练过程、多头注意力、归一化、残差连接或位置编码——只有最基础的机械结构。它的目的是让变换器（transformer）的各个组成部分变得可见且具体。

更改输入词元、序列长度或嵌入大小。向 model.forward() 传递新的序列。观察输出向量是如何更新的。

通过逐行审视这个简单的变换器，你现在可以看到输入词元如何转化为丰富且具有上下文感知的表示。自注意力和前馈层——每一层都只是矩阵乘法和 softmax 的简单组合——构成了每个变换器的核心。现代大型语言模型将这些组件堆叠得更高、更宽，但基础原理保持不变。

关键要点

• 变换器使用自注意力机制，允许词元访问序列中任何其他词元的信息。
• 以前的模型如 RNN 和 CNN 受限于顺序处理和局部性。
• 变换器的关键组件包括输入嵌入、自注意力和前馈网络。
• 自注意力计算所有词元对之间的注意力分数，以混合上下文信息。
• 仅使用矩阵乘法和 softmax 操作即可构建极简主义的变换器。

自己动手试一试

修改提供的迷你 Transformer 代码，使其使用包含 6 个 token ID 的输入序列，而不是之前的 5 个。打印出新序列的注意力矩阵（代码中的 attn_weights）。用一两句话描述模型如何在各个位置重新分配注意力，并说明你在权重中观察到的模式。

进一步资源

• 🎥 在 PyTorch 中从零开始编码 Transformer，附带完整解释、训练和推理
• 🎥 如何从零开始编码 Transformer！
• 📘 从零开始的迷你 Transformer（PyTorch 中的仅解码器 GPT 风格）
• 📘 迷你 Transformer：端到端逐步指南（NumPy 实现）
• 📄 从零构建微型 Transformer：“乐高”式深度学习方法

Day 34 预告：什么是预训练？NLP 中的迁移学习