CS336: Language Models From Scratch (Spring 2025)

本节主要讲了模型的架构设计和超参数选择。

1.Architecture

1.1Norm

pre-norm, post-norm, ‘double’-norm

自从GPT之后大都采用pre-norm，把layernorm层放到FFN、MHA层之前。

3-p1

prenorm和postnorm的效果一样好，而且不需要warm。更好的梯度反向传播，更少的spike。

现在有的模型还使用’double’-norm，即FFN、MHA层之前之后都有layernorm。


LayerNorm, RMSNorm

原始的transformer和早期模型使用LN，现在都改为使用RMSN。

LN：$y = \frac{x - \text{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} * \gamma + \beta$

其中 $\text{E}[x]$ 是均值，$\text{Var}[x]$ 是方差，$\epsilon$ 是防止分母为 0 的小量，$\gamma$（缩放因子）和 $\beta$​​（偏移因子）是可学习参数。


RMSN：$y = \frac{x}{\sqrt{\|x\|_2^2 + \epsilon}} * \gamma$​ 其中 $\|x\|_2^2$​ 是输入 x 的二范数平方，$\epsilon$​ 是防止分母为 0 的小量，$\gamma$​ 是可学习的缩放参数。

不减去均值，也不添加偏置项$\beta$​。


RMSN效果和LN一样好，而且更快。操作更少（无需计算平均值），参数更少（没有偏置项）。

3-p2

曾有研究表明，在模型运算中，矩阵乘法占用的flops达到99.8%，正则化的运算量只占到0.17%。从计算性能的角度看，norm没必要优化。但是内存开销也是一个重要的考量，该研究指出正则化所占的运行时间达到25.5%，在内存搬运上花了相当一部分时间，因此值得优化。


现有的大部分transformer模型都没有bias项，只进行矩阵乘法。reason：更稳定（原因未知）


1.2 Activations

ReLU、GeLU、SwiGLU、GeGLU

GLU（门控线性单元）现在得到广泛使用


ReLU（Rectified Linear Unit，修正线性单元）是深度学习中最常用的激活函数之一。

ReLU 的函数形式非常简单，数学定义为：$\text{ReLU}(x) = \max(0, x)$即：

  • 当输入 $x \geq 0$ 时，输出等于输入本身 $x$；
  • 当输入 $x < 0$ 时，输出为 0。
3-p3

ReLU的优点

ReLU 之所以能取代 Sigmoid、Tanh 等传统激活函数，成为主流选择，核心原因在于以下几点：

  1. 计算效率极高：传统激活函数（如 Sigmoid 的 $\sigma(x)=1/(1+e^{-x})$）需要指数运算，计算成本高；而 ReLU 仅需一个简单的阈值判断（x 是否大于 0），几乎没有计算开销，能显著加速模型训练。
  2. 缓解梯度消失问题：Sigmoid 的导数在 $|x|$ 较大时趋近于 0，导致梯度难以传递到浅层。ReLU 的导数在 (x>0) 时为 1，梯度可以 “无损” 地反向传播，避免了深层网络中梯度被逐层衰减的问题。
  3. 稀疏激活特性：ReLU 会将所有负输入 “抑制” 为 0，只有正输入被保留，这种 “稀疏性” 使得神经网络在处理数 据时，仅激活部分神经元（类似生物大脑的工作模式），减少了冗余计算，同时增强了模型的泛化能力（避免过拟合）。
梯度消失
反向传播时，若梯度经过多层传递后逐渐衰减至接近 0，导致浅层参数几乎无法更新，称为梯度消失。
典型场景：使用 Sigmoid 激活函数时，其导数在输入绝对值较大时接近 0（最大值仅 0.25），多层相乘后梯度快速趋近于 0。
影响：深层网络的浅层参数难以优化，模型无法学习到有效特征。

梯度爆炸
反向传播时，若梯度经过多层传递后急剧增大至非常大的值，导致参数更新幅度过大，模型训练不稳定（如损失值跳变、NaN 等），称为梯度爆炸。
典型场景：循环神经网络（RNN）处理长序列时，梯度可能随时间步累积而指数级增长；或权重初始化过大，导致梯度被放大。
影响：参数更新失控，模型难以收敛。

ReLU 的缺点

尽管优势显著，ReLU 仍存在一些不容忽视的缺陷：

  1. 死亡 ReLU 问题（Dead ReLU Problem）：当输入 (x < 0) 时，ReLU 的输出为 0，且导数也为 0。若在训练中，某个神经元的输入长期为负（例如，权重更新不当导致输入始终小于 0），该神经元的梯度将永远为 0，无法通过反向传播更新参数，最终 “永久死亡”，失去学习能力。
    • 常见诱因：学习率过大（导致权重更新幅度过大，输入被 “推到” 负区间）、数据预处理不当（输入分布偏移）等。
  2. 输出非对称，均值偏移：ReLU 的输出范围是 $[0, +\infty)$，非零均值（输出均值偏向正方向）。这可能导致后续层的输入分布逐渐偏向正值，影响梯度下降的稳定性（例如，梯度方向可能被 “挤压” 到特定区域）。

swish

Swish 由 Google 提出，结合了 ReLU 和 Sigmoid 的特性，定义为：$\text{Swish}(x) = x \cdot \sigma(\beta x)$

其中 $\sigma$ 是 Sigmoid 函数，$\beta$是可学习参数（或固定为 1）。

  • 优势：处处可导（无 ReLU 的 “折点”），负输入区域输出非零（避免死亡神经元），在深层网络（如 Transformer）中表现优异；
  • 缺陷：计算复杂度略高于 ReLU。

SwigLU

SwigLU 是一种结合了 Swish 激活函数和 GLU（Gated Linear Unit）门控机制的激活函数。

SwigLU 数学定义为：$\text{SwigLU}(x) = \text{Swish}(x_1) \otimes x_2$

其中：

  1. x 是输入向量，被拆分为两个等长的子向量 $x_1$ 和 $x_2$（通常通过将输入维度切分或用两个线性层映射得到）；
  2. $\text{Swish}(x_1) = x_1 \cdot \sigma(\beta x_1)$ 是激活函数（$\sigma$ 为 Sigmoid，$\beta$ 为超参数，通常取 1），用于计算 “门控值”；
  3. $\otimes$ 表示逐元素相乘（Hadamard 乘积），即门控值与 $x_2$ 相乘，实现信息筛选。
3-p4

1.3 serial、parallel layers

经典的transformer是serial串行结构。

公式：$y = x + \text{MLP}(\text{LayerNorm}(x + \text{Attention}(\text{LayerNorm}(x))))$


parallel并行结构最早在GPT-J中提出.

公式：$y = x + \text{MLP}(\text{LayerNorm}(x)) + \text{Attention}(\text{LayerNorm}(x))$

输入x经过层归一化后，同时送入MLP和注意力模块，实现并行。

现在大多数采用serial结构。


1.4 Position Embedding

transformer的自注意力机制无法感知序列顺序，不能处理文本序列等需要考虑位置的问题。所以使用位置编码告诉模型输入之间的先后顺序。

transformer中位置编码位于embedding层之后。

3-p5

位置编码为每个位置生成独特的编码向量，将其与输入的词嵌入向量相加，让模型学习到每个元素在序列中的位置。

3-p6

sine embeddings（正弦位置嵌入）

sine Embeddings 是一种经典的绝对位置嵌入方法，在论文《Attention Is All You Need》中提出，用于 Transformer 模型中。利用不同频率的正弦和余弦函数为每个位置生成一个独特的向量。


$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$$$PE_{(pos, 2i + 1)} = \cos\left(\frac{pos}{10000^{2i+1/d_{model}}}\right)$$

​

  • pos 是 token 在序列中的位置索引（从0开始）
  • i 是维度索引（不是维度），取值范围：$0 ≤ i ＜ d_{model} / 2 - 1$
  • $d_{\text{model}}$ 是模型的嵌入维度
  • $PE_{(pos, j)}$表示位置 pos 的编码向量中第 j 维的值

正弦余弦交替使用 对于每对相邻维度，偶数维度使用正弦函数，奇数维度使用余弦函数。


波长的递增序列 指数$10000^{2i/d_{model}}$为不同维度创建了不同的波长：

  • 当$i = 0$时：$10000^{0}=1$，波长为$2\pi$，波长短，函数值变化快。有助于捕捉近距离关系，具有高敏感性。
  • 当$i = d_{model}/2 - 1$时：$10000^{(d_{model}-2)/d_{model}}\approx10000$，波长为$10000\cdot2\pi$，在很长的序列范围内，高维度的编码值几乎恒定，有助于学习长距离依赖关系，提供稳定的、粗粒度的 “区域” 或 “上下文” 信号。

为什么选择10000？ 波长依据维度形成从$2\pi$到$10000\cdot2\pi$的几何级数，保证位置编码的频率覆盖范围足够大，同时兼顾长序列与短序列。


sine embeddings的特性

  1. 唯一性和有界性

    1. 各位置的编码向量互不相同
    2. 所有值都在[-1, 1]范围内
    3. 显示引入位置信息，不需要学习

  2. 无限外推能力

    1. sin 和 cos 对于任何实数都有定义，可以对任意位置生成编码，不受训练数据长度限制
    2. 编码模式沿着序列位置平滑变化，没有突变点

  3. 线性表示相对位置

    任意固定位置偏移的相对位置关系可以被表达为线性变换。

    • 对于任意固定偏移k，$PE_{pos + k}$可以表示为$PE_{pos}$的线性函数。

    • 对于任意位置$pos$和偏移$k$，考虑特定维度$i$：

      令$\alpha=\frac{pos}{10000^{2i/d_{model}}}$和$\beta=\frac{k}{10000^{2i/d_{model}}}$，则

      $PE_{(pos,2i)}=\sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)=\sin\alpha$，$PE_{(pos,2i + 1)}=\cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)=\cos\alpha$​

      通过积化和差公式：

      $PE_{(pos + k,2i)}=\sin(\alpha + \beta)=\sin(\alpha)\cos(\beta)+\cos(\alpha)\sin(\beta)$

      $PE_{(pos + k,2i + 1)}=\cos(\alpha + \beta)=\cos(\alpha)\cos(\beta)-\sin(\alpha)\sin(\beta)$

    这表明$PE_{pos + k}$可以通过$PE_{pos}$的线性组合得到，线性系数仅与偏移$k$相关，独立于位置$pos$。

    能高效推断相对关系，但编码的仍然是每个 token 的绝对位置。有没有直接编码相对位置关系的方式呢？


