PaperlayerView exact revision on arXiv

Source-linked AI summary

PLATON: Pruning Large Transformer Models with Upper Confidence Bound of Weight Importance

Qingru Zhang, Simiao Zuo, Chen Liang, Alexander Bukharin, Pengcheng He, Weizhu Chen, Tuo Zhao

arXiv:2206.12562v1cs.LG

TL;DR

The rewinding process yields models with worse generalization performance, especially in the high sparsity regime (Liang et al., 2021). PLATON smooths sensitivity with an exponential moving average and uses uncertainty in importance estimation to retain weights with low scores but high uncertainty. PLATON significantly outperforms existing approaches, particularly at high sparsity; at 90% sparsity on MNLI, it improves accuracy by 1.2% over state-of-the-art approaches.

  • Problem

    The rewinding process yields models with worse generalization performance, especially in the high sparsity regime (Liang et al., 2021).

    Source: §1 · p. 2

  • Method

    PLATON smooths sensitivity with an exponential moving average and uses uncertainty in importance estimation to retain weights with low scores but high uncertainty.

    Source: §7 · p. 15

  • Results

    PLATON significantly outperforms existing approaches, particularly at high sparsity; at 90% sparsity on MNLI, it improves accuracy by 1.2% over state-of-the-art approaches.

    Source: §1 · p. 4, §7 · p. 15

  • Takeaways & Limitations

    The reported gains span natural language understanding, question answering, and image classification, with particular effectiveness in high-sparsity settings.

    Source: §7 · p. 15

Abstract

from arXiv · show

Large Transformer-based models have exhibited superior performance in various natural language processing and computer vision tasks. However, these models contain enormous amounts of parameters, which restrict their deployment to real-world applications. To reduce the model size, researchers prune these models based on the weights' importance scores. However, such scores are usually estimated on mini-batches during training, which incurs large variability/uncertainty due to mini-batch sampling and complicated training dynamics. As a result, some crucial weights could be pruned by commonly used pruning methods because of such uncertainty, which makes training unstable and hurts generalization. To resolve this issue, we propose PLATON, which captures the uncertainty of importance scores by upper confidence bound (UCB) of importance estimation. In particular, for the weights with low importance scores but high uncertainty, PLATON tends to retain them and explores their capacity. We conduct extensive experiments with several Transformer-based models on natural language understanding, question answering and image classification to validate the effectiveness of PLATON. Results demonstrate that PLATON manifests notable improvement under different sparsity levels. Our code is publicly available at https://github.com/QingruZhang/PLATON.

1 Introduction

The introduction frames pruning as a way to ease the deployment costs of large Transformers, while warning that uncertain importance estimates can discard useful weights. PLATON accounts for both importance and uncertainty, and the introduction reports gains over existing approaches, especially at high sparsity.

1 Introduction

  • Mini-batch sampling and training dynamics make importance scores uncertain, so weights can alternate between pruning and activation and cause instability or divergence.The introduction illustrates sampled weights whose importance scores vary by orders of magnitude during training.Source: §1 · pp. 2–3
  • PLATON smooths importance scores with an exponential moving average and uses local temporal variation as an uncertainty estimate in its importance metric.It tends to retain weights with low importance but high uncertainty; the introduction reports lower variability for the proposed metric than for sensitivity.Source: §1 · p. 3
  • The introduction evaluates PLATON with BERT-base and DeBERTaV3-base on NLU and question answering, and ViT-B16 on image classification.The tasks include GLUE, SQuAD, CIFAR100, and ImageNet.Source: §1 · p. 4
  • At 90% sparsity on MNLI, PLATON achieves a 1.2% accuracy improvement over state-of-the-art approaches.The introduction also reports that PLATON significantly outperforms baselines, especially at high sparsity.Source: §1 · p. 4

2 Background

This background reviews pre-trained Transformers, pruning importance scores, and iterative pruning. It introduces magnitude and sensitivity as scoring approaches and describes the repeated update-and-prune procedure.

