End-to-End Training of Back-Translation Framework with Categorical Reparameterization Trick

Learning discrete representations in neural networks has several advantages. First, it is proper to represent discrete variables such as characters or words in natural language. Second, it can be efficiently implemented at the hardware-level. Lower memory cost and faster matrix multiplication than those of continuous representations are attractive properties [21, 33].

However, because of the non-differentiable property of the discrete representation, a gradient could not be backward propagated through it. To overcome this challenge, the straight-through estimator (STE) was proposed [4], which estimates the gradient for discretizing operations (e.g., sampling or argmax operations) as 1 if the output is 1, otherwise it estimates 0. Although STE imposes a bias on the lower layer’s gradient estimation, it is empirically demonstrated as an effective gradient estimator in the training of existing binary neural networks [21, 33]. The STE estimator can be implemented with automatic differentiation tools such as Pytorch [32, 29]. With a smart reparameterization technique, it outputs a differentiable binary representation as follows.

where $p$ is a normalized probability of the binary variable. $Bern(p)$ and $s$ are the Bernoulli distribution with a parameter $p$ and its random variable. $\text{detach}(c)$ is the operation that detaches its input, $c$, from the gradient computation graph in backward propagation stage. As a result, the final output, $z$, is a binary variable, $z\in\{0,1\}$, and it can flow the gradient through the first term, $p$. As a variant of STE, it is easy to implement and work with other network architectures [34].

The distribution of a sentence is based on a sequence of categorical distributions of words, and a sampled sentence from the distribution is non-differentiable. To make backpropagation feasible through the sentence, we propose CRT, a reparameterization trick for categorical distribution, that is inspired by the binary reparameterization trick as in Section 2.2. Because we handle the categorical distribution (also called as Multinoulli distribution), the Bernoulli distribution in Eq. (1) is replaced by the Multinoulli distribution with the probability $\boldsymbol{p}$ computed by the inference model. Then, a one-hot vector for one word, $\boldsymbol{s}\in\{0,1\}^{|V|}$ s.t. $\sum_{i=1}^{|V|}{\boldsymbol{s}_{i}}=1$, is sampled from the distribution, where $V$ is the vocabulary set. Instead of stochastic sampling based on a distribution, we can also use arbitrarily selection to determine $\boldsymbol{s}$. It makes the CRT can output a one-hot vector for a class that does not have the highest probability. The CRT process is formulated as follow:

where $\odot$ is the element-wise multiplication. $Mult(\boldsymbol{p})$ is the Multinoulli distribution given the normalized probability vector. Based on $\boldsymbol{p}$, we compute the non-differentiable conjugate part, $\boldsymbol{c}$, that is determined by the sample, $\boldsymbol{s}$. Finally, it outputs a one-hot encoded vector $\boldsymbol{z}$ which consists of $\boldsymbol{p}$ and $\boldsymbol{c}$ parts. $\boldsymbol{c}$ is detached from the computation graph, so that backpropagation can flow gradient into the lower layers through $\boldsymbol{p}$. We multiply the scalar $\lambda$ by $\boldsymbol{p}$ in Eqs.(5) and (6) to control the amount of gradient that flows through $\boldsymbol{p}$ while ensuring that the output $\boldsymbol{z}$ remains a one-hot vector. Fig. 1 illustrates the whole process of this trick in the BT framework.

Given the Multinoulli distribution, $Mult(\boldsymbol{p})$, the output process of the one-hot vector, $\boldsymbol{s}$, is implemented by a sampling operator (e.g., Categorical sampling or argmax) and the one-hot encoding function that manually maps an integer value to a $|V|$-dimensional one-hot vector.

In this section, we describe our objective functions for our E2E BT training. The final objective function can be decomposed into two terms with two types of processes: bilingual and monolingual processes. The first objective term for the bilingual process is cross-entropy as follow:

where $x^{b}$ and $y^{b}$ are a source-language and its paired target-language sentences in the bilingual corpus, $D^{b}$, respectively. $\theta$ is the parameter of the ST-NMT model. Likewise, the objective function for the TS-NMT model, $\mathcal{J}_{TS}(\phi)$, is computed in the same way by switching the source and target sentences and replacing the $\theta$ with the TS-NMT model’s parameter, $\phi$.

On the other hand, the objective function for the monolingual process is defined by the negative log likelihood of the monolingual sentence as follows:

where $y^{m}$ is a target-language monolingual sentence in its monolingual corpus, $D^{y}$. The formulations are provided only for the TST process, but the STS process is simply symmetric to the TST process.

As usual, $\log{p(y^{m})}$ in Eq. 11 can be marginalized with latent sentences as follows:

where $\hat{x}$ is the inferred latent sentence that is the aligned representation of $y^{m}$ in the source-language domain, and $p(\hat{x})$ is a prior distribution of $\hat{x}$. By introducing an approximated posterior distribution, $q(\hat{x}|y)$, which is easy to sample $\hat{x}$ given $y^{m}$, we can derive the evidence lower bound objective (ELBO) of the marginal probability, $\log{p(y^{m})}$, based on Jensen’s inequality as follows:

where $D_{KL}$ is Kullback-Liebler divergence (KL). In this TST process of the BT framework, the inference model, $q(\hat{x}|y^{m})$, is modeled by the TS-NMT model with parameter $\phi$. Also, the generation model, $p(y^{m}|\hat{x})$, is modeled by the ST-NMT with parameter $\theta$.

As we mentioned in the introduction section 1, the latent sentence, $\hat{x}$, should be a valid sentence in the source language domain. Therefore, we use a language model’s output distribution [3, 29] as the prior distribution in the KL term, $p(\hat{x})$, instead of isotropic Gaussian in the conventional VAE. In practice, ELBO of the TST monolingual process can be formulated as follows:

where $p_{\psi_{x}}(\hat{x})$ is the pre-trained fixed source language model, and $\alpha_{x}$ is a hyperparameter to control the effect of the KL term in the final objective function. Importantly, to make backpropagation feasible through the latent sentence, we apply the CRT to the latent sentence, $\hat{x}$. In order to input the latent sentence, which comprises a sequence of $|V|$-dimensional one-hot vectors, into the ST-NMT model, we implemented the conventional embedding layer as follows. While, in general, this layer accepts an integer value and returns the corresponding embedding vector from the embedding lookup table, we adapted the embedding layer to output a weighted sum of embedding vectors using the latent sentence vector’s (one-hot) elements as the weight values.

Finally, the total objective functions for $\theta$ and $\phi$ are as follows:

The parameters of the both translation models are updated by the two objective functions.

In this section, we propose an additional regularization technique related to the sampling method of the latent sentence. It would improve performance of our proposed original approach in Section 3.2. As argued in [10], using stochastic sampling in the inference method gives more chances to find a better solution than the greedy method. However, adding more noise to the inference from the beginning of the training can cause the generation model to learn undesirable translation mapping, especially when the bilingual corpus is small. Considering the pro and con, we suggest to anneal the ratio of stochastic sampling during the training process as in the scheduled sampling approach in NMT [2]. As an example of semi-online IBT framework, we newly infer a latent sentence given a monolingual sentence if we do not load the previous latent sentence from the memory. In that time, we infer the latent sentence by only the greedy method at the beginning and slowly increase the ratio of stochastic sampling as the training goes on.

In this section, we demonstrate the experimental results of our internal ablation study to understand the advantages of our proposed E2E BT training and the AS regularization on top of the bilingual NMT and the basic semi-online IBT (not E2E training) with the greedy selection method for latent sentence inference. In addition to the ablation study of ‘En-Tr’ experiment, we also provide comparisons between our proposed CRT and the ST-GST as different reparameterization methods in E2E BT training. To evaluate translation quality, we used case-sensitive SacreBLEU [31] with ‘13a’ tokenizer.

Table 1 presents the BLEU score results on each testset of each model. ‘Transformer’ indicates the NMT model that was trained with only bilingual corpus. ‘BasicBT’ is the basic semi-online IBT method that does not train the two NMT models (ST-NMT and TS-NMT) in an end-to-end fashion. On top of ‘BasicBT’, we applied AS first, and then E2E BT, because the original VAE uses stochastic sampling to infer the latent variable during training, but ‘BasicBT’ is based on greedy selection. Therefore, before applying E2E BT, we applied AS to give appropriate amount of stochasticity during inference at each iteration. Our proposed AS regularization technique shows a slight improvement in En-to-Tr translation. Interestingly, when we apply E2E BT training to the ‘BasicBT+AS’ model, it noticeably increased BLEU scores in general.

Table 2 demonstrates the comparison between several reparameterization methods. First, we find that the ST-GST reparameterization method, ‘E2E BT w/ ST-GST’, is inappropriate in this E2E BT training, leading to significantly low BLEU scores. We interpret the absences of the last two advantageous properties of the CRT, namely ‘(2) Controllable Gradient’ and ‘(3) Flexible Output (FO)’ which were discussed in Section 3.1.1, as significant factors contributing to the poor performances. To check the feasibilities of these two properties, we modified ST-GST to acquire those properties. First, as in CRT, we multiplied $\lambda$ to the last two terms of Eq. (9), ‘E2E BT w/ ST-GST+$\lambda$’, whose $\lambda$ is set to 0.0005 as in ‘E2E BT w/ CRT’. We find that the gradient controllable property is the most crucial factor in the E2E BT training, because it can prevent the model to learn the degenerating solution (refer Section 3.1.1). However, the performances of ‘E2E BT w/ ST-GST+$\lambda$’ are still lower than ‘E2E BT w/ CRT’. Therefore, we further modified the ST-GST to ensuring the FO property with adding the value instead of $g_{i}$ in Eq. (8) that ensures the arbitrarily selected word has the highest value regardless of the output probability. We modeled the value as follow:

where $\epsilon$ is a small value of scalar. Finally, when we applied the final modified version of ST-GST to E2E BT, ‘E2E BT w/ ST-GST+$\lambda$+FO’, it achieved similar performances with ‘E2E BT w/ CRT’. We believe that the current ST-GST needs such modifications to work in E2E BT training, while our proposed CRT is faster and more appropriate to the E2E BT training.

Table 3 presents the ablation study of ‘En-De’ experiments. Macroscopically, our proposed AS regularization and E2E BT training shows similar tendencies of performance gains with ‘En-Tr’ experiments. However, AS regularization, in this experiment, shows noteworthy performance gain in ‘De-En’ more than ‘En-De’. Based on these experiments, we argue that our proposed approaches are advantageous in several benchmark datasets.

In this section, we compare our approaches with other recent BT works that are based on base Transformer architectures [43, 13, 40, 30]. However, there are many differences on testing such as different testsets, different BLEU score metrics, etc. For fair comparisons, we report the average BLEU scores of the previous work and our method on their testsets with the same BLEU metric. Also, we report the size of monolingual corpus that they used for BT training (note that we used 5M monolingual corpora for ‘En-De’ experiments).

Table 4 demonstrates the comparisons between a previous work’s average BLEU score and the same one of our proposed E2E BT model with AS regularization based on the previous work’s BLEU score metric and testsets. We found that our work steadily outperforms the previous works. Especially, the comparison with (H. Pham et al. 2021) indicates an astonishing efficiency of our proposed approach when we consider the amount of monolingual data they used.