Mixture of Tokens

Mixture of Tokens is a Transformer architecture, presented at NeurIPS 2024, designed as a more stable alternative to the sparse Mixture of Experts (MoE) approach used to scale large language models. Like MoE, it lets a model carry a very large parameter count without a proportional increase in computation, keeping most of the network dormant on any given token. Unlike MoE, it avoids making discrete routing decisions during training, which removes the instability and load-imbalance problems that plague sparse routing.

How it differs from Mixture of Experts

In a standard MoE, a router sends each token to a small number of specialized expert sub-networks — a hard, discrete choice that can be unstable to train and can leave some experts overloaded while others sit idle. Mixture of Tokens instead has each expert process a weighted average, a soft mixture, of tokens drawn from across a group of examples. Because the mixing is a continuous, differentiable operation with no hard selection, the model trains as stably as a vanilla Transformer, and every expert receives the same number of tokens, so load imbalance disappears by construction rather than being fought with auxiliary penalties. In effect, the architecture replaces a scheduling problem with a smooth mathematical operation that the optimizer already knows how to handle.

Why the stability matters

The discrete routing at the heart of MoE is what makes it both powerful and fragile. Small changes in the gating network can flip a token from one expert to another, producing noisy gradients, and expert “collapse” — where the router learns to favour a handful of experts — wastes most of the model's capacity. Practitioners fight these problems with load-balancing losses, capacity limits, and careful tuning, all of which add complexity and can still fail at scale. Mixture of Tokens sidesteps both by never forcing a choice. The reported result is competitive quality with a vanilla Transformer at a several-fold reduction in compute and a meaningful wall-clock speedup, while remaining compatible with standard masked and causal language-model training, so it slots into existing pipelines rather than demanding a new one.

The trade-off worth naming

Nothing is free. Mixing tokens across examples raises questions about how information leaks between sequences in a batch, and the technique has to be designed carefully so that causal language modelling — where a token must never see the future — still holds. The authors address this within the method, but it is the kind of subtlety a self-hoster should keep in mind when evaluating any soft-routing architecture rather than assuming the stability comes with no strings attached.

Why it matters for sovereign AI

Efficiency techniques like this are central to making capable models cheaper to train and run, which directly affects whether strong open-weight models become small and efficient enough to self-host on accessible hardware rather than remaining locked behind data-center-scale infrastructure. Every architecture that squeezes more quality out of a fixed compute budget widens the set of people who can own their own model instead of renting access to someone else's.

It is worth placing Mixture of Tokens in context: it is one entry in a fast-moving line of research into softer, more trainable alternatives to hard routing, and it has not displaced sparse Mixture of Experts in the largest deployed models, which still lean on discrete routing despite its difficulties. The value of the idea is as much conceptual as practical. It demonstrates that the compute savings of conditional computation do not strictly require a hard, discrete gate, which reframes the design space and invites hybrids that keep some of MoE's specialization while borrowing the stability of soft mixing. For a self-hoster, the takeaway is to judge an architecture by measured tokens-per-second and quality on real hardware, not by its parameter count alone.

Mixture of Tokens is one of several efficiency strategies behind modern foundation models, and the compute savings it targets echo the cost concerns addressed by Chinchilla-optimal training. It is a direct counterpoint to the sparse routing that expert parallelism was built to scale.