2.1 Pre-trained Transformer-based Models

  • Pre-trained Transformers learn semantic knowledge from large unlabeled datasets and can transfer it to downstream tasks through fine-tuning.The section situates these models in both natural language processing and computer vision.Source: §2.1 · p. 4

2.2 Importance Scores for Pruning

  • Magnitude scores use each weight’s absolute value, while sensitivity estimates the loss change when a parameter is zeroed; small magnitude need not mean small influence on model output.The section describes sensitivity as derived from a first-order Taylor expansion of the loss.Source: §2.2 · p. 4

2.3 Iterative Pruning

  • Iterative pruning ranks weights by importance and removes low-scoring weights as training proceeds, following a stochastic gradient descent step at each iteration.The target percentage of nonzero weights is set by a predefined parameter.Source: §2.3 · p. 5

3 Method

PLATON addresses variability in weight-importance estimates with sensitivity smoothing and uncertainty-aware scoring. It tends to retain weights whose estimated importance is low but uncertain, allowing their capacity to be explored.

3 Method

  • PLATON combines sensitivity smoothing with uncertainty qualification to address variability in importance estimates.The paper attributes variability to mini-batch sampling and training dynamics, including optimization settings such as dropout.Source: §3 · pp. 5–6
  • Exponential moving averages smooth sensitivity scores, limiting the effect of abrupt drops and producing a more stable sparsity pattern.The averaging also emphasizes recent sensitivity scores while dropping stale information.Source: §3 · p. 6
  • PLATON estimates uncertainty from local temporal variation between current sensitivity and its historical average.A large deviation indicates that the current mini-batch sensitivity may not yet reliably indicate a weight’s importance; uncertainty estimates are also smoothed with an exponential moving average.Source: §3 · p. 6
  • The importance score combines sensitivity and uncertainty in an upper-confidence-bound-style score, with a logarithmic transformation to spread highly skewed uncertainty values.The method sets c = 1 for the logarithmic term and describes the score as following the UCB policy.Source: §3 · pp. 6–7
  • Weights with low estimated importance but high uncertainty are retained so PLATON can explore their capacity.The paper frames this as a conservative choice because current sensitivity may be unreliable when uncertainty is high.Source: §3 · p. 7
  • PLATON prunes weights after each gradient update and summarizes the resulting training procedure in Algorithm 1.The algorithm outputs the pruned model at the final training iteration.Source: §3 · pp. 7–8

4 Experiments

The experiments evaluate PLATON across natural-language understanding, question answering, and image classification, and analyze its performance across sparsity levels, importance-score variants, and EMA settings.

4 Experiments

  • Experiments fine-tune BERT-base, DeBERTaV3-base, and ViT-B16 for GLUE, SQuAD v1.1, CIFAR100, and ImageNet; all reported gains passed significance tests (p < 0.05).Source: §4 · p. 8

4.1 Natural Language Understanding

  • PLATON achieves 83.1% accuracy on MNLI-m at 80% sparsity, 1.5% higher than the best-performing baseline, soft movement pruning; it performs better or on par with existing approaches across all datasets and sparsity levels.Source: §4.1 · p. 10
  • Movement pruning and its soft variant perform well on larger GLUE datasets but poorly or fail to converge on smaller datasets such as RTE, CoLA, and STS-B.Source: §4.1 · p. 10
  • At 70% sparsity on SST-2, PLATON scores 94.6 versus 92.7 for the baseline; at 90% sparsity, it scores 90.0 versus 83.4, with larger gains at higher sparsity.For DeBERTaV3-base, PLATON significantly outperforms the baseline across the reported datasets and sparsity levels.Source: Table 2 · p. 10, §4.1 · p. 10

