Reinforcement Learning with Large Action Spaces for Neural Machine Translation

Asaf Yehudai, Leshem Choshen, Lior Fox, Omri Abend · COLING 2022 · 2022

In one sentence

Reinforcement learning helps neural machine translation little because the action space is the whole vocabulary; shrinking it — by using a 1K target vocabulary, or by initializing the decoder's final layer with BERT embeddings and freezing it during RL — recovers 1.5 BLEU on average over standard RL.

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.

Questions this paper answers

why does reinforcement learning barely improve a translation model that was already trained normally?
what limits the effectiveness of policy gradient fine-tuning in neural machine translation, and is the token-level action space the bottleneck?
how do I get a meaningful gain out of RL fine-tuning on top of an MLE-trained translation model?
is RL fine-tuning worth running on my translation model, or will the gain be a fraction of a BLEU point?
Yehudai et al. (2022) identify the size of the action space, spanning all tens of thousands of vocabulary tokens, as a main obstacle to reinforcement learning's effectiveness in neural machine translation.
Holds for: Four low-resource WMT pairs into English, MRT with BLEU as reward; prior work on large action spaces existed outside text generation.
Cutting the target BPE vocabulary from 17K-31K to 1K raises the BLEU gain that RL adds over MLE pretraining from 0.6/0.17/0.23/0.22 to 1.4/0.83/1.1/1.59 points on de-en, cs-en, ru-en and tr-en. That is about 1 BLEU point more improvement on average.
Holds for: Low-resource WMT News Commentary v13 (SETIMES2 for tr-en), gated convolutional encoder-decoder, MRT mixed with token-level loss at alpha=0.3, smoothed BLEU as reward; the 1K models' absolute BLEU is lower.
which research first argued that having tens of thousands of word choices is what makes reinforcement learning weak for translation?
where does the large-action-space explanation for the limited effect of RL in NMT originate?
what should I read first about why the vocabulary-sized action space holds back RL for machine translation?
Yehudai et al. (2022) identify the size of the action space, spanning all tens of thousands of vocabulary tokens, as a main obstacle to reinforcement learning's effectiveness in neural machine translation.
Holds for: Four low-resource WMT pairs into English, MRT with BLEU as reward; prior work on large action spaces existed outside text generation.
does a translation model with a much smaller set of output word pieces benefit more from reinforcement learning?
how do RL gains over MLE pretraining change between a 1K and a 17K-31K target BPE vocabulary?
how do I test whether the number of output tokens is what caps my RL fine-tuning gains in translation?
should I shrink my target subword vocabulary before RL fine-tuning a translation model?
Cutting the target BPE vocabulary from 17K-31K to 1K raises the BLEU gain that RL adds over MLE pretraining from 0.6/0.17/0.23/0.22 to 1.4/0.83/1.1/1.59 points on de-en, cs-en, ru-en and tr-en. That is about 1 BLEU point more improvement on average.
Holds for: Low-resource WMT News Commentary v13 (SETIMES2 for tr-en), gated convolutional encoder-decoder, MRT mixed with token-level loss at alpha=0.3, smoothed BLEU as reward; the 1K models' absolute BLEU is lower.
With a 1K target vocabulary, RL shifts probability mass upward from low-ranked tokens in neural machine translation. Within the first 100 ranks the small-vocabulary model reduces the probability of 83 ranks, against only 2 for the large-vocabulary model.
Holds for: Measured over 700K trials of gold-token rank under teacher-forced contexts, four language pairs into English, MRT training with smoothed BLEU as reward.
when reinforcement learning improves a translation model, does it lift word choices the model previously thought unlikely?
does RL fine-tuning redistribute probability mass toward low-ranked target tokens, or only sharpen already high-probability ones?
how do I tell which output tokens reinforcement learning actually promotes in my translation decoder?
With a 1K target vocabulary, RL shifts probability mass upward from low-ranked tokens in neural machine translation. Within the first 100 ranks the small-vocabulary model reduces the probability of 83 ranks, against only 2 for the large-vocabulary model.
Holds for: Measured over 700K trials of gold-token rank under teacher-forced contexts, four language pairs into English, MRT training with smoothed BLEU as reward.
BERT target-embedding initialization alone improves RL on all four language pairs into English, from 23.19/15.81/17.31/12.66 BLEU for plain RL to 24.44/17.04/18.68/14.37. Freezing MLE-learned target embeddings during RL instead gives only a slight gain.
Holds for: Low-resource WMT setting, gated convolutional NMT, MRT objective; plain RL over MLE is itself near-flat (no change on cs-en and ru-en), which is the baseline the BERT gain is measured against.
can pretrained language model word vectors in a translation decoder's output layer make reinforcement learning work better?
how much BLEU does initializing the decoder's final fully connected layer with BERT target embeddings and freezing it add over standard RL fine-tuning?
how do I shrink the effective action space of my translation decoder without cutting the vocabulary itself?
should I swap in BERT embeddings as my decoder output layer before RL fine-tuning?
Initializing the decoder's final fully connected layer with BERT embeddings and freezing it during RL reaches 24.71/17.37/18.30/14.55 BLEU on de-en/cs-en/ru-en/tr-en, 1.5 BLEU points above regular RL on average.
Holds for: Four low-resource pairs into English, BERT target vocabulary of 30,526, MRT with BLEU reward, MLE pretraining with frozen BERT embeddings; on ru-en freezing does not help (18.30 vs 18.68).
BERT target-embedding initialization alone improves RL on all four language pairs into English, from 23.19/15.81/17.31/12.66 BLEU for plain RL to 24.44/17.04/18.68/14.37. Freezing MLE-learned target embeddings during RL instead gives only a slight gain.
Holds for: Low-resource WMT setting, gated convolutional NMT, MRT objective; plain RL over MLE is itself near-flat (no change on cs-en and ru-en), which is the baseline the BERT gain is measured against.
if the output word vectors of a translation model are left untouched during reinforcement learning, does quality suffer?
what happens to BLEU and to the trainable parameter count when the target embedding layer is frozen during RL fine-tuning of a convolutional NMT model?
how do I cut the number of parameters trained during RL fine-tuning of a translation model without losing quality?
can I freeze my decoder's output embedding layer during RL to save trainable parameters?
Freezing the target embedding layer removes more than 60% of the network's trainable parameters, from 74.8M-77.2M down to 27.9M-30.2M across the four language pairs. BLEU improves rather than degrades under this reduction.
Holds for: Gated convolutional encoder-decoder with hidden size 768 and BERT's 30,526-token target vocabulary, four low-resource pairs into English; the frozen embeddings must be informative rather than random.
Freezing randomly initialized target embeddings during both MLE and RL degrades German-English translation by about 2 BLEU points. The benefit of freezing therefore depends on the quality of the embedding space rather than on parameter reduction.
Holds for: Low-resource NMT setup with gated convolutional encoder-decoder, BERT's 30,526-token target vocabulary and the MRT objective, four pairs into English.
is freezing a translation model's output word vectors helpful because there is less to learn, or because those vectors are already good?
does the benefit of a frozen target embedding layer in RL come from parameter reduction or from the quality of the embedding space?
how do I check whether freezing helps my policy for the right reason before I rely on it?
Freezing randomly initialized target embeddings during both MLE and RL degrades German-English translation by about 2 BLEU points. The benefit of freezing therefore depends on the quality of the embedding space rather than on parameter reduction.
Holds for: Low-resource NMT setup with gated convolutional encoder-decoder, BERT's 30,526-token target vocabulary and the MRT objective, four pairs into English.
In a contextual bandit where 10 real actions are each duplicated 400 times into 4000 policy-level actions, initializing the policy's last layer so duplicated actions share weights speeds up learning. Freezing that informative initialization speeds it up further.
Holds for: Synthetic 10-300-300-4000 feed-forward policy, 50 trials per agent, binary reward with Gaussian noise; the informative initialization encodes only which actions are duplicates and no information about which action is rewarding.
is there a small controlled experiment showing that having many duplicate choices slows down reward-based learning?
how does a contextual bandit with 10 real actions duplicated into 4000 motivate a shared-weight, frozen last layer for the policy?
how do I demonstrate the cost of redundant actions in policy gradient learning outside of machine translation?
In a contextual bandit where 10 real actions are each duplicated 400 times into 4000 policy-level actions, initializing the policy's last layer so duplicated actions share weights speeds up learning. Freezing that informative initialization speeds it up further.
Holds for: Synthetic 10-300-300-4000 feed-forward policy, 50 trials per agent, binary reward with Gaussian noise; the informative initialization encodes only which actions are duplicates and no information about which action is rewarding.
do people actually judge the improved translations as better, not just the automatic score?
was the BLEU gain from the BERT-initialized frozen-embedding RL model confirmed by human adequacy judgments and by a semantic similarity metric?
how do I verify that an RL fine-tuning gain in translation is real and not just reward gaming of BLEU?
should I trust the reported gains from frozen BERT target embeddings, or are they BLEU-only artifacts?
Two professional translators rating 100 translations per language pair on a 0-100 adequacy scale score the BERT-initialized, frozen-embedding RL model above baseline RL on all four language pairs. The Wilcoxon rank sum p-value is 8.5e-5.
Holds for: Both annotators native English speakers, judging how well each translation conveys the reference's information; significance is computed over the pooled score distributions of the two models, not per language pair.
On the SIM semantic-similarity metric the BERT-initialized frozen-embedding RL model scores 72.81/66.44/67.66/63.59 on de-en/cs-en/ru-en/tr-en against 71.17/63.29/66.17/59.99 for plain RL, with the largest gains on cs-en and ru-en.
Holds for: SIM as defined by Wieting et al. (2019), which gives partial credit to lexically different but semantically correct translations; BLEU, not SIM, was the reward optimized during RL.
do the gains from pretrained output word vectors show up on a meaning-based score and not only on the score being optimized?
what semantic similarity scores does the BERT-initialized frozen-embedding RL model reach compared with plain RL on de-en, cs-en, ru-en and tr-en?
how do I measure whether my RL-tuned translation model improved in meaning rather than only in the BLEU reward?
On the SIM semantic-similarity metric the BERT-initialized frozen-embedding RL model scores 72.81/66.44/67.66/63.59 on de-en/cs-en/ru-en/tr-en against 71.17/63.29/66.17/59.99 for plain RL, with the largest gains on cs-en and ru-en.
Holds for: SIM as defined by Wieting et al. (2019), which gives partial credit to lexically different but semantically correct translations; BLEU, not SIM, was the reward optimized during RL.
are pretrained word vectors better than ones a translation model learns itself at putting related words near each other?
how do cosine-similarity distributions for inflection, synonym and random word pairs differ between BERT target embeddings and MLE-learned ones?
how do I check whether my decoder's output embeddings generalize across similar target words?
if I want output embeddings that group related words, should I take BERT's or the ones my MLE training produced?
MLE-learned target embeddings give nearly identical cosine-similarity distributions for inflection pairs, synonym pairs and random word pairs, whereas BERT embeddings separate inflections from random pairs. BERT still places synonyms close to random pairs.
Holds for: Word pair lists compiled from WordNet and spaCy; comparison is between the MLE model's learned target embeddings in this low-resource NMT setup and BERT's embedding layer.
when pretrained word vectors are used as a translation model's output layer, should training be allowed to change them?
during MLE pretraining of an NMT decoder, do BERT target embeddings have to be frozen to give a BLEU gain?
should I keep BERT target embeddings frozen or trainable while pretraining my translation model with cross-entropy?
During MLE pretraining, BERT target embeddings must be frozen to help: frozen BERT embeddings reach 23.46/16.59/18.14/14.15 BLEU versus 22.99/15.32/17.57/12.65 when trainable, a gain of 0.47 to 1.50 points.
Holds for: Four low-resource pairs into English with BERT's 30,526-token vocabulary; the paper attributes the difference to catastrophic forgetting of BERT parameters when the layer is trainable.