参考：

【Transformer】9分钟认识位置编码 | 数学解析 | 旋转位置编码_哔哩哔哩_bilibili


Rope embeddings（旋转位置嵌入）

简单理解：我们需要一种表示相对位置的算法。RoPE在计算注意力时，通过将qk都乘以旋转矩阵，来把注意力运算结果中加入位置信息。

特性：

  1. 相对距离决定注意力：两个 token 之间的注意力分数只依赖于它们的相对位置差，而不是绝对位置
  2. 平移不变性：序列整体平移不会改变内部的注意力模式
  3. 保留长度外推能力：随着相对距离增加，注意力分数自然衰减。这符合靠的近的token注意力高，离得远的token注意力低的要求

原理解析参考，这两个已经解释的较为清楚了：

旋转位置编码RoPE的简单理解_哔哩哔哩_bilibili

十分钟读懂旋转编码（RoPE）


2.Hyperparameters

经过这些年来大模型的迭代发展，人们在一些超参数的选择上逐渐形成了共识。

2.1 Feedforward - model dimension ratio

一般模型： $d_{ff}=4d_{model}$

如果使用GLU及其变体作为激活函数，需要额外引入一个权重矩阵。为了保持与原始 FFN 相同的计算成本和参数规模，GLU 变体普遍将隐藏层维度 $d_{ff}$ 缩小为原来的大约 2/3。所以 $d_{ff}=\frac{2}{3}\times(4d_{model})=\frac{8}{3}d_{model}$

