Expert-Choice Routing

Standard routing computes an affinity matrix $S \in \mathbb{R}^{T \times E}$ and takes the top-$k$ along the expert axis: each token picks its experts. Every problem with capacity follows from that, because nothing constrains how many tokens pick the same one.

Take the top-$k$ along the other axis instead.

Expert $e$ takes the $C$ tokens that score highest for it. Every expert takes exactly $C$ — no more, by construction, and no fewer as long as $T \ge C$. There is no capacity factor, no auxiliary loss, no controller, and no dropping.

Variable compute per token

The interesting consequence is not the balance but its dual. Under token choice, every token gets exactly $k$ experts. Under expert choice, a token may be chosen by many experts, one, or none.

The flip side is that a token chosen by nobody passes through unchanged. This is the same event as a dropped token, arrived at by a different route, and it is untroubling for the same reason — a shared expert or the residual path means the layer is never wholly absent.

Why decoders cannot use it

$\mathrm{top}_C(S_{:,e})$ ranks all $T$ tokens against each other. Whether token 5 is selected depends on how tokens 6 through $T$ scored — so the computation at position 5 depends on the future, which is exactly what causal masking exists to prevent.

The damage is not subtle. During training the model sees a routing that used future information; during autoregressive generation that routing cannot be reproduced, because the future does not exist yet. Train/inference mismatch of this kind does not degrade gracefully.

Zhou et al. evaluated on encoder-style and fine-tuning workloads where the whole sequence is present, which is where the method is sound. For decoder pre-training it is not repairable — a causal variant would have to rank each token against only its predecessors, and then the expert’s top-$C$ changes at every position.

import torch
from torch import Tensor, nn


class ExpertChoiceRouter(nn.Module):
    """Non-causal. Encoders, vision, and fine-tuning on full sequences only."""

    def __init__(self, dim: int, n_experts: int, capacity_factor: float = 2.0):
        super().__init__()
        self.gate = nn.Linear(dim, n_experts, bias=False)
        self.f, self.n_experts = capacity_factor, n_experts

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:    # [T, D]
        t = x.size(0)
        capacity = int(self.f * t / self.n_experts)

        scores = self.gate(x).softmax(-1)                     # [T, E]

        # top-k down the token axis: each expert takes its best `capacity`.
        weight, tokens = scores.t().topk(capacity, dim=-1)    # [E, C]
        return tokens, weight

One line differs from token-choice routing — the .t() — and it changes what the layer is. Worth noting how cheap the balance is once the assignment runs this way: the entire apparatus of capacity factors, auxiliary losses and bias controllers exists to approximate, under token choice, what a transpose gives exactly.

Soft MoE takes the same insight further and drops discreteness altogether: each expert processes a weighted average of all tokens, with the weights from the same affinity matrix. Perfectly balanced, fully differentiable, no dropping — and non-causal for the same reason, since the average is over the whole sequence. Both methods are the same trade: give up autoregression, get balance for free.