基于 PyTorch 从零构建与训练大型语言模型入门指南（上）

步骤 4：输入嵌入和位置编码

这一步进行输入嵌入和位置编码的处理。

首先，输入嵌入层负责将步骤 2 生成的标记 ID 序列转换为词汇表中的索引，并为每个标记生成一个 512 维的嵌入向量。

这个向量能够捕捉标记的深层语义特征。在多维空间中，相似的实体如狗和猫的向量会彼此接近，而与学校、家等不相似实体的向量则相隔较远。

其次，位置编码解决了 Transformer 架构在并行处理序列时可能忽略词序的问题。通过给每个标记的 512 维嵌入向量添加位置信息，保证模型能够理解词序对句子含义的影响。

具体来说，采用正弦和余弦函数对每个维度进行编码，其中正弦应用于偶数维度，余弦应用于奇数维度。这样，每个标记的嵌入向量不仅包含了其语义信息，还包含了其在句子中的位置信息。

import math
import torch.nn as nn

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__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)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

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

    def forward(self, x):
        x = self.embedding(x) * math.sqrt(self.embedding.embedding_dim)
        return self.pos_encoding(x)

步骤 5：多头注意力块

Transformer 模型的精髓在于自注意力机制，它赋予模型动态理解上下文的能力。而多头自注意力则进一步将这一能力细分，让模型能够同时从多个角度捕捉信息，从而更全面地理解句子。

如果熟悉矩阵乘法，掌握多头注意力机制其实相当简单。

首先，我们会从步骤 4 得到的编码输入创建三份副本：Q（查询）、K（键）、V（值）。这些副本将作为自注意力计算的基础。

随后，将 Q、K、V 分别与各自的权重矩阵进行矩阵乘法，这些权重矩阵将初始化为随机值，并在训练过程中不断更新。这一步骤引入了可学习的参数，帮助模型更好地捕捉信息。

按照论文中的设定，我们将使用 8 个头来进行多头注意力的计算。这意味着，每个经过矩阵乘法得到的查询、键、值向量都将被分割成 8 份，每份的维度为 d_k = d_model / num_heads。

接下来，每个查询向量将与序列中所有键向量的转置进行点积运算，得到注意力分数，这些分数反映了标记间的相似度。为避免模型过度关注高分数或忽略低分数，我们通过除以 d_k 的平方根来规范化这些分数。

在应用 softmax 函数之前，如果存在编码掩码，我们会将其与注意力分数结合，确保模型不会受到未来时间步的影响。Softmax 函数将这些分数转换为概率分布，然后这些概率将与相应的值向量相乘，得到每个头的输出。

最终，我们将 8 个注意力头的输出合并，并通过输出权重矩阵 W_o 进一步处理，得到多头自注意力的最终结果。这个结果能够综合考虑单词在句子中的不同上下文含义。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def split_heads(self, x, batch_size):
        x = x.view(batch_size, -1, self.num_heads, self.head_dim)
        return x.permute(0, 2, 1, 3)

    def forward(self, q, k, v, mask=None):
        batch_size = q.shape[0]

        Q = self.split_heads(self.W_q(q), batch_size)
        K = self.split_heads(self.W_k(k), batch_size)
        V = self.split_heads(self.W_v(v), batch_size)

        score = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if mask is not None:
            score = score.masked_fill(mask == 0, -1e9)

        attention = torch.softmax(score, dim=-1)
        output = torch.matmul(attention, V)

        output = output.permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.d_model)
        return self.W_o(output)

以上步骤构成了 Transformer 模型的核心组件。在实际训练中，还需要结合损失函数（如 CrossEntropyLoss）和优化器（如 AdamW）进行端到端的训练。后续章节将详细介绍模型的整体搭建与训练流程。