Balancing Cost and Benefit with Tied-Multi Transformers

1 INTRODUCTION . Neural networks for sequence-to-sequence modeling typically consist of an encoder and a decoder coupled via an attention mechanism . Whereas the very first deep models used stacked recurrent neural networks ( RNN ) ( Sutskever et al. , 2014 ; Cho et al. , 2014 ; Bahdanau et al. , 2015 ) in the encoder and decoder , the recent Transformer model ( Vaswani et al. , 2017 ) constitutes the current state-of-the-art approach , owing to its better context generation mechanism via multi-head self- and cross-attentions . Given an encoder-decoder architecture and its hyper-parameters , such as the number of encoder and decoder layers , vocabulary sizes ( in the case of text based models ) and hidden layers , the parameters of the model , i.e. , matrices and biases for non-linear transformations , are optimized by iteratively updating them so that the loss for the training data is minimized . The hyper-parameters can also be tuned , for instance , through maximizing the automatic evaluation score on the development data . However , in general , it is highly unlikely ( or impossible ) that a single optimized model suffices diverse cost-benefit demands at the same time . For instance , in practical low-latency scenarios , one may accept some performance drop for speed . However , a model used with a subset of optimized parameters might perform badly . A single optimized model can not also guarantee the best performance for each individual input . Although this is practically important , it has drown only a little attention . An existing solution for this problem is to host multiple models simultaneously for flexible choice . However , this approach is not very practical , because it requires an unreasonably large quantity of resources . Furthermore , there are no well-established methods for selecting appropriate models for each individual input . As a more effective solution , we consider training a single model that subsumes multiple models which can be used for decoding with different hyper-parameter settings depending on the input or on the latency requirements . In this paper , we focus on the number of layers as an important hyperparameter that impacts both speed and quality of decoding , and propose a multi-layer softmaxing method , which trains multi-layer neural models referring to the output of all layers during training . Conceptually , as illustrated in Figure 1 , this method involves tying ( sharing ) the parameters of mul- tiple models with different number of layers and is not specific to particular types of multi-layer neural models . On top of the above method , we consider to exploit the model further . For saving decoding time , we design and evaluate mechanisms to choose appropriate number of layers depending on the input . As for further model compression , we leverage other orthogonal types of parameter tying approaches , such as those reviewed in Section 2 . Despite the generality of our proposed method , in this paper , we particularly focus on encoderdecoder models with N encoder and M decoder layers , and compress N ×M models1 by updating the model with a total ofN×M losses computed by softmaxing the output of each of theM decoder layers , where it attends to the output of each of the N encoder layers . The number of parameters of the resultant encoder-decoder model is equivalent to that of the most complex subsumed model with N encoder and M decoder layers . Yet , we can now perform faster decoding using a fewer number of layers , given that shallower layers are better trained . To evaluate our proposed method , we take the case study of neural machine translation ( NMT ) ( Cho et al. , 2014 ; Bahdanau et al. , 2015 ) , where we focus on the numbers of encoder and decoder layers of the Transformer model ( Vaswani et al. , 2017 ) , and demonstrate that a single model with N encoder and M decoder layers trained by our method can be used for flexibly decoding with fewer than N and M layers without appreciable quality loss . We evaluate our proposed method on WMT18 English-to-German translation task , and give a cost-benefit analysis for translation quality vs. decoding speed . The rest of the paper is organized as follows . Section 2 briefly reviews related work for compressing neural models . Section 3 covers our method that ties multiple models by softmaxing all encoderdecoder layer combinations . Section 4 describes our efforts towards designing and evaluating a mechanism for dynamically selecting encoder-decoder layer combinations prior to decoding . Section 5 describes two orthogonal extensions to our model aiming at further model compression and speeding-up of decoding . The paper ends with Section 6 containing conclusion and future work . 2 RELATED WORK . There are studies that exploit multiple layers simultaneously . Wang et al . ( 2018 ) fused hidden representations of multiple layers in order to improve the translation quality . Belinkov et al . ( 2017 ) and Dou et al . ( 2018 ) attempted to identify which layer can generate useful representations for different natural language processing tasks . Unlike them , we make all layers of the encoder and decoder usable for decoding with any encoder-decoder layer combination . In practical scenarios , we can save significant amounts of time by choosing shallower encoder and decoder layers for inference . Our method ties the parameters of multiple models , which is orthogonal to the work that ties parameters between layers ( Dabre & Fujita , 2019 ) and/or between the encoder and decoder within a 1Rather than casting the encoder-decoder model into a single column model with ( N +M ) layers . Algorithm 1 Training a tied-multi model 1 : enc0 = X 2 : for i in 1 to N do 3 : enci = Lenci ( enci−1 ) 4 : for j in 1 to M do 5 : decj = Ldecj ( decj−1 , enci ) 6 : Ŷ = softmax ( decj ) 7 : lossi , j = cross entropy ( Ŷ , Y ) 8 : end for 9 : end for 10 : overall loss = aggregate ( loss1,1 , . . . , lossN , M ) 11 : Back-propagate using overall loss single model ( Xia et al. , 2019 ; Dabre & Fujita , 2019 ) . Parameter tying leads to compact models enabling faster decoding , but they usually suffer from drops in inference quality . In this paper , we counter such drops with knowledge distillation ( Hinton et al. , 2015 ; Kim & Rush , 2016 ; Freitag et al. , 2017 ) . The working principle of this approach is to utilize smoothed data or smoothed training signals instead of the actual training data . A model with a large number of parameters and high performance provides smoothed distributions that are then used as labels for training small models instead of one-hot vectors . As one of the aims in this work is model size reduction , it is related to a growing body of work that addresses the computational requirement reduction . Pruning of pre-trained models ( See et al. , 2016 ) makes it possible to discard around 80 % of the smallest weights of a model without deterioration in inference quality , given it is re-trained with appropriate hyper-parameters after pruning . Currently , most deep learning implementations use 32-bit floating point representations , but 16-bit floating point representations ( Gupta et al. , 2015 ; Ott et al. , 2018 ) or aggressive binarization ( Courbariaux et al. , 2017 ) can be alternatives . Compact models are usually faster to decode ; studies on quantization ( Lin et al. , 2016 ) and average attention networks ( Xiong et al. , 2018 ) address this topic . To the best of our knowledge , none of the above work has attempted to combine multi-model parameter tying , knowledge distillation , and dynamic layer selection for obtaining and exploiting highlycompressed and flexible deep neural models . 3 MULTI-LAYER SOFTMAXING . 3.1 PROPOSED METHOD . Figure 1 gives a simple overview of the concept of multi-layer softmaxing , taking a generic model as an example . The rightmost 4-layer model takes an input and passes it through 4 layers2 before a softmax layer to predict the output . Typically , one would apply softmax to the 4th layer only , compute loss , and then back-propagate gradients in order to update parameters . Instead , we propose to apply softmax to each layer , aggregate the computed losses , and then back-propagate . This enables us to choose any layer combination during decoding instead of only the topmost layer . Extending this to a multi-layer encoder-decoder model is straightforward . In encoder-decoder models , the encoder comprises an embedding layer for the input ( source language for NMT ) and N stacked transformation layers . The decoder consists of an embedding layer and a softmax layer for generating the output ( target language for NMT ) along with M stacked transformation layers . Let X be the input to the N -layer encoder , Y the anticipated output of the M -layer decoder as well as the input to the decoder ( for training ) , and Ŷ the predicted output by the decoder . Algorithm 1 shows the pseudo-code for our proposed method . The line 3 represents the process done by the i-th encoder layer , Lenci , and the line 5 does the same for the j-th decoder layer , L dec j . In simple words , we compute a loss using the output of each of theM decoder layers which in turn is computed using 2We make no assumptions about the nature of the layers . the output of each of the N encoder layers . In line 10 , the N ×M losses are aggregated3 before back-propagation . Henceforth , we will refer to this as the Tied-Multi model . For a comparison , the vanilla model is formulated as follows : decj = Ldecj ( decj−1 , encN ) , Ŷ = softmax ( decM ) , and overall loss = cross entropy ( Ŷ , Y ) . 3.2 EXPERIMENTAL SETUP . We evaluated the following two types of models on both translation quality and decoding speed . Vanilla models : 36 vanilla models with 1 to 6 encoder and 1 to 6 decoder layers , each trained referring only to the last layer for computing loss . Tied-Multi model : A single tied-multi model with N = 6 encoder and M = 6 decoder layers , trained by our multi-layer softmaxing . We experimented with the WMT18 English-to-German ( En→De ) translation task . We used all the parallel corpora available for WMT18 , except ParaCrawl corpus,4 consisting of 5.58M sentence pairs as the training data and 2,998 sentences in newstest2018 as test data . The English and German sentences were pre-processed using the tokenizer.perl and truecase.perl scripts in Moses.5 The true-case models for English and German were trained on 10M sentences randomly extracted from the monolingual data made available for the WMT18 translation task , using the train-truecaser.perl script available in Moses . Our multi-layer softmaxing method was implemented on top of an open-source toolkit of the Transformer model ( Vaswani et al. , 2017 ) in the version 1.6 branch of tensor2tensor.6 For training , we used the default model settings corresponding to transformer base single gpu in the implementation , except what follows . We used a shared sub-word vocabulary of 32k determined using the internal sub-word segmenter of tensor2tensor and trained the models for 300k iterations . We trained the vanilla models on 1 GPU and our tied-multi model on 2 GPUs with batch size halved to ensure that both models see the same amount of training data . We averaged the last 10 checkpoints saved every after 1k updates , and decoded the test sentences , fixing a beam size7 of 4 and length penalty of 0.6 . We evaluated our models using the BLEU metric ( Papineni et al. , 2002 ) implemented in sacreBLEU ( Post , 2018 ) .8 Prior to evaluation , the decoded test set was post-processed using the detokenizer.perl and detruecase.perl scripts in Moses . We also report the time ( in seconds ) consumed to translate the test set , which includes times for the model instantiation , loading the checkpoint , sub-word splitting and indexing , decoding , and sub-word de-indexing and merging , whereas times for detokenization are not taken into account . Note that we did not use any development data for two reasons . First , we train all models for the same number of iterations.9 Second , we use checkpoint averaging before decoding , which does not require a development set unlike early stopping . We adopted this approach , because it is known to give the best results for NMT using the Transformer implementation we use ( Vaswani et al. , 2017 ) . 3We averaged multiple losses in our experiment , but there are a number of options , such as weighted averaging . 4http : //www.statmt.org/wmt18/translation-task.html We excluded ParaCrawl following the instruction on the WMT18 website : “ BLEU score dropped by 1.0 ” for this task . 5https : //github.com/moses-smt/mosesdecoder 6https : //github.com/tensorflow/tensor2tensor 7One can realize faster decoding by narrowing down the beam width . This approach is orthogonal to ours and in this paper we do not insist which is superior to the other . 8https : //github.com/mjpost/sacreBLEU BLEU+case.mixed+lang.en-de+numrefs.1+smooth.exp +test.wmt18+tok.13a+version.1.3.7 9This enables a fair comparison , because it ensures that each model sees roughly the same number of training examples .

Summary: In this paper, the authors propose a new variant of Transformer called Tied-multi Transformer. Given such a model with an N-layer encoder and an M-layer decoder, it is trained with M*N loss functions, where each combination of the nth-layer of the encoder and the mth-layer of the decoder is used to train an NMT model. The authors propose a way to dynamically select which layers to be used when a specific sentence comes.  At last, the authors also try recurrent stack and knowledge to further compress the models.