3-p7

2.2 Head-dim*num-heads to model-dim ratio

Head-dim*num-heads和model-dim的比例基本为1。

3-p8

2.3 Aspect ratios

模型应该更宽还是更深？

3-p9

一般 $d_{model}/n_{layer}=128$​

如下图所示，不同参数量的模型基本都在比率100左右表现最好。

3-p10

2.4 Vocabulary sizes

早期的单语言的模型的词表大小在30-50k。现在的多语言模型在100-250k左右。

3-p11

2.5 Dropout&Regularization

在预训练时，由于数据量非常大，只需要在数据集上训练一遍，不大会出现过拟合。

早期模型会在预训练时采用dropout。现在大多只使用weight decay。

一般来说，使用weight decay是用来防止过拟合的。但实验发现weight decay并不会影响 val loss和train loss（下图左）。weight decay会与cosine模式的optimizer相作用，当学习率减小时，train loss会有效下降（下图中）。

3-p12

3.stability

我们希望模型在训练时，loss稳定下降，而不要波动过大，出现过多spike。

transformer的不稳定成因主要来自softmax：$\sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$

其中的指数运算易出现数值不稳定。其中的除法运算可能会除0错误。

softmax出现在模型的最后和注意力层中，以下分别从这两个位置进行优化。

3.1 z-loss

