架构选择：注意力、位置编码、MoE 与混合模型

交互

Llama 3 70B 如果用多头注意力（MHA），一条 8k token 序列的 KV cache 就要 20 GiB；它实际用的 8 个 KV 头的 GQA 只要 2.5 GiB。而在 1B 规模的对照实验里，比例 2 到 8 的 GQA 和 MHA 的 loss、下游分数几乎一样。现代 LLM 的结构选择大多是这样一笔账：一个组件解决什么问题，代价是什么，消融说它值不值。

这些对照实验都用同一个 1B 消融设置（见从目标到决策）：Llama 3.2 1B 结构，在 FineWeb-Edu、FineMath、Python-Edu 的混合数据上训 45B token。落到具体模型上，答案就是 SmolLM3 的 3B dense 结构。

目录
  1. 先看大家都在用什么
  2. 注意力的头：MHA、MQA、GQA、MLA
  3. 文档掩码：packing 之后谁能看见谁
  4. Embedding 共享
  5. 位置编码：从绝对位置到 RoPE
  6. RoPE 的频率与长上下文：ABF、YaRN
  7. NoPE、RNoPE 与 partial RoPE
  8. 限制注意力范围：分块、滑窗、DCA、attention sink
  9. 稳定性：z-loss、embedding 不做权重衰减、QK-norm
  10. 初始化、激活函数、深度与宽度
  11. MoE：稀疏度、粒度、共享专家、负载均衡
  12. 混合模型：线性注意力与 SSM
  13. Dense、MoE 还是 Hybrid
  14. SmolLM3 的结构
  15. 参考与延伸阅读

先看大家都在用什么

Qwen3、Gemma 3、DeepSeek-V3 看起来各不相同，但底子都是 2017 年的 Transformer。几年来骨架没怎么变，变的是组件：为了降低推理显存、为了大规模训练的稳定性、为了更长的上下文。有的改进已经成为共识（比如从多头注意力换成 GQA），有的还在争论（比如位置编码）。先看几个公开细节比较完整的模型：

模型结构参数训练 token注意力最终上下文位置编码
DeepSeek LLM 7B / 67BDense7B / 67B2TGQA4kRoPE
DeepSeek-V2MoE236B（21B 激活）8.1TMLA128kPartial RoPE
DeepSeek-V3MoE671B（37B 激活）14.8TMLA128kPartial RoPE
MiniMax-01MoE + hybrid456B（45.9B 激活）11.4T线性注意力 + GQA4MPartial RoPE
Kimi K2MoE1T（32B 激活）15.5TMLA128kPartial RoPE
OLMo 2 7BDense7B5TMHA4kRoPE
SmolLM3Dense3B11TGQA128kNoPE（混合）

表里注意力有 MHA、GQA、MLA 三种，位置编码有 RoPE、partial RoPE、NoPE 三种。

注意力的头：MHA、MQA、GQA、MLA

为什么推理时注意力是瓶颈

预训练时算力主要花在前馈层（MLP）上；到了推理，尤其是长上下文，注意力成了瓶颈：它推高计算量和显存，压低吞吐。原因在于自回归生成。先回忆多头注意力（MHA）：输入 \(X\in\mathbb R^{T\times h}\)（\(T\) 个 token、隐藏维度 \(h\)），对第 \(i\) 个头（共 \(n_h\) 个，每头维度 \(d=h/n_h\)）

$$Q_i=XW_Q^{(i)},\quad K_i=XW_K^{(i)},\quad V_i=XW_V^{(i)},\qquad O_i=\operatorname{softmax}\!\Big(\frac{Q_iK_i^\top}{\sqrt d}+M\Big)V_i,$$

其中 \(M\) 是因果掩码（未来位置为 \(-\infty\)）。各头输出拼接后乘 \(W_O\)。直观地说，每个头在做一次「检索」：用当前 token 的 query 去匹配所有历史 token 的 key，把匹配上的 value 取回来。

生成第 \(t\) 个 token 时，只有它自己的 \(q_t\) 是新的；历史 token 的 \(k_j, v_j\)（\(j\lt t\)）在之前的步骤里已经算过，存下来复用即可，这块内存就是 KV cache。它的大小：

$$s_{KV}=2\times n_{\text{bytes}}\times T\times n_{\text{layers}}\times n_{kv}\times d,$$

前面的 2 来自 K 和 V 各存一份。以 Llama 3 的结构为例，假设用 MHA（\(n_{kv}=n_h\)）、序列长 8192、BF16（2 字节）：

$$\text{8B: } 2\times2\times8192\times32\times32\times128 = 4\ \text{GiB},\qquad \text{70B: } 2\times2\times8192\times80\times64\times128=20\ \text{GiB}.$$

这只是单条序列。cache 随序列长度线性增长，而上下文窗口已经长到百万级——这就是压缩 KV cache 的动机。实际的 Llama 3 用的是 8 个 KV 头的 GQA，真实数字是上面的 1/4 和 1/8。

MQA 与 GQA：让多个 query 头共享 K、V

真的需要每个头都有自己的 K、V 吗？多查询注意力（MQA）（Shazeer 2019）让所有 query 头共享同一组 K、V，KV cache 直接除以 \(n_h\)（Llama 3 70B 是 64 倍）；StarCoder 就用了 MQA。但这可能损失太多注意力容量。分组查询注意力（GQA）（Ainslie et al. 2023）取中间：把 \(n_h\) 个 query 头分成 \(n_{kv}\) 组，组内共享一组 K、V。

