Movement Pruning and Dynamic Sparsity

Movement pruning and dynamic sparsity techniques allow the set of pruned weights to evolve during the training or fine-tuning process. This approach helps overcome limitations of static pruning methods, which often assume that weights deemed unimportant early in the process remain unimportant throughout. For example, a common static method, magnitude-based pruning, removes weights with the smallest absolute values. Such static views can be suboptimal because the complex dynamics of training large models might allow initially small weights to gain significance later on.

The Intuition Behind Movement Pruning

Imagine the optimization process as navigating a complex loss surface. Standard magnitude pruning identifies weights close to zero at specific checkpoints and removes them permanently. Movement pruning, however, acknowledges that the gradients flowing through the network might "push" some pruned weights away from zero, suggesting they could become useful again. Conversely, some weights that were initially large might drift towards zero and become candidates for pruning.

The core idea, introduced by Sanh et al. (2020) in the context of BERT, is to allow weights to move into and out of the pruned set during training. This is typically achieved by maintaining a binary mask $M$ for the weights $W$. Unlike static pruning where $M$ is fixed after a pruning step, movement pruning updates $M$ periodically.

Algorithmic Mechanism

Movement pruning often integrates pruning directly into the optimization loop. A common approach involves these steps:

1. Initialization: Start with a dense model or a model pruned using a static method. Define a target sparsity level $s$.
2. Sparsity Application: Apply a sparsity mask $M$ based on weight magnitudes. Only weights $W_{ij}$ where $M_{ij}=1$ are active.
3. Gradient Calculation: Compute gradients $\nabla L$ with respect to the masked weights $W \odot M$ (where $\odot$ is element-wise multiplication).
4. Gradient Accumulation (Important Step): Instead of only updating the active weights, accumulate gradients for all weights, including those currently masked ($M_{ij}=0$). This is often done using the standard optimizer update rule but applying it to the underlying weight tensor before re-applying the mask. For an optimizer like SGD with momentum, the update might look like: $$V_{t+1} = \beta V_t + \nabla L(W_t \odot M_t)$$ $$W_{t+1}^* = W_t - \alpha V_{t+1} \quad \text{// Update applied to all weights}$$
5. Mask Update: Periodically (e.g., every few hundred steps), re-evaluate the mask $M$.
   • Identify the top $(1-s) \times N$ weights in $W_{t+1}^*$ by magnitude (where $N$ is the total number of weights). Set their corresponding entries in $M_{t+1}$ to 1.
   • Set the remaining $s \times N$ entries in $M_{t+1}$ to 0. This allows weights that were previously masked ($M_{t,ij}=0$) but have accumulated significant gradient updates (leading to a larger magnitude in $W_{t+1}^*$) to become unmasked ($M_{t+1,ij}=1$). Conversely, weights that were active but decreased in magnitude can become masked.
6. Weight Update: Apply the new mask to the updated weights: $W_{t+1} = W_{t+1}^* \odot M_{t+1}$.
7. Repeat: Continue the training process from step 3.

This process allows the sparsity pattern to adapt dynamically. Weights are not just pruned; they have a chance to be "revived" if the optimization process deems them important later. Gradual pruning schedules, where the target sparsity $s$ increases over time, are often used in conjunction with movement pruning to stabilize training.

Flow of movement pruning within a training loop. Gradients influence all weights (including masked ones), potentially changing their magnitudes enough to alter the mask in the next regeneration step.

Dynamic Sparsity in Movement Pruning

Movement pruning is a specific instance of dynamic sparsity. The broader concept encompasses any technique where the sparsity pattern is not fixed upfront but can change during training or even adapt at inference time.

• RigL (Rigged Lottery): Proposed by Evci et al. (2020), RigL combines elements of magnitude pruning and random pruning within the training loop. It periodically removes weights by magnitude but also randomly adds new connections, using gradient information to guide which connections are activated. This maintains computational cost while allowing exploration of different sparse structures.
• Dynamic Sparse Training: Some research looks at dynamically adjusting the sparsity mask based on criteria other than just magnitude, potentially incorporating sensitivity analysis or Hessian information, although these methods are often more computationally intensive.
• Input-Dependent Sparsity: A more advanced concept involves sparsity patterns that change based on the input data. For example, in MoE (Mixture-of-Experts) models, only certain "experts" (sub-networks) are activated for a given input token, creating a dynamic, input-dependent form of structured sparsity. While not typically called "pruning," it shares the goal of reducing computation by selectively using parts of the model.

Advantages and Trade-offs

Advantages:

• Improved Accuracy: Dynamic methods, particularly movement pruning, often achieve higher accuracy than static pruning methods for the same level of sparsity. They are better at identifying and preserving truly critical weights that might have been erroneously pruned by static approaches.
• Better Sparse Solutions: By allowing the sparse structure to co-adapt with the weights during training, these methods can discover more effective sparse subnetworks compared to pruning a pre-trained dense model.

Challenges:

• Training Complexity: Implementing and managing dynamic masks adds overhead to the training process. Keeping track of gradients for masked weights and periodically regenerating the mask requires careful implementation.
• Hyperparameter Tuning: Dynamic sparsity introduces new hyperparameters, such as the mask update frequency, the sparsity schedule (how quickly sparsity increases), and potentially parameters related to gradient accumulation or random connection growth (as in RigL). Finding optimal settings can require significant experimentation.
• Potential Instability: The changing sparsity pattern can sometimes lead to instability during training if not managed carefully, especially when combined with high learning rates or aggressive sparsity schedules.
• Hardware Efficiency: While dynamic unstructured sparsity can yield very sparse models, realizing theoretical speedups on hardware can be challenging. Static, structured sparsity patterns are generally easier for current hardware and compilers to accelerate efficiently. However, runtime support for dynamic sparsity is an active area of development.

When to Consider Dynamic Sparsity

Movement pruning and dynamic sparsity are powerful techniques when achieving very high sparsity levels ($> 80\%$) with minimal accuracy loss is the primary goal. They are particularly relevant in research settings or when fine-tuning models where computational resources for training are available, and the focus is on maximizing the performance of the final sparse model. While static structured pruning might be preferred for immediate deployment due to better hardware compatibility, dynamic methods represent a more advanced approach to finding highly optimized sparse LLMs. Understanding these techniques provides insight into the evolving relationship between model parameters, optimization dynamics, and network architecture in the context of efficiency.