MoE Transformer Architecture, in depth: where experts go, router design, shared experts, balancing, capacity and stability

What changes in the block, and what does not

A decoder block has two sublayers: self-attention, which mixes information across tokens, and a feed-forward network (FFN), which transforms each token independently and holds about two thirds of a dense model's parameters. Because the FFN acts per token, it can be split per token without breaking anything, which is exactly what MoE exploits. Attention, residual stream, normalisation, embeddings and output head are unchanged; the FFN's output is replaced by a weighted sum of the chosen experts' outputs.

Placement: which layers become MoE

Not every layer has to be sparse. Early designs such as GShard and Switch Transformer replaced every other FFN with an MoE layer, alternating dense and sparse. Most recent large models make nearly every layer MoE but keep the first few dense. DeepSeek-V3, for example, keeps the first three of its 61 layers as dense FFNs. The usual reasoning is that early layers compute generic, low-level features that every token needs, so routing gains little there, and dense early layers are also reported to make training more stable.

More MoE layers means more total parameters for the same active compute, but also one pair of all-to-all exchanges per MoE layer in distributed training. Start with all but the first one to three layers as MoE, and use alternation to save communication when the interconnect is the bottleneck.

Expert shape: coarse, fine-grained and shared

The first MoE LLMs used a few large experts, each the size of the original FFN: Mixtral has eight experts per layer and routes each token to two. DeepSeekMoE introduced fine-grained experts: split each expert into several smaller ones and route to proportionally more of them, keeping active compute constant. Sixty-four experts of one eighth width with top-16 costs the same FLOPs as eight experts with top-2, but offers vastly more combinations of experts per token, which lets experts specialise more narrowly. DeepSeek-V3 takes this to 256 routed experts with 8 active.

DeepSeekMoE also added shared experts that every token passes through outside the router. They hold common knowledge so routed experts need not each relearn it, and they give gradients a stable path. The cost of fine-grained experts is hardware efficiency: small per-expert matrix multiplies underuse the GPU unless batched into one grouped GEMM, and top-8 routing sends each token to several devices.

Routers: token choice, expert choice, softmax and sigmoid

The router is a linear layer from the hidden state to one score per expert. What happens next defines the router family.

• Token choice, softmax then top-k. Each token takes a softmax over all experts and keeps its top k, optionally renormalising the k weights. This is the GShard and Mixtral style. Every token gets exactly k experts, but nothing stops many tokens from choosing the same expert.
• Top-1 (Switch). Token choice with k = 1: the cheapest routing and the least communication, but each token relies on one expert and balancing becomes critical.
• Expert choice. Each expert picks the tokens it scores highest, up to a fixed capacity. Balance is perfect by construction, but a token may be picked by many experts or by none, and because an expert's choice depends on other tokens in the batch, it does not fit autoregressive decoding cleanly.
• Sigmoid scores. DeepSeek-V3 computes an independent sigmoid affinity per expert rather than a softmax, selects the top k, and normalises the selected scores to form gate weights. Independent scores mean one expert's score does not suppress another's, which becomes useful with hundreds of experts.

The gate weight multiplies the expert's output, so the router receives gradient only through the chosen experts. That is why a router can collapse: the experts that happen to win early get trained, become better and win more.

Balancing: auxiliary loss and bias adjustment

An idle expert is wasted parameters, and under expert parallelism the busiest expert sets everyone's step time. The classic fix is an auxiliary load-balancing loss: for each expert, multiply the fraction of tokens routed to it by its mean router probability, sum over experts and scale by the number of experts. The loss is minimised by a uniform distribution, and its gradient nudges the router toward it. The difficulty is its weight: too small and routing collapses, too large and it fights the language-modelling loss and hurts quality.

DeepSeek-V3 introduced an auxiliary-loss-free alternative. Each expert gets a bias that is added to its score only for top-k selection, not for the gate weight. After each step, overloaded experts' biases are reduced and underloaded experts' biases increased by a fixed small rate. Balance is enforced by a control loop outside the gradient, so the language-modelling loss is not distorted. The report pairs it with a very small sequence-level balance loss as a safeguard against extreme imbalance within individual sequences, and restricts each token to experts on a limited number of nodes to cap communication.

Capacity and dropping, or dropless

Hardware prefers fixed shapes, so early systems gave each expert a capacity of capacity factor x (tokens x k / number of experts), typically 1.0 to 1.25 in training. Tokens routed to a full expert are dropped and pass through the residual only: cheap and predictable, but a dropped token silently loses a layer, and at inference its output depends on the rest of the batch.

Dropless MoE computes every assignment with grouped GEMM kernels that handle uneven expert batches, so step time depends on imbalance instead. Modern stacks mostly run dropless;MoE all-to-all communication covers how tokens are dispatched to experts on other GPUs.

A reference block in PyTorch

The block below combines the modern choices: a shared expert, fine-grained routed experts, sigmoid scores, top-k selection with a balancing bias, gate weights from the unbiased scores, and a dropless combine. It loops over experts for clarity; production kernels sort tokens by expert and run one grouped GEMM. A z-loss helper for softmax routers is included for comparison.

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


