Table of Contents
Fetching ...

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

TL;DR

LatentMoE tackles the practical bottlenecks of deploying mixture-of-experts models by projecting inputs into a low-dimensional latent space for routing and expert computation, thereby reducing memory bandwidth and all-to-all communication costs. By scaling the number of experts and top-k routing in proportion to the compression factor, LatentMoE preserves the model's nonlinear capacity while expanding combinatorial sparsity, achieving higher accuracy per FLOP and per parameter than standard MoEs. Empirical results across scales up to 95B parameters and projections to trillion-parameter regimes show consistent Pareto-front improvements, with the acc variant delivering the best accuracy at fixed cost and the eff variant matching baseline costs. The approach is positioned for practical deployment, as evidenced by measured inference performance and prospective serving advantages in large-scale models like Nemotron-3, along with clear pathways for integration with existing compression and architectural enhancements.

Abstract

Mixture of Experts (MoEs) have become a central component of many state-of-the-art open-source and proprietary large language models. Despite their widespread adoption, it remains unclear how close existing MoE architectures are to optimal with respect to inference cost, as measured by accuracy per floating-point operation and per parameter. In this work, we revisit MoE design from a hardware-software co-design perspective, grounded in empirical and theoretical considerations. We characterize key performance bottlenecks across diverse deployment regimes, spanning offline high-throughput execution and online, latency-critical inference. Guided by these insights, we introduce LatentMoE, a new model architecture resulting from systematic design exploration and optimized for maximal accuracy per unit of compute. Empirical design space exploration at scales of up to 95B parameters and over a 1T-token training horizon, together with supporting theoretical analysis, shows that LatentMoE consistently outperforms standard MoE architectures in terms of accuracy per FLOP and per parameter. Given its strong performance, the LatentMoE architecture has been adopted by the flagship Nemotron-3 Super and Ultra models and scaled to substantially larger regimes, including longer token horizons and larger model sizes, as reported in Nvidia et al. (arXiv:2512.20856).

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

TL;DR

LatentMoE tackles the practical bottlenecks of deploying mixture-of-experts models by projecting inputs into a low-dimensional latent space for routing and expert computation, thereby reducing memory bandwidth and all-to-all communication costs. By scaling the number of experts and top-k routing in proportion to the compression factor, LatentMoE preserves the model's nonlinear capacity while expanding combinatorial sparsity, achieving higher accuracy per FLOP and per parameter than standard MoEs. Empirical results across scales up to 95B parameters and projections to trillion-parameter regimes show consistent Pareto-front improvements, with the acc variant delivering the best accuracy at fixed cost and the eff variant matching baseline costs. The approach is positioned for practical deployment, as evidenced by measured inference performance and prospective serving advantages in large-scale models like Nemotron-3, along with clear pathways for integration with existing compression and architectural enhancements.

Abstract

Mixture of Experts (MoEs) have become a central component of many state-of-the-art open-source and proprietary large language models. Despite their widespread adoption, it remains unclear how close existing MoE architectures are to optimal with respect to inference cost, as measured by accuracy per floating-point operation and per parameter. In this work, we revisit MoE design from a hardware-software co-design perspective, grounded in empirical and theoretical considerations. We characterize key performance bottlenecks across diverse deployment regimes, spanning offline high-throughput execution and online, latency-critical inference. Guided by these insights, we introduce LatentMoE, a new model architecture resulting from systematic design exploration and optimized for maximal accuracy per unit of compute. Empirical design space exploration at scales of up to 95B parameters and over a 1T-token training horizon, together with supporting theoretical analysis, shows that LatentMoE consistently outperforms standard MoE architectures in terms of accuracy per FLOP and per parameter. Given its strong performance, the LatentMoE architecture has been adopted by the flagship Nemotron-3 Super and Ultra models and scaled to substantially larger regimes, including longer token horizons and larger model sizes, as reported in Nvidia et al. (arXiv:2512.20856).
Paper Structure (14 sections, 14 equations, 7 figures, 5 tables)

This paper contains 14 sections, 14 equations, 7 figures, 5 tables.

Figures (7)

  • Figure 1: Standard MoE vs. LatentMoE architectures. In LatentMoE, tokens are projected from the model hidden dimension $d$ into a smaller latent dimension $\ell$ for expert routing and computation, which reduces routed parameter loads and all‑to‑all traffic by a factor of $d / \ell$. We use this efficiency to increase the total number of experts and the top-k active experts per token by the same factor $d/\ell$, which improves the accuracy of the model while keeping overall inference cost approximately constant.
  • Figure 2: Roofline Analysis for serving Qwen3-235B-A22B. Operating points correspond to different per-expert token counts $t_\mathrm{exp}$ (i.e., effective expert batch sizes after MoE routing), mapped to arithmetic intensity $I=\frac{2 \cdot t_\mathrm{exp} \cdot d \cdot m}{d \cdot m + t_\mathrm{exp} \cdot (d+m)}$. At latency-critical batch sizes (low $I$), MoE expert computation is constrained by HBM bandwidth rather than compute, and the operating points lie in the bandwidth-bound regime.
  • Figure 3: Effect of compression ratio on model quality. Validation loss for the 16BT-2BA model using the $\operatorname{\ell-MoE}_\mathrm{eff}$ configuration across varying compression ratios $\alpha=d/\ell$. The total number of experts is scaled by $\alpha$, while the base model configuration follows Table \ref{['tbl:model_arch']}.
  • Figure 4: Impact of expert scaling on model quality. Comparison of validation loss for the 16BT-2BA model using the $\operatorname{\ell-MoE}_\mathrm{eff}$ configuration when the hidden dimension is compressed by $4\times$. The green curve utilizes the proposed expert scaling ($N' = \alpha N$), while the red curve does not. Scaling the expert count effectively mitigates the accuracy loss caused by compression, eliminating the need for extensive hyperparameter retuning.
  • Figure 5: Comparison between LatentMoE variants. Training trajectories for the baseline 16BT-2BA model versus the $\operatorname{\ell-MoE}_\mathrm{eff}$ and $\operatorname{\ell-MoE}_\mathrm{acc}$ ($\ell=512$). $\operatorname{\ell-MoE}_\mathrm{eff}$ matches baseline convergence, while $\operatorname{\ell-MoE}_\mathrm{acc}$ outperforms the baseline.
  • ...and 2 more figures