LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts

发表时间: 2026-01 · arXiv:2601.18089 (NVIDIA)

作者/机构: 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

A1 主要贡献

本文的核心研究目标是解决现有专家混合(MoE)架构在推理成本方面(以每浮点运算(FLOP)和每参数的准确度衡量)的次优问题。

核心问题与研究目标
现有的MoE架构虽然被广泛采用,但它们的设计主要基于高层次的稀疏性论证,并主要针对离线、吞吐量导向的场景进行优化。这种设计忽视了在线部署中严格的延迟、内存带宽和通信约束,导致模型在实际应用中效率低下。作者认为,有效的MoE设计必须同时评估每FLOP的准确度(计算效率)和每参数的准确度(反映内存占用、内存带宽需求、路由通信和分片开销)。忽视后者会导致模型虽然在总计算量上看起来高效,但在实际部署中却存在显著的低效性。因此,研究目标是从硬件-软件协同设计的角度重新审视MoE架构,以在固定的推理成本约束下最大化可实现的准确度。

核心创新点(LatentMoE架构)
为了解决上述问题,本文提出了LatentMoE,一种为最大化每单位计算的准确度而优化的新型MoE架构。其核心思想是:
1. 解耦与潜在空间投影:LatentMoE将专家路由和计算与模型的隐藏维度解耦。它首先通过一个可学习的下投影矩阵,将输入激活从模型的隐藏维度 $h$ 投影到一个共享的低维潜在空间 $l$(如图1所示)。
2. 降低推理成本:这个潜在空间投影直接减少了路由参数的加载量和跨设备的全对全(all-to-all)通信量,减少的比例因子为 $h/l$。这直接解决了在延迟敏感场景下的内存带宽瓶颈和在吞吐量导向场景下的通信瓶颈。
3. 再投资以提升准确度:通过潜在空间投影节省下来的计算、内存和通信资源被“再投资”于提升模型的表达能力。具体而言,LatentMoE将专家总数 $N$ 和每个token的top-k激活专家数 $k$ 同比例增加 $h/l$ 倍。
4. 等成本下的准确度提升:这种设计在保持总体推理成本(FLOPs、内存带宽、通信量)近似不变的情况下,通过增加专家数量和组合稀疏性多样性,提高了模型的有效表达能力和最终准确度。

本文通过在高达95B参数规模和超过1T token的训练范围内的实证设计空间探索,并结合理论分析,证明了LatentMoE在每FLOP和每参数的准确度方面始终优于标准MoE架构。鉴于其强大性能,LatentMoE架构已被旗舰模型Nemotron-3 Super和Ultra采用。

图1 | 标准MoE与LatentMoE架构。在LatentMoE中,token从模型隐藏维度h投影到一个更小的潜在维度ℓ,用于专家路由和计算,这将路由参数负载和all-to-all流量减少了h/ℓ倍。我们利用这种效率,将专家总数和每个token的top-k激活专家数也增加相同的h/ℓ倍,从而在保持总体推理成本近似恒定的同时提高模型准确度。
图1 | 标准MoE与LatentMoE架构。在LatentMoE中,token从模型隐藏维度h投影到一个更小的潜在维度ℓ,用于专家路由和计算,这将路由参数负载和all-to-all流量减少了h/ℓ倍。我们利用这种效率,将专家总数和每个token的top-k激活专家数也增加相同的h/ℓ倍,从而在保持总体推理成本近似恒定的同时提高模型准确度。

A3 LatentMoE核心设计原则

