Transkimmer: Transformer Learns to Layer-wise Skim

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 7275 - 7286 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Transkimmer: Transformer Learns to Layer-wise Skim Yue Guan1,2,⋆, Zhengyi Li1,2,⋆, Jingwen Leng1,2,⋆, Zhouhan Lin1,§ and Minyi Guo1,2,† 1Shanghai Jiao Tong University, 2Shanghai Qizhi Institute ⋆{bonboru,hobbit,leng-jw}@sjtu.edu.cn, §lin.zhouhan@gmail.com, †guo-my@cs.sjtu.edu.cn Abstract Transformer architecture has become the de- facto model for many machine learning tasks from natural language processing and com- puter vision. As such, improving its compu- tational efficiency becomes paramount. One of the major computational inefficiency of Transformer-based models is that they spend the identical amount of computation through- out all layers. Prior works have proposed to augment the Transformer model with the capa- bility of skimming tokens to improve its com- putational efficiency. However, they suffer from not having effectual and end-to-end op- timization of the discrete skimming predictor. To address the above limitations, we propose the Transkimmer architecture, which learns to identify hidden state tokens that are not re- quired by each layer. The skimmed tokens are then forwarded directly to the final output, thus reducing the computation of the successive lay- ers. The key idea in Transkimmer is to add a parameterized predictor before each layer that learns to make the skimming decision. We also propose to adopt reparameterization trick and add skim loss for the end-to-end training of Transkimmer. Transkimmer achieves 10.97× average speedup on GLUE benchmark com- pared with vanilla BERTbase baseline with less than 1% accuracy degradation. 1 Introduction The Transformer model (Vaswani et al., 2017) has pushed the accuracy of various NLP applications to a new stage by introducing the multi-head atten- tion (MHA) mechanism (Lin et al., 2017). Further, the BERT (Devlin et al., 2019) model advances its performances by introducing self-supervised pre-training, and has reached the state-of-the-art accuracy on many NLP tasks. Compared to the recurrent fashion models, e.g. RNN (Rumelhart et al., 1986), LSTM (Hochreiter and Schmidhuber, 1997), the Transformer model leverages the above attention mechanism to process Layer 1 Layer 2 Layer 3 Transkimmer Transformer Layers [CLS] It is a good film . Layer 1 Layer 2 Layer 3 Transformer [CLS] It is a good film . Dynamic Skim Token Embeddings Input Sequence Downstream Classifier Downstream Classifier Downstream Classifier “Positive” “Positive” Output Embeddings Hidden States Embeddings Figure 1: Overview of Transkimmer dynamic token skimming method. Tokens are pruned during the pro- cessing of Transformer layers. Note that actually we don’t need all the tokens given to the downstream clas- sifier in this sequence classification example. We show the full length output embedding sequence to demon- strate the forwarding design of Transkimmer. all the input sequence. By doing so, extremely large scale and long span models are enabled, resulting in a huge performance leap in sequence processing tasks. However, the computation complexity of the attention mechanism is O(N2) with the input length of N, which leads to the high computation demand of the Transformer model. Some prior works (Goyal et al., 2020; Kim and …

Content truncated to 500 words