[Paper Review] Movement Pruning: Adaptive Sparsity by Fine-Tuning

Victor Sanh, Thomas Wolf+1 more|arXiv (Cornell University)|May 15, 2020

Topic Modeling49 references33 citations

TL;DR

The paper proposes movement pruning, a first-order, fine-tuning-aware weight pruning method that uses accumulated weight movement to select prunable weights, showing strong performance in high-sparsity regimes for pretrained language models like BERT, especially with distillation.

ABSTRACT

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.

Motivation & Objective

• Motivate pruning during transfer learning where pretrained weights are fine-tuned on task data.
• Introduce movement pruning as a first-order, adaptive pruning method.
• Provide mathematical foundations and compare to zeroth- and first-order pruning methods.
• Demonstrate strong performance in high-sparsity regimes and with distillation on NLP tasks.

Proposed method

• Define importance scores S and masking M to prune weights based on movement during fine-tuning.
• Use a straight-through estimator to train masks in hard movement pruning; gradient w.r.t. S follows ∂L/∂S_{i,j} = ∂L/∂a_i · W_{i,j} · x_j.
• Offer a soft movement pruning variant with a fixed threshold and a sparsity-regularizing term to encourage decreasing scores over time.
• Relate movement pruning to L0 regularization and describe how gradients propagate in this framework.
• Experiment with BERT-base-uncased on SQuAD, MNLI, and QQP, using a cubic sparsity schedule and task-specific fine-tuning.
• Incorporate distillation to boost performance when pruning.

Experimental results

Research questions

• RQ1 Can movement pruning outperform magnitude pruning in transfer learning scenarios for NLP?
• RQ2 How does first-order movement information during fine-tuning affect sparsity-patterns and performance at high sparsity levels?
• RQ3 Does soft movement pruning with distillation yield better trade-offs between model size and accuracy?
• RQ4 What are the characteristics of pruned models under local vs. global masking strategies?

Key findings

• Movement pruning significantly outperforms magnitude pruning in high-sparsity regimes (<15% weights remaining).
• Soft movement pruning yields the strongest performance among pruning methods in high-sparsity settings, especially with distillation.
• With 3% remaining weights on SQuAD, movement pruning achieves 65.2/76.3 (EM/F1) and soft movement pruning achieves 69.5/79.9.
• For MNLI at 3% remaining weights, soft movement pruning reaches 79.0/79.6 (acc/mm acc).
• For QQP at 3% remaining weights, soft movement pruning reaches 89.3/85.6 (acc/F1).
• Distillation boosts performance across all pruning methods, preserving strong performance at high sparsity (e.g., SQuAD 3% reaches 76.6/84.9 with distillation).