PaperlayerView exact revision on arXiv

Source-linked AI summary

A Fast Post-Training Pruning Framework for Transformers

Woosuk Kwon, Sehoon Kim, Michael W. Mahoney, Joseph Hassoun, Kurt Keutzer, Amir Gholami

arXiv:2204.09656v2cs.CLcs.LG

TL;DR

Transformer pruning is hindered by retraining costs and deployment complexity. This paper introduces a retraining-free structured-pruning framework using Fisher-guided mask search, rearrangement, and tuning; evaluated on BERT-family models, it reduces computation and latency with limited accuracy loss.

  • Problem

    Existing Transformer-pruning methods often require costly additional training and complex deployment pipelines, limiting practical use.

  • Method

    The framework uses a sample dataset and resource constraint to prune heads and filters through Fisher-based mask search, rearrangement, and mask tuning without retraining.

  • Results

    On GLUE and SQuAD, the method achieves up to 50% FLOPs reduction within 1% accuracy degradation and up to 1.56× latency speedup.

  • Takeaways & Limitations

    The end-to-end pipeline provides a fast, deployable alternative to retraining-based Transformer pruning, completing in 39 seconds for GLUE and 135 seconds for SQuAD.

  • Takeaways & Limitations

    The approach assumes the model has converged to a local minimum and treats model parameters as fixed while optimizing mask variables.

Abstract

from arXiv · show

Pruning is an effective way to reduce the huge inference cost of Transformer models. However, prior work on pruning Transformers requires retraining the models. This can add high training cost and high complexity to model deployment, making it difficult to use in many practical situations. To address this, we propose a fast post-training pruning framework for Transformers that does not require any retraining. Given a resource constraint and a sample dataset, our framework automatically prunes the Transformer model using structured sparsity methods. To retain high accuracy without retraining, we introduce three novel techniques: (i) a lightweight mask search algorithm that finds which heads and filters to prune based on the Fisher information; (ii) mask rearrangement that complements the search algorithm; and (iii) mask tuning that reconstructs the output activations for each layer. We apply our method to BERT-base and DistilBERT, and we evaluate its effectiveness on GLUE and SQuAD benchmarks. Our framework achieves up to 2.0x reduction in FLOPs and 1.56x speedup in inference latency, while maintaining < 1% loss in accuracy. Importantly, our framework prunes Transformers in less than 3 minutes on a single GPU, which is over two orders of magnitude faster than existing pruning approaches that retrain the models.

1 Introduction

Transformer deployment is challenging because models are large and slow, while existing structured-pruning pipelines often require costly retraining and complex user intervention. The proposed post-training framework avoids retraining and achieves substantial efficiency improvements with limited accuracy loss.

  • Transformer models are difficult to deploy efficiently because of their large model size and high inference latency.
  • Prior pruning methods often require retraining, additional hyperparameter tuning, and substantial engineering effort.These requirements can increase training time by up to 10×.
  • The proposed framework accepts a Transformer, a sample dataset, and a FLOPs or latency constraint, then produces a deployable pruned model without retraining or user intervention.Its three stages are Fisher-guided mask search, mask rearrangement, and mask tuning.
  • Within 1% of accuracy drop, the framework reduces 30–50% of original FLOPs and achieves up to 1.56× speedup on an NVIDIA V100 GPU.
  • The end-to-end pruning pipeline finishes in 39 seconds for GLUE and 135 seconds for SQuAD on average, over 100× faster than retraining-based methods.

2 Related Work

Transformer pruning includes unstructured and structured approaches, but practical deployment remains difficult because hardware benefits from structured sparsity while existing pipelines impose training and engineering costs. Post-training pruning for Transformers addresses a gap left by methods developed for other architectures.

  • Transformer efficiency techniques include architecture design, hardware-software co-design, distillation, quantization, neural architecture search, and pruning.
  • Unstructured pruning can compress models, but commodity GPUs generally cannot exploit its sparse patterns for inference speedup.
  • Structured pruning removes coarse-grained parameter groups through methods including factorization, block or tile sparsity, head pruning, and layer dropping.
  • Existing structured-pruning methods can require additional training costing up to 10× original training and complex pipelines with extra hyperparameters.
  • Post-training CNN pruning techniques are difficult to extend to Transformers because Transformer multi-head attention lacks the required repeating linear-layer and nonlinearity structure.

