Generating Diverse Translation from Model Distribution with Dropout

For most of machine learning tasks based on neural networks, a model with specific parameters is trained to explain the observed training data. However, there are usually a large number of possible models that can fit the training data well, which leads to model uncertainty. Model uncertainty may result from noisy data, uncertainty in model parameters or structure uncertainty, and it is represented as the confidence which model to choose to predict with. In order to express model uncertainty, we consider all possible models with parameters $\bm{\omega}$ and define a prior distribution $P(\bm{\omega})$ over the model (i.e., the space of the parameters) to denote model uncertainty. Then given a training data set $(\bm{X},\bm{Y})$, the predicted distribution for $\bm{Y}$ can be denoted as $P(\bm{Y}|\bm{X},\bm{\omega})$.

Following Gal et al. (2017), we employ Bayesian neural networks (BNNs) to represent the $P(\bm{\omega}|\bm{X},\bm{Y})$, which is the posterior distribution of the models under $(\bm{X},\bm{Y})$. BNNs offer a probabilistic interpretation of deep learning models by inferring distributions over the models’ parameters which are trained using Bayesian inference. The posterior distribution can be got by invoking Bayes’ theorem as:

Then according to BNNs, given a new test data $\bm{x^{\prime}}$, the predictive distribution of the output $\bm{y^{\prime}}$ is:

The expectations in Equation 7 and 8 are integrated over model distribution $\bm{\omega}$, and the huge space of $\bm{\omega}$ makes it intractable to obtain the results. Therefore, inspired by  Hinton and Van Camp (1993), Graves (2011) proposes a variational approximation method, using a variational distribution $Q(\bm{\omega}|\theta)$ with the parameters $\theta$ to approximate the posterior distribution $P(\bm{\omega}|\bm{X},\bm{Y})$. To this end, the training objective is to minimize the Kullback-Leibler (KL) divergence between the model distribution and the posterior distribution $\mathrm{KL}(Q(\bm{\omega}|\theta)||P({\bm{\omega}}|\bm{X},\bm{Y}))$. With variational inference, the objective is equivalent to maximizing its evidence lower bound (ELBO), so we get

As we can see in Equation 9, the first term on the right side is the expectation of the predicted probability over model distribution on the training set, which can be unbiased estimated with the Monte-Carlo method. And the second term is the KL divergence between the approximate model distribution and the prior distribution. From the perspective of  Hinton and Van Camp (1993) and Graves (2011), with the above objective, we can express model uncertainty under the training data and meanwhile regularize model parameters and avoid over-fitting. Therefore, at inference, we can use the distribution $Q(\bm{\omega}|\theta)$ instead of $P(\bm{\omega}|\bm{X},\bm{Y})$ to evaluate model confidence (i.e., model uncertainty).

To derive the BNN, we need to first decide how to explore for the possible models and then decide the prior distribution and the variational distribution for the models. As in Gal et al., 2017, we can define a simple model with parameters ${\bm{\omega}_{{}_{\bm{W}}}(\bm{W}\in\mathbb{R}_{m\times n})}$ and then drop out some column of ${\bm{\omega}_{{}_{\bm{W}}}}$ to get the possible models. We use matrix Gaussian distribution as the prior model distribution and Bernoulli distribution as the posterior model distribution.

Using ${\bm{W}{.j}}$ to denote the $j$-th column of ${\bm{W}}$, we draw a matrix Gaussian distribution as the probability distribution of dropping out the $j$-th column as

where $l$ is the hyper-parameter.

The above matrix Gaussian distribution is used as the prior distribution of the models got by dropping out the $j$-th column. Then we introduce $\bm{p}\ (\bm{p}\in\mathbb{R}_{1\times n})$ as the probability vector of dropping out the columns of ${\bm{\omega}_{{}_{\bm{W}}}}$, which means dropping out the $j$-th column with the probability of $\bm{p}_{j}$, and keeping the $j$-th column unchanged with the probability of $1-\bm{p}_{j}$. Therefore the posterior model distribution of dropping out the $j$-th column is defined as

where $\bm{W}\in\theta$ and $\bm{p}\in\theta$ are trainable parameters.