Claims and scope

Common misreadings

Terminology in this paper

target embeddings
The rows of the final fully connected layer that maps a decoder's internal d-dimensional representation to the target vocabulary, viewed as embeddings of the output actions in inverse analogy to the input embedding layer.
informative initialization
Initializing a policy network's last layer so that weight vectors projecting to similar or duplicated actions start out identical or close, encoding the structure of the action space without any information about which actions are rewarding.
LTV / STV
Large target vocabulary (17K-31K BPE tokens) versus small target vocabulary (1K BPE tokens) on the target side of a translation model, with the source vocabulary unchanged.
SIM
A semantic similarity metric for translation that assigns partial credit to translations that are semantically correct but lexically different from the reference, introduced by Wieting et al. (2019).

How to cite

@inproceedings{yehudai2022reinforcement,
  author       = {Asaf Yehudai and
                  Leshem Choshen and
                  Lior Fox and
                  Omri Abend},
  editor       = {Nicoletta Calzolari and
                  Chu{-}Ren Huang and
                  Hansaem Kim and
                  James Pustejovsky and
                  Leo Wanner and
                  Key{-}Sun Choi and
                  Pum{-}Mo Ryu and
                  Hsin{-}Hsi Chen and
                  Lucia Donatelli and
                  Heng Ji and
                  Sadao Kurohashi and
                  Patrizia Paggio and
                  Nianwen Xue and
                  Seokhwan Kim and
                  Younggyun Hahm and
                  Zhong He and
                  Tony Kyungil Lee and
                  Enrico Santus and
                  Francis Bond and
                  Seung{-}Hoon Na},
  title        = {Reinforcement Learning with Large Action Spaces for Neural Machine
                  Translation},
  booktitle    = {Proceedings of the 29th International Conference on Computational
                  Linguistics, {COLING} 2022, Gyeongju, Republic of Korea, October 12-17,
                  2022},
  pages        = {4544--4556},
  publisher    = {International Committee on Computational Linguistics},
  year         = {2022},
  url          = {https://aclanthology.org/2022.coling-1.401},
  timestamp    = {Tue, 27 May 2025 07:57:03 +0200},
  biburl       = {https://dblp.org/rec/conf/coling/YehudaiCFA22.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}

References

See the full reference list in the paper.