MHA GQA MQA MLA Q K/V 缓存的部分用橙色表示；MLA 只缓存一个低维潜向量 c，K/V 在用时由它解压
四种注意力的 K/V 共享方式（以 8 个 query 头为例）。

实现上，GQA 的 \(W_K, W_V\) 只输出 \(n_{kv}d\) 维，计算注意力前把每组的 K、V 复制 \(n_h/n_{kv}\) 份（常见的 repeat_kv），或者在 kernel 里直接按组索引。张量形状：\(Q\) 是 [B, T, n_h, d]，\(K,V\) 是 [B, T, n_kv, d]；第 \(i\) 个 query 头读第 \(\lfloor i\,n_{kv}/n_h\rfloor\) 个 KV 头。有的库把 \(n_h/n_{kv}\) 叫「GQA 比例」或 query groups。

MLA：不减少头数，而是压缩 K、V 本身

DeepSeek-V2 提出的多头潜在注意力（MLA）（DeepSeek-V2，V3 沿用）换了一个思路：不减少 K、V 的数量，而是只缓存一个低维的潜向量，用时再「解压」成各头的 K、V：

$$c_t = W_{DKV}\,x_t\in\mathbb R^{d_c},\qquad k_t^{(i)}=W_{UK}^{(i)}c_t,\qquad v_t^{(i)}=W_{UV}^{(i)}c_t.$$

推理时不必真的把 \(k^{(i)}\) 解压出来。注意力分数 \(q^{(i)\top}k^{(i)}= (W_Q^{(i)}x_t)^\top W_{UK}^{(i)}c_j = x_t^\top\big(W_Q^{(i)\top}W_{UK}^{(i)}\big)c_j\)，中间的矩阵可以预先乘好（「吸收」进 query 投影），于是直接拿变换后的 query 和缓存的 \(c_j\) 做点积；\(W_{UV}\) 同理可以吸收进输出投影。这样每层每个 token 只需缓存 \(d_c\) 个数，而且 K、V 共用同一个 \(c\)，所以没有前面那个因子 2。

麻烦出在 RoPE 上（RoPE 是一个依赖位置的旋转，细节见位置编码）：加了 RoPE 之后分数变成 \((W_Q^{(i)}x_t)^\top R_{t-j}\,(W_{UK}^{(i)}c_j)\)，旋转矩阵夹在两个投影中间、还随相对位置变化，就没法预先把两个投影乘成一个固定矩阵了。DeepSeek 的解法是解耦 RoPE：把每个头的 query/key 拆成两段，大段不加 RoPE（照常吸收），另加一小段专门承载位置信息的 \(k^R_t\in\mathbb R^{d_h^R}\)，它在所有头之间共享并单独缓存。DeepSeek-V2 取 \(d_c=4d_h\)、\(d_h^R=d_h/2\)，每层每 token 缓存 \(4.5\,d_h\) 个数，相当于只有 2.25 组的 GQA（\(2\times2.25\times d_h=4.5d_h\)），而论文报告其效果强于 MHA。

机制每 token 的 KV cache 元素数
MHA\(2\times n_h\times n_{\text{layers}}\times d\)
MQA\(2\times 1\times n_{\text{layers}}\times d\)
GQA\(2\times g\times n_{\text{layers}}\times d\)（常见 \(g=2,4,8\)）
MLA（DeepSeek-V2 设定）\(4.5\times n_{\text{layers}}\times d\)
KV cache 对比器
结构
80 GB 能放下多少 token（GQA g）

只算 KV cache 本身，不含权重和激活。层数、头数、头维度来自各模型的公开 config；「DeepSeek-V2 结构」一行的 MHA/GQA/MQA 是假想地把同样的 60 层、128 头、d=128 换成这些机制；MLA 行按 dc=4d、dR=d/2 计算（DeepSeek-V2 的设定）。80 GB 那一格只是「KV cache 独占一张 80GB 卡」的理想上限。

消融：GQA 不输 MHA

1B baseline 有 32 个 query 头、8 个 KV 头，也就是比例 4 的 GQA。把它换成 MHA 会不会更好？换成更少的 KV 头又会差多少？改 KV 头数会改变参数量，尤其是 MHA，所以 MHA 和比例 2 的 GQA 减少了层数来对齐（不对齐会差 100M 以上的参数）：

注意力Query 头KV 头层数参数备注
MQA321161.21B
GQA（比例 16）322161.21B
GQA（比例 8）324161.22B
GQA（比例 4）328161.24Bbaseline
GQA（比例 2）3216151.22B减了 1 层
MHA3232141.20B减了 2 层

结果：MQA 和 GQA 比例 16（只有 1、2 个 KV 头）明显不如 MHA；比例 2、4、8 的 GQA 与 MHA 基本持平。loss 和下游评测一致，HellaSwag、MMLU、ARC 上看得最清楚，OpenBookQA 和 WinoGrande 噪声大一些。结论：GQA 是 MHA 的可靠替代，推理还更省。MLA 没有参与这组消融（当时 Nanotron 还没有实现它）。SmolLM3 用的是 16 个 query 头、4 个 KV 头，即比例 4。

文档掩码：packing 之后谁能看见谁

预训练用固定长度的序列，但文档长短不一：一篇论文可能 10k token，一段代码可能几百。把短文档 padding 到固定长度会在无意义的 pad 上浪费算力，所以通用做法是 packing：把文档打乱、用 EOS token 连起来，再切成固定长度的块：

