Shallow-to-Deep Training for Neural Machine Translation (2020)

Abstract: Deep encoders have been proven to be effective in improving neural machine translation (NMT) systems, but training an extremely deep encoder is time consuming. Moreover, why deep models help NMT is an open question. In this paper, we investigate the behavior of a well-tuned deep Transformer system. We find that stacking layers is helpful in improving the representation ability of NMT models and adjacent layers perform similarly. This inspires us to develop a shallow-to-deep training method that learns deep models by stacking shallow models. In this way, we successfully train a Transformer system with a 54-layer encoder. Experimental results on WMT’16 English-German and WMT’14 English-French translation tasks show that it is 1:4 faster than training from scratch, and achieves a BLEU score of 30:33 and 43:29 on two tasks. The code is publicly available at https://github.com/libeineu/SDT-Training .

## introduction in recent years, neural models have led to state-ofthe-art results in machine translation (mt) @xcite . many of these systems can broadly be characterized as following a multi-layer encoder-decoder neural network design: both the encoder and decoder learn representations of word sequences by a stack of layers @xcite @xcite , building on an interesting line of work in improving such models. the simplest of these increases the model capacity by widening the network, whereas more recent work shows benefits from stacking more layers on the encoder side. for example, for the popular transformer model @xcite , deep systems have shown promising bleu improvements by either easing the information flow through the network @xcite or constraining the gradient norm across layers @xcite @xcite ). an improved system can even learn a 35-layer encoder, which is 5× deeper than that of vanilla transformer @xcite . although these methods have enabled training deep neural mt (nmt) models, questions remain as to the nature of the problem. the main question here is: why and how deep networks help in nmt. note that previous work evaluates these systems in a black-box manner (i.e., bleu score). it is thus natural to study how much a deep nmt system is able to learn that is different from the shallow counterpart. beyond this, training an extremely deep model is expensive although a narrow-anddeep network can speed up training @xcite . for example, it takes us 3× longer time to train the model when we deepen the network from 6 layers to 48 layers. this might prevent us from exploiting deeper models in large-scale systems. in this paper, we explore why deep architectures work to render learning nmt models more effectively. by investigating the change of the hidden states in different layers, we find that new representations are learned by continually stacking layers on top of the base model. more stacked layers lead to a stronger model of representing the sentence. this particularly makes sense in the deep nmt scenario because it has been proven that deep models can benefit from an enriched representation @xcite @xcite . in addition, the finding here inspires us to develop a simple yet efficient method to train a deep nmt encoder: we train model parameters from shallow to deep, rather than training the entire model from scratch. to stabilize training, we design a sparse linear combination method of connecting lower-level layers to the top. it makes efficient pass of information through the deep network but does not require large memory footprint as in dense networks. we experiment with the method in a stateof-the-art deep transformer system. our encoder consists of 48-54 layers, which is almost the deepest transformer model used in nmt. on wmt en-de and en-fr tasks, it yields a 1.4× speedup of training, matching the state-of-the-art on the wmt'16 en-de task. ## background we start with a description of deep transformer. in transformer @xcite , the encoder takes a sequence of words {x 1 , ..., x n } as input. the input is first transformed into a sequence of embeddings {w 1 + p 1 , ..., w n + p n }, where w k is a word embedding and p k is a positional embedding. then, the embedding sequence is fed into a stack of n identical layers. each layer consists of two stacked sub-layers: a multi-head self-attention sublayer and a feed-forward sub-layer. the decoder shares a similar architecture as the encoder but possesses an encoder-decoder attention sub-layer to capture the mapping between two languages. for a deep model, layer normalization networks and layer-wise connections are needed, following the previous work of @xcite and @xcite . ## why do deep models help? the transformer encoder (or decoder) is essentially a representation model @xcite . given a word sequence, a layer generates a distributed representation (i.e., y j ) for each position of the sequence. the representation is a mixture of word+position embedding and context embedding. for a simple implementation, only the top-most representation (i.e., y n ) is used for downstream components of the system. nevertheless, the dense connections can make the lower-level representations directly accessible to the top layers. hence, the representation model is actually encoded by the set of layer outputs {y 1 , ..., y n }. for a stronger model, enlarging the size of each layer can fit the objective function with enough capacity. for example, transformer-big doubles the layer size of the base model and shows consistent improvements on several mt tasks. but the number of parameters increases quadratically with network width, which poses new difficulties in training such systems and the risk of overfitting. alternatively, one can stack layers to strengthen the model because more layers offer more representations for the input sentence. moreover, top-level layers can generate refined representations @xcite by passing the input vectors through more linear and non-linear transformations in different layers. to study how each layer behaves, we evaluate the change of the representation vector in the stack. to do this, we compute the similarity between the outputs of two layers, as below, figure 2 (a) plots sim(i, 0) curves for wmt en-de systems of different encoder depths. here sim(i, 0) measures how similar the output of layer j is to the input embedding of the encoder. we see that the similarity keeps going down when we stack more layers. it indicates that the model can learn new representations by using more stacked layers. but the similarity does not converge for shallow models (e.g., 6-layer and 12-layer encoders). it somehow reflects the fact that the shallow models "want" more layers to learn something new. more interestingly, deeper models (e.g., encoders of 18 layers or more) make the similarity converge as the depth increases, showing the effect of rendering the need of representation learning fulfilled. also, we investigate the similarity between adjacent layers for different systems. figure 2(b ) plots sim(i, i -1) as a function of layer number, which begins with sim(2, 1) rather than sim(1, 0) foot_0 for better visualization. the results show that adjacent layers in converged systems have a high similarity. this agrees with previous work on the similarity of attention weights among layers @xcite though we study a different issue here. the deeper the models, the higher the similarity between adjacent layers. a natural question is whether we can initialize the higher layers by reusing the parameters of previous layers during training procedure? note that the transformer model is doing something like encoding both contextual information and word information for each position of the sequence. here, we design the inter-sim over the sequence to see how much the representation of a position holds to encode the sequence. for layer i, we have ## shallow-to-deep training although we are able to train deep models with the standard methods @xcite , it is obviously a time consuming task to learn a model with too many layers. for example, training a 48-layer system takes us 3× longer time than the 6-layer baseline. as stated in section 3, adjacent layers in a deep network are likely to behave in a similar fashion. this observation inspires us to train upperlevel layers by reusing the learned parameters of lower-level layers. we call this method shallow-todeep training (sdt) because we start with training a shallow model and then train a deeper model on top of it. ## experiments we report the experimental results on two widely used benchmarks -wmt'16 english-german (en-de) and wmt'14 english-french (en-fr). ## related work in this section, we discuss the related work from two aspects as follows: ## conclusions we have investigated the behaviour of the welltrained deep transformer models and found that stacking more layers could improve the representation ability of nmt systems. higher layers share more global information over different positions and adjacent layers behave similarly. also, we have developed a shallow-to-deep training strategy and employ sparse connections across blocks to ease the optimization. with the help of learning rate restart and appropriate initialization we successfully train a 48-layer rpr model by progressive stacking and achieve a 40% speedup on both wmt'16 english-german and wmt'14 english-french tasks. furthermore, our sdt-rpr-24l (big) achieves a bleu score of 30.46 on wmt'16 english-german task, and speeds up the training by 1.5×.