① dense FFN:不管什么 token,全员上岗
第 3 章讲过:注意力负责 token 之间"沟通",FFN 负责每个 token 单独"思考"。
但这个思考很浪费 —— 随便点下面哪个 token ,看它过 FFN 时点亮了多少参数:
↓ 进入同一个 FFN(隐层 4·n_embd = 512 个神经元,下图 1 点 = 4 个)
点一个 token 试试
一个平平无奇的空格 ,和一个信息量很大的字母,动用的是同一套全部 FFN 参数 。
模型越大这笔账越吓人:GPT 这类架构里 FFN 占了参数的大约 2/3 ,而它对每个 token 都是全量计算。
为什么偏偏对 FFN 下手,不拆注意力?
注意力的工作是 token 之间 的交互,拆开会破坏"谁都能看谁"的沟通;
FFN 是逐位置独立 的加工 —— 每个 token 各算各的,谁跟谁都不打招呼。
这种"天生并行、互不依赖"的结构才方便按 token 派给不同专家。
(对注意力做稀疏化的研究也有,但主流 MoE 都是拆 FFN。)
↳ 下一步:把这个 512 隐层的大 FFN 拆成 4 个 256 隐层的小 FFN,每个 token 只走 2 个 —— 谁来决定走哪 2 个?
② 路由器:一个线性层,给每个 token 派单
路由器 只是一个
n_embd → n_expert 的线性层:token 向量进去,4 个分数出来,softmax 变成概率,
只留最高的 top_k=2 个 ,权重重新归一化。选一个 token,看它被派给谁
(下面的概率是 09_moe.py 训练后 layer 0 的真实路由统计,按字符类别汇总):
输出 = 选中专家的加权和 :out = w·Expert_i(x) + w'·Expert_j(x),权重就是归一化后的路由概率。
没被选中的专家这个 token 完全不算 —— 这就是"参数在那儿,但不花算力"的稀疏激活。
top-k 是"挑选",挑选不可导,路由器怎么学?
梯度走的是权重 那条路:被选中的专家,其输出乘的系数 w 来自 softmax,是可导的。
loss 对 w 的梯度会一路传回路由器 —— "这次派给 E1 效果好,下次给 E1 打分再高点"。
没被选中的专家确实拿不到梯度,这正是下一步"强者恒强"问题的根源。
派单的单位是 token 还是整句话?
是 token 。同一句话里,空格可能走 E1、字母 e 走 E2 —— 每个位置独立决策。
代码里先把 (B,T,C) 摊平成 (B·T, C) 再路由,就是在强调"逐 token 派单"。
↳ 代码:09_moe.py 的 MoELayer.forward :router 打分 → topk → 归一化 → 加权求和
↳ 下一步:派单派了 5000 步之后,4 个专家真的形成分工了吗?
③ 分工是训出来的,不是设计出来的
没有任何人规定"E1 管空格、E2 管元音"。但训练 5000 步后,把 val 集的 token 按类别统计
"每类字符被派往哪个专家"(layer 0 真实数据),分工自己浮出来了:
每行加起来 = 100% · 颜色越深 = 该类字符越爱找这个专家 · 鼠标悬停看准确数值
每个专家最常接手的字符(同一次训练的真实统计)
这些分工是涌现的: 换个随机种子,分工方案会不一样,但"有分工"这件事稳定出现 ——
路由器 + 专家在训练里互相适应:专家碰巧对某类 token 加工得好,路由器就更爱把这类 token 派给它,它就练得更专。
专家 = "数学专家 / 法律专家"吗?
不是。大模型里的专家分工大多发生在低层统计模式 上(词形、标点、语言、代码 token 之类),
很少出现人类意义上的"学科专家"。Mixtral 论文自己也承认:看不出专家按主题分工。
本页字符级模型的"空白/元音"分工,就是这种统计分工的微缩版。
↳ 下一步:"接单多 → 练得好 → 接单更多"是个正反馈。它有一个危险的极端……
④ 偏心的尽头:专家塌缩,和治它的 aux loss
正反馈不加管束,路由器会把越来越多 token 塞给少数专家 ——
被冷落的专家拿不到梯度、逐渐报废,这叫专家塌缩 。
要看清它,得用最脆的设定:top-1 路由 (Switch Transformer 同款,aux loss 就是它提出的)。
下面是同一份代码、都用 top-1 跑的两次真实训练 ,唯一区别是负载均衡 开关。点 ▶ 看 layer 0 的负载怎么随训练演化:
top-1 · aux = 0.01(开)
top-1 · aux = 0(关,放飞路由器)
完美均衡 = 每个专家 25%(虚线为参考)· 数据:layer 0,每 500 步采一次 · 更深的层塌得更狠,见下方折叠
▶ 播放训练过程
单步 +500
重置
诚实交代:本关的 top-2 配置,关掉 aux 其实也没塌
实测(同 5000 步):top-2 + 4 专家 就算 --aux 0,最终负载也只是 20/24/27/29% 的轻微偏,
val loss 1.549 和开 aux 打平。原因:top-2 下每个 token 同时养 2 个专家 ,4 选 2 让谁都分得到梯度,天生抗塌。
专家越多、top_k 越小,饿死人的风险越大 —— 所以 top-1 的 Switch 最先撞上这个问题,
每层 256 选 8 的 DeepSeek-V3 也在做负载均衡(STEP 6 另有一组 16 选 6 + 共享专家、完全不做均衡的实测)。上面的 top-1 实验里,偏载对 val loss 的伤害也很小
(1.545 开 / 1.560 关,4 层玩具上差 0.015)—— 塌缩真正的代价是买了的容量白买 ,
以及大规模多机训练时忙闲不均拖慢整批,这两笔账在大模型上才要命。
上图只画了 layer 0,别的层呢?(深层塌得更彻底)
路由是每层各自独立学的,所以每层塌的程度不一样,而且
越深越极端 。
同一次 top-1 训练跑完 5000 步,四层的终局:
关掉 aux,layer 3 有一个专家吃掉
72% 的单、另一个只剩
2% ,基本等于废掉;
开着 aux,四层全部稳在 21–31%。所以"塌缩"不是上图那条曲线的全部,layer 0 反而是塌得最轻的一层。
aux loss 的公式为什么长这样?
aux = n_expert · Σ f_e · P_e:f_e 是实际派给专家 e 的 token 占比(不可导),
P_e 是路由器给 e 的平均概率(可导)。两者都均匀时乘积和最小(=1);
谁的 f 偏大,压它的 P 就最能降 loss —— 梯度顺着可导的 P 把"过热"的专家分数往下压。
这是 Switch Transformer 的经典设计,系数 aux_coef=0.01,轻轻管一下就够。
塌缩之后,MoE 退化成了什么?
退化成一个只有少数专家在干活的小 dense 模型 :参数表上还是 4 个专家,
实际容量缩水一半甚至更多 —— 花了 MoE 的显存,买到的容量白买,最亏的局面。
所以工业级 MoE(Switch / Mixtral / DeepSeek)全都带某种负载均衡手段。
↳ 代码:09_moe.py 的 MoELayer.forward 末尾 aux loss · 复现本组对照:python 09_moe.py --top-k 1 --aux 0
↳ 下一步:分工有了、均衡有了,回头算总账 —— 这笔交易到底赚在哪?
⑤ 算账:用 dense 的算力,买 2 倍的容量
MoE 的本质是一笔交易:总参数(容量)↑,每 token 激活 的计算量不变 。
拨动两个滑块,看这笔账怎么随 n_expert / top_k 变化(模型 = 本关的 4 层 GPT):
MoE 省的是算力,不是显存 —— 总参数一个都不能少,全都得放在(显)存里待命;
省下的是每个 token 的 FLOPs。所以它适合"显存管够、算力金贵"的大规模训练与推理。
真实世界同款账本:Mixtral 8×7B 总 46.7B / 激活 12.9B;DeepSeek-V3 总 671B / 激活 37B。
达标 本关实测(四次训练 · 各 5000 步 · tiny shakespeare)
✓ MoE val loss – ,dense 对照(03)~1.57 :同等每 token 算力下,容量翻倍换来更低的 loss。
✓ 负载均衡有效:top-1 关 aux 实测偏载滚雪球(最深一层单专家吃掉 72% 的单,同层另一个只剩 2% ;layer 2 也有专家被饿到 9.4%),开 aux 则四层全部稳在 21–31%;本关的 top-2 配置天生抗塌(关 aux 也只轻微偏)。
✓ 专家分工涌现且清晰:空白 70.5% 涌向 E1、大写 50.8% 归 E0、元音偏 E2、辅音偏 E3(见 STEP 3 热力图,val 集实测)。
Q&A 常见疑问
MoE 的"专家"和多头注意力的"头",都是"多个小的",区别在哪?
头是全员上岗,专家是按需点名。 多头注意力里,每个 token 都要过所有 的头,只是各头看不同关系;
MoE 里每个 token 只走 top-k 个专家,其余专家对它零计算。
一个是"并行分视角"(算力照付),一个是"稀疏选路"(算力打折),目的完全不同。
现在的 MoE 还是"4 选 2"这么糙吗?
前五步教的是 Mixtral 同款的经典配方(少量大专家 + top-2)。DeepSeekMoE(arXiv 2401.06066)在它上面加了两个旋钮:
细粒度专家 ,把每个专家切窄,个数和每 token 选中的个数同比放大(例如 16 选 2 → 64 选 8,激活的隐层总宽不变,
组合数从 C(16,2)=120 变成 C(64,8)≈44 亿);共享专家 ,留出几个不经路由器、每个 token 必走的专家。
DeepSeek-V3 每层 256 个路由专家选 8 个,外加 1 个共享专家(arXiv 2412.19437 §4.2);
Qwen3-Next-80B-A3B 每层 512 个专家激活 10 个,另有 1 个共享专家(模型卡)。
STEP 6 把这两个旋钮和 V3 的偏置均衡放到本关的 4 层模型上各跑了一遍。
既然容量翻倍这么香,为什么不是所有模型都用 MoE?
代价在工程:显存占用照总参数算 、多机并行时 token 要跨设备找专家(通信开销)、
负载不均会拖慢整批训练、微调也比 dense 娇气。小模型(显存不紧张)用 dense 更省心;
MoE 的甜区是"参数规模顶到显存上限、还想再大"的场景 —— 所以你在千亿级模型上才总看到它。
↳ 跑法:python 09_moe.py · 对照塌缩 python 09_moe.py --aux 0 · 训完自动打印负载 + 分工报告
↳ 下一步:DeepSeek 系把专家切得更碎、留一个人人都过的共享专家,还换了一种不进 loss 的均衡办法。放到这台 1.35M 的小模型上,它们赢过 4 选 2 了吗?
⑥ 切碎专家、留一个常驻专家、换一种均衡
DeepSeek 系 MoE 在经典配方上改了三处:细粒度专家与共享专家 (DeepSeekMoE,arXiv 2401.06066),
以及偏置均衡 (DeepSeek-V3,arXiv 2412.19437)。
三处都不改每 token 激活的 FFN 隐层总宽。点下面的布局 ,看 dense FFN 的 512 个隐层神经元怎么分出去:
布局 每 token 激活的隐层都是 512
dense FFN
经典 4×256 · 选 2
细粒度 16×64 · 选 8
16×64 · 选 6 + 1 个共享专家
全部专家(格子宽度 ∝ 隐层宽度 · 橙 = 这个 token 选中 · 绿 = 共享专家,每个 token 必走)
这个 token 实际用到的隐层
换一个 token
选中哪几个是随机示意;真实训练里由路由器打分决定
四种布局的每 token 激活参数相差不到 1.2 万个:路由器从每层 128×4 变成 128×16,专家的偏置项也多了一些。
均衡 偏置只管挑谁,不改权重
挑选 选中的 k 个 = TopK( pe + be )
权重 we = pe / Σ选中 pj ← 用原始概率 p,不含 b
更新 be ← be − γ · sign( loade − mean(load) )每个训练步一次 · γ = 0.001 · 不进 loss
占位数字:4 个专家选 2,概率 p = 0.40 / 0.30 / 0.20 / 0.10 只为上屏;下面的真实训练是 16 个专家选 6。
DeepSeek-V3 论文原文怎么写这条规则?
挑选:s_{i,t} + b_i ∈ Topk({s_{j,t} + b_j}, K_r)。原句:"the bias term is only used for routing. The gating value, which will be multiplied with the FFN output, is still derived from the original affinity score s_{i,t}"(arXiv 2412.19437 §2.1.2)。
实测 五次训练 · 同一份 09_moe.py · 各 5000 步
RTX 4090 · CUDA · PyTorch 2.4.1,每组一个种子。STEP 1–5 的数字来自早先 Mac MPS 上的训练,硬件不同,两边不逐位比较(STEP 5 的 1.549 和下面经典组的 1.5392 是两台机器各跑一次的结果)。
经典 4×256 · 选 2
细粒度 16×64 · 选 8
共享 · aux
共享 · bias
共享 · 不均衡
val loss @ step 4999(越短越低)· 横轴从 1.52 起,差距被放大 · 点柱子也能切换
各层终局负载
layer 0 训练过程
layer 0
layer 1
layer 2
layer 3
训练步 4999
红 = 最忙 · 蓝 = 最闲 · 虚线 = 均匀负载 · 纵轴上限 = 均匀值的 2 倍
结论: 在 1.35M 参数、tiny shakespeare 字符级、5000 步这个规模上,细粒度切分和共享专家都没有赢过经典 4 选 2
(纯细粒度高 0.023,加共享专家的三组高 0.005–0.011);三种均衡方式之间差不到 0.006,在单种子噪声的量级内。
差 0.006 能说明哪种均衡方式更好吗?
不能。每组只跑了一个种子,val 在同一次训练里也不是单调下降:细粒度组 step 4500 是 1.5604,step 4999 反而是 1.5625。
三个共享专家组的初始化完全相同(step 0 的 val 都是 4.3544,layer 0 负载逐位相同),差别来自均衡方式加上训练中的随机性。
要分出高下,得每组多跑几个种子取平均。
完全不做均衡,为什么没塌?
STEP 4 的 top-1 + 4 专家关掉 aux 后,最深一层有一个专家吃掉 72% 的单。这里是 16 个路由专家选 6 个、外加 1 个共享专家,
关掉所有均衡,四层终局负载仍在 3.1–9.2% (均匀 6.25%);layer 0 的训练过程里,偏载在 step 2000 前后成形,之后最忙的专家从 8.1% 缓慢升到 8.9%。只能说在这个设定、这个训练长度下没塌 。每层 256 选 8 的 DeepSeek-V3 仍然做均衡。
另外,终局负载取自最后一个训练 batch,是一次快照。
偏置都学成了正数,是把所有专家一起抬高了吗?
top-k 只比较 p + b 的相对大小:给所有专家加同一个常数,选中的专家不变。所以要看的是
同一层内偏置之间的差 。
更新规则每步给每个专家 +γ 或 −γ,欠载的专家多于过载的专家时,整层偏置就一起往上漂,绝对值本身不带信息。偏置均衡组的终局:
layer 0 的偏置差 0.027,和平均概率 1/16 ≈ 0.0625 同一量级,这一层的负载也是 16 专家四组里最整齐的;
layer 1–3 的偏置差只有 0.003–0.007,终局负载也没有比 aux 组(4.3–8.0%)更齐。这组数据判断不了原因。
一个没有验证过的猜测:固定步长的 sign 更新让这几层的专家在过载和欠载之间来回翻,偏置差攒不起来。
DeepSeek 在多大的规模上用这套配方?
本页 :4 层,每层 16 个路由专家(隐层 64),训练约 2048 万个字符 token(5000 步 × batch 64 × block 64)。DeepSeek-V3 :61 层、hidden 7168,前 3 层之外全换成 MoE,每层 1 个共享专家 + 256 个路由专家(隐层 2048)选 8 个,
预训练 14.3T + 500B token(arXiv 2412.19437 §4.2)。无 aux 均衡的方法论文 :实验最大 3B 参数、200B token(arXiv 2408.15664)。Qwen3-Next-80B-A3B :每层 512 个专家激活 10 个,另有 1 个共享专家(HF 模型卡)。
↳ 代码:09_moe.py 的 MoELayer.forward :sel_bias 只进 topk 挑选,权重用 probs.gather 取原始概率,训练时按负载加减 γ · 复现(在 phase3-moe/ 下):python 09_moe.py --n-expert 16 --expert-hidden 64 --top-k 6 --shared 1 --shared-hidden 128 --balance bias,把 bias 换成 aux / none 得另两组;细粒度组是 --n-expert 16 --expert-hidden 64 --top-k 8
🎉 手搓 MoE · 通关
你已经亲手把 dense GPT 的 FFN 换成"路由器 + 专家",看到了分工的涌现、塌缩的风险、"参数翻倍算力不变"的账本,也在同一台小模型上试了细粒度专家、共享专家和偏置均衡。下次再看 DeepSeek-V3 的"256 选 8 + 1 个共享专家",你知道每个数字管什么了。