Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643November 7–11, 2021. c©2021 Association for Computational Linguistics

10630

Finetuning Pretrained Transformers into RNNsJungo 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 recurrentneural networks (RNNs) in natural languagegeneration. But this comes with a signifi-cant computational cost, as the attention mech-anism’s complexity scales quadratically withsequence length. Efficient transformer vari-ants have received increasing interest in recentworks. Among them, a linear-complexity re-current variant has proven well suited for au-toregressive generation. It approximates thesoftmax attention with randomized or heuris-tic feature maps, but can be difficult to trainand may yield suboptimal accuracy. This workaims to convert a pretrained transformer intoits efficient recurrent counterpart, improvingefficiency 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 withits linear-complexity recurrent alternative andthen finetune. With a learned feature map,our approach provides an improved tradeoffbetween efficiency and accuracy over the stan-dard transformer and other recurrent variants.We also show that the finetuning process haslower training cost relative to training these re-current variants from scratch. As many modelsfor 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 theexpensive pretraining process.1

1 Introduction

Transformer models (Vaswani et al., 2017) haveadvanced the state of the art beyond recurrent neu-ral network models (e.g., LSTMs, Hochreiter andSchmidhuber, 1997; GRUs, Cho et al., 2014) acrossa 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 betweenfeature vectors over the input sequence throughthe attention mechanism (Bahdanau et al., 2015).However, this comes with significant computationand memory footprint during generation. Since theoutput is incrementally predicted conditioned onthe prefix, generation steps cannot be parallelizedover time steps and require quadratic time complex-ity in sequence length. The memory consumptionin every generation step also grows linearly as thesequence becomes longer. This bottleneck for longsequence generation limits the use of large-scalepretrained transformers, such as GPT-3 (Brownet al., 2020), Image Transformer (Parmar et al.,2018), and DALL-E (Ramesh et al., 2021).

Recent work aims at reducing the overhead ofautoregressive 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., 2020; Peng et al., 2021; Choromanskiet al., 2021; Schlag et al., 2021). Similar to recur-rent neural networks (RNNs), those models repre-sent the context by a recurrent state with a fixedsize, thereby achieving linear time and constantmemory complexity in generation sequence length.When the recurrent state size is smaller than thesequence length, these variants provide substantialspeed and memory advantages over the transformer.A small state size, however, tends to deteriorate thegeneration quality (Peng et al., 2021), leading to atradeoff between efficiency and accuracy.

This work improves the balance between effi-ciency and accuracy by a conversion approach:instead of training a recurrent alternative fromscratch, we develop a method to convert a pre-trained transformer into an efficient RNN thatspeeds up generation and reduces memory foot-

10631

prints. Our conversion proceeds with a swap-then-finetune process. Specifically, we change the expo-nential similarity function in the attention mecha-nism to the dot product after a single-layer MLPfeature mapping. We then finetune the MLP pa-rameters and the other network parameters. Ourexperiments in language modeling and machinetranslation show that the conversion can compressthe context into a much smaller recurrent state thanthe sequence length (e.g., 1/16 of the sequencelength in WikiText-103 language modeling) whileretaining high accuracy. In addition, this conver-sion requires much less GPU time than trainingrandomly initialized models from scratch.

State-of-the-art models in many natural languagetasks are increasingly dependent on large-scale pre-trained transformer models (e.g., GPT-2, Radfordet al., 2019; BERT, Devlin et al., 2019; RoBERTa,Liu et al., 2019; T5, Raffel et al., 2020; BART,Lewis et al., 2020; DeBERTa, He et al., 2021).Converting a large off-the-shelf transformer to alightweight inference model without repeating thewhole training procedure is particularly useful inmany downstream applications. Our work focuseson text generation and presents a viable approachtowards efficient inference with high accuracy.

2 Convert a Transformer into an RNN

The transformer architecture consists of multiheadattention, feedforward, and layer normalizationmodules (Vaswani et al., 2017). When a trans-former is trained for a sequence generation taskwith teacher forcing (Williams and Zipser, 1989),the attention can be parallelized over positions be-cause the target sequence is fully available. Duringgeneration, on the other hand, the output is incre-mentally constructed. As a result, the attention be-comes an inference bottleneck for long sequences.We present a method to eliminate this bottleneck byconverting a pretrained transformer into an efficientRNN of linear time and constant space complexity.We provide a detailed complexity analysis in termsof the sequence length and model dimensions.

2.1 Multihead Attention

The attention module takes as input sequences ofsource and target vectors. The source vectors areused to produce key and value features, while thetarget vectors are mapped to query vectors. Moreformally, denote by {xtgt

i }Ni=1 and {xsrc

j }Mj=1 the

target and source vectors, where xtgti ,xsrc

j ∈ Rh

and h is the model dimensionality. We assumer attention heads of d dimensions (h = dr). Foreach head, the input vectors are first mapped tod dimensional query, key, and value features bylearned affine transformations with W∗ ∈ Rd×hand b∗ ∈ Rd:

qi =Wqxtgti + bq, (1a)

kj =Wkxsrcj + bk, vj =Wvx

srcj + bv. (1b)

The similarities of each query vector qi with allM key vectors are computed and normalized toproduce attention coefficients, which are then usedto output a weighted average of the value vectors(Vaswani et al., 2017):

xouti =

M

∑j=1

sim (qi,kj)

∑M`=1 sim (qi,k`)

vj , (2a)

sim(x,y) = exp (x ⋅ y/√d) . (2b)

Multihead attention runs this procedure for each ofthe r heads in parallel and concatenates r outputvectors to get the final h dimensional vector.2

Generation Speed Overhead Fig. 1 depicts thetransformer computation steps from input vectorsand their time complexity. We assume that thetime complexity of multiplying an n×m matrix byan m × k is O(nmk) as implemented in cuBLAS(NVIDIA, 2014).3 It consists of the following twostages.

• Feature Mapping: computation of {qi}Ni=1,{kj}

Mj=1, and {vj}Mj=1 for all r heads from

input vectors (Eqs. 1a-1b). Time complexityof O(Nh2), O(Mh2), and O(Mh2).

• Attention: weighted average over the valuevectors (Eq. 2a). O(MNh), quadratic in se-quence length (M , N ).

Generation Memory Overhead In autoregres-sive generation, query, key, and value vectors con-sume space complexity of O(h), O(Mh), andO(Mh) in every generation step. Every step’sattention weight (Eq. 2a) spans over M source po-sitions, taking O(Mr) space, linear in sequencelength M .

2.2 Converting Transformers to RNNsTo address this generation bottleneck of quadratictime and linear space, we propose Transformer-to-RNN (T2R), a method to convert a pretrained

