Reinforcement Learning with Large Action Spaces for Neural Machine Translation

Proceedings of the 29th International Conference on Computational Linguistics, pages 4544–4556 October 12–17, 2022. 4544 Reinforcement Learning with Large Action Spaces for Neural Machine Translation Asaf Yehudai, Leshem Choshen, Lior Fox, Omri Abend Department of Computer Science The Hebrew University of Jerusalem {first.last}@mail.huji.ac.il Abstract Applying Reinforcement learning (RL) follow- ing maximum likelihood estimation (MLE) pre-training is a versatile method for enhanc- ing neural machine translation (NMT) perfor- mance. However, recent work has argued that the gains produced by RL for NMT are mostly due to promoting tokens that have already re- ceived a fairly high probability in pre-training. We hypothesize that the large action space is a main obstacle to RL’s effectiveness in MT, and conduct two sets of experiments that lend support to our hypothesis. First, we find that reducing the size of the vocabulary improves RL’s effectiveness. Second, we find that ef- fectively reducing the dimension of the action space without changing the vocabulary also yields notable improvement as evaluated by BLEU, semantic similarity, and human evalua- tion. Indeed, by initializing the network’s final fully connected layer (that maps the network’s internal dimension to the vocabulary dimen- sion), with a layer that generalizes over similar actions, we obtain a substantial improvement in RL performance: 1.5 BLEU points on aver- age.1 1 Introduction The standard training method for sequence-to- sequence tasks, specifically for NMT is to maxi- mize the likelihood of a token in the target sentence, given a gold standard prefix (henceforth, maximum likelihood estimation or MLE). However, despite the strong performance displayed by MLE-trained models, this token-level objective function is lim- ited in its ability to penalize sequence-level er- rors and is at odds with the sequence-level eval- uation metrics it aims to improve. One appealing method for addressing this gap is applying pol- icy gradient methods that allow incorporating non- differentiable reward functions, such as the ones 1https://github.com/AsafYehudai/Reinforcement- Learning-with-Large-Action-Spaces-for-Neural-Machine- Translation often used for MT evaluation (Shen et al., 2016, see §2). For brevity, we will refer to these methods simply as RL. The RL training procedure consists of several steps: (1) generating a translation with the pre- trained MLE model, (2) computing some sequence- level reward function, usually, one that assesses the similarity of the generated translation and a reference, and (3) updating the model so that its fu- ture outputs receive higher rewards. The method’s flexibility, as well as its ability to address the expo- sure bias (Ranzato et al., 2016; Wang and Sennrich, 2020), makes RL an appealing avenue for improv- ing NMT performance. However, a recent study (C19; Choshen et al., 2019) suggests that current RL practices are likely to improve the prediction of target tokens only where the MLE model has al- ready assigned that token a fairly high probability. In this work, we observe that one main differ- ence between NMT and other tasks in which RL methods excel is the size of the action space. Typi- cally, the size of the action space in NMT includes all tokens in the …

Content truncated to 500 words