References & Citations
Computer Science > Computation and Language
Title: Learned Token Pruning for Transformers
(Submitted on 2 Jul 2021 (this version), latest version 2 Jun 2022 (v3))
Abstract: A major challenge in deploying transformer models is their prohibitive inference cost, which quadratically scales with the input sequence length. This makes it especially difficult to use transformers for processing long sequences. To address this, we present a novel Learned Token Pruning (LTP) method that reduces redundant tokens as the data passes through the different layers of the transformer. In particular, LTP prunes tokens with an attention score below a threshold value, which is learned during training. Importantly, our threshold based method avoids algorithmically expensive operations such as top-k token selection which are used in prior token pruning methods, and also leads to structured pruning. We extensively test the performance of our approach on multiple GLUE tasks and show that our learned threshold based method consistently outperforms the prior state-of-the-art top-k token based method by up to ~2% higher accuracy with the same amount of FLOPs. Furthermore, our preliminary results show up to 1.4x and 1.9x throughput improvement on Tesla T4 GPU and Intel Haswell CPU, respectively, with less than 1% of accuracy drop (and up to 2.1x FLOPs reduction). Our code has been developed in PyTorch and has been open-sourced.
Submission history
From: Sehoon Kim [view email][v1] Fri, 2 Jul 2021 09:00:13 GMT (3676kb,D)
[v2] Thu, 23 Sep 2021 19:09:39 GMT (2150kb,D)
[v3] Thu, 2 Jun 2022 19:00:07 GMT (748kb,D)
Link back to: arXiv, form interface, contact.