LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts
Venmugil Elango, Nidhi Bhatia, Roger Waleffe, Rasoul Shafipour, Tomer Asida, Abhinav Khattar, Nave Assaf, Maximilian Golub, Joey Guman, Tiyasa Mitra, Ritchie Zhao, Ritika Borkar, Ran Zilberstein, Mostofa Patwary, Mohammad Shoeybi, Bita Rouhani
cs.LG, cs.AI
2026-01-26
把 MoE 的专家计算从隐藏维投影到更小的潜空间,同时按压缩比放大专家数和 top-k,在 95B/8B 规模上同 FLOP 下精度反超标准 MoE,该架构已被英伟达 Nemotron-3 旗舰模型采用。
MoE 是当下开源和闭源大模型的标配,但「精度每 FLOP」「精度每参数」这两条效率曲线,标准 MoE 离最优还有多远,没人讲清楚。作者从软硬协同的角度拆了三道瓶颈。低延迟、小 batch 时,MoE 主要被「加载权重的显存带宽」卡住,典型部署里每个专家分到的 token 只有几百,远不够把 GPU 跑到算力瓶颈。高吞吐时专家总算力打满了,瓶颈又变成跨卡的 all-to-all 通信,通信和计算的耗时比能到 9 比 1。而要降算力直接砍 top-k 或专家宽度,又会伤模型容量。
LatentMoE 的做法是把专家计算从隐藏维 d 解耦出来:先用下投影把激活压到更小的潜空间(d 除以 ℓ 等于 α,实验里 α=4),专家在潜空间里算,再用上投影投回去。关键不是只压维,而是同时按 α 放大专家数 N 和 top-k:潜空间让每个专家更省带宽和通信,而 (αN 选 αK) 的组合空间相对 (N 选 K) 是指数级增长,非线性预算和专家多样性反而更高。
由此分出两个变体。ℓ-MoEeff:放大 N、top-k 不变,省算力、精度持平。ℓ-MoEacc(推荐版):N 和 top-k 都按 α 放大,推理算力和标准 MoE 持平,精度更高。理论上这跟 Barron 函数的逼近界对得上:单隐层网络误差 O(1/u) 与输入维无关,非线性预算正比于被选中专家的总宽度。
在 95B 总参、8B 激活、训 1T token 上比(300B token 节点):
| 变体 | 激活参 | MMLU Pro | MMLU |
| 标准 MoE 基线 | 8.47B | 29.26 | 58.95 |
| ℓ-MoEacc | 8.44B | 34.91 | 62.23 |
| ℓ-MoEeff | 5.62B | 34.75 | 61.06 |
同算力下 acc 版全线下游任务更高;eff 版用少得多的 FLOP 拿到接近的精度。73B 的 Mamba-Attention 混合 MoE 上结论一致(acc 版 MMLU Pro 52.87 vs 基线 48.30)。推理实测(H100 加 vLLM 加 FP8):高并发下吞吐最多只掉 6%。再往万亿规模外推(对标 Kimi-K2-1T),LatentMoE 相当于约 1.35 倍有效参数,同精度标准 MoE 要多 350B 参数,且在吞吐-延迟前沿慢 1.24 到 3.46 倍。这套架构已被英伟达 Nemotron-3 Super 和 Ultra 旗舰模型采用。
对做大模型基础设施的人,LatentMoE 给的是一条「不靠堆参也能涨点」的路子:在推理成本几乎不变的前提下,把精度天花板抬高,或者同精度下大幅省算力。它直接对准了 MoE 在线推理时被显存带宽和 all-to-all 通信拖累的真问题,而且已被量产旗舰模型验证,不是纸面设计。对关心 Kimi K3 Stable LatentMoE 的人,这篇是这条技术线的源头之一。
作者承认:高并发吞吐最多掉 6%,并指出可用独立 CUDA stream、专门的小矩阵 GEMM kernel 再优化;acc 版用的是和基线一样的超参,再调可能更好。
存疑点:大半结论来自设计空间探索和外推,95B、1T token 的实测规模相对现在的旗舰偏小,「同算力更高精度」在真正万亿级、长 token 上的延续性主要靠 Nemotron-3 的采用背书,论文本身没给那一档的直接数据。推理对比里的标准 MoE 基线是同条件重训还是现成配置,论文表述偏含糊。