Balancing Cost and Benefit with Tied-Multi Transformers (2020)

Abstract: We propose a novel procedure for training multiple Transformers with tied parameters which compresses multiple models into one enabling the dynamic choice of the number of encoder and decoder layers during decoding. In training an encoder-decoder model, typically, the output of the last layer of the N-layer encoder is fed to the M-layer decoder, and the output of the last decoder layer is used to compute loss. Instead, our method computes a single loss consisting of NxM losses, where each loss is computed from the output of one of the M decoder layers connected to one of the N encoder layers. Such a model subsumes NxM models with different number of encoder and decoder layers, and can be used for decoding with fewer than the maximum number of encoder and decoder layers. Given our flexible tied model, we also address to a-priori selection of the number of encoder and decoder layers for faster decoding, and explore recurrent stacking of layers and knowledge distillation for model compression. We present a cost-benefit analysis of applying the proposed approaches for neural machine translation and show that they reduce decoding costs while preserving translation quality.

## 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) @xcite @xcite in the encoder and decoder, the recent transformer model @xcite constitutes the current state-of-the-art approach, owing to its better context modeling 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 models for texts), 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. also, a single optimized model cannot guarantee the best performance for each individual input. an existing solution for these problems is to train multiple models and host them simultaneously. however, this approach is not very practical, because it requires a large number of resources. we also lack a well-established method for selecting appropriate models for each individual input prior to decoding. 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 uses the outputs of all layers during training. conceptually, as illustrated in figure 1 , this is the same as tying (sharing) the parameters of multiple models with different number of layers and is not specific to particular types of multi-layer neural models. despite the generality of our proposed method, in this paper, we focus on encoder-decoder models with n encoder and m decoder layers, and figure 1 : the general concept of multi-layer softmaxing for training multi-layer neural models with an example of a 4-layer model. figure 1a is a depiction of our idea in the form of multiple vanilla models whose layers are tied together. figure 1b shows the result of collapsing all tied layers into a single layer. the red lines indicate the flow of gradients and hence the shallowest layer in the stack receives the largest number of updates. to evaluate our proposed method, we take the case study of neural machine translation (nmt) @xcite , using the transformer model @xcite , 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 englishto-german translation task, and give a cost-benefit analysis for translation quality vs. decoding speed. given a flexible tied model, for saving decoding time, we then design mechanisms to choose, prior to decoding, the appropriate number of encoder and decoder layers depending on the input. we also focus on compact modeling, where we leverage other orthogonal types of parameter tying approaches. compact models are faster to decode and will be useful in cases where a-priori layer prediction might be infeasible. the rest of the paper is organized as follows. section 2 briefly reviews related work for compressing neural models. section 3 covers our 1 rather than casting the encoder-decoder model into a single column model with (n + m ) layers. ## related work there are studies that exploit multiple layers simultaneously. @xcite fused hidden representations of multiple layers in order to improve the translation quality. @xcite and @xcite 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 @xcite and/or between the encoder and decoder within a single model @xcite . parameter tying leads to compact models, but they usually suffer from drops in inference quality. in this paper, we counter such drops with knowledge distillation @xcite @xcite . this approach utilizes smoothed data or smoothed training signals instead of the actual training data. a model with a large number of parameters and high per-formance 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 @xcite 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 @xcite or aggressive binarization @xcite can be alternatives. compact models are usually faster to decode; studies on quantization @xcite and average attention networks @xcite address this topic. none of the above work has attempted to combine multi-model parameter tying, knowledge distillation, and dynamic layer selection for obtaining and exploiting highly-compressed and flexible deep neural models. ## dynamic layer selection motivated by the results shown in figure 2 , we tackled an advanced problem: dynamic selection of one layer combination prior to decoding. foot_4 ## further model compression we examined the combination of our multi-layer softmaxing approach with another parameter-tying method in neural networks, called recurrent stacking (rs) @xcite , complemented by sequence-level knowledge distillation @xcite , a specific type of knowledge distillation @xcite . we demonstrate that these existing techniques help reduce the number of parameters in our model even further. ## conclusion in this paper, we have proposed a novel procedure for training encoder-decoder models, where the softmax function is applied to the output of each of the m decoder layers derived using the output of each of the n encoder layers. this compresses n ×m models into a single model that can be used for decoding with a variable number of encoder (1 ≤ @xmath0. this model can be used in different latency scenarios and hence is highly versatile. we have made a cost-benefit analysis of our method, taking nmt as a case study of encoder-decoder models. we have proposed and evaluated two orthogonal extensions and show that we can (a) dynamically choose layer combinations for slightly faster decoding and (b) further compress models using recurrent stack-ing with knowledge distillation leading to models that also enable faster decoding. for further speed up in decoding as well as model compression, we plan to combine our approach with other techniques, such as those mentioned in section 2. although we have only tested our idea for nmt, it should be applicable to other tasks based on deep neural networks.