CS336 Lecture Notes 2

旋转位置编码（RoPE Embeddings）#

代表模型：GPT-J、PaLM、LLaMA 系列、绝大多数 2024 年后的模型。

RoPE 提出了一个理想的相对位置编码应该满足的条件：

• 左边：是位置 $i$ 的词向量 $x$ 和位置 $j$ 的词向量 $y$，经过位置编码函数 $f$ 后的内积（也就是注意力计算的核心）。
• 右边：这个内积的结果，应该只和两个词的内容 $x$，$y$，以及它们的相对位置差 $i−j$ 有关，而和它们的绝对位置 $i$，$j$ 无关。

简单说：注意力分数只由 “两个词离得有多远” 决定，而不是它们在序列里的绝对位置，这样模型才能真正学到 “相对位置关系”，并且能外推到更长的序列。

现有方法的缺陷

• Sine：它的内积展开式里面的$\langle PE_i​,v_y \rangle​$ 和 $\langle v_x​,PE_j \rangle$ 会引入和绝对位置 $i$，$j$ 相关的项，导致内积不仅和 $i−j$ 有关，还和绝对位置有关，破坏了纯相对位置的特性。

• Absolute：直接给每个位置加一个可学习的向量，内积结果会直接依赖于 $i$ 和 $j$，完全不具备 “相对位置” 的性质。而且无法外推到训练时没见过的序列长度。

• Relative embeddings：$e_{ij} = \frac{x_i W^Q (x_j W^K + a_{ij}^K)^T}{\sqrt{d_z}}$ 让注意力计算不再是纯的内积形式。

RoPE 要解决的问题，就是让注意力计算只依赖相对位置差，而不是绝对位置。它的灵感来自一个关键数学性质：

向量的内积，在整体旋转下是不变的。

也就是说，两个向量一起旋转相同角度，它们的夹角和内积都不会变。

RoPE 就是利用这个性质，给不同位置的向量加上 “旋转角度”，让内积结果天然只和位置差有关。

FFN 维度比例#

几乎所有模型遵循 $d_{\text{ff}} = 4 \cdot d_{\text{model}}$（GLU 变体为 $\frac{8}{3} \cdot d_{\text{model}}$）。

极端案例：T5 11B 曾用 $64 \cdot d_{\text{model}}$，但后续 T5 v1.1 已回归标准值。

Head 维度#

主流选择：$\text{head\_dim} \times \text{num\_heads} = d_{\text{model}}$

大多数模型这个比例为 1，少数 Google 模型（如 LaMDA, PaLM）有例外。

宽深比 (Aspect Ratio)#

模型应该变宽还是变深？多宽以及多深？

大多数模型的 $d_{\text{model}} / n_{\text{layers}}$ 在 100-200 范围：

过深的模型并行化困难、延迟高，但 [Kaplan et al 2020] 显示小模型时深一点更好。

词汇表大小#

• 单语言模型：30k-50k
• 多语言/生产模型：100k-250k（GPT-4 ~100k, Gemma 4 ~262k）

Dropout 与 Weight Decay#

趋势：新模型大多不用 Dropout，仅靠 Weight Decay。

原因：数据量大（万亿 tokens），单遍 SGD 不易过拟合；Weight Decay 主要影响优化动力学而非泛化。

Z-loss#

Transformer 语言模型的输出层用 Softmax 计算 token 概率：

其中归一化项（也叫配分函数）是：

取对数后，我们得到模型的对数概率：

训练时，我们优化的目标就是最大化这个对数概率（交叉熵损失，等价于最小化负对数概率）。

随着模型变大、logits（也就是 $U_r​(x)$）的数值范围变宽，$Z(x)$ 会出现两种极端情况：

1. logits 太大：指数项 $e^{U_r}$​ 爆炸，$Z(x)$ 变得极大，导致 $log(Z(x))$ 也变得非常大，$log(P(x))$ 被压得很小，梯度消失。
2. logits 太小：指数项 $e^{U_r}$​ 接近 0，$Z(x)$ 接近 1，$log(Z(x))$ 接近 0，但此时模型可能过于自信，也容易出现梯度问题。

当 $log(Z(x))$ 出现剧烈波动时，损失和梯度都会变得不稳定，导致训练过程震荡甚至崩溃。

为了解决这个问题，我们在交叉熵损失上额外加一个正则项，也就是 Z-loss：

PaLM 首次大规模验证了 Z-loss 的有效性，后续很多模型，如 Baichuan 2、DCLM、OLMo 2/3 等都沿用了这个做法。

QK-norm#

在 attention softmax 前对 Query 和 Key 做 LayerNorm/RMSNorm。

标准自注意力公式：

QK Norm 定义（对 Q 和 K 分别归一化）：

加入 QK Norm 的最终注意力公式：

