Depth-Adaptive Transformer

1 INTRODUCTION . The size of modern neural sequence models ( Gehring et al. , 2017 ; Vaswani et al. , 2017 ; Devlin et al. , 2019 ) can amount to billions of parameters ( Radford et al. , 2019 ) . For example , the winning entry of the WMT ’ 19 news machine translation task in English-German used an ensemble totaling two billion parameters ( Ng et al. , 2019 ) . While large models are required to do better on hard examples , small models are likely to perform as well on easy ones , e.g. , the aforementioned ensemble is probably not required to translate a short phrase such as `` Thank you '' . However , current models apply the same amount of computation regardless of whether the input is easy or hard . In this paper , we propose Transformers which adapt the number of layers to each input in order to achieve a good speed-accuracy trade off at inference time . We extend Graves ( 2016 ; ACT ) who introduced dynamic computation to recurrent neural networks in several ways : we apply different layers at each stage , we investigate a range of designs and training targets for the halting module and we explicitly supervise through simple oracles to achieve good performance on large-scale tasks . Universal Transformers ( UT ) rely on ACT for dynamic computation and repeatedly apply the same layer ( Dehghani et al. , 2018 ) . Our work considers a variety of mechanisms to estimate the network depth and applies a different layer at each step . Moreover , Dehghani et al . ( 2018 ) fix the number of steps for large-scale machine translation whereas we vary the number of steps to demonstrate substantial improvements in speed at no loss in accuracy . UT uses a layer which contains as many weights as an entire standard Transformer and this layer is applied several times which impacts speed . Our approach does not increase the size of individual layers . We also extend the resource efficient object classification work of Huang et al . ( 2017 ) and Bolukbasi et al . ( 2017 ) to structured prediction where dynamic computation decisions impact future computation . Related work from computer vision includes Teerapittayanon et al . ( 2016 ) ; Figurnov et al . ( 2017 ) and Wang et al . ( 2018 ) who explored the idea of dynamic routing either by exiting early or by skipping layers . We encode the input sequence using a standard Transformer encoder to generate the output sequence with a varying amount of computation in the decoder network . Dynamic computation poses a challenge for self-attention because omitted layers in prior time-steps may be required in the future . We experiment with two approaches to address this and show that a simple approach works well ( §2 ) . Next , we investigate different mechanisms to control the amount of computation in the decoder network , either for the entire sequence or on a per-token basis . This includes multinomial and binomial classifiers supervised by the model likelihood or whether the argmax is already correct as well as simply thresholding the model score ( §3 ) . Experiments on IWSLT14 German-English ∗Work done during an internship at Facebook AI Research . translation ( Cettolo et al. , 2014 ) as well as WMT ’ 14 English-French translation show that we can match the performance of well tuned baseline models at up to 76 % less computation ( §4 ) . 2 ANYTIME STRUCTURED PREDICTION . We first present a model that can make predictions at different layers . This is known as anytime prediction for computer vision models ( Huang et al. , 2017 ) and we extend it to structured prediction . 2.1 TRANSFORMER WITH MULTIPLE OUTPUT CLASSIFIERS . We base our approach on the Transformer sequence-to-sequence model ( Vaswani et al. , 2017 ) . Both encoder and decoder networks contain N stacked blocks where each has several sub-blocks surrounded by residual skip-connections . The first sub-block is a multi-head dot-product self-attention and the second a position-wise fully connected feed-forward network . For the decoder , there is an additional sub-block after the self-attention to add source context via another multi-head attention . Given a pair of source-target sequences ( x , y ) , x is processed with the encoder to give representations s = ( s1 , . . . , s|x| ) . Next , the decoder generates y step-by-step . For every new token yt input to the decoder at time t , the N decoder blocks process it to yield hidden states ( hnt ) 1≤n≤N : h0t = embed ( yt ) , h n t = blockn ( h n−1 ≤t , s ) , ( 1 ) where blockn is the mapping associated with the nth block and embed is a lookup table . The output distribution for predicting the next token is computed by feeding the activations of the last decoder layer hNt into a softmax normalized output classifier W : p ( yt+1|hNt ) = softmax ( WhNt ) ( 2 ) Standard Transformers have a single output classifier attached to the top of the decoder network . However , for dynamic computation we need to be able to make predictions at different stages of the network . To achieve this , we attach output classifiers Cn parameterized by Wn to the output hnt of each of the N decoder blocks : ∀n , p ( yt+1|hnt ) = softmax ( Wnhnt ) ( 3 ) The classifiers can be parameterized independently or we can share the weights across theN blocks . 2.2 TRAINING MULTIPLE OUTPUT CLASSIFIERS . Dynamic computation enables the model to use any of the N exit classifiers instead of just the final one . Some of our models can choose a different output classifier at each time-step which results in an exponential number of possible output classifier combinations in the sequence length . We consider two possible ways to train the decoder network ( Figure 1 ) . Aligned training optimizes all classifiers simultaneously and assumes all previous hidden states required by the self-attention are available . However , at test time this is often not the case when we choose a different exit for every token which leads to misaligned states . Instead , mixed training samples several sequences of exits for a given sentence and exposes the model to hidden states from different layers . Generally , for a given output sequence y , we have a sequence of chosen exits ( n1 , . . . , n|y| ) and we denote the block at which we exit at time t as nt . 2.2.1 ALIGNED TRAINING . Aligned training assumes all hidden states hn−11 , . . . , h n−1 t are available in order to compute selfattention and it optimizes N loss terms , one for each exit ( Figure 1a ) : LLnt = log p ( yt|hnt−1 ) , LL n = |y|∑ t=1 LLnt , Ldec ( x , y ) = − 1∑ n ωn N∑ n=1 ωn LL n . ( 4 ) The compound loss Ldec ( x , y ) is a weighted average of N terms w.r.t . to ( ω1 , . . . ωN ) . We found that uniform weights achieve better BLEU compared to other weighing schemes ( c.f . Appendix A ) . At inference time , not all time-steps will have hidden states for the current layer since the model exited early . In this case , we simply copy the last computed state to all upper layers , similar to mixed training ( §2.2.2 ) . However , we do apply layer-specific key and value projections to the copied state . 2.2.2 MIXED TRAINING . Aligned training assumes that all hidden states of the previous time-steps are available but this assumption is unrealistic since an early exit may have been chosen previously . This creates a mismatch between training and testing . Mixed training reduces the mismatch by training the model to use hidden states from different blocks of previous time-steps for self-attention . We sample M different exit sequences ( n ( m ) 1 , . . . n ( m ) |y| ) 1≤m≤M and evaluate the following loss : LL ( n1 , . . . , n|y| ) = |y|∑ t=1 log p ( yt|hntt−1 ) , Ldec ( x , y ) = − 1 M M∑ m=1 LL ( n ( m ) 1 , . . . , n ( m ) |y| ) . ( 5 ) When nt < N , we copy the last evaluated hidden state hnt to the subsequent layers so that the self-attention of future time steps can function as usual ( see Figure 1b ) . 3 ADAPTIVE DEPTH ESTIMATION . We present a variety of mechanisms to predict the decoder block at which the model will stop and output the next token , or when it should exit to achieve a good speed-accuracy trade-off . We consider two approaches : sequence-specific depth decodes all output tokens using the same block ( §3.1 ) while token-specific depth determines a separate exit for each individual token ( §3.2 ) . We model the distribution of exiting at time-step t with a parametric distribution qt where qt ( n ) is the probability of computing block1 , . . . , blockn and then emitting a prediction with Cn . The parameters of qt are optimized to match an oracle distribution q∗t with cross-entropy : Lexit ( x , y ) = ∑ t H ( q∗t ( x , y ) , qt ( x ) ) ( 6 ) The exit loss ( Lexit ) is back-propagated to the encoder-decoder parameters . We simultaneously optimize the decoding loss ( Eq . ( 4 ) ) and the exit loss ( Eq . ( 6 ) ) balanced by a hyper-parameter α to ensure that the model maintains good generation accuracy . The final loss takes the form : L ( x , y ) = Ldec ( x , y ) + αLexit ( x , y ) , ( 7 ) In the following we describe for each approach how the exit distribution qt is modeled ( illustrated in Figure 2 ) and how the oracle distribution q∗t is inferred . 3.1 SEQUENCE-SPECIFIC DEPTH : . For sequence-specific depth , the exit distribution q and the oracle distribution q∗ are independent of the time-step so we drop subscript t. We condition the exit on the source sequence by feeding the average s of the encoder outputs to a multinomial classifier : s = 1 |x| ∑ t st , q ( n|x ) = softmax ( Whs+ bh ) ∈ RN , ( 8 ) where Wh and bh are the weights and biases of the halting mechanism . We consider two oracles to determine which of the N blocks should be chosen . The first is based on the sequence likelihood and the second looks at an aggregate of the correctly predicted tokens at each block . Likelihood-based : This oracle is based on the likelihood of the entire sequence after each block and we optimize it with the Dirac delta centered around the exit with the highest sequence likelihood . q∗ ( x , y ) = δ ( argmax n LLn ) . We add a regularization term to encourage lower exits that achieve good likelihood : q∗ ( x , y ) = δ ( argmax n LLn−λn ) . ( 9 ) Correctness-based : Likelihood ignores whether the model already assigns the highest score to the correct target . Instead , this oracle chooses the lowest block that assigns the largest score to the correct prediction . For each block , we count the number of correctly predicted tokens over the sequence and choose the block with the most number of correct tokens . A regularization term controls the trade-off between speed and accuracy . Cn = # { t | yt = argmax y p ( y|hnt−1 ) } , q∗ ( x , y ) = δ ( argmax n Cn − λn ) . ( 10 ) Oracles based on test metrics such as BLEU are feasible but expensive to compute since we would need to decode every training sentence N times . We leave this for future work .

Summary: This paper studies using dynamic computation to alter the number of Transformer decoding layers each token uses to translate a given sentence. The paper considers using two variants of losses:  aligned training - same layer wise prediction loss for all tokens, and mixed training: loss on the output of a random layer for each token. For sampling the different layers the paper considers different distributions based on likelihood and the prediction rate.