Reinforced Curriculum Learning on
Pre-trained Neural Machine Translation Models

For the model design of the RL agent, we choose the Deterministic Actor-Critic algorithm.

Actor-critic (?) is a widely used method in reinforcement learning combining both an actor network $\mu_{\phi}$ that outputs an action $a=\mu_{\phi}(s)$ given a state $s$ to maximize the reward, and a critic network $Q_{w}$ that predicts the action value $Q_{w}(s,a)$ of a state-action pair $(s,a)$. The actor learns a near-optimal policy via policy gradient, while the critic estimates the action-value guiding the update direction of the actor. Compared with actor-only methods like REINFORCE (?), the existence of the critic reduces the update variance and accelerates convergence.

As our reward calculation involves evaluating the updated NMT model on a validation set and is thus expensive, we exploit a memory replay buffer to increase sample efficiency. Furthermore, we choose a deterministic policy setting as opposed to stochastic policy due to the fact that the deterministic policy gradient can be calculated much more efficiently as shown in (?).

The update of the framework is illustrated in Figure 2. The critic network $Q_{w}$ takes a state-action pair $(s,a)$, evaluates the action value, and outputs a predicted reward $\tilde{r}=Q_{w}(s,a)$, and updates the parameters supervised by the actual reward $r$ from the environment. Note that the critic network provides an immediate reward per iteration. As a result, we do not need to employ additional techniques, e.g., Temporal-Difference (TD) (?), to approximate the long-term reward. The objective of the critic network is thus to minimize the squared error of $\delta$ between the predicted reward and the actual reward, given by

The update of parameters $w$ is achieved through gradient descent as follows:

where $\alpha_{w}$ corresponds to the learning rate of parameters $w$.

The actor network $\mu_{\phi}$ takes in a state $s$ from the environment, applies the learned policy, and outputs a corresponding action $a$. The objective of the actor network is to learn an optimal policy generating the proper action to maximize the predicted reward $Q_{w}(s,a)$. Policy gradient is used to update the parameters $\phi$, i.e.,

Algorithm 1 summarizes the overall learning process of the proposed framework. In each round of data selection (with $K$ rounds in total), we first train the RL agent for an adequate number of iterations. In each iteration, we derive a state, an action, and a reward in lines 4–7 and update the actor network and the critic network in line 8 and line 9, respectively. After the RL agent is fully trained, we apply the learned policy to select a subset $D_{S}$ and use $D_{S}$ to update the NMT model $p_{\theta_{k}}(y|x)$ and move to the next round. Usually one or two rounds are sufficient.

Figure 3 demonstrates the network structure of the RL agent. A feature network is shared between the Actor Network $\mu_{\phi}$ and the Critic Network $Q_{w}$. It takes in the raw features of the sampled batch of $b$ examples, where each feature of each sample goes through an independent single-layer MLP. The concatenation of the outputs constitutes the state representation $s$. Note that different examples in the sampled batch share the same network weights.

The Actor Network $\mu_{\phi}$ is composed of a two-layer MLP and computes a score for each example in the sampled batch of $b$ examples based on the input state $s$, and outputs the action $a$ as a probability vector representing the probability of each example being selected, by taking a softmax operation over the computed scores of $b$ examples.

The Critic Network $Q_{w}$ also has two layers and calculates the action value of a given state-action pair $(s,a)$, where the action $a$ is the output of the Actor Network, i.e., $a=\mu_{\phi}(s)$, and is concatenated to the second layer of the critic network.

Note that although the feature network is shared between both the actor network $\mu_{\phi}$ and the critic network $Q_{w}$, we only update it with the critic network to reduce training instability.

The state $s$ is meant to be a full summarization of the environment including a data batch of $b$ examples to select from and information about the pre-trained model. However, the number of parameters in the pre-trained model is too large to be included in the state directly at each time step. Thus, we need to find a representation that can represent both the samples to be selected and the existing model using a limited number of parameters. In our state design, we focus on three different dimensions, namely informativeness, uncertainty, and diversity. We use the sentence length as a measure for informativeness, the sentence-level log-likelihood for uncertainty, and the $n$-gram rarity together with NER and POS taggings for diversity.

The most intuitive representation feature of a parallel sentence is the sentence length, i.e. the number of tokens in a sentence. This simple scalar roughly measures the amount of information involved in a sentence and is also used in (?) as a “difficulty” estimate.

Following the intuition that examples yielding large uncertainty can benefit the model performance (?), another feature we use is the sentence-level log-likelihood $L_{p_{\theta}(y|x)}$ calculated by the pre-trained model $p_{\theta}(y|x)$ by summing up the log probability of each word $y_{i}$ of the target sentence:

(?) and (?) suggest that selecting samples that are farthest away from the existing dataset can benefit model training. Incorporating this idea, we further utilize two other feature vectors, $n$-gram rarity and taggings, to represent the similarity between a given sample and the entire training set.

For $n$-gram rarity, we use $n=\{1,2,3,4\}$. Given all the sentences in $D_{T}$, we define the relative frequency for a unique $n$-gram $g^{(n)}_{j}$ in $D_{T}$ as

where $j=1,...,\#\textit{\{unique n-grams\}}$, $s_{k}$ is a sentence from $D_{T}$, $g^{(n)}_{(s_{k})}$ is all $n$-grams in $s_{k}$, and $N^{(n)}$ is the total number of $n$-grams in $D_{T}$. For $n$-gram representation of a sentence $s_{k}$, we form all the $n$-gram frequencies into a vector:

In this paper, we use $n=\{1,2,3,4\}$ and calculate the $n$-gram vectors for both the source and target sentences.

For taggings, we use Named Entity Recognition (NER) and Parts of Speech (POS) and apply ideas similar to our $n$-gram design. As the tagging of a word is dependent of the sentence that the word lies in, a same word may be tagged differently in various sentences. To give an example in POS tagging, the word “Book” can either be tagged as “NOUN” or “VERB” depending on its meaning in the sentence. If most occurrences of the word “Book” are tagged as “NOUN” in the training set, the model may not be able to correctly learn the second meaning of it. Therefore, the model should be fed more with samples in which “Book” is tagged as “VERB”. In order to reflect this phenomenon, we define a tagging value for a word $w$ and its current tag $t$ as

where $N_{w}$ is the number of times the word $w$ has appeared in the training set and $N_{tag(w)=t}$ is the number of times the word $w$ is tagged with tag $t$ among all its occurrences. Similarly, for both NER and POS taggings, the tagging values of words form a vector representation of a sentence:

In the task of curriculum learning, an action represents the process of data selection. (?) assume a stream-based setting where data examples come one by one in a stream, and design their action as making a decision on whether or not a single incoming data example should be selected. We argue that our problem setting is actually pool-based instead of stream-based, where a pool of data exists for selection and deciding on the selection of each individual sample per step would be inefficient. In (?), a dataset is split into several bins according to a noise measure of data samples, and the action determines from which bin the next batch of training data should be selected. However, this method highly depends on an effective heuristic criterion for bin formation and is thus hard to generalize.

Therefore, we propose our action design which samples a batch of $b$ data examples from $D_{T}$, computes a score for each example according to the trained policy, and choose the one with the highest score. Correspondingly, in our design, we can easily control the size of $D_{S}$ by varying the batch size $b$, i.e., $|D_{S}|=|D_{T}|/b$, since we choose one out of $b$ samples for each batch.

For the choice of reward signal, we use the performance improvement of the NMT model evaluated on the validation set after it is updated with the selected sample. Perplexity is used as the performance metric instead of BLEU as it is shown to be more consistent and less noisy (?). We assign a reward of $0$ to unselected samples.

To compare our proposed method with other curriculum learning methods on NMT task, we conduct comprehensive empirical analysis on several zh-en translation datasets:

• •

  MTD is an internal news translation dataset with 1M samples in the training set and 1,892 samples in both the validation set and the test set. The average length of source sentences is $19.55$ and the average length of target sentences is $21.06$.

• •

  CASICTB, CASIA2015, NEU are three independent datasets with 1M, 2M, and 2M examples from different data sources in WMT18 which is a public translation dataset in news area with more than 20M samples. We only use a part of data from WMT18 to evaluate our method. All three datasets share the same validation set newsdev2017 and the test set newstest2017 both composed of 2k samples.

Table 2 summarizes the details of the experimental datasets. The columns titled “Train”, “Val”, and “Test” correspond to the number of examples in training, validation, and test sets. “Src-Len” and “Tgt-Len” correspond to the average sentence length of source language and target language respectively.

For evaluation metrics, we use Perplexity for the reward calculation in RL agent training, and report BLEU (?) for the final performance of the NMT models.

We implement our models in PyTorch 1.1.0 (?) and train the model with a single Tesla P40. We utilize NLTK (?) to perform POS and NER taggings.

For the NMT model, we use the OpenNMT implementation (?) of Transformer (?). It consists of a 6-layer encoder and decoder, with 8 attention heads, and 2,048 units for the feed-forward layers. The multi-head attention model dimension and the word embedding size are both set to 512. During training, we use Adam optimizer (?) with a learning rate of 2.0 decaying with a noam scheduler and a warm-up steps of 8,000. Each training batch contains 4,096 tokens and is selected with bucketing (?). During inference, we employ beam search with a beam size of 5.

