Source-linked AI summary
Movement Pruning: Adaptive Sparsity by Fine-Tuning
Victor Sanh, Thomas Wolf, Alexander M. Rush
TL;DR
Transfer learning limits magnitude pruning because pretrained weights remain close to their original values, reducing pruning adaptivity. The paper proposes deterministic movement pruning, which learns first-order importance scores during fine-tuning, and finds stronger high-sparsity performance, especially with distillation.
Problem
Magnitude pruning is less effective for pruning pretrained models during transfer learning because weights are largely predetermined before end-task fine-tuning.
Method
Movement pruning learns deterministic first-order importance scores during fine-tuning and uses them to select sparse weights, with hard and soft variants.
Results
Movement pruning consistently improves over existing methods in high-sparsity regimes; with distillation, soft movement pruning reaches 95% of BERT-base using about 5% of encoder weights.
Takeaways & Limitations
Movement-derived scores adapt pruning decisions to the fine-tuning objective, producing substantially better sparse pretrained models than absolute-value selection at high sparsity.
Takeaways & Limitations
Preliminary experiments found that increasing the number of pruning steps tended to improve end performance.
Abstract
from arXiv · showhide
Magnitude pruning is a widely used strategy for reducing model size in pure supervised learning; however, it is less effective in the transfer learning regime that has become standard for state-of-the-art natural language processing applications. We propose the use of movement pruning, a simple, deterministic first-order weight pruning method that is more adaptive to pretrained model fine-tuning. We give mathematical foundations to the method and compare it to existing zeroth- and first-order pruning methods. Experiments show that when pruning large pretrained language models, movement pruning shows significant improvements in high-sparsity regimes. When combined with distillation, the approach achieves minimal accuracy loss with down to only 3% of the model parameters.
1 Introduction
Transfer learning makes magnitude pruning less adaptive because pretrained weights remain largely determined before task fine-tuning. The paper introduces deterministic movement pruning and reports stronger results at high sparsity.
- Motivation: Large pretrained models are costly to ship, deploy, and train, motivating sparsity-based compression.Smaller models also reduce energy use on edge devices because parameter movement from storage dominates consumption.
- Existing pruning: Magnitude pruning preserves weights with high absolute values and is widely used across neural network architectures.It has been applied in computer vision, language processing, and lottery-ticket research.
- Problem: Magnitude pruning is less useful for transfer learning because pretrained weights are only fine-tuned rather than primarily determined by end-task data.This limits pruning decisions based on the fine-tuning step.
- Proposed method: Movement pruning selects weights according to changes during fine-tuning, allowing both low- and high-valued weights to be pruned when they shrink.The method shifts selection from zeroth-order weight values toward first-order movement information.
- Results: 95% of original BERT performance is reached with only 5% of encoder weights on MNLI and SQuAD v1.1.In highly sparse regimes with less than 15% of weights remaining, movement pruning improves over magnitude pruning and L0 regularization.
2 Related Work
Related work spans score-based subnetworks, higher-order pruning, and pretrained-language-model compression. Movement pruning differs by updating importance scores during fine-tuning while targeting individual weights without requiring an external teacher model.
- Score-based pruning: Parallel score matrices find sparse subnetworks while keeping model weights fixed, unlike movement pruning.These methods have been applied to convolutional networks and update scores rather than weights.
- Higher-order pruning: Hessian-based methods use second-order information, whereas movement pruning obtains importance scores from standard fine-tuning without costly second-order derivatives.Other cited approaches use absolute or squared values to select weights.
- Language-model compression: Knowledge distillation and structured pruning are alternative approaches for compressing pretrained language models.Movement pruning targets individual weights and can be further improved by adding a teacher model.
3 Background: Score-Based Pruning
Score-based pruning associates importance scores with weights, converts them into a binary mask, and applies that mask during inference. Automated gradual pruning jointly trains the model while increasing sparsity toward a target level.
- Score-based formulation: A parallel score matrix S assigns importance scores to weight matrix W and produces a binary mask M through Topv.Inference uses the masked product (W ⊙ M)x.
- Score-based formulation: Topv retains the v% highest-scoring weights and sets the remaining mask entries to zero.The supplied equation passage is fragmentary, but the surrounding formulation defines Topv as selecting the highest values in S.
- Magnitude pruning: Magnitude pruning uses absolute weight values as importance scores and applies Topv to form the mask.Iterative magnitude pruning repeatedly removes low-magnitude weights and retrains the sparsified model.
- Gradual pruning: Automated gradual pruning allows masked weights to update and increases sparsity from an initial value vi to a final value vf.The model is pruned and trained jointly after a warm-up period using a cubic sparsity schedule.
4 Movement Pruning
Movement pruning uses first-order training information rather than absolute weight magnitude to select connections during fine-tuning. Hard and soft variants update scores alongside weights, enabling task-adaptive sparse subnetworks.
- Core idea: Movement pruning retains weights moving away from zero, whereas magnitude pruning retains weights far from zero.The distinction treats magnitude pruning as zeroth-order and movement pruning as first-order.
- Hard movement pruning: Hard movement pruning computes M = Topv(S) while learning both weights W and importance scores S during training.The forward pass uses the masked weights, and zeroed weights can be discarded.
- Optimization: The straight-through estimator sends gradients through Topv to update importance scores despite the mask’s zero gradient.Masked weights therefore continue receiving score updates during training.
- Soft movement pruning: Soft movement pruning uses a threshold mask M = (S > τ) and a regularizer whose coefficient controls sparsity.The regularization term encourages importance scores to decrease over time.
- Interpretation: Movement pruning’s scores accumulate movement over gradient updates and reflect whether weights shrink toward or move away from zero.The method is related to L0 regularization, while differing in its straight-through formulation.
- Adaptivity: Fine-tuned weights stay close to pretrained values, so movement pruning can select both low- and high-valued weights using task-specific changes.The authors connect this adaptivity to movement pruning’s ability to prune based on end-task data.
5 Experimental Setup
The experiments fine-tune BERT-base on three English transfer-learning tasks and compare automated pruning methods under matched training conditions.
- Model and pruning setup: Experiments use BERT-base-uncased with roughly 84M parameters, freezing embeddings while fine-tuning transformer layers and the task-specific head.
- Datasets and tasks: The evaluation covers SQuAD v1.1, MNLI, and QQP, spanning span extraction and paired sentence classification.
- Training procedure: All pruning methods are fine-tuned for the same number of updates, using cubic sparsity schedules for magnitude, movement, and soft movement pruning.
- Baselines: The study compares against RPP, LayerDrop, and mini-BERT models, using author-reported results for RPP and LayerDrop.
- Pruning comparisons: Unlike post-hoc magnitude pruning, the proposed automated gradual pruning allows masked weights to be updated during training.
6 Results
Across pruning levels, magnitude pruning is strongest at low sparsity, while movement-based methods become superior at high sparsity. Soft movement pruning leads the high-sparsity comparisons, and distillation further improves performance.
- Pruning results: At low sparsity, magnitude pruning outperforms other methods with little or no loss relative to the dense model.
- Pruning results: At high sparsity, first-order methods retain strong performance with less than 15% of the weights remaining, while magnitude-pruning performance drops rapidly.
- High-sparsity results: 79.9 F1: soft movement pruning on SQuAD with 3% of the weights, compared with 54.5 F1 for magnitude pruning and 76.3 F1 for movement pruning.
- High-sparsity results: Soft movement pruning consistently outperforms hard movement pruning and L0 regularization in high-sparsity regimes.
- Baseline comparisons: Movement and soft movement pruning compare favorably with other baselines, although RPP is on par with soft movement pruning on QQP.
- Practical constraint: Sparse-model inference is not optimized on current hardware, so a small dense model may be preferable from an inference-speed perspective.
- Distillation: Distillation strictly increases performance across pruning methods and sparsity levels while preserving their relative comparisons.
- Distillation: With distillation, soft movement pruning reaches 95% of BERT-base performance using approximately 5% of encoder weights on SQuAD and MNLI.
7 Analysis
Movement pruning adapts sparsity to fine-tuning by selecting weights according to movement rather than magnitude alone. Local and global selection perform similarly overall, while global allocation changes layerwise sparsity patterns.
- Movement pruning is adaptive: Movement pruning produces a smoother remaining-weight distribution than magnitude pruning at the same sparsity.Magnitude pruning removes weights close to zero, whereas movement pruning retains weights across the interval except near zero.
- Movement pruning is adaptive: High movement scores systematically correspond to non-zero weights, although importance has no simple relationship with weight magnitude.Both high- and low-absolute-value weights can be considered important when their fine-tuning movement supports retention.
- Local and global masks perform similarly: Local Topv selects weights matrix by matrix, whereas global Topv selects the most important weights across the whole network.The comparison examines whether network-wide selection improves over matrix-level selection at different sparsity levels.
- Local and global masks perform similarly: Local and global selection perform similarly except at high sparsity, where global movement pruning gains 2.3 F1 points with 3% of weights remaining.Distillation boosts pruned-model performance, but the final difference between local and global selection remains marginal.
- Layerwise allocation: Global pruning allocates sparsity non-uniformly across Transformer layers, while global magnitude pruning tends to prune layers uniformly.Global first-order methods allocate more weights to lower layers and heavily prune the highest layers.
8 Conclusion
The paper studies pruning pretrained models during task-specific fine-tuning and compares zeroth- and first-order methods. It finds that straight-through-gradient movement pruning adapts to this regime and improves high-sparsity performance.
- The paper compares zeroth- and first-order pruning methods for pretrained models undergoing task-specific fine-tuning.
- A simple straight-through-gradient pruning method uses a first-order importance score and is effective for task-specific fine-tuning.
- Movement pruning consistently improves over existing methods in high-sparsity regimes for a transformer-based architecture.
- The paper proposes future group-sparsity penalties that could remove entire columns or filters through group-level scores.For Transformers, this would provide a systematic way to perform feature selection and remove embedding-matrix columns.
9 Broader Impact
Reducing NLP model memory could support privacy, security, energy efficiency, and broader access to state-of-the-art capabilities. Extremely compressed models could also let smaller teams study and experiment with large models.
- Running state-of-the-art NLP capabilities on devices could reduce the need to send potentially private data to remote servers.
- Sparse models complement hardware efforts to accelerate inference while reducing energy consumption.
- Extremely compressed models with comparable performance could let smaller teams or individual researchers experiment with large models on a single GPU.
A.1 Guarantees on the decrease of the training loss
Under smooth-loss and small-learning-rate assumptions, the paper proves that swapping in a more important connection decreases training loss. The guarantee extends beyond TopK masking but does not apply to absolute-value score masking.
- Assumptions and setup: Under smoothness and a sufficiently small positive learning rate, the proof analyzes replacing one remaining connection with another.The argument first considers TopK masking with only one active connection, then studies a swap at the next step.
- Loss decrease: When the incoming connection becomes more important than the outgoing one, the swap decreases training loss between consecutive steps.The Taylor expansion has a null first term and a negative second term under the stated inequalities.
- Generalization: The proof generalizes to a set of N simultaneously swapped connections.The text introduces a set E of N swapping connections.
- Masking functions: The guarantee is not specific to TopK masking and remains valid for Threshold masking M := (S >= τ).The paper states that the relevant inequalities remain valid and the proof is unchanged.
- Limitation: The loss-decrease guarantee fails for absolute-value score masking; a negative-threshold case can instead produce an increase in loss.The paper establishes this by contradiction and gives the negative-threshold masking function as the specific case.