3 Overview

The framework prunes encoder-based Transformers by masking attention heads and FFN filters while retaining a dense architecture. Its pipeline uses initialized mask variables, resource-constrained Fisher-based search, rearrangement, and tuning.

  • Transformer Architecture: BERT-style encoder blocks contain a multi-head attention layer followed by a pointwise feed-forward network.The attention layer has independently parameterized heads, while the feed-forward network contains filters.
  • Granularity of Pruning: The framework structurally prunes attention heads and FFN filters but leaves embeddings and the final classifier unpruned.The resulting smaller dense architecture can be accelerated without specialized hardware logic.
  • Mask Variables: Mask variables are associated with head and filter outputs, initialized to 1, and set to zero to prune the corresponding head or filter.
  • Framework Inputs: The framework receives a fine-tuned Transformer, typically 1–2K downstream-task examples, and a FLOPs or target-hardware latency constraint.
  • Compression Pipeline: Its three-stage pipeline performs Fisher-based mask search, mask rearrangement, and mask tuning to select and refine pruned heads and filters.

4 Methodology

The framework formulates post-training Transformer pruning as constrained mask optimization, then uses Fisher-based search, rearrangement, and latency approximations to select structured sparsity without retraining.

  • Problem formulation: Pruning is formulated as minimizing loss over head and filter masks subject to a FLOPs or latency constraint.The cost is generally non-differentiable because it depends on the mask's l0-norm.
  • Fisher-based mask search: The method approximates the loss with a second-order Taylor expansion and replaces the infeasible Hessian with empirical Fisher information computed on a sample dataset.The gradient term is assumed close to zero near a local minimum, and the sample dataset defines the Fisher estimate.
  • Fisher-based mask search: Under a diagonal Fisher approximation, the search prunes heads and filters with the lowest importance scores while satisfying the FLOPs constraint.Algorithm 1 enumerates possible numbers of remaining heads, selects the largest feasible number of filters, and chooses the lowest-scoring units.
  • Latency-constrained pruning: For latency constraints, the framework fits a piece-wise linear latency model with a constant-overhead region and a linear region using hardware lookup-table measurements.This approximation enables the FLOPs-based search procedure to extend to non-linear hardware latency.
  • Fisher-based mask rearrangement: Mask rearrangement captures within-layer interactions ignored by the diagonal approximation through a block-diagonal Fisher model and a warm-start greedy search.The number of units pruned in each layer is fixed from the initial binary mask before reselecting the specific heads and filters.

5 Evaluation

The framework is evaluated on BERTBASE and DistilBERT across GLUE and SQuAD, showing substantial compression and latency improvements with limited accuracy loss. It also compares favorably with prior structured pruning methods while requiring far less pruning time.

  • Performance Evaluation: With only 1% accuracy drop, BERTBASE retains 60–70% of original FLOPs across all evaluated tasks.DistilBERT shows a similar pattern and reaches up to 50% FLOPs reduction on STS-B and MRPC.
  • Performance Evaluation: With batch size 256, BERTBASE achieves 1.47× average latency speedup and up to 1.56× on an NVIDIA V100 GPU.The measurements use latency-constrained pruned models deployed on real hardware.
  • Comparison with the Prior Methods: Our method achieves comparable or better FLOPs-accuracy trade-offs than prior structured pruning methods without model retraining.For high sparsity, the framework with retraining performs comparably to or better than prior methods at the same pruning cost.
  • Ablation Studies: Mask rearrangement and mask tuning both help recover baseline accuracy under a 60% FLOPs constraint.Mask tuning is especially important, recovering up to 2.88% accuracy.
  • Pruning Cost: The pruning pipeline finishes in 39 seconds for GLUE and 135 seconds for SQuAD on average.These times are reported as over 100× faster than retraining-based methods, while a separate cost comparison reports prior methods requiring 5−33 hours and ours less than a minute.