使用模型：DCLM, OLMo 2, Gemma 2, Qwen 3。

Logit Soft-capping#

用 tanh 限制 logits 范围：

1. 先把 logits 除以一个缩放因子 soft_cap（比如 50 或 30）；
2. 再用 tanh 函数做非线性变换，将值限制在 [-1, 1] 之间；
3. 最后乘回 soft_cap，最终 logits 就被限制在 [-soft_cap, soft_cap] 之间。

GQA/MQA：降低推理成本#

推理时的增量计算依赖 KV Cache，内存访问量大。MQA（Multi-Query Attention）和 GQA（Grouped-Query Attention）通过减少 KV heads 来降低内存开销。

问题背景：Attention 的计算与内存特性#

在自回归推理中，每次生成一个新 token 时需要访问之前所有 token 的 Key 和 Value。为了理解为什么 KV Cache 是瓶颈，我们需要分析 Attention 的 FLOPs（浮点运算次数） 和 内存访问量，计算 Arithmetic Intensity（算术强度）。

推理时的 FLOPs 估计

生成第 $t$ 个 token 时，主要张量的形状：

• 当前 Query $Q$：$\mathbb{R}^{h \times d_h}$（$h$ 个头，每头 $d_h$ 维）
• 历史 Key Cache $K$：$\mathbb{R}^{h \times t \times d_h}$（$h$ 个头，$t$ 个历史位置，每头 $d_h$ 维）
• 历史 Value Cache $V$：$\mathbb{R}^{h \times t \times d_h}$

Attention 的主要计算：

1. Query-Key 点积：$Q \times K^\top$ 得到注意力分数矩阵

   • 形状变化：$[h, d_h] \times [h, d_h, t]^\top \rightarrow [h, t]$
   • 每个头：$[d_h] \times [t, d_h]^\top = [t]$，即 $d_h$ 维向量与 $t$ 个 $d_h$ 维向量做点积
   • FLOPs：$t \times d_h$ 次乘法 + $t \times (d_h-1)$ 次加法 $\approx 2td_h$
   • $h$ 个头总计：$2htd_h$

2. Attention 加权 Value：用注意力分数 $[h, t]$ 加权历史 Value $[h, t, d_h]$

   • 形状变化：$[h, t] \times [h, t, d_h] \rightarrow [h, d_h]$
   • 每个头：$[t]$ 个权重加权 $[t, d_h]$ 的 Value，得到 $[d_h]$
   • FLOPs：$t \times d_h$ 次乘法 + $t \times (d_h-1)$ 次加法 $\approx 2td_h$
   • $h$ 个头总计：$2htd_h$

单层 Attention 的总 FLOPs（不含投影）：

$$\text{FLOPs}_{\text{attn}} \approx 4htd_h$$

推理时的内存访问量

设数据类型为 dtype（如 fp16、fp32），字节大小为 $\text{sizeof}(\text{dtype})$。生成第 $t$ 个 token 时：

操作	张量形状	访问量（bytes）
读取当前 Query	$[h, d_h]$	$h \cdot d_h \cdot \text{sizeof}(\text{dtype})$
读取历史 Key Cache	$[h, t, d_h]$	$h \cdot t \cdot d_h \cdot \text{sizeof}(\text{dtype})$
读取历史 Value Cache	$[h, t, d_h]$	$h \cdot t \cdot d_h \cdot \text{sizeof}(\text{dtype})$
写入当前 K/V 到 Cache	$[h, 2, d_h]$	$2 \cdot h \cdot d_h \cdot \text{sizeof}(\text{dtype})$
写入输出	$[h, d_h]$	$h \cdot d_h \cdot \text{sizeof}(\text{dtype})$

总内存访问量：

$$\text{Memory}_{\text{access}} = (2ht + 4h) \cdot d_h \cdot \text{sizeof}(\text{dtype})$$

主要开销来自读取历史 KV Cache：$2ht \cdot d_h \cdot \text{sizeof}(\text{dtype})$。

Arithmetic Intensity（算术强度）

算术强度定义为 FLOPs 与内存访问量的比值：

$$\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{Bytes accessed}}$$

代入上述估计：

$$I = \frac{4htd_h}{(2ht+4h) \cdot d_h \cdot \text{sizeof}(\text{dtype})} = \frac{4t}{(2t+4) \cdot \text{sizeof}(\text{dtype})}$$

当序列长度 $t$ 较大时：

$$I \rightarrow \frac{2}{\text{sizeof}(\text{dtype})} \text{ FLOPs/byte}$$

对于 fp16，$I \rightarrow 1$ FLOP/byte。

关键洞察：

典型 GPU 的峰值算力与内存带宽对比：

• NVIDIA A100：峰值 312 TFLOPS (fp16)，带宽 2 TB/s → 理论峰值强度 156 FLOPs/byte
• NVIDIA RTX 4090：峰值 83 TFLOPS (fp16)，带宽 1 TB/s → 理论峰值强度 83 FLOPs/byte

