KikoCis's picture
Protect recipe: remove heal method (self-distill source, layer-targeting, val-loss); keep results + held-out lesson
1f73c5c verified
|
Raw
History Blame Contribute Delete
6.07 kB

How Small Can a 230B Agent Get? Pruning a MoE Coding Model β€” the Scaling Curve and the Validity Trap

A field report on compressing a 230B-parameter Mixture-of-Experts coding agent by pruning experts, keeping it a working agent in a real tool-using loop. What we found: a clean size↔capability curve, and a memorization trap that only held-out evaluation caught.


TL;DR

We took MiniMax-M2 (230B total, ~10B active, 256 experts/layer, top-8) and tried to make it small enough to run as a local coding agent without it falling apart in a real tool loop. The honest results:

  • 96.7% of the weights are experts. Quantization alone can't halve the model; you must remove experts.
  • Pruning gives a clean curve, no cliff. Keep 192/256 experts β†’ solves 3/6 of an SWE probe set; keep 160 β†’ 1/6; the full quant β†’ 6/6. Capability degrades smoothly with the expert budget β€” and this is untrained (pure pruned quant, no memorization).
  • At a fixed size, prune beats uniform-extreme-quant. 160 experts at ~2.7-bit (46 GB) resolves where all-256-experts at 1.8-bit (48 GB) resolves nothing. Concentrating the bit budget on fewer experts wins.
  • A recovery fine-tune can push higher β€” but only held-out proves it. The gains must be validated on tasks held out of the recovery data, not the ones inside it.
  • Held-out evaluation is non-negotiable. Our most exciting recovery result was memorization β€” the test task had leaked into the recovery data; only hold-out caught it.

The experts are the model

A per-tensor accounting shows the experts are ~97% of the weights; the entire backbone (all attention, routers, embeddings, norms) is ~3%. So halving the footprint requires removing experts, not just shrinking them. Below a certain bit-width the model needs an importance matrix and starts to fall apart anyway.

Pruning by measured activation

Remove experts at random and the model dies. We ranked experts by measured activation over a code/agentic calibration set (an importance-matrix pass recording, per layer, how often each expert fires) and kept the most-used per layer. The surgery is clean: each expert occupies whole contiguous blocks, so keeping a subset is a slice + a rewrite of the expert-count metadata and the router β€” survivors keep their original precision, no re-quantization.

A property fell out of the map: experts are used broadly (almost none are dead), so there's no free lunch β€” every dropped expert costs retained activation mass.

The curve: capability scales smoothly with the expert budget

Evaluated as a real agent β€” a harness driving the model through a Dockerized repo (read, edit, run tests) on the exact SWE instances the un-pruned model resolves 6/6:

Experts kept Size SWE resolved
256 (full) 78 GB 6 / 6
192 ~55 GB 3 / 6
160 46 GB 1 / 6

No cliff. Drop a quarter of the experts and you keep half the solving; drop more and it tapers, but the model stays a coherent, acting agent the whole way down β€” not a babbling wreck. All of this is untrained: pure pruned quant. The resolves are genuine compressed-M2 capability.

Two levers mattered together: more experts, and giving the agent its full context window back (an aggressive context budget we'd inherited from an earlier serving constraint was needlessly starving it).

At a fixed size: prune beats quantize-everything

A fair challenge β€” why prune at all, why not just quantize every expert lower and keep them all? We tested it. At the same size as the 160-expert prune, a uniform extreme-low-bit quant that keeps all experts stays coherent (restates each task) but is a shallow agent: a few tool calls, almost no edits, resolves none. The 160-expert prune, same size, sustains the loop, edits, and resolves. Spreading too few bits across all experts preserves knowledge but destroys execution precision; concentrating the budget on fewer experts keeps execution.

Pushing past the curve: recovery is possible β€” but prove it held-out

Can a recovery fine-tune lift the pruned model further? Yes β€” a recovery pass can restore the decisive editing behavior that pruning alone leaves tentative (a pruned model that explores and understands a bug but won't commit an edit can be brought back to acting). But the size of that gain is exactly what a held-out test exists to check, and where the recovery is applied matters more than how much of it you run. We report the result, not the recipe β€” and the result only counts held-out.

The validity trap

With recovery applied, the model resolved a real task. Exciting β€” and wrong to trust. The test task had leaked into the recovery data, so the "resolve" was memorization: with the test tasks held out, that resolve vanished. Only hold-out caught it. The honest, held-out signal is that agentic behavior (persistence, navigation, editing, format) recovers and generalizes, while the last mile of solving is harder to recover and must be proven held-out β€” never on data the model could have seen.

Reusable takeaways

  1. For a MoE, the experts are the model β€” shrink structurally by pruning experts; quantization alone won't get you there.
  2. Prune by measured activation, per layer β€” random/uniform pruning throws away load-bearing experts.
  3. Capability scales smoothly with the expert budget β€” there's a curve, not a cliff; pick your size/capability point.
  4. At a fixed budget, prune > spread-the-bits-everywhere β€” concentrate precision on fewer experts.
  5. A recovery fine-tune can help β€” but validate it held-out, and don't trust a tool's defaults for where the recovery lands.
  6. Validate held-out, or don't believe it β€” recovery data can smuggle the test answer into training.

A compressed-but-capable agent is reachable β€” as a point on a smooth curve, not a magic small model, and only if you refuse to trust a number you didn't hold out.