6 Conclusion

The paper presents a retraining-free, three-stage post-training pruning framework for Transformers. On BERTBASE and DistilBERT, it achieves substantial FLOPs and latency reductions with limited accuracy degradation and much shorter pruning time.

  • Conclusion: The framework uses Fisher-based mask search, mask rearrangement, and mask tuning to prune heads and filters and recover layer output signals.The three-stage decomposition is the core of the proposed pruning process.
  • Conclusion: Up to 50% FLOPs reduction is achieved within only 1% accuracy degradation on GLUE and SQuAD.The result is reported for experiments using BERTBASE and DistilBERT.
  • Conclusion: The method provides up to 1.56× latency speedup on an NVIDIA V100 GPU.This latency result accompanies the reported FLOPs and accuracy improvements.
  • Conclusion: The end-to-end pruning pipeline takes 39 seconds for GLUE and 135 seconds for SQuAD.The paper reports this as 2−3 orders of magnitude faster than prior methods.

A.1 Proof of Equation 9

The proof constructs alternative masks under the same FLOPs constraint and compares their importance scores. It uses the minimizing property of Algorithm 1 to derive a contradiction, establishing that its output minimizes the target objective.

  • Mask Construction: The proof starts from a mask m∗ produced by Algorithm 1 and defines h as the number of pruned MHA heads.The construction then forms a new mask m′ while preserving the MHA mask.
  • Mask Construction: The constructed mask m′ keeps the MHA mask and zeros out the k least-important FFN filter variables.Here k is chosen to satisfy the cost constraint after accounting for the head contribution.
  • Constraint and Comparison: The proof states that m′ satisfies the FLOPs constraint and that pruning the least-important filters yields the relevant inequalities.These inequalities support comparison with the mask selected by Algorithm 1.
  • Mask Construction: A second mask m⋆ is constructed by zeroing out h least-important head variables from a mask initialized to one.The proof identifies m⋆ with the mask corresponding to the algorithm’s iteration whose remaining-head count is h.
  • Contradiction: Because Algorithm 1 searches across different values of n and selects the minimum, the constructed comparison contradicts the assumption that another feasible mask is better.The contradiction establishes that Algorithm 1 outputs the minimizer of Eq. 8.

A.2 Latency-aware Search Algorithm

The latency-aware search algorithm extends the original mask search to satisfy a latency constraint using approximated layer latency functions and diagonal Fisher information. It separately handles constant latency regions before selecting the highest-scoring feasible mask.

  • Inputs and Output: Algorithm 2 takes a latency constraint, approximated MHA and FFN latency functions, and a diagonal Fisher information matrix as inputs.Its output is a binary mask for MHA heads and FFN filters.
  • Latency Modeling: The algorithm separately considers the constant part where the number of heads or filters is below its corresponding threshold.This distinguishes latency regions that do not change with additional pruning.
  • Search Procedure: It identifies the most important threshold-level heads and filters, then enumerates additional head selections under the remaining latency budget.For each candidate head count, it computes the largest feasible filter count.
  • Search Procedure: For every candidate, the algorithm selects the highest-scoring feasible heads and filters using summed diagonal Fisher information.It records each candidate mask and chooses the candidate with maximum score.

A.3 Derivation of Equation 11

