Movement Pruning: Adaptive Sparsity by Fine-Tuning

Victor Sanh; Thomas Wolf; Alexander M. Rush

Source-linked AI summary

Movement Pruning: Adaptive Sparsity by Fine-Tuning

Victor Sanh, Thomas Wolf, Alexander M. Rush

arXiv:2005.07683v2May 15, 2020cs.CLcs.LG

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.
Source: §1 · p. 1
Method
Movement pruning learns deterministic first-order importance scores during fine-tuning and uses them to select sparse weights, with hard and soft variants.
Source: §1 · pp. 1–2, §3 · p. 2
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.
Source: §6 · p. 7, §8 · p. 9
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.
Source: §6 · p. 6, §8 · p. 9
Preliminary experiments found that increasing the number of pruning steps tended to improve end performance.
Source: §6 · p. 5

Abstract

from arXiv · show

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.Source: §1 · p. 1

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.Source: §1 · p. 1

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.Source: §1 · p. 1

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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §3 · p. 2
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.Source: §3 · p. 2

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.Source: §3 · pp. 2–3

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.Source: §3 · p. 3

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.Source: §4 · pp. 3–4

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.Source: §4 · pp. 3, 5

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.Source: §4 · p. 3

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.Source: §4 · p. 3

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.Source: §4 · p. 4

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.Source: §4 · p. 4, Figure 1 · p. 4

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.Source: §5 · p. 5

Datasets and tasks

The evaluation covers SQuAD v1.1, MNLI, and QQP, spanning span extraction and paired sentence classification.Source: §5 · p. 5

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.Source: §5 · p. 5

Baselines

The study compares against RPP, LayerDrop, and mini-BERT models, using author-reported results for RPP and LayerDrop.Source: §5 · p. 5

Pruning comparisons

Unlike post-hoc magnitude pruning, the proposed automated gradual pruning allows masked weights to be updated during training.Source: §5 · p. 5

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.Source: §6 · p. 6
At high sparsity, first-order methods retain strong performance with less than 15% of the weights remaining, while magnitude-pruning performance drops rapidly.Source: §6 · p. 6

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.Source: §6 · p. 6
Soft movement pruning consistently outperforms hard movement pruning and L0 regularization in high-sparsity regimes.Source: §6 · p. 6

Baseline comparisons

Movement and soft movement pruning compare favorably with other baselines, although RPP is on par with soft movement pruning on QQP.Source: §6 · p. 6

Practical constraint

Sparse-model inference is not optimized on current hardware, so a small dense model may be preferable from an inference-speed perspective.Source: §6 · p. 6

Distillation

Distillation strictly increases performance across pruning methods and sparsity levels while preserving their relative comparisons.Source: §6 · pp. 6–7, Figure 3 · p. 7
With distillation, soft movement pruning reaches 95% of BERT-base performance using approximately 5% of encoder weights on SQuAD and MNLI.Source: §6 · p. 7

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.Source: §7 · p. 8
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.Source: §7 · p. 8

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.Source: §7 · p. 8, Figure 5 · p. 8
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.Source: §7 · p. 8

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.Source: §7 · p. 8, Figure 6 · p. 8

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.Source: §8 · p. 9
A simple straight-through-gradient pruning method uses a first-order importance score and is effective for task-specific fine-tuning.Source: §8 · p. 9
Movement pruning consistently improves over existing methods in high-sparsity regimes for a transformer-based architecture.Source: §8 · p. 9
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.Source: §8 · p. 9

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.Source: §9 · p. 9
Sparse models complement hardware efforts to accelerate inference while reducing energy consumption.Source: §9 · p. 9
Extremely compressed models with comparable performance could let smaller teams or individual researchers experiment with large models on a single GPU.Source: §9 · p. 9

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.Source: §A.1 · p. 13

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.Source: §A.1 · p. 13

Generalization

The proof generalizes to a set of N simultaneously swapped connections.The text introduces a set E of N swapping connections.Source: §A.1 · p. 13

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.Source: §A.1 · p. 13

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.Source: §A.1 · p. 14