Hierarchical, Hybrid, and Omnimodal Models › Modular Experts and Routing

Increasing appropriately trained model capacity often improves performance, but larger dense models increase memory and inference cost and may be impractical to deploy. A sparse mixture-of-experts (MoE) layer can decouple total parameter count from per-example expert FLOPs by keeping many parameter sets but activating only a few per input. Parameter storage, routing, dispatch, communication, and memory traffic can still grow with the number and placement of experts. In world models the pattern is doubly useful: different regions of state space, different modalities, or different task contexts call for different dynamics, and a router can learn to select the right one. The first benefit is computational: a model that would be too slow to run in full can reduce expert computation by evaluating only a fraction of its experts per input. The second is an architectural inductive bias: specialization may reduce interference by assigning different regimes to different experts, although it does not automatically enlarge the representable function class relative to an unrestricted sufficiently large monolithic network.

A mixture of $K$ experts with a learned router $g$ computes

where:

• $p(y \mid x)$: the output distribution over $y$ given input $x$.
• $x$: input to the mixture.
• $y$: output variable.
• $K$: number of experts.
• $k$: expert index running from $1$ to $K$.
• $g_k(x)$: router weight for expert $k$ given input $x$.
• $f_k(y \mid x)$: the predictive density produced by expert $k$ for output $y$ given input $x$.

The output density is a mixture of expert densities with input-dependent weights. Unlike an ordinary finite mixture with input-independent component proportions, an MoE uses a learned gating network whose weights depend on $x$. Influential sparse Transformer MoEs commonly use a lightweight linear router, though more structured routers are possible:

where:

• $x$: input feature vector passed to the router.
• $g(x)$: router probability vector over experts.
• $W_g$: weight matrix of the router.
• $b_g$: bias vector of the router.
• $\text{softmax}$: softmax function.

The softmax turns raw scores into non-negative weights that sum to one, giving a valid mixture density. Keeping the router lightweight limits overhead. Specialization depends more directly on its inputs, objective, load balancing, capacity, and optimization than on router size alone.

Dense MoE computes all experts and combines them, so its expert compute grows with $K$ and it does not provide the conditional-compute scaling advantage of sparse MoE. A common sparse token-choice MoE selects the top $k$ experts (often $k = 1$ or $2$) and may renormalize weights over the selected set; other sparse routing schemes include top-1 Switch routing, expert-choice routing, threshold routing, and fixed or learned alternatives. Sparse MoE grows capacity (number of parameters) without growing per-example expert compute (which scales with $k$, not $K$). This decoupling is the central practical advantage: you can add experts to raise the model's total capacity while keeping the expert cost per example essentially flat, because only the selected experts run.

The router can be trained in several ways, and the choice trades off stability, compute, and gradient quality.

• Learned linear gate. Simplest and most common. Backpropagates through $g(x)$, so the router is trained jointly with the experts. Every expert receives a task gradient weighted by its router probability, giving continuous routing gradients, although saturated probabilities or optimization dynamics can still produce instability or negligible gradients.
• Top-1 or top-k hard routing. The router still produces soft logits, but only the top-$k$ experts are evaluated and only those receive gradients. Sparse but potentially unstable, because hard assignments can starve experts of data: an expert that is never selected never receives a gradient and can never improve enough to be selected.
• Noisy top-k. Add Gaussian noise to the router logits before top-$k$ selection so that different samples can explore different experts (Shazeer et al., 2017). With a nonzero noise scale large enough relative to the logit margins, this can reduce the risk that a small early advantage starves an expert of future updates. It may improve load distribution, but does not guarantee balanced traffic; explicit balancing objectives and capacity mechanisms are commonly used alongside it.
• Soft relaxations. Dense soft routing computes all experts with continuous weights and trades expert compute for smoother routing gradients. Differentiable soft-top-k relaxations can instead be truncated to preserve sparse execution while approximating discrete selection.

