Mixture-of-Experts

A Mixture-of-Experts (MoE) layer replaces one dense feed forward network with many expert networks plus a router that sends each token to only a few of them. This decouples parameter count from compute: the model can hold billions of parameters while each token activates only a small subset.

The Sparse MoE Layer

Given $E$ experts $\{f_1,\dots ,f_E\}$ (usually MLPs) and a router that produces logits $g(x)\in {\mathbb{R}}^E$ per token, top-$k$ gating keeps only the $k$ largest logits:

$$\begin{array}{rl}p(x)& =\text{softmax}(\text{TopK}(g(x),k))\\ y& =\sum _{i\in \text{TopK}}{p}_i(x)f_i(x)\end{array}$$

Non selected experts contribute zero, so the softmax is taken only over the $k$ chosen logits. Typical settings are $k=1$ (Switch Transformer) or $k=2$ (classic MoE, Mixtral). Because routing is sparse, MoE pairs naturally with Expert Parallelism, where different experts live on different devices.

Router Gating in PyTorch

The router is a single linear layer. torch.einsum combines the per token gate weights with the stacked expert outputs.

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
 
class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, k=2):
        super().__init__()
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.experts = nn.ModuleList(
            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),
                          nn.Linear(d_ff, d_model))
            for _ in range(n_experts)
        )
        self.k = k
 
    def forward(self, x):
        # x: (b, t, d) -> flatten tokens
        tokens = rearrange(x, "b t d -> (b t) d")
        logits = self.router(tokens)                  # (n, e)
        topv, topi = logits.topk(self.k, dim=-1)      # (n, k)
        gates = F.softmax(topv, dim=-1)               # (n, k)
 
        # dense expert bank: (n, e, d), then select the k routed experts
        all_out = torch.stack([e(tokens) for e in self.experts], dim=1)
        idx = topi.unsqueeze(-1).expand(-1, -1, tokens.shape[-1])
        chosen = torch.gather(all_out, 1, idx)        # (n, k, d)
 
        # weighted combine of the k experts per token
        out = torch.einsum("nk,nkd->nd", gates, chosen)
        return rearrange(out, "(b t) d -> b t d", b=x.shape[0]), logits

Why einsum here

einsum("nk,nkd->nd", gates, chosen) is the routed combination: it multiplies each token’s $k$ gate weights against its $k$ expert outputs and sums, which is exactly ${\sum }_i{p}_i{f}_i(x)$.

Load Balancing

Left alone, the router collapses onto a few favorite experts while others starve. An auxiliary load balancing loss counteracts this. With $f_i$ the fraction of tokens routed to expert $i$ and $P_i$ the mean router probability for expert $i$ over a batch:

$${\mathcal{L}}_{\text{aux}}=\alpha \cdot E\sum _{i=1}^E{f}_i{P}_i$$

This is minimized when the load is uniform, and it is added to the main loss with a small coefficient $\alpha$.

Capacity

Each expert has a fixed capacity $=\lceil \text{capacity\_factor}\cdot \frac{k\cdot \text{tokens}}{E}\rceil$. Tokens beyond capacity are dropped (skipped or passed through via the residual). Capacity keeps expert buffers a fixed size so dispatch can be batched on hardware.

Common pitfalls

• Top-$k$ selection is not differentiable through the discrete choice; gradients flow only through the softmax over selected logits, so the balancing loss is what actually shapes routing.
• Too small a capacity factor drops many tokens and hurts quality; too large wastes memory and compute.
• Router logits often use fp32 and jitter noise during training to keep routing stable.

MoE vs Dense

MoE gives more parameters per unit of FLOPs, so it is a strong lever for scaling capacity cheaply. The costs are memory (all experts are stored), communication (all to all for Expert Parallelism), and training instability from the router.