Attention 推理的算术强度 $\approx 1$ FLOPs/byte，远低于 GPU 峰值强度，属于 memory-bound（内存受限） 操作：计算单元大部分时间在等待数据加载，而非真正执行计算。

这就是 KV Cache 成为推理瓶颈的根本原因——降低内存访问量比优化计算更关键。

1
def gqa_attention(Q, K, V, h=64, g=8):
2
    """
3
    Q: [h, d_h]    - h 个 query 头
4
    K: [g, d_h]    - g 个 key 头
5
    V: [g, d_h]    - g 个 value 头
6
    """
7
    heads_per_group = h // g
8

9
    outputs = []
10
    for i in range(h):
11
        group_idx = i // heads_per_group  # 确定 query 头属于哪个组
12
        scores = (Q[i] @ K[group_idx]) / math.sqrt(d_h)
13
        weight = softmax(scores)
14
        out = weight * V[group_idx]
15
        outputs.append(out)
16

17
    return outputs  # [h, d_h]

Sliding Window Attention#

标准 Attention 的计算复杂度为 $O(n^2)$，当序列长度 $n$ 很大时（如长文档、长对话），计算和内存开销成为瓶颈。

Sliding Window Attention（滑动窗口注意力） 由 Longformer (Beltagy et al., 2020) 提出，核心思想是限制每个 token 只能 attend 窗口范围内的相邻 token，而非全部历史。后续 Mistral (Jiang et al., 2023) 将 SWA 与滚动缓存结合，实现了高效的无限长度生成。

原理与计算复杂度

设窗口大小为 $w$，每个 token 只与最近的 $w$ 个 token 计算注意力：

其中 $K^\top$ 的形状从 $[n, d_k]$ 变为 $[w, d_k]$（只取窗口内的 Key）。

计算复杂度变化：

当 $w \ll n$ 时，复杂度接近线性。例如 $n=8192, w=4096$ 时，计算量减少一半；$w=512$ 时，减少 16 倍。

实现示意

1
def sliding_window_attention(Q, K, V, window_size):
2
    """
3
    Q: [n, d_k]  - n 个 query
4
    K: [n, d_k]  - n 个 key
5
    V: [n, d_v]  - n 个 value
6
    window_size: 窗口大小 w
7
    """
8
    n = Q.shape[0]
9
    outputs = []
10

11
    for i in range(n):
12
        # 确定窗口范围: [i - window_size + 1, i]
13
        start = max(0, i - window_size + 1)
14
        end = i + 1
15

16
        # 只取窗口内的 K 和 V
17
        K_window = K[start:end]  # [w, d_k]
18
        V_window = V[start:end]  # [w, d_v]
19

20
        # 计算 attention
21
        scores = Q[i] @ K_window.T / math.sqrt(d_k)
22
        weights = softmax(scores)
23
        out = weights @ V_window
24
        outputs.append(out)
25

26
    return outputs  # [n, d_v]

窗口设计的权衡

窗口过小会丢失全局依赖（如段落间引用、跨句语义），窗口过大则计算收益有限。

现代混合策略

单纯使用 SWA 会限制全局信息流动。现代做法是交错使用全注意力和局部注意力层：

关键洞察：混合策略在计算效率与全局信息捕获之间取得平衡。全注意力层充当”信息枢纽”，将局部信息汇聚后传递到全局。

Linear Attention#

核心思想：利用矩阵乘法的结合律重排注意力计算：

$\text{Attn}(Q, K, V) = \rho(QK^\top)V = Q(K^\top V)$

从 $n^2 d_k + n^2 d_v$ 降到 $2n d_v d_k$，实现线性时间复杂度。

递归形式（RNN duality）：

$S_t = S_{t-1} + k_t v_t^\top, \quad y_t = q_t^\top S_t$

这种”对偶性”允许：

• 训练时使用并行二次形式（效率高）
• 推理时使用串行线性形式（内存友好）

若给 $S_{t-1}$ 加权重因子 $\gamma$，得到 RetNet。

实际应用：MiniMax M1 使用 7:1 混合（7 层线性注意力 + 1 层全注意力），在长上下文场景表现良好。

Mamba-2：带门控的线性注意力#

对线性注意力引入逐位置权重：

$S_t = \gamma_t S_{t-1} + k_t v_t^\top, \quad y_t = q_t^\top S_t + v_t^\top D$

其中 $\gamma_t = f(x_t)$ 是数据依赖的门控因子。相比线性注意力，门控机制提供了更强的表达能力，同时保持对偶性质（并行计算 $\gamma$，应用对偶变换）。

实际应用：Nemotron 3 采用 Mamba/Attention 3:1 混合，性能与同规模稠密模型相当或更好。