With Equation 10 as prior, the KL divergence for the $j$-th column of the matrix can be represented as:

where

and

Since the probability distribution among different neural networks and different columns of neural network are independent. For a complex multi-layer neural network $\theta$, the KL divergence between model distribution $Q(\omega|\theta)$ and prior distribution $P({\omega})$ is

Stated in detail in  Vaswani et al. (2017), in Transformer, dropout is commonly used to the output of modules, including the output of embedding, attention layer and feed-forward layer. Also, from Equation 13, we find it’s important to find the network $W$ corresponding to the dropout module. The correspondences in Transformer are as follows:

Embedding module Embedding module works for mapping the words into embedding vectors. The embedding module contains a matrix $W_{E}\in\mathbb{R}_{l_{d}\times d}$, where $l_{d}$ is the length of dictionary and $d$ is the dimension of embedding vector. For the $i$-th word in the dictionary, its embedding vector is the $i$-th column of $W_{E}$. Since dropout the $j$-th dimension of word embedding is equivalent to dropping out the $j$-th row of $W_{E}$, we utilize $W_{E}^{T}$ and its corresponding dropout in Equation 13.

Attention module For attention modules, as we can see in 5, their outputs are generated by concatenating the output of different heads and projecting by matrix $W^{O}$. Since dropout is used in the output of attention module, we take $W^{O}$ and its corresponding dropout in calculating Equation 13.

Feed-forward module As shown in Equation 6, the output is generated through $W_{2}$ with bias $b_{2}$. As we can see, for network $y=xW+b$, we can find that

as we can see, dropout to the output of the feed-forward module can be regraded as dropping out $W_{2}$ and $b_{2}$. So, during training, we use $\mathrm{Concat}(W^{T},b^{T})^{T}$ to calculate Equation 13.

Although dropout is frequently utilized in Transformer, there are some networks in Transformer like $W_{i}^{Q}$, $W_{i}^{K}$, $W_{i}^{V}$ in Equation 4, and their output is not masked by dropout. So, in our implementation, we obtain the model distribution by fine-tuning the pre-trained model, freezing their parameters and only updating dropout probabilities. Moreover, we choose different trained modules to train their output dropout probability, and in calculating Equation 15, we only take those trained dropout probabilities into consideration. By allowing dropout probabilities to change, our method can better represent model distributions under the training dataset than the fixed dropout probabilities. The mini-batch training algorithm is expressed in Equation 1. It is worth to mention that since we train the model distribution with batches of data, we scale the KL divergence with the proposition of the batch in the entire training dataset.

During updating the dropout probability, due to the discrete characteristics of the Bernoulli distribution, we cannot directly calculate the gradient of the first term in Equation 9 to the dropout probability. So, we adopt concrete dropout, which is used in  Gal et al. (2017). As a continuous relaxation of dropout, for its input $\bm{y}$, the output can be expressed as $\bm{y^{\prime}}=\bm{y}\odot\bm{z}$, and vector $\bm{z}$ satisfies:

where $u\sim\mathcal{U}(0,1)$, $p$ is dropout probability.

In the inference stage, we just randomly mask model parameters with trained dropout probabilities, with different random seeds, NMT models with different parameters are sampled. Since diverse translations are demanded, we performed several forward passes through different sampled NMT models, and different translations are generated with different model outputs and beam search.

In this experiment, we train models with different training modules and hyper-parameter $l$ with NIST dataset to evaluate their effects, and some results are shown in Table 1.

As we can see in Table 1, for those training modules, when $l$ is small, with the same hyper-parameter $l$, choosing smaller part of training modules will lead to lower BLEU and Pairwise-BLEU, showing that accuracy of the generate translations increases while diversity decreases. We also find that in the same training modules, with the increase of $l$, the Pairwise-BLEU decreases steadliy, and then increases when the BLEU is close to zero; and the BLEU has similar trends with Pairwise-BLEU, however, when $l$ is relatively low, the BLEU tends to stablize.

