Stable LatentMoE
Summary: Stable LatentMoE is an extreme-sparsity Mixture-of-Experts architecture (16/896 = 1.78% activation) used in Kimi K3. It combines LatentMoE's compressed expert representations with Quantile Balancing for stable routing, Per-Head Muon for adaptive optimization, SiTU activation for expert FFN control, and Gated MLA for attention-expert alignment — enabling stable training at 2.8T parameters.
Overview
Standard MoE (e.g., DeepSeek-V3: 256 experts, 8 active = 3.1%) faces instability at higher sparsity: router collapse, expert imbalance, dead experts, gradient variance. Stable LatentMoE pushes to 896 experts with only 16 active (1.78%), requiring coordinated innovations across routing, optimization, activation, and attention.
Key Components
1. Latent Expert Representation
- Experts share a latent basis; each expert = linear combination of basis vectors
- Reduces parameter redundancy across 896 experts
- Enables smoother interpolation between experts
- Related to: Parameter-efficient MoE, expert merging
2. Quantile Balancing (Router Stability)
- Problem: Standard top-k + aux loss has sensitive hyperparameter (aux loss weight)
- Solution: Expert allocation derived directly from router score quantiles
- Mechanism: Target each expert to fire at fixed quantile of router distribution
- Benefit: No heuristic aux loss weight; no manual balancing hyperparameter
3. Per-Head Muon (Optimization)
- Extends Muon by maintaining separate optimizer state per attention head
- Different heads have different gradient geometries at 3T scale
- Especially important for attention heads interacting with sparse experts
4. SiTU Activation (Expert FFN Control)
- Sigmoid Tanh Unit: bounded, self-gated, smooth
- Mitigates dead-expert gradients and activation explosion
- Quantization-friendly (MXFP4/MXFP8)
5. Gated MLA (Attention-Expert Alignment)
- Per-head gates modulated by expert routing
- Heads specialize for different expert subsets
- Reduces interference in shared attention
Mathematical Formulation
In plain English: Standard MoE: route token → pick top-k experts → compute. Stable LatentMoE: route token → quantile-normalized scores → pick top-k latent experts → decode via shared basis → expert FFN with SiTU → attention heads gated by routing.
Latent Expert Decoding
E = B @ W_expert # E: [num_experts, d_ffn], B: [num_basis, d_ffn], W_expert: [num_experts, num_basis]
expert_out = SiTU(x @ E[active_experts].T) @ E[active_experts]
Quantile Balancing
router_scores = router(x) # [B, L, E]
quantile_targets = linspace(0, 1, E) # uniform target distribution
# Match router score quantiles to uniform via monotonic transport
balanced_scores = quantile_match(router_scores, quantile_targets)
topk_experts = topk(balanced_scores, k=16)
Per-Head Muon State
for layer, head:
M[layer, head] = β M[layer, head] + (1-β) G[layer, head]
Δ = orthogonalize(M[layer, head])
W[layer, head] -= lr * Δ
Variants / Extensions
- Hierarchical LatentMoE: Two-level — coarse experts route to fine experts
- Dynamic k: Variable active experts per token (quantile-balanced)
- Expert Merging: Periodically merge similar latent experts to maintain diversity
Applications
- Kimi K3 (2.8T, 896 experts, 16 active)
- Ultra-sparse MoE (>512 experts, <2% activation)
- Long-context MoE where expert specialization by context region helps
Sparsity Comparison
| Model | Total Experts | Active Experts | Active % | Balancing Method |
|---|---|---|---|---|
| Mixtral 8x7B | 8 | 2 | 25% | Aux loss |
| DeepSeek-V2 | 160 | 6 | 3.75% | Aux loss + bias |
| DeepSeek-V3 | 256 | 8 | 3.125% | Aux loss + bias |
| Kimi K3 | 896 | 16 | 1.78% | Quantile Balancing |
Historical Context
| Year | Work | Contribution |
|---|---|---|
| 2017 | Shazeer et al. | MoE (sparse FFN) |
| 2021 | GShard / Switch Transformer | Large-scale MoE, aux loss |
| 2022 | ST-MoE / Harmony | Stable routing |
| 2023 | DeepSeek-V2 | Fine-grained experts + MLA |
| 2024 | DeepSeek-V3 | 256 experts, 8 active, shared experts |
| 2024 | Various | Latent/PES MoE (parameter-efficient) |
| 2026 | Stable LatentMoE (Kimi K3) | 896 experts, 1.78% activation, quantile balancing, co-designed optimizer/activation/attention |
Related Concepts
- quantile-balancing
- per-head-muon
- situ-activation
- gated-mla
- kimi-delta-attention
- attention-residuals
- kimi-k3
- moe-sparsity-patterns
- latent-moe
Sources
- kimi-k3-doc: Introduction — Stable LatentMoE framework, efficiently activating 16 of 896 experts
- kimi-k3-tech-blog: Architecture and Infrastructure — Stable LatentMoE with Quantile Balancing (no heuristic aux loss), Per-Head Muon, SiTU, Gated MLA; enables 2.8T stable training