大模型学习 （Datawhale_Happy-LLM）笔记5: 搭建一个 Transformer

大模型学习 （Datawhale_Happy-LLM）笔记5: 搭建一个 Transformer

搭建 Transformer 的核心组件总结

1. 基础功能模块

• 自注意力机制：通过 QKV 矩阵计算序列内依赖关系，公式为：
  $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
• 多头注意力：将特征拆分为多个子空间并行计算注意力，增强模型表达能力
• 位置编码：通过正弦余弦函数为序列添加位置信息，解决 Transformer 无序列感知问题

2. 核心网络层

• 层归一化 (LayerNorm)：对每个样本的所有特征维度归一化，公式为：
  $$\text{LayerNorm}(x)=\alpha \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$
• 前馈神经网络 (MLP)：由 Linear+ReLU+Linear 构成，增强模型非线性表达能力

3. 编码器与解码器架构

• 编码器层：由 LayerNorm+MultiHeadAttention+残差连接 和 LayerNorm+MLP+残差连接 组成
• 解码器层：比编码器多一个 掩码多头注意力，用于避免预测时看到未来信息
• 堆叠结构：多个编码器/解码器层堆叠形成深度网络

4. 输入输出处理

• 嵌入层：将离散token转换为连续向量，与位置编码相加后输入网络
• 输出层：通过线性层将特征映射到词表空间，用于生成概率分布

5. 关键技术点

• 残差连接：解决深度网络训练梯度消失问题
• Dropout：随机丢弃神经元，防止过拟合
• 掩码机制：在解码器中屏蔽未来位置，确保预测时的因果关系

完整 Transformer 架构流程图

输入序列 ──→ 嵌入层 + 位置编码 ──→ 编码器序列 ──→ 编码器输出
                                        │
                                        ↓
                                  解码器序列（含掩码）
                                        │
                                        ↓
                                  线性层 + Softmax ─→ 输出序列

多头自注意力模块

class ModelArgs:
    def __init__(self):
        self.n_embed = 256
        self.n_head = 4
        self.head_dim = self.n_embed // self.n_heads  # 显式定义head维度
        self.dropout = 0.1
        self.max_seq_len = 512
        self.n_layers = 6
        self.vocab_size = None  # 添加词表大小
        self.block_size = None  # 添加最大序列长度

class MultiHeadAttention(nn.Module):
    
    def __init__(self, args: ModelArgs, is_causal=False):
        # 构造函数
        # args: 配置对象
        super().__init__()
        # 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵
        assert args.n_embed % args.n_head == 0
        # 模型并行处理大小，默认为 1
        model_parallel_size = 1
        # 本地计算头数， 等于总头数除以模型并行处理大小
        self.n_local_heads = args.n_heads // model_parallel_size
        # 每个头的维度，等于模型维度除以头的总数
        self.head_dim = args.dim // args.n_heads

        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
        # 这⾥通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        # 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
        # 注意力的dropout
        self.attn_dropout = nn.Dropout(args.dropout)
        # 残差连接的 dropout
        self.resid_dropout = nn.Dropout