Finetuning Pretrained Transformers into RNNs

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 10630 Finetuning Pretrained Transformers into RNNs Jungo Kasai♡∗ Hao Peng♡ Yizhe Zhang♣ Dani Yogatama♠ Gabriel Ilharco♡ Nikolaos Pappas♡ Yi Mao♣ Weizhu Chen♣ Noah A. Smith♡♢ ♡Paul G. Allen School of Computer Science & Engineering, University of Washington ♣Microsoft ♠DeepMind ♢Allen Institute for AI {jkasai,hapeng,gamaga,npappas,nasmith}@cs.washington.edu {Yizhe.Zhang, maoyi, wzchen}@microsoft.com dyogatama@google.com Abstract Transformers have outperformed recurrent neural networks (RNNs) in natural language generation. But this comes with a signifi- cant computational cost, as the attention mech- anism’s complexity scales quadratically with sequence length. Efficient transformer vari- ants have received increasing interest in recent works. Among them, a linear-complexity re- current variant has proven well suited for au- toregressive generation. It approximates the softmax attention with randomized or heuris- tic feature maps, but can be difficult to train and may yield suboptimal accuracy. This work aims to convert a pretrained transformer into its efficient recurrent counterpart, improving efficiency while maintaining accuracy. Specif- ically, we propose a swap-then-finetune pro- cedure: in an off-the-shelf pretrained trans- former, we replace the softmax attention with its linear-complexity recurrent alternative and then finetune. With a learned feature map, our approach provides an improved tradeoff between efficiency and accuracy over the stan- dard transformer and other recurrent variants. We also show that the finetuning process has lower training cost relative to training these re- current variants from scratch. As many models for natural language tasks are increasingly de- pendent on large-scale pretrained transformers, this work presents a viable approach to improv- ing inference efficiency without repeating the expensive pretraining process.1 1 Introduction Transformer models (Vaswani et al., 2017) have advanced the state of the art beyond recurrent neu- ral network models (e.g., LSTMs, Hochreiter and Schmidhuber, 1997; GRUs, Cho et al., 2014) across a wide range of natural language processing tasks. In particular, the transformer architecture has been ∗Work was done during an internship at Microsoft. 1https://github.com/jungokasai/T2R/. widely used in autoregressive modeling such as lan- guage modeling (Baevski and Auli, 2019) and ma- chine translation (Vaswani et al., 2017). The trans- former makes crucial use of interactions between feature vectors over the input sequence through the attention mechanism (Bahdanau et al., 2015). However, this comes with significant computation and memory footprint during generation. Since the output is incrementally predicted conditioned on the prefix, generation steps cannot be parallelized over time steps and require quadratic time complex- ity in sequence length. The memory consumption in every generation step also grows linearly as the sequence becomes longer. This bottleneck for long sequence generation limits the use of large-scale pretrained transformers, such as GPT-3 (Brown et al., 2020), Image Transformer (Parmar et al., 2018), and DALL-E (Ramesh et al., 2021). Recent work aims at reducing the overhead of autoregressive transformers (Child et al., 2019; Ki- taev et al., 2020; Beltagy et al., 2020, inter alia). Among them are recurrent alternatives that approx- imate the standard softmax attention (Katharopou- los et al., …

Content truncated to 500 words