在深入探讨LatentMoE的具体细节之前,我们首先从系统层面审视部署一个既准确又具成本效益的MoE模型需要什么。在本节中,我们以Qwen3-235B-A22B作为建模的运行示例,其参数为:$N = 128$ 个专家,每个token激活 $k = 8$ 个专家,隐藏维度 $h = 4096$,前馈网络中间维度 $i = 1536$。为具体起见,我们考虑在由高带宽NVLink fabric互连的NVIDIA GB200 GPU上部署,每个GPU提供约1800 GB/s的双向带宽(即单向带宽 $BW_{NVL} = 900$ GB/s)。为确保专家通信保持在单个NVLink域内,专家通过专家并行(Expert Parallelism, EP)分布在 $EP = 64$ 个GPU上。注意力层在同一组GPU上使用数据并行执行。每个GB200 GPU提供 $P = 10$ PFLOPs的峰值FP4 Tensor Core吞吐量和 $BW_{HBM} = 8$ TB/s的HBM内存带宽【NVIDIA Corporation, NVIDIA Blackwell Architecture Datasheet, 2024, https://nvdam.widen.net/s/wwnsxrhm2w/blackwell-datasheet-3384703】 。

2.1. 内存带宽瓶颈

在小批量、低延迟场景下,MoE计算受内存带宽限制。在高度交互式(即低延迟)且通常使用小批量大小的设置中,MoE计算主要受到内存带宽的瓶颈限制。图2提供了性能与算术强度(arithmetic intensity)的高级roof-line分析。

GB200系统的算术强度阈值。对于一个GB200系统,只有当其算术强度(即每字节的FLOPs)超过以下阈值时,计算才会变为计算密集型:

专家处理的token数量。设 $T_{total}$ 为在MoE路由之前,EP个GPU上的总token数。假设token在专家之间均匀分布,分配给单个专家的token数为:$T_{exp} := T_{total} \cdot \frac{k}{N}$。在Qwen3-235B-A22B的例子中,$N = 128$,$EP = 64$,每个GPU承载 $N/EP = 2$ 个专家;因此每个GPU在每个MoE层处理大约 $2 \cdot T_{exp}$ 个专家token。

单个专家的计算与内存开销。单个专家的FP4计算成本为 $F_{exp} = 2 \cdot T_{exp} \cdot h \cdot i$,相应的FP4精度下的内存流量(包括权重、输入和中间激活)为 $M_{exp} = h \cdot i + T_{exp} \cdot (h + i)$。

算术强度的计算。由于在我们的例子中每个GPU处理两个专家,算术强度 $I$ 由总计算量与总内存流量之比给出:

计算密集型所需的token数。为了在计算密集型区域操作,我们需要 $I \ge 1250$。代入Qwen3-235B-A22B的参数,得到条件:

结论:低延迟场景为内存密集型。在典型的延迟关键部署中,有效批量大小很小,导致 $T_{exp}$ 仅为几百个token的量级,远低于1418的阈值。因此,MoE专家在roofline曲线的内存密集型区域运行(图2),其性能受限于权重加载而非计算能力。

图2 | Qwen3-235B-A22B服务的Roofline分析。操作点对应于不同的每专家token数 T_exp(即MoE路由后的有效专家批量大小),映射到算术强度 I = 2*T_exp*h*i / (h*i + T_exp*(h+i))。在延迟关键的批量大小(低T_exp)下,MoE专家计算受HBM带宽而非计算能力的限制,操作点位于带宽限制区域。
图2 | Qwen3-235B-A22B服务的Roofline分析。操作点对应于不同的每专家token数 T_exp(即MoE路由后的有效专家批量大小),映射到算术强度 I = 2*T_exp*h*i / (h*i + T_exp*(h+i))。在延迟关键的批量大小(低T_exp)下,MoE专家计算受HBM带宽而非计算能力的限制,操作点位于带宽限制区域。

设计原则 I
在低延迟服务场景中,MoE推理通常由加载模型权重的内存带宽成本主导。因此,对于具有高交互性要求的应用,最大化每参数的准确度至关重要。

2.2. 通信瓶颈

在吞吐量导向场景下,通信成为瓶颈。在吞吐量导向的设置中,一旦专家变为计算密集型,通信就成为分布式环境中端到端执行时间的重要贡献者。专家并行要求跨设备进行token的全对全(all-to-all)路由,这会产生可能控制端到端执行时间的开销。回想一下,在我们的示例配置中,每个GPU承载两个专家。由于 $T_{exp}$ 表示每个专家处理的token数,因此每个GPU在其本地专家上每个MoE层总共处理 $2T_{exp}$ 个token。

