Pruning Basics › Pruning Schedules

Michael BrenndoerferFebruary 27, 20267 min read
Language AI HandbookAll chapters
Part 41 of 67Model Compression
Chapter 309Pruning Basics › Pruning Schedules
Reading section 5 of 11 · Pruning Schedules · 7 minSection 5 of 11 · 7 min

Book contents

Section 5 of 11 · 7 min read

Pruning Basics › Pruning SchedulesOne section · Sections in this chapter5 / 11 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

How and when you remove weights matters as much as which weights you remove. The pruning schedule controls the trajectory from a dense, fully trained model to a sparse, pruned model. Different schedules make different tradeoffs between compute cost, final accuracy, and implementation complexity.

One-Shot Pruning

The simplest schedule is one-shot pruning: train the model to convergence, prune to the target sparsity in a single step, then fine-tune on the training data to recover accuracy. This is the most conceptually straightforward approach and is a useful baseline for evaluating other schedules.

One-shot pruning works reasonably well at moderate sparsity levels, up to around 50 to 60 percent, but becomes increasingly problematic at high sparsity. When a large fraction of weights are removed simultaneously, the model's ability to recover through fine-tuning is limited. The network can no longer represent many of the functions it previously encoded, and fine-tuning on the remaining sparse network may not be sufficient to compensate.

Think of it like removing pieces from a jigsaw puzzle and asking someone to complete it with the remaining pieces. If you remove 20% of the pieces, completion may still be feasible; the remaining pieces constrain the solution well. If you remove 80% of the pieces, the remaining pieces are insufficient to reconstruct the original image, and the best you can do with fine-tuning is to fill in the gaps as creatively as possible.

One-shot pruning shines when computational cost is the primary constraint. A single fine-tuning run is inexpensive relative to the full original training. If moderate sparsity is sufficient for the deployment target, one-shot pruning achieves a good accuracy-efficiency tradeoff with minimal additional compute.

Iterative Pruning

Iterative pruning alternates between pruning and fine-tuning across multiple cycles. A typical schedule looks like:

  1. Train the model to convergence on the full dataset.
  2. Prune a fraction of the target sparsity (for example, prune 10% of weights to reach 10% sparsity).
  3. Fine-tune the pruned model for a few thousand steps.
  4. Repeat steps 2 and 3 until the target sparsity is reached.

Iterative pruning consistently outperforms one-shot pruning at the same final sparsity level. The reason is that each fine-tuning phase allows the remaining weights to adjust and compensate for what was removed before the next round of pruning. The model adapts incrementally rather than absorbing a single catastrophic reduction. At each step, the model is at a working, recoverable state before the next pruning is applied.

The accuracy advantage of iterative over one-shot pruning grows with the target sparsity. At 50% sparsity, iterative pruning may achieve one or two additional percentage points of accuracy. At 90% sparsity, the gap can be much larger, with one-shot pruning sometimes failing to recover meaningful accuracy while iterative pruning maintains reasonable performance.

The main disadvantage of iterative pruning is cost. Each prune-then-finetune cycle requires additional compute. For very large models, even a single fine-tuning run can take days on a cluster of GPUs, and iterative pruning might require five to twenty such cycles. The compute budget for pruning a large language model via iterative pruning is substantial, sometimes comparable to the original pretraining cost of a smaller model. This is one reason why the field has invested heavily in one-shot second-order methods like SparseGPT: they achieve high-quality pruning without the repeated fine-tuning cycles.

Gradual Pruning

Gradual pruning (also called polynomial decay pruning) provides a continuous schedule that increases sparsity smoothly from zero to the target over the course of a single training run. Rather than alternating discrete prune-and-finetune phases, gradual pruning applies small incremental pruning steps at regular intervals during training, so pruning and learning happen simultaneously.

A common schedule increases the target sparsity s(t)s(t) at training step tt according to:

s(t)=sf(1−(1−t−t0n⋅Δt)3)s(t) = s_f \left(1 - \left(1 - \frac{t - t_0}{n \cdot \Delta t}\right)^3\right)

where:

  • sfs_f: the final target sparsity (for example, 0.9 for 90% sparsity)
  • t0t_0: the training step at which pruning begins (often after an initial warmup period)
  • nn: the total number of pruning steps planned
  • Δt\Delta t: the interval between consecutive pruning steps
  • tt: the current training step

The cubic polynomial factor (1−(1−⋅)3)\left(1 - (1 - \cdot)^3\right) means sparsity increases rapidly at first and then slows as it approaches the target. This front-loading of pruning gives the model more fine-tuning time at high sparsity levels, where recovery is most needed. The logic is that the first pruning steps remove the clearest redundancies, which the model can compensate for quickly, while the later pruning steps at high sparsity are more damaging and require more recovery time.

Gradual pruning integrates cleanly with standard training loops. You do not need separate training and fine-tuning phases; instead, you insert a pruning step every Δt\Delta t gradient updates. The model never experiences a large discontinuous change in its parameter values, which keeps optimization stable. It is the default in many production pruning pipelines, particularly those based on the TensorFlow Model Optimization Toolkit and the Hugging Face optimum library.

One important detail of gradual pruning is the choice of t0t_0, the step at which pruning begins. Starting pruning from the very first step is generally suboptimal, because the model has not yet learned a useful representation, and the importance scores computed at initialization are not meaningful. Starting pruning after a brief warmup period, typically 10 to 20% of total training, gives the model time to develop stable weight magnitudes before the pruning criterion is applied.

Rewinding and Lottery Ticket Pruning

Motivated by the Lottery Ticket Hypothesis, a more elaborate schedule involves weight rewinding. After pruning the trained model to identify which weights to keep, instead of continuing from the current weights, the surviving weights are reset to their values from an early checkpoint (for example, after only 1,000 to 10,000 training steps). Only the pruning mask, indicating which weights are zero, is retained from the trained model. The model is then retrained from this early checkpoint with the pruning mask in place.

The intuition is that the winning ticket subnetwork trains better when started from its original early-training weights rather than the fully trained weights accumulated over the entire training run. Full training can cause surviving weights to partially compensate for the weights that will eventually be pruned, encoding information in ways that are hard to disentangle. By rewinding to an early checkpoint and retraining, the subnetwork gets a cleaner start: it develops its representations from scratch, knowing from the beginning which connections are available to it.

In practice, rewinding to step 0 (true initialization rewinding) works best on small networks and fails on large ones. Rewinding to a small positive step, such as after 1,000 steps (around 1 to 5% of total training), works reliably on deep networks including transformers. The key insight is that a few training steps are enough to find a starting point that is significantly better than pure initialization, while still being "early enough" that the subnetwork can develop useful representations from that starting point.

Weight rewinding is more expensive than standard one-shot pruning: after identifying the pruning mask, you must retrain from the early checkpoint, which is nearly as expensive as the original training run. The payoff is that rewound lottery ticket subnetworks sometimes generalize better than subnetworks fine-tuned from the pruned endpoint, particularly when the fine-tuning data distribution differs from the pretraining distribution.