Finetuning Pretrained Transformers into RNN s (2021)

Abstract: Transformers have outperformed recurrent neural networks (RNNs) in natural language generation. But this comes with a signifi- cant computational cost, as the attention mechanism’s complexity scales quadratically with sequence length. Efficient transformer variants have received increasing interest in recent works. Among them, a linear-complexity recurrent variant has proven well suited for autoregressive generation. It approximates the softmax attention with randomized or heuristic 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. Specifically, we propose a swap-then-finetune procedure: in an off-the-shelf pretrained transformer, 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 standard transformer and other recurrent variants. We also show that the finetuning process has lower training cost relative to training these recurrent variants from scratch. As many models for natural language tasks are increasingly dependent on large-scale pretrained transformers, this work presents a viable approach to improving inference efficiency without repeating the expensive pretraining process.

## introduction transformer models @xcite have advanced the state of the art beyond recurrent neural network models (e.g., lstms, @xcite across a wide range of natural language processing tasks. in particular, the transformer architecture has been widely used in autoregressive modeling such as language modeling @xcite and machine translation @xcite . the transformer makes crucial use of interactions between feature vectors over the input sequence through the attention mechanism @xcite . 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 complexity 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 @xcite , image transformer @xcite , and dall-e @xcite . recent work aims at reducing the overhead of autoregressive transformers @xcite @xcite . among them are recurrent alternatives that approximate the standard softmax attention @xcite @xcite . similar to recurrent neural networks (rnns), those models represent the context by a recurrent state with a fixed size, thereby achieving linear time and constant memory complexity in generation sequence length. when the recurrent state size is smaller than the sequence length, these variants provide substantial speed and memory advantages over the transformer. a small state size, however, tends to deteriorate the generation quality @xcite , leading to a tradeoff between efficiency and accuracy. this work improves the balance between efficiency and accuracy by a conversion approach: instead of training a recurrent alternative from scratch, we develop a method to convert a pretrained transformer into an efficient rnn that speeds up generation and reduces memory foot-prints. our conversion proceeds with a swap-thenfinetune process. specifically, we change the exponential similarity function in the attention mechanism to the dot product after a single-layer mlp feature mapping. we then finetune the mlp parameters and the other network parameters. our experiments in language modeling and machine translation show that the conversion can compress the context into a much smaller recurrent state than the sequence length (e.g., 1/16 of the sequence length in wikitext-103 language modeling) while retaining high accuracy. in addition, this conversion requires much less gpu time than training randomly initialized models from scratch. state-of-the-art models in many natural language tasks are increasingly dependent on large-scale pretrained transformer models (e.g., gpt-2, @xcite @xcite @xcite . converting a large off-the-shelf transformer to a lightweight inference model without repeating the whole training procedure is particularly useful in many downstream applications. our work focuses on text generation and presents a viable approach towards efficient inference with high accuracy. ## convert a transformer into an rnn the transformer architecture consists of multihead attention, feedforward, and layer normalization modules @xcite . when a transformer is trained for a sequence generation task with teacher forcing @xcite , the attention can be parallelized over positions because the target sequence is fully available. during generation, on the other hand, the output is incrementally constructed. as a result, the attention becomes an inference bottleneck for long sequences. we present a method to eliminate this bottleneck by converting a pretrained transformer into an efficient rnn of linear time and constant space complexity. we provide a detailed complexity analysis in terms of the sequence length and model dimensions. ## experiments we present extensive experiments on standard benchmarks for language modeling and machine translation. our results show that t2r achieves efficient autoregressive generation while retaining high accuracy. ## analysis and ablations we presented t2r, a method to convert a pretrained transformer into an efficient rnn. in this section, we analyze our conversion approach by examining the impact of the feature size and induced attention weight distributions. our analysis shows that t2r implicitly learns attention distributions similar to the original transformer. ## further related work in addition to the work we already discussed, we highlight related methods from prior work that make transformer models efficient. ## conclusion and future work we present t2r, a method that converts a pretrained transformer to a recurrent neural network that reduces the time and memory cost of autoregressive generation. our experiments in language modeling and machine translation demonstrated that our model produces an improved tradeoff between efficiency and accuracy over randomly initialized training and previous models with lightweight attention. our work provides further support for the claim that large-scale pretrained models can be compressed into efficient inference models that facilitate downstream applications.