2Layer normalization (Ba et al., 2016), residual connection(He et al., 2016), and projection are suppressed for brevity.

3If the batch size is small enough, parallelization can speedup matrix multiplication.

10632

Pretrained Transformer T2R

RNN State

Figure 1: Attention computation steps and their time complexity in pretrained transformer and T2R models duringinference generation. Features φ(qi) and φ(kj) are directly computed from input vectors, and qi and kj are neverconstructed. M : source length; N : target length; h: model dimensions; k: feature size; r: # heads.

transformer to an RNN inference model of lineartime and constant memory complexity in sequencelength (Fig. 1). T2R follows a swap-then-finetuneprocedure that modifies the attention computationof a pretrained transformer, and finetunes the modelwith the task objective.

We first replace the dot-then-exponential similar-ity function in a pretrained transformer (Eq. 2b) by

sim (x,y) = φ (x) ⋅φ (y) , (3a)

φ (x) = relu (Wφx + bφ) . (3b)

Here Wφ ∈ Rk×d and bφ ∈ Rk are learned pa-rameters of a single-layer MLP. They map a d di-mensional vector to a k dimensional kernel featurespace. The relu activation (Fukushima, 1980) en-sures that the features are non-negative.4 DifferentMLP parameters are used for different attentionheads, and thus we add a total of rk(d + 1) learn-able parameters per layer (less than 0.2% parame-ter increase in our language model, §3). We thenfinetune all parameters in this modified network,including the MLP parameters, with the originaltask objective.5

During inference generation, we reformulate theattention computation (Eq. 2a) as

xouti =

M

∑j=1

sim (qi,kj)

∑M`=1 sim (qi,k`)

vj

=⎛

⎝

φ (qi) ⋅∑Mj=1φ (kj)⊗ vj

φ (qi) ⋅∑M`=1φ (k`)

⎞

⎠

⊺ (4)

4We found that relu stabilized training by prohibiting nega-tive similarities φ(q) ⋅φ(k). Other activation functions, suchas cos, tanh, and elu, did not improve performance.

5We tried training the MLP parameters only, but this set-ting resulted in degraded development performance.

by the associativity of matrix multiplication. Thisformulation lends itself to recurrent computation.In causal attention where each query only attendsto its prefix to predict the next word (M = i), definestates:

Si =i

∑j=1

φ (kj)⊗ vj , zi =i

∑j=1

φ (kj) (5)

where Si,zi ∈ Rk×d,Rk. These states can be com-puted recurrently (Katharopoulos et al., 2020):

Si = Si−1 +φ (ki)v⊺

i zi = zi−1 +φ (ki) (6)

In the self-attention or encoder-to-decoder (cross)attention of a sequence-to-sequence model, Si andzi are constant with respect to i and only need tobe computed once. Given the two states at positioni, we can obtain the output vector:

xouti = (

φ (qi)⊺ Si

φ (qi)⊺ zi)

⊺

(7)

This avoids quadratic computation with respect tothe input sequence length. We also speed up in-ference by merging the MLP feature map with theaffine feature maps that produce queries and keys.

φ (qi) = relu (Wqxtgti + bq) , (8a)

φ (kj) = relu (Wkxsrcj + bk) , (8b)

where Wq =WφWq, Wk =WφWk, (8c)

bq = bφ +Wφbq, bk = bφ +Wφbk. (8d)

After the model is trained, Eqs. 8c–8d are computedonce before generation; the intermediate featuresof qi and kj are never computed during inference.

Generation Speed Overhead The time com-plexity of each step in a T2R model is shown inFig. 1. Similar to the transformer, it proceeds over

10633

two stages.• Feature Mapping: computation of{φ(qi)}

Ni=1, {φ(kj)}

Mj=1, and {vj}

Mj=1

for all r heads (Eqs. 8a–8b). Time complexityof O(Nhkr), O(Mhkr), and O(Mh2).

• Attention: the RNN states and the outputsfor r heads (Eqs. 5–7) are computed withO(Mhk) and O(Nhk).

Comparing this with the pretrained transformer, wesee that if the feature size is much smaller thaninput sequence lengths (k ≪M,N ), the change inthe attention stage from O(MNh) to O(hk(M +

N)) in T2R brings a substantial speedup.

Generation Memory Overhead T2R onlyneeds to store the RNN state, and thus its spacecomplexity is O(hk), constant in sequence length.This implies reduction in memory footprint whenk ≪M , compared to the transformer’s O(Mh).

2.3 Autoregressive Linear Transformers

In principle, any kernel function can be used asthe similarity function in Eq. 2a (Tsai et al., 2019).Previous work proposed several untrainable fea-ture map functions φ and developed autoregressivetransformer variants with linear time and constantspace complexity in sequence length (Katharopou-los et al., 2020; Peng et al., 2021; Choromanskiet al., 2021). While those models follow similarcomputation steps to T2R, there are several differ-ences in generation efficiency. Since the featuremap in Katharopoulos et al. (2020) preserves inputdimensions, the feature size is always the same asthe head dimensions (k = d). This means that thespeedup and memory savings from using a smallfeature size are restricted by design. In our exper-iments (§3.3), our T2R models gain further effi-ciency by using a feature size that is even smallerthan the head dimensions (k = 32 and d = 128 forlanguage modeling). Peng et al. (2021) and Choro-manski et al. (2021) scale query and key vectorsby their norms before the random approximation tobound the error. Consequently, the feature mappingstage needs additional steps of producing interme-diate q and k and scaling them. T2R suppressesthese steps and speeds up generation further (§3.3).

3 Experiments

We present extensive experiments on standardbenchmarks for language modeling and machinetranslation. Our results show that T2R achieves

efficient autoregressive generation while retaininghigh accuracy.

3.1 Baselines and ComparisonWe compare performance with previous trans-former models for autoregressive generation withlinear time and constant space complexity in in-put sequence length.6 As discussed in §2.3, thoseprior methods correspond to two different untrain-able feature maps φ. We experiment with twotypes of feature maps for comparisons: ELU(φ (x) = elu (x) + 1, Katharopoulos et al., 2020);RFA (random feature approximation with softmaxtemperature reparameterization, Peng et al., 2021).Each feature map is evaluated in two settings: ran-dom initialization and pretrain. Random initializa-tion is our reimplementation of the experiments inKatharopoulos et al. (2020) and Peng et al. (2021).The pretrain setting follows the same protocol asT2R except that we use different feature mapsφ than our proposed one-layer MLP with reluactivation. Positive orthogonal random features(Performer, Choromanski et al., 2021) providesimilar random approximation to RFA and wereevaluated in the biology domain, but we found thatthis method caused training divergence in the lan-guage modeling task.7

