Source-linked AI summary
Parameter-Efficient Transfer Learning with Diff Pruning
Demi Guo, Alexander M. Rush, Yoon Kim
TL;DR
Large pretrained networks make separate task-specific finetuning difficult to deploy when many tasks must be stored. The paper introduces diff pruning, which learns a sparse task-specific parameter difference over a fixed shared pretrained model, and reports matching fully finetuned GLUE baselines while modifying 0.5% of pretrained parameters per task.
Problem
Separate full finetuned models are difficult to store for multi-task, memory-constrained settings because pretrained networks are large.
Method
Diff pruning learns a task-specific diff vector over fixed pretrained parameters and encourages sparsity with a differentiable approximation to the L0-norm.
Results
0.5% of pretrained parameters per task lets diff pruning match fully finetuned BERT baselines on GLUE.
Takeaways & Limitations
Diff pruning stores only nonzero task-specific diff positions and weights while amortizing the shared pretrained model across tasks.
Takeaways & Limitations
Training requires more memory and is approximately 1.5× to 2× slower per minibatch than regular finetuning.
Abstract
from arXiv · showhide
While task-specific finetuning of pretrained networks has led to significant empirical advances in NLP, the large size of networks makes finetuning difficult to deploy in multi-task, memory-constrained settings. We propose diff pruning as a simple approach to enable parameter-efficient transfer learning within the pretrain-finetune framework. This approach views finetuning as learning a task-specific diff vector that is applied on top of the pretrained parameter vector, which remains fixed and is shared across different tasks. The diff vector is adaptively pruned during training with a differentiable approximation to the L0-norm penalty to encourage sparsity. Diff pruning becomes parameter-efficient as the number of tasks increases, as it requires storing only the nonzero positions and weights of the diff vector for each task, while the cost of storing the shared pretrained model remains constant. It further does not require access to all tasks during training, which makes it attractive in settings where tasks arrive in stream or the set of tasks is unknown. We find that models finetuned with diff pruning can match the performance of fully finetuned baselines on the GLUE benchmark while only modifying 0.5% of the pretrained model's parameters per task.
1 Introduction
Diff pruning addresses the storage and scalability costs of task-specific finetuning by learning sparse task-specific extensions over a shared pretrained model. It matches fully finetuned BERT performance on GLUE while modifying only 0.5% of pretrained parameters per task.
- Motivation: Task-specific finetuning is effective but difficult to scale because each task requires storing a full model.This is especially problematic in multi-task and memory-constrained settings.
- Related approaches: Existing parameter-efficient approaches either retain substantial nonzero parameters per task or require task sets and joint training assumptions.Compressed models may retain 10%-30% of parameters, while multi-task methods train shared models across tasks.
- Approach: Diff pruning extends a pretrained model with a task-specific difference vector instead of modifying its architecture.This preserves task-specific extensions while avoiding architectural changes to the base model.
- Approach: The pretrained parameters remain fixed while the task-specific diff vector is finetuned and regularized with a differentiable approximation to the L0-norm.The regularizer encourages the diff vector to be sparse.
- Results: 0.5% of pretrained parameters per task is sufficient for diff pruning to match fully finetuned BERT baselines on GLUE.Storage becomes increasingly favorable as more tasks share the fixed pretrained model.
2 Background: Transfer Learning
The paper frames NLP transfer learning as pretraining followed by task-specific optimization, then examines the challenge of storing separate large models across potentially unknown or streaming task sets. Large pretrained models make independent finetuning increasingly parameter-inefficient.
- Pretrain-finetune paradigm: NLP transfer learning commonly initializes task models from pretrained parameters and finetunes them on task-specific objectives.Pretraining objectives include context prediction, autoencoding, machine translation, and language modeling.
- Formal setting: For each task, the objective is to produce model parameters θτ that minimize empirical risk.The model function and loss are task-indexed, while an optional regularizer may be included.
- Storage challenge: Independent task-specific parameters follow naturally from pretrain-finetune transfer but create substantial storage demands.The paper notes that BERTBASE and BERTLARGE contain 110M and 340M parameters, respectively.
- Alternative: Joint multi-task learning shares a model but usually requires the task set to be known in advance to prevent catastrophic forgetting.This conflicts with settings where tasks may arrive in a stream or remain unknown.
3 Diff Pruning
Diff pruning reparameterizes each task model as a fixed pretrained model plus a sparse learned diff, using differentiable L0 regularization and subsequent sparsity control. A structured variant groups parameters so pruning can reflect local architectural regions.
- 3 Diff Pruning: Diff pruning represents each task model as the pretrained parameters θ plus a task-specific diff vector δτ, while θ remains fixed.This makes the shared pretrained model’s storage cost amortized across tasks.
- 3 Diff Pruning: Sparse diff vectors reduce marginal storage because only task-specific differences need to be stored for new tasks.The method is motivated by making ∥δτ∥0 much smaller than ∥θ∥0 as the number of tasks increases.
- 3.1 Differentiable approximation to the L0-norm: The training objective uses a relaxed mask and dense weight vector so gradient-based optimization can approximate L0 sparsity.A continuous relaxation of the binary mask is multiplied by the dense vector, producing many exactly zero dimensions after clamping.
- 3.2 L0-ball projection with magnitude pruning for sparsity control: After training, the method can project the diff vector onto a target L0-ball by retaining only the largest-magnitude diff values and finetuning with fixed nonzero masks.Pruning is based on diff-vector magnitudes rather than the original model parameters.
- 3.3 Structured Diff Pruning: The structured extension partitions parameter indices into groups and uses one group-level gate so parameters can be removed jointly.This is intended to let pruning adapt to local regions of the model architecture.
4 Experiments
The experiments evaluate structured and non-structured diff pruning against full finetuning, last-layer finetuning, Adapters, and non-adaptive pruning on GLUE and SQuAD. They use BERTLARGE for the main comparisons and specify task metrics, implementation details, and training settings.
- Datasets and evaluation: The evaluation covers GLUE and the SQuAD extractive question answering dataset, using BERTLARGE for the main experiments.GLUE includes MNLI, QNLI, SST-2, CoLA, STS-B, MRPC, and RTE, with task-specific metrics.
- Pruning procedure: Magnitude pruning retains the top t% × d values of the diff vector, followed by finetuning with the nonzero mask fixed.The pruning criterion uses diff-vector magnitudes rather than pretrained model parameter magnitudes.
- Structured variant: Structured diff pruning groups dimensions to model dependence between parameters and allow modifications in local regions.The implementation defines groups using each matrix or bias vector of the pretrained model.
- Datasets and evaluation: GLUE performance is measured with Matthew’s correlation, Spearman correlation, F1, or accuracy depending on the task.CoLA uses Matthew’s correlation; STS-B uses Spearman; MRPC and QQP use F1; MNLI, QNLI, SST-2, and RTE use accuracy.
- Baselines: The baselines include full finetuning, last-layer finetuning, Adapters, and non-adaptive diff pruning.Both structured and non-structured variants of diff pruning are compared against these baselines.
- Training settings: The experiments initially train for 3 epochs and search over batch sizes of 5, 8, 12, or 16 and learning rates of 1×10^-5, 2×10^-5, or 5×10^-5.Fixed-mask finetuning uses 3 epochs generally and 5 epochs for MRPC, STS-B, RTE, and SST-2.
5 Results
On GLUE, structured diff pruning matches fully finetuned BERTLARGE while modifying only 0.5% of parameters per task. It also compares favorably with non-structured and non-adaptive variants, and achieves comparable or better SQuAD performance with 1.0% additional parameters.
- GLUE: 0.5% of pretrained parameters per task is sufficient for structured diff pruning to match fully finetuned BERTLARGE performance on GLUE.The main GLUE results are reported in Table 1.
- GLUE: Structured diff pruning performs better than non-structured diff pruning, while non-adaptive magnitude pruning performs significantly worse.The comparison indicates that learning the binary mask is important for performance.
- GLUE: Diff pruning obtains similar performance to Adapters while requiring many fewer parameters per task.This comparison applies to the GLUE experiments.
- SQuAD: 1.0% additional parameters yields comparable or better SQuAD performance for diff pruning.The approach measurably improves over the full finetuning baseline while modifying fewer parameters.
- SQuAD: Diff pruning’s SQuAD improvement over full finetuning is described as a useful regularization effect in addition to parameter-efficiency.The result is reported on SQuAD v1.1.
6 Analysis
The analysis examines how sparsity targets, structure, task-specific allocation, and implementation choices affect diff pruning’s efficiency and performance. It also compares diff pruning with adapters and compressed BERT models while noting storage and training caveats.
- 6.1 Varying the target sparsity: Structured diff pruning consistently outperforms non-structured and non-adaptive variants across target sparsity rates.The advantage of adaptive methods is more pronounced at extreme sparsity rates.
- 6.2 Structured vs. Non-structured Diff Pruning: Structured diff pruning makes entire parameter groups more likely to remain unchanged than non-structured diff pruning.Groups are formed from matrix and bias vectors, and the structured mask encourages pruning complete groups.
- 6.3 Task-specific Sparsity: Different tasks modify different parts of the pretrained network, with some qualitative similarities across task-specific layer distributions.The observed sparsity patterns show some consistency across repeated runs of the same task, despite randomness.
- 6.4 Exact sparsity control: Magnitude pruning followed by fixed-mask finetuning achieves exact sparsity targets with little performance loss.Fixed-mask finetuning is important for maintaining performance after projection onto the L0-ball.
- 6.5 Comparison against BERT compression: Diff pruning is more parameter-efficient than direct BERT compression across GLUE tasks while maintaining better performance.Compression methods can still offer faster inference, such as TinyBERT4 being 9.4× faster than BERTBASE.
- 6.6 Storage comparison: Diff pruning remains more storage-efficient than adapters in the reported Python implementation despite storing both nonzero positions and weights.Storage comparisons are implementation-specific and assume float32 weights and int32 positions.
- 6.7 Discussion and caveats: Diff pruning requires more training memory and is approximately 1.5× to 2× slower per minibatch than regular finetuning.The authors note that increasing model sizes could make the additional training-memory requirement problematic.
7 Related Work
The paper situates diff pruning among parameter-efficient transfer methods, model compression, masking, and regularization approaches. It contrasts these methods by their task-specific parameter costs, training settings, and performance or storage trade-offs.
- Diff pruning comparisons: Table 4 reports sparsity and validation performance for structured diff pruning, including settings without magnitude pruning and with 0.5% target sparsity.
- Diff pruning comparisons: Across GLUE tasks, diff pruning is compared with compressed BERT variants using BERTBASE as the 109M-parameter baseline, while file-size comparisons account for weights and positions.
- Multi-task learning: Multi-task learning and feature-based transfer reduce per-task parameters by sharing a model, but feature-based methods are generally outperformed by fully finetuned models.
- Multi-task learning: Adapters add smaller task-specific modules without requiring access to all tasks during training, providing a modular alternative to full finetuning.
- Learning to mask and regularization: Diff pruning is related to differentiable masking methods and to regularization toward pretrained or shared models, but the latter typically target goals other than parameter efficiency.
8 Conclusion
The paper concludes that diff pruning enables parameter-efficient transfer learning by modifying only a small task-specific parameter set while retaining finetuning performance. It also identifies structured sparsity and future combinations with pretraining, adapters, and compression as directions for further improvement.
- Diff pruning matches fully finetuned baselines on standard NLP benchmarks while requiring only a few additional parameters per task.
- Diff pruning can sometimes improve upon regular finetuning through a regularization effect.
- The structured variant provides further improvements over the basic diff pruning approach.
- Future work includes incorporating parameter-efficiency objectives into pretraining and combining diff pruning with adapters or model compression.
A.1 Hyperparameters
The SQuAD v1.1 experiments used distributed training across eight GPUs with specified batch, sequence, stride, learning-rate, and epoch settings.
- SQuAD v1.1 training used 8 GPUs with per-GPU batch size 3, maximum sequence length 384, and document stride 128.
- The SQuAD v1.1 learning rate was 3 × 10^-5, with 2 initial training epochs and 2 finetuning epochs.
A.2 Consistency of Nonzero Parameters
The SST-2 runs show both variation and consistency in which layers receive modified parameters under 0.5% target sparsity. The first layer is modified considerably more than other layers across all five runs.
- Across 5 SST-2 runs, the percentage of modified parameters varies across layers but retains a degree of consistency.
- The first layer is modified considerably more than other layers across all runs.
- Figure 3 orders layers from earlier to later, with the embedding layer at the top and each x-axis spanning 0% to 20%.