Reinforcement Learning with Large Action Spaces for Neural Machine Translation (2022)

Abstract: Applying Reinforcement learning (RL) following maximum likelihood estimation (MLE) pre-training is a versatile method for enhancing neural machine translation (NMT) performance. However, recent work has argued that the gains produced by RL for NMT are mostly due to promoting tokens that have already received 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 effectively reducing the dimension of the action space without changing the vocabulary also yields notable improvement as evaluated by BLEU, semantic similarity, and human evaluation. Indeed, by initializing the network’s final fully connected layer (that maps the network’s internal dimension to the vocabulary dimension), with a layer that generalizes over similar actions, we obtain a substantial improvement in RL performance: 1.5 BLEU points on average.

## introduction the standard training method for sequence-tosequence tasks, specifically for nmt is to maximize 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 limited in its ability to penalize sequence-level errors and is at odds with the sequence-level evaluation metrics it aims to improve. one appealing method for addressing this gap is applying policy gradient methods that allow incorporating nondifferentiable reward functions, such as the ones often used for mt evaluation @xcite . 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 pretrained mle model, (2) computing some sequencelevel reward function, usually, one that assesses the similarity of the generated translation and a reference, and (3) updating the model so that its future outputs receive higher rewards. the method's flexibility, as well as its ability to address the exposure bias @xcite , makes rl an appealing avenue for improving nmt performance. however, a recent study (c19; @xcite suggests that current rl practices are likely to improve the prediction of target tokens only where the mle model has already assigned that token a fairly high probability. in this work, we observe that one main difference between nmt and other tasks in which rl methods excel is the size of the action space. typically, the size of the action space in nmt includes all tokens in the vocabulary, usually tens of thousands. by contrast, common rl settings have either small discrete action spaces (e.g., atari games @xcite ), or continuous action spaces of low dimension (e.g., mujoco @xcite and similar control problems). intuitively, rl takes (samples) actions and assesses their outcome, unlike supervised learning (mle) which directly receives a score for all actions. therefore, a large action space will make rl less efficient, as individual actions have to be sampled in order to assess their quality. accordingly, we experiment with two methods for decreasing the size of the action space and evaluate their impact on rl's effectiveness. we begin by decreasing the vocabulary size (or equivalently, the number of actions), conducting experiments in low-resource settings on translating four languages into english, using bleu both as the reward function and the evaluation metric. our results show that rl yields a considerably larger performance increase (about 1 bleu point on average) over mle training than is achieved by rl with the standard vocabulary size. moreover, our findings indicate that reducing the size of the vocabulary can improve upon the mle model even in cases where it was not close to being correct. see §4. however, in some cases, it may be undesirable or unfeasible to change the vocabulary. we therefore experiment with two methods that effectively, reduce the dimensionality of the action space without changing the vocabulary. we note that generally in nmt architectures, the dimensionality of the decoder's internal layers (henceforth, d) is significantly smaller than the target vocabulary size (henceforth, |v t |), which is the size of the action space. a fully connected layer is generally used to map the internal representation to suitable outputs. we may therefore refer to the rows of the matrix (parameters) of this layer, as target embeddings, mapping the network's internal low-dimensional representation back to the vocabulary size, the actions. we use this term to underscore the analogy between the network's first embedding layer, mapping vectors of dimension |v t | to vectors of dimension d, and target embeddings that work in an inverse fashion. indeed, it is often the case (e.g., in bert, @xcite that the weights of the source and target embeddings are shared during training, emphasizing the relation between the two. using this terminology, we show in simulations ( §5.1) that when similar actions share target embeddings, rl is more effective. moreover, when target embeddings are initialized based on highquality embeddings (bert's in our case), freezing them during rl yields further improvement still. we obtain similar results when experimenting on nmt. indeed, using bert's embeddings for target embeddings improves performance on the four language pairs, and freezing them yields an additional improvement on both mle and rl as reported by both automatic metrics and human evaluation. both initialization and freezing are novel in the context of rl training for nmt. moreover, when using bert's embeddings, rl's ability to improve performance on target tokens to which the pre-trained mle model did not assign a high probability, is enhanced ( §5.2). ## methodology architecture. we use a similar setup as used by @xcite , adapting their fairseq-based @xcite codebase to our purposes. 2 similar to their transformer architecture we use gated convolutional encoders and decoders @xcite . we use 4 layers for the encoder and 3 for the decoder, the size of the hidden state is 768 for all layers, and the filter width of the kernels is 3. additionally, the dimension of the bpe embeddings is set to 768. ## reducing the vocabulary size we begin by directly testing our hypothesis that the size of the action space is a cause for the long convergence time of rl for nmt. to do so, we train a model with target-side bpe taken from a much smaller vocabulary than is typically used. we begin by training two mle models, one with a large (17k-31k) target vocabulary (ltv) and another with a target vocabulary of size 1k (stv). the source vocabulary remains unchanged. we start with the mle pretraining and then train each of the two models with rl. results ( we next turn to analyze what tokens are responsible for the observed performance gain. specifically, we examine whether reducing the vocabulary size resulted in rl being able to promote target tokens that received a low rank by the pretrained mle model. figure 1 presents the probability difference per rank for ltv and stv. we can see that for the first rank the probability shift due to rl training with stv is more than twice the shift caused by rl training with ltv. consequently, the probability shift for the following ranks is usually more negative for small vocabulary settings. the figure indicates that indeed the shift of probability mass to higher positions occurs substantially more when we apply rl using a smaller action space. moreover, the stv training was able to shift probability mass from lower ranks upwards compared to ltv. an indication for that is that, within the first one hundred ranks, stv reduces the probability of 83 of them, whereas ltv of only 2. ## reducing the effective dimensionality of the action space finding that reducing the number of actions improves rl's performance, we propose a method for reducing the effective number of actions, without changing the actual output. the vocabulary size might be static, as in pre-trained models @xcite , and reducing it might help rl but be sup-optimal for mle (gowda and may, 2020), or introduce out-of-domain words @xcite . we propose to do so by using target embeddings that generalize over tokens that appear in similar contexts. we explore two implementations of this idea, one where we initialize the target embeddings with high-quality embeddings, and another where we freeze the learned target embeddings during rl. we also explore a combination of the two approaches. freezing the target embeddings (decoder's last layer) can be construed as training the network to output the activations of the penultimate layer, where a fixed function then maps it to the dimension of the vocabulary. we note that although freezing is a common procedure @xcite @xcite , as far as we know, it has never been applied in the use of rl for sequence to sequence models. denote the function that the network computes with f θ . f θ can be written as h θ 2 • g θ 1 , where θ = (θ 1 , θ 2 ), g maps the input -source sentence x and model translation prefix @xmath0, and using this notation, we can formulate the method as loading pre-trained mle target embeddings to h θ 2 or freezing it (or both). as for many encoderdecoder architectures (including the transformer), it holds that d |v t |, this can be thought of as constraining the agent to select a d-dimensional continuous action, where h θ 2 is a known transformation performed by the environment. the importance of target embedding is that they allow for better generalization over actions. the intuition is as follows. assume two tokens have the same embedding, and similar semantics, i.e., they are applicable in the same contexts (synonyms). since they have the same target embeddings, during training the network will perform the same gradient updates when encountering either of them, except for in the last layer (since they are still considered different outputs). if the target embeddings are not frozen, encountering either of them during training will lead to very similar updates (since they have the same target embeddings), but their target embeddings may drift slightly apart, which will cause a subsequent drift in the lower layers. if the target embeddings are frozen, the gradient updates they will yield will remain the same and expedite learning. we hypothesize a similar effect during training, where tokens that have similar (but not identical) embeddings, and a similar (but not identical) distribution would benefit in training from each other. this motivates us to explore a combination of informative initialization and parameter freezing. (see formal proof in supp. §e). ## human evaluation we perform human evaluation, comparing the baseline rl with our proposed model. we selected 100 figure 3 : comparison of the change in the rank distribution of the target token following rl in two settings, one where rl training with frozen bert embeddings is used (blue) and the second when we used basic rl training (red). the gain in probability in the first rank indicates that the model is more probable to be correct (which is reflected in its superior performance over the pre-trained mle model). the negative values in the following places demonstrate how rl with frozen highquality target embeddings can improve not only when the mle model is initially close to being correct. ## comparing target embedding spaces the previous section discussed how bert's target embeddings improve rl performance, compared to target embeddings learned by mle. we now turn to directly analyze the generalization ability of the two embeddings. we do so by comparing the embeddings of semantically related words. we use wordnet @xcite and spacy 8 to compile three lists of word pairs: inflections (e.g., 'documentaries' / 'documentary', 'boxes' / 'box', 'stemming' / 'stem'), synonyms (e.g., 'luckily' / 'fortunately', 'amazement' / 'astonishment', 'purposely' / 'intentionally'), and random pairs, and compare the embeddings assigned to these pairs using bert and mle embeddings. figure 5 presents the distributions of the cosine similarity of the pairs in the three lists for both embedding spaces. results show that mle embeddings for the different lists have almost identical distributions, demonstrating the limited informativeness of these target embeddings. in contrast, bert embeddings only display a small overlap between the similarity disthe ability of the embeddings to distinguish between these three groups is informative of their ability to map semantically related words to similar embeddings. the better discrimination ability of bert embeddings is thus likely related to their superiority as target embeddings over mle embeddings. ## conclusion in this paper, we addressed the limited effectiveness of rl for nmt, seeking to understand its origins and offer means for tackling it. we hypothesized that this limitation arises from the size of the action spaces used in nmt and examined two ways of reducing their effective dimension. in the first method, we experiment with smaller vocabularies, showing improved rl effectiveness. while this method constrains the size of the vocabulary, which may be limiting in some settings @xcite , it motivates further research along these lines. the second approach introduces a new method of using informative target embeddings and potentially freezing them during rl. we find that this method may be beneficial as well, but its effectiveness crucially depends on the quality of the employed embeddings. indeed, we find using both simulations and nmt experiments that freezing in itself results in some improvement in rl performance, but that combined with target embeddings that generalize over words with a similar distribution, it may yield substantial gains as shown by bleu, semantic similarity, and human evaluation. we compare the target embeddings produced by mle and those by bert, finding the latter to be considerably stronger. those results in low resources settings, encourage further research aiming to address the problem of large action space for tg in richer data settings by adapting and extending our methods. future work will increase the exploration ability of rl training in nmt. a promising line of research towards this goal is using off-policy methods. off-policy methods, in which observations are sampled from a different policy than the one we currently optimize, are prominent in rl @xcite , and were also studied in the context of policy gradient methods @xcite . we believe that the adoption of such methods to enhance exploration, combined with our proposed method for using target embeddings, can be a promising path forward for the application of rl in nmt, and more generally in tg. a different line of future work will focus on changing the network's architecture to predict a d dimension continuous action, instead of discrete actions. such an approach may directly reduce the size of the action space without limiting the number of words that can be predicted.