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.

提出了 Transkimmer architecture，通过添加参数预测器和 reparameterization trick 来实现 Transformer 模型中的 token 跳过，从而提高计算效率，并在 GLUE 基准测试中取得了 10.97x 的平均加速比。