LatentMoE: Efficient Latent Mixture of Experts

• LatentMoE is a parameter-efficient mixture of experts variant that leverages a lower-dimensional latent space for sparse routing and expert computation.
• It projects full-dimensional activations into a compact latent space, reducing bandwidth requirements while allowing for increased expert capacity and effective nonlinear expressivity.
• Integrated in models like NVIDIA’s Nemotron 3, LatentMoE demonstrates enhanced throughput and model quality with minimal latency impact, making it ideal for large-scale language models.

LatentMoE, also known as Mixture of Latent Experts (MoLAE), is a parameter-efficient and hardware-aware variant of the Mixture-of-Experts (MoE) paradigm for large neural networks. LatentMoE introduces a key architectural shift: the sparse routing, expert computation, and all communication occur in a lower-dimensional latent space rather than in the full model space. This approach fundamentally restructures the compute and communication patterns of MoE layers, leading to significant gains in throughput, efficiency, and model quality, particularly in large-scale LLMs such as NVIDIA’s Nemotron 3 Super and Ultra (NVIDIA et al., 24 Dec 2025), and in parameter-efficient MoE variants for LLMs (Liu et al., 29 Mar 2025).

1. Motivation and Architectural Distinctions

Traditional MoE designs route input activations with dimension $d$ directly through top-$K$ out of $N$ full-capacity experts, each being a dense $d \times m$ feed-forward block. This design leads to practical scaling bottlenecks:

• Memory-bound regime: For latency-sensitive inference (e.g., batch size 1, short sequences), DRAM bandwidth for expert weights becomes limiting.
• Communication-bound regime: For throughput-optimized inference (e.g., large batches or long sequences), all-to-all cross-device communication for routed activations dominates.

LatentMoE addresses both by first projecting each token’s hidden state $x \in \mathbb{R}^d$ into a latent space of dimension $\ell \ll d$. All routing, gating, expert computation, and communication then occur within this smaller latent space. The savings in bandwidth (scaling as $d/\ell$) are reinvested to increase the number of experts ($N' = N \cdot (d/\ell)$) and active fan-out ($K' = K \cdot (d/\ell)$), effectively multiplying the expressivity and nonlinear capacity without additional runtime overhead (NVIDIA et al., 24 Dec 2025, Liu et al., 29 Mar 2025).

2. Mathematical Formulation

Given an input $x \in \mathbb{R}^d$, LatentMoE implements the following sequence of operations within a block:

1. Projection to Latent Space:

$$z_i = W_{\mathrm{down}} x_i, \quad z_i \in \mathbb{R}^\ell$$

2. Gating (performed in the full $d$-dimensional space):

$u_i = W_{\mathrm{gate}} x_i + b_{\mathrm{gate}}$

$$g_i = \mathrm{softmax}(u_i) \xrightarrow[\text{Top-}K']{} \{g_{i,j}\}_{j \in S_i}$$

where $S_i$ indexes the top-$K'$ experts.

3. Expert Application: Each selected expert $j$ computes an independent FFN in the latent space:

$$e_{i,j} = \mathrm{FFN}_j(z_i), \quad e_{i,j} \in \mathbb{R}^\ell$$

4. Mixture/Aggregation:

$\tilde z_i = \sum_{j \in S_i} g_{i,j} e_{i,j}$

5. Projection Back to Model Space:

$$y_i = W_{\mathrm{up}} \tilde z_i, \quad y_i \in \mathbb{R}^d$$

The output is added residually: $x_{i+1} = x_i + y_i$.

This block structure can be alternately summarized using the notation of (Liu et al., 29 Mar 2025):

$$\text{MoLAE}(x) = \sum_{k=1}^{K} g_k(x)\,P\,\bigl(E_k\,P^T x\bigr)$$

where $P \in \mathbb{R}^{d \times \ell}$ is the shared projection, $E_k \in \mathbb{R}^{\ell \times \ell}$ are expert-specific transforms, and $g(x)$ is the sparse gating mask.

3. Integration with Model Backbones

In NVIDIA Nemotron 3 (NVIDIA et al., 24 Dec 2025), LatentMoE is integrated into a hybrid Mamba–Transformer backbone comprising:

• Mamba-2 state-space model layers for $O(d)$ memory efficiency
• Sparse self-attention layers for long-range dependencies
• LatentMoE blocks interleaved at the same locations as conventional MoE blocks

For the Super and Ultra models, every standard sparse FFN MoE block is replaced by a LatentMoE block, optimizing both the compute pathway and the communication footprint, while retaining the depth and overall network topology. The gating network and all non-MoE layers remain in the full hidden dimension to preserve attention flow and residual information.

4. Training Objective, Regularization, and Conversion

The principal optimization objective is the autoregressive cross-entropy loss for language modeling:

$$\mathcal{L}_{\mathrm{CE}} = -\sum_{t=1}^{T} \log p(x_t \mid x_{<t})$$

LatentMoE inherits auxiliary balancing losses from prior MoE practice, including:

• Importance loss: Balances the cumulative routing weight across experts.
• Load loss: Balances the count of tokens assigned to each expert.

$$\mathcal{L}_{\mathrm{imp}} = \lambda_{\mathrm{imp}} \sum_{j=1}^{N'} \left[ \mathrm{Importance}(j) - \frac{1}{N'} \right]^2$$

