PaperlayerView exact revision on arXiv

Source-linked AI summary

Learned Token Pruning for Transformers

Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, Kurt Keutzer

arXiv:2107.00910v3cs.CL

TL;DR

Transformer deployment is difficult because self-attention inference cost grows quadratically with input length, and not all tokens are necessary for successful inference. The paper introduces Learned Token Pruning, which removes tokens below learned layer-specific attention thresholds while adapting to each example and avoiding top-k selection. Across GLUE and SQuAD, LTP improves efficiency and robustness while maintaining accuracy relative to prior pruning methods.

  • Problem

    Transformer deployment is challenging because self-attention cost scales quadratically with sequence length under practical resource constraints.

  • Method

    LTP adaptively prunes tokens whose attention-based importance scores fall below learned layer-specific thresholds using a differentiable soft binarized mask.

  • Results

    LTP outperforms prior token-pruning methods by up to ~2.5% accuracy at the same FLOPs and achieves up to 2.10× FLOPs reduction with less than 1% accuracy degradation.

  • Takeaways & Limitations

    LTP provides adaptive token pruning with minimal inference complexity and stronger robustness to variations in input sentence lengths.

  • Takeaways & Limitations

    Reported speedups are compared with unpadded baselines rather than baselines using extra padding to a task-level sequence length.

Abstract

from arXiv · show

Deploying transformer models in practice is challenging due to their inference cost, which scales quadratically with input sequence length. To address this, we present a novel Learned Token Pruning (LTP) method which adaptively removes unimportant tokens as an input sequence passes through transformer layers. In particular, LTP prunes tokens with an attention score below a threshold value which is learned for each layer during training. Our threshold-based method allows the length of the pruned sequence to vary adaptively based on the input sequence, and avoids algorithmically expensive operations such as top-k token selection. We extensively test the performance of LTP on GLUE tasks and show that our method outperforms the prior state-of-the-art token pruning methods by up to ~2.5% higher accuracy with the same amount of FLOPs. In particular, LTP achieves up to 2.1x FLOPs reduction with less than 1% accuracy drop, which results in up to 1.9x and 2.0x throughput improvement on Intel Haswell CPUs and NVIDIA V100 GPUs, respectively. Furthermore, we demonstrate that LTP is more robust than prior methods to variations on input sentence lengths. Our code has been developed in PyTorch and has been open-sourced.

1 INTRODUCTION

Transformer inference is costly because self-attention scales quadratically with sequence length, while token lengths vary substantially across examples and datasets. LTP addresses this by learning threshold-based, example-adaptive token pruning and achieves strong efficiency, accuracy, and robustness results.

  • Motivation: Self-attention has quadratic complexity in input sequence length, limiting efficient deployment on long sequences.Transformer models also face large size and constrained energy, compute, and memory resources.
  • Motivation: Structured token pruning progressively removes unimportant tokens during inference, exploiting that not every input word is required for successful inference.Unlike unstructured pruning, structured pruning is generally easier for commodity hardware to deploy.
  • Problem with prior methods: Input sequence lengths vary greatly within and between datasets, so one pruning configuration can under-prune short inputs or over-prune long inputs.Prior fixed-configuration methods may therefore mismatch individual examples and dataset distributions.
  • LTP: LTP adapts to each example’s length and content using learned layer-specific thresholds and avoids computationally expensive top-k operations.A differentiable soft binarized mask automates threshold learning for different layers and tasks.
  • Results: 2.10× FLOPs reduction with less than 1% accuracy degradation produced up to 1.93× GPU and 1.97× CPU throughput improvements over the unpruned FP16 baseline.LTP also outperformed SpAtten and LAT in most cases at comparable accuracy loss.
  • Results: Up to 16.4% accuracy gaps over LAT demonstrate LTP’s robustness across varying sentence-length distributions.The result concerns comparisons across different sentence-length distributions.

2 RELATED WORK