4.2 Question Answering

  • PLATON significantly outperforms the baseline on DeBERTaV3-base SQuAD v1.1 and is more effective in the high-sparsity regime.Source: §4.2 · p. 11
  • PLATON consistently outperforms existing approaches on BERT-base SQuAD v1.1 across six sparsity settings in both exact match (EM) and F1.The paper reports that movement and soft movement pruning are more efficient at high sparsity, but can perform on par with or worse than magnitude pruning and ℓ0 regularization at low sparsity; PLATON is effective throughout.Source: Table 3 · p. 11, §4.2 · p. 11
  • 0.2% higher F1 than the best-performing baseline, magnitude pruning, is reported for PLATON at 50% sparsity on BERT-base SQuAD v1.1; the gain is 2.5% at 85% sparsity.Source: Table 3 · p. 11, §4.2 · p. 11

4.3 Image Classification

  • 76.3% accuracy on ImageNet at 80% sparsity compares with 66.5% for magnitude pruning and 66.6% for movement pruning; PLATON also significantly outperforms methods on CIFAR100 and ImageNet at all sparsity levels.The reported gains are more significant at high sparsity.Source: Table 4 · p. 12, §4.3 · p. 12

4.4 Analysis

  • Across QQP, MNLI-m, and SQuAD v1.1, PLATON consistently improves performance over the baseline at all sparsity levels, with gains more significant when sparsity is above 50%.Source: §4.4 · p. 12
  • Using only sensitivity (S = I) or only uncertainty (S = U) deteriorates model performance by over 1.0% on all three datasets; PLATON’s importance score is the product of sensitivity and uncertainty (S = I ⊙U).Source: §4.4 · p. 13
  • With β1 = β2 = 0.85, PLATON is robust to these EMA hyperparameters and still outperforms the baseline when the selected values are not optimal.The experiments vary β1 and β2 on BERT-base fine-tuned on MNLI-m.Source: §4.4 · p. 13

5 Extension to Structured Pruning

PLATON extends to structured pruning by computing importance group-wise rather than entry-wise, a more implementation-friendly route to inference speedup. On BERT-base SQuAD v1.1, its performance drops by about 1.5 and 2.0 at 50% and 60% sparsity, respectively.

5 Extension to Structured Pruning

  • Structured pruning removes columns or blocks and is more implementation-friendly for inference speedup than entry-wise unstructured pruning.Unstructured pruning can require low-level implementation not supported by existing deep-learning libraries.Source: §5 · p. 13

5.1 Uncertainty Adjusted Structured Pruning

  • PLATON computes importance group-wise rather than entry-wise, extending its pruning method to structured parameter groups.For a weight matrix, a group can represent one column.Source: §5.1 · p. 14

5.2 Experimental Results

  • About 1.5 and 2.0 performance drops are reported for structured PLATON at 50% and 60% sparsity, respectively.The experiment prunes BERT-base on SQuAD v1.1.Source: Table 6 · p. 14, §5.2 · p. 14
  • Structured pruning performs lower than unstructured pruning because it offers less control over individual weights, which can group important and unimportant weights together.The authors explain that weights in the same column may be pruned together.Source: §5.2 · p. 14

6 Discussion

The discussion attributes instability in movement pruning to variable mini-batch importance estimates and weights that can become stale while masked. PLATON accounts for uncertainty and reactivates pruned weights from zero.

6 Discussion

  • Mini-batch sensitivity estimates vary widely, causing weights to alternate between retained and pruned and making training unstable.The discussion reports empirical instability and even divergence for movement pruning.Source: §6 · p. 15
  • PLATON’s importance-score variability is drastically smaller, yielding more stable training and better model performance.The discussion attributes the reduction to considering uncertainty in importance estimation.Source: §6 · p. 15
  • Weights reactivated after several iterations can have stale values that strongly affect the model; PLATON instead restarts them at zero.The rationale given is that pruned weights remain part of training dynamics as zero-valued weights.Source: §6 · p. 15

7 Conclusion

PLATON combines weight sensitivity with uncertainty in importance estimation, using smoothed mini-batch sensitivity and local temporal variation. Across natural language understanding, question answering, and image classification, it outperforms existing approaches, particularly at high sparsity.

