CS336 LecTure 4 笔记补充

Scaling Laws（缩放定律）

我们现在所需要研究的问题是，给定固定的计算预算 \(C\)（FLOPs），如何分配模型参数量 \(N\) 和数据量 \(D\) 才能使 Loss \(L\) 最小？

三大 Scaling Laws 公式

参数缩放律 (Parameter Scaling)

当数据量 \(D \to \infty\)（假设数据无限），Loss 仅随参数量 \(N\) 下降：

\[L(N) = \left( \frac{N_c}{N} \right)^{\alpha_N} + L_\infty \]

• \(N_c \approx 8.8 \times 10^9\) (特征参数量)
• \(\alpha_N \approx 0.076\)
• \(L_\infty\)：数据的不可压缩损失（下限）

现实中数据 \(D\) 是有限的，盲目增加 \(N\) 会导致严重过拟合，Loss 不再下降。

数据缩放律 (Dataset Scaling)

当模型足够大（\(N \to \infty\)），Loss 仅随数据量 \(D\) 下降：

\[L(D) = \left( \frac{D_c}{D} \right)^{\alpha_D} + L_\infty \]

• \(D_c \approx 5.4 \times 10^{13}\) (特征 Token 数)
• \(\alpha_D \approx 0.095\)

联合缩放律 (Joint Scaling Law)

同时考虑 \(N\) 和 \(D\) 的影响，物理含义是：模型容量不足和数据不足对 Loss 的伤害是独立且可加的：

\[L(N, D) = \left( \frac{N_c}{N} \right)^{\alpha_N} + \left( \frac{D_c}{D} \right)^{\alpha_D} + L_\infty \]

Compute-Optimal Training (Chinchilla 法则) ️

计算量公式

训练总 FLOPs \(C\) 与 \(N, D\) 的关系：

\[C \approx 6 N D \]

前向传播 \(2ND\) + 反向传播 \(4ND\) = \(6ND\)，忽略 Embedding 和 LayerNorm 的少量计算

最优配比推导

将 \(D = \frac{C}{6N}\) 代入联合缩放律，对 \(N\) 求导并令 \(\frac{\partial L}{\partial N} = 0\)，解得最优参数量 \(N^*\)：

\[N^* \propto C^{\frac{\alpha_D}{\alpha_N + \alpha_D}} \]

代入经验常数 \(\alpha_N \approx 0.076, \alpha_D \approx 0.095\)，可得：

\[N^* \propto C^{0.56}, \quad D^* \propto C^{0.44} \]

Chinchilla 经验法则

在常用范围内，上述幂律近似为线性关系：

\[N_{opt} \approx \frac{C}{12} \quad \text{or} \quad D_{opt} \approx 20 N \]

结论： 1B 参数的模型，应该用 20B Tokens 训练。

• LLaMA-1 (7B) 只用了 1T Tokens -> 欠训练 (Compute-Inefficient)
• LLaMA-3 (8B) 用了 15T+ Tokens -> 接近 Compute-Optimal

损失函数的幂律形式 (Power Law)

在最优配比下，Loss 与计算量 \(C\) 的关系：

\[L(C) = \left( \frac{C_c}{C} \right)^{\alpha_C} + L_\infty \]

• \(\alpha_C = \frac{\alpha_N \alpha_D}{\alpha_N + \alpha_D} \approx 0.042\)

\(\alpha_C\) 极小，意味着增加计算量的边际收益递减极快。计算量翻 10 倍，Loss 仅下降约 9%。这解释了为什么业界开始卷数据质量，而不是无脑堆算力。

---

Attention Alternatives

标准 Attention 的痛点在于其二次复杂度

标准 Attention 的复杂度瓶颈

• 注意力分数矩阵： \(A = Q K^T \implies\) 计算复杂度 \(O(N^2 D)\)，显存占用 \(O(N^2)\)
• 加权求和： \(O = A V \implies\) 计算复杂度 \(O(N^2 D)\)
• 总复杂度： \(O(N^2 D)\)

线性 Attention (Linear Attention) 的数学变换

为了将 \(O(N^2)\) 降为 \(O(N)\)，核心思想是改变矩阵乘法的结合律。

• 引入核函数 (Kernel Function)： 将 Softmax 替换为逐元素的非负核函数 \(\phi\)。
  \[\text{Attention}(Q, K, V) = \phi(Q) (\phi(K)^T V) \]

• 结合律变换：
  • 标准计算顺序：\((\phi(Q) \phi(K)^T) V \implies\) 复杂度 \(O(N^2 D)\)
  • 线性计算顺序： \(\phi(Q) (\phi(K)^T V) \implies\) 复杂度 \(O(N D^2)\)