文件1 "燕麦棒的做法…"(400) <EOS> 文件2 "def hello_world()…"(300) <EOS>
文件3 "气候变化的影响…"(1000) <EOS> 文件4 "import numpy as np…"(3000) <EOS> …
切成 4k 序列：
序列1: [文件1][文件2][文件3][文件4 的前一部分]
序列2: [文件4 的其余部分][文件5][文件6]…

（也有在文档开头加 BOS 的做法，那时模型配置里会有一个不同的 bos_token_id。）用标准因果掩码，第 4 个文件里的 Python 代码能「看见」前面的燕麦棒菜谱和气候文章。这有多普遍？de Vries 的分析发现 Common Crawl 和 GitHub 里 80–90% 的文件短于 2k token；对 FineWeb-Edu、DCLM、FineMath、Python-Edu 做同样的统计，也是 80% 以上短于 2k。所以在 2k 或 4k 的训练序列里，绝大多数 token 在花算力关注不相干的文档。（例外是 PDF 来源的数据：FinePDFs 的平均长度是网页文本的两倍，和 FineWeb-Edu、DCLM 混合能提升效果。）

Zhao et al. 2024 发现跨文档注意力除了浪费算力，还会引入无关内容的噪声、损害效果，提出文档内掩码（intra-document masking）：每个 token 只能看同一文档里在它之前的 token。SkyLadder 也观察到了好处，但给了另一种解释：较短的上下文对训练更好，而文档内掩码实际上降低了平均上下文长度。Llama 3 也用了文档内掩码，发现短上下文预训练阶段影响有限，但对长上下文扩展帮助很大；ProLong 发现在继续预训练扩展 Llama 3 8B 上下文时用文档掩码，长短上下文的基准都受益。

实现上，文档内掩码是一个块对角的因果掩码。FlashAttention 的 varlen 接口用累计序列长度 cu_seqlens 描述每个文档的边界，直接跳过块对角之外的计算，所以它不仅是「改个 mask」，而是真的省了算力。下面把几种掩码画在一个 32 token 的序列上，最后两种是限制注意力范围用的分块和滑窗掩码：

注意力掩码可视化（32 个 token）
被计算的 (q,k) 对
相对全因果
跨文档的对
平均可见 token

行是 query、列是 key。实心格 = 允许注意；橙色实心 = 允许但跨了文档（packing 带来的噪声）。上边和左边的紫/灰色条交替标出各个文档。文档长度总和超过 32 时截断，不足时剩余部分当作一个文档。

1B 上的消融（Nanotron 里把 _use_doc_masking 设为 true）：loss 曲线和下游评测与标准因果掩码几乎一样，只在 PIQA 上有一点提升——和 Llama 3 的观察一致，短上下文阶段影响不大。但在扩展到长序列时（SmolLM3 从 4k 扩到 64k），文档掩码对加速训练至关重要，所以 SmolLM3 全程使用。

Embedding 共享

LLM 有两个和词表相关的矩阵：输入 embedding（查表，\(V\times h\)）和输出投影（LM head，把隐藏状态映射成词表上的 logits，\(h\times V\)）。两者分开时共 \(2Vh\) 个参数。共享（tying）就是让输出投影直接用输入 embedding 的转置：

$$\text{logits}_t = W_E\,\mathrm{norm}(x_t^{(L)})\in\mathbb R^{V},\qquad W_E\in\mathbb R^{V\times h}.$$

这样 \(W_E\) 同时从两头收到梯度：作为查找表（只有出现的 token 的行被更新）和作为输出层（每一步所有行都被更新）。小模型里 embedding 占比很大，共享是自然的优化：1B baseline 里 embedding 约 262.7M、占两成多。大模型一般不共享，因为占比小：Llama 3.1 8B 不共享时 embedding 只占约 13%（\(2\times128256\times4096\approx1.05\text{B}\)），70B 只占约 3%。

MobileLLM 在 125M 规模的消融显示，共享减少了 11.8% 的参数，精度损失很小。1B 上的消融设计得很干净，三组对比：

结果：共享的 1.2B 与不共享的 1.46B 除 WinoGrande 外基本持平，参数却少 18%；不共享、12 层的 1.2B 两者都不如，loss 更高、下游更低。也就是说，在同样的参数预算下，把参数花在深度上比花在解除共享上更值。SmolLM3 保留了共享 embedding。

位置编码：从绝对位置到 RoPE

为什么需要位置信息

注意力本身对 token 的顺序不敏感。把输入的行任意打乱（置换矩阵 \(P\)），没有掩码时 \(\operatorname{softmax}(PQ(PK)^\top)PV = P\operatorname{softmax}(QK^\top)V\)：输出只是跟着一起打乱，每个 token 得到的结果和它在哪个位置无关。在模型眼里，「Adam 打败了 Muon」和「Muon 打败了 Adam」长得差不多。所以要给每个 token 一个「地址」。

最早的 Transformer 用绝对位置编码：一张（学习的或正弦的）表，把位置 1、2、3… 映射成向量加到 token embedding 上。问题是模型只见过训练长度以内的位置，超出就不会了。于是人们转向相对位置编码：两个词相距 3 个位置，比它们在 (5,8) 还是 (105,108) 更重要。ALiBi 是其中简洁的一种：直接在注意力分数上加一个与距离成正比的惩罚 \(-m_i\,(t-j)\)，每个头的斜率 \(m_i\) 不同，越远越被压低。但近几年主导 LLM 的是 RoPE。

RoPE：把位置变成旋转

