SlowGuess/NAACL_2022_OCR huggingface dataset:
Atijit Anuchitanukul<sup>1</sup>, Julia Ive<sup>2</sup>
Imperial College London<sup>1</sup>, Queen Mary University of London<sup>2</sup>
atijit.anuchitanukul20@imperial.ac.uk
j.ive@qmul.ac.uk
The performance of Reinforcement Learning (RL) for natural language tasks including Machine Translation (MT) is crucially dependent on the reward formulation. This is due to the intrinsic difficulty of the task in the high-dimensional discrete action space as well as the sparseness of the standard reward functions defined for limited set of ground-truth sequences biased towards singular lexical choices. To address this issue, we formulate SURF, a maximally dense semantic-level unsupervised reward function which mimics human evaluation by considering both sentence fluency and semantic similarity. We demonstrate the strong potential of SURF to leverage a family of Actor-Critic Transformer-based Architectures with synchronous and asynchronous multi-agent variants. To tackle the problem of large action-state spaces, each agent is equipped with unique exploration strategies, promoting diversity during its exploration of the hypothesis space. When BLEU scores are compared, our dense unsupervised reward outperforms the standard sparse reward by  $2\%$  on average for in- and out-of-domain settings.
Reinforcement Learning (RL) has shown promise in the field of text generation. This is mainly due to the fact that it allows the usage of nondifferentiable evaluation functions fit for the discrete natural language tasks. It also serves as a solution for bridging the gap between training and inference time regimes ("exposure bias") that arises from the fact that the model is never exposed to its own errors as only ground-truth labels are used to condition the generation during training (Wang and Sennrich, 2020). One of the essential components of the RL framework is the reward function, which is used to provide agents with indicative signals in terms of the effectiveness of chosen actions.
The usage of RL in Neural Machine Translation (NMT) and language generation however has been doubted largely due to the difficulty of exploration in the high-dimensional discrete action space combined with the sparse reward signal. The latter comes from the typical metrics used as rewards (e.g., BLEU (Papineni et al., 2002)). These rewards evaluate text in a shallow way by measuring the string similarity between generated and ground-truth sequences, making them extremely sparse and biased towards singular lexical choices. As the RL policy is usually initialised with some pre-trained distribution over words, suspicion has been raised that in this situation, those words already most likely gain probability mass regardless of the rewards (Choshen et al., 2020). Thus, the current sparse rewards are not beneficial for rigorous exploration of different words during training. Recent studies suggest that the main benefit for NMT from RL is in performing domain adaptation when using proper hypothesis space exploration along with special emphasis on reward scaling and normalisation (Kiegeland and Kreutzer, 2021).
To address the problem posed by sparse or biased rewards, we propose SURF, a formulation of the unsupervised reward function that evaluates machine-generated texts in the semantic space by factoring in different qualitative aspects. Furthermore, we introduce an additional scaling and normalisation mechanism which ensures fairness and uniformity of the reward function regardless of the complexity of the natural language task.
Our main contributions are thus threefold: (a) the proposal of SURF, an unsupervised dense reward assessing both sentence fluency and adequacy (Section 4). We demonstrate this reward leads to a translation quality favourably comparable to the traditional sparse BLEU reward both in automatic and human evaluation; (b) the proposal of an additional normalisation using reward shaping mechanisms for the unsupervised reward; (c) demonstra
Human Language Technologies, pages 4508 - 4522
tion of the strong potential of the proposed reward to elicit benefits of various RL architectures. We experiment with multi-agent synchronous and asynchronous Actor-Critic (AC) architectures as applied to the problem of MT (Section 3.3). Each of the parallel agents in those architectures is trained using different segments of the training dataset which has their own unique exploration strategy. When BLEU scores are compared, our dense unsupervised reward outperforms the standard sparse reward by  $2\%$  on average for both in- and out-of-domain settings. To the best of our knowledge, this formulation of the unsupervised reward for a range of multi-agent architectures is the first of its kind for MT.
Our datasets and settings are described in Section 5, and results of our experiments are described in Section 6.
The following section describes the work related to ours in the subfields of Machine Translation and RL for language generation.
Reinforcement Learning Algorithms for Neural Machine Translation REINFORCE (Williams, 1992) algorithm and its variants have so far been the most widely used RL algorithms in MT (Ranzato et al., 2015; Rennie et al., 2017; Paulus et al., 2018; Hu et al., 2018). The fact that REINFORCE-based approaches suffer from high variance in general and in MT in particular has stimulated attempts to apply Actor-Critic (AC) models to the task. The first attempt of the kind was the one of Bahdanau et al. (2016). More advanced AC models with Q-Learning are rarely applied to language generation problems. However, there are exceptions (e.g., entropy-regularised AC models that promote exploration of actions (Dai et al., 2018; Ive et al., 2021)). This could be explained by the difficulty of approximating the Q-function for large action space. In this work we explore a series of multi-agent AC architectures which to the best of our knowledge have never been applied to MT before.
Unsupervised Rewards for Language Generation Tasks Recent work on unsupervised rewards in NLP has explored both dynamic (Ive et al., 2021) and static rewards (Gao et al., 2020; Garg et al., 2021). For example, Ive et al. (2021) introduces a dynamic distribution over latent frequency classes as a reward signal. This distribution is shaped to promote more rare words in the policy search space.
Static rewards are very often designed to assess generated text in terms of its fluency and adequacy. Fluency judgment assesses how a hypothesis satisfies the grammatical norms of a language. Adequacy judgment assesses how well a hypothesis conveys the meaning of the source sentence. The recent research performs both evaluations as semantic similarity assessments using the pre-trained contextualised embeddings such as BERT (Zhang et al., 2020; Mathur et al., 2019). For MT, semantic similarity assessment could be carried out using monolingual pre-trained embeddings against a reference, as in Gao et al. (2020), or using multilingual pre-trained embeddings in the unsupervised reference-less approach by considering the semantic similarity to source sentences (Wei et al., 2019; Song et al., 2021). We adopt the latter approach to measure adequacy.
We start by formulating MT using RL, then introduce the Actor-Critic architectures and the reward functions used.
3.1 Neural Machine Translation (NMT)
A typical Neural Machine Translation (NMT) system is a Seq2Seq architecture (Sutskever et al., 2014; Bahdanau et al., 2014), where each source sentence  $(\mathbf{X})$  is encoded by the encoder into a sequence of hidden states. At each decoding step  $t$ , a target word  $y_{t}$  is generated according to  $p(y_{t}|\mathbf{y}_{< t},\mathbf{X})$  conditioned on the input sequence  $\mathbf{X}$  and decoded sequence  $\mathbf{y}_{< t} = (y_1,\dots ,y_{t - 1})$  up to the  $t$ -th time step:
$$
\mathcal {L} _ {m l e} = \log p (y _ {t} | \mathbf {y} _ {<   t}, \mathbf {X}) \tag {1}
$$
3.2 Reinforcement Learning for NMT
In the RL framework, a Seq2Seq model is viewed as an agent and its parameters define the agent's policy  $(\pi)$ . At each timestep  $(t)$ , the agent observes the current state  $(s_t)$  of the environment, which is essentially the sequence of generated words from previous timesteps  $(\hat{\mathbf{y}}_{1:t-1})$ . Then, the agent's policy, which is based on the conditional probability  $p(\hat{y}_t|\hat{\mathbf{y}}_{1:t-1},\mathbf{X})$ , is used to select an action  $(a_t)$  at each timestep. In this context, an action is the candidate word  $(\hat{y}_t)$  chosen from the vocabulary. Subsequently, the environment adds the chosen word to the generated sequence, transitioning it to the next state  $(\hat{\mathbf{y}}_{1:t})$ . It also returns a reward  $(r_{t+1})$  to the agent as an indication of how effective the chosen
word is. Hence, one possible training objective of the policy is to maximise the discounted sum of expected rewards from all timesteps:
$$
\pi_ {*} = \max _ {\pi} \sum_ {t = 1} ^ {T} \gamma^ {t - 1} \underset {\hat {y} _ {t} \sim \pi (\cdot | \hat {\mathbf {y}} _ {1: t - 1}, \mathbf {X})} {\mathbb {E}} [ r _ {t + 1} ] \quad (2)
$$
where  $\pi_*$  denotes the optimal policy and  $\gamma$  is a constant discount factor. Under the policy  $\pi$ , one can formulate two functions: the state value function  $(V_{\pi}(s_t))$  and the state-action value function  $(Q_{\pi}(s_t, a_t))$ . The former,  $V_{\pi}(s_t)$ , determines the effectiveness of the agent being in a particular state while the latter,  $Q_{\pi}(s_t, a_t)$ , indicates the effectiveness of selecting a certain action in that state:
$$
V _ {\pi} (s _ {t}) = \underset {\hat {y} _ {t} \sim \pi} {\mathbb {E}} \left[ Q _ {\pi} \left(s _ {t} = \hat {\mathbf {y}} _ {1: t - 1}, a _ {t} = \hat {y} _ {t}\right) \right] \quad (3)
$$
$$
Q _ {\pi} (s _ {t}, a _ {t}) = \mathbb {E} _ {\pi} \left[ \sum_ {k = 1} ^ {T - t} \gamma^ {k - 1} r _ {t + k} | s _ {t}, a _ {t} \right] \tag {4}
$$
Hence the definition of the advantage function is:
$$
A _ {\pi} (s _ {t}, a _ {t}) = Q _ {\pi} (s _ {t}, a _ {t}) - V _ {\pi} (s _ {t}) \qquad (5)
$$
A training objective can aim to maximise the advantage function  $\max_{a} A_{\pi}(s_t, a_t)$ .
Alternatively, considering the definition of  $A_{\pi}(s_t,a_t)$  in Equation (5), it implies that we can directly maximise the  $Q$  function:
$$
\max  _ {a} A _ {\pi} \left(s _ {t}, a _ {t}\right)\rightarrow \max  _ {a} Q _ {\pi} \left(s _ {t}, a _ {t}\right) \tag {6}
$$
The first objective (Equation (2)) has been used in REINFORCE-based methods (Sutton et al., 2000) such as the MIXER architecture (Ranzato et al., 2015). These methods sample trajectories, series of consecutive states, actions and rewards, and use their true returns to update the policy. As they use the true returns, they are considered to be unbiased. However, as an action in a certain state can be part of many trajectories with different returns, these methods are considered to have high variance (Sutton and Barto, 2014). To address this issue, Actor-Critic algorithms (Konda and Tsitsiklis, 2001) adopt the Temporal Difference (TD) learning method which performs bootstrapping by using only the immediate reward and estimated values to guide future action selection. The training objective used is Equation (6).
3.3 Actor-Critic Architectures
3.3.1 Actor-Critic with Q-Learning (ACQ) Model
An Actor-Critic model usually consists of an actor and a critic (Konda and Tsitsiklis, 2001). The two networks are neural networks parameterised by  $\theta$  and  $\phi$ , respectively. The actor acts as the policy of the model while the critic is a function approximation network. One simple variant (ACQ) of the AC architecture is trained by maximising the  $Q$  function. In this variant, the critic is defined as a  $Q$  network approximating the true  $Q$  function.
The actor's training objective is to maximise the probability of actions that yield high  $Q$  values. Using  $Q$  value estimates  $(Q_{\phi}(\hat{\mathbf{y}}_{1:t - 1,i},w))$  computed by the main critic, the actor's policy loss  $(\mathcal{L}_{policy})$  at each training timestep can be expressed as follows:
$$
\mathcal {L} _ {\text {p o l i c y}} = - \left[ \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {t} \sum_ {w \in \mathcal {W}} \right. \tag {7}
$$
$$
\pi_ {\pmb {\theta}} (w | \hat {\mathbf {y}} _ {1: t - 1, i}) Q _ {\phi} (\hat {\mathbf {y}} _ {1: t - 1, i}, w) ]
$$
where  $N$  denotes the training batch size. The loss is calculated by summing over all the possible actions  $(w)$  in the entire vocabulary  $(\mathcal{W})$ . Following (Bahdanau et al., 2017), to avoid early policy determination and gradient vanishing issues, the final actor loss  $(\mathcal{L}_{ACQ - actor})$  consists of the policy loss  $(\mathcal{L}_{policy})$  and the Maximum Likelihood Estimation (MLE) loss  $(\mathcal{L}_{mle})$  from cross-entropy training (XENT) (weighted by  $\lambda_{mle}$ ). In other words, the addition of XENT is to address the problem of training collapse<sup>1</sup>, commonly encountered when applying RL in language tasks.
$$
\mathcal {L} _ {A C Q - a c t o r} = \mathcal {L} _ {p o l i c y} + \lambda_ {m l e} \mathcal {L} _ {m l e} \tag {8}
$$
The TD learning method, as mentioned previously, is used to train the critic network. It adopts the bootstrapping methodology which performs estimation based on other known estimates. The critic's training objective is to minimise the mean squared difference, called TD error, between all estimated Q values and their corresponding target values in each timestep. Intuitively, the critic is trained to be as good of a  $Q$  function approximator
<sup>1</sup>As pointed out by (Bahdanau et al., 2017), the MLE loss can help prevent early policy determination and vanishing gradient problems.
as possible:
$$
\begin{array}{l} \mathcal {L} _ {T D} = \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {t} \left(Q _ {\phi} \left(\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t}\right) - \right. \tag {9} \\ \hat {Q} _ {\bar {\phi}} (\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t})) ^ {2} \\ \end{array}
$$
Each target  $\hat{Q}_{\bar{\phi}}(\hat{\mathbf{y}}_{1:t - 1,i},\hat{y}_t)$ , expressed below, is defined as the sum of the immediate reward after generating  $\hat{y}_t$  and the expected  $Q$  value of the proceeding timestep, which is computed using another  $Q$  network named the target critic:
$$
\begin{array}{l} \hat {Q} _ {\bar {\phi}} (\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t}) = r _ {t + 1} + \\ \sum_ {w \in \mathcal {W}} \pi_ {\boldsymbol {\theta}} (w | \hat {\mathbf {y}} _ {1: t, i}) Q _ {\bar {\phi}} (\hat {\mathbf {y}} _ {1: t, i}, w) \tag {10} \\ \end{array}
$$
To ensure stability, the weights of the target critic  $(\bar{\phi})$  are updated more slowly than the main critic with the linear interpolation between the current weights of the main and target critics. Also, following (Bahdanau et al., 2017), in addition to the TD error, the critic's loss  $(\mathcal{L}_{ACQ - \text{critic}})$  contains an additional term, weighted by  $\lambda_{var}$ , which aims to minimise the variance in  $Q$  value estimation.
$$
\begin{array}{l} \mathcal {L} _ {A C Q - c r i t i c} = \mathcal {L} _ {T D} + \lambda_ {v a r} \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {w \in \mathcal {W}} (11) \\ \left(Q _ {\phi} (\hat {\mathbf {y}} _ {1: t - 1, i}, w) - \bar {Q} _ {\phi} (\hat {\mathbf {y}} _ {1: t - 1, i})\right) ^ {2} \\ \bar {Q} _ {\phi} \left(\hat {\mathbf {y}} _ {1: t - 1, i}\right) = \frac {1}{| \mathcal {W} |} \sum_ {w ^ {\prime} \in \mathcal {W}} Q _ {\phi} \left(\hat {\mathbf {y}} _ {1: t - 1, i}, w ^ {\prime}\right) (12) \\ \end{array}
$$
Figure 1: High-level structure of the Asynchronous Actor-Critic with Q-Learning Model (Async-ACQ) and the Asynchronous Advantage Actor-Critic (A3C) Model: multiple parallel agents and critics are trained independently. Their weights are used to update the weights of the global agent one by one.
Synchronous and Asynchronous ACQ Both the asynchronous and synchronous versions of the
ACQ model can be easily constructed by deploying N actors with respective N critics on multiple threads. Each of the parallel actors is trained using different segments of the training dataset. For the synchronous variant, the weights of each of the actors are averaged to update the weights of the global agent. For the asynchronous variant, the global agent is updated by the local weights of each agent one by one. That is, during training, each thread-specific agent generates output sequences by sampling from its policy. Then, it performs loss computation and gradient accumulation until it reaches the pre-defined number of timesteps, in which it transfers the accumulated gradients to the global model. The global model subsequently performs an update on its parameters. As the last step of the asynchronous update, the parameters of the thread-specific agent invoking the update are syncsed with the parameters of the global model.
3.3.2 Advantage Actor-Critic (A2C) Model
Another variant of the AC model is the Advantage Actor-Critic (A2C) architecture (Konda and Tsitsiklis, 2001). In this model, the critic is defined as a function approximator, parameterised by  $\psi$ , of the true  $V$  function. Compared to the first variant, the A2C model applies a different training objective to ACQ (Equation (6)).
Given the state space in language tasks is massive, calculating the expectation term would be computationally expensive or even impossible. Therefore, the advantage function can be approximated by sampling once.
$$
A _ {\pi} \left(s _ {t}, a _ {t}\right) \approx r _ {t + 1} + \gamma V _ {\pi} \left(s _ {t + 1}\right) - V _ {\pi} \left(s _ {t}\right) \tag {13}
$$
The actor network in the A2C model is trained in a similar fashion to that of ACQ. Here, the critic estimates the state values (i.e.,  $V_{\psi}(\hat{\mathbf{y}}_{1:t-1,i})$  and  $V_{\psi}(\hat{\mathbf{y}}_{1:t,i})$ ) which are used by the actor to calculate the advantage value. The actor loss function  $(\mathcal{L}_{A2C - actor})$  can be outlined as follows:
$$
\mathcal {L} _ {\text {p o l i c y}} = - \left[ \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {t} \log \right. \tag {14}
$$
$$
\pi_ {\boldsymbol {\theta}} (\hat {y} _ {t} | \hat {\mathbf {y}} _ {1: t - 1, i}) A _ {\boldsymbol {\psi}} (\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t}) ]
$$
$$
A _ {\psi} \left(\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t}\right) = r _ {t + 1} + \gamma V _ {\psi} \left(\hat {\mathbf {y}} _ {1: t, i}\right) - \tag {15}
$$
$$
V _ {\psi} (\hat {y} _ {1: t - 1, i})
$$
$$
\mathcal {L} _ {A 2 C - a c t o r} = \mathcal {L} _ {p o l i c y} + \lambda_ {m l e} \mathcal {L} _ {m l e} \tag {16}
$$
Furthermore, compared to the first variant, the A2C critic is trained to minimise the TD error between its estimation and ground-truth data. The ground-truth data is essentially the true discounted reward-to-go  $(v_{t})$ .
$$
\mathcal {L} _ {A 2 C - c r i t i c} = \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {t} \left(V _ {\psi} \left(\hat {\mathbf {y}} _ {1: t - 1, i}\right) - v _ {t}\right) ^ {2} \tag {17}
$$
Synchronous and Asynchronous A2C The model setups for the synchronous and asynchronous A2C variants are analogous to the ACQ variants, in which N pairs of actors and critics are deployed.
In the asynchronous A2C architecture (A3C) (Mnih et al., 2016), the model also employs n-step TD Learning which uses the true returns from multiple steps in the advantage function to reduce the model bias (Sutton and Barto, 2014). The term  $n$  defines the number of steps to use the real rewards before bootstrapping (using the critic). The standard TD Learning would just use the immediate reward (1-step TD).
$$
A _ {\psi^ {\prime}} \left(\hat {\mathbf {y}} _ {1: t - 1, i}, \hat {y} _ {t}\right) = \sum_ {\tau = 0} ^ {n - 1} \gamma^ {\tau} r _ {t + \tau} + \tag {18}
$$
$$
\gamma^ {n} V _ {\boldsymbol {\psi} ^ {\prime}} (\hat {\mathbf {y}} _ {1: t + n - 1, i}) - V _ {\boldsymbol {\psi} ^ {\prime}} (\hat {\mathbf {y}} _ {1: t - 1, i})
$$
where  $\theta^{\prime}$  and  $\psi^\prime$  represent the thread-specific parameters of each actor and critic, respectively.
4 Semantic-level Unsupervised Reward Function (SURF)
Our semantic-level unsupervised reward, SURF, is based on two scores: Sentence Fluency and Sentence-level Semantic Similarity (SLSS) (Song et al., 2021). Each score assesses translation quality of generated sequences from different aspects and is computed using a pre-trained model. To prevent reward sparsity, the reward function introduces a score normalisation mechanism which normalises scores of a generated sequence (from all timesteps) with respect to the score of its full target sequence. This subsequently yields an unsupervised reward function that is maximally dense. The Sentence Fluency score  $(F(\hat{\mathbf{y}}_{1:t}))$  is defined as the average log-likelihood of the generated sequence tokens  $(\hat{\mathbf{y}}_{1:t})$  as defined by a pre-trained large LM.
The SLSS score measures the overall semantic similarity between the entire generated sequence
and its source sequence calculated as the cosine similarity between the two sentence cross-lingual embeddings.
Score Normalisation From the RL perspective, the MT task does not define the environment component that the agent operates in. That is, unlike the classical RL setting where the environment is relatively fixed, the 'environment' in the MT task is mostly dependent on the source sequence, in terms of its sophistication, structure, length, etc. As a result, a valid and good translation of a source sequence would receive a relatively high score but is not directly comparable to other sequences due to the difference in source sentence complexity. Therefore, it is important to ensure that the reward function is uniform and generalised across all source sentences.
In order to do this for each source sequence, the reward function uses the corresponding target sequence as a 'soft' upper bound for what a machine-generated sequence could achieve. That is, for each of the two score metrics outlined above, the scores from all timesteps received by a generated sequence is normalised to the range 0 to 1 with respect to the score of the full target sequence. To demonstrate the normalisation method, let us consider the formulation below which uses the Sentence Fluency score metric as an example. Given a pair of source (X) and target (Y) sequences and a candidate sequence  $(\hat{\mathbf{y}}_{1:t})$ , the fluency scores from all timesteps of  $\hat{\mathbf{y}}_{1:t}$  would be  $\{F(\hat{\mathbf{y}}_1), F(\hat{\mathbf{y}}_{1:2}), \dots, F(\hat{\mathbf{y}}_{1:t-1}), F(\hat{\mathbf{y}}_{1:t})\}$  while the fluency score for the entire reference target sequence (Y) would be  $F(\mathbf{Y})$ . Using the fluency scores of the candidate and that of the reference, the normalised candidate scores can be calculated as follows:
$$
F _ {n o r m} \left(\hat {\mathbf {y}} _ {1: i}\right) = \left\{ \begin{array}{c c c} \frac {F \left(\hat {\mathbf {y}} _ {1 : i}\right) - m i n}{m a x - m i n} & \text {i f} & m a x \neq m i n \\ 0. 5 & \text {i f} & m a x = m i n, \end{array} \right. \tag {19}
$$
$$
\max  = \max  \left(\left\{F \left(\hat {\mathbf {y}} _ {1}\right), \dots , F \left(\hat {\mathbf {y}} _ {1: t}\right), F (\mathbf {Y}) \right\}\right) \tag {20}
$$
$$
\min  = \min  (\left\{F \left(\hat {\mathbf {y}} _ {1}\right), \dots , F \left(\hat {\mathbf {y}} _ {1: t}\right), F (\mathbf {Y}) \right\}) \tag {21}
$$
Considering the example formulation above, one can observe that the normalisation with respect to the reference score is considered as a 'soft' upper bound as it allows for candidate scores to be higher than the reference score (i.e., allowing the possi
bility that candidate sequences can be better than their references).
The Pay-off Function and Reward Shaping Given a generated sequence  $\hat{\mathbf{y}}_{1:t}$  at timestep  $t$ , its quality can be formulated as the Pay-off Function  $(PO(\hat{\mathbf{y}}_{1:t},\mathbf{X}))$ . The Pay-off Function, as expressed below, is based on the Sentence Fluency and SLSS scores described above.
$$
P O \left(\hat {\mathbf {y}} _ {1: t}, \mathbf {X}\right) = w _ {F} \times e ^ {F \left(\hat {\mathbf {y}} _ {1: t}\right)} + w _ {S} \times e ^ {S L S S \left(\hat {\mathbf {y}} _ {1: t}, \mathbf {X}\right)} \tag {22}
$$