End-to-End Training for Back-Translation with Categorical Reparameterization Trick

TL;DR

The paper tackles semi-supervised neural machine translation by addressing the gradient-blocking discreteness of latent sentences in back-translation. It introduces a categorical reparameterization trick (CRT) that renders latent sentences differentiable, enabling end-to-end training of both BT models within a VAE-like objective that couples reconstruction terms with a language-model prior and KL regularization. The approach is augmented with Annealing Stochasticity (AS) to balance exploration and stability during training. Empirical results on WMT En-De and En-Tr show BLEU improvements over baselines and prior BT methods, with CRT offering faster reparameterization and controllable gradients, suggesting significant gains for NMT when large supervised data are unavailable.

Abstract

Back-translation (BT) is an effective semi-supervised learning framework in neural machine translation (NMT). A pre-trained NMT model translates monolingual sentences and makes synthetic bilingual sentence pairs for the training of the other NMT model, and vice versa. Understanding the two NMT models as inference and generation models, respectively, the training method of variational auto-encoder (VAE) was applied in previous works, which is a mainstream framework of generative models. However, the discrete property of translated sentences prevents gradient information from flowing between the two NMT models. In this paper, we propose the categorical reparameterization trick (CRT) that makes NMT models generate differentiable sentences so that the VAE's training framework can work in an end-to-end fashion. Our BT experiment conducted on a WMT benchmark dataset demonstrates the superiority of our proposed CRT compared to the Gumbel-softmax trick, which is a popular reparameterization method for categorical variable. Moreover, our experiments conducted on multiple WMT benchmark datasets demonstrate that our proposed end-to-end training framework is effective in terms of BLEU scores not only compared to its counterpart baseline which is not trained in an end-to-end fashion, but also compared to other previous BT works. The code is available at the web.

Figures (2)

• Figure 1: Examples of English-to-German-to-English monolingual process of (a) previous BT and (b) BT with our proposed CRT. Given the English sentence, 'I need backprop to learn.', they infer the German sentence, 'Ich brauche Backprop, um zu lemen.' (only the first three words are illustrated for brevity). The solid black lines are forward propagations, and the thin red lines are backward propagations. While the sampling operation blocks the gradient flow in the previous works (a), the gradient flow detours to the inference model (English-to-German) in (b).
• Figure 2: Example of the BT process with different reparameterization tricks, our CRT or the GST, where the third word's distribution at the third step (green) changed at the next step (pink). Our CRT can output the previous word, $w_c$, while the GST only outputs the maximally probable word, $w_a$, as the final output.