IOT: Instance-wise Layer Reordering for Transformer Structures

1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) has been the dominant architecture in deep learning models ( Hassan et al. , 2018 ; Ng et al. , 2019 ; Carion et al. , 2020 ; Radford et al. , 2019 ; Dai et al. , 2019 ; Lee et al. , 2019 ; Devlin et al. , 2018 ; Yang et al. , 2019 ; Cai & Lam , 2019 ) . A Transformer model is stacked by several identical blocks , and each block consists of sequentially ordered layers : the self-attention ( SA ) , encoder-decoder attention ( ED ) ( decoder only ) and feed-forward ( FF ) layer . Recently , various modifications have been proposed , where the focus is on replacing or inserting some components ( e.g. , attention layer/layer norm/position encoding ) in standard Transformer ( Wu et al. , 2019 ; Lu et al. , 2019 ; Shaw et al. , 2018 ; So et al. , 2019 ; Ahmed et al. , 2017 ) . Despite these Transformer alternatives have achieved improved performances , one critical element is almost neglected in current models , which is how to arrange the components within a Transformer network , i.e. , the layer order also matters . As pointed by He et al . ( 2016b ) , different orders of ReLU , batch normalization and residual connection significantly affect the performance of ResNet ( He et al. , 2016a ) . Therefore , we ask : What if we reorder the sequential layers in Transformer ( e.g. , SA→FF or FF→SA of encoder , SA→FF→ED or FF→ED→SA of decoder ) ? What is the best order for these different layers ? ∗Equal contribution and corresponding authors . 1https : //github.com/instance-wise-ordered-transformer/IOT We first conduct preliminary experiments . We vary the three layers in decoder with all six variants ( each with a unique order of the three layers ) and train these models . Results on IWSLT14 German→English translation are reported in Table 1 . As we can see , their performances are similar and no one is outstanding . The corpus BLEU variance is only 0.0045 , which means that simply reordering the layers and training over the whole corpus impacts little . Press et al . ( 2019 ) also reported this for machine translation , but they stopped here . This seems to be a negative answer . However , we take a further step and ask one more question : Does different data favor different ordered layers ? That is , we investigate whether each specific data has its own preference for one particular order . Intuitively , putting various data patterns in one order should not be the best choice . For example , harder samples may favor a particular order while easier ones favor another one . Thus , for each order , we count the ratio of samples that achieve the best score with that order . In Table 1 , we find they almost lie on a uniform distribution ( e.g. , 17.9 % samples achieve the best BLEU with order SA→ED→FF ) . Besides , we calculate the BLEU variance for each sample , and average all these variances , the result is 114.76 , which is much larger than above corpus variance ( 0.0045 ) . These both mean the data indeed has its own preference to different orders . In Table 2 , we present translations from all decoders for on example with BLEU and TER score to give an evidence . Motivated by above observations , in this work , we present Instance-wise Ordered Transformer ( IOT ) , in which the layer order is determined by the specific data through instance-wise learning . To achieve this , we utilize a light predictor to predict the confidence for each order , given the corresponding classification losses as training signals . However , directly training the predictor with conventional ( i.e. , NMT ) loss tends to quickly converge to a bad order , and ignore explorations on others . Thus , we introduce an exploration loss and an exploitation loss to make an effective training while keeping an unambiguous prediction for each data so that the best order can be decided during inference . We evaluate our approach on 3 sequence generation tasks , including neural machine translation ( NMT ) , abstractive summarization ( ABS ) and code generation ( CG ) . For NMT , we work on 8 IWSLT and 2 WMT tasks , both on low-resource and rich-resource scenarios . Our method can consistently obtain 1.0 BELU score improvements over Transformer . For ABS , IOT also outperforms Transformer and other baselines on Gigaword dataset . For CG tasks , the results on 2 large-scale real-world code datasets ( Java and Python ) collected from Github surpass the state-of-the-art performances . These all demonstrate the effectiveness of our IOT . Furthermore , we provide detailed studies to verify that the instance-wise learning and order selection make a reasonable and necessary modeling . The contributions of this work can be summarized as follows : • We are the first to leverage instance-wise learning for layer order selection in a Transformer model ( with shared parameters ) , and we demonstrate the instance-wise learning is critical . • We demonstrate our learning approach can be universally applied to other structures beside Transformer ( e.g. , Dynamic Convolutions ) , as long as there are multiple different layers . • Experiments on 3 sequence generation tasks and 9 datasets verify the effectiveness of IOT with consistent performance improvements . 2 RELATED WORK . Architecture Exploration Inventing novel architectures by human designing or automatic searching plays an important role in deep learning . Specific to Transformer structures , various modifications have been proposed . For example , human knowledge powered designs include DynamicConv ( Wu et al. , 2019 ) , Macaron Network ( Lu et al. , 2019 ) , Reformer ( Kitaev et al. , 2020 ) and others ( Fonollosa et al. , 2019 ; Ahmed et al. , 2017 ; Shaw et al. , 2018 ) . As for automatic searching , neural architecture Feed Forward 𝑋 ! 𝑋 '' 𝑋 # 𝑌 '' 𝑌 # 𝑌 ! L x Pred L x Weight Tying Encoder DecoderInput Output Self Attention Self Attention Enc-Dec Attention Feed Forward Self Attention Enc-Dec Attention Feed Forward Feed Forward Self Attention Pred Weight Tying Figure 1 : The IOT framework . Pred means the light predictor introduced in 3.1 for order selection . We show two ordered encoders/decoders here . After taking X1 , X2 , X3 , the selected order for Y2 , Y3 is the lower encoder and upper decoder , while for Y1 is the upper encoder and lower decoder . search can discover networks with state-of-the-art performances but always with complicated computation , i.e. , Evolved Transformer ( So et al. , 2019 ) . The underlying principle is to add or replace some components of Transformer . For instance , Wu et al . ( 2019 ) replace self-attention with dynamic convolution , So et al . ( 2019 ) add a separate convolution layer in a new branch . Different from them , we , instead , only focus on the selection of layer orders for each data sample so as to improve the model performance , without a heavy modification . Besides , our approach is structure agnostic , which can be universally applied to other structures , only if multiple different layers exist . Instance-wise Learning Deep learning models are trained over large-scale datasets , and data samples are often treated equally without modeling the difference between them . Some works attempt to weight each data with different importance ( Ren et al. , 2018 ; Hu et al. , 2019 ; Chang et al. , 2017 ) or feed data with curriculum learning according to its difficulty ( Bengio et al. , 2009 ; Fan et al. , 2018 ) . However , they often explicitly manipulate the data during training only , while no distinction exists in inference , and under one fixed model . Elbayad et al . ( 2020 ) take a step further and propose the depth-adaptive Transformer , which can forecast different depths of the network by predicting the required computation for a particular data . Similarly , Liu et al . ( 2020 ) propose a sample-wise adaptive mechanism to dynamically calculate the number of required layers . They both aim at reducing the computation cost and speed up the inference . Schwartz et al . ( 2020 ) , Bapna et al . ( 2020 ) and Shazeer et al . ( 2017 ) all leverage conditional computation for each sample to control the computation and accuracy tradeoff during inference . Instead , we pay attention to the variant modeling functions and perform instance-wise order selection in order to boost the Transformer performance . The most related work is Press et al . ( 2019 ) , which manually generates randomly ordered Transformer encoders and finds the Sandwich Transformer can slightly reduce the perplexity of language modeling . However , they find that Sandwich Transformer pattern has no effect on NMT task . Besides , it still performs over the whole corpus without considering each specific data . We , instead , investigate on various sequence-to-sequence generation tasks and greatly improve the task performances through instance-wise learning , so as to discover the optimal ordered Transformer for each particular data . 3 INSTANCE-WISE ORDERED TRANSFORMER . The overall framework of IOT is presented in Figure 1 . In comparison with the standard Transformer , IOT only incorporates light-weighted predictors and reorders the encoder/decoder with weight tying , under the constraint of almost same number of parameters and exempt from heavy modifications . In this section , we introduce the details of IOT , including training , inference and discussions . Notations Sequence-to-sequence learning aims to map one sequence x = [ x1 , x2 , ... , xTx ] into another sequence y = [ y1 , y2 , ... , yTy ] , where xi , yj denotes the i-th and j-th token of x and y , Tx and Ty are the corresponding lengths . Given one sentence pair ( x , y ) and a learning modelM , we can define the training objective as minimizing the cross-entropy loss LM = − ∑Ty j=1 logP ( yj |y < j , x ) . Besides , DKL ( P‖Q ) denotes the Kullback-Leibler ( KL ) divergence between distributions P and Q . 3.1 INSTANCE-WISE ENCODER/DECODER . IOT intends to break the fixed order of layers in Transformer . As shown in introduction , simply reordering the layers w.r.t the whole corpus impacts little , while each data has its own preference to orders . Therefore , IOT incorporates instance-wise learning to adjust the favorable order for each data . As shown in Figure 1 , both encoder and decoder in IOT consist of several blocks of SA , ED , FF layer with dynamic order , and we assume there are M ( e.g. , M = 2 ) ordered encoders and N ( e.g. , N = 6 ) ordered decoders ( with shared weights ) . Inspired by the fact that lower training loss implies the higher proficiency confidence for candidate orders , we utilize the cross-entropy loss as signals to learn the confidence . That is , we calculate confidence γm and λn for each encoder encm , decoder decn ( resulted modelMm , n ) , and use them to weight the training loss LMm , n . To calculate the confidence , we add a simple and light predictor to help distinguish the orders . Training Given one source sequence x = [ x1 , x2 , ... , xTx ] , we first map each token into word embedding e = [ e1 , e2 , ... , eTx ] , where ei ∈ Rd , and then apply one light encoder predictor πenc to predict the confidence of encoder orders using sentence embedding se = 1Tx ∑Tx i=1 ei . Concretely , πenc takes se as input and predicts γm for encm by Gumbel-softmax ( Jang et al. , 2016 ) : γm = exp ( ( log ( πencm ) + gm ) /τe ) ∑M k=1 exp ( ( log ( πenck ) + gk ) /τe ) , πenc = softmax ( seWe ) , ( 1 ) where gm is sampled from Gumbel distribution : gm = − log ( − logUm ) , Um ∼ Uniform ( 0 , 1 ) , We ∈ Rd×M is the weight matrix , τe is a constant temperature to control the distribution to be identical approximation with categorical distribution . Simultaneously , the token embeddings e will feed to the encoders to get hidden states h = [ h1 , h2 , ... , hTx ] , then we can calculate decoder order confidence λn by one predictor πdec in the same way as πenc : λn = exp ( ( log ( πdecn ) + gn ) /τd ) ∑N k=1 exp ( ( log ( πdeck ) + gk ) /τd ) , πdec = softmax ( sdWd ) , ( 2 ) where sd = 1Tx ∑Tx i=1 hi and Wd is the weight matrix . For each ordered path through encm and decn , we can obtain the training loss LMm , n , and the final cross-entropy loss is weighted by confidence γm and λn with LMm , n , formulately as : LC = M∑ m=1 N∑ n=1 ( γm · λn ) LMm , n . ( 3 ) Inference During inference , we directly replace the Gumbel-softmax used in training with argmax , in order to choose the most capable encoder and decoder for each sequence x : enc = argmax ( seWe ) , dec = argmax ( sdWd ) . ( 4 ) Discussion The decoding process is almost the same as standard Transformer , with only little overhead for order predictions . One may concern the training cost is increased through our training . As we present in Section 5.1 , the cost is actually affordable with a fast convergence . Currently , we reorder the layers of the encoder/decoder block and stack the same ordered block L times ( see Figure 1 ) . A complex extension is to reorder all L blocks of encoder/decoder and we take it as future work .

Summary: This work explores instance-wise layer re-ordering in transformers. The key idea is to incorporate classifiers that predict the ordering of sub-layers (self-attention, cross-attention, feed-forward) from the averaged input sequence representation, one classifier each for the encoder and the decoder. During training the model uses a soft Gumbel-noised output of the classifier to combine the outputs from stacks with differently ordered sub-layers. During inference the argmax of the classifier prediction is used to generate the output sequence. The model is trained with two auxiliary losses: (i) A loss to ensure the expected output of the classifiers is uniform and (ii) A loss to ensure the classifier output for each individual sample is distant from uniform.