Mixture of experts and speculative decoding

Arjun’s evaluation spreadsheet had a clear winner and a clear problem. On Crumb’s test set of 300 real customer chats, a 70-billion-parameter open model answered Marathi questions far better than their 8-billion one: 91% judged correct against 74%. Meera wanted it. Then Arjun estimated what it would take to run it.

“In 16-bit it’s 140 GB of weights. That’s two big GPUs only to hold it. And every token it writes has to read all 140 GB from memory. Our 8B model reads 16. So roughly nine times slower per token, on hardware that costs four or five times as much.”

“So better answers cost ten times as much,” said Meera. “That’s not a trade I can make for a bakery.”

“There are models that know about as much as a big one but run closer to the speed of a small one,” said Priya. “And there’s a trick that makes any big model write faster without changing a single word it writes. Both are about not doing work you don’t need.”

Why bigger models are slower

Recall from Module 30 that each transformer block has two parts: attention, where tokens share information, and a feed-forward network (FFN), a small two-layer network each token passes through on its own. The FFNs hold about two thirds of a model’s weights, and researchers believe much of a model’s stored knowledge lives there.

In an ordinary, dense model, every token goes through every weight. A rough rule: generating one token takes about 2 arithmetic operations per parameter. So an 8-billion-parameter model does about 16 billion operations per token, and a 70-billion one about 140 billion. And as you saw two lessons ago, decoding is mostly limited by reading weights from memory, so the 70B model must also move about nine times more bytes per token.

Arjun’s naive plan, “run the bigger model”, pays that full cost on every single token, including tokens like “the” and “your” that need no deep knowledge at all.

Priya’s kitchen of specialists

“How does the Crumb central kitchen handle the Diwali rush?” Priya asked.

“Stations. Bread, cakes, mithai, savouries. And Ravi at the pass reads each ticket and sends it to the right stations.”

“Does a kaju katli order go through the bread station?”

“No. Mithai station, maybe packing. Two stations, not all of them.”

“So the kitchen as a whole knows how to make hundreds of things, but each order only uses the stations it needs. A big kitchen’s knowledge, a small kitchen’s work per order. Now replace every FFN in a transformer with a set of FFN ‘stations’, called experts, and put a tiny Ravi in front of them, called a router, who sends each token to its best two.”

That is a mixture of experts, or MoE. The model’s total parameter count (all experts in all layers) can be huge, but the active parameters (the ones actually used for a given token) stay small.

Where the analogy breaks: Ravi sends orders by what they are. A learned router sends tokens by whatever patterns training found useful, and those rarely line up with human categories like “Marathi” or “baking”. Researchers who studied the Mixtral model’s routing found experts specialised more by kinds of token and syntax than by obvious topics.

Following one token through an MoE layer

Take the token “mawa” arriving at one MoE layer with 4 experts, where the router picks the top 2.

Step 1: the router scores every expert. The router is a single small linear layer. For this token it outputs scores (logits) of:

Expert	1	2	3	4
Router score	2.0	0.5	1.5	-1.0

Step 2: keep the top 2. Experts 1 and 3 win. Experts 2 and 4 are skipped entirely: their weights are never read, their arithmetic never happens.

Step 3: turn the winners’ scores into weights. Softmax over the two kept scores, [2.0, 1.5], gives 0.622 and 0.378. (Some models instead apply softmax over all experts, giving 0.532 and 0.323, and then renormalise; the result is the same 0.622 and 0.378.)

Step 4: run the two experts and blend. Suppose expert 1’s output for this token is [1.0, 0.2] and expert 3’s is [0.4, 0.8] (real outputs have thousands of numbers):

$$0.622\times [1.0,0.2]+0.378\times [0.4,0.8]=[0.773,0.427]$$

That blend is the layer’s output for “mawa”, and it flows on to the next block. The next token, perhaps “cake”, might go to experts 2 and 3. Each token gets its own route, at every MoE layer.

Count the work: this token used 2 of 4 experts, so half the FFN compute of a dense layer the same total size. A real model with 8 experts and top-2 uses a quarter; with 256 experts and 8 active, about 3%.

Opening the hood: the MoE layer precisely

For a token vector $x$, the router computes one score per expert, $r=x{W}_r$. Let $\text{TopK}$ be the indices of the $k$ highest scores. The layer’s output is:

