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: This work proposes a way to reduce the latencies incurred in inference for neural machine translation. Basic idea is to train a model with softmax attached to each output of decoder layers, and computes a loss by aggregating the cross entropy losses over the softmaxes. During inference, it could either use one of the softmax or train an additional model which dynamically selects softmaxes given an input string. Experimental results show that it is possible to reduce latencies by trading off the translation qualities measured by BLEU. Dynamically selection did not show any gains in latencies, though, this work empirically shows potential gains in oracle studies. This work further shows that the model could be compressed further by knowledge distillation.