RoPE（旋转位置编码）的做法是：不把位置向量加到 embedding 上，而是按位置旋转 query 和 key。把每个头的 \(d\) 维向量两两分组成 \(d/2\) 对，第 \(k\) 对 \((x_{2k}, x_{2k+1})\) 看成平面上的一个点，在位置 \(m\) 处把它旋转角度 \(m\theta_k\)：

$$\begin{pmatrix}x'_{2k}\\x'_{2k+1}\end{pmatrix}=\begin{pmatrix}\cos m\theta_k & -\sin m\theta_k\\ \sin m\theta_k & \cos m\theta_k\end{pmatrix}\begin{pmatrix}x_{2k}\\x_{2k+1}\end{pmatrix},\qquad \theta_k=b^{-2k/d},\ k=0,\dots,\tfrac d2-1.$$

\(b\) 是基频（base，也叫 rope theta），常用 10000。不同的维度对转速不同：\(k=0\) 每个位置转 1 弧度，最快；\(k\) 越大越慢。举个例子：1B baseline 每头 64 维，分成 32 对；「The quick brown fox」里 fox 在位置 3，第一对的旋转角是 \(3\times b^{0}=3\) 弧度 \(\approx 172°\)——\(k=0\) 时指数为 0，基频不起作用。所以小 \(k\) 的维度对转得快，分辨近处的细节；大 \(k\) 的维度对转得慢，承载远距离的信息。

为什么旋转能编码相对位置？把第 \(k\) 对写成复数 \(z=x_{2k}+i\,x_{2k+1}\)，旋转就是乘以 \(e^{im\theta_k}\)。两个二维实向量的点积等于 \(\operatorname{Re}(a\bar b)\)，于是位置 \(m\) 的 query 和位置 \(n\) 的 key 在这一对上的贡献是

$$\operatorname{Re}\!\big[q_k e^{im\theta_k}\,\overline{k_k e^{in\theta_k}}\big]=\operatorname{Re}\!\big[q_k\bar k_k\,e^{i(m-n)\theta_k}\big].$$

只依赖 \(m-n\)！把所有对加起来，注意力分数是 \(q,k\) 和相对距离 \(m-n\) 的函数，与绝对位置无关。用矩阵说：\(R_m^\top R_n=R_{n-m}\)，所以 \((R_mq)^\top(R_nk)=q^\top R_{n-m}k\)。展开实部：

$$s(m,n)=\sum_{k}\Big[(q_{2k}k_{2k}+q_{2k+1}k_{2k+1})\cos\big((m-n)\theta_k\big)-(q_{2k+1}k_{2k}-q_{2k}k_{2k+1})\sin\big((m-n)\theta_k\big)\Big].$$

当 \(q=k\) 时 sin 项消失，分数只剩每对的幅度乘以 \(\cos((m-n)\theta_k)\) 再求和，下面演示里的第二张图画的就是这种情形。不管哪种情形，距离为 5 的两个 token 永远有同样的角度关系，所以模型学到的是基于距离的模式，在任何绝对位置都适用。实现上 RoPE 在 Q、K 投影之后、点积之前施加，V 不旋转。

每一对的「波长」是转一整圈需要的位置数 \(\lambda_k=2\pi/\theta_k=2\pi b^{2k/d}\)。\(d=64, b=10^4\) 时最快的一对波长约 6.3 个 token，最慢的一对（\(k=31\)）约 \(2\pi\times10^{4\cdot62/64}\approx4.7\times10^4\) 个 token。高频对分辨近距离，低频对编码远距离。

RoPE 可视化：旋转与相对位置
q·k（随机 q、k，全部维度对）
相对距离 m − n
最慢一对的波长

上图：第一对、中间一对、最后一对维度上，query（紫）在位置 m 转了 mθk，key（橙）在位置 n 转了 nθk。拖「同时平移」时两个向量一起转，夹角不变，所以上面的 q·k 读数不变。下图：q = k 时点积随相对距离（横轴对数刻度）的变化 \(\frac{2}{d}\sum_k\cos(\Delta\theta_k)\)，竖线是当前的 m − n。q、k 是固定的随机向量，只为演示。

RoPE 的频率与长上下文：ABF、YaRN

实践中，大部分预训练用较短的上下文（2–4k token）、RoPE 基频 10k 或 50k。从一开始就用长序列很贵（注意力随长度平方增长），长文档也少（超过 4k 的样本不多，见上面的文档长度统计），研究还表明会伤害短上下文表现（SkyLadder）。模型一开始学的是词与词之间的短程关联，长序列帮不上忙。所以常规做法是大部分时间用短序列，最后几百 B token 或在继续预训练里换长序列。

问题是：序列变长，旋转角 \(m\theta_k\) 随位置线性变大，远距离 token 的注意力分数会衰减得太快（Xiong et al. 2023、Code Llama）。在下面的曲线里能看到：基频小时，点积在较短距离上就振荡并衰减到零附近。对策是在拉长上下文时调大基频：

两个实际例子：Qwen3 在把上下文从 4k 扩到 32k 时用 ABF 把基频从 10k 提到 1M，再用 YaRN 外推到 131k（4 倍）。SmolLM3 在中训练阶段 4k→32k 时把基频调到 2M、32k→64k 时调到 5M，再用 YaRN 外推到 128k。最优值没有强共识，最好在扩展阶段针对你的设置和基准做实验。在演示里把基频从 \(10^4\) 拖到 \(5\times10^6\)，可以看到曲线在远距离上「撑」得久得多。

NoPE、RNoPE 与 partial RoPE

