1.2 线性层的数学本质

线性层的数学公式极其简洁：

$\mathbf{y} = \mathbf{x} \mathbf{W}^\top + \mathbf{b}$

其中：

• $\mathbf{x}$ 是输入向量（维度 d_in）
• $\mathbf{W}$ 是权重矩阵（维度 d_out × d_in）
• $\mathbf{b}$ 是偏置向量（维度 d_out）

如果省略偏置（代码中的实现），则简化为： $\mathbf{y} = \mathbf{x} \mathbf{W}^\top$

从信息流的角度看，线性层做的事情是：将输入的每个特征，通过加权求和的方式，映射到输出空间的每个维度。$W_{ij}$ 控制了输入特征 $j$ 对输出特征 $i$ 的贡献强度。

1.3 权重矩阵的形状

在PyTorch标准实现中，线性层的权重形状为 (out_features, in_features)：

1
weight_shape = (out_features, in_features)
2
self.weight = nn.Parameter(torch.empty(weight_shape, device=device, dtype=dtype))

这意味着在矩阵乘法中，输入x的形状是 (..., d_in)，权重W的形状是 (d_out, d_in)，输出为 (..., d_out)。

前向传播时，采用einsum表示这个操作：

1
def forward(self, x: torch.Tensor) -> torch.Tensor:
2
    return einsum(x, self.weight, '... d_in, d_out d_in -> ... d_out')

einsum（爱因斯坦求和约定）的写法清晰地表达了：输入在最后一维d_in与权重的第二维d_in进行点乘，结果输出到d_out维度。

1.4 权重初始化

模型参数的初始化方式对训练的稳定性有决定性影响。代码中使用了截断正态分布（Truncated Normal），并采用了特定于“mishmash”架构的初始化策略：

1
sigma = math.sqrt(2.0 / (in_features + out_features))
2
torch.nn.init.trunc_normal_(self.weight, mean=0.0, std=sigma, a=-3 * sigma, b=3 * sigma)

• 使用截断正态的原因：普通正态分布会产生极端的离群值（如超过3倍标准差的权重），这可能导致前向传播时激活值过大或过小。截断正态将权重限制在 [-3σ, 3σ] 范围内，有效控制了这种风险。
• 标准差是 sqrt(2/(d_in + d_out))的原因：这是针对“mishmash”架构设计的初始化策略。对于标准线性层，常见的是Xavier初始化（1/sqrt(d_in)或sqrt(2/(d_in + d_out))），这里使用的是后者的一种变体。核心思想是保持输入和输出的方差一致，从而避免梯度消失或爆炸。

1.5 从线性层到整个Transformer

线性层是Transformer中最密集的参数来源。以GPT-2 124M为例，其线性层参数占总参数的95%以上。理解线性层的参数规模计算方法如下：

• 一个 d_in × d_out 的权重矩阵有 d_in × d_out 个参数
• 在Transformer中，注意力模块的Q、K、V、O投影各贡献 d_model × d_model 参数
• FFN模块的线性层贡献 d_model × d_ff + d_ff × d_model 参数

2.1 嵌入的本质是查表

词嵌入（Word Embedding） 做的事情很简单：将离散的token ID映射为稠密的连续向量。在代码层面，嵌入层本质上是一个查找表：

1
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
2
    return self.embedding_matrix[token_ids]

给定一个形状为 (batch_size, seq_len) 的token ID张量，返回形状为 (batch_size, seq_len, embedding_dim) 的向量张量。

2.2 嵌入矩阵的维度设计

嵌入矩阵的形状为 (num_embeddings, embedding_dim)：

1
embedding_shape = (num_embeddings, embedding_dim)
2
self.embedding_matrix = nn.Parameter(torch.empty(embedding_shape, device=device, dtype=dtype))

• num_embeddings = 词汇表大小（vocab_size），即有多少种不同的token
• embedding_dim = 嵌入向量的维度，也就是Transformer中的 d_model

比如第一篇博客训练的BPE分词器有18,017个token，如果 d_model=768，嵌入矩阵就是 (18017, 768)，约1,380万参数。

2.3 嵌入层的初始化

代码中用标准差为1的截断正态分布初始化嵌入矩阵：

1
std = 1.0
2
torch.nn.init.trunc_normal_(self.embedding_matrix, std=std, a=-3 * std, b=3 * std)

这是一个相对温和的初始化——标准差为1，截断在±3范围内。在实际训练中，嵌入层通常配合学习率调度器使用（通常使用稍高的学习率），因为嵌入层需要快速适应训练数据中的token分布。

2.4 嵌入层与线性层的联系与区别