$$y=\sum _{i\in \text{TopK}}{g}_i(x)E_i(x),g(x)=\text{softmax}\left(r_{\text{TopK}}\right)$$

$E_i$ is expert $i$, an ordinary FFN; $g_i$ is its gate weight. Attention layers are usually left dense and shared; only the FFNs become experts.

The best-known open example is Mistral’s Mixtral 8x7B (2023): 8 experts per layer, top-2 routing. Mistral reported about 47 billion total parameters and about 13 billion active per token. (Not 8 × 7 = 56, because attention and embeddings are shared, not duplicated.) It matched or beat much larger dense models of its time on many benchmarks, at roughly the compute per token of a 13B dense model.

Keeping the experts busy

Left alone, routers develop favourites. A few experts get most tokens, train faster, get chosen even more, and the rest wither. Worse, in large deployments each expert may sit on a different GPU, so an overloaded expert becomes a traffic jam.

The classic fix (from Google’s Switch Transformer, 2021) adds a small load-balancing loss during training. For $N$ experts, if $f_i$ is the fraction of tokens sent to expert $i$ and $P_i$ the average router probability for it:

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

This is smallest when tokens are spread evenly, so it nudges the router to share work. Systems also set an expert capacity: if too many tokens in a batch pick the same expert, the extras are dropped or rerouted. The DeepSeek case study later in this module shows a newer approach that avoids the extra loss.

A toy MoE layer in PyTorch

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

class Expert(nn.Module):
    """One expert: an ordinary feed-forward network."""
    def __init__(self, d_model, d_hidden):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_model, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_model))
    def forward(self, x):
        return self.net(x)

