LLM TREND

大模型技术趋势分享：LLM、RAG、Prompt、Agent、Fine-tuning等

LLM

从零开始构建大型语言模型：Transformer架构详解

By llmtrend #Transformer

大型语言模型（LLM）的崛起，离不开Transformer架构的突破性创新。本文将深入探讨Transformer架构的核心组件，并通过从零开始构建一个简化的Transformer模型，帮助读者理解这些卓越模型背后的运作机制。本文将深入探讨Transformer架构，并重点剖析其自注意力机制、残差连接、层归一化等关键概念，以及decoder-only架构的原理和优势。

Transformer架构：LLM成功的基石

Transformer架构由Vaswani等人在2017年的论文“Attention Is All You Need”中提出，彻底改变了自然语言处理领域。相较于传统的循环神经网络（RNN）和卷积神经网络（CNN），Transformer避免了循环依赖，能够并行处理序列数据，极大地提高了训练效率。此外，其核心的自注意力机制能够捕捉序列中任意两个位置之间的依赖关系，更好地理解上下文信息，使得Transformer成为构建大型语言模型（LLM）的首选架构。Google的BERT、OpenAI的GPT系列以及百度的文心一言等知名LLM，均基于Transformer架构。

Decoder-Only架构：自回归生成的利器

现代LLM，如GPT系列，普遍采用decoder-only架构。这种架构专为自回归生成任务设计，即根据之前的token逐个预测下一个token。

自回归的本质

自回归生成过程可以简单概括为：

  1. 输入: “今天天气”
  2. 预测: “晴朗” (基于上下文)
  3. 下一个输入: “今天天气晴朗”
  4. 下一个预测: “，” (基于扩展的上下文)
  5. 重复: 直到生成结束或满足停止条件

这种序列化的预测方式需要decoder架构，及其特有的因果掩码 (causal masking)——模型只能关注之前的token，而不能“偷看”未来的token。这保证了模型在生成文本时，严格依据上下文进行预测，避免了信息泄露。

序列长度与并行预测

虽然推理过程是序列化的，但训练过程中，decoder-only Transformer能够并行地做出多个预测。例如，如果序列长度为10个token，那么在一次前向传播中，可以同时做出10个预测：

  • 位置 0 → 预测 token 1
  • 位置 1 → 预测 token 2
  • …
  • 位置 9 → 预测 token 10

这种并行预测极大地提高了训练效率，使得训练大规模LLM成为可能。

核心组件详解

从零开始构建Transformer架构，需要理解其核心组件。下面我们逐一介绍，并结合文章提供的代码进行分析。

1. Embedding层：文本到向量的桥梁

Embedding层负责将离散的token转换为连续的向量表示。文章中InputEmbeddings类实现了这个功能：

class InputEmbeddings(nn.Module):
    def __init__(self, embedding_dim: int, vocab_size):
        # d_model: dimension of vector representing each word
        # vocab_size: how many words there are in
        super(InputEmbeddings, self).__init__()
        self.embedding_dim = embedding_dim
        self.vocab_size = vocab_size
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.pred_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
    def forward(self, x):
        pred_token = self.pred_token.expand(x.size(0), 1, self.embedding_dim)
        x =self.embedding(x)
        #x = torch.cat([x, pred_token], dim=1)
        return x * math.sqrt(self.embedding_dim ** 0.5)
        # math.sqrt(self.d_model**0.5): The last sentence of 3.4 Embeddings and Softmax paragraph

nn.Embedding是PyTorch提供的Embedding层，它将每个token映射到一个embedding_dim维的向量。vocab_size指定了词汇表的大小，决定了Embedding层需要学习的向量数量。forward函数接收token序列作为输入，返回对应的向量表示。文章代码中还包括对Embedding向量进行缩放的操作，这有助于稳定训练过程。

例如，如果vocab_size为10000，embedding_dim为512，那么Embedding层将学习10000个512维的向量，每个向量对应一个唯一的token。

2. 位置编码：序列信息的补充

由于Transformer架构没有循环结构，无法直接感知token在序列中的位置信息。位置编码 (Positional Encoding) 用于将位置信息添加到Embedding向量中。文章中PositionalEncoding类实现了位置编码的功能：