The derivation reduces the constrained optimization under a block-diagonal Fisher-information assumption to independent per-layer optimization problems. A contradiction argument shows that layerwise minimizers jointly solve the global problem.

  • Layerwise decomposition: The proof assumes that solving each layer’s optimization problem yields a mask that satisfies the warm-start constraint.The layerwise mask is denoted ˆm = (ˆm1, . . . , ˆmL).
  • Layerwise decomposition: Under the block diagonal assumption, the global mask optimization can be reduced to independent optimization problems for each layer.The l-th diagonal block of the Fisher information matrix defines each layer’s subproblem.
  • Contradiction argument: If another feasible mask strictly improved the reduced objective, at least one layer would have to improve over its layerwise minimizer.This follows while retaining the constraint in Eq. 21.
  • Contradiction argument: Such an improvement contradicts the assumption that the corresponding layerwise mask minimizes its individual optimization problem.Therefore, no globally better feasible mask exists.
  • Conclusion: Consequently, independently optimized layer masks form an optimal solution to the reduced global problem.The conclusion follows from the contradiction argument under the stated block-diagonal assumption.

A.4 Formulating Equation 12 as a Linear Least Squares Problem

The mask-tuning formulation reconstructs a multi-head-attention layer’s original output while keeping rearranged zero entries fixed. This constrained reconstruction becomes a linear least-squares problem solved numerically for stability.

  • Reconstruction objective: Mask tuning minimizes reconstruction error for a multi-head-attention layer using the mask produced by mask rearrangement.The formulation explicitly takes ˆm as the rearranged mask.
  • Constraint: The zero-valued entries in the rearranged mask are fixed at zero during tuning to preserve the FLOPs/latency constraint.The zero-entry index set is denoted Z(m).
  • Least-squares formulation: With pruned mask entries excluded from reconstruction, the tuning objective is rewritten as a linear least-squares problem.This formulation preserves the selected pruning pattern while adjusting remaining mask variables.
  • Least-squares formulation: The mask values for pruned heads are multiplied into matrix A so their output activations cannot reconstruct the original output.This enforces the pruning pattern inside the reconstruction system.
  • Solution method: Although a closed-form solution is available, the implementation uses CuPy’s numerical solver for higher stability.The stated closed form is (A⊺A)−1A⊺B.

A.5 Experimental Details

The experiments evaluate the framework on fine-tuned BERTBASE and DistilBERT models using GLUE and SQuAD, with analyses of data size, sparsity, importance metrics, and pruning stages. Results cover accuracy–FLOPs trade-offs, retraining comparisons, and runtime breakdowns.

  • Setup: The baseline models are fine-tuned with standard recipes, then pruned using 2K training examples from GLUE and SQuAD.GLUE uses accuracy except STS-B, which uses Spearman correlation; SQuAD uses F1 score.
  • Setup: The GLUE evaluation covers sentence similarity, sentiment classification, and textual entailment or inference tasks, while SQuAD includes two question-answering datasets.CoLA and WLNI are excluded because of unstable behavior.
  • Sample dataset size: Using more samples can improve accuracy by up to 0.4%, but requires 2–4× longer pruning time, which remains below 3 minutes.The sample-size study compares 2K, 4K, and 8K examples under a 60% FLOPs constraint.
  • Baseline comparison: The framework’s FLOPs–accuracy trade-off is compared with Flop, SLIP, Sajjad et al., DynaBERT, EBERT, and BMP using reported results without additional distillation or augmentation.Accuracy-drop comparisons account for differences in baseline accuracy.
  • Retraining extension: For high sparsity, the framework can be combined with retraining after skipping mask tuning and fixing the binary mask.The BERTBASE experiment retrains for 3 epochs on the full training dataset with learning rate 2e−5.
  • High sparsity: At high sparsity, the method performs comparably to DynaBERT and better than other methods on SST-2, while consistently outperforming all baselines on MRPC.This comparison uses retrained pruned models without data augmentation or knowledge distillation.
  • High sparsity: The high-sparsity results indicate that mask search and mask rearrangement find more optimal binary masks than other methods.The authors connect this result to the corresponding analysis in Figure 7.
  • Runtime: DynaBERT, EBERT, and BMP incur substantial training costs through multi-stage or long retraining procedures, whereas the framework reports an end-to-end pruning time breakdown on one NVIDIA V100 GPU.The cited table partitions runtime into Gradient Computation, Mask Search, Mask Rearrangement, and Mask Tuning.
Loading 2204.09656v2…