A Block Metropolis-Hastings Sampler for Controllable Energy-based Text Generation (2023)

Abstract: Recent work has shown that energy-based language modeling is an effective framework for controllable text generation because it enables flexible integration of arbitrary discriminators. However, because energy-based LMs are globally normalized, approximate techniques like Metropolis-Hastings (MH) are required for inference. Past work has largely explored simple proposal distributions that modify a single token at a time, like in Gibbs sampling. In this paper, we develop a novel MH sampler that, in contrast, proposes re-writes of the entire sequence in each step via iterative prompting of a large language model. Our new sampler (a) allows for more efficient and accurate sampling from a target distribution and (b) allows generation length to be determined through the sampling procedure rather than fixed in advance, as past work has required. We perform experiments on two controlled generation tasks, showing both downstream performance gains and more accurate target distribution sampling in comparison with single-token proposal techniques.

## introduction controllable text generation has many important downstream applications, ranging from reducing bias in generated text to increasing factuality @xcite @xcite @xcite . while traditional autoregressive language models (lms) can produce highly fluent text, controlling their output and generating text which satisfies specific desired attributes remains a hard problem for all but the largest industrial lms. one line of past work has made progress on controllable text generation by integrating discriminators-e.g. pretrained text classifiers that directly measure control attributes-into the scoring function for text gener-ation @xcite @xcite . these techniques provide a flexible interface for exerting control: a user can combine discriminators and heuristic scoring functions together with likelihoods from traditional lms to form a product of experts, guiding outputs to satisfy target criteria. while these techniques enable effective control, they present a new challenge for decoding. the scoring functions introduced by discriminators are not autoregressive: they are global potential functions that take the entire utterance as input. this means that the overall model is not autoregressive and exact sampling is intractable. past work has developed various heuristic or approximate decoding strategies @xcite @xcite @xcite @xcite . one of the more principled inference techniques treats the product of experts as an energy-based lm-that is, a globally normalized language model @xcite @xcite -and introduces a metropolis-hastings (mh) sampler for decoding. more specifically, @xcite use bert @xcite to propose a change to a single token of the current sequence at each step of the mh chain (like a traditional gibbs sampler) and the energy lm exerts its influence through mh's accept/reject step, correcting the bias of the proposal distribution. while principled, this approach has serious limitations. first, since only a single token can be changed at each step, inference is extremely slow. second, since the proposal distribution does not alter the length of the current sequence, the length of the desired output must be specified in advance. we conduct experiments on two downstream text style transfer tasks that have been used in past work as benchmarks for controllable generation @xcite . specifically, we study style transfer performance on two challenging datasets: the shakespeare author imitation dataset @xcite and the gyafc formality corpus @xcite . across experiments, we find that our novel sampler is able to make substantially faster progress towards highscoring samples per forward-pass of the target energy lm in comparison with the single-token resampling mh procedure from past work. further, for most downstream tasks, our novel sampler also leads to improvements in the output text in terms of fluency, style transfer accuracy, and semantic similarity to the desired ground truth generations. ## background the m&m approach @xcite defines an mcmc sampling procedure for language models that are globally normalized, which are often called energy-based lms. explicitly, an energy-based sequence model defines a globally normalized probability distribution over the space of possible finite-length sequences x as: θ) , where e(x; θ) corresponds to the scalar energy of a sequence x that is judged by some model parameterized by θ. lower energy corresponds to higher likelihood of x. unlike popular autoregressive techniques, there is no general tractable method of sampling from energy models formulated in this way -even the likelihood function is intractable to compute due to the global normalization constant. however, their high flexibility and compatibility with black-box experts make energy models highly attractive, warranting research into this problem. ## methodology in this section, we will describe and motivate our approach. similar to m&m, we frame controlled generation as a sampling problem where our goal is to get samples from a specific energy-based sequence model. however, m&m has important limitations in the sampling procedure that should be noted: limitations of token-level sampling the m&m masking process destroys important information that is often relevant to the task at hand: for example, if a name is masked out, it is unlikely to be predicted again; this means m&m can largely not restructure sentences and instead prefers minimal edits which achieve the end goal. importantly, editing a single token at a time also significantly slows mixing. for example, if we want to make the sentence "how are you?" to be more shakespearean, the single-token edit "how art you?" is not fluent or grammatical and is likely to be rejected, but is a necessary step to achieve the end goal of "how art thou?"; this important issue is illustrated in figure 1. using a block mh sampler sidesteps this issue by allowing the proposal distribution to select which parts of the sentence to edit and to propose changes to multiple tokens simultaneously. furthermore, m&m uses bert to calculate p(x| x) and p( x|x). importantly, since bert was trained on a dataset of modern english, samples from this distribution will also be. in figure 1 bert is unlikely to propose the token "art" in the first place, this is not a modern english token and bert has no information about the task. prompting an llm with information about the task guides the model towards making more impactful changes. finally, m&m is a fixed-length sampling method: the output is always the same length as the input. the freedom to add or delete tokens is very valuable for many downstream tasks. our sampling procedure, detailed below, targets these weak-points and improves upon past work. @xcite , we devise a metropolis-hastings (mh) chain that iteratively edits text in order to produce samples from the target energy model. we begin with a set seed text and progressively edit this sentence, forming a long markov chain in the process. the acceptance or rejection of these edits is a function of both the expert blackbox models and sample probability as judged by the proposal model. unlike previous work where the proposal function was replacement of a single token, we instead choose to prompt flan-t5-xxl @xcite to edit the sentence; this allows for arbitrary-length generation and makes our approach a block-level mh sampler (similar to blocked gibbs sampling) as multiple variables (tokens) are updated every proposal step. ## intrinsic evaluation of sampler in this section, we aim to conduct an intrinsic evaluation of the proposed sampler, which we refer to as mh-block, separate from the downstream controllable generation tasks we consider in section 5. specifically, we would like to evaluate how well mh-block approximates exact sampling from a complex target distribution relative to the baseline token-level sampling procedure, which we refer to as m&m. to accomplish this, we need to define a target energy model for which exact sampling is actually tractable so that we can draw exact samples and compare. for this purpose, we treat a prompted conditional distribution of llama-7b @xcite as our target 'energy' model by setting e(x) in equation 1 to llama-7b's negative log-likelihood. specifically, we prompt llama-7b to paraphrase a fixed input sentence (randomly sampled from the shakespeare dataset mentioned in section 5, consisting of 13 tokens) and treat the resulting conditional over text sequences as our target. we produce 100 samples using mh-block, 100 samples using m&m, and 1000 exact samples using ancestral sampling and compare the distribution of resulting energy values under the target in figure 2 . for mh-block, we run 100 separate mh chains consisting of exactly 10 proposal steps each, and take the final step's sequence as the output sample. for m&m, the setup is the same, except that we run 130 proposal steps per chain to account for m&m's limitation to a single token change per step. this means that while mh-block only requires 10 forward passes of llama-7b per sample, m&m requires 130. in figure 2 , we see that the distribution of samples from mh-block has a mean energy closer to that of the exact samples than m&m does. this indicates that even with an order of magnitude fewer forward passes in the target model, mh-block is able to produce more accurate samples than the baseline m&m. ## downstream task evaluation controllable generation is a relatively wide field with many tasks. we focus on one of particular importance: style transfer. style transfer is the task of taking text written in one "style" and rewriting it in a different "style" while preserving semantic meaning or "content". for this paper, we focus on the two datasets: the shakespearean author imitation dataset @xcite which provides shakespearean sentences and their modern english counterparts, and the gyafc formality corpus @xcite which contains informal sentences and paired formal versions. data condition following past work in style transfer, our evaluation setup relies on having parallel data of the form (s, t), where @xmath0. we will evaluate our models according to several criteria, some of which only evaluate t (e.g., fluency) and some of which evaluate t with respect to s (e.g., semantic similarity). note that this similarity between model output and the target domain is only used during evaluation and not during model inference. for training and baseline purposes, we assume access to unpaired data belonging to s and t. that is, despite our evaluation requiring paired data, our training setup does not. due to computational constraints, the shakespeare test set was sub-sampled to 100 entries and the gyafc corpus to 300; when evaluating mh-block, we run the markov chain for 20 steps for the shakespeare dataset and 10 for the gyafc dataset. the shakespeare dataset itself contains 31,444 entries, 29,982 of which can be used for training. the gyafc dataset contains 112,890 entries, 105,169 of which can be used for training. ## style transfer results in this section, we will present results of the proposed method on downstream style transfer tasks. quantitative performance is reported in as seen in lo, here upon thy cheek the stain doth sit of an old tear that is not washed off yet. i can tell thee, but young romeo shall be older when thou findest him than when thou first began to look for him. however with lower fluency and significantly lower transfer rate. the grounding of flan-t5-xxl by the expert black-box models shows gains in efficacy especially when compared to prior work investigating the use of mh sampling for style transfer. despite not having an explicit fluency measure, we see our approach has high levels of fluency in all directions. looking at to qualitatively illustrate the difference between the methods, we have also included ## related work controllable generation methods that rely on energy-based constraints are the ones closest to our work @xcite @xcite . mix and match @xcite in particu-lar, is the work closest to ours. their approach relies on single token sampling and masking, rendering the method unable to (1) change the sequence length or (2) perform block sampling of multiple tokens at the same time. our work solves this by enabling block-sampling of multiple tokens through the use of instruction-tuned models. there is also literature exploring free-form or constrained editing of inputs. @xcite follows an editing procedure, with the goal of correcting errors in incorrect code. @xcite uses editing of random sentences sampled from a corpus in place of autoregressive lms to generate fluent natural language text. mallinson et al. (2022) also uses t5 for editing, this time in a 'semi-autoregressive' manner with the goal of combining the quality of autoregressive generation and the speed of non-autoregressive methods. there are a slew of other methods related to ours, where the goal is to steer generation, without the need to re-train models from scratch. in these other approaches, however, there is often the need to use gradients or train auxiliary models to better guide the decoding. one technique guides a large model using smaller discriminator networks with the goal of sampling from an implicitly defined model, an idea explored in plugand-play lm @xcite . in this approach stepwise discriminators are applied to the top-level hidden state to modify the posterior distribution formed by the lm by guiding it to fullfill the desired attributes at each autoregressive generation step by gradient ascent. another work, fudge @xcite , explores a similar idea with reranking the stepwise generations, but additionally explicitly trains the future discriminators on incomplete generations. another set of gradient based methods @xcite view this task as optimizing the generative model's likelihood subject to global differentiable attribute-based constraints by gradient descent. there are also approaches that involve finetuning a backbone language model on domainspecific data @xcite @xcite @xcite or even training from scratch @xcite @xcite @xcite @xcite , to do controllable generation. approaches specifically for style transfer have also been explored by prior work. @xcite frames style transfer as a paraphrasing problem and solves it in an unsupervised way, @xcite has a similar methodology rooted in machine translation. @xcite attempts to model the problem using variational autoencoders. more recently, llms have shown strong efficacy when used for these tasks. chatgpt and gpt3 @xcite are particularly strong performers, able to solve many creative writing tasks in the zero-shot or one-shot regime @xcite . flan-t5 has also shown great fewshot performance despite being less than 1/10th the size of these models @xcite . ## limitations our approach was designed to be as general as possible, however, it is not suitable for all settings. our method relies on having accurate energy models that can model the desired probability distribution. in situations where no such models are available, mh-block is not particularly applicable. additionally, it is best if the desired distribution can be easily described in text, as we must prompt an llm to perform the task; if this is not possible, mixing could be greatly slowed and performance could suffer. however, this issue could be minimized by providing examples of the desired target style to the llm. ## conclusion while we have demonstrated empirically that our novel block mh sampler benefits controllable generation tasks by producing more accurate samples from energy-based lms, our approach may have broader applications in other areas of nlp that use globally normalized models. our approach highlights the utility of separating modeling concerns from inference challenges, potentially paving the way for further approaches that can use llms to impactfully edit text while still giving the system developer fine-grained control of the output.