SlowGuess/NAACL_2022_OCR huggingface dataset:
Antoine Chaffin $^{1,2}$ , Vincent Claveau $^{1}$ , Ewa Kijak $^{1}$
$^{1}$ CNRS, IRISA, Univ. Rennes 1, Campus de Beaulieu, 35000 Rennes, France
$^{2}$ IMATAG, 13 Rue Dupont-des-Loges, 35000 Rennes, France
{antoine.chaffin, vincent.claveau, ewa.kijak}@irisa.fr
Large language models (LM) based on Transformers allow to generate plausible long texts. In this paper, we explore how this generation can be further controlled at decoding time to satisfy certain constraints (e.g. being nontoxic, conveying certain emotions, using a specific writing style, etc.) without fine-tuning the LM. Precisely, we formalize constrained generation as a tree exploration process guided by a discriminator that indicates how well the associated sequence respects the constraint. This approach, in addition to being easier and cheaper to train than fine-tuning the LM, allows to apply the constraint more finely and dynamically. We propose several original methods to search this generation tree, notably the Monte Carlo Tree Search (MCTS) which provides theoretical guarantees on the search efficiency, but also simpler methods based on reranking a pool of diverse sequences using the discriminator scores. These methods are evaluated, with automatic and human-based metrics, on two types of constraints and languages: review polarity and emotion control in French and English. We show that discriminator-guided MCTS decoding achieves state-of-the-art results without having to tune the language model, in both tasks and languages. We also demonstrate that other proposed decoding methods based on re-ranking can be really effective when diversity among the generated propositions is encouraged.
Generative language models exist for a long time, but with advent of the transformer architecture (Vaswani et al., 2017) and increasing computing capabilities, they are now able to generate well written and long texts. In particular, large models, such as the well known GPT-2 (Radford et al., 2019) and GPT-3 (Brown et al., 2020), have been used successfully for various applications: assisting writers, summarizing, augmentating data for
subsequent NLP tasks, generating fake news (Kumar et al., 2020; Papanikolaou and Pierleoni, 2020; Zellers et al., 2019). Yet, beside the prompt used to initiate the generation process, there are few options to have control on the generation process. Being able to add some constraints on the generated texts is useful for various situations. For example, it allows to create texts that follow a certain writing style, convey a certain emotion or polarity or to ensure that a generated summary contains correct information. More critically, it can be used to prevent the inherent toxicity of language models trained on the internet, or to not reproduce gender or race stereotypes. So far, most methods necessitate to fine-tune the LM, so that it specifically learns to model this constraint, i.e. the constraint is -hopefully- incorporated in the LM. This finetuning approach has several drawbacks. It implies to train multiple specific LMs (one per constraint), which is costly, when even possible given the size of current state-of-the-art LM, and results in several models.
In this paper, we propose new approaches to add such additional constraints on the texts but at decoding time. We exploit a discriminator that is trained to determine if a text follows a given constraint or not; its output provides information to guide the generation toward texts that satisfy this expected constraint. In order to make the most of the discriminator information, we propose an original method based on the Monte Carlo Tree Search (MCTS) algorithm (Coulom, 2006), namely Plug and Play Language - Monte Carlo Tree Search (PPL-MCTS). We also propose simpler methods based on re-ranking to fulfil this goal. Both approaches do not require fine-tune the LM; adding a new constraint can thus simply be done by providing a discriminator verifying if a text complies with what is expected. More precisely, our main contributions are the following ones:
1. we propose to use MCTS as a decoding strat
Human Language Technologies, pages 2953 - 2967
egy to implement constrained generation and we show, on 3 datasets and 2 languages, that it yields state-of-the-art results while offering more flexibility;
2. we also explore simpler generation methods based on re-ranking and show that this kind of approach, with low computational costs, can also be competitive if the diversity within propositions to re-rank is encouraged;
3. we provide a fully functional code implementing a batched textual MCTS $^1$  working with the popular HuggingFace's Transformers library (Wolf et al., 2020)
The goal of constrained textual generation is to find the sequence of tokens  $x_{1:T}$  which maximises  $p(x_{1:T} \mid c)$ , given a constraint  $c$ . Few methods address the constrained textual generation.
Class-conditional language models. Class-conditional language models (CC-LMs), as the Conditional Transformer Language (CTRL) model (Keskar et al., 2019), train or fine-tune the weights  $\theta$  of a single neural model directly for controllable generation, by appending a control code in the beginning of a training sequence. The control code indicates the constraint to verify and is related to a class containing texts that satisfy the constraint. For the sake of simplicity, we will denote without distinction the class, the constraint verified by its texts and the associated control code by  $c$ . Trained with different control codes, the model learns  $p_{\theta}(x_{1:T} \mid c) = \prod_{t=1}^{T} p_{\theta}(x_t \mid x_{1:t-1}, c)$ . The constraint can then be applied during generation by appending the corresponding control code to the prompt. While this method gives some kind of control over the generation, the control codes need to be defined upfront and the LM still needs to be trained specifically for each set of control codes. This is an important limitation since the current trend in text generation is the use of large pre-trained models which can hardly be fine-tuned (for instance, the last version of GPT, GPT-3, cannot be fine-tuned without access to very large hardware resources).
Discriminator-based methods The general idea of discriminator-guided generation is to combine
a discriminator  $D$  with a generative LM. The discriminator explicitly models the constraint by calculating the probability  $p_{D}(c \mid x_{1:T})$  of the sequence  $x_{1:T}$  to satisfy the constraint  $c$ . This probability is directly related to  $p(x_{1:T} \mid c)$  through Bayes' rule:  $p(x_{1:T} \mid c) \propto p_{D}(c \mid x_{1:T}) p_{\theta}(x_{1:T})$ . Discriminator-based methods alleviate the training cost problem, as discriminators are easier to train than a LM. Moreover, any additional constraint can be defined a posteriori without tuning the LM, only by training another discriminator. The discriminators have been used in different ways to explore the search space. In the work of (Holtzman et al., 2018; Scialom et al., 2020), the space is first searched using beam search to generate a pool of proposals with a high likelihood  $p_{\theta}(x_{1:T})$ , and then the discriminator is used to re-rank them. However, in addition that beam search can miss sequences with high likelihood, it is biased towards the likelihood, while the best sequence might only have an average likelihood, but satisfies the constraint perfectly.
Hence, it might be more suitable to take the discriminator probability into account during decoding rather than after generating a whole sequence. In this case, the discriminator is used at each generation step  $t$  to get the probability  $p_D(c \mid x_{1:t})$  for each token of the vocabulary  $\mathcal{V}$ , and merge it to the likelihood  $p_\theta(x_{1:t})$  to choose which token to emit. In order to reduce the cost of using a discriminator on every possible continuation, GeDi (Krause et al., 2020) proposes to use CC-LMs as generative discriminators. The method relies on the fact that the CC-LM computes  $p_\theta(x_t \mid x_{1:t-1}, c)$  for all tokens of the vocabulary which can be used to get  $p_\theta(c \mid x_{1:t})$  for all tokens using Bayes' equation. This approach is thus at the intersection of tuning the LM and using a discriminator: it tunes a small LM (the CC-LM) to guide a bigger one.
In Plug And Play Language Model (PPLM) (Dathathri et al., 2020), the discriminator is used to shift the hidden states of the pre-trained transformer-based LM towards the desired class at every generation step. PPLM can be used on any LM and with any discriminator. However, PPLM needs to access the LM to modify its hidden states, while our approach only requires the output logits. As some LM can only be used through access to logits (e.g. GPT-3 API), this makes our approach more plug and play than PPLM.
A common drawback of all these approaches is their lack of a long-term vision of the generation.
https://github.com/NohTow/PPL-MCTS
Indeed, the discriminator probabilities become necessarily more meaningful as the sequence grows and might only be trustable to guide the search when the sequence is (nearly) finished. When used in a myopic decoding strategy, classification errors will cause the generation process to deviate further and further. Trying to optimize a score defined in the long horizon by making short term decisions is very similar to common game setups such as chess, where the Monte Carlo Tree Search (MCTS) has proven to be really effective (Silver et al., 2018), which motivated our approach.
3 PPL-MCTS method
The approach that we propose is in line with methods using a discriminator to guide a large LM decoding, without the need to re-train it. Also, it can be applied to any LM with any discriminator, following the plug and play paradigm. Unlike previous approaches, it is able to have a long term vision on what is generated. Being able to make a short-term decision (choice of the next token  $x_{t}$  at time step  $t$ ) that is promising in the long run is based on the exploration of the search space. We propose here to use the Monte Carlo Tree Search (MCTS) for an efficient exploration of this space.
MCTS is very well suited for this problem for three reasons. First, it allows to get a local score (i.e., a score for the next token to emit) using finished sequences. Hence, this score is more meaningful than scores based only on the next step. Second, it allows to explicitly define the compromise between exploitation of promising sequences (with a high likelihood), and exploration of other potentially promising sequences (to not miss better sequences with a lower likelihood). The fact that regret, i.e. the number of simulations done on a suboptimal sequence, has a theoretical upper bound in MCTS (Rosin, 2011) is a nice guarantee that the computation time is not wasted and the search is efficient. Finally, it outputs a solution at each iteration and so can fit our computational budget by allowing to adjust the quality of the solution to calculation spent.
Text generation as tree exploration process. The search space of the text generation corresponds to a tree: its root is the prompt and the child of a node is its father's sequence with one of the  $|\mathcal{V}|$  possible tokens appended. In the case of constrained generation, the goal is thus to find the path, and therefore the sequence  $x$ , with the highest  $p(x \mid c)$
Figure 1: Illustration of the constrained generation process as a tree exploration from the prompt The cat. Classification probabilities are only represented on completed sequences.
possible without exploring the whole tree in width and depth. As mentioned previously, this probability can be computed as the product of the likelihood  $p_{\theta}(x)$  and the probability given by a discriminator  $p_D(c \mid x)$ . An illustration of such a tree can be found in Fig. 1, where the likelihood of  $x$  is forged by multiplying corresponding conditional probabilities along the path, and the classification probability is calculated at the terminal node.
Applying MCTS to text generation. MCTS is a heuristic based iterative algorithm that uses randomness to solve deterministic problems that cannot be solved using traditional approaches, often because the search space is too large to be entirely explored. Each iteration consists in four consecutive steps. In the particular context of applying MCTS to text generation, we made some adaptations:
1. Selection Recursively choose children from the root to a node that has not been expanded yet. To only explore viable sequences, the probability  $p_{\theta}(x_i \mid x_{1:t-1})$  of a given token  $x_i$  given by the LM is used during the selection phase. To this end, the children chosen are those maximizing the Polynomial Upper Confidence Trees (PUCT) (Rosin, 2011) as defined in (Silver et al., 2017):
$$
P U C T (i) = \frac {s _ {i}}{n _ {i}} + c _ {p u c t} p _ {\theta} \left(x _ {i} \mid x _ {1: t - 1}\right) \frac {\sqrt {N _ {i}}}{1 + n _ {i}} \tag {1}
$$
with  $s_i$  the aggregated score of the node  $i$ ,  $n_i$  the number of simulations played after this node,  $N_i$  the number of simulations played
after its parent, and  $c_{puct}$  a constant defining the compromise between exploration and exploitation. In the task of constrained generation, we define the score of a sequence as its probability knowing the class  $p(x \mid c)$ .
2. Expansion If the selected node is not terminal, use the LM to expand it by creating its children.
3. Simulation (roll-out) Sample one of these children according to  $p_{\theta}(x_i \mid x_{1:t-1})$ , and go to a terminal node through a random walk or another pattern.
4. Backpropagation Aggregate the final score obtained at the terminal node  $(p(x \mid c))$  to each parent until root. There are different strategies to aggregate scores, as computing the average between the actual score and the one being backpropagated, or taking the maximum of the two. We take the aggregated score  $s_i$  associated to the node  $i$  as the averaged probability over all simulations played after this node.
When the number of iterations has reached the allocated budget, the building of the tree stops. The token  $x_{i}$  selected for the current decoding step can be selected as the most played node amongst the root's children nodes, or the one with the highest aggregated score. We chose the most played one.
These adaptations of MCTS to constrained generation are summarized in Fig. 2. Note that any language model can be used for defining the probability  $p_{\theta}(x_i \mid x_{1:t-1})$  and any discriminator for scoring sequences, hence the name of our approach: Plug and Play Language - Monte Carlo Tree Search (PPL-MCTS). MCTS has been very recently used for machine translation (Leblond et al., 2021), question generation and summarization (Scialom et al., 2021). The differences with these concurrent studies are discussed in Appendix A.5.
Model improvements. In order to allow a finer control on how the constraint is applied, we introduce a parameter  $\alpha \in [0,1]$  to control the compromise between likelihood and constraint strength, modifying Bayes' equation:  $p(x\mid c)\propto p_D(c\mid x)^{\alpha}p_{\theta}(x)^{1 - \alpha}$ . Note that PUCT (1) already considers the likelihood of the sequence, favoring the selection of nodes with high likelihoods. Hence, even sequences generated with  $\alpha = 1$  are correctly
Figure 2: MCTS application to text generation.
written. Setting  $\alpha < 1$  forces the algorithm to explore solutions even closer to the language model. In our experiments, we set  $\alpha = 1$  to strengthen the constraint.
To avoid expensive roll-outs, one may also assign a value to unfinished sequences at the cost of a less precise evaluation that may be not as meaningful as when doing roll-outs. Indeed, the discriminator can be trained on sequences with variable numbers of tokens, allowing it to be used at each node without the need of simulations. In this setup, the MCTS is used as an efficient compromise between exploration and exploitation, losing part of its long view property but allowing to skew the exploration toward promising solutions.
Finally, during our first experiments, we observed that PPL-MCTS leads to repetitive patterns. This is very similar of what happens with greedy search, where a single sequence with a high likelihood is dominating the search. If such sequences also have a pretty high discriminator scores, they will be repeated often. CTRL (Keskar et al., 2019) offers a simple yet very powerful method to avoid noisy repetitions. It applies a scalar factor  $I(i)$  to the temperature parameter  $\tau$  of a given token  $x_{i}$  that penalizes this token if it is already in the input sequence. The probability of a given token becomes:
$$
p _ {\theta} ^ {\prime} \left(x _ {i} \mid x _ {1: t - 1}\right) = \frac {\exp \left(z _ {i} / (\tau \cdot I (i))\right)}{\sum_ {v} \exp \left(z _ {v} / (\tau \cdot I (v))\right)} \tag {2}
$$
with the repetition penalty  $I(i) > 1$  if  $x_{i}$  is already in the prompt and 1 otherwise, and  $z$  the neural LM predicted logits over the vocabulary  $\mathcal{V}$ . Thus, probabilities of already emitted tokens are penalized, but if the language model gives a really high score to one token (hence, it is very confident that this should be the token to emit), it may still be selected as the output token.
4.1 Performance assessment
The goal of constrained generation is to generate samples that 1) belong to a specific class while 2) keeping the language quality of the original LM, and 3) with enough diversity across samples. We chose three different metrics to evaluate each of these aspects: 1) accuracy, which is verified by an external "oracle" discriminator trained on a dataset disjoint from the one used to guide the generation; 2) perplexity, which is computed using an "oracle" LM, i.e. an unconstrained LM trained on different data than the one used to train the constrained generator; 3) Self-BLEU score (Zhu et al., 2018), which is the BLEU score (Papineni et al., 2002) of a sample using the other samples as references: a high Self-BLEU score means that there is a lot of overlap between generated samples, and thus that the diversity is low. Such automatic metrics have known limitations (Caccia et al., 2020) but results of human evaluation on the CLS dataset, detailed in Section 4.6, confirm that they provide a good overview of the performance.
In practice, the studied dataset (see below) is split into two parts, each part being sub-divided in train/val/test sets. The first part serves to train models used for the generation (LM and discriminator), while the second is used to train oracles which serve to compute the automatic evaluation metrics. The test set of this second part will also be used to forge prompts for the generation. Further details on data splits are given in Appendix A.1. Each metric is evaluated on a pool of 900 generated samples.
Three different datasets are used in the experiments presented hereafter: amazon_polarity (Zhang et al., 2015), CLS (from the FLUE (Le et al., 2020) dataset) and emotion (Saravia et al., 2018). The first two are Amazon reviews which have been labeled as positive or negative, so the intended task is to study the possibility of applying polarity to the generation. As CLS is in French, these two datasets will serve to ensure that the methods have the same behaviour for different languages. Emotion is a collection of tweets classified under eight basic emotions: anger, anticipation, disgust, fear, joy, sadness, surprise and trust. This dataset is supposed to be more challenging since there are more classes and texts are smaller (only composed of one sentence), hence
the model needs to precisely generate the target emotion with few tokens. It is worth noting that the 3 datasets have different sizes: 4,000,000 instances in total for amazon_polarity, 20,000 for emotion and 6,000 for CLS. They are available at https://huggingface.co/datasets/.
We adapted prompts used to start the generation for each datasets depending on the data format. Amazon_polarity comes with a "title" column which corresponds to the title the user gave to the review. This field is directly used as prompt. For the two other datasets, the prompts are the very first tokens of the text field. Because texts from emotion and CLS have different lengths, the size of prompts are also different: it is arbitrarily set to 6 tokens for CLS and 4 for emotion.
4.3 Methods and baselines
Baselines. Beside PPL-MCTS, we propose several baselines and simple techniques. Most studies on re-ranking create proposals using beam search and then re-rank them using the product of likelihood and discriminator probability, limiting the diversity in the proposals pool. We propose re-ranking with different variations, in the way sequences to re-rank are produced, and in the way the final sequence is chosen in an attempt to improve such approaches. Three methods are tested to generate propositions: beam search (Dept., 2018) (with a beam size of 3), nucleus (top-p) sampling (Holtzman et al., 2020) (with  $p = 0.9$ ), as well as beam sampling (as described in (Caccia et al., 2020)). For the final choice, we also propose three different methods: argmax, where the sequence that has the highest  $p(x|c)$  is chosen; first true, where propositions are sorted by descending likelihood and the first sequence that belongs to the correct class according to the guiding discriminator is chosen; and sampling, where the distribution of  $p(x|c)$  for the propositions is normalized and the chosen sequence is sampled following this distribution. Similarly to PPL-MCTS, the likelihood part of  $p(x|c)$  is omitted (i.e.,  $\alpha = 1$ ) since sequences in the pool of propositions already have an high likelihood.
It should be noted that in our setting, a generated sequence corresponds to a document (e.g. a whole review). This choice makes sense for our datasets, but re-ranking at a smaller level (after each sentence, after x tokens...) would also be possible and might produce different results.
Figure 3: Example of two constrained generations using PPL-MCTS, one on the negative class, one on the positive class, using the same prompt (in bold) from amazon_polarity.
Methods from the literature We compare our results with methods from the literature. In particular, we test CC-LMs trained on the target task, similarly as CTRL. We tested this method using greedy search as well as sampling for decoding. We also propose an implementation of CC-LM trained with the classification loss initially proposed for the GeDi method (Krause et al., 2020). These CC-LMs are further used to implement the state-of-the-art GeDi model. In the experiments reported below, we report results for GeDi models trained with and without the classification loss. Finally, we report results of PPLM. For a fair comparison, the same discriminator and LM are used for our PPL-MCTS approach, the re-ranking approaches (baselines), and PPLM.
4.4 Experimental setting
For each method, a number of tokens equal to the average length of sequences of the dataset are generated: 98 tokens for amazon_polarity, 23 for emotion and 137 for CLS. Fixing the number of generated tokens ensures fair comparisons between the tested methods. Indeed, even though perplexity and Self-BLEU metrics are normalized by the length of the sequence, these metrics can tend to penalize a model producing longer sequences: such model has more risk to deviate and repeat itself, which would results in higher values compared to a model producing shorter sequences. An example of generation from amazon_polarity is given in Fig. 3.
To run all of these methods, three different models are needed: one discriminator, a "vanilla" LM used as generator, and the CC-LM used in the CTRL and GeDi approaches. For the discriminator used to guide the generation, we rely on BERT-base-cased (Devlin et al., 2019) for the English datasets and FlauBERT-large-cased (Le et al., 2020) for CLS. As vanilla LM, we use GPT-2 small models, relying on OpenAI's pre-trained model for the English datasets and on belgpt2 for
Figure 4: Accuracy according to the roll-out size; CLS dataset