上下文继续变长（百万级）时，即使是调过频率的 RoPE，在比「大海捞针」（NIAH）更难的长上下文基准（RULER、HELMET）上也开始吃力。于是有人反过来问：显式位置编码真的必要吗？

NoPE 完全不用位置编码，让模型通过因果掩码和注意力模式隐式地学位置。为什么可行？因果掩码本身就打破了置换对称：第 \(t\) 个 token 恰好能看到 \(t\) 个 token，比如对所有可见 token 做均匀注意力，得到的输出就和 \(1/t\) 相关，模型可以由此恢复位置（论文给了第一层就能表示绝对位置的构造）。论文发现 NoPE 的长度外推好于 ALiBi 和 RoPE；但实践中 NoPE 在短上下文的推理和知识任务上往往弱于 RoPE。

Yang et al.（Cohere）因此提出混合：RNoPE，在层间交替 RoPE 和 NoPE。RoPE 层提供显式位置和对近处的偏好（recency bias），NoPE 层改善远距离检索。他们的模型里 RoPE 层还配了 4096 的滑动窗口，NoPE 层是全局注意力，比较了 1:1、1:3、1:7 的交替比例，1:3 最好。Llama 4、Command A、SmolLM3 都用了这类设计。下文为简洁，把这种混合也直接叫 NoPE。

1B 消融对比了三种：纯 RoPE baseline；每 4 层去掉一层的位置编码（NoPE）；NoPE + 文档掩码。三者 loss 和评测都相近——短上下文能力没丢，长上下文的基础打好了。SmolLM3 采用 NoPE + 文档掩码，config 里是 no_rope_layer_interval: 4，即第 4、8、12… 层不加 RoPE。

Partial RoPE 则是在同一层里混合：只对每个头的一部分维度施加 RoPE，其余维度不旋转。GLM-4.5、MiniMax-01 用了它，更早的 GPT-J 也是。所有用 MLA 的模型都用 partial RoPE，原因就是 MLA 的「吸收」：把头维度拆成 \(d_k=d_{\text{nope}}+d_{\text{rope}}\)，大的那块不旋转、照常吸收成 \((x_tU^{(h)}_{\text{nope}})^\top c_i\)，只有小块加 RoPE。

限制注意力范围：分块、滑窗、DCA、attention sink

上面都在改位置怎么编码。另一条互补的路是限制哪些 token 之间能互相注意。设想模型在 8 个 token 的序列上预训练，推理时要处理 16 个：位置 8–15 对位置编码来说是分布外的。限制注意力范围的方法让注意力模式保持在熟悉的范围内，同时还降低计算和显存。

前两种在上面的掩码演示里可以切换查看。

Attention sink：长上下文模型会给序列开头的几个 token 异常高的注意力，即使它们语义上不重要（Xiao et al.）。softmax 要求注意力权重之和为 1，当一个头「没什么可看」时，权重需要一个去处，开头的 token 就成了这个「水槽」。实际意义：KV cache 装不下时，保留开头几个 token 加最近的滑动窗口，就能基本保持效果，无需微调。实现方式各不相同：原论文建议预训练时加一个专门的占位 token；gpt-oss 把它做成每个头一个可学习的 bias logit，拼进 softmax 的分母里（相当于一个永远存在的「空」选项）。gpt-oss 还在注意力层用了 bias 项，这在 GPT-2 之后很少见，一般认为多余（Dehghani et al. 的结果显示对 loss 影响很小），但它能起到实现 sink 的作用。

稳定性：z-loss、embedding 不做权重衰减、QK-norm

训练不稳定（典型表现是 loss spike：loss 突然跳升）在大规模训练里尤其常见，轻则浪费几百步，重则要回滚到之前的 checkpoint。OLMo 2、Qwen3 等近期的大规模训练常用三个简单技巧来压住它。

Z-loss

softmax 交叉熵里，记输出 logits 为 \(z\)，配分函数 \(Z=\sum_j e^{z_j}\)，交叉熵 \(=-z_y+\log Z\)。logits 可以整体平移而不改变概率，所以 \(\log Z\) 原则上可以漂到很大，带来数值问题。PaLM 的 z-loss 加一个惩罚把它拉回 0 附近：

$$\mathcal L_{z}=\lambda\,\log^2 Z,\qquad \frac{\partial\mathcal L_z}{\partial z_j}=2\lambda\log Z\cdot\operatorname{softmax}(z)_j.$$

梯度把所有 logits 往 \(\log Z=0\) 的方向推，PaLM 用的系数是 \(10^{-4}\)。1B 消融显示加 z-loss 不影响 loss 和下游；SmolLM3 最终没用，因为他们的实现带来了还没来得及优化的训练开销。

Embedding 不做权重衰减

权重衰减通常作用于所有参数。OLMo 2 发现把 embedding 排除在外能提升稳定性。理由（Takase et al.）：权重衰减会让 embedding 的范数在训练中逐渐变小，而 LayerNorm/RMSNorm 的雅可比和输入范数成反比——对 \(y=x/\mathrm{rms}(x)\)，有

$$\frac{\partial y}{\partial x}=\frac{1}{\mathrm{rms}(x)}\Big(I-\frac{x x^\top}{d\,\mathrm{rms}(x)^2}\Big),$$

输入越小，经过第一层 norm 的梯度被放得越大，前几层容易出 spike。消融有三组：baseline、embedding 不做衰减、再叠加 NoPE + 文档掩码（检查组合没有负面交互），三者 loss 和评测几乎一样，于是三个改动都用进了 SmolLM3。

QK-norm