从数学角度看，嵌入层是线性层的特例：当输入是独热（one-hot）向量时，线性变换 one_hot × Embedding_Matrix 就是查表操作。但直接使用独热向量在计算上不可行（词汇表高达数万），所以PyTorch用查表作为高效的替代方案。

两者在梯度流上的表现相同：在反向传播中，token ID对应的那一行嵌入向量会接收到梯度，其他行不变。

3.1 归一化层的作用

深度神经网络在训练中会遇到一个棘手的问题：内部协变量偏移（Internal Covariate Shift） —— 网络各层输入的分布会随着前面层参数的变化而变化，导致训练不稳定，收敛缓慢。

归一化层的作用就是将每层的输入拉回到一个稳定的分布范围，从而让梯度流动更加平稳。在Transformer中，归一化层通常放在残差连接之后、每个子层之前（Pre-Norm结构）。

3.2 LayerNorm vs RMSNorm

Layer Normalization（LayerNorm） 的公式为：

$\text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$

其中 $\mu$ 和 $\sigma^2$ 分别是均值和方差，$\gamma$ 和 $\beta$ 是可学习的缩放和偏移参数。

RMSNorm 的公式为：

$\text{RMSNorm}(x) = \frac{x}{\sqrt{\text{RMS}(x)^2 + \epsilon}} \cdot \gamma$

其中 $\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2}$。

核心区别在于：

1. RMSNorm去掉了均值中心化（$\mu$）。它只除以根均方（RMS），不减去均值。这意味着RMSNorm对输入的整体偏移更敏感，但好处是计算开销更低。
2. RMSNorm去掉了可学习的偏置参数（$\beta$）。相比LayerNorm，RMSNorm少了一倍的参数。

3.3 现代LLM使用RMSNorm的原因

RMSNorm在LLM领域的流行始于Llama系列模型。其优势在于：

1. 计算效率更高：计算均方根比计算均值和方差的组合要少一次求和操作。在长序列训练中，这带来的加速不可忽视。
2. 效果相当：大量实验表明，在Transformer架构中，RMSNorm与LayerNorm的下游表现几乎没有差异。即使有细微差距，也可以通过调整模型宽度或训练步长来弥补。
3. 参数更少：少了$\beta$参数，模型更紧凑。

3.4 代码实现

1
def forward(self, x: torch.Tensor) -> torch.Tensor:
2
    input_dtype = x.dtype
3
    x = x.to(torch.float32)  # 先转为float32，避免数值不稳定
4

5
    variance = x.pow(2).mean(-1, keepdim=True)
6
    x = x * torch.rsqrt(variance + self.eps)
7

8
    return (self.weight * x).to(input_dtype)

关键细节：计算归一化时先将输入转换为float32。这是因为在混合精度训练（如bfloat16或float16）中，中间计算可能产生下溢或溢出。rsqrt（倒数平方根）在低精度下对极大值或极小值特别敏感，所以先在float32下计算，再转回原始精度。

4.1 经典FFN结构回顾

Transformer的原始设计中，前馈网络（FFN）是一个简单的两层线性网络加ReLU激活：

$\text{FFN}(x) = \text{ReLU}(x W_1 + b_1) W_2 + b_2$

这个结构简单有效，但存在一个问题：ReLU在负半区输出恒为零，导致“神经元死亡”现象——部分神经元在训练中可能永远不被激活。

4.2 SwiGLU的设计哲学

SwiGLU（Swish-Gated Linear Unit）是近年来LLM架构中最重要的改进之一，被Llama、Mistral、PaLM等模型广泛采用。

SwiGLU的公式为：

$\text{SwiGLU}(x) = W_2 \cdot (\text{Swish}(W_1 x) \odot W_3 x)$

其中 $\text{Swish}(x) = x \cdot \sigma(x)$，$\sigma$ 是sigmoid函数，$\odot$ 是逐元素相乘。

这个结构包含三个权重矩阵：$W_1$、$W_2$、$W_3$。这里的 $W_1$ 和 $W_3$ 并行作用在输入上，$W_1$ 的输出经过Swish激活后与 $W_3$ 的输出逐元素相乘，再通过 $W_2$ 映射回 d_model。

4.3 SwiGLU vs FFNSiLU：参数量对比

为了公平比较，在参数量相当的前提下对比两种设计：

• 经典FFN（ReLU）：两个权重矩阵，如果隐藏维度为 d_ff，参数量为 d_model × d_ff + d_ff × d_model = 2 × d_model × d_ff
• SwiGLU：三个权重矩阵，参数量为 3 × d_model × d_ff（$W_1$、$W_2$、$W_3$各一个）

