Mixture-of-Experts Routing

A router produces logits over experts, keeps the top $k$, renormalises them into weights and sums the selected experts’ outputs. Everything interesting about MoE is a consequence of the routing being discrete: the gate is not differentiable through the selection, only through the weights of what was selected.

Load balancing

Left alone, routers collapse: a few experts win early, receive more gradient, and win harder. The standard remedy is an auxiliary loss on the product of the dispatch fraction $f_i$ and the mean gate probability $p_i$, which is minimised when both are uniform.

import torch
import torch.nn.functional as F
from torch import Tensor, nn


class TopKRouter(nn.Module):
    def __init__(self, d_model: int, n_experts: int, k: int = 2):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.k = k

    def forward(self, x: Tensor):
        logits = self.gate(x)                 # [T, E]
        probs = logits.softmax(dim=-1)
        weights, idx = probs.topk(self.k, dim=-1)
        weights = weights / weights.sum(-1, keepdim=True)

        # auxiliary balance loss: dispatch fraction x mean gate prob
        n_experts = probs.size(-1)
        one_hot = F.one_hot(idx[..., 0], n_experts).float()
        f = one_hot.mean(0)
        p = probs.mean(0)
        aux = n_experts * (f * p).sum()
        return weights, idx, aux

The balance loss here is computed over the top-1 assignment only, which is what the Switch formulation specifies; summing over all $k$ selected experts changes the fixed point. In a distributed setting $f$ must be all-reduced across data-parallel ranks, or each rank balances its own shard and nothing balances globally.