The Router: How the Model Decides Which Experts to Activate

Given that MoE only activates a subset of experts per token, an obvious question follows: how does the model actually decide WHICH experts are the right ones for a given token? The answer is the router.

The router is like a multi-class classifier that produces softmax scores over all the available experts, for each token independently. Based on these scores, the top K experts are selected — K being a hyperparameter (e.g. top-2 out of 8 total experts is a common configuration in real MoE models). The scores aren't hand-designed or hardcoded — the router is trained ALONGSIDE the rest of the network, learning over the course of training to select the experts that are actually most useful for each specific token, based purely on the training signal the whole model receives.

This means the router's behavior emerges from training rather than being specified by an engineer — nobody manually decides 'expert 3 should specialize in code, expert 5 should specialize in math.' Instead, the router and the experts co-adapt during training: as certain experts become better at handling certain kinds of tokens (whether that specialization ends up being linguistic, syntactic, domain-specific, or something less humanly interpretable), the router learns to route more of those tokens toward them.

This training-time co-adaptation is exactly what makes the router genuinely non-trivial to get right — and it's precisely the source of the two specific training challenges covered in the next subtopic. Because the router's selections directly influence which experts get gradient updates (only SELECTED experts receive training signal for a given token), the router's own learning dynamics and the experts' learning dynamics are tightly coupled and can reinforce each other in unhelpful ways if left unaddressed.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Router(nn.Module):
    """A multi-class classifier over experts — trained alongside
    the rest of the network, learning which experts fit which tokens."""
    def __init__(self, d_model: int, num_experts: int):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)

    def forward(self, x: torch.Tensor, top_k: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
        # x: (batch, seq_len, d_model) — one routing decision PER TOKEN
        logits = self.gate(x)                          # (batch, seq_len, num_experts)
        scores = F.softmax(logits, dim=-1)              # multi-class classifier over experts
        top_k_scores, top_k_indices = scores.topk(top_k, dim=-1)  # select top-K experts
        return top_k_scores, top_k_indices

d_model, num_experts = 4096, 8
router = Router(d_model, num_experts)
tokens = torch.randn(1, 5, d_model)  # 5 tokens in this example
top_k_scores, top_k_experts = router(tokens, top_k=2)
print("Selected experts per token:", top_k_experts.squeeze(0).tolist())