Gated Delta Net (GDN)#

进一步泛化：门控输入 + 选择性擦除状态：

$S_t = \gamma_t(I - \beta_t k_t k_t^\top)S_{t-1} + \beta_t k_t v_t^\top, \quad y_t = q_t^\top S_t$

关键特性：

• $\beta_t$ 控制”无输入操作”（$\beta = 0$时不更新）
• $(I - \beta_t k_t k_t^\top)$ 除当前 key 方向的信息（选择性遗忘）

与 fast weight programming / test-time training 思想密切相关。

实际应用：Qwen 3.5 / Qwen Next 采用 3:1 GDN/Attention 混合，推理效率良好。

Hybrid 性能总结#

混合架构（线性注意力/SSM + 全注意力）的主流配置：

实证表明：较小的混合比例（少量全注意力层）通常能保持接近全注意力的性能，同时显著降低长上下文成本。

稀疏注意力 (DSA)#

替代混合方案：不 attend 所有 token，使用稀疏注意力（DeepSeek Sparse Attention, v3.2 / GLM5）。

特点：

• Indexer 可很轻量，显著降低开销
• 可在稠密短上下文预训练后”事后适配”

为什么 MoE 受欢迎？#

1. 相同 FLOPs，更多参数 → 更好性能：实证表明增加专家数量不改变推理成本，但提升质量
2. 训练更快：相同数据量下，MoE 达到相同 loss 需更少步数
3. 竞争性强：与稠密等价模型相比，MoE 性能优异
4. 易于并行化：每个专家可放置在不同设备

MoE 架构#

典型设计：替换 MLP 层为 MoE 层。较少见：对 attention heads 使用 MoE（ModuleFormer, JetMoE）。

关键变量：

• 路由函数（Routing function）
• 专家规模
• 训练目标

路由机制#

几乎所有 MoE 使用 Top-K 路由：token 选择 top-k 个专家。

Top-K 路由细节：

Gate 由 logistic regressor 选择（DeepSeek V1-2, Grok, Qwen）。Mixtral、DBRX、DeepSeek v3 在 TopK 后做 softmax。

近期变体（DeepSeek / Qwen）：

• 更小、更多专家 + 少量”共享专家”（始终激活）
• 共享专家来自 DeepSpeed MoE，确保通用知识始终可用

专家配置对比：

消融实验结论：

• 更多专家 + 共享专家通常有益（DeepSeek 消融）
• OlMoE：细粒度专家有增益，共享专家无增益（不同架构结论可能不同）

MoE 训练挑战#

核心问题：稀疏路由决策不可微分。

解决方案对比：

平衡损失（Auxiliary Loss）：

来自 Switch Transformer，惩罚专家使用频率不均：

$\mathcal{L}_{\text{aux}} = \alpha \sum_i f_i P_i$

其中 $f_i$ 是专家 $i$ 接收的 token 比例，$P_i$ 是路由概率。

DeepSeek v1-2：额外添加设备级平衡损失（按设备聚合）。

DeepSeek v3 变体：per-expert bias + online learning，称为 “auxiliary-loss-free balancing”（实际仍有辅助损失）。

系统并行性#

MoE 天然支持并行：每个 FFN/专家可放置在单独设备。

挑战：

• Token dropping（batch 级路由可能导致其他请求的 token 被丢弃）
• 稀疏矩阵操作复杂

解决方案：MegaBlocks 等库使用智能稀疏矩阵乘法；Nemotron 3 下投影激活以减少通信。

训练稳定性#

MoE 路由器可能导致 loss spike。

解决方案：对路由器使用 Float32 + auxiliary z-loss（类似注意力中的 z-loss，防止 logits 过大）。

微调问题#

稀疏 MoE 在小规模微调数据上易过拟合。

解决方案：

• Zoph 等：微调非 MoE MLP
• DeepSeek：使用大量 SFT 数据（1.4M）

Upcycling：从稠密模型初始化 MoE#

用预训练稠密模型初始化 MoE，避免从头训练。

成功案例：

• MiniCPM MoE：基于 MiniCPM，top-k=2, 8 专家，~4B 激活参数，520B tokens
• Qwen MoE：基于 Qwen 1.8B，top-k=4, 60 专家 + 4 共享，首批确认的 upcycling 成功案例

DeepSeek MoE 演进#

DeepSeek v3 关键技术：

1. MLA (Multi-head Latent Attention)：将 Q/K/V 表示为低维”潜在激活”的函数。KV-cache 只需存储 $c_t^{KV}$，显著压缩。RoPE 与 MLA 缓存冲突的解决方案：保留少量非潜在 key 维度用于旋转。

2. MTP (Multi-Token Prediction)：小型轻量模型预测多步 ahead（类似 EAGLE），v3 只用单步预测。