Deciding which experts to drop

Every mixture-of-experts model has experts that barely earn their parameters. The tempting move is to rank experts by how often the router picks them and delete the bottom few. That ranking is wrong in a specific, recoverable way, and fixing it takes a second number. This post is the whole method: where the utilisation figure comes from, why weight-space cosine similarity is a usable proxy for functional redundancy, how the two combine into an asymmetric keep-score, and which of the resulting numbers you should refuse to believe.

Utilisation alone gets it backwards

Consider two experts in the same layer. Expert A takes 4% of routed tokens and is the only expert that handles a particular kind of input — say, code in a model that mostly sees prose. Expert B takes 22% and is nearly identical to a sibling that takes 30%. Utilisation ranks A last and says drop it. Capability says drop B, whose work is already covered, while A is the only thing standing between you and a hole in the model.

This is not a hypothetical failure mode. Load-balancing losses during MoE training actively push the utilisation distribution flat, which means the surviving variance in utilisation is weakly informative at best — and the experts that escape the balancing pressure are often the specialists you least want to lose. So the pruner computes two numbers per expert, and the plan comes from their interaction rather than from either alone.

Number one: utilisation, from a seeded router

The first number is share of routed tokens. Getting it honestly would mean running the model, which the browser cannot do. Instead the page runs a seeded simulation: synthetic token representations produce logits through a per-expert bias plus structured variation, softmax over the top-k, tally the winners across every MoE layer.

MK.simulateRouting = function (model, opts) {
  var E = model.moe.expertCount, K = Math.min(opts.topk || 2, E);
  var rng = MK.rng(opts.seed || 1);                       // deterministic
  var bias = []; for (var e = 0; e < E; e++) bias.push((rng() - 0.35) * 1.7);

  for (var li = 0; li < layers.length; li++)
    for (var t = 0; t < nTok; t++) {
      var logits = [];
      for (var e = 0; e < E; e++)
        logits.push(bias[e] + (rng() * 2 - 1) * 1.3 + Math.sin((t+1)*(e+1)*0.7 + li) * 0.5);
      var idx = logits.map((_, i) => i).sort((a, b) => logits[b] - logits[a]).slice(0, K);
      idx.forEach(i => loads[i]++);                        // utilisation
      for (var a = 0; a < idx.length; a++)                 // co-activation matrix
        for (var b = a + 1; b < idx.length; b++) { coact[idx[a]][idx[b]]++; coact[idx[b]][idx[a]]++; }
    }
  return { E, K, loads, util: loads.map(x => x / (nTok * layers.length)), coact, … };
};

The sinusoidal term is not decoration: without it the top-k selection is pure noise and every expert converges to 1/E, which produces a board with nothing on it. The term creates a token-dependent preference structure — different tokens systematically prefer different experts — which is the qualitative property real routers have and the one the visualisation is about.

The simulation is seeded, so the same checkpoint always produces the same board. That matters more than it sounds: an explainability tool whose numbers move between refreshes teaches you nothing, because you can never tell a real change from a re-roll.

The number the co-activation matrix gives you for free

Because top-k picks k experts per token, the simulation also records which pairs fire together. Two experts with high co-activation are being used jointly — they compose rather than substitute — which is an argument against pruning either, independent of similarity. High similarity plus high co-activation is the genuinely ambiguous case, and the tool shows both rather than collapsing them.

Number two: redundancy, from weight signatures

The second number asks a different question: if this expert vanished, is there another that does approximately the same thing?

Answering it properly means comparing full weight matrices — tens of millions of parameters per expert, which is neither loadable in a tab nor necessary. Instead each expert is reduced to a fixed-length signature by folding its strided samples into a small accumulator, using the same machinery the visualiser uses for heatmaps:

MK.expertSignatures = function (model, dim) {
  dim = dim || 48;
  var sig = [];
  for (var e = 0; e < E; e++) {
    var acc = new Float32Array(dim), cnt = 0;
    model.moe.layers.forEach(function (Lr) {
      (Lr.experts[e] || []).forEach(function (t) {
        var s = model.sample(t, dim * 8);                 // 384 strided values
        for (var i = 0; i < s.values.length; i++) acc[i % dim] += s.values[i];
        cnt += s.values.length;
      });
    });
    sig.push(acc);                                        // one 48-vector per expert
  }
  return sig;
};

MK.cosine = function (a, b) {                             // then all-pairs cosine
  var d = 0, na = 0, nb = 0;
  for (var i = 0; i < a.length; i++) { d += a[i]*b[i]; na += a[i]*a[i]; nb += b[i]*b[i]; }
  return (na && nb) ? d / Math.sqrt(na * nb) : 0;
};

