~/satyajit

A6B: k-expansion, or what breaks when you force a top-8 MoE router to fire 32 experts

mdjsonmcp

2026-07-08 · 11 min · explainer · llm · architecture · inference-optimization

a6b-k-expansion is a single-author research log with a sharp, testable question: a Mixture-of-Experts model already stores far more experts than it fires per token — so what happens if you just fire more of them at inference? The author takes Qwen3.6-35B-A3B (their stated base: 35B total, ~3B active, 256 routed experts per layer, native top-8 routing) and overrides the router to keep top-32 instead of top-8. Active compute goes from ~3B to ~6.6B parameters per token — the regime they call "A6B" — at zero new weights. Then they measure what it costs, and try to heal it.

The interesting part is that it does not work for free, and the repo says so in the title of its own findings: the damage from naive k-expansion is smooth and monotonic — there is no free sweet spot. This is a measurement-first log with paired A/B tests, exact McNemar significance, and the negative results kept in. That honesty is the reason it's worth reading.

The one-line edit

Every MoE layer scores all 256 experts with a softmax router (in fp32), keeps the top-k, and — this is the load-bearing detail — renormalizes the selected gate weights so they sum to 1. Renormalization is baked into the architecture. So raising k is not a matter of appending a few experts on the side; it redistributes the entire gate mass across four times as many slots.

Two side-by-side MoE blocks. Left: stock A3B, router keeps top-8 of 256 experts (blue cells), ≈3B active per token, top-8 carry 46% of gate mass. Right: A6B, same weights but router keeps top-32; the 8 blue experts stay and 24 orange experts (ranks 9-32) are added, ≈6.6B active per token, the orange ranks carrying 54% of the renormalized gate mass. Both feed a renormalized combine plus an always-on shared expert.
Identical weights; one routing constant changed. The added experts (orange) are load-bearing — 54% of the renormalized gate mass — but were never trained to collaborate (a6b-k-expansion, Fig 1).

The experts are stored as packed 3D tensors — gate_up_proj [256, 1024, 2048] and down_proj [256, 2048, 512] per layer — so widening k adds compute but no parameters. If the routing here is unfamiliar, I built the top-k gate up from nothing in Mixture of Experts, from scratch; this article assumes that machinery and pokes at one constant inside it.

Formally, for the selected set the renormalized gate weight of expert ii is

g~i=gijtop-kgj,itop-k.\tilde{g}_i = \frac{g_i}{\sum_{j \in \text{top-}k} g_j}, \qquad i \in \text{top-}k .

The denominator grows with k. So every already-selected expert's g~i\tilde{g}_i shrinks as k rises, and the freed mass flows to the newcomers. The author profiled the router at k=32 and found the newcomers are not a rounding error: ranks 9-32 carry 54.0% of the renormalized gate mass. More than half the block's output now comes from experts the model never trained to fire together.

Slide k below and watch the mass move off the trained top-8 and onto the untrained tail — with the measured accuracy for each k next to it:

router gate mass · renormalized to top-32illustrative profile
rank 8ranks 9-321256 experts →gate mass · sums to 1 after renormtop-8 · 46%untrained · 54%
top-k
trained top-8 untrained tail
active params
~6.6B
untrained mass
54%
MMLU (measured)
0.8067
GSM8K (measured)
0.8650

The router scores all 256 experts, keeps the top-32, and renormalizes the selected gates to sum to 1. So raising k does not just append experts — it shrinks every existing expert's weight and hands the difference to the newcomers. At k=32 the added ranks 9-32 carry 54% of the mass, yet the model was never trained to co-activate them. Zero new weights are added — only compute. Accuracy falls at every step.

No sweet spot

You might hope for a lucky k — 2-3× the native width, where the extra experts add capacity before they add noise. There isn't one. Changing only the inference-time top-k on the frozen base model, both a knowledge benchmark and a math benchmark decline at every single step:

Benchmarkk = 8k = 16k = 24k = 32
MMLU0.84330.82830.81500.8067
GSM8K0.89330.88830.87830.8650

The shape is the whole point: monotonic, sweet-spot-free. Every step down is paid the moment more untrained expert combinations switch on, and the noise grows with the mass those combinations carry — a mass that is real (54.0% on ranks 9-32). This is what motivates healing the model rather than searching for a lucky k. There is nothing to search for.

Healing without touching the router

The fix is deliberately surgical, in the spirit of ESFT — Expert-Specialized Fine-Tuning (Wang et al., 2024). The recipe:

  1. Profile. Run the target corpus through the model at k=32, count token-level routing frequency for every expert across all 40 layers.
  2. Select. Keep experts by cumulative routing frequency up to top-p 0.2 — the ones actually carrying the work. That is 833 of the 10,240 layer-experts (40 × 256).
  3. Train residual deltas. Add trainable delta tensors to the selected experts' FFN slices and train only those deltas2.62B of 35B params (7.5%). The router and everything else stay frozen.
  4. Toggle. Because nothing but the deltas moved, turning them off returns the exact stock model. The deltas ship as one 5.2 GB artifact, patched on or off at load time.

Toggle the deltas below to see the surgical footprint and the gap-to-baseline flip from "significant loss" to "statistical tie":

selective expert fine-tune · router frozen833 / 10,240 layer-experts
deltas on · mixed ESFT (Gen 2)
trained delta frozen
trainable
2.62B / 35B · 7.5%
artifact
5.2 GB · toggleable
MMLU vs base@k8
−2.0 pt
p=0.141 · statistical tie
GSM8K vs base@k8
−2.2 pt
p=0.263 · statistical tie