Without effective mitigation, learned sparse token-choice routers can develop severe load imbalance or collapse, with a few experts absorbing most traffic. The collapse can be self-reinforcing, because experts that start slightly ahead receive more data and improve faster, which makes them even more likely to be chosen next. A common mitigation is an auxiliary load-balancing loss. For the common top-1 Switch-style form, let $n_k$ be the fraction of tokens routed to expert $k$ in a batch and let $p_k$ be the average router probability for expert $k$. Because each token receives one hard assignment, both vectors sum to one, and the auxiliary loss is

where:

• $\mathcal{L}_{\text{balance}}$: auxiliary load-balancing loss used to discourage concentrated routing.
• $K$: number of experts.
• $k$: expert index running from $1$ to $K$.
• $n_k$: fraction of tokens assigned to expert $k$ by top-1 routing in the batch.
• $p_k$: average router probability for expert $k$ over the batch.

Uniform top-1 routing is a useful reference point. If $n_k = p_k = 1/K$ for every expert, then

The loss is added to the task objective with a small coefficient. A concentrated router instead makes the hard token fractions and average soft probabilities large for the same experts, increasing their dot product.

The auxiliary loss has a simple shape as routing concentrates. The next cell computes the top-1 load-balancing loss for a four-expert router as the highest-weight expert receives an increasing fraction of the traffic.

This auxiliary objective is a heuristic rather than a guarantee of uniform use because $n_k$ and $p_k$ are different statistics; implementations also rely on routing-capacity constraints or alternative objectives. If its coefficient is too small, concentrated routing may persist. If it is too large, the router can be pushed toward uniformity even when the data favors specialization.

Routing is not only about scaling. A decomposed world model can capture qualitatively different dynamics. Different regions of state space, different modes of a hybrid system, different tasks, or different modalities may benefit from distinct dynamics. A router that learns to send samples to the right expert effectively performs an implicit soft discretization of state space. Even when conditional-compute savings are not the main concern, explicit routing can provide a useful modular inductive bias and may improve optimization or sample efficiency by assigning different regimes to different experts. It does not automatically enlarge the representable function class relative to an unrestricted sufficiently large monolithic network.

This is related to the switched systems of classical control, where $\dot{s} = f_{\sigma(t)}(s, a)$ and $\sigma(t)$ identifies an active discrete mode. Here $\dot{s}$ is the state derivative, $a$ is the control, and $f_{\sigma(t)}$ is mode-specific dynamics. A hard-routed MoE can be interpreted as learning a state/action-dependent mode selector; a soft MoE instead learns a weighted blend. Under a hard, mode-specialized interpretation, router chattering near a boundary is one relevant control concern.

Expert capacity should suit the complexity of the regime it handles. Near mode boundaries, routing and expert outputs should also avoid spurious jumps or rapid chattering while preserving real discontinuities in the system. Noise encourages routing exploration; entropy regularization discourages overly sharp probabilities; soft kernel-based routers can interpolate continuously. Output compatibility, hysteresis, or explicit continuity regularization may still be needed. The right treatment depends on whether the physical transition is continuous at that boundary.

• Multi-embodiment robots. Different bodies can share perception while using embodiment-specific dynamics. If embodiment information is omitted or aliased in the represented input, a point-loss shared model can average incompatible body-specific transitions. Experts routed by embodiment are one remedy; conditioning a sufficiently expressive dense model on morphology is another.
• Multi-task control. Navigation, manipulation, and speech respond to inputs differently; separate experts reduce interference. Task interference is the reason a single network trained on many tasks can lose accuracy on each of them, and routing sidesteps it.
• Heterogeneous modalities. Different sensors have different statistics and noise; separate experts per modality concentrate compute where it is needed.
• Regime-switching physical systems. Contact, free-fall, and sliding dynamics have qualitatively different structure. A learned soft boundary can track gradual transitions when the dynamics are continuous and the router is trained on the boundary region; discontinuous transitions may require hard, hysteretic, or event-aware treatment.