Learned Token Pruning for Transformers

By Sehoon Kim, Sheng Shen, David Thorsley, A. Gholami, Joseph Hassoun, Kurt Keutzer. 2021

225 citations recorded by the source corpus.

Efficient deployment of transformer models in practice is challenging due to their inference cost including memory footprint, latency, and power consumption, which scales quadratically with input sequence length. To address this, we present a novel token reduction method dubbed Learned Token Pruning (LTP) 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, whose value 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 and SQuAD 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.1× FLOPs reduction with less than 1% accuracy drop, which results in up to 1.9× and 2.0× throughput improvement on Intel Haswell CPUs and NVIDIA V100 GPUs. Furthermore, we demonstrate that LTP is more robust than prior methods to variations in input sequence lengths. Our code has been developed in PyTorch and open-sourced

Subfields: Computer Networks and Communications, Computer Vision and Pattern Recognition, Hardware and Architecture

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

In Computer Vision and Pattern Recognition

In Hardware and Architecture

Judy Hoffman

Citations in the corpus

DyTox: Transformers for Continual Learning with DYnamic TOken eXpansion (2021, background): “Finally, while the computational overhead is small, it may be important for thousand of tasks: future works should investigate token pruning [36] or token merging.”
Token Merging: Your ViT But Faster (2022, background): “…any number of tokens, several recent works have attempted to prune the tokens from transformers in both NLP (Goyal et al., 2020; Kim & Cho, 2020; Kim et al., 2021; Lassance et al., 2021) and Vision (Meng et al., 2022; Yin et al., 2022; Kong et al., 2022; Song et al., 2022; Rao et al., 2021;…”
Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (2023, methodology): “Nevertheless, these methods can only be applied to non-autoregressive models and typically require an additional re-training phrase, making them less suitable for auto-regressive LLMs like ChatGPT and Llama.”
A Survey on Dynamic Neural Networks for Natural Language Processing (2022, background): “;) or (.!?) PoWER (Goyal et al., 2020) attention drop tokens TR-BERT (Ye et al., 2021) hidden states forward tokens LAT (Kim and Cho, 2021) attention forward tokens LTP (Kim et al., 2022) attention drop tokens Transkimmer (Guan et al., 2022) hidden states forward tokens VCRNN (Jernite et al., 2017) ”
Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer (2021, methodology): “…for selection in all layers via twice inference ( last class attention ); taking the column mean of the attention matrix as the score of each token as proposed in (Kim et al. 2021) tention outperforms the other selection strategies and common sub-sampling methods on both accuracy and efﬁciency.”
ASPEN: High-Throughput LoRA Fine-Tuning of Large Language Models with a Single GPU (2023, unknown): “Approaches such as sorting inputs by length [53] and token pruning [23] aim to reduce padding tokens.”