Forty-eight floats per expert instead of tens of millions, and an all-pairs comparison that is E²·48 multiply-adds — microseconds for any E you'll meet.

Why this works at all, and where it stops working

The honest framing: cosine similarity between strided, position-folded weight samples is a coarse proxy for functional similarity, and it earns its place on one property — it is sensitive to shared structure and insensitive to scale. Two experts trained toward the same function end up with weights pointing in a similar direction even when their magnitudes differ, and the normalisation in cosine discards exactly the magnitude difference you don't care about.

The permutation blind spot is the one to keep in mind. It means the tool's similarity scores are a lower bound on true redundancy: what it flags is real, what it misses may still be duplicated. Fixing it properly needs an assignment solve (Hungarian matching between neuron sets) before comparison, which is a real algorithm and not a 120-line-toolkit algorithm. Absent that, folding by i % dim at least makes the signature depend on many neurons rather than the first 48, which softens the failure without curing it.

The keep-score, and why it's asymmetric

The two numbers combine into a single ranking: normalised utilisation, minus a penalty proportional to how similar an expert is to its most-used twin.

score(e) = util(e) / max_util
         − λ · max over j≠e of [ sim(e, j) · 1{ util(j) > util(e) } ]

The indicator is the whole design. An expert's penalty comes only from its similarity to a more heavily used twin, so of a near-identical pair the busier one is penalised by nothing and survives, while the quieter one absorbs the full penalty and becomes the candidate. Penalise both symmetrically and you happily delete both halves of a pair — losing a capability entirely while congratulating yourself on double the parameter savings. Every naive implementation of this idea has that bug.

The plan, and the number you should distrust

A slider sets aggressiveness. The page takes the lowest keep-scores up to that fraction of experts and reports experts pruned, parameters cut, share of expert weight removed, and an estimated capability retained:

function planAt(pct) {
  var k = Math.round(A.E * pct / 100);
  var pruned = {}; A.ranked.slice(A.E - k).forEach(function (x) { pruned[x.e] = true; });
  A.experts.forEach(function (x) {
    totalUtil += x.util;
    if (pruned[x.e]) { removedParams += x.params; prunedUtil += x.util; }
  });
  // retained ≈ 1 − (share of routed tokens whose preferred expert is gone),
  //            discounted by how redundant those experts were
  …
}

What the tool is genuinely good for is the shape of the answer: which experts cluster, whether the router's load is lopsided (the header prints an imbalance ratio — max load over mean load, where 1.0× is perfectly balanced and anything past ~2× is worth investigating), and roughly how much of the model is duplicated effort. Those conclusions survive the approximations. The retention percentage does not.

Using it yourself

The honest workflow this tool fits into. Browser analysis is step one of four, and it is the cheap one:

1. Screen here. Load the checkpoint, read the similarity matrix, note the candidate pairs and the imbalance ratio. Download the report. Five minutes, no GPU.
2. Measure real utilisation. Hook the router in your inference stack and count expert selections over a few thousand batches of your traffic. This replaces the simulated number with a measured one and is the single highest-value upgrade to everything above.
3. Prefer merging to deleting. For a flagged pair, a utilisation-weighted average of the two experts' weights usually retains more than dropping one does, because it keeps the contribution of the quieter expert's distinct directions.
4. Evaluate, then heal. Measure on a real eval; a short fine-tune afterwards recovers a surprising amount, and its size tells you how much you actually broke.

The measured-utilisation step in step two is about ten lines against a torch model:

# count real expert selections, then feed them back in place of the simulation
from collections import Counter
counts = Counter()

def hook(module, inputs, output):            # output: router logits [tokens, E]
    topk = output.topk(2, dim=-1).indices
    counts.update(topk.flatten().tolist())

for layer in model.model.layers:
    layer.block_sparse_moe.gate.register_forward_hook(hook)

for batch in your_real_traffic:              # a few thousand is plenty
    model(**batch)

util = {e: c / sum(counts.values()) for e, c in counts.items()}

With that dict in hand, the keep-score formula above is the same three lines and the danger quadrant becomes a measurement rather than a guess.

The report

The download button generates a standalone HTML report — the ranking, the plan, the parameters, the caveats — assembled as a string in the page and handed to the browser as a Blob:

var blob = new Blob([html], { type: "text/html" });
var url  = URL.createObjectURL(blob);
a.href = url; a.download = "prune-report.html"; a.click();
setTimeout(function () { URL.revokeObjectURL(url); }, 4000);

Nothing is uploaded to produce it. You can hand the file to a colleague and it opens with no dependencies, which felt like the right shape for an artefact whose entire argument is "here is what I measured, here is what I inferred, and here is which of the two each number is."

Related: how the checkpoint gets read in the first place — reading a checkpoint in the browser.