Apple’s IDEA Prune cuts 2.8B models to 1.3B and lifts MMLU to 46.4%

Apple has published IDEA Prune, a pipeline that pretrains a language model bigger than the one it intends to ship, then cuts it back down to size. The paper compresses 2.8B parameter models to 1.3B using up to 2T pretraining tokens. Apple posted it to its machine learning research site this August, while the arXiv version is dated 7 March 2025.

The mechanism is a single cosine annealing learning rate schedule. A naive pipeline runs enlarged pretraining, pruning and recovery as three separate stages, each with its own decay and warm up. Every restart pushes the learning rate back up, which the authors say triggers a sharp rise in the loss curve and forgetting. IDEA Prune keeps one decay across all three stages and strips feed forward neurons out gradually rather than in one shot.

That design matters because the enlarged model is thrown away afterwards. Apple built the 2.8B version by widening only the feed forward layers, taking the 1.3B model’s feed forward hidden dimension of 6528 up to 2048 by 8. The authors note that on device deployments typically restrict models to 3 billion parameters, so the artefact you ship is small even when the artefact you trained isn’t. That constraint is the same one that shapes what fits on your own hardware.

Here’s Table 1 from the paper. Every row is a 1.3B model trained on DCLM, which the paper describes as 4T unique tokens. Perplexity is measured on OpenWebText and MMLU accuracy is 5 shot.

Method	OpenWebText perplexity	MMLU
1.3B, 1T tokens, from scratch	9.10	28.9%
1.3B, 2T tokens, from scratch	8.95	45.7%
One shot random pruning	8.98	32.5%
Minitron pruning, naive pipeline	8.97	31.4%
Sheared LLaMA pruning, naive pipeline	8.96	33.4%
IDEA Prune	8.88	46.4%

The MMLU column carries the result. IDEA Prune reaches 46.4% against 31.4% to 33.4% for the three naive pipelines, and its OpenWebText perplexity of 8.88 beats the best baseline at 8.96. Measured against the 1.3B model trained from scratch on the same 2T tokens, though, the margin shrinks to 8.88 versus 8.95 and 46.4% versus 45.7%. The paper doesn’t hide that.

This implies that the enlarge-and-prune pipeline does not always increase token efficiency compared to training target-size models from scratch given the same training tokens, highlighting the need for careful pipeline selection.

Li et al., IDEA Prune, via arXiv

So the honest reading is narrower than “pruning beats training from scratch”. Pruning is cheap when somebody else already paid for the large model, which is the case the earlier work was built for. Sheared LLaMA reports LLaMA2-7B cut to 1.3B and 2.7B at 3% of the compute of training from scratch. Minitron reports 8B and 4B models derived from a pretrained 15B model with up to 40x fewer training tokens each.

Neither figure transfers directly to Apple’s table, because Apple reimplemented both as pruning methods inside its own pipeline. It ran Minitron without the distillation and Sheared LLaMA without the dynamic batch loading that each original method uses. Count the enlarged model’s own training bill and the arithmetic changes, which is the whole argument of the paper and a familiar problem to anyone costing a frontier training run.

Adding distillation didn’t reorder the results. Using a 2.8B model trained on 2T tokens as the teacher, pruning with distillation reached 46.6% on MMLU at 8.862 perplexity. A 1.3B student trained from scratch on 1T tokens with distillation managed 44.5% and 8.949.

Neither Apple’s page nor the arXiv listing points to weights or code, so nobody outside Apple can rerun any of this. Two ablation findings are worth watching even so. A 300M target scored better when pruned from a 600M model than from a 1B one, and the paper puts the best feed forward width enlargement factor at 2.6x. That’s a narrow band, and it’s the number a team with a fixed token budget would test first.