class PositionalEncoding(nn.Module):
    def __init__(self, embedding_dim: int, seq_len: int, dropout: float):
        super(PositionalEncoding, self).__init__()
        #learnable positional encoding. Add +1 to seq_len because of prediction token
        self.pe = nn.Embedding(seq_len, embedding_dim)
    def forward(self, x):
        # x: (batch_size, seq_len+1, embedding_dim)
        positions = torch.arange(0, x.size(1), device=x.device).unsqueeze(0)
        pos_embed = self.pe(positions)
        return x + pos_embed

nn.Embedding在这里被用来学习位置信息。seq_len指定了序列的最大长度，决定了位置编码层需要学习的向量数量。forward函数接收Embedding向量作为输入，将对应位置的位置编码向量加到Embedding向量上，从而将位置信息注入到模型中。

例如，如果seq_len为1024，embedding_dim为512，那么位置编码层将学习1024个512维的向量，每个向量对应一个序列中的一个位置。

3. 层归一化：稳定训练的保障

层归一化 (Layer Normalization) 是一种常用的正则化技术，用于稳定训练过程，尤其是在深层网络中。文章中LayerNormalization类实现了层归一化的功能：

class LayerNormalization(nn.Module):
    def __init__(self, embedding_dim: int, eps: float = 1e-6):
        # eps: epsilon. Its for preventiong vanishing the denomiator while normalization.
        super(LayerNormalization, self).__init__()
        # alpha and bias: to introduce some fluctuations in data
        self.eps = eps
        self.alpha = nn.Parameter(torch.ones(embedding_dim))  # per-feature scale
        self.bias = nn.Parameter(torch.zeros(embedding_dim))  # per-feature bias
    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True, unbiased=False)
        return self.alpha * (x - mean) / (std + self.eps) + self.bias

层归一化对每个样本的每个特征进行归一化，使其均值为0，方差为1。eps是一个很小的数值，用于防止除以0的情况。alpha和bias是可学习的参数，用于调整归一化后的数值范围。forward函数接收输入向量，返回归一化后的向量。

4. 前馈网络：特征的非线性变换

前馈网络 (Feed Forward Network, FFN) 用于对自注意力机制的输出进行非线性变换，增强模型的表达能力。文章中FeedForwardNetwork类实现了前馈网络的功能：

class FeedForwardNetwork(nn.Module):
    def __init__(self, embedding_dim: int = 512, ffn_hidden_dim: int = 3072, dropout: float = 0):
        super(FeedForwardNetwork, self).__init__()
        self.embedding_dim = embedding_dim
        self.ffn_hidden_dim = ffn_hidden_dim
        self.layer1 = nn.Linear(self.embedding_dim, self.ffn_hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()
        self.layer2 = nn.Linear(self.ffn_hidden_dim, self.embedding_dim)
    def forward(self, x):
        # (batch, seq_len, d_model) -> (batch, seq_len, dff) -> (batch, seq_len, d_model)
        return self.layer2(self.dropout(self.activation(self.layer1(x))))

前馈网络由两个线性层组成，中间有一个激活函数。embedding_dim是输入和输出的维度，ffn_hidden_dim是隐藏层的维度。dropout是一种常用的正则化技术，用于防止过拟合。GELU (Gaussian Error Linear Unit) 是一种常用的激活函数，在Transformer架构中表现良好。forward函数接收输入向量，返回经过非线性变换后的向量。

5. 多头自注意力机制：捕捉序列依赖

多头自注意力 (Multi-Head Self-Attention) 机制是Transformer架构的核心创新。它能够并行地计算序列中任意两个位置之间的依赖关系，更好地理解上下文信息。文章中MultiHeadAttentionBlock类实现了多头自注意力机制的功能：

class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, embedding_dim: int, num_heads: int, dropout: float):
        # d_model: Dimension of word embedding vector
        # h: number of heads which will divide the embedding dimension
        super(MultiHeadAttentionBlock, self).__init__()
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        assert (embedding_dim % num_heads == 0), "Word embedding must be divisible by number of heads (embedding_dim / num_heads)"
        self.d_k = self.d_q = self.d_v = self.embedding_dim // self.num_heads
        # To make it more readable.
        self.W_q = nn.Linear(embedding_dim, embedding_dim, bias=True)# , bias=False)
        self.W_k = nn.Linear(embedding_dim, embedding_dim, bias=True)# , bias=False)
        self.W_v = nn.Linear(embedding_dim, embedding_dim, bias=True)# , bias=False)
        # Only weights, biases werent mentioned in the paper
        # also OK to add bias. Your choice
        self.W_o = nn.Linear(embedding_dim, embedding_dim, bias=True)# , bias=False)
        self.dropout = nn.Dropout(dropout)
    @staticmethod
    def attention(q, k, v, mask, d_k, embedding_dim, dropout):
        attention_scores = torch.einsum("nqhd,nkhd->nhqk", [q, k])
        if mask:
            mask_tmp = torch.triu(torch.ones(q.size(1), q.size(1)), diagonal=1).to(q.device)
            attention_scores = attention_scores.masked_fill(mask_tmp == 0, float("-1e20"))
        attention = torch.softmax(attention_scores / (d_k ** (1 / 2)), dim=3)
        if dropout is not None:
            attention = dropout(attention)
        attention_result = torch.einsum("nhql,nlhd->nqhd", [attention, v]).reshape(
            v.size(0), v.size(1), embedding_dim
        )
        return attention_result
    def forward(self, q, k, v, mask):
        query = self.W_q(q)  # (batch, seq_len, d_model)
        key = self.W_k(k)  # (batch, seq_len, d_model)
        value = self.W_v(v)  # (batch, seq_len, d_model)
        # Split into heads
        query = query.reshape(q.size(0), q.size(1), self.num_heads, self.d_q)  # (batch, seq_len, h, d_q)
        key = key.reshape(k.size(0), k.size(1), self.num_heads, self.d_k)  # (batch, seq_len, h, d_k)
        value = value.reshape(v.size(0), v.size(1), self.num_heads, self.d_v)  # (batch, seq_len, h, d_v)
        out = self.W_o(MultiHeadAttentionBlock.attention(query, key, value, mask, self.d_k, self.embedding_dim, self.dropout))
        return out