class Expert(nn.Module):
    """A SwiGLU feed-forward network: the same shape as a dense FFN, only narrower."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class MoEFFN(nn.Module):
    """Shared expert + fine-grained routed experts, sigmoid scores, bias-based balancing."""
    def __init__(self, d_model, d_expert, n_routed, n_shared, top_k, bias_rate=1e-3):
        super().__init__()
        self.n_routed, self.top_k, self.bias_rate = n_routed, top_k, bias_rate
        self.router = nn.Linear(d_model, n_routed, bias=False)
        self.register_buffer("route_bias", torch.zeros(n_routed))   # not trained by SGD
        self.experts = nn.ModuleList(Expert(d_model, d_expert) for _ in range(n_routed))
        self.shared = Expert(d_model, d_expert * n_shared) if n_shared else None

    def forward(self, x):                                  # x: [tokens, d_model]
        logits = F.linear(x.float(), self.router.weight.float())   # fp32 even in bf16
        scores = torch.sigmoid(logits)
        _, idx = torch.topk(scores + self.route_bias, self.top_k, dim=-1)  # bias picks...
        gate = scores.gather(-1, idx)                      # ...raw scores weight
        gate = gate / gate.sum(-1, keepdim=True)

        flat = idx.flatten()                               # [tokens * top_k]
        order = flat.argsort()                             # group assignments by expert
        token_of = torch.arange(x.size(0), device=x.device).repeat_interleave(self.top_k)
        counts = torch.bincount(flat, minlength=self.n_routed)

        out = torch.zeros_like(x)
        start = 0
        for e, n in enumerate(counts.tolist()):            # real kernels: one grouped GEMM
            if n == 0:
                continue
            sel = order[start:start + n]
            toks = token_of[sel]
            w = gate.flatten()[sel].unsqueeze(-1).to(x.dtype)
            out.index_add_(0, toks, self.experts[e](x[toks]) * w)
            start += n

        if self.shared is not None:
            out = out + self.shared(x)
        if self.training:
            self.update_bias(counts)
        return out, counts

    @torch.no_grad()
    def update_bias(self, counts):
        # In data-parallel training, all-reduce counts first so every rank agrees.
        load = counts.float()
        self.route_bias += self.bias_rate * torch.sign(load.mean() - load)


def router_z_loss(logits):
    """ST-MoE's stabiliser for softmax routers: penalise large router logits."""
    return torch.logsumexp(logits.float(), dim=-1).pow(2).mean()

Details that are easy to get wrong: the bias affects selection only, never gate weights; it is a buffer updated under no_grad, not a parameter; token counts must be summed across data-parallel ranks before the update; and the router matmul is done explicitly in fp32 so it stays fp32 after the module is cast to bf16.

Worked example: tracing one batch through the router

Take 8 tokens, 4 routed experts and top-2: 16 assignments, ideally 4 per expert. Suppose the router produces counts of 7, 5, 3 and 1. The mean load is 4, so with a bias rate of 0.001 the update gives experts 0 and 1 a bias change of minus 0.001 and experts 2 and 3 plus 0.001. After a few hundred steps of consistent imbalance, the accumulated biases are large enough to flip marginal top-2 decisions toward experts 2 and 3, while each token's gate weights still come from its raw scores.

Now compare execution modes with one expert per GPU. Dropless, the GPU holding expert 0 does 7 units of work while the others wait: the step takes 7 / 4 = 1.75 times the balanced time. With capacity = 1.25 x 16 / 4 = 5, every expert runs a padded batch of 5, so the step costs 1.25 times, paid for by dropping 2 of expert 0's assignments; those tokens get only their other expert's output. That multiplier, not the FLOP count, is what imbalance costs.

Training stability

• Router precision. Compute router logits and softmax or sigmoid in fp32 even when the rest of the model runs in bf16. Small rounding differences flip top-k choices and cause noisy routing.
• Router z-loss. ST-MoE found that large router logits cause instability with softmax routers and added a small penalty on the squared log-sum-exp of the logits, shown in the code.
• Initialisation. Start router weights small so early routing is near uniform and every expert receives gradient.
• Loss spikes. MoE models spike more than dense ones. Watch per-layer expert load, router entropy and the fraction of dropped tokens alongside the loss, so a spike can be traced to a layer.
• Fine-tuning. Keep balancing on, or freeze the router, when adapting on small datasets.

How the architecture meets hardware

Total parameters decide memory, since every expert must be resident somewhere; active parameters decide compute; top-k and expert count decide communication: underexpert parallelism, each MoE layer dispatches tokens to their experts' devices and gathers results back, two all-to-all exchanges per layer per step. Fine-grained experts make the GEMMs small and the dispatch wide. At inference, small batches mean each expert sees few tokens, so decoding becomes bound by memory bandwidth for expert weights; serving MoE models covers how engines batch and place experts to cope.

Failure modes

• Routing collapse. A few experts take most tokens and the rest go idle. Symptoms: falling router entropy and flat load histograms turning spiky.
• Over-regularised routing. An auxiliary loss weight so high that routing becomes nearly uniform and experts fail to specialise; quality matches a smaller dense model.
• Silent token dropping. A capacity factor that is too low drops a few percent of tokens in some layers, degrading quality without errors. Log the drop rate.
• Biased gate weights. Using bias-adjusted scores as weights in an aux-loss-free implementation.
• Straggler experts. One device hosting hot experts sets the cluster's step time.

Trade-offs

For small models, where memory is the scarce resource, the balance shifts again; small MoE architecture works through that case.

What to do next

1. Read the config of an open MoE model and label each field: layers, dense layers, routed and shared experts, top-k, gating function, balancing method.
2. Implement the reference block above, run it on random data and plot per-expert counts with and without the bias update.
3. Train a small MoE and a dense model with the same active parameters on the same data; compare loss curves and memory.
4. Log router entropy, per-expert load, drop rate and max-to-mean load ratio for every MoE layer from the first step.
5. Keep router computations in fp32 and start with small router weights.
6. Choose between capacity and dropless execution based on your kernels and whether your inference must be batch-independent.
7. Before deploying, estimate total-parameter memory and the all-to-all volume per layer for your target hardware.