Pruning Schedule

The core of a schedule is the sparsity ramp function, which defines how the target sparsity increases from an initial value (often 0%) to a final target over the pruning period. Common functions include:

• Linear Ramp: A simple, constant increase in sparsity per step.
• Cubic Ramp: A slower initial increase that accelerates, often preserving early training stability.
• Exponential Ramp: Rapid early pruning followed by a slow approach to the target. The choice of function balances aggressive compression with the network's ability to adapt.

This defines when pruning updates occur. Key strategies are:

• One-Shot Pruning: Apply the entire final sparsity target in a single step, typically post-training. This is simple but often causes severe accuracy loss.
• Iterative Pruning: The schedule alternates between pruning steps (removing parameters) and recovery phases (retraining the sparse network). This is the standard for Gradual Magnitude Pruning (GMP).
• Continuous Pruning: Sparsity is updated at every training step or at a very high frequency, allowing for seamless adaptation. This is common in sparse training paradigms.

The schedule must be coupled with a pruning criterion that determines which parameters to remove at each step. The schedule dictates the amount to prune, while the criterion provides the ranking. Common integrations include:

• Magnitude-Based Schedules: Use weight magnitude (L1 norm) as the criterion. The schedule defines the percentile threshold that increases over time.
• Gradient-Based Schedules: Used in Movement Pruning, where the schedule controls the pruning of weights based on how their value changes during fine-tuning.
• Regularization Schedules: In pruning-aware training, an L1 or L0 regularizer's strength is scheduled to gradually induce sparsity.

The total duration of the pruning phase is critical. It is often defined as a percentage of the total training steps or epochs. A key component is the pruning warm-up period, where no pruning occurs for a set number of steps at the start of training or fine-tuning. This allows the model to learn stable representations before compression begins, which is essential for maintaining accuracy. For example, a schedule might specify: Train for 10% of total steps (warm-up), then linearly increase sparsity from 0% to 80% over the next 50% of steps, followed by final fine-tuning.

Advanced schedules incorporate rewinding, a technique where, after a pruning step, the network's weights are reset to values from an earlier checkpoint in training (but not to initial randomness) before retraining continues. The schedule defines the rewind point. After the target sparsity is reached, a final fine-tuning phase is almost always required. The schedule may define a separate learning rate, duration, and data regimen for this phase to recover maximum accuracy from the sparse architecture.

For production deployment, schedules can be designed to produce hardware-efficient sparsity patterns. Instead of a simple global sparsity target, the schedule may enforce:

• Per-Layer Sparsity Targets: Different layers are pruned at different rates based on their sensitivity or contribution to FLOPs/memory reduction.
• Structured Sparsity Targets: The schedule aims for patterns like N:M sparsity (e.g., 2:4), where in every block of 4 weights, 2 are zero. This pattern can be executed efficiently on modern GPU sparse tensor cores.
• Block-Wise Sparsity: Increases sparsity in coarse-grained blocks to align with memory fetch sizes on the target accelerator.

Within framework APIs, a pruning schedule is defined by its type, intensity progression, and timing. Key configurable hyperparameters include:

• Schedule Type: PolynomialDecay, ConstantSparsity, ExponentialDecay.
• Initial Sparsity: The starting sparsity level (often 0%).
• Final Sparsity: The target sparsity at the end of the schedule (e.g., 80%, 90%).
• Begin Step/Epoch: The training step or epoch at which pruning starts.
• End Step/Epoch: The point at which the final sparsity is reached and pruning stops.
• Update Frequency (pruning_frequency): How often (in steps or epochs) the pruning mask is recalculated and applied. A frequency of 100 steps is common.
• Power (for Polynomial): The exponent controlling the decay shape; a value of 3 (cubic) is standard for GMP.

Frameworks handle the integration of the pruning schedule into the training loop differently, but the core mechanics are similar:

1. Mask Application: Before the forward pass, a binary mask (defined by the current schedule step) is element-wise multiplied with the weights, zeroing out pruned parameters.
2. Forward-Backward Pass: Training occurs normally. The gradients for the masked (zero) weights are typically also zero.
3. Optimizer Step: The optimizer updates all weights, including those currently masked.
4. Mask Update: At the scheduled frequency, the pruning criterion (e.g., magnitude) is re-evaluated on the updated weights. A new mask is generated based on the current target sparsity from the schedule, potentially pruning different weights.
5. Permanent Pruning: After training, a final strip() function (TFMOT) or remove() (PyTorch) is called to permanently remove the pruned weights and the mask infrastructure, creating a final, smaller model file.

The pruning rate is the percentage or fraction of a model's parameters targeted for removal during a single pruning operation. It is the primary control knob for the final sparsity level of the network.

• A schedule like Gradual Magnitude Pruning dynamically increases this rate from 0% to a target (e.g., 90%) over many training iterations.
• Setting the rate too high too quickly can cause irreversible accuracy loss, as the network cannot adapt to the drastic structural change.

A pruning criterion is the heuristic or metric used to score parameters for removal. The schedule determines when to apply this criterion.

• Magnitude-Based: Removes weights with the smallest absolute value (e.g., |weight| < threshold).
• Gradient-Based: Methods like Movement Pruning remove weights based on how much their value changes during fine-tuning.
• Second-Order: Algorithms like Optimal Brain Damage (OBD) use Hessian approximations to estimate parameter saliency. The criterion generates the scores; the schedule uses these scores to decide the timing and extent of removal.

Pruning granularity defines the structural unit of removal, which a schedule must be designed to accommodate. It ranges from fine to coarse:

• Unstructured/Fine-Grained: Individual weights are pruned anywhere in a tensor.
• Structured/Coarse-Grained: Entire structural components are removed (e.g., neuron pruning, channel pruning, filter pruning).
• Block-Wise: Patterns like N:M sparsity (e.g., 2:4), where 2 of every 4 consecutive weights are non-zero. A schedule for structured pruning must account for the impact on layer dimensions and may require different per-layer policies.

Iterative Magnitude Pruning (IMP) is a canonical algorithm that embodies a specific, cyclical pruning schedule. It does not gradually increase sparsity within a single training run. Instead, it follows a strict, repeated cycle:

1. Train a dense network to convergence.
2. Prune a fixed percentage (e.g., 20%) of the smallest-magnitude weights.
3. Reset the remaining weights to their values from an earlier training checkpoint (rewinding).
4. Retrain the pruned network from the rewound state. Steps 2-4 are repeated until the target sparsity is reached. This schedule is central to finding Lottery Ticket Hypothesis subnetworks.

Sparse training is an alternative paradigm that eliminates the need for a traditional pruning schedule. Instead of training dense then pruning, the network is initialized with a fixed, sparse connectivity pattern and trained with it from the start.

• The sparsity pattern (e.g., Erdős–Rényi random graph) is defined a priori.
• There is no gradual increase of sparsity during training; the sparsity level is static.
• This approach avoids the costly dense pre-training phase and can be more memory-efficient throughout, but finding optimal initial sparse architectures remains a research challenge.

Hardware-aware pruning is an optimization objective that must inform the design of the pruning schedule and the chosen sparsity pattern. The goal is to maximize inference efficiency on target silicon.

• Schedules may prune to specific structured patterns (like N:M) that map efficiently to hardware sparse tensor cores in modern GPUs.
• Unstructured sparsity, while offering high theoretical compression, often requires specialized libraries to achieve speedups on general-purpose hardware.
• The schedule must balance the pursuit of high sparsity with the enforcement of hardware-friendly patterns to ensure actual latency and energy gains in deployment.