Prior transformer pruning work includes architectural, weight, and token pruning methods, with token pruning progressively removing less important input tokens. Existing token-pruning approaches differ in whether configurations are fixed, length-based, heuristic, or dynamically learned.

  • Pruning approaches: Transformer compression methods include efficient architectures, knowledge distillation, quantization, and pruning.The paper focuses its related-work discussion on pruning methods.
  • Weight pruning: Unstructured pruning removes arbitrary parameter or feature-map patterns, whereas structured pruning removes organized parameter sets more amenable to hardware deployment.Examples include attention-head, layer, matrix, block, and filter pruning.
  • Token pruning: Token pruning progressively removes less important input tokens during inference rather than pruning model parameters.PoWER-BERT and LAT learn or search pruning configurations, while SpAtten assigns configurations proportional to input length.
  • Limitations of prior token pruning: Fixed or length-proportional configurations apply the same pruning behavior to sequences sharing a task or length, regardless of their contents.This can be suboptimal when token importance differs across individual inputs.
  • Adaptive token pruning: TR-BERT dynamically adapts pruning to sequence length and content with reinforcement learning, but its large search space makes training difficult and costly.The method uses imitation-learning and action-sequence-sampling heuristics to mitigate this issue.

3 METHODOLOGY

The method progressively removes unimportant tokens using attention-derived importance scores and learned, layer-specific thresholds. A differentiable soft mask enables threshold learning before hard-pruning fine-tuning.

  • 3.1 Background: Transformer layers combine multi-head attention with point-wise feed-forward blocks and residual connections.Multi-head attention uses independently parameterized heads to measure pairwise token importance.
  • 3.1 Background: Token pruning reduces later-layer computation by removing unimportant tokens as sequences pass through transformer layers.Attention computation costs O(d^2n + n^2d), so reducing sequence length n targets its quadratic component.
  • 3.2 Threshold Token Pruning: Token importance is computed from the attention probability each token receives across all heads and other tokens.A token is considered important when it receives more aggregate attention.
  • 3.2 Threshold Token Pruning: LTP prunes a token when its importance score falls below a layer-specific threshold, avoiding computationally expensive top-k selection.Pruned tokens remain excluded from calculations in succeeding layers, progressively reducing computation.
  • 3.3 Learnable Threshold for Token Pruning: Learnable thresholds use a sigmoid-based soft mask during training, then binarize the mask and fine-tune the model with hard pruning.With sufficiently small temperature, the soft mask closely approximates hard masking while preserving gradient flow into thresholds.
  • 3.3 Learnable Threshold for Token Pruning: L1 regularization penalizes retaining tokens, and larger regularization values produce higher pruning ratios.The regularization gradient pushes thresholds upward when many tokens lie near the pruning boundary.

4 EXPERIMENTS

Experiments evaluate LTP on GLUE and SQuAD 2.0, comparing accuracy, computation, throughput, robustness to sequence-length shifts, and compression extensions. LTP generally preserves accuracy while reducing computation and outperforms prior pruning methods under matched FLOPs.

  • Experiment Setup: LTP evaluation covers six GLUE tasks and SQuAD 2.0, using task-appropriate accuracy, F1, Pearson, and Spearman metrics.The GLUE tasks include sentence similarity, sentiment classification, textual entailment, and natural language inference.
  • Performance Evaluation: 1.96× average speedup and up to 2.10× speedup are achieved within 1% accuracy degradation on GLUE.GFLOPs are averaged over development-set inputs because sequence lengths can vary across examples and layers.
  • Performance Evaluation: 1.89× speedup is achieved on SQuAD 2.0 with less than 1% F1-score drop, while baseline-level F1 is reached at 0.58 relative FLOPs.As pruning increases, F1 decreases for answer-containing examples but increases for no-answer examples.
  • Robustness to Sequence Length Variation: LTP outperforms LAT by up to 16.44% on QNLI and 9.20% on QQP for evaluation sequences above the third quantile.Both methods are trained using sequences shorter than the evaluation median; Table 3 reports results across evaluation-length ranges.
  • Ablation Studies: Learned thresholds consistently outperform manually assigned linearly rising thresholds at the same FLOPs.The manual approach sets each layer threshold from the final-layer threshold, whereas LTP learns thresholds during training.

5 CONCLUSIONS