$$\mathcal{L}_{\mathrm{load}} = \lambda_{\mathrm{load}} \sum_{j=1}^{N'} \left[ \mathrm{Load}(j) - \frac{1}{N'} \right]^2$$

Conversion from a pre-trained MoE to a MoLAE/LatentMoE block leverages SVD-based low-rank factorizations to approximate each expert’s weight matrix as $W_k \approx P E_k$. A two-step algorithm aligns the shared projection and expert-specific transforms to minimize the Frobenius norm reconstruction error (Liu et al., 29 Mar 2025).

5. Computational and Empirical Advantages

LatentMoE provides pronounced efficiency gains by reducing both the parameter count and distributed communication:

Model	Params (FFN)	Downstream Perplexity	# Experts	Experts/Latent
Standard MoE	151M	75.86	32	1
MoLAE (Latent)	94M	81.57	32	8

For Nemotron 3-scale models (NVIDIA et al., 24 Dec 2025):

• With $d{=}4096$, $\ell{=}1024$, $N{=}128$, $K{=}6$, a standard MoE yields 48.30% on MMLU-Pro.
• LatentMoE with $(N', K') = (512, 22)$ boosts MMLU-Pro to 52.87%, code to 55.14% (+3.19%), and achieves similar runtime and FLOP count.
• Bandwidth and routed activations per expert drop by a factor of $d/\ell$; experiment uses $d/\ell=4$.
• <1% additional end-to-end inference latency, due to small overhead from projection layers.

MoLAE conversions on large LLMs (e.g., Qwen1.5-MoE 2.7B) show that with the right choice of latent rank, >98% of task performance is retained while reducing FFN parameter count by ≈40% (Liu et al., 29 Mar 2025).

6. Comparison to Conventional MoE and Related Methods

Conventional MoE layers always operate in the unreduced model space ($d$), causing communication and DRAM access to scale linearly with the number of experts and routing fan-out. This limits either the total number of experts (N), their capacity, or the active fan-out (K) before hitting hardware bottlenecks.

LatentMoE breaks this trade-off by relocating the bottleneck to a much smaller space ($\ell$), allowing both N and K to scale up by approximately $d/\ell$, thus increasing overall model capacity and effective nonlinear expressivity without increasing the critical bandwidth or runtime costs (NVIDIA et al., 24 Dec 2025, Liu et al., 29 Mar 2025).

Parameter and runtime scaling comparisons:

Model Variant	Param. Count	Memory/Comm.	Routing Dim.	Throughput/Growth
Standard Sparse MoE	$N d m$	$d$	$d$	Limited by $N, K$
LatentMoE / MoLAE	$d \ell + N \ell m$	$\ell$	$\ell$	Scales as $d/\ell$

A plausible implication is that, with appropriate selection of latent dimension $\ell$ (e.g., $\ell \approx d/4$), one attains vastly improved capacity and balanced efficiency without affecting core modeling dynamics in the rest of the network.

7. Limitations, Trade-Offs, and Future Directions

While LatentMoE delivers strong efficiency and quality improvements, its structural constraints introduce potential downsides:

• Excessive sharing (i.e., too many experts per latent space) can degrade task expressivity and performance.
• Fixed latent dimension $\ell$ may not optimally suit all tokens or layers.
• Factorizing “down” projections may increase approximation error; in practice, some architectures only factorize “up” and “gate” components.
• Expressivity is ultimately bounded by the shared latent basis; low-rank approximations may not capture all fine-grained expert specialization if $r$ is chosen too small.

Open directions include dynamic latent ranks, extension of latent factorization to attention weights, and end-to-end fine-tuning post-conversion (Liu et al., 29 Mar 2025). The success of LatentMoE in Nemotron 3 demonstrates its ability to support agentic reasoning, tool-use, and ultra-long context extrapolation while maintaining hardware efficiency and high-quality outputs (NVIDIA et al., 24 Dec 2025).