如果希望SwiGLU的参数量与经典FFN持平，需要将 d_ff 缩小为 2/3。在FFN中通常取 d_ff = 4 × d_model，因此在SwiGLU中通常取 d_ff = 8/3 × d_model，确保能被64整除。

代码中提供了一个可选开关：

1
class SwiGLU(nn.Module):
2
    def __init__(self, d_model, d_ff, use_swiglu=True):
3
        self.w1 = LinearModule(d_model, d_ff)
4
        self.w2 = LinearModule(d_ff, d_model)
5
        if use_swiglu:
6
            self.w3 = LinearModule(d_model, d_ff)

当 use_swiglu=True 时，$W_3$ 被创建，前向传播变为：

1
return self.w2(self.silu(self.w1(x)) * self.w3(x))

当 use_swiglu=False 时，回退为标准的 FFNSiLU（即用SiLU替代ReLU）：

1
return self.w2(self.silu(self.w1(x)))

4.4 SwiGLU效果更好的原因

SwiGLU的核心优势在于门控机制（Gating）。$W_1 x$ 经过Swish激活后充当“门”，$W_3 x$ 充当“值”，门控选择性地让哪些信息通过。这类似于LSTM中的“输入门”设计。

这种结构赋予了FFN更强的非线性表达能力。在LLM的scaling law实验中，SwiGLU被证明能在大约相同参数量下实现比ReLU-FFN更低的困惑度，这也是当今几乎所有前沿开源LLM都采用SwiGLU的原因。

5.1 位置编码的必要性

Transformer的自注意力机制本身是置换不变（Permutation Invariant）的 —— 把“我 爱 你”三个词打乱顺序，得到的注意力计算结果和原顺序是一样的。这显然不符合语言的基本特性：顺序决定意义。

因此，模型需要一种方式来注入位置信息。传统的位置编码方案包括：

• 绝对位置编码（APE）：为每个位置赋予一个固定的编码向量，与词嵌入相加
• 相对位置编码（RPE）：在注意力计算时，直接编码两个位置之间的相对距离

5.2 RoPE的核心思想

RoPE（Rotary Positional Embedding） 由苏剑林提出，是目前最主流的位置编码方案，被Llama、PaLM、Qwen等模型采用。

RoPE的核心思想是 —— 不额外添加位置向量，而是通过旋转矩阵“转动”词向量，使得词向量的方向携带了位置信息。

1. 从二维平面看本质

在二维平面上，将向量旋转角度 $\theta$ 的矩阵为：

$\text{Rot}(\theta) = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}$

对于第 $m$ 个位置的token，其查询向量 $q_m$ 被旋转角度 $m\theta$：$q'_m = \text{Rot}(m\theta) q_m$，键向量同理 $k'_n = \text{Rot}(n\theta) k_n$。

2. 核心性质推导（相对距离的自动生成）

RoPE编码相对位置的关键在于旋转矩阵具有正交同态性质：$\text{Rot}(a)^\top \text{Rot}(b) = \text{Rot}(b - a)$。

计算旋转后查询和键的注意力分数（点积）：$(q'_m)^\top k'_n = [\text{Rot}(m\theta) q_m]^\top [\text{Rot}(n\theta) k_n]$

由于转置性质 $(AB)^\top = B^\top A^\top$，且旋转矩阵满足 $\text{Rot}(m\theta)^\top = \text{Rot}(-m\theta)$，代入得：

$(q'_m)^\top k'_n = q_m^\top \text{Rot}(m\theta)^\top \text{Rot}(n\theta) k_n = q_m^\top \text{Rot}((n-m)\theta) k_n$

最终结果中，旋转角度直接变成了 $(n-m)\theta$。也就是说，注意力分数只依赖于两个token之间的相对距离 $n-m$，与它们各自的绝对位置 $m$、$n$ 无关。这就把绝对位置编码巧妙地转化成了相对位置编码。

3. 扩展到高维空间（复数视角）

现实中的词向量维度 $d_k$ 远大于2。RoPE的处理方式是将 $d_k$ 维空间划分为 $d_k/2$ 个二维子空间，每个子空间使用不同的旋转频率 $\theta_i$。

从复数视角看，这相当于把每对维度 $(x_{2i}, x_{2i+1})$ 视为一个复数 $z = x_{2i} + i x_{2i+1}$，旋转操作即是乘以单位复数 $e^{i m \theta_i}$：

$z'_i = z_i \cdot e^{i m \theta_i}$

两个维度组成一个二维子空间，构成一个“复数对”。如果 $d_k$ 是奇数，最后一维无法配对成复数，就无法进行旋转操作。最终，整个高维旋转矩阵 $\mathcal{R}_m$ 是一个分块对角矩阵（Block Diagonal Matrix）：