2025年10月15日 ai-systems

从 Happy-LLM 剥离：一份极简 Transformer 实现指南

本文从 Datawhale 的 Happy-LLM 项目出发，提供一个最小化的 Transformer 模型实现教程。通过 PyTorch 代码，我们将一步步构建从输入嵌入到输出概率的完整数据流，帮助读者在代码层面深入理解 Transformer 的核心工作原理。

内容加载中...

引言：源于 Happy-LLM 的探索精神

由 Datawhale 开源社区精心打造的 Happy-LLM 项目，为广大人工智能爱好者和从业者提供了一条从零开始系统学习大语言模型（LLM）的清晰路径。该项目不仅涵盖了从自然语言处理（NLP）基础到前沿应用的全面理论，更强调通过动手实践来深化理解。其第二章深入剖析了作为现代 LLM 基石的 Transformer 架构。

本文旨在继承 Happy-LLM 的“先理解、再实现”的精神，但将焦点进一步收窄，目标是从这个庞大的知识体系中“剥离”出一个核心的、最小化的 Transformer 模型实现。我们将不依赖任何高级封装，仅使用 PyTorch 的基础组件，来搭建一个完整的、可运行的 Transformer 模型。这趟旅程将带领你穿越从输入文本到最终输出概率的每一个关键步骤，让你在代码层面建立对 Transformer 数据流的直观且深刻的认知。

最小化实现的核心目标

在着手编写代码之前，明确我们的“最小化”原则至关重要。这并不意味着牺牲核心逻辑，而是：

聚焦核心数据流：我们将专注于一个完整的正向传播过程，即从接收原始输入序列（token IDs）开始，依次经过词嵌入、位置编码、Transformer 编码器层（包括多头自注意力和前馈网络），最后通过一个线性层得到预测的词汇概率分布。
简化模型结构：我们将仅实现一个编码器（Encoder）模块，这足以展示 Transformer 的核心思想。一个完整的编码器-解码器（Encoder-Decoder）结构，虽然对于机器翻译等任务是必需的，但其核心组件（如自注意力）与单独的编码器是相通的。
教学导向的代码：每一部分代码都将附有详尽的注释，解释其功能和在整个模型中的作用。我们将避免复杂的优化或技巧，以保持代码的清晰和易读性。

构建 Transformer 的基石：代码实现

现在，让我们开始用 PyTorch 构建这个最小化的 Transformer。

1. 词嵌入与位置编码

一切始于输入。计算机无法直接理解文本，因此我们首先需要将输入的词元（tokens）转换为高维向量。这一步由**词嵌入（Word Embedding）**层完成。

然而，Transformer 架构本身并不包含任何关于序列顺序的信息。为了让模型理解单词的位置关系（例如，“我爱你”与“你爱我”的区别），我们需要引入位置编码（Positional Encoding）。这是一种特殊设计的、与词嵌入维度相同的向量，它被加到词嵌入之上，为模型注入了序列的语序信息。

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        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 的形状是 [seq_len, batch_size, d_model]
        x = x + self.pe[:x.size(0), :]
        return x

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, tokens):
        # tokens 的形状是 [seq_len, batch_size]
        return self.embedding(tokens.long()) * math.sqrt(self.d_model)

2. Transformer 编码器层

这是模型的核心。一个编码器层由两个主要部分组成：一个多头自注意力（Multi-Head Self-Attention）机制和一个前馈神经网络（Feed-Forward Network）。

多头自注意力：允许模型在处理一个词时，同时“关注”到输入序列中的所有其他词，并根据相关性计算出该词的上下文感知表示。通过“多头”，模型可以从不同的表示子空间学习信息。
前馈网络：这是一个简单的全连接网络，对自注意力层的输出进行进一步的非线性变换，以增强模型的表达能力。

此外，每个子层（自注意力和前馈网络）的周围都包裹着残差连接（Residual Connection）和层归一化（Layer Normalization），这对于训练深度模型至关重要，可以有效防止梯度消失和爆炸。