多头自注意力机制将输入向量分别投影到Query (Q)、Key (K) 和 Value (V) 三个不同的空间。然后，计算Q和K之间的注意力分数，通过softmax函数归一化，得到注意力权重。最后，将注意力权重应用于V，得到自注意力的输出。num_heads指定了attention头的数量，每个头负责学习不同的依赖关系。forward函数接收Q、K和V作为输入，返回自注意力的输出。

6. 残差连接：深度网络的福音

残差连接 (Residual Connection) 是一种用于训练深层网络的技巧，它允许梯度直接从输出层传播到输入层，避免梯度消失的问题。文章中ResidualConnection类实现了残差连接的功能：

class ResidualConnection(nn.Module):
    def __init__(self, embedding_dim: int, dropout: float):
        super(ResidualConnection, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(embedding_dim)
    def forward(self, x, sublayer):
        # return self.norm(x + self.dropout(sublayer(x)))
        return x + self.dropout(sublayer(self.norm(x)))

残差连接将输入向量直接加到子层的输出上。dropout用于防止过拟合。norm是层归一化层，用于稳定训练过程。forward函数接收输入向量和子层作为输入，返回经过残差连接后的向量。

7. Decoder Block：组件的有机组合

Decoder Block将自注意力机制、前馈网络、层归一化和残差连接组合在一起，形成一个完整的decoder层。文章中DecoderBlock类实现了Decoder Block的功能：

class DecoderBlock(nn.Module):
    def __init__(self,
                 self_attention_block: MultiHeadAttentionBlock,
                 feed_forward_network: FeedForwardNetwork,
                 embedding_dim: int,
                 dropout: float):
        super(DecoderBlock, self).__init__()
        self.self_attention_block = self_attention_block
        self.feed_forward_network = feed_forward_network
        self.residual_connections = nn.ModuleList([ResidualConnection(embedding_dim, dropout) for _ in range(2)])
    def forward(self, x, input_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, input_mask))
        x = self.residual_connections[1](x, self.feed_forward_network)
        return x

Decoder Block首先对输入向量进行自注意力计算，然后进行残差连接和层归一化。接着，对输出向量进行前馈网络计算，再次进行残差连接和层归一化。forward函数接收输入向量和掩码作为输入，返回经过decoder层处理后的向量。

8. Decoder：多层Decoder Block的堆叠

Decoder由多个Decoder Block堆叠而成，用于逐步提取输入序列的特征。文章中Decoder类实现了Decoder的功能：