3-p13

增加一个惩罚项，使Z的值稳定在1。


3.2 QK norm

看不懂，跳过


3.3 Soft-capping

看不懂，跳过


4.Attention

4.1 GQA/MQA

当我们生成文字时，需要不断地预测下一个token，这一过程不能并行。

通过KV cache可以减少开销。

KV cache

KV Cache 是大语言模型（LLM）推理阶段的核心优化技术，通过缓存注意力机制中重复计算的键（Key）和值（Value），显著提升生成效率。

在 Transformer 解码器（如 GPT 系列模型）中，每一步生成新 token 时，都需要对已生成序列和当前输入计算自注意力（Self-Attention）。假设已生成序列长度为 n，新生成 token 为第 (n+1) 个，则自注意力的核心公式为：

(\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V)

其中：

  • Q（Query）：当前输入的查询向量（维度 (d_k)）；
  • K（Key）、V（Value）：已生成序列中所有 token 的键和值向量（维度均为 (d_k)）。

未优化时的问题：生成第 (n+1) 个 token 时，需要重新计算前 n 个 token 的 K 和 V，导致计算量随序列长度增长呈 (O(n^2)) 增长，推理速度极慢。

原理：

KV Cache 的本质是空间换时间：将已生成序列的 K 和 V 向量缓存起来，避免重复计算。

  1. 首次计算（生成第 1 个 token）：
    • 输入：初始序列（如提示词）；
    • 计算：为每个 token 生成对应的 (K_1, V_1)，并缓存到内存中。
  2. 后续计算（生成第 t 个 token，(t > 1)）：
    • 输入：前 (t-1) 个 token 生成的序列 + 新输入（第 t 个 token 的 Query）；
    • 计算：仅需为第 t 个 token 生成 (Q_t)，直接复用缓存中的 (K_1, K_2, …, K_{t-1}) 和 (V_1, V_2, …, V_{t-1})；
    • 更新缓存：将新生成的 (K_t, V_t) 追加到缓存中，供下一步使用。

通过这一机制，每次生成新 token 时，仅需计算当前 token 的 Q 以及与历史 (K/V) 的注意力得分，计算量从 (O(n^2)) 降至 (O(n))。


还讲了GQA、MQA、sparse attn、sliding window attn、interleave’full’ and ‘LR’ attn

慢慢补吧……