7 Conclusion

  • PLATON combines weight sensitivity and uncertainty, smoothing mini-batch sensitivity with an exponential moving average and quantifying uncertainty through local temporal variation.The conclusion describes the resulting importance score as smoothed sensitivity times uncertainty, which reduces variability and stabilizes training.Source: §7 · p. 15
  • More than 5.0% improvement is reported when pruning DeBERTaV3-base on SQuAD v1.1, illustrating PLATON’s effectiveness at high sparsity.The conclusion also reports significant outperformance across natural language understanding, question answering, and image classification.Source: §7 · p. 15
  • The conclusion reports that training is more stable and less sensitive to hyperparameters when 90% of weights are pruned.It also notes that performance does not change dramatically across different hyperparameter setups.Source: §7 · p. 15

A Sparsity Ratio Schedule

  • The appendix applies the cubic sparsity-ratio schedule used by existing methods to both baseline methods and PLATON.The cited schedule includes initial and final warmups.Source: A Sparsity Ratio Schedule · p. 21

B GLUE Dataset Statistics

This section presents dataset statistics for GLUE.

B GLUE Dataset Statistics

C.1 Training Details

The training-details section describes implementation bases and hyperparameter choices for PLATON on GLUE, and reports limited sensitivity to its β parameters.

C.1 Training Details

  • The BERT-base implementation is based on the MT-DNN codebase, while the DeBERTaV3-base GLUE implementation uses Huggingface Transformers.The section cites Liu et al. (2019a, 2020) for MT-DNN and He et al. (2021a) and Wolf et al. (2019) for the latter implementation.Source: §C.1 · pp. 21–22
  • PLATON’s GLUE implementation uses β1 = 0.85 across experiments after testing values from 0.75 to 0.90, and tests β2 values from 0.850 to 0.975.The section reports that β1 = 0.85 generally performed best.Source: §C.1 · p. 22
  • The implementation uses AdamW, selects learning rates from 3 × 10−5 to 1 × 10−4, and chooses batch sizes from 8, 16, or 32 for each dataset.Table 8 lists the detailed setup for each dataset.Source: §C.1 · p. 22, Table 8 · p. 22
  • PLATON’s performance is reported as not sensitive to β1 and β2 or dramatically altered by different hyperparameter setups.Baseline hyperparameters are set as reported by Sanh et al. (2020).Source: §C.1 · p. 22

D.1 Dataset

This appendix specifies the question-answering evaluation and training settings, alongside the ViT model configuration used for image classification.

D.1 Dataset

  • SQuAD v1.1 is used to evaluate PLATON on question answering, following Sanh et al. (2020).The dataset is attributed to Rajpurkar et al. (2016b).Source: §D.1 · p. 22

D.2 Training Details

  • For QA fine-tuning, all algorithms use batch size 16, 10 epochs, AdamW, and learning rate 3 × 10−5.Source: §D.2 · p. 23
  • PLATON uses a cubic sparsity schedule with ti = 5400 and tf = 22000; baselines follow Sanh et al. (2020).Source: §D.2 · p. 23

E Image Classification

E.1 Datasets

The image-classification evaluation uses CIFAR-100 and ILSVRC-2012 ImageNet, with ViT fine-tuning and PLATON training settings specified for the two datasets.

E.1 Datasets

  • CIFAR-100 has 100 classes and 600 images per class, split into 500 training and 100 testing images, at resolution 224.Source: §E.1 · p. 23
  • ILSVRC-2012 ImageNet contains 1.2 million images across 1,000 categories and is evaluated at resolution 384 × 384.Source: §E.1 · p. 23

E.2 Training Details

  • ViT fine-tuning uses SGD with momentum, learning rate 0.03, batch sizes 512 for CIFAR-100 and 150 for ImageNet, gradient clipping at norm 1, and cosine decay.Source: §E.2 · p. 23
  • PLATON uses EMA 0.85; CIFAR-100 runs 20,000 steps with 2,000-step warmup and 8,000-step final warmup, while ImageNet runs 60,000 steps with ti = 7,000 and tf = 10,000.Source: §E.2 · p. 24

Related papers

Loading 2206.12562v1…