QK-norm（Dehghani et al.，ViT-22B）在计算注意力前对 query 和 key 各做一次 LayerNorm，防止注意力 logits 过大，很多新模型用它提升稳定性。但 RNoPE 那篇论文发现它损害长上下文：放到「针」上的注意力更少，无关上下文上更多。原因是归一化去掉了 query–key 点积中的幅度信息，logits 彼此更接近、分布更平，注意力就不够「尖」。所以 SmolLM3 没用 QK-norm（3B 的小模型本来也没那么容易不稳定）。

初始化、激活函数、深度与宽度

MoE：稀疏度、粒度、共享专家、负载均衡

基本结构

混合专家（MoE）的直觉是：预测每个 token 不需要整个模型。做翻译时，学会了代码语法的那部分参数可以不参与。技术目标很简单：增加总参数，而不增加每个 token 的激活参数。粗略地说，总参数决定学习容量，激活参数决定训练成本和推理速度。DeepSeek-V3、Kimi K2 等前沿模型都是 MoE。

从 dense 到 MoE，改动只有一处：把每层（或部分层）的单个 MLP 换成 \(E\) 个 MLP（「专家」），前面加一个可学习的路由器。对 token 表示 \(x\)：

$$p=\operatorname{softmax}(W_rx)\in\mathbb R^{E},\qquad \mathcal T=\operatorname{TopK}(p,K),\qquad y=\sum_{i\in\mathcal T}g_i\,\mathrm{FFN}_i(x),$$

\(g_i\) 是门控权重（通常是被选中专家的 \(p_i\)，有的实现再在 top-K 内归一化；DeepSeek-V3 用 sigmoid 打分再归一化）。每个 token 只跑 \(K\) 个专家，所以有「总参数」和「激活参数」之分。设计 MoE 层要回答两类问题：专家的形状和稀疏度（多少个、多大、每个 token 激活几个、要不要始终激活的共享专家）；以及利用率与专门化（怎么路由才能既让专家都被用上、又鼓励它们分工）。这里关心的是：给定计算预算，什么 MoE 配置让 loss 最低（不是系统吞吐）。

怎么衡量一个 MoE 配置「划不划算」？蚂蚁集团的 MoE scaling law 论文（Tian et al. 2025）定义了一个好用的指标，效率杠杆（Efficiency Leverage, EL）：要达到某个 MoE 配置的 loss，dense 模型需要多少倍的 FLOPs：

$$\mathrm{EL}=\frac{C_{\text{dense}}}{C_{\text{MoE}}}\quad(\text{同一 loss 下}).$$

EL 越高，说明这个 MoE 每单位算力换来的 loss 下降越多。

稀疏度

