Gradient-based Constrained Sampling from Language Models (2022)

Abstract: Large pretrained language models are successful at generating fluent text but are notoriously hard to controllably sample from. In this work, we study constrained sampling from such language models, i.e., generating text that satisfies user-defined constraints, while maintaining fluency and model’s performance in a downstream task. We propose MuCoLa—a sampling procedure that combines the log-likelihood of the language model with arbitrary (differentiable) constraints in a single energy function, and then generates samples in a non-autoregressive manner. Specifically, it initializes the entire output sequence with noise and follows a Markov chain defined by Langevin Dynamics using the gradients of this energy. We evaluate MuCoLa on text generation with soft and hard constraints as well as their combinations, obtaining significant improvements over competitive baselines for toxicity avoidance, sentiment control, and keyword-guided generation.

## introduction transformer-based language models (lms) trained on web-scale corpora @xcite @xcite are generating impressively realistic texts. despite having human-level fluency, they are far from reaching human-level communication abilities and are hard to control for content, context, and intent in communication. this results in unreliable models that lack basic knowledge, hallucinate facts, and discriminate users @xcite @xcite . controlled text generation-sampling text from lms to satisfy constraints on the properties of generated text-aims to address these issues. prior works incorporate constraints in existing decoding figure 1 : mucola, our proposed method, stylized as µcola. given a language model, a prompt/input x, and desired constraints defined as thresholds on differentiable functions, we perform langevin dynamics updates to generate the entire output sequence y nonautogressively. we show experiments highlighting both hard and soft constraints ( §4). we propose an algorithm to sample text nonautoregressively from a conditional or unconditional lm trained to perform any language generation task-translation, summarization, dialog, prompt completion-while controlling for multiple, potentially competing constraints, and without sacrificing the base model quality. combining the lm likelihood with constraints into a single "energy" function, our algorithm follows a markov chain @xcite to iteratively transform an output sequence initialized randomly into a desired output with low energy. since common monte carlo markov chain (mcmc) sampling methods can be intractably slow @xcite , we propose to define this markov chain using gradients of the energy function with respect to token embeddings of the output sequence. additionally, we introduce stochasticity in the process in order to generate diverse samples, by modifying the gradients with additive noise, a process referred to as langevin dynamics @xcite @xcite @xcite . finally, we operationalize the energy function by defining each constraint to be smaller than threshold, and writing it as a lagrangian-with language model likelihood as the primary objective @xcite . besides allowing us to combine any number of constraints of varying scales without a need of tuning their weights, we show that low-energy solutions under this definition are true samples from the lm distribution. we call this algorithm mucola for sampling with multiple constraints from lms using langevin dynamics ( §3; also see figure 1 in the appendix). we show the efficacy and generality of mu-cola on a variety of tasks, lms and constraints from prior work ( §4), including soft constraints ( §5) defined using auxiliary models (e.g., classifiers or smaller lms), as well as hard rule-based constraints ( §6) defined by presence or absence of certain keyphrases. we conduct experiments on (a) toxicity avoidance, sentiment control in gpt2-large, and (b) keyword controlled generation with gpt2-large and xl, and (c) entity control to improve fidelity in translation. through both automatic metrics and human evaluation, we show versatility of this method through improved or matched performance with competitive baselines, in terms of quality and diversity. finally, we present preliminary results on new tasks, showing promising directions for future research. ## background: constrained sampling from language models let p (y|x; θ) model the conditional probability distribution of an output token sequence @xmath0, . . . , y n ), given an optional input token sequence @xmath1, . . . , x m ) where x m , y @xmath2. we are interested in constrained sampling from p -finding output sequences y that have a high probability under p while minimizing a given set of constraint functions: {f 1 , . . . , f c }. we assume that each f i : @xcite , y) → r is defined such that a lower value of f i implies that the output better satisfies the constraint. for example, to constrain the outputs to only non-toxic continuations for a given prompt x, we define a classifier p toxic (y) which predicts the output toxicity probability, with lower probability implying lower toxicity. we assume all f i are differentiable. enforcing these constraints in an autoregressive (i.e., left-to-right) decoding strategy like beam search or sampling is challenging, since the constraints are defined conceptually on the whole output sequence and are hard to evaluate accurately only on the generated prefix @xcite . with multiple constraints, their satisfaction/balancing becomes challenging. recent work thus explored non-autoregressive controlled generation @xcite , using constrained optimization over y-finding a single output y which maximizes p given the constraints by performing gradient descent on the outputs y. this involves (1) representing the constrained optimization problem as a single objective e(y) (often referred to as an energy function, discussed in §3.3), and (2) relaxing the discrete outputs y to continuous approximations such that gradient descent is feasible. in previous works, the latter is achieved by creating a soft-representation of y, ỹ = (ỹ 1 , . . . , ỹn ) where each ỹn ∈ r |v| is a simplex (or "logits" which are converted to a simplex using softmax) over the target vocabulary v, representing the probability of the n-th token in the sequence. we refer to these methods as gradientbased decoding. representing the decoding objective as min ỹ e(ỹ) and initializing ỹ with ỹ0 , it is updated as ## constrained sampling via langevin dynamics in embedding space to enable efficient gradient-based sampling from lms, we introduce mucola which modifies the non-autoregressive framework in §2 to (a) generate multiple samples instead of optimizing for only one deterministic output, (b) optimize for much smaller intermediate token representations as opposed to their distribution on the entire vocabulary. ## experimental setup we evaluate mucola on four constrained generation tasks. these tasks are selected based on defining different kinds of constraints for which prior work designed specialized training or decoding mechanisms which cannot be generalized beyond those tasks or language models. the main contribution of mucola is generating diverse samples which conform to the language model p as well as can satisfy user defined arbitrary combination of constraints for which fine-tuning is generally infeasible and tuning weights of each constraint is cumbersome. for a pre-defined sentence length l, we initialize the token representation for each step ẽ1 , . . . , ẽl using token embeddings randomly sampled from the target vocabulary v. 7 for all our experiments, we run the langevin dynamics simulation for a maximum of 250 iterations unless specified otherwise. we describe additional implementation details including noise schedule, stopping criterion and multiplier update schedule in appendix c. ## text generation with soft constraints first, we evaluate mucola with real valued constraint functions defined via auxiliary models such as classifiers or lms. given an lm gpt2-large @xcite , and a prompt x, we generate continuations y. we conduct experiments with: toxicity avoidance and sentiment control. each of the tasks define a binary constraint. let the desired label be denoted by label 1 , and other one with label 0 (label 1 is non-toxic in toxicity avoidance and positive in sentiment control). for both setups, we assume availability of corpora to train the constraint functions. 8 baselines in addition to decoding without any constraints (which we simply call gpt2), we consider the following baselines which decode from left-to-right: constraint functions each of these baselines can be adopted as constraint functions for mu-cola as follows: ## decoding with hard constraints in the previous two tasks, we explored how mu-cola can be applied on soft constraints, defined via real valued functions like probabilities of classifiers or language models. now, we consider a ruledbased constraint that a specific word or phrase must appear in the generated text. existing autoregressive solutions to this task have explored various strategies either based on explicitly modifying probabilities to up-weight desired words @xcite , or search-based strategies based on beamsearch @xcite . in this work, we define a differentiable distance function d(w, ẽ) which measures overlap between desired word (w) and the output token embeddings ẽ (we use the notation w to refer to as the word itself and its index in the vocabulary interchangeably). we then propose a simple criterion to define a threshold ϵ that guarantees that if d(w, ẽ) < ϵ, then w's embedding appears in ẽ (and by extension w appears in y). inspired from @xcite ; @xcite , this distance is computed in three steps (1) define a "probability distribution" for each output token, π @xmath3, . . . , -∥ẽ ne |v| ∥ 2 2 ) where {e 1 , . . . , e |v| } ∈ e, (2) define @xmath4, . . . , g n /τ ), where g @xmath5,w , gumbel softmax provides a way to differentiably sample @xcite and τ is a hyperparameter, and finally, (3) d(w, ẽ) = n @xmath6. intuitively, this function minimizes the euclidean distance between one of the output embeddings (chosen with stochasticity) and w's embedding, e w . this function can be easily extended for phrases, @xmath7, . . . , w l ) by defining based on this definition, for each desired keyword, we define a threshold ϵ w aslog π w,w .we provide an intuitive explanation of the distance function, and a semi-formal and empirical proof of hard satisfaction of this threshold in appendix e. tasks we formally evaluate this setup on two tasks: evaluation following prior work, we measure the performance on two axes, (1) coverage, measured by (a) count average number of keywords appearing in the output; and (b) percent, measuring the fraction of outputs which contain all the desired keywords. (2) fluency, as measured by gpt2-xl perplexity and human evaluation, where on a sample of 200 outputs, we ask 3 annotators to rate each output on a 3-point likert scale. as reported in ## discussion and analysis speed and memory generating a sequence of length l using mucola requires maintaining l×d parameters. in contrast, performing langevin dynamics in the vocabulary space requires l × |v| parameters (|v| >> d). in this analysis, we empirically verify the benefits of our setup. with gpt2-large as underlying lm, we sample sequences of varying lengths with various constraints, on different commercially available gpus using both our approach and an ablation with vocabulary sized representations (logits+softmax; more details in appendix h). as summarized in ## conclusion we present mucola, a sampling algorithm from language models that flexibly combines pretrained lms with any differentiable constraints. our primary contributions are a (1) gradient based mcmc sampling method (langevin dynamics) performed on (2) intermediate representation of tokens (embeddings). with experiments on both soft and hard constraints with different pretrained lms, we show that this approach generates diverse outputs which better conform both to desired constraints as well as the underlying lm distribution. despite the observed improvements, we believe we have barely scratched the surface. in future work, we will explore ways to improve the convergence properties of this algorithm using more sophisticated mcmc algorithms @xcite and develop constraints to improve performance on longer sequences. furthermore, since we perform updates on embeddings rather than vocabulary distributions, future work may also study ways to expand vocabularies at decoding time.