All-to-all通信量计算。假设token均匀分布,每个GPU在每个MoE层的all-to-all通信量为:

这里的2.5倍因子考虑了混合精度流量(FP4分派占0.5字节,BF16聚合占2字节)。

本地专家计算量。在计算方面,两个本地专家的总FLOP计数为:

通信时间与计算时间的比值。相应的计算时间为 $t_{comp} = \frac{F_{comp}}{P} = \frac{4 \cdot T_{exp} \cdot h \cdot i}{P}$。类似地,all-to-all通信时间为 $t_{comm} = \frac{C_{comm}}{BW_{NVL}} = \frac{5 \cdot T_{exp} \cdot h}{BW_{NVL}}$,其中 $BW_{NVL} = 900$ GB/s是有效的单向NVLink带宽。因此,通信时间与计算时间之比为:

结论:通信开销主导吞吐量场景。代入GB200 NVL72和Qwen3-235B-A22B的参数,得到的比率约为9。这表明在吞吐量导向的场景中,MoE层严重受到all-to-all通信开销的主导。

设计原则 II
在吞吐量导向的MoE部署中,提升性能需要最小化all-to-all操作的数据量。该数据量与以下成正比:

因此,可以通过减少被路由的隐藏维度 $h$ 或激活专家的数量 $k$ 来减轻通信开销。注意,修改中间维度 $i$ 不会影响token大小,因此不会直接带来改善。

2.3. 模型质量

保留模型质量至关重要。除了优化推理速度,保留模型质量至关重要。为此,我们借鉴了关于神经网络表达能力和组合稀疏性的理论见解。关于Barron函数的经典结果【Barron, Andrew R., Universal approximation bounds for superpositions of a sigmoidal function, 1993, IEEE Transactions on Information Theory】指出,具有 $n$ 个非线性单元的单隐藏层网络可实现与输入维度 $d$ 无关的均方误差 $O(1/n)$。

有效非线性预算。在MoE层中,每个token的有效非线性预算与所选专家的总宽度成正比:

直接降低k或i会损害模型质量。这意味着减少激活专家数 $k$ 或中间维度 $i$ 会直接惩罚有效容量 ($W_{eff}$),从而有降低模型质量的风险。

设计原则 III
保持模型质量需要保留有效的非线性预算 $k \cdot i$。因此,为了在不牺牲模型质量的情况下缓解内存和通信瓶颈,我们应保持激活专家的数量和中间维度不变。

隐藏维度的下限。每个推理任务都具有一个内在的特征秩 $h_{eff}$,对应于保留任务相关信息所需的最小自由度。将隐藏维度 $h$ 减少到此阈值以下必然会丢弃此类信息,导致准确度下降。因此,$h_{eff}$ 作为 $h$ 的一个与任务相关的下限。

设计原则 IV
存在一个任务特定的特征秩 $h_{eff}$,它对 $h$ 的缩减施加了下限。将 $h$ 减少到此限制以下会导致模型质量急剧下降。

组合稀疏性的好处。此外,MoE架构受益于组合稀疏性,为每个token提供 $\binom{N}{k}$ 种可能的专家组合【Dai, Damai, et al., Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024, CoRR, https://arxiv.org/abs/2401.06066】。增加专家总 数 $N$ 会扩展这个专业化空间。

同时扩展N和k可指数级增加多样性。进一步地,将 $N$ 和 $k$ 都按比例因子 $\gamma$ 进行扩展,会指数级地增加专家混合的多样性:

设计原则 V
同时扩展专家数量 $N$ 和每个token的top-k数量 $k$ 能通过指数级扩展专家组合空间来提升模型质量。