LTP is a fully automated transformer token-pruning framework that adds minimal inference complexity through threshold comparisons. It improves accuracy-efficiency trade-offs and robustness across input sequence lengths.

  • LTP determines unimportant tokens by comparing token importance scores with learned threshold values.The thresholds are learned per layer through a differentiable soft binarized mask.
  • Up to ~2.5% higher accuracy is achieved than state-of-the-art token-pruning methods at the same FLOPs.
  • Up to 2.10× FLOPs reduction is achieved with only 1% accuracy degradation relative to the baseline model.
  • Preliminary implementation results show up to 1.9× CPU and 2.0× GPU throughput improvements.The reported hardware is an Intel Haswell CPU and an NVIDIA V100 GPU.
  • LTP exhibits significantly better robustness and consistency across different input sequence lengths.

A.1 Training Details

LTP training separates soft pruning from hard pruning and jointly learns model parameters and layer-specific thresholds during downstream-task training. The procedure uses specified initialization, temperature, regularization, and optimizer-search settings.

  • LTP training consists of two stages: soft pruning followed by hard pruning.
  • During soft pruning, model parameters and thresholds are trained on downstream tasks for 1 to 10 epochs.The epoch count depends on dataset size.
  • Thresholds are initialized with linearly rising values while the final layer uses a fixed threshold.
  • The temperature is searched over {1, 2, 5, 10, 20}e-4, while λ varies from 0.001 to 0.4.
  • SpAtten and LAT use their paper-based training settings with specified searches over token-retain ratios and length-drop probabilities.SpAtten varies its final token retain ratio from 1.0 to -1.0, while LAT tests 0.25, 0.3, 0.35, and 0.4.

A.2 Computation Efficiency Comparison

The computation-efficiency comparison measures top-k and threshold operations under matched batch and retention settings. Threshold pruning avoids the latency growth observed for top-k selection and can be substantially faster on Intel Haswell CPUs.

  • Latency is averaged over 1000 runs with batch size 32 across five token retain ratios from 10% to 50%.The ratios are evaluated for each sequence length.
  • Top-k latency increases as token retain ratios and sequence lengths become larger, unlike threshold pruning.Threshold pruning requires only a comparison operation.
  • Top-k operation is up to 7.4× and 33.4× slower than threshold operation on an Intel Haswell CPU.

A.3 Discussion

Figure A.2 compares how LAT and LTP reduce sequence lengths for inputs of varying lengths. LAT applies a fixed pruning configuration, whereas LTP can adapt pruning aggressiveness to sequence content.

  • Figure A.2 shows pruned sequence length as input sequence length varies.
  • LAT uses a fixed token-pruning configuration for every sequence in the dataset.
  • LTP varies pruning aggressiveness according to sequence content.

A.3.1 Example Sequence Length Trajectories.

LTP reduces the number of tokens processed per layer while maintaining a controlled 1% accuracy drop, and its accuracy is not significantly correlated with input sequence length.

  • A prior report found only 1.1× speedup for GPT-2 without a dedicated top-k hardware engine.This further confirms the inefficiency of top-k operations.
  • 25.86% fewer tokens per layer for MNLI-m and 12.08% fewer for SST-2 were calculated by LTP than LAT.Both methods were trained for a 1% accuracy drop compared with baseline.
  • No significant correlation between sequence length and accuracy was observed across multiple tasks.The result suggests LTP is not biased toward longer or shorter input sequences.

A.3.2 Unbiased Token Pruning for Various Sequence Length.

The appendix compares token-length trajectories, length distributions, and operation latency across input sequences. These analyses include LTP, LAT, and top-k or threshold operation comparisons on classification tasks and an Intel Haswell CPU.

  • TR-BERT does not report results on the described GLUE tasks, and the authors were unable to obtain high-accuracy convergence in their attempted runs.The attempts varied learning rates and the length-penalty parameter α.
  • Threshold-operation latency is independent of sequence length on an Intel Haswell CPU across token-retain ratios.The comparison is against the top-k operation.
  • LTP and LAT trajectories are compared on SST-2 and MNLI-m using samples selected across initial sequence lengths at a matched 1% accuracy drop.LTP mean sequence length and LAT baseline are shown as dotted and dashed lines, respectively.
  • Figure A.3 groups correctly and incorrectly classified sentences by their pruned sequence lengths across transformer blocks.The histogram uses green for correct classifications and red for incorrect classifications.
Loading 2107.00910v3…