开源 AI 论文复现实验与代码解读 · 第五轮 · 070 检索与实验日期：2026-09-11。本文提供可运行的 CPU 机制实验，区分论文报告、本次实际验证与作者推断。
摘要
训练日志中的辅助损失下降，专家就一定工作正常吗？本文从 Switch Transformer 的单专家路由出发，用三十二个参数的混合专家模型，把软概率、离散选择、容量截断和任务误差拆开检查。最小实验既展示“损失等于一但全部选择同一专家”的反例，也记录九组实际训练。目标是让读者能解释一张路由热力表，而不只是画出颜色好看的图。
目录
- 复现价值与证据边界
- 核心思想与公式
- 官方代码阅读路线
- 最小实验与实际结果
- 评测协议
- 失败排查
- 后续科研问题
- 总结与参考资料
复现价值：先知道自己在数什么
MoE，即 Mixture of Experts（混合专家），让输入只调用部分专家参数。这里的“专家”是一组可训练函数，并不自动等于数学、编程或某个语种的知识模块。Router（路由器）依据输入给专家打分；top-1 表示每个 token 只选择分数最高的一位。Token 在本文实验里是一行合成特征，不是自然语言词元。
**论文报告：**Fedus、Zoph、Shazeer 的 Switch Transformers 发表于 JMLR 2022，研究以单专家路由扩展稀疏模型，并讨论容量、通信与训练稳定性。论文的预训练加速结论依赖其模型和计算条件，不能从激活专家数直接推导本机速度。论文页面
**本次实际验证：**训练的是单层、标量输出的合成回归模型；没有预训练权重、语言任务或分布式通信。**作者推断：**对科研新人而言，先复现“选择是否正确、溢出如何统计、梯度流向哪里”，比先复刻大模型参数量更容易得到可信证据。
核心思想与公式：三种负载不能混为一谈
令输入为 $X\in\mathbb{R}^{T\times d}$，路由参数为 $W\in\mathbb{R}^{d\times E}$，则：
$$ Z=XW,\quad p_{ti}=\operatorname{softmax}(Z_t)i,\quad a_t=\arg\max_i p{ti}. $$
$T$ 是路由组内 token 数，$d$ 是输入维度，$E$ 是专家数。$Z$ 是归一化前分数，$p$ 的形状为 [T,E]，每行和为一；$a$ 是长度为 $T$ 的整数向量。本实验分别取六十四、四、四，输入最后一列恒为一，使路由器能够学习偏置。
先数容量限制前的选择比例，再统计平均概率：
$$ f_i=\frac{1}{T}\sum_t\mathbf{1}[a_t=i],\quad P_i=\frac{1}{T}\sum_t p_{ti},\quad L_{bal}=E\sum_i f_iP_i. $$
$f$ 是硬选择频率，$P$ 是软概率质量，均为长度 $E$ 的向量。反向传播通过 $P$，不通过离散的 $f$。本文把系数 $\alpha$ 放在总损失外层；原文公式四把系数写进辅助项，比较日志时必须确认是否已经乘过。Switch 正文公式四至六
当两向量都均匀时，未加权的 $L_{bal}=1$；但反过来不成立。把六十四行概率都设成 (0.25,0.25,0.25,0.25)，代码的并列规则全部选第零位，于是 $f=(1,0,0,0)$，损失仍为一。这是代入公式得到并实际执行的反例，不是对训练最终状态的假设。NumPy 明确规定并列最大值返回第一次出现的位置。argmax 文档
容量 $C=\lceil cT/E\rceil$ 限制每个专家实际接收多少 token，$c$ 是容量系数，向上取整是本文的实现约定。设 $m_t$ 表示是否被接收，本实验输出为：
$$ \hat y_t=0.25X_{t0}+m_t p_{t,a_t}v_{a_t}^{\top}X_t. $$
$v_i\in\mathbb{R}^{d}$ 是第 $i$ 个线性专家；第一项是固定残差路径。溢出时仅跳过专家分支，样本仍参与任务误差。容量截断后的接收计数之和可能小于 $T$，不能冒充前面的选择频率。
ST-MoE 另提出 router z-loss，本文写为 $L_z=T^{-1}\sum_t(\log\sum_i e^{Z_{ti}})^2$。变量仍是路由分数矩阵；它惩罚对数归一化常数，论文用它研究数值稳定性。它不是“专家均衡”的另一种写法，也不等价于给每个分数逐项做平方惩罚。ST-MoE 全文第三节
官方代码阅读路线：追踪掩码的生命周期
先读原论文指向的 Mesh TensorFlow moe.py，从 _switch_gating 依次定位概率、单专家选择、density_1、density_1_proxy 和累计位置。辅助负载项先计算，之后才按容量修改掩码。其主分支是本次访问的快照，未固定提交；需要逐位复现实验时，应另存提交号。作者代码
再读固定 v4.57.1 标签的 Transformers 实现，这是官方框架移植，并非作者的原始训练程序。沿 SwitchTransformersTop1Router.forward 查看 argmax、cumsum 与容量掩码；这里按序列位置累计，没有执行类注释所说的按概率排序。读源码时应核对实际运算，而不只转述注释。
随后追踪 SwitchTransformersSparseMLP.forward：它对截断后的 router_mask 再取 argmax，全零溢出行会得到零号索引。若把这个返回值直接当作原始选择，日志可能混入错误计数。本文用“原本全选二号”的构造检查了这种信息损失；这不是宣称运行了完整框架。另需留意 _unpack_router_logits 会沿序列维拼接层的记录，诊断单层失衡时应在拼接前取数据。固定版本源码
图是诊断设置示意，珠子和柱高不对应实测样本。上方读出原始选择，下方分别记录截断前后负载和溢出；同色代表同一专家。被截断的柱子只能变短，图中没有宣称专家经过训练变得均衡。真实概率热力表由脚本另行输出。
最小实验：反例、训练与真实热力表
进入本文 code/，使用 Python 3.12 执行：
python -m pip install -r requirements.txt
python minimal_moe.py --check-only
python minimal_moe.py
python audit_results.py
依赖固定为 NumPy 2.3.5。实际运行环境为 Python 3.12.14、CPU、双精度。数据种子七千，训练、验证、测试各有二百五十六、一百二十八、二百五十六行，集合互不重叠。前三列独立采样标准正态；前两列符号划分四个区域，区域决定第三列的斜率和截距。区域标签只生成目标，不作为路由监督。
路由矩阵和四个专家共三十二个参数。每个训练种子使用相同初值和抽样序列比较三种损失：任务均方误差、加 $0.1L_{bal}$、再加 $0.001L_z$。初始化把零号路由偏置加二，刻意制造拥挤；Adam 学习率零点零三，每步无放回抽六十四行，共六百步，训练容量系数一点二五。三个种子为七十至七十二，不因结果好坏删掉任何一组。
**本次实际验证：**完全均匀概率反例中，六十四项全选零号，容量二十，只接收二十项，溢出率为 68.75%。另一构造把原始负载 [28,12,12,12] 截到每专家八项，得到 [8,8,8,8]；均匀的接收柱状图隐藏了原始拥挤。全分数加一百后概率和选择不变，但 z-loss 改变，进一步说明它不是负载指标。
九组训练的固定终点测试结果如下；每项为三个种子的均值与样本标准差，均衡损失、z-loss 均未乘权重：
MSE 是预测值与真实值之差的平方平均，越小越好；溢出比例则以全部有效输入为分母。两者回答不同问题：一个衡量任务误差，一个衡量专家分支被跳过的频率。实验不计运行速度，因此也没有据此报告吞吐提升。
| 训练目标 | 测试 MSE | 溢出比例 | 均衡损失 | z-loss |
|---|---|---|---|---|
| 仅任务 | 0.982 ± 0.166 | 0.154 ± 0.065 | 1.295 ± 0.144 | 41.59 ± 8.18 |
| 加均衡项 | 0.933 ± 0.229 | 0.133 ± 0.101 | 1.298 ± 0.170 | 37.83 ± 5.74 |
| 再加 z 项 | 0.978 ± 0.048 | 0.118 ± 0.030 | 1.213 ± 0.041 | 21.09 ± 2.81 |
这组观测不支持“加均衡项就必然降低均衡损失”。均衡组的种子七十一在测试中选择计数为 [32,2,116,106]，二号与三号拥挤，一号几乎不用；但不能把有限窗口内低频使用直接称为永久塌缩。加 z 项后均值降低，也不证明复现了大模型低精度稳定性收益：本次没有出现或测试那种硬件数值故障。
打开 实际路由热力表，可展开九组运行，逐行查看第一个固定测试批的四个概率、原始选择与是否接收。色深等于概率，数字保留四位小数；热力表只展示六十四行，汇总表则用全部测试行。完整精度概率、预测、权重和训练曲线均在 results/，颜色不能代替这些数据。
评测协议：同时报告任务、分配和预算
本文预先固定六百步终点；验证集只记录曲线，不用于挑选检查点，测试集也不用于调系数。每六十四行独立执行容量限制，再平均批指标。尤其最大选择比例和变异系数，应先逐批计算；先汇总全测试集再计算可能掩盖局部拥挤。标准差只描述初始化和抽样变化，不是置信区间。
还要分开记录每个专家的原始计数、接收计数、零负载次数以及平均 token 熵。熵高代表单个输入的路由不确定；它可能与固定并列造成的硬选择集中同时出现。本文把塌缩视为待诊断的持续失衡现象，没有设一个通用阈值就宣布专家死亡。
保存权重后，本次额外扫了五种测试容量，保持输入及顺序不变。均衡组种子七十一的系数从一点二五改到四时，溢出从 24.22% 降为零，MSE 从 1.1139 降为 0.5358。这隔离了该模型对容量的敏感性，代价是允许更多专家计算；不能作为同预算方法改进，也不等价于重新训练。
失败排查：先排除观测错误
遇到全部选同一专家，先检查分数并列、初始偏置、有效 token 遮罩和数据顺序，再调整辅助系数。本文没有 padding；迁移到语言批次时，填充位置必须从分母和计数中排除，并保持不同方法的有效 token 集一致。
遇到均衡曲线下降但任务变差，应同时看未加权损失和加权贡献。回归误差的数值尺度与语言交叉熵不同，不能直接照搬论文系数。容量越紧，越可能让序列后部跳过专家；重排 token 可以改变受影响的样本，因此评测必须保存批内顺序。
代码通过全部参数的中心有限差分检查，最大绝对误差 2.79×10⁻¹⁰，扰动前后离散选择与容量掩码保持一致；这验证局部梯度，不宣称 argmax 可导。独立标量审计又从保存权重重算四十五组容量评测及一万一千五百二十个预测。官方框架运行、语言模型训练、专家并行、低精度和论文基准均为待人工核验。
后续科研问题
第一，专家真的分工了吗？在固定数据分组上统计路由关联后，再交换专家或屏蔽专家，观察任务变化。仅凭某种颜色集中不能推出语义专长；路由差异也可能由位置或词频驱动。
第二，容量造成的错误是否集中在某些样本？可以预先定义区域和位置切片，比较溢出样本与接收样本的误差，再做固定预算下的顺序扰动；需要保留配对样本，避免用不同测试集解释机制。
第三，局部均衡与全局均衡是否冲突？把同一批数据分到不同路由组，比较逐组统计和合并统计。在进入多机实验前，这个小改动就能暴露统计口径问题；是否改善通信效率仍须另行实测。
总结
理解 MoE 路由，需要把“想送给谁”“实际送了多少”“样本最后表现如何”连在一起。本文的反例和负结果说明，任何单一辅助损失都不足以担保专家健康。可复用的产物是可回放的选择、掩码、权重和评测协议；接入真实模型时，应首先守住这些检查点。
参考资料
检索日期：2026-09-11。以下一手链接均已实际打开；正文说明各自支持的论点。
- William Fedus、Barret Zoph、Noam Shazeer，Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity，JMLR 23(120)，2022。出版页面；出版全文。
- Barret Zoph 等八位作者，ST-MoE: Designing Stable and Transferable Sparse Expert Models，2022，arXiv v2，修订于 2022-04-29。版本记录；全文。
- TensorFlow Mesh 贡献者，Switch 路由原始代码，访问时 master，未固定提交。moe.py。
- Hugging Face 贡献者，Transformers v4.57.1，固定标签源码。Switch 实现。
- NumPy 贡献者，NumPy 2.3 版本文档。argmax 并列规则。