Per-prompt FFN pruning activates 3B of a 9B LLM, matching the dense 9B with up to 1.8x speedup

2026-08-04

Apple's IFPruning uses a small predictor to select an FFN subnetwork per prompt; a 9B LLM at 3B active params nearly matches the dense 9B, with 57% lower time-to-first-token.

What problem this solves

Structured pruning is the standard way to shrink a large model: learn a binary mask, drop entire rows and columns of the feed-forward layers, and ship a single fixed smaller model. The mask is static, so the same pruned model has to handle coding, math, and domain questions at once, even though those tasks rely on different parameters and knowledge. A one-size-fits-all pruned model trades accuracy for speed, or speed for accuracy.

The real question is whether a model can pick, from its own parameters, the subset best suited to the current instruction. This ICML 2025 paper from Apple's AI/ML team says yes, and the predictor it builds to do the picking is both small and accurate.

Method

At the center is a sparsity predictor, far smaller than the LLM (about 300M parameters). It reads the user's instruction, scores every feed-forward layer's rows and columns, and a SoftTopK operator turns those scores into a differentiable mask that keeps exactly the target number of dimensions. The mask drops the i-th column of the first FFN matrix and the i-th row of the second, pruning a 9B model down to 3B active parameters.

Three design choices matter:

• Parameters are chosen once, then frozen for the whole decode. This differs from mixture-of-experts and contextual sparsity, which reload a fresh set of weights at every generated token and pay a high weight-loading cost. IFPruning fixes the subnetwork before decoding starts and reuses it for the entire response, caching it as a dense model. It gets static-pruning speed with task-tailored accuracy. The authors frame this for on-device use (phones, laptops), where each query typically samples only a few responses and the activated parameters can be reused.
• Two-stage training. Continued pre-training splits the input into equal chunks and uses each chunk's predicted mask to compute the next-token loss for the following chunk, co-optimizing predictor and masked LLM. Then comes supervised fine-tuning. An ablation shows that skipping continued pre-training hurts every benchmark.
• The predictor is a 302M LLM plus a two-layer MLP mask head. Only FFN layers are pruned in this work, though the method extends to attention heads.

Results

Three source models (6B, 9B, 12B) are all pruned to 3B active and compared against three baselines: a dense 3B trained with roughly twice the pretraining tokens, a 3B pruned then distilled from a 12B teacher, and an unpruned dense 9B as an upper bound.

At the same 3B active budget, IFPruning 9B to 3B beats dense 3B by 8.7 points on coding, 5.3 on MATH, and 4.3 on MMLU. Against the full 9B it is close to flat: HumanEval 43.3 versus 46.5, MATH 37.1 versus 37.3. Larger source models (6B to 12B) keep helping, most clearly on math and coding. The predictor also routes an unseen instruction dataset, GPTeacher, to a sensible subnetwork.

It can also select per task rather than per input: one task instruction fixes a subnetwork shared across all of that task's instances, with no extra fine-tuning, and the gap versus per-input selection is usually under 1%.

The selection is interpretable. Low layers activate near-identical subnetworks across inputs; higher layers split by domain. The MMLU computer-science subset shares parameters with Code-Alpaca; math, physics, and GSM8K cluster together; history, law, and general instructions form separate groups.

On-device latency (batch 1, four generations per query): pruning 9B to 3B cuts time-to-first-token by up to 57% and decode time by up to 41%, a 1.8x overall speedup. Selection and loading together take under 0.1s, 1-2% of generation time. Runtime matches static-pruned 3B while adding input-specific adaptivity.

Why it matters

On-device inference is the target. Phones and laptops are tight on memory and compute, and cannot run a full 9B. IFPruning stores a 9B weight bank but activates only 3B per instruction, running at 3B dense speed while approaching 9B accuracy. Because MoE gets more expensive at small batch sizes from weight loading, the pick-once-run-frozen design wins for low-batch on-device serving.

A side finding is worth stating plainly: a large model carries many parameters irrelevant to any single task, and a 300M predictor is enough to find the ones that matter.

Limitations

The authors list a few: only FFN layers are pruned so far; end-to-end training may underuse the data, and a contrastive loss could push similar inputs toward similar subnetworks; server-side batched serving mixes tasks in one batch with different subnetworks, which breaks batching.

Two more gaps stand out. All latency numbers are GPU simulations of on-device use (batch 1); there is no real phone or laptop measurement. And while capability tasks like coding and math hold up, AlpacaEval-style open instruction following gains little, so instruction-conditioned selection is not equally effective across task types.

Terms

• Structured pruning: dropping entire rows and columns of weights (not scattered individual weights), yielding a smaller dense model that runs faster on real hardware.
• FFN: the feed-forward network inside a transformer, usually the bulk of the parameters and the main pruning target.
• SoftTopK: a differentiable keep-the-top-k operator that turns continuous scores into a mask preserving exactly k dimensions, trainable by gradient descent.
• Contextual sparsity: predicting a sparse subnetwork at every decode step; accurate but pays weight-loading cost each step.
• TTFT (time-to-first-token): time from feeding input to emitting the first token, a key on-device metric.

Source

• Paper

What people are saying

• AFM3 20B paper proposes prompt-level pruning that activates only 20% of MLP layers Aaaaaaaaaeeeee

All paper explainers