class Decoder(nn.Module):
    def __init__(self, layers: nn.ModuleList, embedding_dim: int):
        super(Decoder, self).__init__()
        self.layers = layers
        self.norm = LayerNormalization(embedding_dim)
    def forward(self, x, input_mask):
        for layer in self.layers:
            x = layer(x, input_mask)
        return self.norm(x)

Decoder将输入向量依次传递给每个Decoder Block，最后进行层归一化。forward函数接收输入向量和掩码作为输入，返回经过decoder处理后的向量。

9. Projection Layer：输出概率的映射

Projection Layer用于将decoder的输出映射到词汇表空间，得到每个token的概率。文章中ProjectionLayer类实现了Projection Layer的功能：

class ProjectionLayer(nn.Module):
    def __init__(self, embedding_dim: int, vocab_size: int):
        super(ProjectionLayer, self).__init__()
        self.linear = nn.Linear(embedding_dim, vocab_size, bias=True)
    def forward(self, x):
        # Shape: (batch, seq_len, d_model) → (batch, seq_len, vocab_size)
        return self.linear(x)

Projection Layer由一个线性层组成，将decoder的输出向量映射到vocab_size维的向量，表示每个token的logits。forward函数接收decoder的输出向量作为输入，返回每个token的logits。通常，还需要对logits进行softmax归一化，得到每个token的概率。

10. Transformer：整体架构的集成

Transformer将Embedding层、位置编码层、Decoder和Projection Layer组合在一起，形成一个完整的Transformer模型。文章中Transformer类实现了Transformer的功能：

class Transformer(nn.Module):
    def __init__(self,
                 decoder_module: Decoder,
                 input_embedding: InputEmbeddings,
                 input_positional_encoding: PositionalEncoding,
                 projection_layer: ProjectionLayer
                 ):
        super(Transformer, self).__init__()
        self.decoder_module = decoder_module
        self.input_embedding = input_embedding
        self.input_positional_encoding = input_positional_encoding
        self.projection_layer = projection_layer
    def forward(self, input, input_mask):
        input = self.input_embedding(input)
        input = self.input_positional_encoding(input)
        input = self.decoder_module(input, input_mask)
        return self.projection_layer(input)

Transformer首先对输入序列进行Embedding和位置编码，然后将其传递给Decoder进行特征提取，最后通过Projection Layer得到每个token的概率。forward函数接收输入序列和掩码作为输入，返回每个token的概率。

关键设计决策

文章中提到了一些关键的设计决策，这些决策对Transformer模型的性能和训练稳定性至关重要。

1. Pre-Layer Normalization

文章采用Pre-Layer Normalization，即在自注意力机制和前馈网络之前应用层归一化。相比于在之后应用层归一化（Post-LN），Pre-LN具有以下优点：

  • 稳定深层网络的训练
  • 减少梯度消失问题
  • 往往可以避免学习率预热

2. GELU激活函数

文章使用GELU激活函数，而不是ReLU。GELU提供更平滑的梯度，在Transformer架构中表现更好，是现代LLM的标准选择。

3. 因果掩码

自注意力机制使用因果掩码，防止模型在训练过程中“作弊”，看到未来的token。这保证了模型只能基于过去的上下文进行预测。代码中通过torch.triu创建上三角矩阵作为掩码。

4. 残差连接

每个主要组件都包含残差连接，使得训练非常深的网络成为可能，提供梯度流动路径，提高收敛稳定性。

5. 权重绑定 (Weight Tying)

文章使用了权重绑定技术，将Projection Layer的权重与Embedding层的权重绑定。这可以减少模型参数的数量，并提高模型的泛化能力。权重绑定基于Embedding层和输出层都与词汇表相关，因此共享权重可以提高效率。

多头自注意力机制的深入理解

多头自注意力机制是Transformer架构的核心，值得深入理解。

自注意力如何工作

自注意力机制计算序列中所有token对之间的注意力权重：

  1. Query, Key, Value: 每个token被投影到三个不同的表示空间。
  2. Attention Scores: 计算Query和Key的转置矩阵乘积。
  3. Attention Weights: 将注意力分数通过softmax函数。
  4. Output: 对Value进行加权求和，权重为注意力权重。

为什么需要多头

多头自注意力机制并行地运行多个自注意力机制，每个头关注不同的方面。例如：

  • 头 1: 可能关注句法关系
  • 头 2: 可能捕捉语义相似性
  • 头 3: 可能识别长距离依赖
  • 头 4: 可能跟踪实体引用

