Unfolding and Shrinking Neural Machine Translation Ensembles

Abstract: Ensembling is a well-known technique in neural machine translation (NMT) to improve system performance. Instead of a single neural net, multiple neural nets with the same topology are trained separately, and the decoder generates predictions by averaging over the individual models. Ensembling often improves the quality of the generated translations drastically. However, it is not suitable for production systems because it is cumbersome and slow. This work aims to reduce the runtime to be on par with a single system without compromising the translation quality. First, we show that the ensemble can be unfolded into a single large neural network which imitates the output of the ensemble system. We show that unfolding can already improve the runtime in practice since more work can be done on the GPU. We proceed by describing a set of techniques to shrink the unfolded network by reducing the dimensionality of layers. On Japanese-English we report that the resulting network has the size and decoding speed of a single NMT network but performs on the level of a 3-ensemble system.

The top systems in recent machine translation evaluation campaigns on various language pairs use ensembles of a number of NMT systems BIBREF0 , BIBREF1 , BIBREF2 , BIBREF3 , BIBREF4 , BIBREF5 , BIBREF6 . Ensembling BIBREF7 , BIBREF8 of neural networks is a simple yet very effective technique to improve the accuracy of NMT. The decoder makes use of INLINEFORM0 NMT networks which are either trained independently BIBREF9 , BIBREF2 , BIBREF3 , BIBREF4 or share some amount of training iterations BIBREF10 , BIBREF1 , BIBREF5 , BIBREF6 . The ensemble decoder computes predictions from each of the individual models which are then combined using the arithmetic average BIBREF9 or the geometric average BIBREF5 .
Ensembling consistently outperforms single NMT by a large margin. However, the decoding speed is significantly worse since the decoder needs to apply INLINEFORM0 NMT models rather than only one. Therefore, a recent line of research transfers the idea of knowledge distillation BIBREF11 , BIBREF12 to NMT and trains a smaller network (the student) by minimizing the cross-entropy to the output of the ensemble system (the teacher) BIBREF13 , BIBREF14 . This paper presents an alternative to knowledge distillation as we aim to speed up decoding to be comparable to single NMT while retaining the boost in translation accuracy from the ensemble. In a first step, we describe how to construct a single large neural network which imitates the output of an ensemble of multiple networks with the same topology. We will refer to this process as unfolding. GPU-based decoding with the unfolded network is often much faster than ensemble decoding since more work can be done on the GPU. In a second step, we explore methods to reduce the size of the unfolded network. This idea is justified by the fact that ensembled neural networks are often over-parameterized and have a large degree of redundancy BIBREF15 , BIBREF16 , BIBREF17 . Shrinking the unfolded network leads to a smaller model which consumes less space on the disk and in the memory; a crucial factor on mobile devices. More importantly, the decoding speed on all platforms benefits greatly from the reduced number of neurons. We find that the dimensionality of linear embedding layers in the NMT network can be reduced heavily by low-rank matrix approximation based on singular value decomposition (SVD). This suggest that high dimensional embedding layers may be needed for training, but do not play an important role for decoding. The NMT network, however, also consists of complex layers like gated recurrent units BIBREF18 and attention BIBREF19 . Therefore, we introduce a novel algorithm based on linear combinations of neurons which can be applied either during training (data-bound) or directly on the weight matrices without using training data (data-free). We report that with a mix of the presented shrinking methods we are able to reduce the size of the unfolded network to the size of the single NMT network while keeping the boost in BLEU score from the ensemble. Depending on the aggressiveness of shrinking, we report either a gain of 2.2 BLEU at the same decoding speed, or a 3.4 INLINEFORM1 CPU decoding speed up with only a minor drop in BLEU compared to the original single NMT system. Furthermore, it is often much easier to stage a single NMT system than an ensemble in a commercial MT workflow, and it is crucial to be able to optimize quality at specific speed and memory constraints. Unfolding and shrinking address these problems directly.
The first concept of our approach is called unfolding. Unfolding is an alternative to ensembling of multiple neural networks with the same topology. Rather than averaging their predictions, unfolding constructs a single large neural net out of the individual models which has the same number of input and output neurons but larger inner layers. Our main motivation for unfolding is to obtain a single network with ensemble level performance which can be shrunk with the techniques in Sec. SECREF3 .
Suppose we ensemble two single layer feedforward neural nets as shown in Fig. FIGREF1 . Normally, ensembling is implemented by performing an isolated forward pass through the first network (Fig. SECREF2 ), another isolated forward pass through the second network (Fig. SECREF3 ), and averaging the activities in the output layers of both networks. This can be simulated by merging both networks into a single large network as shown in Fig. SECREF4 . The first neurons in the hidden layer of the combined network correspond to the hidden layer in the first single network, and the others to the hidden layer of the second network. A single pass through the combined network yields the same output as the ensemble if the output layer is linear (up to a factor 2). The weight matrices in the unfolded network can be constructed by stacking the corresponding weight matrices (either horizontally or vertically) in network 1 and 2. This kind of aggregation of multiple networks with the same topology is not only possible for single-layer feedforward architectures but also for complex networks consisting of multiple GRU layers and attention.
For a formal description of unfolding we address layers with indices INLINEFORM0 . The special layer 0 has a single neuron for modelling bias vectors. Layer 1 holds the input neurons and layer INLINEFORM1 is the output layer. We denote the size of a layer in the individual models as INLINEFORM2 . When combining INLINEFORM3 networks, the layer size INLINEFORM4 in the unfolded network is increased by factor INLINEFORM5 if INLINEFORM6 is an inner layer, and equal to INLINEFORM7 if INLINEFORM8 is the input or output layer. We denote the weight matrix between two layers INLINEFORM9 in the INLINEFORM10 -th individual model ( INLINEFORM11 ) as INLINEFORM12 , and the corresponding weight matrix in the unfolded network as INLINEFORM13 . We explicitly allow INLINEFORM14 and INLINEFORM15 to be non-consecutive or reversed to be able to model recurrent networks. We use the zero-matrix if layers INLINEFORM16 and INLINEFORM17 are not connected. The construction of the unfolded weight matrix INLINEFORM18 from the individual matrices INLINEFORM19 depends on whether the connected layers are inner layers or not. The complete formula is listed in Fig. FIGREF5 .
Unfolded NMT networks approximate but do not exactly match the output of the ensemble due to two reasons. First, the unfolded network synchronizes the attentions of the individual models. Each decoding step in the unfolded network computes a single attention weight vector. In contrast, ensemble decoding would compute one attention weight vector for each of the INLINEFORM0 input models. A second difference is that the ensemble decoder first applies the softmax at the output layer, and then averages the prediction probabilities. The unfolded network averages the neuron activities (i.e. the logits) first, and then applies the softmax function. Interestingly, as shown in Sec. SECREF4 , these differences do not have any impact on the BLEU score but yield potential speed advantages of unfolding since the computationally expensive softmax layer is only applied once.
After constructing the weight matrices of the unfolded network we reduce the size of it by iteratively shrinking layer sizes. In this section we denote the incoming weight matrix of the layer to shrink as INLINEFORM0 and the outgoing weight matrix as INLINEFORM1 . Our procedure is inspired by the method of Srinivas and Babu sparsify-datafree. They propose a criterion for removing neurons in inner layers of the network based on two intuitions. First, similarly to Hebb's learning rule, they detect redundancy by the principle neurons which fire together, wire together. If the incoming weight vectors INLINEFORM2 and INLINEFORM3 are exactly the same for two neurons INLINEFORM4 and INLINEFORM5 , we can remove the neuron INLINEFORM6 and add its outgoing connections to neuron INLINEFORM7 ( INLINEFORM8 ) without changing the output. This holds since the activity in neuron INLINEFORM14 will always be equal to the activity in neuron INLINEFORM15 . In practice, Srinivas and Babu use a distance measure based on the difference of the incoming weight vectors to search for similar neurons as exact matches are very rare.
The second intuition of the criterion used by Srinivas and Babu sparsify-datafree is that neurons with small outgoing weights contribute very little overall. Therefore, they search for a pair of neurons INLINEFORM0 according the following term and remove the INLINEFORM1 -th neuron. DISPLAYFORM0 
Neuron INLINEFORM0 is selected for removal if (1) there is another neuron INLINEFORM1 which has a very similar set of incoming weights and if (2) INLINEFORM2 has a small outgoing weight vector. Their criterion is data-free since it does not require any training data. For further details we refer to Srinivas and Babu sparsify-datafree.
Srinivas and Babu sparsify-datafree propose to add the outgoing weights of INLINEFORM0 to the weights of a similar neuron INLINEFORM1 to compensate for the removal of INLINEFORM2 . However, we have found that this approach does not work well on NMT networks. We propose instead to compensate for the removal of a neuron by a linear combination of the remaining neurons in the layer. Data-free shrinking assumes for the sake of deriving the update rule that the neuron activation function is linear. We now ask the following question: How can we compensate as well as possible for the loss of neuron INLINEFORM3 such that the impact on the output of the whole network is minimized? Data-free shrinking represents the incoming weight vector of neuron INLINEFORM4 ( INLINEFORM5 ) as linear combination of the incoming weight vectors of the other neurons. The linear factors can be found by satisfying the following linear system: DISPLAYFORM0 
where INLINEFORM0 is matrix INLINEFORM1 without the INLINEFORM2 -th column. In practice, we use the method of ordinary least squares to find INLINEFORM3 because the system may be overdetermined. The idea is that if we mix the outputs of all neurons in the layer by the INLINEFORM4 -weights, we get the output of the INLINEFORM5 -th neuron. The row vector INLINEFORM6 contains the contributions of the INLINEFORM7 -th neuron to each of the neurons in the next layer. Rather than using these connections, we approximate their effect by adding some weight to the outgoing connections of the other neurons. How much weight depends on INLINEFORM8 and the outgoing weights INLINEFORM9 . The factor INLINEFORM10 which we need to add to the outgoing connection of the INLINEFORM11 -th neuron to compensate for the loss of the INLINEFORM12 -th neuron on the INLINEFORM13 -th neuron in the next layer is: DISPLAYFORM0 
Therefore, the update rule for INLINEFORM0 is: DISPLAYFORM0 
In the remainder we will refer to this method as data-free shrinking. Note that we recover the update rule of Srinivas and Babu sparsify-datafree by setting INLINEFORM0 to the INLINEFORM1 -th unit vector. Also note that the error introduced by our shrinking method is due to the fact that we ignore the non-linearity, and that the solution for INLINEFORM2 may not be exact. The method is error-free on linear layers as long as the residuals of the least-squares analysis in Eq. EQREF10 are zero.
The terminology of neurons needs some further elaboration for GRU layers which rather consist of update and reset gates and states BIBREF18 . On GRU layers, we treat the states as neurons, i.e. the INLINEFORM0 -th neuron refers to the INLINEFORM1 -th entry in the GRU state vector. Input connections to the gates are included in the incoming weight matrix INLINEFORM2 for estimating INLINEFORM3 in Eq. EQREF10 . Removing neuron INLINEFORM4 in a GRU layer means deleting the INLINEFORM5 -th entry in the states and both gate vectors.
Although we find our data-free approach to be a substantial improvement over the methods of Srinivas and Babu sparsify-datafree on NMT networks, it still leads to a non-negligible decline in BLEU score when applied to recurrent GRU layers. Our data-free method uses the incoming weights to identify similar neurons, i.e. neurons expected to have similar activities. This works well enough for simple layers, but the interdependencies between the states and the gates inside gated layers like GRUs or LSTMs are complex enough that redundancies cannot be found simply by looking for similar weights. In the spirit of Babaeizadeh et al. sparsify-noiseout, our data-bound version records neuron activities during training to estimate INLINEFORM0 . We compensate for the removal of the INLINEFORM1 -th neuron by using a linear combination of the output of remaining neurons with similar activity patterns. In each layer, we prune 40 neurons each 450 training iterations until the target layer size is reached. Let INLINEFORM2 be the matrix which holds the records of neuron activities in the layer since the last removal. For example, for the decoder GRU layer, a batch size of 80, and target sentence lengths of 20, INLINEFORM3 has INLINEFORM4 rows and INLINEFORM5 (the number of neurons in the layer) columns. Similarly to Eq. EQREF10 we find interpolation weights INLINEFORM6 using the method of least squares on the following linear system. DISPLAYFORM0 
The update rule for the outgoing weight matrix is the same as for our data-free method (Eq. EQREF12 ). The key difference between data-free and data-bound shrinking is the way INLINEFORM0 is estimated. Data-free shrinking uses the similarities between incoming weights, and data-bound shrinking uses neuron activities recorded during training. Once we select a neuron to remove, we estimate INLINEFORM1 , compensate for the removal, and proceed with the shrunk network. Both methods are prior to any decoding and result in shrunk parameter files which are then loaded to the decoder. Both methods remove neurons rather than single weights.
The data-bound algorithm runs gradient-based optimization on the unfolded network. We use the AdaGrad BIBREF20 step rule, a small learning rate of 0.0001, and aggressive step clipping at 0.05 to avoid destroying useful weights which were learned in the individual networks prior to the construction of the unfolded network.
Our data-bound algorithm uses a data-bound version of the neuron selection criterion in Eq. EQREF8 which operates on the activity matrix INLINEFORM0 . We search for the pair INLINEFORM1 according the following term and remove neuron INLINEFORM2 . DISPLAYFORM0 
The standard attention-based NMT network architecture BIBREF19 includes three linear layers: the embedding layer in the encoder, and the output and feedback embedding layers in the decoder. We have found that linear layers are particularly easy to shrink using low-rank matrix approximation. As before we denote the incoming weight matrix as INLINEFORM0 and the outgoing weight matrix as INLINEFORM1 . Since the layer is linear, we could directly connect the previous layer with the next layer using the product of both weight matrices INLINEFORM2 . However, INLINEFORM3 may be very large. Therefore, we approximate INLINEFORM4 as a product of two low rank matrices INLINEFORM5 and INLINEFORM6 ( INLINEFORM7 ) where INLINEFORM8 is the desired layer size. A very common way to find such a matrix factorization is using truncated singular value decomposition (SVD). The layer is eventually shrunk by replacing INLINEFORM9 with INLINEFORM10 and INLINEFORM11 with INLINEFORM12 .
The individual NMT systems we use as source for constructing the unfolded networks are trained using AdaDelta BIBREF21 on the Blocks/Theano implementation BIBREF22 , BIBREF23 of the standard attention-based NMT model BIBREF19 with: 1000 dimensional GRU layers BIBREF18 in both the decoder and bidrectional encoder; a single maxout output layer BIBREF24 ; and 620 dimensional embedding layers. We follow Sennrich et al. nmt-bpe and use subword units based on byte pair encoding rather than words as modelling units. Our SGNMT decoder BIBREF25 with a beam size of 12 is used in all experiments. Our primary corpus is the Japanese-English (Ja-En) ASPEC data set BIBREF26 . We select a subset of 500K sentence pairs to train our models as suggested by Neubig et al. sys-neubig-wat15. We report cased BLEU scores calculated with Moses' multi-bleu.pl to be strictly comparable to the evaluation done in the Workshop of Asian Translation (WAT). We also apply our method to the WMT data set for English-German (En-De), using the news-test2014 as a development set, and keeping news-test2015 and news-test2016 as test sets. En-De BLEU scores are computed using mteval-v13a.pl as in the WMT evaluation. We set the vocabulary sizes to 30K for Ja-En and 50K for En-De. We also report the size factor for each model which is the total number of model parameters (sum of all weight matrix sizes) divided by the number of parameters in the original NMT network (86M for Ja-En and 120M for En-De). We choose a widely used, simple ensembling method (prediction averaging) as our baseline. We feel that the prevalence of this method makes it a reasonable baseline for our experiments.
The idea of pruning neural networks to improve the compactness of the models dates back more than 25 years BIBREF15 . The literature is therefore vast BIBREF28 . One line of research aims to remove unimportant network connections. The connections can be selected for deletion based on the second-derivative of the training error with respect to the weight BIBREF15 , BIBREF16 , or by a threshold criterion on its magnitude BIBREF29 . See et al. sparsify-nmt confirmed a high degree of weight redundancy in NMT networks.
In this work we are interested in removing neurons rather than single connections since we strive to shrink the unfolded network such that it resembles the layout of an individual model. We argued in Sec. SECREF4 that removing neurons rather than connections does not only improve the model size but also the memory footprint and decoding speed. As explained in Sec. SECREF9 , our data-free method is an extension of the approach by Srinivas and Babu sparsify-datafree; our extension performs significantly better on NMT networks. Our data-bound method (Sec. SECREF14 ) is inspired by Babaeizadeh et al. sparsify-noiseout as we combine neurons with similar activities during training, but we use linear combinations of multiple neurons to compensate for the loss of a neuron rather than merging pairs of neurons.
Using low rank matrices for neural network compression, particularly approximations via SVD, has been studied widely in the literature BIBREF30 , BIBREF31 , BIBREF32 , BIBREF33 , BIBREF34 . These approaches often use low rank matrices to approximate a full rank weight matrix in the original network. In contrast, we shrink an entire linear layer by applying SVD on the product of the incoming and outgoing weight matrices (Sec. SECREF18 ).
In this paper we mimicked the output of the high performing but cumbersome ensemble by constructing a large unfolded network, and shrank this network afterwards. Another approach, known as knowledge distillation, uses the large model (the teacher) to generate soft training labels for the smaller student network BIBREF11 , BIBREF12 . The student network is trained by minimizing the cross-entropy to the teacher. This idea has been applied to sequence modelling tasks such as machine translation and speech recognition BIBREF35 , BIBREF13 , BIBREF14 . Our approach can be computationally more efficient as the training set does not have to be decoded by the large teacher network.
Junczys-Dowmunt et al. averaging2,averaging1 reported gains from averaging the weight matrices of multiple checkpoints of the same training run. However, our attempts to replicate their approach were not successful. Averaging might work well when the behaviour of corresponding units is similar across networks, but that cannot be guaranteed when networks are trained independently.
We have described a generic method for improving the decoding speed and BLEU score of single system NMT. Our approach involves unfolding an ensemble of multiple systems into a single large neural network and shrinking this network by removing redundant neurons. Our best results on Japanese-English either yield a gain of 2.2 BLEU compared to the original single NMT network at about the same decoding speed, or a INLINEFORM0 CPU decoding speed up with only a minor drop in BLEU.
The current formulation of unfolding works for networks of the same topology as the concatenation of layers is only possible for analogous layers in different networks. Unfolding and shrinking diverse networks could be possible, for example by applying the technique only to the input and output layers or by some other scheme of finding associations between units in different models, but we leave this investigation to future work as models in NMT ensembles in current research usually have the same topology BIBREF0 , BIBREF1 , BIBREF2 , BIBREF3 , BIBREF4 , BIBREF6 .
This work was supported by the U.K. Engineering and Physical Sciences Research Council (EPSRC grant EP/L027623/1).
Data-free and data-bound shrinking can be interpreted as setting the expected difference between network outputs before and after a removal operation to zero under different assumptions.
For simplicity, we focus our probabilistic treatment of shrinking on single layer feedforward networks. Such a network maps an input INLINEFORM0 to an output INLINEFORM1 . The INLINEFORM2 -th output INLINEFORM3 is computed according the following equation DISPLAYFORM0 
where INLINEFORM0 is the incoming weight vector of the INLINEFORM1 -th hidden neuron (denoted as INLINEFORM2 in the main paper) and INLINEFORM3 the outgoing weight matrix of the INLINEFORM4 -dimensional hidden layer. We now remove the INLINEFORM5 -th neuron in the hidden layer and modify the outgoing weights to compensate for the removal: DISPLAYFORM0 
where INLINEFORM0 is the output after the removal operation and INLINEFORM1 are the modified outgoing weights. Our goal is to choose INLINEFORM2 such that the expected error introduced by removing neuron INLINEFORM3 is zero: DISPLAYFORM0 

Figure and table captions:
Figure 1: Unfolding mimics the output of the ensemble of two single layer feedforward networks.
Figure 2: General formula for unfolding weight matrices. The set InnerLayers := [2, D− 1] includes all layers except the input, output, and bias layer.
Table 1: Shrinking layers of the unfolded network on Ja-En to their original size.
Table 2: Compensating for neuron removal in the data-bound algorithm. Row (d) corresponds to row (f) in Tab. 1.
Table 3: Time measurements on Ja-En. Layers are shrunk to their size in the original NMT model.
Table 5: Our best models on Ja-En.
Figure 3: Impact of shrinking on the BLEU score.
Table 4: Layer sizes of our setups for Ja-En.
Table 6: Our best models on En-De.

Questions about this paper:
What were the performance results of their network?

Answers:
For the test set a BLEU score of 25.7 on Ja-En and 20.7 (2014 test set), 23.1 (2015 test set), and 26.1 (2016 test set) on En-De