Reinforcement Learning for Edit-Based Non-Autoregressive Neural Machine Translation (2024)

Abstract: Non-autoregressive (NAR) language models are known for their low latency in neural machine translation (NMT). However, a performance gap exists between NAR and autoregressive models due to the large decoding space and difficulty in capturing dependency between target words accurately. Compounding this, preparing appropriate training data for NAR models is a non-trivial task, often exacerbating exposure bias. To address these challenges, we apply reinforcement learning (RL) to Levenshtein Transformer, a representative edit-based NAR model, demonstrating that RL with self-generated data can enhance the performance of edit-based NAR models. We explore two RL approaches: stepwise reward maximization and episodic reward maximization. We discuss the respective pros and cons of these two approaches and empirically verify them. Moreover, we experimentally investigate the impact of temperature setting on performance, confirming the importance of proper temperature setting for NAR models’ training.

## introduction non-autoregressive (nar) language models @xcite generate translations in parallel, enabling faster inference and having the potential for real-time translation applications. however, despite their computational efficiency, nar models have been observed to underperform autoregressive (ar) models due to the challenges posed by the large decoding space and difficulty in capturing dependency between target words accurately @xcite . to bridge the performance gap, many nar architectures and training methods have been proposed, including edit-based models like insertion transformer @xcite and levenshtein transformer @xcite . prior research has also explored knowledge distilla-* work done during internship at cyberagent ai lab. unlike ar models, preparing teacher data and designing appropriate training objectives have always been challenging for nar models @xcite . teacher forcing with inappropriate teacher data may exacerbate the exposure bias problem @xcite , affecting model performance. reinforcement learning (rl) is known for its ability to tackle the exposure bias @xcite and alleviate the object mismatch issue @xcite . despite its importance, explorations of rl for nar are still scarce. @xcite proposed a method for reducing the estimation variance. however, this method is only applicable to nar models with a fixed output length, which is unsuitable for edit-based models. in this paper, we empirically analyze conditions for performance improvement in applying rl to edit-based nar models in neural machine translation (nmt). specifically, we focus on levenshtein transformer (levt) @xcite , a prominent edit-based nar architecture that has shown promise in reducing decoding latency and flexible length adjustment. we demonstrate that rl with self-generated data significantly improves levt's performance. importantly, our methods are orthogonal to existing research on nar architectures, indicating potential for widespread applicability. we explore two rl approaches: stepwise reward maximization, which computes rewards after each edit operation, and episodic reward maximization, which only computes rewards after all generations are completed. we analyze these two approaches' respective advantages and disadvantages and empirically verify them. furthermore, through a series of experiments, we investigate the impact of temperature settings on softmax sampling, aiming to identify the optimal temperature that strikes a balance between exploration and exploitation during the rl training process. ## background reinforcement learning reinforcement learning has been widely applied to improve the performance of ar nmt models @xcite @xcite because its ability to train models to optimize nondifferentiable score functions and tackle the exposure bias problem @xcite . in practice, reinforce (williams, 1992) with a baseline is commonly used for estimating the policy gradient, which can be computed as follows: softmax with temperature in the domain of rl, we need to consider the explorationexploitation trade-off @xcite , where temperature τ is an important parameter. τ is used to control the softness of the softmax distribution, a larger τ leads to a more uniform distribution, promoting exploration, while a smaller τ creates a more peaky distribution, emphasizing exploitation. @xcite shows that training with an increased temperature can mitigate the peakiness effect due to rl @xcite , indicating that a suitable temperature is significant for rl training in nmt. rl for nar compared to ar methods, studies of reinforcement learning for nar remain unexplored. @xcite proposed a method to reduce the estimation variance of reinforce by fixing the predicted word at position t and sampling words of other positions for n times. however, this method is only applicable to models with a fixed length, which is unsuitable for edit-based models. levenshtein transformer levenshtein transformer @xcite is an nar model based on three edit operations: delete tokens, insert placeholders, and replace placeholders with new tokens. it uses a supervised dual-policy learning algorithm to minimize the levenshtein distance @xcite for training and greedy sampling for decoding. the decoding stops when two consecutive refinement iterations return the same output or a max-figure 1 : the illustration of levenshtein transformer's decoding process @xcite . in each decoding iteration, three edit operations are performed sequentially: delete tokens, insert placeholders, and replace placeholders with new tokens. levt's dual-policy learning generates teacher data by corrupting the ground truth and reconstructing it with its adversary policy. this mechanism not only offers a unique approach to data generation but also underscores the inherent difficulty in preparing teacher data. this introduces concerns regarding the exposure bias, particularly whether the training process can maintain consistency with the text during decoding. to address this issue, we employ rl approaches that use self-generated data for training. ## approaches in this section, we present our reinforcement learning approaches in detail. we train a levenshtein transformer model as our baseline using the dualpolicy learning algorithm. based on it, we introduce two distinct rl approaches within the rein-force framework: stepwise reward maximization and episodic reward maximization. moreover, we present our methods for temperature control. stepwise reward maximization general rl training methods for ar nmt models are all episodic foot_0 , as it is difficult to calculate bleu @xcite when the sentence is not fully generated. in contrast, nar models can calculate bleu on outputs at each decoding step. from the perspective of estimating a more accurate gradient, we propose stepwise reward maximization, which calculates reward for each edit operation foot_1 using score differences from one previous edit. since every step's reward is calculated separately, this approach should be easier to learn than episodic approaches @xcite . however, it is also more prone to learning bias since the editing process is inherently multi-step. this drawback should not be emphasized since maximizing the reward for each step will likely maximize the episodic reward in nar models' training. we use a leave-one-out baseline (luo, 2020) for b(s) in equation 1 instead of the greedy baseline proposed in scst @xcite because the greedy decoding is too strong in levt, which makes gaining positive rewards in scst difficult and may reduce learning efficiency. for each edit, we sample k actions from the policy at this point. then, we calculate the baseline as follows: in a straightforward implementation, one might consider applying sampling again to all k samples from the last edit. however, this will cause a combination explosion when the number of edit operations increases. practically, we randomly choose a sample from the previous edit to perform the subsequent operations. we show an illustration of the sampling process in (a) of figure 2 and pseudo code of our algorithm in appendix a. ## conclusion and future work this paper explored the application of reinforcement learning to edit-based non-autoregressive neural machine translation. by incorporating rl into the training process, we achieved a significant performance improvement. by empirically comparing stepwise and episodic reward maximization, we analyzed the advantages and disadvantages of these rl approaches. we plan to have a deeper exploration of stepwise reward maximization and find a way to alleviate training inequality for multiple edit operations in the future. our investigation of temperature settings in nar softmax sampling provided insights into striking a balance between exploration and exploitation during training. although our annealing methods perform well, they are not optimal and still depend on manually adjusting the parameters such as total training steps. in the future, we plan to develop a self-adaption temperature control method using various indicators like entropy and advantage sd. the experiments in this paper focused on the basics, and we plan to do more study for practical applications in future work. as our methods are orthogonal to existing research on nar architectures, our next step involves exploring the methods' applicability across a broader spectrum, including state-of-the-art models. additionally, we plan to investigate how to effectively apply rl to the distillation model, the impact of different baseline designs on performance, and the impact of rl on output diversity. applying rl to nar is a massive and complex research question. we look forward to more researchers joining this topic.