从零实现Transformer：核心模块详解与优化技巧

1. Transformer 架构概述

Transformer 架构自2017年由Vaswani等人提出以来，已成为自然语言处理领域的基石模型。与传统RNN和CNN不同，Transformer完全基于自注意力机制，能够并行处理序列数据并捕获长距离依赖关系。在CS336课程作业中，我们需要从零开始实现一个完整的Transformer模型，这要求我们深入理解每个组件的数学原理和实现细节。

现代大型语言模型（如GPT、LLaMA等）都基于Transformer架构，但通常会进行一些改进。本次作业实现的版本包含以下核心创新：

• 使用RMSNorm替代传统LayerNorm
• 采用SwiGLU作为前馈网络激活函数
• 引入旋转位置编码(RoPE)
• 使用预归一化(Pre-Norm)的残差连接

这些改进使得模型在训练稳定性和表现力方面都有显著提升。下面我们将逐一拆解每个模块的技术细节。

2. 基础线性与嵌入模块

2.1 线性变换层（无偏置版本）

线性层是神经网络中最基础的组件，其数学表示为：

$$ y = xW^\top $$

其中：

• $x \in \mathbb{R}^{..., d_{in}}$ 是输入张量
• $W \in \mathbb{R}^{d_{out}, d_{in}}$ 是权重矩阵
• $y \in \mathbb{R}^{..., d_{out}}$ 是输出张量

在实现时需要注意：

1. 权重初始化采用截断正态分布：
   $$ W \sim \mathcal{N}(0, \frac{2}{d_{in} + d_{out}}) $$
   并截断到$[-3\sigma, 3\sigma]$范围内

2. 现代Transformer实现通常省略偏置项，这可以：

   • 减少参数数量
   • 提高计算效率
   • 与LayerNorm/RMSNorm更好地配合

提示：在实际编码时，可以使用nn.Linear并设置bias=False来创建无偏置线性层。初始化权重时要注意保持方差稳定，防止梯度爆炸或消失。

2.2 词嵌入层

词嵌入层将离散的token ID映射到连续的向量空间：

$$ \text{Embedding}(i) = W_e[i,:] $$

其中：

• $i$ 是token ID（整数）
• $W_e \in \mathbb{R}^{vocab_size, d_{model}}$ 是嵌入矩阵
• 输出维度为$d_{model}$（通常256-4096）

关键细节：

1. 嵌入矩阵通常与最后的线性输出层共享权重（节省参数）
2. 需要对嵌入进行缩放（乘以$\sqrt{d_{model}}$），防止初始阶段注意力分数过大
3. 现代模型通常不对嵌入层使用偏置

2.3 RMSNorm（均方根归一化）

RMSNorm是LayerNorm的简化版本，计算更高效：

$$ \text{RMSNorm}(a_i) = \frac{a_i}{\text{RMS}(a)} \cdot g_i $$

其中：
$$ \text{RMS}(a) = \sqrt{\frac{1}{d_{model}}\sum_{i=1}^{d_{model}} a_i^2 + \epsilon} $$

与LayerNorm相比：

1. 去除了均值中心化
2. 仅使用均方根进行缩放
3. 保留可学习的增益参数$g_i$

优势：

• 计算量减少约20-30%
• 在大多数情况下性能相当
• 更适合大规模模型训练

实现要点：

python复制class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.scale = dim ** -0.5
        self.eps = eps
        self.g = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
        return x / (norm + self.eps) * self.g

3. 前馈网络模块

3.1 SiLU与GLU激活函数

SiLU（Sigmoid Linear Unit）激活函数：

$$ \text{SiLU}(x) = x \cd