3.2 Setup and ImplementationsWe apply our method to causal attention in lan-guage models and both cross and causal attentionin machine translation. For language modeling, weuse a 32-dimensional feature map function. We donot modify the encoder in machine translation asits generation speed overhead is much less signif-icant than the decoder (Kasai et al., 2021). Ourexploration showed that reducing the feature sizeof causal attention tends to have less impact onthe final translation accuracy as opposed to crossattention; we use feature sizes of 32 and 4 for crossand causal attention, respectively. This observation

6See §5 for our discussion on more transformer variantswith linear time complexity, but most of those variants needmodifications for autoregressive modeling and have yet to beempirically evaluated in autoregressive generation tasks.

7Our implementation closely follows the code releasedby the authors (https://github.com/lucidrains/performer-pytorch/blob/main/performer_pytorch/performer_pytorch.py#L75-L81), butdoes not subtract the maximum logit; otherwise it woulddisallow the linear complexity in causal attention. Weconjecture that this is the reason why Performer becomes lessstable in our experiments. We suspect that some techniquesare necessary to improve numerical stability in languagemodeling and machine translation.

10634

is consistent with previous work that showed thatcausal attention can be more drastically simplifiedthan cross attention in transformer machine transla-tion models (You et al., 2020; Tay et al., 2021).

3.2.1 Language ModelingWe use the WikiText-103 benchmark, whichconsists of 103M tokens sampled from EnglishWikipedia (Merity et al., 2017). We choose similarhyperparameters to prior work (Baevski and Auli,2019; Fan et al., 2020): 32 layers, 8 heads, 128head dimensions, 1024 model dimensions, 4096fully connected dimensions and dropout (Srivas-tava et al., 2014) and layer dropout rates of 0.2.We partition the training data into non-overlappingblocks of 512 contiguous tokens ignoring docu-ment boundaries and train the model to predicteach token from left to right (Baevski and Auli,2019). Validation and test perplexity are measuredby predicting the last 256 words out of the input of512 consecutive words to avoid evaluating tokensin the beginning with limited context (early tokencurse, Press et al., 2021). We generally follow theoptimization method from Baevski and Auli (2019),but some hyperparameters, such as the learning ratefor the T2R finetuning, are adjusted for better con-vergence than randomly initialized training. SeeAppendix A.1 for more details.

3.2.2 Machine TranslationWe experiment with 3 translation benchmarks:WMT14 EN-DE (4.5M train pairs, Bojar et al.,2016), WMT14 EN-FR (36M, Bojar et al., 2014),and WMT17 ZH-EN (20M, Bojar et al., 2017). Wefollow the preprocessing and data splits by previ-ous work (EN-DE: Vaswani et al., 2017; EN-FR:Gehring et al., 2017; EN-ZH: Hassan et al., 2018).We use the hyperparameters of the large sized trans-former (Vaswani et al., 2017): 6 layers, 16 attentionheads, 1024 model dimensions, and 4096 hiddendimensions for both the encoder and decoder. Weapply dropout with 0.3 and label smoothing withε = 0.1. Following Ott et al. (2018), we use anincreased batch size of approximately 460K to-kens. Each randomly initialized model is trainedfor 30K (60K for the large EN-FR dataset) stepsusing Adam with a learning rate of 5 ⋅ 10−4 andβ = (0.9,0.98) (Kingma and Ba, 2015). We ob-served that convergence of the T2R conversion canbe achieved with 20K (40K for EN-FR) steps anda reduced learning rate of 2 ⋅ 10−4. We average thecheckpoints from the last five epochs to obtain the

final model (Vaswani et al., 2017). In inference, weapply beam search with size 5 and length penalty0.6. Consistent with previous practice, we evalu-ate with tokenized BLEU (Papineni et al., 2002).Further details are described in Appendix A.1.

ppl. train

Model k dev. test timeELU + Random Init. 128 22.0 22.8 470hRFA + Random Init. 32 20.4 21.3 512hT2R + Random Init. 32 20.1 20.8 474hELU + Pretrain 128 21.5 22.2 97hRFA + Pretrain 32 20.8 21.6 104hT2R + Pretrain 32 19.0 19.6 98hT2R 75% + Pretrain 32 17.9 18.5 95hPretrained Transformer – 17.9 18.5 –Baevski and Auli (2019) – – 18.7 –

Table 1: WikiText-103 language modeling results (per-plexity). Train time is measured in GPU hours. Thetop two rows are our reimplementations of Katharopou-los et al. (2020) and Peng et al. (2021). Pretrain in-dicates initialization with a pretrained transformer forlanguage modeling. T2R 75% indicates a model whereevery fourth layer from the top is kept as the originaltransformer layer. Perplexity (ppl.) is measured by pre-dicting the last 256 words out of the input of 512 con-secutive words. All models use 128 head dimensions.We assume access to a pretrained transformer modeland measure the finetuning time in GPU hours.

3.3 ResultsLanguage Modeling Seen in Table 1 are lan-guage modeling results in perplexity. We observethat T2R with the learnable MLP feature map out-performs the other two linear transformer modelsby more than 2.0 perplexity points in the pretrainsetting. Unlike the other linear transformer models,T2R greatly benefits from pretraining (T2R + Pre-train: 19.6 vs. T2R + Random Init.: 20.8 test per-plexity points). We attribute this advantage of T2Rto the fact that the MLP feature map is able to learnattention patterns that are similar to those of the pre-trained transformer, as evidenced in §4. Notice alsothat the T2R conversion is ∼5x faster (measuredin GPU hours) than training a model from scratch.These results illustrate that a lightweight modelcan be obtained without repeating the expensivetraining of large-scale pretrained language modelssuch as GPT-2 and GPT-3 (Radford et al., 2019;Brown et al., 2020). T2R’s generation speedup(∼4x when producing 512 consecutive words) and

10635

Feature Size k WMT14 WMT17 Train Time

Model cross causal EN-DE EN-FR ZH-EN (GPU hours)ELU + Random Init. 64 64 28.4 * 23.4 120hRFA + Random Init. 32 4 28.1 41.7 23.4 135hT2R + Random Init. 32 4 27.5 39.8 23.1 123hELU + Pretrain 64 64 28.4 41.8 23.8 80hRFA + Pretrain 32 4 27.6 41.8 23.2 90hT2R + Pretrain 32 4 28.7 42.1 23.8 82hPretrained Transformer Large – – 28.9 42.2 24.2 –Vaswani et al. (2017) – – 28.4 41.8 – –

Table 2: Machine translation test results in BLEU scores. The top two rows are our reimplementations ofKatharopoulos et al. (2020) and Peng et al. (2021). Pretrain indicates initialization with a trained transformer-large model. *: diverged even when running with multiple random seeds and smaller learning rates. We assumeaccess to a pretrained transformer model and measure the finetuning time in GPU hours.

memory savings are later benchmarked with vary-ing sequence lengths. There remains a gap of 1.1perplexity points between the T2R and pretrainedtransformer models (19.6 vs. 18.5). However, thegap can be closed when every fourth layer fromthe top is kept as the original transformer layer andthe model is finetuned in the same way (T2R 75%).This suggests that keeping a small fraction of thequadratic attention layers can provide an effectivemiddle ground between efficiency and accuracy.8

Machine Translation Seen in Table 2 are ma-chine translation results in BLEU from various con-figurations. Departing from the language modelingexperiments, the T2R model underperforms theother two linear transformer models when initial-ized randomly. However, consistent with languagemodeling, the T2R model substantially benefitsfrom pretraining (e.g., 28.7 vs. 27.5 BLEU pointsin EN-DE). As a result, the T2R model achievessimilar BLEU scores to the original transformeracross all language pairs. ELU trained from thepretrained transformer yields comparable perfor-mance to T2R, but the feature size is much larger(64 vs. 32 and 64 vs. 4 in cross and causal attention),thus leading to increased overhead, as shown later.Note that the T2R finetuning time is only mod-erately smaller than that of randomly initializedtraining here, but further speedup in conversioncan be potentially achieved with more extensivehyperparameter tuning.9

8Concurrent work (Lei, 2021) also explores reducing thenumber of attention layers for efficiency.

9We found that the batch size could be reduced for T2Rconversion without hurting accuracy, while randomly initial-ized models deteriorate with small batch sizes. This suggests

8 16 32 64 128 256 512 1024 2048

1K

2K

3K

4K

5K

6K

7K

Sentence Length

DecodingSpeed(T

oken

s/s)

Transformer

ELU 64-64

RFA 32-4

T2R 32-4

Figure 2: Machine translation speed of various models.Speed is measured on a single TPU v2 accelerator withbatch size 16 and beam size 1, following Peng et al.(2021). 32-4 indicates the feature sizes of 32 and 4 forcross and causal attention, respectively.

Speedup and Memory Savings in GenerationWe run a conditional generation experiment to com-pare the decoding speed of the models in Table 2(Fig. 2). Here we assume the input and output se-quences are of the same length. All models aretested using greedy decoding with the same batchsize of 16 on a TPU v2 accelerator.10 We see thatindeed the linear transformer models can generatean almost constant number of tokens per secondregardless of the sequence length and outpace thetransformer model dramatically as the sequencebecomes longer. The T2R model achieves a 15%+

that the computational cost for conversion can be much lighterthan training from scratch, and T2R is advantageous whenonly a limited number of GPUs are available.

10https://opensource.google/projects/jax.

10636

8 16 32 64 128 256 51222

24

26

28

Sentence Length

Attention

Mem

ory

(MB)

Transformer

ELU 64-64

T2R/RFA 32-4

Figure 3: Memory consumption from the attentioncomputation of various machine translation models ininference with batch size 16 and beam size 1.

speedup over ELU and RFA due to its smaller fea-ture sizes and faster feature mapping respectively;this confirms our analysis on T2R’s speed advan-tage over them (§2.3). Fig. 3 plots memory con-sumption from the attention computation duringdecoding for machine translation. Since the T2R,RFA, and ELU models compress keys and valuesinto a k × d matrix S and a k dimensional vectorz (§2.2), the required memory at each decodingstep is constant over varying sequence lengths. Itis also roughly proportional to the feature size k.The MLP feature map in the T2R model allows forsmall feature dimensions than the ELU feature ofthe head dimensions, resulting in a 70% memory re-duction. The attention computation in the standardtransformer, on the other hand, consumes memorylinearly in sequence length at each decoding stepbecause all previous key and value vectors have tobe stored. We also found a similar speedup andmemory savings in unconditional generation withthe T2R language model (∼4x speedup in generat-ing 512 consecutive words over the transformer).

4 Analysis and Ablations

We presented T2R, a method to convert a pretrainedtransformer into an efficient RNN. In this section,we analyze our conversion approach by examiningthe impact of the feature size and induced attentionweight distributions. Our analysis shows that T2Rimplicitly learns attention distributions similar tothe original transformer.

Feature Size and Pretraining We saw that T2Rbenefits substantially from transformer pretraining.Fig. 4 compares T2R with pretraining and random

8 16 32 64 128 256 512

18

20

22

24

Transformer

Feature Size k

ValidationPerplexity

T2R + Random Init.

T2R + Pretrain

Figure 4: WikiText-103 validation perplexity with vary-ing feature sizes.

8 16 32 64 128 256 5120

0.1

0.2

0.3

Feature Size k

Average

EuclideanDistance

T2R before Finetuning

T2R MLP Frozen

T2R

Figure 5: Average Euclidean distance of T2R mod-els from the transformer attention weights with vary-ing feature sizes. The distances are computed onthe Wikitext-103 validation data for predicting a wordgiven the preceding 512 words. All models are initial-ized with a pretrained transformer model.

initialization in terms of the relation between thevalidation perplexity from WikiText-103 and thefeature sizes. We see that as the feature size (RNNstate size) becomes smaller, pretraining becomesparticularly important to achieve low perplexity.Transformer pretraining achieves a Pareto improve-ment over random initialization in the tradeoff be-tween efficiency (small feature size) and accuracy(low perplexity).

Attention Distribution T2R is not explicitlytrained to mimic the original attention distributions,and there is no guarantee that the MLP feature mapapproximates the exponential similarity function,unlike previous approximation approaches (Penget al., 2021; Choromanski et al., 2021). Here, weanalyze the properties of the attention weight dis-

10637

tributions that are induced by finetuning. We usethe validation data from WikiText-103 and run lan-guage models to predict the next word given theinput of 512 contiguous words. We compute theattention weight distribution over the 512 wordsfor each attention head in the model layers.

Fig. 5 compares the attention distributions fromT2R in various configurations. T2R MLP frozenindicates a model that is finetuned with the MLPparameters frozen. Euclidean distances in attentiondistributions between the original transformer andeach model are averaged across validation samples,model layers, and attention heads.11 ComparingT2R before finetuning and the full T2R model, wesee that the finetuning process induces much moresimilar attention distributions, and the distance di-minishes as the feature size increases (and the per-plexity approaches the original transformer, Fig. 4).We also observed that when the MLP parametersare not trained (T2R MLP frozen), the distancefrom the original attention distributions increases.These results suggest that finetuning of the wholenetwork in T2R implicitly develops similar atten-tion distributions to the original transformer eventhough the training supervision comes solely fromlanguage modeling.

5 Further Related Work

In addition to the work we already discussed, wehighlight related methods from prior work thatmake transformer models efficient.

5.1 Knowledge Distillation

Knowledge distillation (Hinton et al., 2015) isclosely related to our T2R conversion and usesa similar pipeline: a teacher model with large ca-pacity is first trained and is used to generate silvertraining data for a new lightweight inference model.It has been successfully applied to machine trans-lation (e.g., Kim and Rush, 2016; Gu et al., 2018)to make generation efficient. In particular, severalprior works distill a transformer translation modelto an RNN (Senellart et al., 2018; Kim et al., 2019).We share the same motivation toward fast genera-tion with light memory, but our approach differs intwo ways: the original training data are used forfinetuning an RNN model, and its model parame-ters are initialized with the “teacher” transformer.

11We do not consider random initialization baselines herebecause random initialization makes it impossible to alignattention heads and layers between models.

Our method does not use the computationally ex-pensive teacher model to generate new training data.While data generation is a one-time computationalcost, it becomes expensive as the teacher modelsize and training data increase. Moreover, sincethe pretrained parameters can be directly used, con-version requires fewer GPU hours than training abrand new lightweight model from scratch (§3.3).

5.2 Efficient TransformersPrior work suggested many other strategies to im-prove efficiency in transformers, such as weightsharing and factorization (Dehghani et al., 2019;Lan et al., 2020), weight and layer pruning (Michelet al., 2019; Fan et al., 2020), quantization (Zafriret al., 2019; Shen et al., 2020), and modifying thecombination of sublayers (Press et al., 2020; Man-dava et al., 2020). Some of these methods presentorthogonal design choices and can be integratedinto our T2R model to gain further efficiency. For amore comprehensive survey, see Tay et al. (2020b).Below we describe several prior works along twomajor strategies: compressing the attention contextand sparsifying the attention patterns.

Attention Context Compression This strand ofmethods compresses the context that is attended to,thereby reducing the time and memory overheadin the attention. RNN models that we convertedpretrained transformers into compress the contextinto a recurrent state. Other approaches include lowrank approximation of the attention computation(Wang et al., 2020; Tay et al., 2021) and adding amemory module that can access multiple tokens atonce (Liu et al., 2018; Dai et al., 2019; Lee et al.,2019; Ainslie et al., 2020; Rae et al., 2020; Beltagyet al., 2020; Zaheer et al., 2020).

Sparse Attention Patterns Another approach toreducing the time and memory overhead from theattention computation is to limit the tokens that areattended to by sparsifying the attention patterns.These patterns can be set in advance or learned dur-ing training (Tay et al., 2020b). For example, priorworks introduced fixed patterns of blockwise atten-tion (Qiu et al., 2020) and strided attention (Childet al., 2019; Beltagy et al., 2020; Zaheer et al.,2020). Other previous works presented methodsto learn attention patterns from data (Sukhbaataret al., 2019; Roy et al., 2020; Tay et al., 2020a).

It should be noted that significant modificationsare necessary to apply many of these methods toautoregressive generation tasks such as language

10638

modeling and machine translation, and their em-pirical evaluation in these generation settings hasyet to be conducted (Peng et al., 2021). This workpresents extensive empirical evaluation in autore-gressive generation settings.

6 Conclusion and Future Work

We present T2R, a method that converts a pre-trained transformer to a recurrent neural networkthat reduces the time and memory cost of au-toregressive generation. Our experiments in lan-guage modeling and machine translation demon-strated that our model produces an improved trade-off between efficiency and accuracy over ran-domly initialized training and previous models withlightweight attention. Our work provides furthersupport for the claim that large-scale pretrainedmodels can be compressed into efficient inferencemodels that facilitate downstream applications.

Acknowledgments

We thank Ofir Press, Bill Dolan, Lei Li, and theanonymous reviewers for their valuable feedbackand discussion on this work. Nikolaos Pappas wassupported by the Swiss National Science Founda-tion grant P400P2_183911.

ReferencesJoshua Ainslie, Santiago Ontanon, Chris Alberti, Va-

clav Cvicek, Zachary Fisher, Philip Pham, AnirudhRavula, Sumit Sanghai, Qifan Wang, and Li Yang.2020. ETC: Encoding long and structured inputs intransformers. In Proc. of EMNLP.

Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton.2016. Layer normalization.

Alexei Baevski and Michael Auli. 2019. Adaptive in-put representations for neural language modeling. InProc. of ICLR.

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Ben-gio. 2015. Neural machine translation by jointlylearning to align and translate. In Proc. of ICLR.

Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020.Longformer: The long-document transformer.

Ondrej Bojar, Christian Buck, Christian Federmann,Barry Haddow, Philipp Koehn, Johannes Leveling,Christof Monz, Pavel Pecina, Matt Post, HerveSaint-Amand, Radu Soricut, Lucia Specia, and AlešTamchyna. 2014. Findings of the 2014 workshop onstatistical machine translation. In Proc. of WMT.

Ondrej Bojar, Rajen Chatterjee, Christian Federmann,Yvette Graham, Barry Haddow, Shujian Huang,Matthias Huck, Philipp Koehn, Qun Liu, Varvara Lo-gacheva, Christof Monz, Matteo Negri, Matt Post,Raphael Rubino, Lucia Specia, and Marco Turchi.2017. Findings of the 2017 conference on machinetranslation (WMT17). In Proc. of WMT.

Ondrej Bojar, Rajen Chatterjee, Christian Federmann,Yvette Graham, Barry Haddow, Matthias Huck, An-tonio Jimeno Yepes, Philipp Koehn, Varvara Lo-gacheva, Christof Monz, Matteo Negri, AurélieNévéol, Mariana Neves, Martin Popel, Matt Post,Raphael Rubino, Carolina Scarton, Lucia Spe-cia, Marco Turchi, Karin Verspoor, and MarcosZampieri. 2016. Findings of the 2016 conferenceon machine translation. In Proc. of WMT.

Tom B. Brown, Benjamin Mann, Nick Ryder, MelanieSubbiah, Jared Kaplan, Prafulla Dhariwal, ArvindNeelakantan, Pranav Shyam, Girish Sastry, AmandaAskell, Sandhini Agarwal, Ariel Herbert-Voss,Gretchen Krueger, Tom Henighan, Rewon Child,Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu,Clemens Winter, Christopher Hesse, Mark Chen,Eric Sigler, Mateusz Litwin, Scott Gray, BenjaminChess, Jack Clark, Christopher Berner, Sam Mc-Candlish, Alec Radford, Ilya Sutskever, and DarioAmodei. 2020. Language models are few-shot learn-ers. In Proc. of NeurIPS.

Rewon Child, Scott Gray, Alec Radford, and IlyaSutskever. 2019. Generating long sequences withsparse transformers.

Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bah-danau, and Yoshua Bengio. 2014. On the propertiesof neural machine translation: Encoder–decoder ap-proaches. In Proc. of SSST-8.

Krzysztof Choromanski, Valerii Likhosherstov, DavidDohan, Xingyou Song, Andreea Gane, Tamás Sar-lós, Peter Hawkins, Jared Davis, Afroz Mohiuddin,Lukasz Kaiser, David Belanger, Lucy Colwell, andAdrian Weller. 2021. Rethinking attention with Per-formers. In Proc. of ICLR.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Car-bonell, Quoc Le, and Ruslan Salakhutdinov. 2019.Transformer-XL: Attentive language models beyonda fixed-length context. In Proc. of ACL.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals,Jakob Uszkoreit, and Lukasz Kaiser. 2019. Univer-sal transformers. In Proc. of ICLR.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, andKristina Toutanova. 2019. BERT: Pre-training ofdeep bidirectional transformers for language under-standing. In Proc. of NAACL.

Angela Fan, Edouard Grave, and Armand Joulin. 2020.Reducing transformer depth on demand with struc-tured dropout. In Proc. of ICLR.

10639

Kunihiko Fukushima. 1980. Neocognitron: A self-organizing neural network model for a mechanismof pattern recognition unaffected by shift in position.Biological Cybernetics.

Jonas Gehring, Michael Auli, David Grangier, DenisYarats, and Yann Dauphin. 2017. Convolutional se-quence to sequence learning. In Proc. of ICML.

Jiatao Gu, James Bradbury, Caiming Xiong, Vic-tor O. K. Li, and Richard Socher. 2018. Non-autoregressive neural machine translation. In Proc.of ICLR.

Hany Hassan, Anthony Aue, Chang Chen, VishalChowdhary, Jonathan Clark, Christian Feder-mann, Xuedong Huang, Marcin Junczys-Dowmunt,William Lewis, Mengnan Li, Shujie Liu, Tie-YanLiu, Renqian Luo, Arul Menezes, Tao Qin, FrankSeide, Xu Tan, Fei Tian, Lijun Wu, Shuangzhi Wu,Yingce Xia, Dongdong Zhang, Zhirui Zhang, andMing Zhou. 2018. Achieving human parity on au-tomatic Chinese to English news translation.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and JianSun. 2016. Deep residual learning for image recog-nition. In Proc. of CVPR.

Pengcheng He, Xiaodong Liu, Jianfeng Gao, andWeizhu Chen. 2021. DeBERTa: Decoding-enhanced BERT with disentangled attention.

Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean.2015. Distilling the knowledge in a neural network.In Proc. of NeurIPS Deep Learning and Representa-tion Learning Workshop.

Sepp Hochreiter and Jürgen Schmidhuber. 1997. Longshort-term memory. Neural Computation.

Hakan Inan, Khashayar Khosravi, and Richard Socher.2017. Tying word vectors and word classifiers: Aloss framework for language modeling. In Proc. ofICLR.

Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross,and Noah A. Smith. 2021. Deep encoder, shallowdecoder: Reevaluating non-autoregressive machinetranslation. In Proc. of ICLR.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap-pas, and François Fleuret. 2020. Transformers areRNNs: Fast autoregressive transformers with linearattention. In Proc. of ICML.

Yoon Kim and Alexander M. Rush. 2016. Sequence-level knowledge distillation. In Proc. of EMNLP.

Young Jin Kim, Marcin Junczys-Dowmunt, Hany Has-san, Alham Fikri Aji, Kenneth Heafield, RomanGrundkiewicz, and Nikolay Bogoychev. 2019. Fromresearch to production and back: Ludicrously fastneural machine translation. In Proc. of WNGT.

Diederik P. Kingma and Jimmy Ba. 2015. Adam:A method for stochastic optimization. In Proc. ofICLR.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.2020. Reformer: The efficient transformer. In Proc.of ICLR.

Zhenzhong Lan, Mingda Chen, Sebastian Goodman,Kevin Gimpel, Piyush Sharma, and Radu Soricut.2020. ALBERT: A lite BERT for self-supervisedlearning of language representations. In Proc. ofICLR.

Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Ko-siorek, Seungjin Choi, and Yee Whye Teh. 2019.Set transformer: A framework for attention-basedpermutation-invariant neural networks. In Proc. ofICML.

Tao Lei. 2021. When attention meets fast recurrence:Training language models with reduced compute.

Mike Lewis, Yinhan Liu, Naman Goyal, Mar-jan Ghazvininejad, Abdelrahman Mohamed, OmerLevy, Veselin Stoyanov, and Luke Zettlemoyer.2020. BART: Denoising sequence-to-sequence pre-training for natural language generation, translation,and comprehension. In Proc. of ACL.

Peter J. Liu, Mohammad Saleh, Etienne Pot, BenGoodrich, Ryan Sepassi, Lukasz Kaiser, and NoamShazeer. 2018. Generating Wikipedia by summariz-ing long sequences. In Proc. of ICLR.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man-dar Joshi, Danqi Chen, Omer Levy, Mike Lewis,Luke S. Zettlemoyer, and Veselin Stoyanov. 2019.RoBERTa: A robustly optimized BERT pretrainingapproach.

Ilya Loshchilov and Frank Hutter. 2017. SGDR:stochastic gradient descent with restarts.

Swetha Mandava, Szymon Migacz, and Alex Fit Florea.2020. Pay attention when required.

Stephen Merity, Caiming Xiong, James Bradbury, andRichard Socher. 2017. Pointer sentinel mixture mod-els. In Proc. of ICLR.

Paul Michel, Omer Levy, and Graham Neubig. 2019.Are sixteen heads really better than one? In Proc. ofNeurIPS.

Paulius Micikevicius, Sharan Narang, Jonah Alben,Gregory Diamos, Erich Elsen, David Garcia, BorisGinsburg, Michael Houston, Oleksii Kuchaiev,Ganesh Venkatesh, and Hao Wu. 2018. Mixed pre-cision training. In Proc. of ICLR.

NVIDIA. 2014. The NVIDIA CUDA basic linear alge-bra subroutines (CUBLAS).

Myle Ott, Sergey Edunov, Alexei Baevski, AngelaFan, Sam Gross, Nathan Ng, David Grangier, andMichael Auli. 2019. fairseq: A fast, extensibletoolkit for sequence modeling. In NAACL Demon-strations.

10640

Myle Ott, Sergey Edunov, David Grangier, andMichael Auli. 2018. Scaling neural machine trans-lation. In Proc. of WMT.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. BLEU: a method for automatic eval-uation of machine translation. In Proc. of ACL.

Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, LukaszKaiser, Noam Shazeer, Alexander Ku, and DustinTran. 2018. Image transformer. In Proc. of ICML.

Adam Paszke, Sam Gross, Francisco Massa, AdamLerer, James Bradbury, Gregory Chanan, TrevorKilleen, Zeming Lin, Natalia Gimelshein, LucaAntiga, Alban Desmaison, Andreas Kopf, EdwardYang, Zachary DeVito, Martin Raison, Alykhan Te-jani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang,Junjie Bai, and Soumith Chintala. 2019. PyTorch:An imperative style, high-performance deep learn-ing library. In Proc. of NeurIPS.

Hao Peng, Nikolaos Pappas, Dani Yogatama, RoySchwartz, Noah A. Smith, and Lingpeng Kong.2021. Random feature attention. In Proc. of ICLR.

Ofir Press, Noah A. Smith, and Omer Levy. 2020. Im-proving transformer models by reordering their sub-layers. In Proc. of ACL.

Ofir Press, Noah A. Smith, and Mike Lewis. 2021.Shortformer: Better language modeling usingshorter inputs. In Proc. of ACL.

Ofir Press and Lior Wolf. 2017. Using the output em-bedding to improve language models. In Proc. ofEACL.

Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih,Sinong Wang, and Jie Tang. 2020. Blockwise self-attention for long document understanding. In Proc.of EMNLP.

Alec Radford, Jeff Wu, Rewon Child, David Luan,Dario Amodei, and Ilya Sutskever. 2019. Languagemodels are unsupervised multitask learners.

Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar,Chloe Hillier, and Timothy P. Lillicrap. 2020. Com-pressive transformers for long-range sequence mod-elling. In Proc. of ICLR.

Colin Raffel, Noam Shazeer, Adam Roberts, KatherineLee, Sharan Narang, Michael Matena, Yanqi Zhou,Wei Li, and Peter J. Liu. 2020. Exploring the limitsof transfer learning with a unified text-to-text trans-former. JMLR.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, ScottGray, Chelsea Voss, Alec Radford, Mark Chen, andIlya Sutskever. 2021. Zero-shot text-to-image gener-ation.

Aurko Roy, Mohammad Saffar, Ashish Vaswani, andDavid Grangier. 2020. Efficient content-basedsparse attention with routing transformers. TACL.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber.2021. Linear transformers are secretly fast weightmemory systems. In Proc. of ICML.

Jean Senellart, Dakun Zhang, Bo Wang, GuillaumeKlein, Jean-Pierre Ramatchandirin, Josep Crego,and Alexander Rush. 2018. OpenNMT system de-scription for WNMT 2018: 800 words/sec on asingle-core CPU. In Proc. of WNG.

Rico Sennrich, Barry Haddow, and Alexandra Birch.2016. Neural machine translation of rare words withsubword units. In Proc. of ACL.

Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, ZheweiYao, Amir Gholami, Michael W. Mahoney, and KurtKeutzer. 2020. Q-BERT: hessian based ultra lowprecision quantization of BERT. In Proc. of AAAI.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky,Ilya Sutskever, and Ruslan Salakhutdinov. 2014.Dropout: A simple way to prevent neural networksfrom overfitting. JMLR.

Sainbayar Sukhbaatar, Edouard Grave, Piotr Bo-janowski, and Armand Joulin. 2019. Adaptive at-tention span in transformers. In Proc. of ACL.

Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan,Zhe Zhao, and Che Zheng. 2021. Synthesizer: Re-thinking self-attention in transformer models. InProc. of ICML.

Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. 2020a. Sparse sinkhorn attention. InProc of ICML.

Yi Tay, M. Dehghani, Dara Bahri, and Donald Metzler.2020b. Efficient Transformers: A survey.

Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada,Louis-Philippe Morency, and Ruslan Salakhutdinov.2019. Transformer dissection: An unified under-standing for transformer’s attention via the lens ofkernel. In Proc. of EMNLP.

Ashish Vaswani, Noam Shazeer, Niki Parmar, JakobUszkoreit, Llion Jones, Aidan N. Gomez, ŁukaszKaiser, and Illia Polosukhin. 2017. Attention is allyou need. In Proc. of NeurIPS.

Sinong Wang, Belinda Z. Li, Madian Khabsa, HanFang, and Hao Ma. 2020. Linformer: Self-attentionwith linear complexity.

Ronald J. Williams and David Zipser. 1989. A learn-ing algorithm for continually running fully recurrentneural networks. Neural Computation.

Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin,and Michael Auli. 2019. Pay less attention withlightweight and dynamic convolutions. In Proc. ofICLR.

Weiqiu You, Simeng Sun, and Mohit Iyyer. 2020.Hard-coded Gaussian attention for neural machinetranslation. In Proc. of ACL.

10641

Ofir Zafrir, Guy Boudoukh, Peter Izsak, and MosheWasserblat. 2019. Q8BERT: quantized 8bit BERT.In Proc. of EMC2.

Manzil Zaheer, Guru Guruganesh, Kumar AvinavaDubey, Joshua Ainslie, Chris Alberti, Santiago On-tanon, Philip Pham, Anirudh Ravula, Qifan Wang,Li Yang, and Amr Ahmed. 2020. Big Bird: Trans-formers for longer sequences. In Proc. of NeurIPS.

A Appendix

A.1 Hyperparameters and Setting

All training is implemented in fairseq (Ott et al.,2019) and run with PyTorch 1.7.1 (Paszke et al.,2019), 8 Telsa V100 GPUs, and CUDA 11.0. Weused mixed precision and distributed training over8 GPUs (Micikevicius et al., 2018; Ott et al., 2018).Apart from EN→ZH where we used separate BPEoperations and only tied the decoder input and out-put embeddings, we tie all embeddings (Press andWolf, 2017; Inan et al., 2017). We experimentedwith feature sizes of [16, 32, 64] and [4, 8, 16, 32]for language modeling and machine translation re-spectively, and chose the smallest feature sizes thatretained the development performance comparedto the standard transformer.

A.1.1 Language ModelingWe generally follow the optimization method fromBaevski and Auli (2019). For optimizing a modelfrom random initialization, the learning rate is lin-early warmed up from 10−7 to 1 for the initial 16Ksteps and then annealed using a cosine learningrate schedule with cycles (Loshchilov and Hutter,2017). Each period lasts for twice the number ofupdates than the previous cycle, and we lower themaximum and minimum learning rates by 25%compared to the previous cycle. The initial mini-mum and maximum learning rates are 10−5 and 1respectively (Baevski and Auli, 2019). We train themodel with a batch size of about 74K tokens witha total of 286K steps (Baevski and Auli, 2019).When we convert a pretrained transformer to anRNN model by finetuning, we found that we couldspeed up training by reducing the warm-up steps,total update steps, maximum and minimum rates,and batch size to 8K steps, 142K steps, 5 ⋅ 10−6,0.5, and 25K tokens without loss in validation per-plexity.

Randomly Initialized Training We generallyfollow the hyperparameters chosen in Baevski andAuli (2019); Fan et al. (2020). Specifically, we listthe hyperparameters in Table 3 for easy replication.All other hyperparameter options are left as defaultvalues in fairseq.

Finetuning Pretrained Transformer Seen inTable 4 are the hyperparameters for finetuning apretrained transformer to RNN models. The learn-ing rates, the max number of updates, and the learn-ing period length are all reduced.

10642

architecture transformer_lm_wiki103criterion adaptive_losstokens-per-sample 512sample-break-mode none# max tokens 3072dropout rate 0.2layer dropout rate 0.2decoder embed dim 1024decoder ffn dim 4096# decoder attn heads 8optimizer naglr-scheduler cosinelr-period-updates 270Klr-shrink 0.75t-mult 2max-lr 1min-lr 1e-9lr 1e-4clip-norm 0.1warm-up lr 1e-7# warmup updates 16K# max updates 286K# GPUs 8update-freq 3

Table 3: Language modeling hyperparameters whenrandomly initialized in the fairseq library.

A.1.2 Machine TranslationWe experiment with 3 translation benchmarks:WMT14 EN-DE (4.5M train pairs, Bojar et al.,2016), WMT14 EN-FR (36M, Bojar et al., 2014),and WMT17 ZH-EN (20M, Bojar et al., 2017). Wefollow the preprocessing and data splits by previ-ous work (EN-DE: Vaswani et al., 2017; EN-FR:Gehring et al., 2017; EN-ZH: Hassan et al., 2018;Wu et al., 2019). These datasets are all encodedinto subwords by BPE (Sennrich et al., 2016). Werun joint BPE on all language pairs except EN-ZH.We use the hyperparameters of the large sized trans-former (Vaswani et al., 2017): 6 layers, 16 attentionheads, 1024 model dimensions, and 4096 hiddendimensions for both the encoder and decoder. Weapply dropout with 0.3, weight decay with 0.01 andlabel smoothing with ε = 0.1. Following Ott et al.(2018), we use an increased batch size of approx-imately 460K tokens by accumulating gradientswithout updating parameters.

Randomly Initialized Training We generallyfollow the hyperparameters chosen in Vaswani et al.(2017); Ott et al. (2018). Specifically, we list thehyperparameters in Table 5 for easy replication. Allother hyperparamter options are left as default val-ues in fairseq. The parameters from the last fiveepochs were averaged to obtain the final model.

Finetuning Pretrained Transformer Seen inTable 6 are the hyperparameters for finetuning a

architecture transformer_lm_wiki103criterion adaptive_losstokens-per-sample 512sample-break-mode none# max tokens 3072dropout rate 0.2layer dropout rate 0.2decoder embed dim 1024decoder ffn dim 4096# decoder attn heads 8optimizer naglr-scheduler cosinelr-period-updates 135Klr-shrink 0.75t-mult 2max-lr 0.5min-lr 1e-9lr 5e-5clip-norm 0.1warm-up lr 1e-7# warmup updates 8K# max updates 142K# GPUs 8update-freq 1

Table 4: Finetuning language modeling hyperparame-ters in the fairseq library. The learning rates aresmaller than randomly initialized training.

pretrained transformer to RNN models. The learn-ing rate and the max number of updates are reduced.The parameters from the last five epochs were againaveraged to obtain the final model.

A.2 Attention DistributionPeakiness of Attention Fig. 6 plots the averageentropy of the T2R models with and without pre-training. Entropy is averaged across validation sam-ples, layers, and attention heads. Comparing Figs.4 and 6, we see that there is strong correlationbetween validation perplexity and entropy. The en-tropy decreases (and thus the attention distributiongets peakier) when a large feature size is used or thetransformer pretraining is applied. This observa-tion hints at potential future improvement of lineartransformer models by introducing an inductivebias towards peaky attention distributions.

10643

architecture transformer_vaswani_en_de_bigcriterion label_smoothed_cross_entropylabel smoothing 0.1# max tokens 3584dropout rate 0.3weight decay 0.0encoder embed dim 1024encoder ffn dim 4096# encoder attn heads 16# encoder layers 6decoder embed dim 1024decoder ffn dim 4096# decoder attn heads 16# decoder layers 6max source positions 1024max target positions 1024Adam lrate 5e-4, 3e-4 (T2R)*Adam β1 0.9Adam β2 0.98lr-scheduler inverse squarewarm-up lr 1e-7# warmup updates 4000# max updates 30K, 60K (EN-FR)length penalty 0.6beam size 5# GPUs 8update-freq 16

Table 5: Machine translation hyperparameters whenrandomly initialized in the fairseq library. *: wereduced the learning rate for T2R to avoid training di-vergence.

architecture transformer_vaswani_en_de_bigcriterion label_smoothed_cross_entropylabel smoothing 0.1# max tokens 3584dropout rate 0.3weight decay 0.0encoder embed dim 1024encoder ffn dim 4096# encoder attn heads 16# encoder layers 6decoder embed dim 1024decoder ffn dim 4096# decoder attn heads 16# decoder layers 6max source positions 1024max target positions 1024Adam lrate 2e-4Adam β1 0.9Adam β2 0.98lr-scheduler inverse squarewarm-up lr 1e-7# warmup updates 4000# max updates 20K, 40K (EN-FR)length penalty 0.6beam size 5# GPUs 8update-freq 16

Table 6: Finetuning machine translation hyperparame-ters. The learning rate is smaller than randomly initial-ized training.

8 16 32 64 128 256 512

3.5

4

4.5

5

Transformer

Feature Size k

AverageAttentionEntrop

yT2R + Random Init.

T2R + Pretrain

Figure 6: Average entropy of the attention weights.They are computed on the Wikitext-103 validation datafor predicting a word given the preceding 512 words.