Info
Last Execution: 2026-08-19
| Package | Version |
|---|---|
| nnsight | 0.7.1 |
| Python | 3.12.0 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Analysing Routing in a Mixture-of-Experts Model¶
📗 You can find an interactive Colab version of this tutorial here.
A mixture-of-experts (MoE) layer replaces a transformer block's single MLP with many parallel MLPs — the experts — and a small linear router that picks a few of them per token. Only the picked experts run, so the model can hold billions of parameters while spending the compute of a much smaller one.
That makes routing an unusually attractive target for interpretability. Unlike a direction in the residual stream, the routing decision is discrete, named and countable: for every token you get a short list of expert ids. It is tempting to read those lists as a map of what the model has specialised in, and to test that map by deleting an expert and watching the loss.
Both halves of that plan have a trap in them, and this tutorial is mostly about the traps. We will:
- find the router and work out what its output actually is, without being told;
- anchor the reading — check that the experts we think were selected really are the top-k of the logits we read, and that the weights really are the softmax of those logits;
- show that the obvious expert ablation — masking a router logit to
-inf— is a silent no-op on thetransformersbackend, understand why, and replace it with one that works; - show that even a working logit mask is biased, because this model does not renormalise its top-k weights;
- measure what routing depends on (domain, language, paraphrase), each against a control;
- handle the padded-batch trap, which has an MoE-specific twist: the router's output has no batch axis;
- see how far a single generic script gets across other MoE families.
Our model is allenai/OLMoE-1B-7B-0924: 6.9B total parameters, 1.3B active, 16 layers, 64 experts per layer, top-8 routing. It is small enough to run on CPU, which is what we do here — it needs about 14 GB of RAM in bfloat16 and no GPU at all.
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight
from IPython.display import clear_output
import torch
import torch.nn.functional as F
import nnsight
from nnsight import TransformersModel
import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
We load the model on CPU in bfloat16. Two loading notes worth knowing before you copy this line:
- Use
device="cpu", notdevice_map="cpu"— the latter is currently ignored and the model lands oncuda:0. - Do not reach for
dtype="nf4"to shrink an MoE. Ontransformers5 the experts are stored as 3-Dnn.Parameterstacks rather thannn.Linearlayers, andbitsandbytesonly swapsnn.Linear. On this model that quantises the 64 attention projections and none of the 6.4B parameters of expert weights — measured at 12.53 GiB against 12.89 GiB forbfloat16. Worse for us, it would perturb the router's input while leaving the router exact, which is the one confound a routing study cannot afford.
model = TransformersModel(
"allenai/OLMoE-1B-7B-0924",
task="text-generation",
dtype="bfloat16",
device="cpu",
dispatch=True,
)
clear_output()
print(model)
OlmoeForCausalLM(
(model): OlmoeModel(
(embed_tokens): Embedding(50304, 2048, padding_idx=1)
(layers): ModuleList(
(0-15): 16 x OlmoeDecoderLayer(
(self_attn): OlmoeAttention(
(q_proj): Linear(in_features=2048, out_features=2048, bias=False)
(k_proj): Linear(in_features=2048, out_features=2048, bias=False)
(v_proj): Linear(in_features=2048, out_features=2048, bias=False)
(o_proj): Linear(in_features=2048, out_features=2048, bias=False)
(q_norm): OlmoeRMSNorm((2048,), eps=1e-05)
(k_norm): OlmoeRMSNorm((2048,), eps=1e-05)
)
(mlp): OlmoeSparseMoeBlock(
(gate): OlmoeTopKRouter()
(experts): OlmoeExperts(
(act_fn): SiLUActivation()
)
)
(input_layernorm): OlmoeRMSNorm((2048,), eps=1e-05)
(post_attention_layernorm): OlmoeRMSNorm((2048,), eps=1e-05)
)
)
(norm): OlmoeRMSNorm((2048,), eps=1e-05)
(rotary_emb): OlmoeRotaryEmbedding()
)
(lm_head): Linear(in_features=2048, out_features=50304, bias=False)
(generator): Generator(
(streamer): Streamer()
)
)
1. Finding the router¶
The printout above ends each layer with an OlmoeSparseMoeBlock instead of the usual MLP. Let's look at one on its own.
cfg = model.config
E, K = cfg.num_experts, cfg.num_experts_per_tok
N_LAYERS = cfg.num_hidden_layers
print(f"{N_LAYERS} layers, {E} experts per layer, top-{K} routing")
LAYER = 8
mlp = model.model.layers[LAYER].mlp
gate = mlp.gate
print(mlp)
16 layers, 64 experts per layer, top-8 routing
OlmoeSparseMoeBlock(
(gate): OlmoeTopKRouter()
(experts): OlmoeExperts(
(act_fn): SiLUActivation()
)
)
Two children: a gate (the router) and an experts module. Note what is not there — there is no experts[0], experts[1], ... to hook. OlmoeExperts keeps all 64 experts in two 3-D parameter stacks and loops over slices of them, so an individual expert is not an nn.Module and has no envoy of its own.
for name, p in mlp._module.named_parameters():
print(f"{name:28s} {tuple(p.shape)} {p.dtype}")
gate.weight (64, 2048) torch.bfloat16 experts.gate_up_proj (64, 2048, 2048) torch.bfloat16 experts.down_proj (64, 2048, 1024) torch.bfloat16
So the router is a 2048 -> 64 projection and every expert lives inside those two (64, ...) stacks. That tells us where the routing decision is made, but not what the module returns. Rather than guess, we can print the module's forward with .source, which renders the actual Python of forward with a label for every traceable operation inside it. This works without weights loaded, so you can do it before committing to a 14 GB download.
Start with the parent, not the router. .source is usually introduced as a way to reach values inside a forward, but it is just as useful for the opposite question: what does the caller do with a child's return value?
print(mlp.source)
* def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
0 batch_size, sequence_length, hidden_dim = hidden_states.shape
hidden_states_view_0 -> 1 hidden_states = hidden_states.view(-1, hidden_dim)
self_gate_0 -> 2 _, top_k_weights, top_k_index = self.gate(hidden_states)
self_experts_0 -> 3 final_hidden_states = self.experts(hidden_states, top_k_index, top_k_weights).reshape(
reshape_0 -> + ...
4 batch_size, sequence_length, hidden_dim
5 )
6 return final_hidden_states
7
Line 2 is the whole story:
_, top_k_weights, top_k_index = self.gate(hidden_states)
The router returns three things, and the block throws the first one away. Keep that _ in mind — it is the entire subject of section 3.
Line 1 matters too: hidden_states.view(-1, hidden_dim) flattens batch and sequence together before the router is called, so everything the router sees and returns is indexed (batch * seq, ...). That is section 6.
Now the router's own forward.
print(gate.source)
* def forward(self, hidden_states):
hidden_states_reshape_0 -> 0 hidden_states = hidden_states.reshape(-1, self.hidden_dim)
F_linear_0 -> 1 router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts)
torch_nn_functional_softmax_0 -> 2 router_probs = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)
torch_topk_0 -> 3 router_top_value, router_indices = torch.topk(router_probs, self.top_k, dim=-1) # (seq_len, top_k)
4 if self.norm_topk_prob:
router_top_value_sum_0 -> 5 router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
router_top_value_to_0 -> 6 router_top_value = router_top_value.to(router_logits.dtype)
7 router_scores = router_top_value
8 return router_logits, router_scores, router_indices
9
OlmoeTopKRouter.forward is six lines: a linear projection to router_logits, a softmax over all 64 experts, a topk, an optional renormalisation guarded by self.norm_topk_prob, and a return of (router_logits, router_scores, router_indices).
That is everything we need, and none of it requires .source to read — the three values are the module's ordinary output. Let's take them.
prompt = "The capital of France is Paris and the capital of Germany is"
with torch.no_grad():
with model.trace(prompt):
router_in = gate.input.save() # .input BEFORE .output!
router_logits, routing_weights, expert_idx = gate.output
router_logits = router_logits.save()
routing_weights = routing_weights.save()
expert_idx = expert_idx.save()
print("gate.input ", tuple(router_in.shape), router_in.dtype)
print("gate.output[0] logits ", tuple(router_logits.shape), router_logits.dtype)
print("gate.output[1] weights", tuple(routing_weights.shape), routing_weights.dtype)
print("gate.output[2] indices", tuple(expert_idx.shape), expert_idx.dtype)
print()
n_tokens = len(model.tokenizer(prompt)["input_ids"])
print(f"the prompt is {n_tokens} tokens long")
gate.input (12, 2048) torch.bfloat16 gate.output[0] logits (12, 64) torch.bfloat16 gate.output[1] weights (12, 8) torch.bfloat16 gate.output[2] indices (12, 8) torch.int64 the prompt is 12 tokens long
Two things to notice.
First, the shapes are (T, E), (T, K), (T, K) — there is no batch axis. A 12-token prompt gives (12, 64), not (1, 12, 64), because of that view(-1, hidden_dim) in the parent. With one sequence this is harmless; with a padded batch it is not, and section 6 shows what goes wrong.
Second, we asked for gate.input before gate.output. .input and .output of the same module are two ordered events in the forward pass, and requesting them out of order raises OutOfOrderError naming the exact value. The same rule bites later: read mlp.gate.output before mlp.output, since the router runs first.
We can also see the routing decision for one token directly.
tokens = [model.tokenizer.decode(t) for t in model.tokenizer(prompt)["input_ids"]]
print(f"routing for the final token {tokens[-1]!r}:")
order = routing_weights[-1].float().argsort(descending=True)
for rank, slot in enumerate(order.tolist()):
e = int(expert_idx[-1][slot])
print(f" rank {rank}: expert {e:2d} weight {routing_weights[-1][slot].item():.4f}"
f" logit {router_logits[-1][e].item():+.3f}")
print(f" these {K} weights sum to {routing_weights[-1].float().sum():.4f}")
[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer GPTNeoXTokenizer. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.
routing for the final token ' is': rank 0: expert 18 weight 0.0742 logit +0.475 rank 1: expert 62 weight 0.0688 logit +0.395 rank 2: expert 15 weight 0.0474 logit +0.024 rank 3: expert 59 weight 0.0304 logit -0.422 rank 4: expert 0 weight 0.0288 logit -0.475 rank 5: expert 40 weight 0.0281 logit -0.500 rank 6: expert 60 weight 0.0270 logit -0.539 rank 7: expert 33 weight 0.0254 logit -0.602 these 8 weights sum to 0.3301
2. Anchoring the wiring¶
We now have three tensors that look like a routing decision. Before we build any statistics on them we should check that they are wired together the way we think, because a plausible-looking tensor read from the wrong place produces plausible-looking science. Two claims are worth testing, and both are one line:
Anchor A — the experts in gate.output[2] are the top-8 of the logits in gate.output[0].
Anchor B — the weights in gate.output[1] are the softmax of those logits, gathered at those indices.
probs = torch.softmax(router_logits.float(), dim=-1)
my_idx = probs.topk(K, dim=-1).indices
print("A: topk(softmax(logits)) == gate.output[2], exactly and in rank order :",
torch.equal(my_idx, expert_idx))
gathered = probs.gather(-1, expert_idx)
err = (routing_weights.float() - gathered).abs().max().item()
print(f"B: max|weights - softmax(logits)[idx]| = {err:.3e}")
print(f" (full 64-way softmax row sums: {probs.sum(-1).min():.6f} .. {probs.sum(-1).max():.6f})")
A: topk(softmax(logits)) == gate.output[2], exactly and in rank order : True B: max|weights - softmax(logits)[idx]| = 8.268e-04 (full 64-way softmax row sums: 1.000000 .. 1.000000)
Anchor A passes exactly — not just as a set, but in the same rank order. Anchor B passes to 8e-4, which is the resolution of bfloat16 around weights of order 0.05: the router computes its softmax in float32 and casts the result back down on the way out (router_top_value.to(router_logits.dtype), line 6 of the source above), so this is rounding, not a wiring error. Load the model in float32 and the residual shrinks accordingly.
Now a third check that is not a wiring test but a configuration test, and which turns out to matter more than either of the above.
mass = routing_weights.float().sum(-1)
print(f"B': top-{K} routing weights sum to {mass.mean():.4f} on average "
f"(min {mass.min():.4f}, max {mass.max():.4f})")
print(f" config.norm_topk_prob = {cfg.norm_topk_prob}")
renormalised = gathered / gathered.sum(-1, keepdim=True)
print(f" max|weights - RENORMALISED softmax| = "
f"{(routing_weights.float() - renormalised).abs().max():.3f}")
B': top-8 routing weights sum to 0.4412 on average (min 0.3254, max 0.7766)
config.norm_topk_prob = False
max|weights - RENORMALISED softmax| = 0.190
The top-8 weights do not sum to 1. Because norm_topk_prob=False, this model uses the raw 64-way softmax probabilities as its mixing weights, so the eight chosen experts carry only about 44% of the probability mass and the MoE block's output is scaled down accordingly. Many MoE models (Mixtral, Qwen3-MoE, GPT-OSS) set this to True and their weights sum to 1.0; this one does not.
That is a two-line config detail, and it is the reason the second ablation trap in section 4 exists. Check norm_topk_prob before you ablate anything.
One more anchor, cheap and worth having: recompute the logits ourselves from the router's input and weight matrix.
W = gate.weight # (E, hidden)
recomputed = F.linear(router_in.float(), W.float())
print(f"C: max|W @ h - gate.output[0]| = "
f"{(recomputed - router_logits.float()).abs().max():.3e}"
f" (logit std = {router_logits.float().std():.3f})")
C: max|W @ h - gate.output[0]| = 9.927e-03 (logit std = 1.031)
Consistent with bfloat16 output rounding on logits whose spread is about 1.0. The router is a plain linear map of the pre-MoE hidden state, and we are reading the right tensor.
3. Ablating an expert: the trap¶
The natural causal experiment is: remove one expert, see what breaks. And the natural way to remove an expert from a top-k router is to set its logit to -inf so it can never be picked. nnsight's own vLLM documentation says exactly this — "To ablate an expert, mask its router logit to -inf in mlp.gate.output" — and on vLLM it is correct, because there mlp.gate is a ReplicatedLinear whose output is the logits the block consumes.
On TransformersModel it does nothing at all. Let's not ablate one expert; let's set all 64 logits to -inf and watch the model not care.
with torch.no_grad():
with model.trace(prompt):
_, _, base_idx = gate.output
base_idx = base_idx.save()
base_mlp = mlp.output.save()
base_lgts = model.output.logits.save()
with torch.no_grad():
with model.trace(prompt):
gate.output[0][:, :] = -float("inf") # every logit, every token
noop_mlp = mlp.output.save()
noop_lgts = model.output.logits.save()
print(f"max|delta| on the MoE block output : {(noop_mlp - base_mlp ).abs().max():.6f}")
print(f"max|delta| on the model's logits : {(noop_lgts - base_lgts).abs().max():.6f}")
max|delta| on the MoE block output : 0.000000 max|delta| on the model's logits : 0.000000
Every router logit in layer 8 is negative infinity and nothing moved. No exception, no warning, no clue.
The reason is the line we read in section 1:
_, top_k_weights, top_k_index = self.gate(hidden_states)
Element [0] of the router's return is a vestigial value that the block discards. nnsight did exactly what it promises — it handed us the live tensor and our write landed on it — but no downstream operation ever reads that tensor again.
This is the most dangerous failure mode in intervention work, because it does not look like a failure. Had we skipped this check and gone straight to "ablate expert 18 and measure the loss", we would have reported that expert 18 does not matter, with a clean zero to prove it, and been wrong about all 64 experts.
The general rule: an intervention that lands on a value the caller discards is indistinguishable from an intervention that had no effect. Before trusting an ablation, read the parent's
.sourceand confirm the value you are writing is one the caller actually consumes.
Three sites in this block are consumed. Two are the other elements of the router's output — the weights [1] and the indices [2]. The third is the pre-softmax logits inside the router's forward, which we can reach with .source at the F_linear_0 operation we saw earlier. Let's compare all of them on the same target.
TARGET = int(base_idx[-1, 0]) # the top-ranked expert at the last token
print(f"target expert: {TARGET}\n")
def delta(mlp_out, logits, label):
print(f"{label:52s} d_mlp {(mlp_out[0, -1] - base_mlp[0, -1]).float().norm():.4f}"
f" d_logits {(logits - base_lgts).float().norm():7.2f}")
# (a) the vLLM recipe: mask the returned logits
with torch.no_grad():
with model.trace(prompt):
gate.output[0][:, TARGET] = -float("inf")
m = mlp.output.save(); l = model.output.logits.save()
delta(m, l, "(a) gate.output[0] logit <- -inf")
# (b) zero the routing weight wherever this expert was selected
with torch.no_grad():
with model.trace(prompt):
_, w, i = gate.output
w[i == TARGET] = 0.0
m = mlp.output.save(); l = model.output.logits.save()
delta(m, l, "(b) gate.output[1] weight <- 0")
# (c) selection ablation: re-run the top-k with this expert excluded,
# but gather the weights from the UNTOUCHED softmax
with torch.no_grad():
with model.trace(prompt):
lo, w, i = gate.output
p = torch.softmax(lo.float(), dim=-1)
p_masked = p.clone()
p_masked[:, TARGET] = -1.0
new_i = p_masked.topk(K, dim=-1).indices
new_w = p.gather(-1, new_i).to(w.dtype)
hit = (i == TARGET).any(-1, keepdim=True) # touch only affected tokens
gate.output = (lo, torch.where(hit, new_w, w), torch.where(hit, new_i, i))
m = mlp.output.save(); l = model.output.logits.save()
delta(m, l, "(c) gate.output = (logits, new weights, new indices)")
# (d) mask the logit inside the router, before the softmax
with torch.no_grad():
with model.trace(prompt): # warm-up trace for .source
pass
with torch.no_grad():
with model.trace(prompt):
gate.source.F_linear_0.output[:, TARGET] = -float("inf")
m = mlp.output.save(); l = model.output.logits.save()
delta(m, l, "(d) gate.source.F_linear_0.output <- -inf")
# reference: delete the whole MoE block
with torch.no_grad():
with model.trace(prompt):
mlp.output[:] = 0.0
m = mlp.output.save(); l = model.output.logits.save()
delta(m, l, " (reference) zero the entire MoE block")
target expert: 18
(a) gate.output[0] logit <- -inf d_mlp 0.0000 d_logits 0.00
(b) gate.output[1] weight <- 0 d_mlp 0.1687 d_logits 125.58
(c) gate.output = (logits, new weights, new indices) d_mlp 0.1715 d_logits 125.28
(d) gate.source.F_linear_0.output <- -inf d_mlp 0.1706 d_logits 123.77
(reference) zero the entire MoE block d_mlp 0.3133 d_logits 876.74
Sites (b), (c) and (d) all move the model; site (a) is flat zero. Note that the .source intervention in (d) needed a throwaway trace first — the first .source access on a module raises OutOfOrderError when the block has already read something from it, once per module, and a warm-up trace clears it.
(b), (c) and (d) land within a couple of percent of each other on this prompt, but they are not the same experiment, and section 4 shows where they come apart. They are three different counterfactuals:
- (b) zero the weight. The expert's contribution is removed and nothing replaces it, so the block's total mixing mass drops from ~0.44 to ~0.44 − w.
- (c) selection ablation. The expert is removed from the shortlist and the 9th-ranked expert takes its slot, with the weight it would have had. Mixing mass stays where it was. This is "what would this layer have done with 63 experts".
- (d) mask the logit pre-softmax. Same selection change as (c) — plus the softmax denominator shrinks, so every surviving weight gets slightly larger.
Which one is right depends on the question, but only one of them survives the control in the next section.
4. The second trap: which ablation is honest?¶
Any expert-ablation study needs a null. The cleanest one available is: an expert that no token in the corpus routes to must produce exactly zero change. If a method reports a nonzero effect for an expert the model never uses, then whatever number it reports for a real expert includes that much fiction.
We run each of the three working ablations on four experts at layer 8 — the most used, the median, the rarest that is used at all, and one that is never selected — over a short prose corpus. The effect size is the mean KL divergence between the original and ablated next-token distributions, which is zero if and only if nothing changed.
PROSE_ABL = "\n\n".join([
"The harbour at Bergen has been in continuous use since the eleventh century, when merchants "
"from the Hanseatic League established a trading post on the eastern shore. Timber, dried cod and "
"iron ore moved south through the fjords, while grain and cloth arrived from the Baltic ports. The "
"wooden warehouses that line the quay were rebuilt many times after fires, most recently in the "
"eighteenth century, and they now house museums and small restaurants rather than cargo.",
"Photosynthesis converts light energy into chemical energy stored in carbohydrates. In plants the "
"reaction takes place inside chloroplasts, organelles that contain stacked membranes called "
"thylakoids. Chlorophyll molecules embedded in those membranes absorb photons, chiefly in the blue "
"and red parts of the spectrum, and transfer the excitation to a reaction centre where water is "
"split and oxygen is released as a by-product.",
"The trial of Galileo before the Roman Inquisition in 1633 turned less on astronomy than on "
"authority. Galileo had published a dialogue in which the Copernican account of the heavens was "
"defended by the more persuasive of two speakers, and the church held that he had violated an "
"earlier injunction. He was compelled to abjure, and spent the remainder of his life under house "
"arrest at Arcetri.",
])
abl_ids = model.tokenizer(PROSE_ABL, return_tensors="pt")["input_ids"]
print("corpus:", abl_ids.shape[1], "tokens")
with torch.no_grad():
with model.trace(abl_ids):
_, _, corpus_idx = gate.output
corpus_idx = corpus_idx.save()
corpus_lgts = model.output.logits.save()
base_logp = F.log_softmax(corpus_lgts[0].float(), dim=-1)
def mean_kl(new_logits):
lp = F.log_softmax(new_logits[0].float(), dim=-1)
return float((base_logp.exp() * (base_logp - lp)).sum(-1).mean())
counts = torch.bincount(corpus_idx.reshape(-1), minlength=E)
print(f"{int(counts.sum())} routing slots, {int((counts > 0).sum())} of {E} experts used at least once")
corpus: 259 tokens
2072 routing slots, 62 of 64 experts used at least once
def ablate_expert(target, method):
with torch.no_grad():
with model.trace(abl_ids):
if method == "logit": # mask pre-softmax, via .source
gate.source.F_linear_0.output[:, target] = -float("inf")
else:
lo, w, i = gate.output
if method == "weight": # zero the weight
w[i == target] = 0.0
else: # selection ablation
p = torch.softmax(lo.float(), dim=-1)
pm = p.clone()
pm[:, target] = -1.0
ni = pm.topk(K, dim=-1).indices
nw = p.gather(-1, ni).to(w.dtype)
hit = (i == target).any(-1, keepdim=True)
gate.output = (lo, torch.where(hit, nw, w),
torch.where(hit, ni, i))
out = model.output.logits.save()
return mean_kl(out)
order = counts.argsort(descending=True)
rarest_n = min(c for c in counts.tolist() if c > 0)
picks = [("most used", int(order[0])),
("median", int(order[E // 2])),
(f"rarest ({rarest_n} sel)", [j for j in range(E) if counts[j] == rarest_n][0]),
("NEVER selected", int(order[-1]))]
with torch.no_grad():
with model.trace(abl_ids[:, :4]): # warm up .source
pass
print(f"{'expert':22s} {'id':>4s} {'sel':>5s} {'logit mask':>12s} {'weight->0':>12s} {'selection':>12s}")
for label, e in picks:
row = [ablate_expert(e, m) for m in ("logit", "weight", "selection")]
print(f"{label:22s} {e:4d} {int(counts[e]):5d} " + " ".join(f"{v:12.6f}" for v in row))
with torch.no_grad():
with model.trace(abl_ids):
mlp.output[:] = 0.0
out = model.output.logits.save()
print(f"\n{'zero the whole MoE block':22s} {'':4s} {'':5s} {mean_kl(out):12.6f}")
expert id sel logit mask weight->0 selection
most used 14 109 0.008330 0.008020 0.008198
median 2 29 0.001261 0.001031 0.001273
rarest (1 sel) 16 1 0.000462 0.000160 0.000159
NEVER selected 27 0 0.000379 0.000000 0.000000
zero the whole MoE block 0.069978
Read the bottom row of the table first. Zeroing the weight and ablating the selection both give exactly 0.000000 for an expert that no token routes to. They pass the null.
Masking the logit does not. It reports 0.00038 for expert 27, which the model never once selects. That is not a small error in absolute terms — it is more than twice the real effect of the rarest expert that is used (0.00016), and about a third of the median expert's. An expert-importance ranking built from logit masks is noise in its entire lower half.
The cause is norm_topk_prob=False from section 2. Because the routing weights are raw softmax probabilities, removing one expert from the softmax shrinks the denominator and inflates every surviving weight, whether or not the removed expert was ever in the top-8. The whole MoE block simply gets slightly louder.
For scale, note the last line of the table: deleting the whole MoE block costs 0.070, and the single most-used expert accounts for 0.0082 of that — 12% of the block's causal contribution from 5% of its routing slots. Slightly super-proportional, but nothing like a bottleneck. No one expert of the 64 is load-bearing at this layer.
never_e = picks[-1][1]
with torch.no_grad():
with model.trace(abl_ids):
lo, w, i = gate.output
lo, w, i = lo.save(), w.save(), i.save()
p = torch.softmax(lo.float(), dim=-1)
print(f"expert {never_e} is selected by {int(counts[never_e])} of {int(counts.sum())} routing slots,")
print(f"but carries {p[:, never_e].mean():.4f} of the softmax mass on the average token.")
p_masked = p.clone()
p_masked[:, never_e] = 0.0
p_masked = p_masked / p_masked.sum(-1, keepdim=True)
ratio = (p_masked.gather(-1, i) / p.gather(-1, i)).mean()
print(f"masking its logit rescales every surviving weight by x{ratio:.6f}")
expert 27 is selected by 0 of 2072 routing slots, but carries 0.0113 of the softmax mass on the average token. masking its logit rescales every surviving weight by x1.011404
A 1.1% uniform gain on the MoE block, applied at every token, purely from deleting an expert that was never used. That is the whole spurious signal — and it is a systematic one, not noise: it is larger for experts that hold more softmax mass, which correlates with how often they are selected, so it inflates exactly the numbers you were hoping to compare.
The rule: ablate by selection, not by logit — recompute the top-k with the expert excluded, and gather the weights from the untouched softmax. Then verify with a never-selected expert, which must give exactly zero.
Zeroing the weight passes the null too, and it is a legitimate experiment; it just answers a different question. It removes the expert's contribution without letting the router replace it, so it measures "how much did this expert contribute" rather than "what would this layer do without it". Report which one you ran.
(If your model sets norm_topk_prob=True, the logit mask is much better behaved, because the renormalisation cancels the denominator change for the surviving experts. It is still worth running the null.)
A footnote on that hit guard¶
The selection ablation above contains one line whose purpose is not obvious:
hit = (i == target).any(-1, keepdim=True)
gate.output = (lo, torch.where(hit, nw, w), torch.where(hit, ni, i))
We rebuild the top-k for every token but only install it where the target expert was actually selected. That looks redundant — for a token that never picked the expert, our recomputed top-k should equal the model's. It very nearly does, and "very nearly" is enough to break the control.
with torch.no_grad():
with model.trace(abl_ids):
lo, w, i = gate.output
p = torch.softmax(lo.float(), dim=-1)
pm = p.clone()
pm[:, never_e] = -1.0
ni = pm.topk(K, dim=-1).indices
set_differs = (ni.sort(-1).values != i.sort(-1).values).any(-1).sum().save()
order_differs = (ni != i).sum().save()
nw = p.gather(-1, ni).to(w.dtype)
weight_err = (nw.float() - w.float()).abs().max().save()
print(f"recomputing the top-{K} with a NEVER-selected expert masked out, "
f"{abl_ids.shape[1]} tokens:")
print(f" tokens whose selected SET changed : {int(set_differs)}")
print(f" slots whose rank ORDER changed : {int(order_differs)}")
print(f" max |our weights - model weights| : {float(weight_err):.3e}")
recomputing the top-8 with a NEVER-selected expert masked out, 259 tokens: tokens whose selected SET changed : 3 slots whose rank ORDER changed : 13 max |our weights - model weights| : 0.000e+00
Three tokens out of 259 come back with a different set of experts, even though the expert we masked was never in any of them. Our softmax and the model's are not bit-identical — F.softmax(x, dtype=torch.float) and torch.softmax(x.float()) can disagree in the last bit — and on those tokens the 8th and 9th experts are close enough that the last bit decides. Reordering within the top-k is harmless (the same experts get the same weights), but a changed set is a real intervention on a token we did not mean to touch, and it leaks straight into the never-selected control. Without the guard the control reports a nonzero effect and the whole point of the null is lost.
The general lesson is the same one as section 3, from the other side: when you replace a value the model computed, replace only the part you mean to change, and derive everything else from the model's own tensors rather than from your reimplementation of them.
5. What does routing actually depend on?¶
With reading anchored and ablation understood, we can ask the descriptive question: do different kinds of text use different experts?
We need a corpus with several clearly different domains. We write it out inline so nothing depends on a downloaded dataset: English encyclopaedic prose, source code, and the same kind of prose in French and German.
EN_PROSE = [
"The harbour at Bergen has been in continuous use since the eleventh century, when merchants from the Hanseatic League established a trading post on the eastern shore. Timber, dried cod and iron ore moved south through the fjords, while grain and cloth arrived from the Baltic ports. The wooden warehouses that line the quay were rebuilt many times after fires, most recently in the eighteenth century, and they now house museums and small restaurants rather than cargo.",
"Photosynthesis converts light energy into chemical energy stored in carbohydrates. In plants the reaction takes place inside chloroplasts, organelles that contain stacked membranes called thylakoids. Chlorophyll molecules embedded in those membranes absorb photons, chiefly in the blue and red parts of the spectrum, and transfer the excitation to a reaction centre where water is split and oxygen is released as a by-product. The resulting electrons drive a chain of carriers that pumps protons across the membrane.",
"The trial of Galileo before the Roman Inquisition in 1633 turned less on astronomy than on authority. Galileo had published a dialogue in which the Copernican account of the heavens was defended by the more persuasive of two speakers, and the church held that he had violated an earlier injunction. He was compelled to abjure, and spent the remainder of his life under house arrest at Arcetri, where he nonetheless completed his most durable work on the motion of falling bodies.",
"Monetary policy operates through several channels. When a central bank raises its policy rate, short-term borrowing becomes more expensive, which discourages investment and durable consumption. The exchange rate tends to appreciate as capital flows towards higher yields, dampening exports. Asset prices fall, reducing household wealth and the collateral available to firms. Each channel acts with a lag of several quarters, which is why forecasters disagree so persistently about timing.",
"Coral reefs occupy less than one per cent of the ocean floor yet shelter roughly a quarter of all marine species. The reef structure is built by colonial polyps that deposit calcium carbonate and host symbiotic algae in their tissues. When water temperatures remain elevated for several weeks the polyps expel the algae, losing both their colour and their principal source of nutrition. Repeated bleaching events have reduced coverage on many reefs by more than half since the nineteen eighties.",
"The printing press did not create literacy, but it changed what literacy was for. Manuscript culture had treated a book as a scarce object to be annotated and copied; print made identical copies cheap enough that readers could own several, compare them, and disagree with them in the margins. Standardised page numbers and indexes appeared within two generations, and with them the habit of citing a passage rather than remembering it.",
]
CODE = [
'''def merge_intervals(intervals):
if not intervals:
return []
intervals = sorted(intervals, key=lambda x: x[0])
merged = [list(intervals[0])]
for start, end in intervals[1:]:
if start <= merged[-1][1]:
merged[-1][1] = max(merged[-1][1], end)
else:
merged.append([start, end])
return merged
class LRUCache:
def __init__(self, capacity):
self.capacity = capacity
self.store = {}
self.order = []
def get(self, key):
if key not in self.store:
return -1
self.order.remove(key)
self.order.append(key)
return self.store[key]
''',
'''import numpy as np
def kmeans(X, k, iters=100, seed=0):
rng = np.random.default_rng(seed)
idx = rng.choice(len(X), size=k, replace=False)
centers = X[idx].copy()
for _ in range(iters):
d = ((X[:, None, :] - centers[None, :, :]) ** 2).sum(-1)
labels = d.argmin(1)
new = np.stack([X[labels == j].mean(0) if (labels == j).any() else centers[j]
for j in range(k)])
if np.allclose(new, centers):
break
centers = new
return centers, labels
''',
'''SELECT customer_id,
COUNT(*) AS order_count,
SUM(total_amount) AS lifetime_value
FROM orders
WHERE order_date >= '2023-01-01'
AND status NOT IN ('cancelled', 'refunded')
GROUP BY customer_id
HAVING COUNT(*) > 3
ORDER BY lifetime_value DESC
LIMIT 100;
CREATE INDEX idx_orders_customer_date ON orders (customer_id, order_date);
UPDATE inventory
SET quantity = quantity - 1
WHERE sku = 'AB-1042'
AND quantity > 0;
''',
]
FR = [
"Le port de Bergen est utilise sans interruption depuis le onzieme siecle, lorsque des marchands de la Ligue hanseatique y etablirent un comptoir sur la rive orientale. Le bois, la morue sechee et le minerai de fer descendaient les fjords vers le sud, tandis que le grain et le drap arrivaient des ports de la Baltique. Les entrepots de bois qui bordent le quai ont ete reconstruits bien des fois apres les incendies.",
"La photosynthese transforme l'energie lumineuse en energie chimique stockee dans les glucides. Chez les plantes, la reaction se deroule dans les chloroplastes, des organites qui contiennent des membranes empilees appelees thylakoides. Les molecules de chlorophylle absorbent les photons, surtout dans le bleu et le rouge, et transferent l'excitation vers un centre reactionnel ou l'eau est scindee et l'oxygene libere comme sous-produit.",
"La politique monetaire agit par plusieurs canaux. Lorsque la banque centrale releve son taux directeur, l'emprunt a court terme devient plus couteux, ce qui decourage l'investissement. Le taux de change tend a s'apprecier, ce qui freine les exportations. Les prix des actifs baissent, reduisant la richesse des menages. Chaque canal agit avec un decalage de plusieurs trimestres.",
]
DE = [
"Der Hafen von Bergen wird seit dem elften Jahrhundert ununterbrochen genutzt, als Kaufleute der Hanse am oestlichen Ufer eine Handelsniederlassung gruendeten. Holz, Trockenfisch und Eisenerz gelangten durch die Fjorde nach Sueden, waehrend Getreide und Tuch aus den Ostseehaefen eintrafen. Die hoelzernen Lagerhaeuser am Kai wurden nach Braenden mehrfach wieder aufgebaut und beherbergen heute Museen und kleine Gaststaetten statt Waren.",
"Die Photosynthese wandelt Lichtenergie in chemische Energie um, die in Kohlenhydraten gespeichert wird. Bei Pflanzen laeuft die Reaktion in den Chloroplasten ab, Organellen mit gestapelten Membranen, den Thylakoiden. Chlorophyllmolekuele nehmen Photonen auf, vor allem im blauen und roten Bereich, und uebertragen die Anregung auf ein Reaktionszentrum, in dem Wasser gespalten und Sauerstoff freigesetzt wird.",
"Die Geldpolitik wirkt ueber mehrere Kanaele. Erhoeht eine Zentralbank ihren Leitzins, wird die kurzfristige Kreditaufnahme teurer, was Investitionen und langlebige Konsumgueter daempft. Der Wechselkurs neigt zur Aufwertung, was die Ausfuhren bremst. Die Vermoegenspreise fallen und mindern das Haushaltsvermoegen. Jeder Kanal wirkt mit einer Verzoegerung von mehreren Quartalen.",
]
GROUPS = {"en_prose": EN_PROSE, "code": CODE, "fr": FR, "de": DE}
Now a capture function. We want the expert indices at every layer for every token, which is one trace per sequence with 16 saves in it.
Two idioms are worth pointing out. First, we wrap the whole thing in torch.no_grad() — a read-only trace otherwise runs with autograd on, which costs roughly 18x the memory for nothing. Second, a list comprehension of .save() calls bound to a local name inside the trace body does not survive the block; use nnsight.save([...]) on the whole list instead.
One deliberate omission: normally you would put these sequences into a single trace with several tracer.invoke calls, which runs them in one forward instead of fifteen and is one of the clearest wins nnsight offers over raw hooks. We do not, because our sequences have different lengths and batching them means padding them — and for this statistic padding is exactly the thing that goes wrong. Section 6 shows what it costs and how to batch safely once you know.
layers = model.model.layers
def capture_routing(text):
# selected expert ids at every layer -> (n_layers, n_tokens, K)
with torch.no_grad():
with model.trace(text):
per_layer = nnsight.save([layers[l].mlp.gate.output[2]
for l in range(N_LAYERS)])
return torch.stack(per_layer)
captured = {}
for name, texts in GROUPS.items():
captured[name] = torch.cat([capture_routing(t) for t in texts], dim=1)
print(f"{name:9s} {captured[name].shape[1]:5d} tokens "
f"expert ids {tuple(captured[name].shape)}")
en_prose 547 tokens expert ids (16, 547, 8)
code 534 tokens expert ids (16, 534, 8)
fr 360 tokens expert ids (16, 360, 8)
de 410 tokens expert ids (16, 410, 8)
5a. Expert usage by domain¶
The simplest summary of routing is a usage histogram: for a given layer, how often is each of the 64 experts selected, as a fraction of all routing slots. Let's look at English prose against code.
def usage(idx_LTK, layer):
c = torch.bincount(idx_LTK[layer].reshape(-1), minlength=E).float()
return c / c.sum()
SHOW = 12
rows, labels = [], []
for name in GROUPS:
rows.append(usage(captured[name], SHOW).numpy())
labels.append(name)
fig = px.imshow(
rows,
y=labels,
labels={"x": "expert", "y": "", "color": "usage fraction"},
color_continuous_scale="Viridis",
title=f"Expert usage at layer {SHOW} (fraction of routing slots)",
aspect="auto",
)
fig.show()