Learned Token Pruning for Transformers

TL;DR

This paper introduces Learned Token Pruning (LTP), a method that adaptively removes unimportant tokens in transformer models during inference, reducing computational cost while maintaining high accuracy.

Contribution

LTP is a novel threshold-based token pruning method that adaptively varies sequence length and outperforms prior methods in accuracy and efficiency on GLUE tasks.

Findings

LTP achieves up to 2.5% higher accuracy than previous methods.

LTP reduces FLOPs by up to 2.1x with less than 1% accuracy loss.

LTP improves throughput by up to 1.9x on CPUs and 2.0x on GPUs.

Abstract

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.