SlowGuess/ACL_2020_OCR huggingface dataset:
Ruiyi Zhang<sup>1</sup>, Changyou Chen<sup>2</sup>, Zhe Gan<sup>3</sup>, Wenlin Wang<sup>4</sup>
Dinghan Shen<sup>3</sup>, Guoyin Wang<sup>1</sup>, Zheng Wen<sup>5</sup>, Lawrence Carin<sup>1</sup>
$^{1}$  Duke University  $^{2}$  University at Buffalo  $^{3}$  Microsoft Dynamics 365 AI
4 Citadel LLC 5 DeepMind
ryzhang@cs.duke.edu
Auto-regressive text generation models usually focus on local fluency, and may cause inconsistent semantic meaning in long text generation. Further, automatically generating words with similar semantics is challenging, and hand-crafted linguistic rules are difficult to apply. We consider a text planning scheme and present a model-based imitation-learning approach to alleviate the aforementioned issues. Specifically, we propose a novel guider network to focus on the generative process over a longer horizon, which can assist next-word prediction and provide intermediate rewards for generator optimization. Extensive experiments demonstrate that the proposed method leads to improved performance.
Text generation is an important area of investigation within machine learning. Recent work has shown excellent performance on a number of tasks, by combining reinforcement learning (RL) and generative models. Example applications include image captioning (Ren et al., 2017; Rennie et al., 2016), text summarization (Li et al., 2018b; Paulus et al., 2017; Rush et al., 2015), and adversarial text generation (Guo et al., 2017; Lin et al., 2017; Yu et al., 2017; Zhang et al., 2017; Zhu et al., 2018). The sequence-to-sequence framework (Seq2Seq) (Sutskever et al., 2014) is a popular technique for text generation. However, models from such a setup are typically trained to predict the next token given previous ground-truth tokens as input, causing what is termed exposure bias (Ranzato et al., 2016). By contrast, sequence-level training with RL provides an effective means of solving this challenge, by treating text generation as a sequential decision-making problem. By directly optimizing an evaluation score (cumulative rewards) (Ranzato et al., 2016), state-of-the-art results have been ob
tained in many text-generation tasks (Paulus et al., 2017; Rennie et al., 2016). However, one problem in such a framework is that rewards in RL training are particularly sparse, since a scalar reward is typically only available after an entire sequence has been generated. Furthermore, the recurrent models focus more on local fluency, and may cause inconsistent semantic meanings for long text generation.
For RL-based text generation, most existing works rely on a model-free framework, which has been criticized for its high variance and poor sample efficiency (Sutton and Barto, 1998). On the other hand, while model-based RL methods do not suffer from these issues, they are usually difficult to train in complex environments. Further, a learned policy is usually restricted by the capacity of an environment model. Recent developments on model-based RL (Gu et al., 2016; Kurutach et al., 2018; Nagabandi et al., 2017) combine the advantages of these two approaches, and have achieved improved performance by learning a model-free policy, assisted by an environment model. In addition, model-based RL has been employed recently to solve problems with extremely sparse rewards, with curiosity-driven methods (Pathak et al., 2017).
In this paper, we propose a model-based imitation-learning method to overcome the aforementioned issues in text-generation tasks. Our main idea is to employ an explicit guider network to model the generation environment in the feature space of sentence tokens, used to emit intermediate rewards by matching the predicted features from the guider network and features from generated sentences. The guider network is trained to encode global structural information of training sentences, and thus is useful to guide next-token prediction in the generative process. Within the proposed framework, to assist the guider network, we also develop a new type of self-attention mechanism to provide high-level planning-ahead information
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2516-2531 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
and maintain consistent semantic meaning. Our experimental results demonstrate the effectiveness of proposed methods.
Text Generation Model Text generation models learn to generate a sentence  $Y = (y_{1},\ldots ,y_{T})$  of length  $T$  possibly conditioned on some context  $X$  .Here each  $y_{t}$  is a token from vocabulary  $\mathcal{A}$  Starting from the initial state  $s_0$  , a recurrent neural network (RNN) produces a sequence of states  $(s_{1},\dots,s_{T})$  given an input sentence-feature representation  $(e(y_{1}),\ldots ,e(y_{T}))$  , where  $e(\cdot)$  denotes a word embedding function mapping a token to its  $d$  dimensional feature representation. The states are recursively updated with a function known as the cell:  $s_t = h(s_{t - 1},e(y_t))$  . One typically assigns the following probability to an observation  $y$  at location  $t$  ..  $p(y|Y_{< t}) = [\mathrm{softmax}(g(s_t))]_y$  . Together  $(g,h)$  specifies a probabilistic model  $\pi$  , i.e.,
$$
\log \pi (Y) = \sum_ {t} \log p \left(y _ {t} \mid Y _ {<   t}\right). \tag {1}
$$
To train the model  $\pi$ , one typically uses maximum likelihood estimation (MLE), via minimizing the cross-entropy loss, i.e.,  $J_{\mathrm{MLE}}(\pi) = -\mathbb{E}[\log \pi(Y)]$ . In order to generate sentence  $Y^{s}$  from a (trained) model, one iteratively applies the following operations:
$$
y _ {t + 1} ^ {s} \sim \operatorname {M u l t i} (1, \operatorname {s o f t m a x} (g (\boldsymbol {s} _ {t}))), \tag {2}
$$
$$
\boldsymbol {s} _ {t} = h \left(\boldsymbol {s} _ {t - 1}, e \left(y _ {t} ^ {s}\right)\right), \tag {3}
$$
where  $\operatorname{Multi}(1, \cdot)$  denotes one draw from a multinomial distribution.
Model-Based Imitation Learning Text generation can be considered as an RL problem with a large number of discrete actions, deterministic transitions, and deterministic terminal rewards. It can be formulated as a Markov decision process (MDP)  $\mathcal{M} = \langle S, A, P, r, \gamma \rangle$ , where  $S$  is the state space,  $A$  is the action space,  $P$  is the deterministic environment dynamics,  $r(s, y)$  is a reward function, and  $\gamma \in (0, 1)$  is the discrete-time discount factor. The policy  $\pi_{\phi}$ , parameterized by  $\phi$ , maps each state  $s \in S$  to a probability distribution over  $A$ . The objective is to maximize the expected reward:
$$
J (\pi) = \sum_ {t = 1} ^ {\infty} \mathbb {E} _ {P, \pi} \left[ \gamma^ {t - 1} \cdot r \left(\boldsymbol {s} _ {t}, y _ {t}\right) \right]. \tag {4}
$$
In model-based imitation learning (Baram et al., 2017; Cheng et al., 2019), a model is built to make
predictions for future state  $s_{t + \triangle t}$  conditioned on the current state<sup>1</sup>, which can be used for action selection, e.g., next-token generation. This model is typically a discrete-time system, taking the current state-action pair  $(s_t, y_t)$  as input, and outputting an estimate of the future state  $s_{t + \triangle t}$  at time  $t + \triangle t$ . At each step  $t$ ,  $y_t$  is chosen based on the model, and the model will re-plan with the updated information from the dynamics. This control scheme is different from a standard model-based method, and is referred to as model-predictive control (MPC) (Nagabandi et al., 2017). Note that in our setting, the state in RL typically corresponds to the current generated sentences  $Y_{1,\dots,t}$  instead of the RNN state of generator (decoder).
The model is illustrated in Figure 1, with an autoencoder (AE) structure for sentence feature extraction and generation. The encoder is shared for sentences from both training data and generated data, as explained in detail below. Overall, text generation can be formulated as an imitation-learning problem. At each timestep  $t$ , the agent, also called a generator (which corresponds to the LSTM decoder), takes the current LSTM state as input, denoted as  $s_t$ . The policy  $\pi_{\phi}(\cdot | s_t)$  parameterized by  $\phi$  is a conditional generator, to generate the next token (action) given  $s_t$ , the observation representing the current generated sentence. The objective of text generation is to maximize the total reward as in (4). We detail the components for our proposed model in the following subsections.
3.1 The Guider Network
The guider network, implemented as an RNN with LSTM units, is adopted to model environment dynamics to assist text generation. The idea is to train a guider network such that its predicted sentence features at each time step are used to assist next-word generation and construct intermediate rewards, which in turn are used to optimize the sentence generator. Denote the guider network as  $G^{\psi}(s_{t-1}^{G}, f_{t})$ , with parameters  $\psi$  and input arguments  $(s_{t-1}^{G}, f_{t})$  at time  $t$ , to explicitly write out the dependency on the guider network latent state  $s_{t-1}^{G}$  from the previous time step. Here  $f_{t}$  is the input to the LSTM guider, which represents the feature of the current generated sentence extracted
$^1 \triangle t > 1$ ; the model predicts future states based on the collected trajectories.
2517
Figure 1: Model overview of text generation with a guider network. Solid lines mean gradients are backpropagated in training; dash lines mean gradients are not backpropagated. CNN is the feature extractor, and MLP outputs the parameters of the Gaussian density which is compatible with the initial state of the LSTM Guider and Decoder.
by an encoder network. Specifically, let the current generated sentence be  $Y_{1\dots t}$  (encouraged to be the same as parts of a training sentence in training), with  $f_{t}$  calculated as:  $f_{t} = \operatorname{Enc}(Y_{1\dots t})$ . The initial state of the guider network is the encoded feature of a true input sentence by the same convolutional neural network (CNN), i.e.,  $s_0^G = \operatorname{Enc}(X)$ , where  $\operatorname{Enc}(\cdot)$  denotes the encoder transformation, implemented with a CNN (Zhang et al., 2017). Importantly, the input to the guider network, at each time point, is defined by features from the entire sentence generated to that point. This provides an important "guide" to the LSTM decoder, accounting for the global properties of the generated text.
Text Generation with Planning We first explain how one uses the guider network to guide next-word generation for the generator (the LSTM decoder in Figure 1). Our framework is inspired by the MPC method (Nagabandi et al., 2017), and can be regarded as a type of plan-ahead attention mechanism. Given the feature  $\pmb{f}_t$  at time  $t$  from the current input sentence, the guider network produces a prediction  $G^{\psi}(s_{t-1}^{G}, \pmb{f}_t)$  as a future feature representation, by feeding  $\pmb{f}_t$  into the LSTM guider. Since the training of the guider network is based on real data (detailed in the next paragraph), the predicted feature contains global-structure information of the training sentences. To utilize such information to predict the next word, we combine the predicted feature with the output of the decoder by constructing an attention-like mechanism. Specifically, we first apply a linear transformation  $\varphi$  on the predicted feature  $G^{\psi}(s_{t-1}^{G}, \pmb{f}_t)$ , forming a weight vector  $\pmb{w}_t \triangleq \varphi(G^{\psi}(s_{t-1}^{G}, \pmb{f}_t))$ . The weight  $\pmb{w}_t$  is applied to the output  $O_t$  of the LSTM decoder by an element-wise multiplication operation. The result is then fed into a softmax layer to generate the next token  $y_t$ . Formally, the generative process
is written as:
$$
\boldsymbol {O} _ {t} = g \left(\boldsymbol {s} _ {t - 1}\right), \quad \boldsymbol {w} _ {t} = \varphi \left(G ^ {\psi} \left(\boldsymbol {s} _ {t - 1} ^ {G}, \boldsymbol {f} _ {t}\right)\right), \tag {5}
$$
$$
y _ {t} \sim \operatorname {M u l t i} (1, \operatorname {s o f t m a x} \left(\boldsymbol {O} _ {t} \cdot \boldsymbol {w} _ {t}\right)), \tag {6}
$$
$$
\boldsymbol {s} _ {t} ^ {G} = h ^ {G} \left(\boldsymbol {s} _ {t - 1} ^ {G}, \boldsymbol {f} _ {t}\right), \quad \boldsymbol {s} _ {t} = h \left(\boldsymbol {s} _ {t - 1}, e \left(y _ {t}\right)\right). \tag {7}
$$
Guider Network Training Given a sentence of feature representations  $(\pmb{f}_1, \pmb{f}_2, \dots, \pmb{f}_T)$  for a training sentence, we seek to update the guider network such that it is able to predict  $\pmb{f}_{t + c}$  given  $\pmb{f}_t$  where  $c > 0$  is the number of steps that are looked ahead. We implement this by forcing the predicted feature,  $G^{\psi}(s_t^G, f_t)$ , to match both the sentence feature  $\pmb{f}_{t + c}$  (first term in (8)) and the corresponding feature-changing direction (second term in (8)). This is formalized by maximizing an objective function of the following form at time  $t$ :
$$
\begin{array}{l} J _ {G} ^ {\psi} = \mathcal {D} _ {\cos} \left(\boldsymbol {f} _ {t + c}, G ^ {\psi} \left(\boldsymbol {s} _ {t - 1} ^ {G}, \boldsymbol {f} _ {t}\right)\right) \tag {8} \\ + \mathcal {D} _ {\mathrm {c o s}} \left(\pmb {f} _ {t + c} - \pmb {f} _ {t}, G ^ {\psi} (\pmb {s} _ {t - 1} ^ {G}, \pmb {f} _ {t}) - \pmb {f} _ {t}\right), \\ \end{array}
$$
where  $\mathcal{D}_{\mathrm{cos}}(\cdot ,\cdot)$  denotes the cosine similarity. By maximizing (8), an ideal guider network should be able to predict the true next words conditioned on the current word in a sentence. As a result, the prediction is used to construct an intermediate reward, used to update the generator (the LSTM decoder), as described further below.
3.2 Feature-Matching Rewards and Generator Optimization
As in many RL-based text-generation methods, such as SeqGAN (Yu et al., 2017) and LeakGAN (Guo et al., 2017), the generator is updated based on policy-gradient methods. As a result, collecting rewards in the generation process is critical.
2We found that the cosine similarity worked better than the  $l_{2}$  -norm.
Though SeqGAN (Yu et al., 2017) has proposed to use rollout to get rewards for each generated word, the variance of the rewards is typically too high to be useful practically. In addition, the computational cost may be too high for practical use. We below describe how to use the proposed guider network to define intermediate rewards, leading to a definition of feature-matching reward.
Feature-Matching Rewards We first define an intermediate reward to generate a particular word. The idea is to match the ground-truth features from the CNN encoder in Figure 1 with those generated from the guider network. Equation (8) indicates that the further the generated feature is from the true feature, the smaller the reward should be. To this end, for each time  $t$ , we define the intermediate reward for generating the current word as:
$$
\begin{array}{l} r _ {t} ^ {g} = \frac {1}{2 c} \sum_ {i = 1} ^ {c} (\mathcal {D} _ {\cos} (\boldsymbol {f} _ {t}, \hat {\boldsymbol {f}} _ {t}) + \\ \mathcal {D} _ {c o s} (\pmb {f} _ {t} - \pmb {f} _ {t - i}, \hat {\pmb {f}} _ {t} - \pmb {f} _ {t - i}), \\ \end{array}
$$
where  $\hat{\pmb{f}}_t = G^{\psi}(s_{t - c - 1}^G,\pmb{f}_{t - c})$  is the predicted feature. Intuitively,  $\pmb {f}_t - \pmb {f}_{t - i}$  measures the difference between the generated sentences in feature space; the reward is high if it matches the predicted feature transition  $\hat{\pmb{f}}_t - \pmb {f}_{t - i}$  from the guider network. At the last step of text generation, i.e.,  $t = T$ , the corresponding reward measures the quality of the whole generated sentence, thus it is called a final reward. The final reward is defined differently from the intermediate reward, discussed below for both the unconditional- and conditional-generation cases.
Note that a token generated at time  $t$  will influence not only the rewards received at that time but also the rewards at subsequent time steps. Thus we propose to define the cumulative reward,  $\sum_{i = t}^{T}\gamma^{i}r_{i}^{g}$  with  $\gamma$  a discount factor, as a feature-matching reward. Intuitively, this encourages the generator to focus on achieving higher long-term rewards. Finally, in order to apply policy gradient to update the generator, we combine the feature-matching reward with the problem-specific final reward, to form a  $Q$ -value reward specified below.
Similar to SeqGAN, the final reward is defined as the output of a discriminator, evaluating the quality of the whole generated sentence, i.e., the smaller the output, the less likely the generation is a true sentence. As a result, we combine the adversarial reward  $r^f \in [0,1]$  by the discriminator (Yu et al.,
Algorithm 1 Model-based Imitation Learning for Text Generation
Require: generator policy  $\pi^{\phi}$ ; guider network  $G^{\psi}$ ; a sequence dataset  $\{X_{1\dots T}\}$  by some expert policy.
1: Initialize  $G^{\psi}$ ,  $D^{\theta}$  with random weights.
2: while Imitation Learning phase do
3: Update generator  $\pi^{\phi}$ , guider  $G^{\psi}$  with MLE loss.
4: end while
5: while Reinforcement Learning phase do
6: Generate a sequence  $Y_{1\dots T}\sim \pi^{\phi}$
7: Compute  $Q_{t}$ , and update  $\pi^{\phi}$ .
8: end while
2017) with the guider-matching rewards, to define a  $Q$ -value reward as  $Q_{t} = (\sum_{i = t}^{T}\gamma^{i}r_{i}^{g})\times r^{f}$ .
Generator Optimization The generator is initialized by pre-training on sentences with an autoencoder structure, based on MLE training. After that, the final  $Q$ -value reward  $Q_{t}$  is used as a reward for each time  $t$ , with standard policy gradient optimization methods to update the generator. Specifically, the policy gradient is
$$
\nabla_ {\phi} J = \mathbb {E} _ {(\boldsymbol {s} _ {t - 1}, y _ {t}) \sim \rho_ {\pi}} \left[ Q _ {t} \nabla_ {\phi} \log p (y _ {t} | \boldsymbol {s} _ {t - 1}; \phi , \varphi) \right],
$$
$$
\nabla_ {\varphi} J = \mathbb {E} _ {\left(\boldsymbol {s} _ {t - 1}, y _ {t}\right) \sim \rho_ {\pi}} \left[ Q _ {t} \nabla_ {\varphi} \log p \left(y _ {t} \mid \boldsymbol {s} _ {t - 1}; \phi , \varphi\right) \right],
$$
where  $p(y_{t}|s_{t - 1};\phi ,\varphi)$  is the probability of generating  $y_{t}$  given  $s_{t - 1}$  in the generator. Algorithm 1 describes the proposed model-based imitation learning framework for text generation.
Model-based or Model-free Text generation seeks to generate the next word (action) given the current (sub-)sentence (state). The generator is considered as an agent that learns a policy to predict the next word given its current state. In previous work (Ranzato et al., 2016), a metric reward is given and the generator is trained to only maximize the metric reward by trial, thus this is model-free learning. In the proposed method, the guider network models the environment dynamics, and is trained by minimizing the cosine similarity between the prediction and the ground truth on real text. For generator training, it maximizes the reward which is determined by the metric and guider network, and thus is model-free learning with model-based boosting (Gu et al., 2016). The model predictive control scheme is included in our method, where the guider network is used to help next-word selection at each time-step.
Figure 2: Guided style transfer: the Guider network controls the sentiment in the higher level, and the Generator focuses on preserving content in the lower level.
4 Extension to Non-parallel Text Style Transfer
As illustrated in Figure 2, our framework naturally provides a way for style transfer, where the guider network plays the role of style selection, and the generator only focuses on maintaining content without considering the styles. To make the guider network focus on the guidance of styles, we assign the label  $l$  as the initial state  $s_0^G$  of the guider network. Specifically, at each step  $t$ , we feed the current sentence representation  $f_t$  and label  $l$  into the guider network:
$$
\boldsymbol {O} _ {t} = g \left(\boldsymbol {s} _ {t - 1}\right), \quad \boldsymbol {w} _ {t} = \varphi \left(G ^ {\psi} \left(\boldsymbol {s} _ {t - 1} ^ {G}, [ \boldsymbol {f} _ {t}, l ]\right)\right), \tag {9}
$$
$$
y _ {t} \sim \operatorname {M u l t i} (1, \operatorname {s o f t m a x} \left(\boldsymbol {O} _ {t} \cdot \boldsymbol {w} _ {t}\right)). \tag {10}
$$
For the generator, we put an adversarial regularizer on the encoded latent  $s_0(X)$  and penalize it if it contains the sentiment information, by maximizing the entropy, i.e.,  $\max \sum_{l} p(l | s_0(X)) \log p(l | s_0(X))$ , where  $p$  is a pre-trained classifier. Intuitively, the generator gives candidate words represented by  $O_t$ , while the guider makes a choice implicitly by  $w_t$  based on the sentiment information. The sentiment information is contained in  $w_t$ , while the content of the original sentence is represented by  $O_t$ . To achieve style-transfer, one feeds the original sentence  $X$  with the target style label  $l$  to get the transferred sentence  $Y$  with style  $l$ . Following previous work (Hu et al., 2017; Yang et al., 2018; Cheng et al., 2020), we adopt a classifier as the discriminator and the soft-argmax approach (Kusner and Miguel, 2016) for the update of generator instead of policy gradient (Sutton and Barto, 1998).
We first review related works that combine RL and GAN for text generation. As one of the most rep
representative models in this direction, SeqGAN (Yu et al., 2017) adopts Monte-Carlo search to calculate rewards. However, such a method introduces high variance in policy optimization. There are a number of works proposed subsequently to improve the reward-generation process. For example, RankGAN (Lin et al., 2017) proposes to replace the reward from the GAN discriminator with a ranking-based reward, MaliGAN (Che et al., 2017) modifies the GAN objective and proposes techniques to reduce gradient variance, MaskGAN (Fedus et al., 2018) uses a filling technique to define a  $Q$ -value reward for sentence completion, RelGAN (Nie et al., 2019) uses a relational memory based generator for the long-distance dependency modeling, FMGAN (Chen et al., 2018) uses a feature mover distance to match features of real and generated sentences inspired by optimal transport (Chen et al., 2019; Zhang et al., 2018), and LeakGAN (Guo et al., 2017) tries to address the sparse-reward issue for long-text generation with hierarchical RL by utilizing the leaked information from a GAN discriminator. One problem of LeakGAN is that it tends to overfit the training data, yielding generated sentences that are often not diverse. By contrast, by relying on a model-based imitation learning approach, our method learns global-structure information, which generates more-diverse sentences, and can be extended to conditional text generation. Zhang et al. (2020) designed a differentiable nested Wasserstein distance for semantic matching, which can be applied for further improvement.
RL techniques can also be used in other ways for text generation (Bachman and Precup, 2015). For example, Ranzato et al. (2016) trained a Seq2Seq model by directly optimizing the BLEU/ROUGE scores with the REINFORCE algorithm. To reduce variance of the vanilla REINFORCE, Bahdanau et al. (2017) adopted the actor-critic framework for sequence prediction. Furthermore, Rennie et al. (2016) trained a baseline algorithm with a greedy decoding scheme for the REINFORCE method. Note that all these methods can only obtain reward after a whole sentence is generated. Planning techniques in RL have also been explored to improve text generation (Gulcehre et al., 2017; Serdyuk et al., 2018). Zhang et al. (2020) introduced the self-imitation scheme to exploit historical high-quality sentences for enhanced exploration. Compared to these related works, the proposed guider network can provide a planning mechanism and intermediate rewards.