PyTorch 的 nn.TransformerEncoderLayer 为我们提供了现成的实现，我们直接使用它来保持我们实现的“最小化”和“教学化”。

3. 组装最小化 Transformer

现在我们有了所有的部件，可以将它们组装成一个完整的模型了。我们的模型将依次包含：

TokenEmbedding 层：将输入的 token 序列转换为嵌入向量。
PositionalEncoding 层：为嵌入向量添加位置信息。
一个 nn.TransformerEncoder：它由我们前面定义的多个 nn.TransformerEncoderLayer 堆叠而成。为了最小化，我们只使用一层。
一个最终的线性层（Generator）：将 Transformer 编码器的输出映射回词汇表的大小，从而得到每个位置上每个词的预测概率。

class MinimalTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, dim_feedforward, num_encoder_layers):
        super(MinimalTransformer, self).__init__()
        self.token_embedding = TokenEmbedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=nhead, 
            dim_feedforward=dim_feedforward
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, 
            num_layers=num_encoder_layers
        )
        
        self.generator = nn.Linear(d_model, vocab_size)

    def forward(self, src):
        # src 的形状是 [seq_len, batch_size]
        src_emb = self.token_embedding(src)
        src_pos_emb = self.positional_encoding(src_emb)
        memory = self.transformer_encoder(src_pos_emb)
        logits = self.generator(memory)
        return logits

运行与验证

让我们用一些虚拟数据来测试我们的模型，看看数据是如何在其中流动的。

# 定义模型超参数
VOCAB_SIZE = 1000  # 假设我们的词汇表大小为1000
D_MODEL = 512      # 模型的维度
NHEAD = 8          # 多头注意力的头数
DIM_FEEDFORWARD = 2048 # 前馈网络的维度
NUM_ENCODER_LAYERS = 1 # 编码器层数，保持最小化

# 实例化模型
model = MinimalTransformer(VOCAB_SIZE, D_MODEL, NHEAD, DIM_FEEDFORWARD, NUM_ENCODER_LAYERS)
model.eval()

# 创建虚拟输入数据
SEQ_LEN = 10
BATCH_SIZE = 4
input_tokens = torch.randint(0, VOCAB_SIZE, (SEQ_LEN, BATCH_SIZE))

# 正向传播
with torch.no_grad():
    output_logits = model(input_tokens)

# 检查输出形状
# 输出的形状应为 [seq_len, batch_size, vocab_size]
# 这代表了在每个序列位置上，模型预测的词汇表中每个词的原始分数
print(f"Input shape: {input_tokens.shape}")
print(f"Output logits shape: {output_logits.shape}")

# (10, 4, 1000)
# 这意味着对于一个长度为10的序列，批处理大小为4，模型为每个位置都预测了一个1000维的向量

这段代码清晰地展示了数据从 [10, 4] 的整数 ID 序列，经过嵌入、位置编码和 Transformer 核心模块的处理，最终转换为 [10, 4, 1000] 的浮点数张量，代表了模型对下一个词的预测分布。

结论与展望

通过遵循 Happy-LLM 项目的实践精神，我们成功地从零开始构建了一个虽小但功能完备的 Transformer 编码器模型。这个过程揭示了 Transformer 内部清晰而强大的数据处理流程：从离散的符号输入，到丰富的、包含上下文和位置信息的连续表示，再到最终的概率输出。

这个最小化的实现是你深入理解更复杂模型（如 BERT、GPT）的绝佳起点。你可以基于此进行扩展，例如：

构建解码器：添加一个解码器部分，实现一个完整的编码器-解码器模型，用于机器翻译或文本摘要任务。
增加层数：通过增加 NUM_ENCODER_LAYERS 来加深模型，观察其性能变化。
训练与微调：在一个真实的数据集上从头开始训练这个模型，或者加载预训练权重进行微调，来解决实际的 NLP 问题。

希望这篇指南能为你点燃探索 Transformer 内部世界的火花，并激励你回到 Happy-LLM 的广阔天地中，继续你的大语言模型学习之旅。