9 Dynamic routing with Mixture of Experts
Mixture of Experts is presented as a practical way to make language models more adaptive by routing different tokens to different specialized components at runtime. Rather than forcing every input through the same dense path, the model learns where to send each token so computation is used more selectively. The chapter situates this approach among other dynamic inference methods, contrasting it with early exiting and token dropping, and emphasizes that MoE is the most mature and widely adopted of the three, especially in modern frontier systems and increasingly in smaller-scale models.
The first project builds a two-expert MoE on SmolLM2-1.7B-Instruct by replacing selected MLP modules in the deeper layers with a custom block containing a frozen general expert, a trainable clinical expert, and a router. The model is trained with a mix of clinical extraction prompts and general conversation data so the router can learn to distinguish contexts without harming the base model’s overall behavior. Soft routing blends both experts during training, while hard routing switches to a binary choice at inference time, showing that the same architecture can either combine or sparsify expert usage depending on how the forward pass is implemented.
The second project reuses a previously fine-tuned clinical model as an expert, generalizes the MoE to support configurable top-k routing, and trains mostly the router before lightly adapting the expert tail to match the base model’s attention layers. This staged approach restores full performance on the clinical task while keeping general benchmark behavior stable or slightly improved, underscoring that expert specialization does not have to come at the cost of general capability. The chapter closes by connecting these experiments to sparse upcycling and Mixtral-style architectures, showing how the same underlying ideas scale from small instructional setups to large production MoEs.
The MoE block replaces the MLP module in selected transformer layers. A router receives the hidden states and assigns routing weights to two experts that process in parallel — expert_general, which holds the original MLP weights (frozen), and expert_clinical, the trainable copy. Attribute names in parentheses correspond directly to the model's code structure.
Two routing behaviors applied to the same token. Soft routing (left, top-k = 2) runs both experts in parallel and combines their outputs in a weighted sum — × 0.38 for expert_general, × 0.62 for expert_clinical. Hard routing (right, top-k = 1) selects expert_clinical via argmax and passes its output through directly.
Benchmark comparison across five general-knowledge tasks for the base model, MoE Soft, and MoE Hard. Delta values expressed in percentage points (pp) relative to the baseline.
Two views of the same top-k routing decision. Left: the router produces one logit score per expert for each token; the two highest scores per row (highlighted) are selected. Right: the same assignments reorganized by expert, showing which tokens each expert will process and their renormalized weights — for each token, those weights sum to 1 across its assigned experts. The example uses four experts rather than two so that top_k = 2 leaves half the experts idle per token. The loop in TrainableTopKMoE.forward performs exactly this regrouping.
Two-phase training strategy for the upcycled MoE. Phase 1 trains only the router across all 24 layers (lr = 1e-3) for three epochs — 98,304 parameters, 0.0034% of the model — reaching 90% schema compliance. Phase 2 reduces the router learning rate to lr = 1e-6 and unfreezes the expert_clinical tail layers 18–23 at lr = 1e-5 for one epoch, bringing trainable parameters to 10.35% and compliance to 100%.
Benchmark comparison across five general-knowledge tasks for the upcycled MoE, baseline vs top_k=2 vs top_k=1.
Summary
- Dynamic inference techniques, like early exiting, token dropping, and Mixture of Experts, allow the model to make runtime decisions about what to compute. Of the three, MoE is the most mature and the only one with growing support in standard inference libraries.
- An MoE architecture is built around three elements. Experts process information independently, the router decides which expert acts on each token, and the combined forward pass merges their outputs into a single response.
- Building an MoE starting from a copy of the base model weights allows initializing the experts with already-acquired knowledge, instead of starting from scratch. This strategy, known as sparse upcycling, speeds up convergence.
- Soft routing and hard routing are two inference behaviors that can be obtained on the same trained weights, by modifying only the forward method. The model's behavior depends both on its weights and on how the experts' outputs are combined.
- The router does not work on raw tokens, but on their hidden states, which accumulate the context of everything that has occurred before in the sequence. That is why the same token can be sent to different experts depending on the context in which it appears.
- The TrainableTopKMoE class generalizes the architecture to accept a configurable number of experts and a top-k parameter that controls how many of them are activated per token.
- An MLP Module from a pre-trained model can be transplanted directly as an expert into an MoE architecture, without the need to retrain it from scratch. This opens the door to reusing existing specialized models as components of a larger system.
- If the transplanted expert comes from a model with a different Attention Module than the base model, it might need a light adjustment to adapt to the tensors it receives. A light training of the layers closest to the output is usually enough to recover full performance.
- Specializing an expert in a specific domain does not necessarily penalize the general capabilities of the model. General benchmarks of the MoE model built in this chapter are maintained or improved compared to the base model.
Rearchitecting LLMs ebook for free