$$\text{激活比}=\frac{\#\text{激活专家}}{\#\text{总专家}},\qquad \text{稀疏度}=\frac{1}{\text{激活比}}.$$

两个极端都不好：全部激活就退回 dense；激活参数少到极端（比如只有一个参数）显然不够用。计算成本只由激活参数决定：固定激活专家的数量和大小、增加总专家数，FLOPs 基本不变，但容量增加了，只要训得够久，模型通常更好。各家论文的经验结论一致：固定激活专家，增加总专家数（提高稀疏度）能降低 loss，但在很高的稀疏度上收益递减；Kimi K2 和蚂蚁的图都显示了这一点，后者还发现高稀疏度的 MoE 从更多算力中获益更多。看几个公开 MoE 的稀疏度：

模型总专家每 token 激活（含共享）稀疏度
Mixtral-8×7B824.0
OLMoE-1B-7B6488.0
gpt-oss 20b3248
Qwen3-30B-A3B / 235B-A22B128816.0
DeepSeek-V2160+26+220.25
DeepSeek-V3256+18+128.6
gpt-oss 120b128432
Kimi K2384+18+142.8
Qwen3-Next-80B-A3B512+110+146.6

趋势很清楚：越来越稀疏。但最优稀疏度还取决于硬件和端到端效率：Step-3 为了特定硬件和带宽约束，故意没把稀疏度拉满；gpt-oss-20b 稀疏度低是因为端侧内存约束——不激活的专家也要占内存。

粒度

除了多少个，还要决定每个专家多大。常用蚂蚁论文里的定义：

$$G=\frac{\alpha\,d_{\text{model}}}{d_{\text{expert}}},\qquad \alpha=2\ \text{或}\ 4.$$

\(G\) 越大，专家越多越小（总参数固定时）。dense 模型常取 \(d_{\text{ff}}=4d_{\text{model}}\)，若取 \(\alpha=4\)（Krajewski et al. 的约定），\(G\) 可以粗略理解为「拼出一个 dense MLP 的宽度需要几个专家」。蚂蚁论文取 \(\alpha=2\)，下面也用它。几个例子：Mixtral \(2\times4096/14336=0.57\)；DeepSeek-V3 和 Kimi K2 \(2\times7168/2048=7.0\)；Qwen3-Next \(2\times2048/512=8.0\)。注意 \(G\) 随 \(d_{\text{model}}\) 缩放，跨模型比较要小心。蚂蚁的结论：粒度不是 EL 的主要驱动因素，高于 2 时有帮助，存在一个甜点，再往上收益变平。它是个有用的旋钮，近期发布的模型都在往高粒度走，但不应孤立地优化。

共享专家

共享专家是每个 token 都会经过的专家。它吸收数据中基础的、反复出现的模式，让路由专家可以更激进地专门化。一般一个就够，最多两个。粒度越高（从 Qwen3 式往 Qwen3-Next 式走），共享专家越有用；但蚂蚁的图显示总体影响不大。简单的经验法则：用一个共享专家，和 DeepSeek-V3、Kimi K2、Qwen3-Next 的选择一致。

负载均衡

负载均衡没做好，其他设计全白搭。设想 4 张卡、每卡一个专家：如果路由坍缩、所有 token 都去专家 1，只有四分之一的卡在干活；而且没被用到的专家学不到东西，模型的有效容量也下降了。标准做法是给路由器加一个辅助损失。Switch Transformer（Fedus et al.）的形式是

$$\mathcal L_{\text{Bal}}=\alpha\,N_r\sum_{i=1}^{N_r}f_i\,P_i,\qquad f_i=\frac{1}{T}\sum_{t}\mathbb 1[\,i\in\mathcal T_t\,],\quad P_i=\frac{1}{T}\sum_t p_{i}(x_t).$$

乘 \(N_r\) 是为了让完美均衡时损失恰好等于 \(\alpha\)，与专家数无关（top-K 路由时，有的实现再把 \(f_i\) 除以 \(K\)）。\(f_i\) 是实际分到专家 \(i\) 的 token 比例，它来自 argmax/top-K，不可导；\(P_i\) 是路由概率的平均，可导。两者相乘：梯度通过 \(P_i\) 流回路由器，权重是 \(f_i\)——越过载的专家，其概率被压得越狠。完美均衡时 \(f_i=P_i=1/N_r\)。\(\alpha\) 太小引导不够；太大则均匀性压过了语言建模本身。

无辅助损失的均衡（DeepSeek-V3）：给每个专家一个偏置 \(b_i\)，只在选 top-K 时加到亲和度上：选专家看 \(s_{i,t}+b_i\)，但乘到专家输出上的门控值仍用原始的 \(s_{i,t}\)。每步结束后看整个 batch 的负载，过载的专家 \(b_i\) 减 \(\gamma\)，欠载的加 \(\gamma\)。它不往 LM loss 里掺任何别的梯度。

还有一个关键细节：\(f_i\) 和 \(P_i\) 在什么范围统计？每个 worker 的局部 micro-batch，还是全局 batch？Qwen 团队（Qiu et al. 2025）发现，局部 batch 的 token 多样性不够时（比如一个 micro-batch 全是代码），按局部统计会强迫每个小 batch 内部都均匀，反而损害专家的专门化和整体效果。应尽量用全局统计（至少跨设备聚合）。论文发表时，包括 Megatron 在内的很多框架默认是局部统计的。

负载均衡模拟：无均衡 / 辅助损失 / 偏置更新
步数
最大负载 / 平均负载

玩具模型：16 个专家、top-2、每步 1024 个 token。每个 token 对专家 i 的打分 = 固定偏好 skew·ui + 标准正态噪声（ui 在 [−1,1] 均匀铺开，模拟「天生热门」的专家）。辅助损失模式按 Switch 损失对路由 logits 的解析梯度更新一个可学习的修正项；偏置模式只在选 top-2 时加 bi，每步按负载符号 ±γ 更新。没有模拟语言建模损失本身，所以这里两种方法都能拉平负载；真实训练中两者的区别在于辅助损失的梯度会和语言建模的梯度混在一起，偏置更新则不会。

MoE 的结构消融很难做，因为因素之间有交互（比如共享专家是否有用取决于粒度）。还有一些值得深入的方向：零计算专家、MoE 层重缩放和训练监控（LongCat-Flash）；正交损失式的均衡（ERNIE 4.5）；在训练中调度均衡系数；优化器排名在 MoE 上是否改变；怎么把 μP 用到 MoE；专家每个 batch 看到的 token 数不同，学习率该怎么调；开头放几层 dense 层。

混合模型：线性注意力与 SSM

另一个趋势是在 dense 或 MoE 里加入状态空间模型（SSM）或线性注意力层（Falcon-H1、MiniMax-01）。它们想解决 Transformer 处理超长上下文时的根本低效。循环模型能以线性代价处理任意长的上下文，但未必能充分利用上下文；Transformer 能很好地利用上下文中的模式，但在长上下文上代价很高。Waleffe et al. 研究 Mamba 发现它在很多基准上表现不错，但 MMLU 等少数任务较弱，推测是缺少上下文学习能力。把 SSM 和注意力块组合起来，就是混合模型。

线性注意力：交换乘法顺序

推理时 token \(t\) 的注意力输出是

$$o_t=\sum_{j=1}^{t}\frac{\exp(q_t^\top k_j)\,v_j}{\sum_{l=1}^{t}\exp(q_t^\top k_l)}.$$

去掉 softmax，再交换求和与乘法的顺序：

$$o_t=\sum_{j\le t}(q_t^\top k_j)\,v_j=\Big(\sum_{j\le t}v_jk_j^\top\Big)q_t=S_tq_t,\qquad S_t=S_{t-1}+v_tk_t^\top\in\mathbb R^{d\times d}.$$

左边的写法是「对每个历史 token 算一个标量、缩放它的 \(v\)、再加起来」，第 \(t\) 步要 \(O(td)\)，生成 \(T\) 个 token 共 \(O(T^2d)\)。右边维护一个状态矩阵 \(S_t\)，它已经总结了所有历史的 \((k,v)\)；每来一个新 token 做一次外积更新（\(O(d^2)\)），输出是一次矩阵–向量乘（\(O(d^2)\)），共 \(O(Td^2)\)。用对序列长度的依赖换来对维度的依赖，而且推理时状态大小固定，不再有随长度增长的 KV cache。训练时（不考虑因果掩码）同样的道理就是矩阵乘法结合律：

$$\underbrace{(QK^\top)}_{T\times T}V=Q\underbrace{(K^\top V)}_{d\times d}.$$

（带因果掩码时不能直接这样换，实际实现用分块：块内用二次形式，块间传递状态。）这看起来就是一个 RNN。但 softmax 起着重要的稳定作用，朴素的线性注意力不加归一化会不稳定。

Lightning Attention

NormAttention（Qin et al. 2022）用基于范数的缩放代替 softmax 归一化：\(\mathrm{RMSNorm}(Q(K^\top V))\)。Lightning Attention 在此基础上着重高效实现，并做了几处结构调整：Q、K、V 过 SiLU，门 \(G\) 过 sigmoid；状态带衰减 \(KV_t=\lambda KV_{t-1}+k_t^\top v_t\)，\(o_t=q_tKV_t\)；最后 \(Y=G\odot\mathrm{RMSNorm}(O)\)。MiniMax-01 报告混合了 Lightning Attention 的模型在大多数任务上与 softmax 注意力相当，在 NIAH 这类检索任务上反而好很多，暗示 softmax 层和线性层之间有某种互补。

但也有反例：MiniMax-M2 没有用混合或线性注意力。据其预训练负责人的说法，M1 早期在小规模、当时流行的基准（MMLU、BBH、MATH）上看起来不错，但到大规模发现它在复杂多跳推理上有明显缺陷；他们还提到 RL 训练中的数值精度问题和基础设施成熟度。他们也承认，随着 GPU 算力增长放缓而数据长度持续增长，线性和稀疏注意力的好处会逐渐显现。这说明结构消融的复杂性，以及研究和生产之间的距离。

带门控的线性注意力

从循环模型学到的一课是：状态要能偶尔「忘掉」过去。几乎所有新的线性注意力都加了一个对旧状态的门 \(G_t\)：

$$S_t=G_t\odot S_{t-1}+v_tk_t^\top.$$

不同模型的区别在于 \(G_t\) 怎么参数化（Yang et al., GLA 整理过一张表）：Mamba-2 用一个依赖输入的标量 \(\gamma_t\)（整个状态统一衰减）；mLSTM 和 Gated Retention 也是标量门（\(\sigma(x_tW_\gamma)\) 及其幂）；GLA、HGRN-2、RWKV-6 用沿 key 维的向量门 \(\alpha_t^\top\mathbf 1\)（不同维度不同衰减）；Mamba 的门由 \(\exp(-\alpha_t\odot\exp(A))\) 给出。Mamba-2 被用在 Nemotron-H、Falcon-H1、Granite-4.0-h 等混合模型里。Qwen3-Next 用了 gated DeltaNet，报告长上下文推理更快、训练更快、标准基准也更好；Kimi 也提出了 Kimi Delta Attention。另一条路是稀疏注意力，只对选中的块或 query 计算注意力：Native Sparse Attention、DeepSeek Sparse Attention、InfLLM v2。

Dense、MoE 还是 Hybrid

优点缺点什么时候用
Dense支持最广、理解最透、训练稳定、单位参数性能好计算随规模线性增长：70B 的成本约是 3B 的 23 倍内存受限、或第一次训 LLM 的默认选择
MoE训练和推理的单位算力性能更好内存需求高（所有专家都要加载）；训练更复杂；分布式要处理专家放置、负载均衡、all-to-all 通信不受内存约束、想要单位算力的最佳性能
Hybrid长上下文可能更好；超长序列更高效不如前两者成熟，验证过的配方少，框架支持有限要扩到超长上下文、同时降低推理开销
模型在哪里运行？ 端侧 / 内存受限多数情况：Dense有经验的团队：可考虑 Hybrid 内存充足 → 团队经验？ 第一次训Dense，把基础打牢 熟悉 dense→ 时间线？ 非常有经验MoE 或MoE + Hybrid 紧：需要验证过的路线Dense 宽松：可以探索MoE 或 MoE + Hybrid
结构选择的决策流程：先看部署位置，再看团队经验和时间线。

（也有团队在探索文本扩散模型，目前还太早期。）SmolLM3 的情况：端侧部署、约三个月工期、过去主要训 dense。内存约束排除了 MoE，工期短和「dense 已能支撑 128k 目标上下文」排除了 hybrid，所以选了 Llama 式 dense。

SmolLM3 的结构

SmolLM3 以 SmolLM2 1.7B 的结构为起点，用 Qwen2.5-3B 的布局在 100B token 上训了一个 3B 消融模型作 baseline，逐个测试改动。每个改动要么提升英文基准上的 loss 和下游表现，要么在不降质量的前提下带来可测量的收益（比如推理速度）。最终进入正式训练的：

对照它的 config.json：hidden_size 2048、num_hidden_layers 36、num_attention_heads 16、num_key_value_heads 4、intermediate_size 11008、tie_word_embeddings true、no_rope_layer_interval 4、vocab_size 128256，和上面一一对应。

结构上的经验可以浓缩成一句话：用例决定选择。具体来说：让部署目标指导结构决策；在创新和务实之间平衡——不能忽视已被验证的进步（有 GQA 还用 MHA 是糟糕的选择），也不要追逐每篇承诺边际收益的论文；系统胜过直觉，每个结构改动都要验证；规模效应是真实的，算力允许就在目标尺寸上重新消融。

参考与延伸阅读