class MoELayer(nn.Module):
    def __init__(self, d_model=64, d_hidden=256, n_experts=8, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList(Expert(d_model, d_hidden) for _ in range(n_experts))
        self.router = nn.Linear(d_model, n_experts, bias=False)   # one score per expert
        self.top_k = top_k

    def forward(self, x):                                  # x: (tokens, d_model)
        scores = self.router(x)                            # (tokens, n_experts)
        top_scores, top_idx = scores.topk(self.top_k, dim=-1)
        gates = F.softmax(top_scores, dim=-1)              # weights over the chosen experts
        out = torch.zeros_like(x)
        for slot in range(self.top_k):
            for e, expert in enumerate(self.experts):
                chosen = top_idx[:, slot] == e             # tokens sending this slot to expert e
                if chosen.any():                           # skipped experts do no work at all
                    out[chosen] += gates[chosen, slot, None] * expert(x[chosen])
        return out, top_idx

topk picks the best k experts per token. The loops send each group of tokens only to the experts that chose them; an expert no token picked never runs. Real implementations replace the loops with fast batched kernels, but the logic is the same.

torch.manual_seed(0)
moe = MoELayer()
out, chosen = moe(torch.randn(10, 64))                     # 10 tokens
total = sum(p.numel() for p in moe.parameters())
active = 2 * sum(p.numel() for p in moe.experts[0].parameters()) + moe.router.weight.numel()
print(f"total parameters: {total:,}   active per token: {active:,}")
print("tokens per expert:", torch.bincount(chosen.flatten(), minlength=8).tolist())

total parameters: 265,216   active per token: 66,688
tokens per expert: [1, 1, 3, 2, 4, 4, 4, 1]

A quarter of the parameters do the work for any one token. And even in this untrained toy, the load is uneven: three experts got 4 tokens each, three got 1. Training without a balancing mechanism makes that worse, not better.

The second trick: guess ahead, check in bulk

MoE makes each token cheaper. Speculative decoding attacks a different waste. Remember from the KV-cache lesson that decoding writes one token per forward pass, and each pass is limited by reading the weights from memory, not by arithmetic. The GPU’s arithmetic units sit mostly idle, waiting for data.

But checking several tokens at once is almost as cheap as writing one. Feeding the model five tokens in one pass, as it does during prefill, reads the weights once and computes five sets of predictions in parallel. So if someone could guess the next five tokens, the big model could verify all five for roughly the price of writing one.

“It’s piping rosettes on the Christmas cakes,” said Priya. “The apprentice is quick but sometimes wrong. The senior baker is slow but always right. So the apprentice pipes five rosettes fast, and the senior looks at all five in one glance. Keep every rosette up to the first bad one, the senior fixes that one, and the apprentice carries on.”

“And if the apprentice is good, the senior hardly ever has to pipe.”

“Exactly. And the cake ends up exactly as if the senior had done every rosette.”

Following one reply through speculative decoding

The prompt so far: “Your order is out for”. The big target model is Crumb’s 70B model. The small draft model is a 1B model from the same family, sharing the same tokenizer.

Step 1: the draft guesses 4 tokens, one at a time, cheaply: delivery, and, will, arrive.

Step 2: the target checks them in one pass. It reads “Your order is out for delivery and will arrive” and, at each of the four new positions plus one beyond, says what it would have written:

Position	Draft guessed	Target would write	Verdict
1	delivery	delivery	accept
2	and	and	accept
3	will	will	accept
4	arrive	reach	reject
5	(not checked)	-	-

Step 3: keep the agreed tokens, plus the target’s correction. Output: “delivery and will reach”. Four tokens from a single target pass, instead of one. The draft’s “arrive” is thrown away, along with anything after it.

Step 4: repeat from “Your order is out for delivery and will reach”.

If every guess had been accepted, the target’s pass would also give a fifth token for free, the one at position 5. If the very first guess had been wrong, you would still get one token from the target, exactly what you would have got without speculation. You never get less than normal decoding, apart from the draft’s own cost.

How many tokens per target pass, on average? If the draft’s guesses are accepted with probability $\alpha$ each and it drafts $k$ tokens, the expected number is:

$$\frac{1-{\alpha }^{k+1}}{1-\alpha }$$

With $\alpha =0.8$ and $k=4$, that is about 3.4 tokens per target pass. Subtract the draft model’s running cost and published systems typically report speedups of around 2 to 3 times.

Opening the hood: why the output is exactly the target’s

With greedy decoding (always pick the most likely token), exactness is easy to see: you only keep tokens the target itself would have picked.

With sampling (Module 01’s temperature), it needs a cleverer rule, introduced in 2023 by researchers at Google and, independently, DeepMind. Let $p$ be the target’s probability for the drafted token and $q$ the draft’s. Accept the token with probability $\min (1,p/q)$. If rejected, sample a replacement from the leftover distribution $\max (0,p-q)$, renormalised. The proof is short and the result is remarkable: the final text follows exactly the target model’s probability distribution. Speculative decoding is a pure speed trick, never a quality trade.

Here is a minimal greedy version, using GPT-2 Medium (355M parameters) as the target and DistilGPT-2 (82M) as the draft. It skips the KV cache to stay readable.

# uv add transformers torch
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("gpt2")
target = AutoModelForCausalLM.from_pretrained("gpt2-medium").eval()   # careful senior baker
draft = AutoModelForCausalLM.from_pretrained("distilgpt2").eval()     # quick apprentice

@torch.no_grad()
def speculative_generate(prompt, n_new=40, k=4):
    ids = tok(prompt, return_tensors="pt")["input_ids"]
    stop_at = ids.shape[1] + n_new
    target_passes = 0
    while ids.shape[1] < stop_at:
        # 1) The draft guesses k tokens, one at a time (cheap)
        guess = ids
        for _ in range(k):
            next_id = draft(guess).logits[0, -1].argmax()
            guess = torch.cat([guess, next_id.view(1, 1)], dim=1)
        drafted = guess[0, ids.shape[1]:]
        # 2) The target checks all k guesses in ONE forward pass
        logits = target(guess).logits[0, ids.shape[1] - 1:]
        target_passes += 1
        target_choice = logits.argmax(dim=-1)        # what the target would pick at each slot
        # 3) Accept guesses up to the first disagreement, then take the target's token
        n_ok = 0
        while n_ok < k and drafted[n_ok] == target_choice[n_ok]:
            n_ok += 1
        ids = torch.cat([ids, drafted[:n_ok].view(1, -1), target_choice[n_ok].view(1, 1)], dim=1)
    return tok.decode(ids[0]), target_passes

text, passes = speculative_generate("Thank you for your order. Your eggless chocolate cake")
print(passes, "target passes")
print(text)

On Arjun’s laptop this wrote 40 tokens in 18 target passes instead of 40, and the text was character-for-character identical to the target generating alone. Hugging Face has this built in: pass assistant_model=draft to target.generate(...), and vLLM supports several speculative methods in production.

Where it breaks

Small speedups, or slowdowns, on the wrong hardware. Here is the honest part of Arjun’s experiment. On his laptop CPU, the target alone took 2.0 seconds; Hugging Face’s assisted generation took 2.8 seconds; his cache-free toy took 5.2. Fewer target passes, yet slower. Speculation pays off when the target is much bigger than the draft and decoding is memory-bound, as with large models on GPUs. With two small models on a CPU, the draft’s own cost and the extra bookkeeping eat the gain. Measure on your real setup.

Draft and target must agree often. Acceptance depends on how similar the two models are and on the text. Formulaic replies (“Your order is out for delivery”) get high acceptance; creative writing gets low acceptance. Variants exploit this: prompt lookup decoding drafts by copying spans from the prompt itself, which works brilliantly when answers quote a retrieved document, as in RAG. Other variants (Medusa, EAGLE, and the multi-token prediction heads in DeepSeek-V3) bolt small guessing heads onto the target itself, so no separate draft model is needed.

Same tokenizer required. A classic draft-and-target pair must share a vocabulary, which usually means the same model family.

MoE serving is harder than it looks. Memory holds every expert, uneven routing creates hot spots, and with many GPUs, tokens must be shuffled to whichever GPU holds their expert. MoE models also tend to be touchier to fine-tune than dense ones.

Try it yourself

1. A model card says “236B total, 21B active parameters”. Roughly how many operations per token does it need, and how much memory for weights in 16-bit?
2. Using the formula, how many tokens per target pass do you expect with $\alpha =0.6$ and $k=4$? With $\alpha =0.9$?

Show the answer

1. About $2\times 21=42$ billion operations per token (it computes like a 21B model), but about $236\times 2=472$ GB of memory for weights in 16-bit (it must store everything). That is why such models are served across several GPUs, often quantised.
2. With $\alpha =0.6$: $(1-0.6^5)/(1-0.6)=(1-0.078)/0.4\approx 2.3$ tokens per pass. With $\alpha =0.9$: $(1-0.9^5)/0.1=(1-0.590)/0.1\approx 4.1$. A better draft matters a lot.

Common mistakes

• Treating active parameters as memory needs. MoE cuts compute, not memory. Budget for the total.
• Assuming experts are topic specialists. Routing is learned and rarely maps to human categories. Do not expect a “Marathi expert”.
• Thinking speculative decoding changes answers. With the proper acceptance rule it does not, so you can test it without re-evaluating quality, only speed.
• Benchmarking speculation on a laptop and generalising. Its gains depend on model sizes, hardware and the kind of text.
• Pairing models with different tokenizers. Classic speculative decoding needs the same vocabulary.

Recap

• Dense models pay for every parameter on every token; mixture-of-experts models route each token through only a few expert FFNs.
• A small router scores experts, keeps the top $k$, and blends their outputs with softmax weights.
• MoE gives large total knowledge at small active compute, but needs memory for all experts and mechanisms to balance the load.
• Speculative decoding has a cheap draft model guess several tokens and the big model verify them in one pass, keeping everything up to the first disagreement.
• With the right acceptance rule, speculative decoding reproduces the target model’s output exactly; its speedup depends on acceptance rate, model sizes and hardware.

Arjun took the idea to Meera with a new shortlist: an open MoE model with about 47 billion total parameters, quantised to 4 bits, and an 8B dense model as the fallback.

“Think of the central kitchen at Diwali,” he said. “Lots of specialist stations, but each order only visits two. This model knows a lot but only does a small model’s work per word. It still needs a big enough fridge to hold every station, so we’d need a bigger card than today’s. And for the replies, a small assistant model guesses ahead and the big one approves several words at a glance. Same words, written faster.”

“Same words?” Meera checked.

“Exactly the same. It’s a speed trick, not a quality trick.”

“Then what makes the words good? Not only correct, but polite, the way we’d want? And why do some models now ‘think’ before they answer?”

Priya smiled. “That’s not architecture. That’s how they’re trained after pretraining. Next lesson: RLHF, and the step-by-step thinking tricks.”

Go deeper

• vLLM documentation, including mixture-of-experts serving and speculative decoding.
• Hugging Face docs, see the Transformers guides on assisted generation.
• Attention Is All You Need, for the dense transformer that MoE modifies.