这种并行处理允许模型同时捕捉多种类型的关系。

缩放和效率

注意力计算的复杂度随序列长度呈平方增长 (O(n²))，这就是为什么上下文长度是LLM的一个限制因素。上下文长度每翻一倍，注意力计算成本就会增加四倍。

模型训练动态和输出处理

在Transformer层处理后，我们得到一个隐藏状态序列。为了将其转换为预测，需要以下步骤：

  1. 提取最后一个隐藏状态: 从最后一层提取输出。
  2. 语言模型头: 应用线性投影到词汇表大小。
  3. Logits: 词汇表中每个token的原始预测分数。
  4. 概率分布: 应用softmax函数获得概率分布。

采样策略

在推理过程中，我们可以通过不同的采样方法来控制生成：

  1. 贪婪采样: 始终选择概率最高的token。简单但可能导致重复输出。
  2. Top-k采样: 从前k个最可能的token中进行采样。平衡质量和多样性。
  3. 温度采样: 在softmax之前，按温度缩放logits：
    • 低温 (< 1.0): 更集中，保守的输出
    • 高温 (> 1.0): 更随机，创造性的输出
  4. Nucleus (Top-p) 采样: 从累积概率超过p的最小token集中进行采样。提供动态词汇选择。

文章中使用了温度采样，通过调节温度参数来控制生成文本的随机性。

架构测试

文章中提供了测试Transformer架构的代码，包括参数数量检查和推理过程测试。通过测试，可以验证模型的实现是否正确，并检查输出是否符合预期。

参数数量检查

文章使用了count_parameters和parameter_report函数来检查模型的参数数量。这可以确保模型的参数数量与预期一致，避免实现错误。

推理过程检查

文章使用GPT2Tokenizer对文本进行编码，然后将其输入到Transformer模型中进行推理。通过观察模型的输出，可以验证模型是否能够生成文本，并检查生成文本的质量。

decoder-only架构的优势

decoder-only架构可能看起来比encoder-decoder模型更有限，但它实际上更通用：

  • 统一框架: 同一个架构处理理解和生成。
  • 缩放优势: 更简单的架构可以更好地扩展规模。
  • 迁移学习: 预训练模型可以在不同的任务中使用。
  • 涌现能力: 大型decoder-only模型会产生意想不到的能力。

现代LLM表明规模很重要：

  • 更多参数: 增加模型容量和能力。
  • 更多数据: 提高泛化和知识。
  • 更多计算: 支持训练更大、更强大的模型。

这种扩展带来了涌现现象，如少样本学习、推理，甚至一些常识理解。

挑战与考虑

二次注意力复杂度带来了挑战：

  • 内存使用: 随序列长度呈二次增长。
  • 计算成本: 限制了实际的上下文长度。
  • 推理速度: 自回归生成本质上是串行的。

研究人员正在开发解决方案：

  • 稀疏/线性/Flash Attention: 减少瓶颈部分的计算。
  • 混合专家模型: 增加模型容量，而无需成比例的计算增加。
  • 模型并行: 将大型模型分布在多个设备上。

总结与展望

本文详细介绍了Transformer架构的核心组件，以及从零开始构建一个简化的Transformer模型的过程。通过学习本文，读者可以深入理解Transformer架构的运作机制，并为后续学习更复杂的LLM打下坚实的基础。下一步，我们将重点关注Transformer模型的训练过程，包括优化策略、训练稳定性技术和评估指标。随着技术的不断发展，Transformer架构及其变体将在自然语言处理领域发挥越来越重要的作用。

By llmtrend

Related Post

LLM

本地部署 DeepSeek-R1：使用 Ollama 轻松驾驭大模型

llmtrend
LLM

Spring AI赋能：利用工具调用（Tool Calling）构建更强大的大语言模型应用

llmtrend
LLM

从预训练到策略优化：大型语言模型 (LLM) 如何实现与人类意图对齐

llmtrend

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

You Missed

LLM

从预训练到策略优化：大型语言模型 (LLM) 如何实现与人类意图对齐

LLM

Spring AI赋能：利用工具调用（Tool Calling）构建更强大的大语言模型应用

AI Tools

AI 内容捷径：为什么过度依赖 AI 写作会损害你的 WordPress 网站 SEO？

Agent

2025年值得关注的五大 Agentic AI框架：迎接自主智能新时代