For the above-mentioned experimental results, we can interpret as follows: in training modules, since Equation 15 is the sum of the training modules’ KL divergence, with the training modules increase, the KL divergence accordingly increases, pushing the dropout probability higher and making translations diverse. In terms of hyper-parameter $l$, as we can see in Equation 10, when $l$ increases, the prior distribution is squeezed to zero matrix; thus during training, the dropout probabilities will get higher to make the model distribution close to prior distribution. However, when the $l$ is too high to make most of dropout probabilities close to 1, uncertainty of model parameters decreases, making Pairwise-BLEU increases.

From the previous section, we can see that by selecting different training modules and hyper-parameter, translations with different accuracy and diversity can be obtained. Then we conduct experiments to generate 5 groups of different translations on NIST Zh-En dataset and WMT’14 En-De dataset, and compare the diversity and accuracy of the translations generated by our method and the following baseline approaches:

• Beam Search: we choose the optimal 5 results directly generated by beam search in this paper.

• Diverse Beam Search (DBS)  (Vijayakumar et al., 2016): it works by grouping the outputs and adding regularization terms in beam search to encourage diversity. In our experiment, the number of output translations of groups are all 5.

• HardMoE (Shen et al., 2019): it trains model with different hidden states and obtains different translations by controlling hidden state. In our experiment, we set the number of hidden states is 5.

• Head Sampling (Sun et al., 2019): it generate different translations by sampling different encoder-decoder attention heads according to their attention weight, and copying the samples to other heads in some conditions. Here, we set the parameter $\mathrm{K}=3$.

The results are shown in Figure 1. In Figure 1 we plot the BLEU versus Pairwise-BLEU, the scattered points show the results of baseline approaches, and the points on the curves are results in the same training modules with different hyper-parameter $l$. From Figure 1, firstly, we can verity that choosing different training modules can lead to different balance of translation diversity and accuracy, for NIST Zh-En and WMT’14 En-De, training dropout probabilities in the full model can get better translations.

Also, we suggest that in NIST Zh-En task, by adjusting training modules and hyper-parameter $l$, our results which has higher BLEU and lower Pairwise-BLEU values than baselines without HardMoE, even for HardMoE itself, our method is comparable with proper $l$ while training the whole model. In WMT’14 En-De, we also find that our method exceeds the baseline approach except HardMoE. For the gap in performace with HardMoE, we interpreted that since our models are randomly sampled from the model distribution, it could be hard for our models to represent such distinguishable characteristics like HardMoE, which trains multiple different latent variables.

Also, to intuitively display the improvement of diversity in our translations, we choose a case from NIST Zh-En task, the results are shown in Table 2. The case shows that compared with beam search, which only varies in few words, our method can obtain more diverse translations while ensuring the accuracy of translation, and diversities are not only shown in words, but also reflected in lexical characteristic.

Some researches  (Voita et al., 2019; Michel et al., 2019; Fan et al., 2019) found that a well-trained Transformer model is over-parameterized. Useful information gathers in some parameters and some modules and layers can be pruned to improve the efficiency during test time. Since dropout can play the role of regularization and there are differences in the trained dropout probabilities of different neuron, we conjecture that the trained dropout probability and the importance of each module are correlated. To investigate this, we choose the model in which dropout probabilities of the full model is trained with $l^{2}=400$ in NIST Zh-En task, and separately calculate the average dropout probability $\bar{p}_{dropout}$ of different attention modules. Also, we manually pruned the corresponding modules of the model, obtained translations and calculated its BLEU. the more the BLEU drops, the more important the module is to translation. To quantify their relevance, we calculate the Pearson correlation coefficient (PCC) $\rho$ in different kinds of training modules, and highlights the highest and lowest results.

Results of our experiment are shown in Table 3. Firstly, we can see that the average dropout probabilities and BLEU are not fully positively correlated, which might be explained by the contingency of sampling from model distribution during training. But from the maximum and minimum of $\bar{p}_{dropout}$ and BLEU, we can find that the dropout probabilities $\bar{p}_{dropout}$ and the BLEU of translations show some similar information in module importance. Also, we quantify the correlation between the $\bar{p}_{dropout}$ and BLEU, finding that it is highly correlated in self-attention module in encoder and in E-D attention in decoder, since its correlation coefficient $\rho>0.8$, and the $\bar{p}_{dropout}$ and BLEU is also correlated in self-attention in decoder.