For the RL agent, we use an Deterministic Actor-Critic architecture and build our system based on (?). In our framework, we use several tricks proven to be effective to RL training including a memory replay buffer of size 2,500, a warm-up phase of 500 steps, and a target network which is updated by mixing weights with the on-line network with a mix factor of $0.1$. For calculating rewards, we train the NMT model with the single selected sample using SGD and a learning rate of 1e-4.

The feature network maps data features of sentence length, sentence-level log-likelihood, taggings, and $n$-gram rarity to vectors of size 1, 8, 16, and 32 respectively with a FC layer, and concatenates them together as a shared state representation. The actor network is composed of two FC layers with hidden sizes of 300 and 400. The critic network is designed the same as the actor network except that the output action from actor network is concatenated to the second layer. Relu is used as the activation function in each FC layer in this network.

In experiments, we conduct two rounds of RL agent training and data selection. In each round, the RL agent is trained for 20k steps and the best model with the highest sum of rewards during the last 1,000 steps is used for data selection. The RL agent keeps selecting data given randomly sampled batches from the training set $D_{T}$ and feed the selected data to the NMT model until no performance improvement is observed. For the training process on selected data, we keep the NMT model’s optimizer and learning rate setting unchanged.

We use different batch sizes of the sampled batch $b$ for the two rounds of training and selection with $b_{1}=16$ and $b_{2}=128$ indicating in the first round we select 1 sample from 16 and in the second round, from 128 samples we select the best one. This is because we think in order to achieve improvement further on the basis of the first round, a stricter data selection criterion must be applied.

To make comparisons with other existing curriculum methods, we have conducted several baseline experiments.

We take the core ideas of existing curriculum learning methods of training on data samples with gradually increasing difficulty (?) and gradually decreasing noise (?) and apply them to our setting with pre-trained models. We evaluate the following three baseline methods along with our proposed method.

• •

  Denoising is a curriculum learning method of training an NMT model in a noise-annealing fashion (?). They propose to measure NMT data noise with the help of a trusted dataset which contains generally cleaner data compared to the training dataset. (?) also utilize data noise in their curriculum design and achieve similar performance as (?). For the choice of the trusted dataset, we choose a subset of 500 sentences from the validation set newsdev2017 of CASICTB, CASIA2015 and NEU.

• •

  Sentence Length is an intuitive difficulty measure used in (?), since longer sentences tend to contain more information and more complicated sentence structure.

• •

  Word Rarity is another metric for measuring the sample difficulty, as rare words appear less frequently in the training process and should be presented to the learning system more. The formula for calculating the word rarity of a sentence can be found in (?).

For baseline experiments, the pre-trained NMT model is further trained on 20% of the original data, which are selected by the above criteria, i.e., the least noisy, the longest, and the highest word rarity, respectively.

Table 3 compares the performance of our method with other baseline methods on different datasets evaluated using BLEU. The result shows that our proposed method significantly out-performs other baseline methods by a great margin. We conduct two rounds of training and update in our experiments. While the result of the first round surpasses almost all the baseline methods, our second round further improves the performance and achieves a final BLEU improvement of $+0.90$, $+0.60$, $+0.59$, and $+0.71$ on MTD, CASICTB, CASIA2015, and NEU respectively over the pre-trained model.

The reason of our success is due to our utilization of an RL framework to proactively select data samples that are potentially beneficial to the training of the NMT model. First, we formulate the task of curriculum learning on pre-trained NMT models as a reinforcement learning problem. Second, we construct an effective design of state, action and reward. Our state representation includes features of different dimensions of informativeness, uncertainty and diversity. Third, we propose a Deterministic Actor-Critic framework that learns a policy to select the best samples from the training set to improve the pre-trained model. By incorporating all these designs together, our proposed framework is able to achieve a significant performance enhancement on the pre-trained NMT model.

We evaluate the impact of different modules and methods by ablation test on MTD dataset. Table 4 list the performance of our model variants with different features included.

We incrementally accommodate different features of examples to the state by first starting from sentence length and sentence-level log-likelihood as they are both scalars. The performance increased slightly by $0.26$ BLEU compared with the pre-trained base model. Then we further accommodate $n$-gram rarity and POS and NER taggings to the state vectors, and observe a larger increase of performance of $0.42$ and $0.60$ respectively. Finally, we incorporate the second round of RL agent training and data selection on the basis of the result of first round, and achieve the best performance with a $0.90$ BLEU increase. Note that a stricter selection policy is applied to the second round (128 choose 1) compared with the first round (16 choose 1).