• 状态空间递推 (State Space Model / Mamba 基础)：
  将 \(\phi(K)^T V\) 视为一个随时间步 \(t\) 累积的状态 \(S_t\)：
  \[S_t = S_{t-1} + k_t v_t^T \]
  \[o_t = q_t^T S_t \]
  (这种递推形式将序列维度的复杂度从 \(O(N)\) 降为了常数 \(O(1)\)，总复杂度变为 \(O(N D^2)\)

Hybrid Attention

结合全局与局部的优势：

\[\text{Output} = \text{GlobalAttention}(Q, K, V) + \text{LocalAttention}(Q, K, V) \]

• Global Attention: 捕获长距离依赖（计算量大，但窗口小或步长大）。
• Local Attention: 捕获局部语法结构（滑动窗口，计算量 \(O(N \times W)\)，其中 \(W\) 为窗口大小）。

---

混合专家架构 MoE

MoE 在干什么

• 核心思想： 模型拥有巨大的参数容量，但在每次前向传播时，只激活其中一小部分专家。实现了“参数量”与“计算量”的解耦。
• 性能收益： 在相同的计算预算（FLOPs）下，MoE 模型的训练 Loss 下降速度显著快于密集模型。
• 系统优势： 专家天然适合专家并行 (Expert Parallelism)，极大提升千卡集群的吞吐量。

为什么只替换 FFN，不替换 Attention

Attention 涉及全局 Token 交互，替换为 MoE 会导致通信开销爆炸且难以稳定训练；而 FFN 是逐位置（Position-wise）独立计算的，天然适合稀疏化。

路由机制 (Routing Function)

路由器决定每个 Token 去哪个专家。主流策略为 Token Choice (Top-K)。

Top-K 路由公式

\[G(x) = \text{Softmax}(x \cdot W_g) \]

\[\hat{G}(x)_i = \begin{cases} G(x)_i & \text{if } i \in \text{TopK}(G(x)) \\ 0 & \text{otherwise} \end{cases} \]

\[y = \sum_{i=1}^{E} \hat{G}(x)_i \cdot E_i(x) \]

负载均衡 (Load Balancing) ️

如果没有约束，Router 会“偷懒”，把所有 Token 都发给同一个最好学的专家，导致其他专家“饿死”。

辅助损失 (Auxiliary Loss)

在总 Loss 中加入一项负载均衡惩罚：

\[L_{aux} = \alpha \sum_{i=1}^{E} f_i \cdot P_i \]

• \(f_i = \frac{1}{T} \sum_{x \in \text{Batch}} \mathbb{I}(\text{argmax}(G(x)) = i)\)：专家 \(i\) 被选中的频率
• \(P_i = \frac{1}{T} \sum_{x \in \text{Batch}} G(x)_i\)：专家 \(i\) 的平均概率
• 数学直觉： 只有当频率 \(f_i\) 和概率 \(P_i\) 都高时，Loss 才大。这迫使 Router 均匀分配概率，且实际选择也均匀。

DeepSeek-V3 ：Expert Biases

引入可学习的专家偏置 (Per-Expert Biases)，在 Router 的 Logits 上加上 Bias。如果某个专家负载过高，Bias 会自动降低，迫使 Router 选择其他专家，比传统的 Auxiliary Loss 更稳定。

Softmax 稳定性与 Z-Loss (工程数学)

在 MoE 训练中，由于路由器 logits 容易发散导致 Softmax 溢出，通常会引入 Z-Loss：

\[L_{z} = \lambda_z \log^2 \left( \sum_{i=1}^{E} e^{s_i} \right) \]

• 作用： 惩罚 Softmax 的分母（即配分函数 \(Z\)）过大，强制将路由器的 Logits 限制在一个合理的数值范围内，防止 NaN 崩溃。

专家设计变体与训练技巧

• 细粒度专家 (Fine-Grained)： 专家数量更多，单个专家更小（如 DeepSeek-V2 使用 160 个专家，每次激活 6 个）。更细粒度的知识分工使路由更灵活。
• 共享专家 (Shared Experts)： 保留 1 个（或几个）始终被激活的专家，存储通用知识（如语法、基础逻辑），防止路由专家因负载限制而丢失基础能力。
• Token Drop： 为控制显存，限制每个专家最多处理的 Token 数。超出的 Token 会被丢弃，导致即使在 temperature=0 时输出也可能不确定。
• Upcycling (密集转稀疏)： 拿一个已训练好的 Dense 模型，将 FFN 复制 N 份作为 N 个专家的初始化，然后继续训练 MoE。节省大量预训练成本，且收敛极快。

MoE 的计算量与显存公式

• 计算量 (FLOPs)： 仅与激活的专家有关
  \[\text{FLOPs}_{MoE} \approx 6 \cdot N_{active} \cdot D_{tokens} \]
  (其中 \(N_{active} = K \times N_{expert\_size}\))
• 显存占用 (Memory)： 与所有专家有关
  \[\text{Memory}_{MoE} \propto N_{total} = E \times N_{expert\_size} \]
  (即使不激活，所有专家的权重也必须加载到 GPU 显存中)

---

Dense 与 MOE 的比较

这种比较还是扔给 DeepSeek 来写，一般 DeepSeek 生成的比较表格方面挺全面的。

维度	密集模型 (Dense)	混合专家 (MoE)
计算效率	高（所有参数都参与计算）	低（有路由开销、通信开销、显存碎片）
参数效率	低（增加参数 = 增加计算）	极高（增加参数 ≠ 增加计算）
训练稳定性	高	中（需处理负载均衡、Token Drop）
显存占用	仅模型参数	模型参数 + 所有专家参数（即使不激活也要加载）
适用场景	中小模型、低延迟推理	超大模型、高吞吐服务、长上下文

MoE 的本质是用“系统复杂度”换取“参数效率”。在单卡时代，MoE 是累赘；但在千卡集群时代，MoE 是通往万亿参数的唯一路径。如果计算预算只够训练 7B Dense，用 MoE 只会得到更差的 7B 模型。