SlowGuess/NAACL_2022_OCR huggingface dataset:
Minghao Zhu and Junli Wang and Chungang Yan
Key Lab of Embedded System and Service Computing (Tongji University),
Ministry of Education, Shanghai 201804, China.
National (Province-Ministry Joint) Collaborative Innovation Center for
Financial Network Security, Tongji University, Shanghai 201804, China.
{mzhu, junliwang, yanchungang}@tongji.edu.cn
Variational Autoencoder (VAE) is an effective framework to model the interdependency for non-autoregressive neural machine translation (NAT). One of the prominent VAE-based NAT frameworks, LaNMT, achieves great improvements to vanilla models, but still suffers from two main issues which lower down the translation quality: (1) mismatch between training and inference circumstances and (2) inadequacy of latent representations. In this work, we target on addressing these issues by proposing posterior consistency regularization. Specifically, we first perform stochastic data augmentation on the input samples to better adapt the model for inference circumstance, and then conduct consistency training on posterior latent variables to construct a more robust latent representations without any expansion on latent size. Experiments on  $\mathrm{En} < - >\mathrm{De}$  and  $\mathrm{En} < - >\mathrm{Ro}$  benchmarks confirm the effectiveness of our methods with about 1.5/0.7 and 0.8/0.3 BLEU points improvement to the baseline model with about  $12.6\times$  faster than autoregressive Transformer.
Neural Machine Translation (NMT) achieves great success in recent years, and typical sequence-to-sequence frameworks like Transformer (Vaswani et al., 2017) achieved state-of-the-art performance on the task of NMT. In this framework, source sentences are translated in an autoregressive (AT) manner where each token is generated depending on previously generated tokens. Inevitably, such sequential decoding strategy results in a high inference latency. To alleviate this issue, Non-autoregressive translation (NAT; Gu et al., 2018) was proposed to speed-up decoding procedure by generating target tokens in parallel. However, the translation quality of vanilla NAT is compromised, one of the most significant problems is multi-modality and it usually results in multiple translation results, duplicate or missing words in
target sentences of NAT models (Gu et al., 2018). This situation results from the conditional independence proposed by NAT, since models are trained to maximize the log-probability of target tokens at each position while the interdependency is omitted.
The key to alleviate the multi-modality issue is to model the dependency information of targets implicitly or explicitly so decoder can easily learn and capture the information between target tokens and generate more accurate translations. For example, Ghazvininejad et al. (2019) and (2020b) model the target dependency by providing observed target tokens in training and performing iterative inference. Ran et al. (2021) generates intermediate representations by permuting the source sentences in the target order. Libovický and Helcl (2018), Shao et al. (2020) and Ghazvininejad et al. (2020a) model the target dependency by introducing objective functions with alignment. Guo et al. (2020) and Wang et al. (2019) improve the translation quality by proposing additional regularizations. Zhou and Keung (2020) and Liu et al. (2021) utilize the external information like monolingual data or semantic structure to help the training.
Previous studies have validated the effectiveness of applying VAE on AT (Zhang et al. 2016; McCarthy et al. 2019; Su et al. 2018) and NAT (Kaiser et al. 2018; Shu et al. 2020) frameworks. A prominent NAT model is  $\mathrm{LaNMT}^{1}$  (Shu et al., 2020) which encodes the source and target tokens into intermediate Gaussian distribution latent variables and outperforms vanilla NAT with about 5.0 BLEU points on WMT14 En-De task with  $12.5\times$  speedup to base Transformer. However, there exists a slight lag behind the state-of-the-art fully NAT models. It may be attributed to two reasons: (1) The inadequate representations of latent variables. Figure 1 shows the effect of latent size to translation qualities. Obviously, the optimal capacity of latent variables is significantly lower than the model's hid
'https://github.com/zomux/lanmt
Human Language Technologies, pages 607 - 617
Figure 1: Effect of latent size. Values are obtained by our experiments on WMT14 En-De benchmark
den size (512) while high-capacity latent variables conversely deteriorate the performance because the minimization between prior and posterior becomes difficult (Shu et al., 2020). (2) The mismatch between training and inference circumstances that the posterior module receives the gold sentence as inputs during training but imperfect initial translation instead during inference. Thus, in this paper, we aim to improve the robustness of the latent representation and move the training circumstance close to inference circumstance.
To this end, we propose a consistency regularized posterior network for better latent representations and closer training-inference circumstance with no extended latent size. Specifically, it can be split into two main steps: we first apply stochastic data augmentation methods to inject stochastic noises in posterior inputs  $x$  and  $y$  to get two different views  $(x_{1},y_{1})$  and  $(x_{2},y_{2})$ , which are then transformed to corresponding latent variables  $z_{1}$ ,  $z_{2}$  by the posterior network. Secondly, the consistency regularization step tries to minimise the gap between  $z_{1}$  and  $z_{2}$  since they are derived from the same pair of input  $x$  and  $y$ . These two steps enable the sequences better represented by the size-limited latent variables which is trained to be more robust to the noises in input samples. Meanwhile, posterior module receives noisy views instead of gold samples during training, it is more adaptive to the inference time which receives imperfect initial translations as inputs.
We verified the performance and effectiveness of our methods on WMT14 En<->De and WMT16 En<->Ro benchmarks. Our methods outperform the latent variable baseline with about
1.5/0.7 and 0.8/0.3 BLEU points improvement on four benchmarks. With these improvements, we achieve the comparable performance to the state-of-the-art fully NAT approaches: 25.65/30.23 and 31.56/31.20 of BLEU scores with similar decoding speed, and it can be improved further with latent search. The contributions of our work can be summarized as follows:
- For better latent representations, we propose consistency regularization optimized posterior network, which improves the translation quality by training more robust latent variables to noises.
- We apply four data augmentation methods to cooperate with posterior consistency regularization, where all of them are also benefit to the translation quality by alleviating the mismatch between training and inference circumstances.
- We show our strategy is capable of improving the translation quality of the base latent-variable NAT model to be comparable with the state-of-the-art fully NAT frameworks.
2.1 Non-Autoregressive Translation
Traditional sequence-to-sequence NMT models generate target sentences in an autoregressive manner. Specifically, given a source sentence  $x$ , AT frameworks model the conditional probability of  $y = \{y_{1},y_{2},\dots ,y_{l_{y}}\}$  by the following form:
$$
\log p (y | x) = \sum_ {i = 1} ^ {l _ {y}} \log p \left(y _ {i} \mid y _ {<   i}, x\right) \tag {1}
$$
where  $y_{< i}$  indicates the target tokens already generated before  $y_{i}$ . Hence, the target tokens are generated sequentially which results in a high decoding latency. To alleviate this issue, vanilla NAT (Gu et al., 2018) breaks the conditional dependency by conditional independence assumption so that all tokens can be generated independently. Following its probability form:
$$
\log p (y | x) = \sum_ {i = 1} ^ {l _ {y}} \log p \left(y _ {i} | x\right) \tag {2}
$$
where each target token  $y_{i}$  now only depends on the source sentence  $x$ . Benefit from the parallel
computing capability of hardware accelerators like GPU or TPU, all tokens can be generated with one iteration in an ideal circumstance.
2.2 Latent-Variable Model
We mainly focus on performing optimization on the variational NAT framework proposed by Shu et al. (2020). The network architecture is constructed by four main components. An encoder  $p_{\omega}(z|x)$  encodes the source representation of input  $x$  and computes the prior latent variable. An approximate posterior network  $q_{\phi}(z|x,y)$  accepts both the source sentence  $x$  and target sentence  $y$  as the input and computes the posterior latent variable. A length predictor  $p(l_y|z)$  predicts the length of target sentence  $y$ , and finally a decoder  $p_{\theta}(y|x,z,l_y)$  with a length transform module to transform the latent variables  $z$  to the target length  $l_y$  at first and reconstruct  $y$  from  $z$  with the source representations of  $x$ . Note that the  $l_y$  here is the gold length in training. Hence, the training objective is aiming to maximize the evidence lowerbound (ELBO):
$$
\begin{array}{l} \mathcal {L} (x, y) = \mathbb {E} _ {z \sim q _ {\phi}} [ \log p _ {\theta} (y | x, z) ] \\ - \mathrm {K L} [ q _ {\phi} (z | x, y) | | p _ {\omega} (z | x) ] \tag {3} \\ \end{array}
$$
$$
p _ {\theta} (y | x, z) = p _ {\theta} (y | x, z, l _ {y}) p (l _ {y} | z)
$$
where the latent variables  $z$  is constrained with the same length as  $x$  and the value is modeled as spherical Gaussian distribution. KL denotes Kullback-Leibler divergence.
2.3 Consistency Regularization
Consistency regularization is considered as an effective method on semi-supervised learning to capture the potential features from unlabeled samples (Sajjadi et al., 2016; Laine and Aila, 2017; Tarvainen and Valpola, 2017; Xie et al., 2020). It is also utilized as a complementary regularization tool with other regularization methods to prevent model from overfitting (Liang et al., 2021). In a nutshell, consistency regularization assumes a well trained model should be robust enough to any small changes in the input samples or hidden states and generate invariant outputs (Xie et al., 2020). To this end, it regularizes model's final outputs to be invariant to input samples with small stochastic noises injected by minimizing the gap between two augmented views of one sample.
In this paper, we focus on the posterior module of the variational framework and apply consistency regularization on it instead of the whole network.
Figure 2: The overall pipeline of training with posterior consistency regularization
Along with data augmentation for noise injection, consistency regularization is capable to improve the representation of this module and result in better translation quality.
The posterior module is considered to train with consistency regularization and data augmentation for better translation quality. In this section, we will introduce the details of our method, including the overall network architecture, the objective and procedure of training with consistency regularization, four data augmentation methods and three decoding strategies applied for inference.
We follow the variational model architecture proposed by Shu et al. (2020) with four main components: encoder, posterior, length predictor and decoder module. Since we apply consistency regularization on the posterior, an additional stochastic data augmentation module is added for noise injection on posterior input samples. With two augmented views derived from one sample, each sam
ple thus appears twice in a training batch. Figure 2 shows the brief model architecture and training pipeline of our work. The part in the dashed box is the major difference to the base model.
3.2 Posterior Consistency Regularization
As discussed above, consistency regularization is applied on the posterior module to improve its robustness. Given a training sample with a pair of source sentence  $x$  with  $l_{x}$  tokens and target sentence  $y$  with  $l_{y}$  tokens, we first apply data augmentation on both  $x$  and  $y$  twice to inject stochastic noises and obtain two different views (denoted as  $x_{1}, y_{1}$  and  $x_{2}, y_{2}$ ). Both views are forwarded to the posterior network  $q_{\phi}(z|x,y)$  to predict the mean and variance vectors of two latent variables  $z_{1}$  and  $z_{2}$ . Since the latent variables derive from the same input sample, the consistency regularization method tries to minimize the difference between these two latent variables by measuring bidirectional KL-divergence as follows:
$$
\begin{array}{l} \mathcal {L} _ {c o n s} = \frac {1}{2} (\mathrm {K L} (z _ {1} | | z _ {2}) + \mathrm {K L} (z _ {2} | | z _ {1})), \\ z _ {1} = q _ {\phi} (Z | X = x _ {1}, Y = y _ {1}), \tag {4} \\ z _ {2} = q _ {\phi} (Z | X = x _ {2}, Y = y _ {2}) \\ \end{array}
$$
Combining with the basic negative log-likelihood (NLL) objective on the decoder, since there are two different  $z$  for the same sample, it is evaluated by averaging them:
$$
\mathcal {L} _ {n l l} = - \frac {1}{2} \left(\log p _ {\theta} (y | x, z _ {1}, l _ {y}) + \log p _ {\theta} (y | x, z _ {2}, l _ {y})\right) \tag {5}
$$
Note that the gold length  $l_{y}$  of target sentence  $y$  is used which is known during training. Similarly, the objective of the length predictor is calculated by:
$$
\mathcal {L} _ {l e n} = - \frac {1}{2} \left(\log p \left(l _ {y} \mid z _ {1}\right) + \log p \left(l _ {y} \mid z _ {2}\right)\right) \tag {6}
$$
To back propagate the gradient information from the decoder and length predictor to posterior, reparameterization trick is applied to sample  $z$  from  $q_{\phi}$  where  $z = \mu + \theta * \mathcal{N}(0,1)$  in Eq.(6) and (5). Here,  $\mu$  and  $\theta$  indicate mean and variance vector. For encoder, it not only generates representations of source sentence  $x$  but also computes the prior latent variables. Thus, we close the KL-divergence between prior and two posterior latent variables by:
$$
\mathcal {L} _ {p r i o r} = \frac {1}{2} (\mathrm {K L} (z _ {1} | | z _ {p}) + \mathrm {K L} (z _ {2} | | z _ {p})) \tag {7}
$$
where  $z_{p} = p_{\omega}(Z|X = x)$ ,  $z_{1}$  and  $z_{2}$  are obtained from (4). Finally, to achieve the similar goal of maximizing (3), we minimize the loss function by combining (4), (5), (6) and (7) as follows:
$$
\mathcal {L} _ {\text {l o s s}} = \mathcal {L} _ {\text {n l l}} + \mathcal {L} _ {\text {l e n}} + \mathcal {L} _ {\text {p r i o r}} + \alpha \mathcal {L} _ {\text {c o n s}} \tag {8}
$$
where  $\alpha$  here is the only hyperparameter to weight the consistency regularization loss.
3.3 Data Augmentation Methods
Given an embedding matrix  $\mathbb{R}^{L\times d}$  with  $L$  tokens embedded into  $d$ -dimensions vectors, to generate different views for the posterior network inputs and perform posterior consistency regularization, as well as to close the gap between training and inference circumstances, we explore four data augmentation methods for this purpose including dropout, feature cutoff, token cutoff and replacement as presented in Figure 3.
Dropout Dropout (Srivastava et al., 2014) is widely used as a regularization method to prevent neural networks from overfitting. But in this paper, we found that it is also an effective data augmentation method for noise injection. Specifically, we randomly choose values on token embeddings by a specific proportion and force them to zero.
Cutoff This is a simple but effective augmentation method proposed by Shen et al. (2020). The cutoff methods we adopt include token cutoff and feature cutoff. For token cutoff, a specific proportion of tokens are chosen from the token dimension  $L$  and dropped by setting the vectors to zero. For feature cutoff, the dropped values are chosen from feature dimension  $d$  instead.
Replacement This is similar to the token replacement adopted by BERT pre-training (Devlin et al., 2019) where the chosen token vectors are replaced by the embedding of new tokens that randomly selected from the vocabulary instead of setting them to zero or any special tokens directly.
3.4 Decoding Strategies
Non-refinement For this strategy, we completely follow the original design (Shu et al., 2020) where the posterior network is discarded since the target sentence  $y$  is unknown during inference. The foremost step is to obtain the representations of  $x$  and the prior latent variable  $z$  from encoder with source input  $x$ . The latent variable is then used to determine the target length and generate target sentence.
Figure 3: Four stochastic data augmentation methods we used for noise injection
Note that to avoid randomness during inference,  $z$  is set to its mean value  $\mu$  instead of reparameterization sampling. This can be summarized as follows:
$$
\mu_ {0} = \mathbb {E} _ {p _ {\omega} (z | x)} [ z ],
$$
$$
l _ {y _ {0}} = \underset {l _ {y}} {\arg \max } p \left(l _ {y} \mid z = \mu_ {0}\right), \tag {9}
$$
$$
y _ {0} = \underset {y} {\arg \max } p _ {\theta} (y | x, l _ {y _ {0}}, z = \mu_ {0})
$$
Deterministic Refinement The posterior network  $q_{\phi}$  can be reused to take refinement on the initial output  $y_0$  above. According to Shu et al. (2020), its original design allows multi-step iterative refinement for more precise translations, but sacrifices huge cost in decoding speed for a tiny quality improvement. Thus, we consider one-step only refinement in this paper:
$$
\mu_ {1} = \mathbb {E} _ {q _ {\phi} (z | x, y _ {0})} [ z ],
$$
$$
l _ {y _ {1}} = \underset {l _ {y}} {\arg \max } p \left(l _ {y} \mid z = \mu_ {1}\right), \tag {10}
$$
$$
y _ {1} = \underset {y} {\arg \max } p _ {\theta} (y | x, l _ {y _ {1}}, z = \mu_ {1})
$$
Here the  $y_{1}$  is the final output after refinement.
Latent Search Since reparameterization is disabled in above two strategies to generate deterministic results, it is also able to search the best latent variable from Gaussian distribution. Specifically,  $m$  prior latent variables are sampled by reparameterization and decoded in parallel, result in  $m$  target candidates for each source sentence. To get the best result, we select the candidate with the highest score by averaging the log-probability of tokens as the final output. This is different from Shu et al. (2020) or Noisy Parallel Decoding (NPD;
Gu et al. 2018) which rescore the candidates by autoregressive teacher and at least cuts the decoding speed by half, our no-rescoring strategy is still effective and much faster.
In this section, we will introduce the settings of our experiments, report the main results and compare our model to the representative NAT frameworks. Our experiments mainly focus on (1) the improvement benefit from our optimization to former VAE-based NAT model. (2) The effectiveness of consistency regularization and different data augmentation methods.
Dataset Four of the commonly used machine translation benchmarks are adopted to evaluate our proposed method: WMT14 English<-German2 (En-De and De-En, 4.5M) and WMT16 English<-Romanian3 (En-Ro and Ro-En, 610K). We follow previous works' data preprocessing configurations to preprocess the data (En-De: Shu et al., 2020, En-Ro: Ghazvininejad et al. 2019). To learn the subword vocabulary, we apply SentencePiece (Kudo and Richardson, 2018) to generate joint subword vocabulary of 32K tokens for each dataset respectively.
Knowledge Distillation Following previous studies on NAT that models are trained on distilled data generated by autoregressive teacher, we also apply sentence-level knowledge distillation for all datasets to obtain less noisy and more deterministic data. In this work, Transformer (Vaswani et al.,
$^{2}$ https://www.statmt.org/wmt14/
<sup>3</sup>https://www.statmt.org/wmt16/
Models	Iter.	WMT14
En-De	WMT14
De-En	WMT16
En-Ro	WMT16
Ro-En	Speed
Transformer (Vaswani et al., 2017)	N	27.30	/	/	/	/
Transformer (ours)	N	27.74*	31.28*	33.73*	34.38*	1.0×
NAT-IR (Lee et al., 2018)	10	21.61	25.48	29.32	30.19	1.5×
CMLM (Ghazvininejad et al., 2019)	10	27.03	30.53	33.08	33.31	1.7×
LevT (Gu et al., 2019)	Adv.	27.27	/	/	33.26	4.0×
JM-NAT (Guo et al., 2020)	10	27.69	32.24	33.52	33.72	5.7×
Vanilla-NAT (Gu et al., 2018)	1	17.69	21.47	27.29	29.06	15.6×
Imitate-NAT (Wei et al., 2019)	1	22.44	25.67	28.61	28.90	18.6×
FlowSeq (Ma et al., 2019)	1	23.72	28.39	29.73	30.72	1.1×
NAT-DCRF (Sun et al., 2019)	1	23.44	27.22	/	/	10.4×
NAT-REG (Wang et al., 2019)	1	20.65	24.77	/	/	27.6×
BoN (Shao et al., 2020)	1	20.90	24.61	28.31	29.29	10.7×
AXE (Ghazvininejad et al., 2020a)	1	23.53	27.90	30.75	31.54	/
GLAT (Qian et al., 2021)	1	25.21	29.84	31.19	32.04	15.3×
Reorder-NAT (Ran et al., 2021)	1	22.79	27.28	29.30	29.50	16.1×
SNAT (Liu et al., 2021)	1	24.64	28.42	32.87	32.21	22.6×
LT (Kaiser et al., 2018)	/	19.80	/	/	/	3.8×
LaNMT (Shu et al., 2020)	1	22.20	26.76*	29.21*	28.89*	22.2×
+ refinement	2	24.10	29.47*	30.76*	30.86*	12.5×
+ latent search w/ rescoring	2	25.10	/	/	/	6.8×
Ours, decode w/o refinement	1	23.92	27.39	29.90	29.04	25.6×
+ latent search (m=9) w/o rescoring	1	25.59	30.11	31.40	31.63	21.1×
decode w/ refinement	2	25.65	30.23	31.56	31.20	12.6×
+ latent search (m=9) w/o rescoring	2	26.23	31.23	32.50	32.14	11.0×
Table 1: BLEU scores and speedup rates for performance comparison on WMT14 En<>De and WMT16 En<>Ro benchmarks without rescoring. We report the best scores here among all tested combinations of data augmentation methods with consistency regularization. Iter. denotes the number of iterations during inference. Adv. means adaptive. / denotes the value is not reported, * denotes the results are obtained by our implementation.
2017) with base settings is adopted and reproduced as the teacher model for data distillation.
Implementation Details The model is trained by the objective function illustrated on Eq.(8). To avoid posterior collapse, freebits annealing (Chen et al., 2017) is applied on KL terms in Eq.(7) to keep a distance between prior and posterior. Its threshold is fixed to 1 for the first half training steps, and linearly decay to 0 on the second half. For both datasets, we train the model with a batch of approximate 40K tokens for overall 100K steps on four Tesla V100 GPUs and conduct fine-tuning for additional 20K steps with annealing disabled.
For network settings, we use 6 layers encoder and decoder with  $d_{model} / d_{feedforward} = 512 / 2048$ . Following Shu et al. (2020), the posterior network contains 3 transformer layers and the dimension of latent variable is set to 8. We
set L2 weight decay of 0.01 and dropout between attention layers with rate of  $0.1 / 0.3$  for WMT14  $\mathrm{En} < - >\mathrm{De}$  and WMT16  $\mathrm{En} < - >\mathrm{Ro}$  respectively as well as label smoothing rate  $\epsilon = 0.1$  on target tokens. Models are trained by AdamW (Loshchilov and Hutter, 2017) with settings of  $\beta = (0.9, 0.98)$  and  $\epsilon = 1e - 4$ . The learning rate is warmed up for first 4000 steps to  $1.4\mathrm{e} - 3$  and decayed by inverse-square-root scheduler. To obtain the final model, we average 5 best checkpoints chosen by validation BLEU score. For the rate of noise injection, all four augmentation strategies are evaluated with rates from 0.1 to 0.3 on WMT14 En-De benchmark and the settings for other experiments are determined according to these results.
Evaluation For all benchmarks, we use sacre-BLEU $^4$  (Post, 2018) to evaluate BLEU score of
<sup>4</sup>https://github.com/mjpost/sacrebleu
translation results. Following Lee et al. (2018) and Shu et al. (2020), repetition tokens are removed before generating the final outputs for evaluation. The results of latent search is obtained by the mean score of 5 independent runs on the test set of each benchmark to get more precise measures since reparameterization causes randomness in decoding.
To evaluate the decoding speed, following previous works (e.g. Gu et al. 2018, Lee et al. 2018), models are run on WMT14 En-De test set with batch size of 1 under the environment with one GPU only. The mean value of decoding latency among all samples is collected and represent as the decoding speed. Meanwhile, base Transformer is reproduced and evaluated on the same machine to obtain the speed up rates.