mlx-optiq
Workflow · Expert pruning

Expert pruning for MoE quants

A Mixture-of-Experts model stores far more parameters than it uses. Qwen3.6-35B-A3B keeps 256 experts in every layer and routes each token to eight of them, so the expert bank is most of what sits on disk and almost none of what runs. optiq prune-experts removes the experts that contribute least and slices the router to match. Active parameters per token do not change — top-k routing is preserved — so the model gets smaller without getting slower.

prune.shbash
$ optiq prune-experts mlx-community/Qwen3.6-35B-A3B-OptiQ-4bit \
    --target-retention 0.5 --output ./reap

It implements REAP (Cerebras Research, ICLR 2026) in the quantized domain. There is no BF16 parent and nothing is dequantized: every retained expert is copied out of the parent quant bit-for-bit, so the pruned model inherits the parent’s mixed-precision bit map exactly.

What it costs

Published OptiQ variants all target 50 % retention, the same way every quant targets 4-bit: one number across the lineup, chosen because that is where the saving is worth having rather than because it is free.

Factual knowledge goes first. On Qwen3.6-35B-A3B-OptiQ-4bit at 50 %, MMLU falls 21.4 points while procedural ability barely moves — GSM8K +2.6, IFEval +4.3, tool calling −1.0, HumanEval −1.3 — for a Capability Score of 80.03 to 76.57 and 22.1 GB down to 13.9 GB. Gemma-4 pays more at the same fraction: 75.76 to 68.13, with HumanEval dropping too. Check MMLU against your own workload before relying on a pruned variant.

Which models are worth pruning

Not all of them, and the expert count is a poor guide. Laguna-XS routes over 256 experts, exactly as many as Qwen3.6, and loses far more at the same fraction. The command measures the damage directly instead of guessing: before writing anything it reports the KL divergence between the pruned and unpruned model, and warns when that number crosses 1.0.

ModelExpertsKL at 50 %Result
Qwen3.5-35B-A3B2560.07Published
Qwen3.6-35B-A3B2560.13Published
Gemma-4-26B-A4B1280.56Published, larger loss
Hy31920.39Rejected — arithmetic regressed
LFM2.5-8B-A1B321.25Rejected
Laguna-XS-2.12561.90Rejected

Every checkpoint measured above 1.0 wrote fluent, confident, factually wrong text while still handling code and arithmetic correctly — the kind of failure that survives a quick read. Hy3 is the instructive one: it passed the KL threshold and still regressed, which is why the published variants are also checked by running the unpruned parent on the same prompts. If a model lands above the threshold, keep more experts; --retain takes any count.

How experts are ranked

Experts are scored by router weight times expert output norm, measured on calibration data. How that score is aggregated depends on the router, and it is decided per checkpoint rather than assumed: the command scores both candidate rules against the model’s own unpruned output and keeps the winner, overriding the family default only when it wins by a clear margin.

Where the router softmaxes over every expert, the routing weight expresses a genuine preference, and the score is a mean over the times each expert fired — a rarely-chosen but strong expert outranks a common weak one. Where the router renormalizes over the selected top-k, every chosen expert gets roughly 1/k whatever the router thought of it, and the preference survives only in how often it was chosen; there the score is weighted by selection frequency instead. On Gemma-4 that distinction is worth a factor of three in divergence.

Models larger than memory

Slicing streams one tensor at a time, so its peak is a single expert tensor rather than the model — disk, not RAM, is the limit. Profiling needs a forward pass, and there OptiQ streams expert weights off SSD when the checkpoint will not fit: Qwen3.5-122B-A10B, 43 GB on disk, profiles in about 6 GB of memory. A 122B model can be pruned on a 24 GB Mac.

What is not supported

Diffusion language models are refused rather than pruned. DiffusionGemma and LLaDA2 denoise a canvas instead of predicting a next token, so the autoregressive calibration this command profiles on measures the wrong thing entirely. Dense models have no routed experts to remove.

Pruned variants are published alongside their parents on the models page. Full flag reference in the CLI reference.