Healing trains only residual deltas on the 60 highest routing-frequency experts shown here (the real run selects 833 of 10,240 by top-p 0.2) — 2.62B of 35B params, router and everything else frozen. The deltas ship as one 5.2 GB artifact you patch on or off at load time; toggle them off and the exact stock model returns. On heals a significant loss into a statistical tie — a repair, not a win. The point estimate stays negative.

Freezing the router is not just a cost decision. The deltas are learned relative to a specific routing distribution — freeze it and every delta keeps meaning the same thing at inference; let it move and the coordinate system shifts underneath the deltas. Second, a router that drifts under SFT is the classic path to routing collapse, where a handful of experts capture all the traffic. A frozen router removes that failure mode entirely.

Each healing generation is a broader training corpus, measured as the gap to base@k8 on its own machine (negative = still below native top-8), with the verdict from exact McNemar at p = 0.05:

GenerationCorpusMMLU ΔMMLU pGSM8K ΔGSM8K p
Gen 0naive k32 (no training)−3.7 pt0.002 · loss−2.8 pt0.016 · loss
Gen 1agentic-only ESFT−3.0 pt0.010 · loss−0.8 pt0.487 · tie
Gen 2mixed + replay ESFT−2.0 pt0.141 · tie−2.2 pt0.263 · tie

Each broader corpus removes more of the misalignment: the agentic-only patch already heals math to a tie, and adding coding, tool-calling, math and a small general/knowledge replay is what finally pulls MMLU into a tie too. But read the header honestly — a tie is a repair, not a gain. The MMLU point estimate is still −2.0 pt; it is no longer statistically distinguishable from base, but it is not zero.

And it does not get better with more steps. A checkpoint trajectory reads MMLU 0.825 / 0.823 / 0.820 at steps 1200 / 2100 / 3150 — healing saturates by ~38% of training. The author reads this as a capacity ceiling of the selective deltas, not a data-volume problem: more steps on this delta set will not close the gap; a larger trainable surface or a different objective would be needed.

Where training does buy something beyond healing is code. On HumanEval (n = 164), the agentic patch reaches 0.902 — above both base@k8 and naive k32 — while compressing median generation to about a third of the tokens:

HumanEval (%) — paired, same-machine (Workstation A)
agentic patch
90.2
base@k8
86.6
naive k32
84.1
coding-only patch
76.2
050100

The hazard: style transfers before knowledge

That fourth bar is the most useful result in the whole repo. A coding-only patch — same recipe, single-domain corpus — taught the model a style (terse code) faster than it taught it to stay correct. Median generation crashed to 186 tokens, and accuracy collapsed with it: HumanEval 0.762 (the worst of every arm) and GSM8K 0.820, which is below even naive k32, at p = 0.002. You can train a confident, compact, and wrong model this way.

The mixed patch is not immune to the same pressure, only more resistant. It still compressed MBPP generations (median 531 vs base's 2852 tokens) and lost MBPP by −10.2 pt (p < .0001) — even while HumanEval stayed at parity. Two benchmarks that both "test coding" diverged sharply, and the difference is how much each rewards long, explicit generation, which the patch has learned to suppress. Because this compression grows with training length, an earlier checkpoint may beat the final one on generation-heavy tasks; that evaluation is still running.

Why I trust the numbers

The measurement discipline is where this log earns its credibility, and it's the part most self-reported results skip:

Training data also passed a hard decontamination gate against every benchmark used — exact-match, word-13-gram, short-question containment, HumanEval signature purge, Terminal-Bench instruction match — with the knowledge-replay slice further screened by embedding similarity against the full MMLU test set. Reported result: 0 residual hits. You can disagree with the conclusions, but the instrument is honest about what it measured.

The take

k-expansion is a clean idea with a clean negative result. Firing 4× the experts at inference is free in weights but not in accuracy, because renormalization is not a bystander — it hands 54% of every MoE block's output to 24 experts per layer that never learned to work together, and the model degrades monotonically for it. There is no lucky k. Router-frozen selective deltas — 2.62B trainable params, toggleable, 833 of 10,240 experts — heal the knowledge loss back to a statistical tie, which is a genuine result and an honest one: a tie, saturating at ~38% of training, with the point estimate still negative. The coding axis actually improves (HumanEval 0.902), but generation-length compression is a live hazard, and a single-domain corpus is a fast way to train a compact, confident, wrong model.

What I'd take from it, beyond A6B: the failure modes generalize. Renormalized top-k is not free to widen. SFT teaches style before substance. Freeze the router or watch the deltas lose their coordinate system. And measure in pairs on one machine, because a −0.5 pt result that is really cross-machine drift will fool you. Whether A6B ever nets out positive after rejection-FT and GRPO is unsettled — the repo says so plainly — but the instrumentation is the part I'd copy tomorrow.


Built on the a6b-k-expansion research log (hikarioyama, 2026) — README, METHOD.md, and the two HTML reports under docs/, MIT-licensed. All numbers are the author's own paired, McNemar-tested measurements on Qwen3.6-35B-A3B; I report them as stated and have not independently reproduced them. The interactive diagrams are my illustration of the mechanism, not measured traces; the architecture figure is reproduced from the repo for commentary.

Cite this article

For attribution, please use the following reference or BibTeX:

Satyajit Ghana, "A6B: k-expansion, or what breaks when you force a top-8 MoE router to fire 32 experts", ai.thesatyajit.com, July 2026.

bibtex
@misc{ghana2026a6bkexpansion,
  author = {Satyajit Ghana},
  title  = {A6B: k-expansion, or what breaks when you force a top-8 MoE router to fire 32 experts},
  url    = {https://ai.thesatyajit.com/articles/a6b-k-expansion},
  year   = {2026}
}
share