Transkimmer: Transformer Learns to Layer-wise Skim (2022)

Abstract: Transformer architecture has become the de-facto model for many machine learning tasks from natural language processing and computer vision. As such, improving its computational efficiency becomes paramount. One of the major computational inefficiency of Transformer based models is that they spend the identical amount of computation throughout all layers. Prior works have proposed to augment the Transformer model with the capability of skimming tokens to improve its computational efficiency. However, they suffer from not having effectual and end-to-end optimization 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 required by each layer. The skimmed tokens are then forwarded directly to the final output, thus reducing the computation of the successive layers. 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.97x average speedup on GLUE benchmark compared with vanilla BERT-base baseline with less than 1% accuracy degradation.

## introduction the transformer model @xcite has pushed the accuracy of various nlp applications to a new stage by introducing the multi-head attention (mha) mechanism @xcite . further, the bert @xcite 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. @xcite , lstm @xcite , 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" ## input skimming search space in this section, we discuss the challenges of dynamic input skimming idea in details. moreover, we compare techniques and design decisions from prior works described in tbl. 1. ## conclusion input skimming or dynamic input reduction is an emerging transformer model acceleration method studied by many works recently. this idea utilizes the semantic structure of language and the syntactic information of the input context for inference acceleration. compared to static model weight compression methods, input skimming explores the redundancy in the input and hidden state tensors. as such, it is orthogonal and compatible with those model compression algorithms with its dynamic feature. in this work, we propose an accurate and efficient transformer inference acceleration method by teaching it how to skim input contents. the proposed transkimmer method is trained with an easy and end-to-end paradigm. furthermore, transkimmer is also generally applicable to various transformer-based model structures. it is even compatible with the static model compression methods like knowledge distillation and weight sharing. we believe that the above features guarantee the transkimmer method a wide range of applicable production scenarios.