综合设计原则。设计原则I和II指出,提高推理速度需要降低内存带宽和通信成本。内存带宽成本与 $h$ 和 $i$ 成正比,而通信成本与 $h$ 和 $k$ 成正比。然而,原则III警告不要减少 $k$ 或 $i$,因为这样做可能会降低模型质量。这使得 $h$ 成为最有希望减少的维度,从而在吞吐量和延迟导向的场景中都能实现性能提升,而不会显著损失准确度。原则IV进一步为 $h$ 设定了一个下限 ($h_{eff}$),以防止质量崩溃。此外,原则V表明增加 $N$ 和 $k$ 可以提高模型质量。由于内存带宽和通信成本与 $h$ 呈线性关系,我们可以同时将 $N$ 增加一个因子 $\gamma$,并将 $h$ 减少相同的因子 $\gamma$(前提是 $h/\gamma \ge h_{eff}$)。我们假设,并在后续章节中通过经验验证,这种变换(也如图1b所示)能在保持内存带宽和通信成本不变的同时,提高网络表达能力和组合稀疏性,从而在每FLOP和每参数上获得更高的准确度。

A2 方法细节

LatentMoE架构概述。根据第2节中概述的设计原则,我们引入了LatentMoE,这是一种为高效扩展而设计的新型MoE架构。LatentMoE首先使用一个可学习的下投影矩阵 $W_{\downarrow} \in R^{\ell \times h}$ 将每个输入token $x \in R^h$ 投影到一个低维潜在空间 $R^\ell$ 中。然后,得到的压缩表示被路由到选定的专家。每个专家 $E_j(\cdot; \ell)$ 完全在潜在空间内操作,并由权重 $W^{(j)}_{FC1}, W^{(j)}_{gate} \in R^{i \times \ell}$ 和 $W^{(j)}_{FC2} \in R^{\ell \times i}$ 参数化。专家计算后,输出被聚合,并使用一个可学习的上投影矩阵 $W_{\uparrow} \in R^{h \times \ell}$ 投影回原始输入维度。

l-MoE_eff的设计动机。由于我们仅将输入维度 $h$ 压缩到 $\ell$,同时保持中间维度 $i$ 不变,有效非线性预算 $W_{eff}$ 保持不变。虽然设计原则III表明这在理论上应能保持准确度,但在实践中,更大的模型通常更容易训练,且对超参数变化更具鲁棒性【Frankle, Jonathan, and Michael Carbin, The lottery ticket hypothesis: Finding sparse, trainable neural networks, 2019, International Conference on Learning Representations, https://openreview.net/forum?id=rJl-b3RcF7 】; 【Novak, Roman, et al., Sensitivity and generalization in neural networks: an empirical study, 2018, International Conference on Learning Representations, https://openreview.net/forum?id=HJC2SzZCW 】; 【Taylor, Rhian, et al., Sensitivity analysis for deep learning: Ranking hyper-parameter influence, 2021, 2021 IEEE 33rd International Conference on Tools with Artificial Intelligence (ICTAI)】。为了避免为压缩后的模型进行大量的超参数调整,我们利用设计原则V,将专家总数 $N$ 按比例因子 $\gamma = h/\ell$ 进行扩展,从而扩大组合专业化空间。关键的是,由于内存带宽成本(在延迟导向场景中)和通信成本(在吞吐量导向场景中)都不依赖于 $N$,这种扩展遵循设计原则I和II,不会产生额外的推理开销。此后,我们将这种架构修改称为 ℓ-MoE_eff,其正式定义如下:

l-MoE_eff架构详解。在此公式中,$N' = N \cdot \gamma$ 表示扩展后的被路由专家集合。被路由的专家 $E_j(\cdot; \ell)$ 在潜在空间中运作,而共享专家 $E_s(\cdot; h)$ 在原始输入空间中运作。路由权重 $g' = \text{Softmax}(W'_g \cdot x)$ 是使用一个可学习的权重矩阵 $W'_g \in R^{N' \times h}$ 从原始token $x \in R^h$ 计算得出的,而 $\text{TOP}_{k,N'}$ 表示根据路由分数从总共 $N'$ 个专家中选出的top-$k$个专家的索引。为简单起见,所有在被路由专家之外的操作——包括MoE路由机制和共享专家——继续在原始隐藏维度 $h$ 中进行,因为它们对已识别的内存和通信瓶颈没有显著贡献。

l-MoE_eff的成本降低。在下投影 $W_{\downarrow}$ 之后,token的分派和聚合发生在潜在空间 $R^\ell$ 中。与标准MoE相比,这使通信量减少了 $\gamma$ 倍。同样,因为专家权重位于潜在空间中($R^{i \times \ell}$ 和 $R^{\ell \times i}$),加载权重的内存带宽成本也减少了 $\gamma$ 倍。

l-MoE_acc的设计动机。设计原则V进一步表明,将 $N$ 和 $k$ 都按比例因子 $\gamma$ 进行扩展,会指数级地增加专家多样性,从而提高模型质量。遵循这一原则,默认的LatentMoE配置(也称为 ℓ-MoE_acc)定义如下:

l-MoE_acc架构详解。其中 $k' = k \cdot \gamma$。此公式与 ℓ-MoE_eff 的唯一不同在于激活专家的数量,它使用了 top-k' 选择函数 $\text{TOP}_{k',N'}$。

l-MoE_acc的成本与收益。由于 $k$ 增加了 $\gamma = h/\ell$ 倍,该变体相对于标准MoE保持了通信成本和内存带宽要求不变。然而,增加的专家多样性和每个token的非线性预算,在同等推理成本下带来了更优的模型准确度,从而将模型的帕累托前沿推向了一个新水平。表1总结了 ℓ-MoE_eff 和 ℓ-MoE_acc 这两种配置的成本和收益。为完整起见,我们在第4节中对两种设置都进行了评估。

A4 实验环境

本文通过在不同规模上预训练Transformer MoE模型来进行设计空间探索,以验证LatentMoE架构的有效性。

表1 | 每GPU的渐进通信和内存带宽成本比较。成本已通过硬件常数归一化。T_exp表示每个专家的平均token数,EP是专家并行度。箭头表示改进(↑)、维持(→)或基线(–)。

表2 | 用于设计空间探索的基线模型架构规范。对于混合模型的Mamba层,我们使用具有128个头、64个头维度、128个状态维度和8个组的Mamba-2块。

A4 实验结果

4.1. LatentMoE 消融研究

在16BT-2BA模型上进行的消融实验验证了LatentMoE的核心设计选择。

PROTECTED_IMAGE_19____PROTECTED_IMAGE_20
图5 | LatentMoE变体之间的比较。基线16BT-2BA模型与ℓ-MoEeff和ℓ-MoEacc (ℓ = 512)的训练轨迹。ℓ-MoEeff与基线收敛情况相当,而ℓ-MoEacc优于基线。

4.2. LatentMoE 规模化研究

图6 | 95B模型训练收敛情况。95BT-8BA基线、ℓ-MoEeff和ℓ-MoEacc配置(ℓ = 1024, γ = 4)的验证损失曲线。ℓ-MoEeff与基线收敛情况相当,而ℓ-MoEacc优于基线。
表3 | 95BT-8BA模型有无LatentMoE的准确度比较。与基线相比,在参数量相同的情况下,ℓ-MoEacc在所有下游任务上提供更高的准确度,而ℓ-MoEeff在FLOPs少得多的情况下提供相当或更好的准确度。

表4 | 混合Mamba-Attention MoE有无LatentMoE的准确度比较。与基线相比,在参数量相同的情况下,ℓ-MoEacc在所有下游任务上提供更高的准确度,而ℓ-MoEeff在FLOPs少得多的情况下提供相当或更好的准确度。

4.3. 推理性能

表5 | LatentMoE与标准MoE性能指标比较。
PROTECTED_IMAGE_21____PROTECTED_IMAGE_22

A7 补充细节

与现有工作的关系。专家混合(MoE)模型已成为最先进大语言模型服务的基石。本文首次挑战了原始MoE的设计范式,并引入了一种在等参数和等FLOPs约束下实现更高准确度的替代架构。

与模型压缩技术的正交性。与此同时,社区开发了丰富的模型压缩技术以降低推理成本,包括量化【Rouhani, Bita Darvish, et al., With shared microexponents, a little shifting goes a long way, 2023a, Proceedings of the 50th Annual International Symposium on Computer Architecture, ISCA ’23】、【Rouhani, Bita Darvish, et al., Microscaling data formats for deep learning, 2023b, https://arxiv.org/abs/2310.10537】和稀疏化 【Xie, Yanyue, et al., Moe-pruner: Pruning mixture-of-experts large language model using the hints from its router, 2024, https://arxiv.org/abs/2410.12013】。在专家层面,也提出了剪枝 【Lu, Xudong, et al., Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models, 2024, https://arxiv.org/abs/2402.14800】、 【Lasby, Mike, et al., Reap the experts: Why pruning prevails for one-shot moe compression, 2025, https://arxiv.org/abs/2510.13999】、 【Chen, Tianyu, et al., Task-specific expert pruning for sparse mixture-of-experts, 2022, https://arxiv.org/abs/2206.00277】和合并 【Li, Pingzhi, et al., Merge, then compress: Demystify efficient smoe with hints from its routing policy, 2024, https://arxiv.org/abs/2310.01334】等方法。这些方法与LatentMoE的设计是正交的,可以组合使用以获得进一步的效率提升 。

与MoLAE的比较。与本文最相关的工作可能是MoLAE【Liu, Zehua, et al., Molae: Mixture of latent experts for parameter-efficient language models, 2025, https://arxiv.org/abs/2503.23100】。MoLAE是一种基于潜在空间中专家权重的低秩近似的后训练压缩方法。尽管两种方法在表面上看起来相似,但LatentMoE通过将专家压缩与增加网络表达能力和组合稀疏性相结合,做出了根本不同的设计权衡。相比之下,为了弥补潜在空间投影造成的准确度损失,MoLAE引入了分组潜在投影,并仅将压缩限制在部分专家(FC2)上。这些设计选择反过来放弃了token分派时的通信节省,并限制了内存带宽的减少,最终限制了可实现的效率增益。如第2节所讨论,高效的MoE服务并非受限于FLOPs;仅减少FLOPs不足以改善准确度-吞吐量-延迟的帕累托前沿 。

与mHC的比较。同期的工作通过修改残差连接而非专家路径来提高固定计算量下的模型质量。流形约束超连接(mHC)【Xie, Zhenda, et al., mhc: Manifold-constrained hyper-connections, 2026, https://arxiv.org/abs/2512.24880】通过加宽残差流和增加残差路径连接性,在等计算量下提高了质量。实现这一点需要一种截然不同的残差拓扑(多流残差状态)和一种学习的连接生成机制(用于连接图 的RMSNorm → 线性 → tanh门控,以及用于稳定性的约束残差混合)。我们认为LatentMoE和mHC是互补的,可以叠加使用。进一步的探索留给未来的工作。

A5 结论

本文提出了LatentMoE,一种修正的专家混合(MoE)架构,旨在通过明确考虑现代推理系统中的主要内存带宽和通信瓶颈,来最大化每FLOP和每参数的准确度。通过将token投影到一个低维潜在空间,LatentMoE减少了路由过程中的all-to-all通信量,以及每个专家所需的内存带宽和计算量。这些节省下来的资源被重新投入到扩展专家数量和路由多样性中,而无需增加推理成本。在高达95B参数的广泛实验、混合架构以及预测的万亿参数服务场景中,LatentMoE在准确度-效率的帕累托前沿上始终优于标准的MoE模型。