COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics
Lianhui Qin, Sean Welleck, Daniel Khashabi, Yejin Choi
Introduction
Many text generation applications require producing text that is not only fluent, but also satisfies various constraints which control the semantics or style of the generated text. For example (Figure 1), for knowledge-grounded or keyword-guided generation, we might want to ensure that certain keywords are included in the generated output as hard lexical constraints . For other types of text generation, we often wish to incorporate soft topical constraints to contextualize the desired output, e.g., abductively reasoning about what happened in the middle of a story given the past and the future story context . Yet another class of text generation applications requires revising an input based on a new counterfactual condition , which simultaneously requires semantic coherence as well as minimal-edit constraints with respect to the input text .
The dominant paradigm to various text generation applications has been supervised learning with task-specific training data. However, different applications require varied and potentially evolving constraints, and annotating a large amount of task-specific training data for each different combination of constraints can be costly. Recent work has explored incorporating constraints through energy-based text modeling that alleviates the need of supervised data . Yet those approaches still require expensive training of specific generation models. In addition, training might not even be feasible with recent models that are extreme in scale, like GPT-3 . This motivates the need to enrich decoding algorithms that can work directly with pretrained language models without task-specific fine-tuning, and support complex combinations of hard and soft constraints to control the generated text on the fly.
We propose a new constrained decoding approach that formulates decoding as sampling from an energy-based model (EBM) . Constrained generation with our approach amounts to specifying an energy function by plugging in arbitrary constraint functions that are suitable for the task at hand, then sampling from its induced distribution. In particular, to overcome the longstanding challenges of sampling discrete text from EBMs, we for the first time introduce Langevin dynamics to text-based EBMs for efficient gradient-based sampling. As a result, our approach, Constrained Decoding with Langevin Dynamics (Cold), performs sampling by iteratively updating a continuous relaxation of text using gradients of the energy function. The resulting continuous text samples are then mapped back to the discrete space with a simple guided discretization approach, yielding text sequences that are fluent and adhere to the constraints.
Our work makes unique contributions to a recent line of research investigating decoding algorithms for incorporating different constraints in three distinct aspects. First, our formulation unifies various constrained generation scenarios that involve hard lexical constraints and/or soft contextual constraints: specifying an energy function, then sampling from its induced distribution. Second, we propose a sampling method, which complements decoding algorithms that look for a single optimal solution. Finally, we provide new empirical insights into the strengths and weaknesses of existing approaches to discrete search and differentiable reasoning.
To test the flexibility and empirical performance of Cold decoding, we experiment with three challenging text generation tasks: lexically constrained generation , abductive reasoning , and counterfactual story generation . Cold achieves better lexical coverage than NeuroLogic , a beam-based discrete decoding algorithm specifically designed for lexically constrained generation, while producing more coherent and higher quality text than Delorean , a state-of-the-art gradient-based generation method for abductive reasoning and counterfactual reasoning. Cold supports all three constrained generation settings under a unified framework – specifying an energy function using a collection of fluency and task-specific constraints, then sampling from its induced distribution and achieves strong performance on both automatic and human evaluation.
Background
Neural text generation. Neural text generation typically involves two stages: modeling a distribution over text sequences, and using a decoding algorithm to generate sequences with the model. Let denote a discrete sequence where each is a token from a vocabulary . Common neural language models (e.g., GPT-2/3 ) factorize the probability of a sequence into the product of per-token conditionals in left-to-right order, , with each conditional parameterized by a shared neural network, such as transformer . Popular decoding algorithms, ranging from beam search or greedy decoding to sampling methods such as top- or nucleus sampling, produce text sequences using the model , often conditioned on a prompt .
Constrained text generation. We view text generation as the problem of finding a sequence that satisfies a collection of constraints. For instance, the scenario above amounts to generating a sequence subject to a soft constraint that the continuation should be fluent and logically coherent with the prompt . Other constrained generation problems impose additional constraints, such as text infilling where coherence constraints move beyond a left-hand prefix, lexically constrained generation in which hard constraints require the output to contain given tokens, and various forms of semantically-constrained generation in which the output is softly constrained to be similar to another sequence. Since common decoding algorithms generate text monotonically, relying on for determining the next token, it is challenging to enforce these diverse constraints.
Despite the flexibility, however, sampling from an EBM is particularly challenging, as computing is intractable. Common gradient-free Markov chain Monte Carlo (MCMC) methods such as Gibbs sampling can be used, but they are often prohibitively slow . Langevin dynamics , a gradient-based MCMC method, offers more efficient sampling by using the gradient of the energy function , enabling sampling in domains such as image generation . However, since text is discrete, the gradient is not well-defined, making it non-trivial to apply Langevin dynamics for sampling text from an EBM. Our approach bridges this gap with continuous relaxation of text, differentiable constraints, and guided discretization, as described below.
COLD Decoding with Langevin Dynamics
To enable flexible and diverse constrained generation in off-the-shelf language models, we develop Constrained Decoding with Langevin Dynamics (Cold), a decoding approach that treats text generation as sampling from an energy-based distribution, allowing for flexibly composing constraints based on the task at hand. Cold decoding generates text by sampling from an EBM defined over a sequence of “soft” tokens using Langevin dynamics, then maps the continuous sample into discrete, fluent text. We provide our formulation of constrained text generation (§3.1), present differentiable constraints that can be composed into energy functions (§3.2) along with our discretization method (§3.3), and discuss practical details of Cold decoding (§3.4). Figure 2 provides an overview.
The set of constraints induces a distribution over text, written in an energy-based form as:
where is the weight of the th constraint, is the normalizing factor. Here is the energy function. This energy-based form is flexible, as one can plug in any constraint functions required for a task of interest. Generating text under the constraints can then be seen as sampling from the energy-based distribution . One can also draw multiple samples and pick the best if only one sample is needed, as discussed later (§3.4).
As mentioned above, for efficient sampling from we want to use Langevin dynamics, which makes use of the gradient . However, in our case is a discrete sequence and the gradient is not well-defined. As a result, we perform Langevin dynamics with an energy defined on a sequence of continuous token vectors, described below.
2 A Collection of Cold Constraints
Cold provides a flexible framework for plugging in a wide range of constraint functions for a task of interest. We describe constraint functions that are useful in various constrained generation problems, such as those we consider in the experiments (§4). The constraints include language model-based fluency constraints, along with lexical and semantic constraints on the sequence content. More generally, any differentiable function that outputs a goodness score of (soft) text can be used as a constraint function, as long as it reflects the requirements of the target task.
Soft fluency constraint. Fluency is a common requirement for generated text. To promote fluency, we use a constraint which favors soft sequences that receive high probability according to the underlying left-to-right LM (e.g., GPT2):
Future-token prediction constraint. Applications such as text infilling involve future input tokens that remain fixed, but should contribute to updating past positions. For instance, consider updating the second position of The __ has eight legs. A sample should be coherent with the tokens on the right (i.e., has eight legs).
To this end, we use a constraint that adjusts soft tokens to maximize the likelihood of input tokens ,
N-gram similarity constraint. Many constrained generation scenarios pose requirements on the wording and expression of generated text sequences. For instance, lexically constrained generation tasks require certain keywords to be presented in the text samples, while counterfactual reasoning or text editing tasks require the text to retain the essence of a reference sequence.
We formulate these requirements as an -gram similarity constraint which favors sequences that overlap with a reference at the -gram level,
3 From Soft to Discrete and Fluent Text
We refer to this method as “top- filtering”. The resulting text tends to be fluent because each token is among the top- most probable tokens from the LM . In practice, to ease the satisfaction of certain constraints (e.g. -gram similarity), we expand the candidate set to include constraint tokens (e.g., in the tasks of abductive reasoning §4.1 and lexically constrained decoding §4.3).
Figure 2 illustrates the decoding procedure to get one output from Cold decoding. Algorithm 1 summarizes the algorithm. Next, we move to practical considerations of applying Cold.
4 Implementation of Cold Decoding
Noise schedule. Each iteration of Langevin dynamics adds noise to the gradient (Eq. 2). We gradually decrease across iterations, which intuitively transitions the decoding procedure from exploration to optimization. In our experiments, we typically used the schedule which sets/reduces to at iterations , respectively.
Long sequences. Cold decoding produces a fixed-length sequence . To produce longer sequences, e.g. in cases where is not the end of a sentence, we use to produce a continuation of using greedy decoding.
Experiments
We evaluate Cold on three constrained generation tasks. Using Cold for each task amounts to specifying a set of task-specific constraints (instances of those in §3.2). Our focus is enabling constrained generation for settings in which fine-tuning is infeasible, through changing the decoding method. Thus, our experiments (i) use off-the-shelf LMs without fine-tuning, and (ii) compare Cold primarily against alternative decoding methods. As our base LM, we use GPT2-XL .
We study a specific formulation of abductive reasoning as a language generation challenge. Specifically, given a beginning sentence and an ending sentence , the abductive language generation (NLG) problem consists of generating a bridge sentence that fills in between the two sentences and forms a coherent full story (see Figure 1 for example). The task is particularly challenging for conventional monotonic left-to-right LMs (such as GPT-2 and GPT-3) since it requires non-monotonic reasoning that not only conditions on the past context (, on the left), but also the future story ending (, on the right).
Cold decoding can readily accommodate the abductive reasoning task by simply plugging in appropriate constraints to specify an energy function. Specifically, the generated text needs to be (1) fluent and consistent with the left context , and (2) coherent with the right context . Accordingly, we compose an energy using relevant constraints from §3.2:
That is, we combine (a) a soft fluency constraint (Eq. 3) conditioning on the left sentence to enforce fluency and consistency with the left context, and a reverse fluency constraint with a right-to-left LM conditioning on to encourage coherence with the right context; (b) a future-token prediction constraint (Eq. 4) that enforces consistency between the generation and the story ending ; (c) a -gram similarity constraint (Eq. 5) between the generation and keywords (non-stopwords) in (excluding those in ), i.e., , which intuitively promotes a ‘smooth transition’ between , , and .
For the energy function in Eq.(7), we select the constraint weights on the dev set. Throughout the experiments, we set the number of Langevin dynamics steps to , with a step size (Eq. 2). We discuss more details of the configurations in the appendix.
Baselines. We compare with previous decoding approaches for this task. In particular, we compare with Delorean which outperformed a wide range of supervised and unsupervised methods on the abductive reasoning task in Qin et al. . Following Qin et al. , we also compare with a Left-only method that generates the continuation of without considering the right-side , i.e., .
Evaluation. We perform both automatic and human evaluation. We adopt the standard automatic metrics on the task that measure the minimal edit between the generated text and the human-written references on the test set, including BLEU , ROUGE , CIDEr , and BERTScore . For the human evaluation, we follow and let crowdworkers from Amazon Mechanical Turk rate the generations on 200 test examples. Workers were presented a pair of observations ( and ) and a generated hypothesis , and asked to rate the coherence of the hypothesis with respect to the observation (i.e., ), the observation (i.e., ), and both (i.e., ), as well as the grammaticality of the hypothesis itself, on a 5-point Likert scale. The average ordinal Krippendorff alpha () is 0.36, indicating a fair inner-annotator agreement.
1.2 Results
Table 1 shows the evaluation results on the abductive reasoning task. Under automatic evaluation (the left panel), Cold consistently outperforms the previous best unsupervised decoding algorithm Delorean, as well as the Left-only method, in terms of both the lexical overlap metrics (BLEU, ROUGE and CIDEr) and semantic similarity metric BERTScore. The human evaluation (the right panel) provide more fine-grained insights. Cold achieves the best overall coherence, meaning that the generated from Cold fits best with both the left-side context and the right-side context compared to the other methods. In contrast, Delorean excels only in terms of the left-side coherence (with ), with inferior right-coherence (with ). We speculate this is because of Delorean’s complex interleaving of forward and backward decoding passes that make it difficult to balance the left- and right-coherence constraints. In terms of grammaticality, unsurprisingly, Left-only obtains the best score as it ignores any other constraints (and fails this task with low coherence scores). More importantly, Cold achieves a high grammaticality score along with its high coherence, substantially improving over Delorean. Example generations in Appendix Table 7 show how Cold can reason with the right-hand context (e.g. ‘no heels’), while Delorean’s generations are contradictory (‘red shoes’ vs. ‘white pair’) or equivalent to those from Left-only.
2 Counterfactual Story Rewriting
Next, we consider counterfactual story rewriting . Given a story context with ending , the task is to generate a new story ending that is (i) similar to the original ending , yet (ii) consistent with a new story context (see Figure 1 for example). The task is challenging as it requires capturing the aspects of future events that are invariant under the new (counterfactual) context, while only making necessary edits for coherence.
To tackle this task, we use Cold with an energy composed of constraint functions that promote coherence with the new context , and minimal edits to the original ending :
Baselines. Similar to the setup in §4.1, we compare with DeLorean , a recent state-of-the-art decoding algorithm. As a reference, we also report the performance of a trivial solution, Left-only, that generates a continuation of without considering the minimal edit constraint with the original ending . Thus the method is expected to generate a coherent ending which however does not necessarily resemble the original ending. Finally, we compare with Mix-and-Match , a recent energy-based decoding method with discrete MCMC sampling, using BERT-base and BERT-Large.
2.2 Results
Table 3 shows the results of automatic and human evaluation in terms of both minimal-edit and coherence. As expected, the reference method Left-only that completely ignores the minimal edit constraint can easily generate a new ending that is coherent with the new context . Compared to the baseline approach Delorean, our method Cold achieves overall superior performance, with substantially improved coherence score and comparable minimal-edit score by human evaluation. Mix-and-Match, based on discrete MCMC sampling, performs poorly. Intuitively, its discrete sampling tends to get stuck in a mode of the target distribution (i.e., the region surrounding the original story ending), and struggles to explore further to find samples of interest. Cold’s gradient-based sampling with continuous approximation leads to more efficient and effective exploration and mixing, as evidenced by samples that better meet the task requirements. See Appendix for examples.
3 Lexically Constrained Decoding
Next, we use Cold for lexically constrained decoding. Given a set of words , the task aims to generate a coherent sentence that contains these words (Figure 1). The task is challenging as it requires proper planning to coherently include the constraint words.
We specify an energy function of the following form:
Baselines. We compare with a recent state-of-the-art method NeuroLogic , a beam-search variant specifically designed for lexically constrained generation which outperformed many supervised and unsupervised approaches in Lu et al. . We also report the results of TSMH as another recent baseline which uses Monte-Carlo Tree Search .
Evaluation. We use the set of constraint words from the CommonGen corpus , but adopt the canonical setting that the generated text must contain the exact constraint words (e.g., write) instead of their variants (e.g., wrote) . Following previous works , we report a measure of constraint words coverage as well as language fluency by evaluating the perplexity of the text . We also ask crowdworkers to rate the text fluency on a 3-point Likert scale on 200 test examples. The average ordinal Krippendorff alpha is 0.29, indicating a fair inner-annotator agreement.
3.2 Results
Table 3 shows the evaluation results for the lexically constrained decoding task. Cold, a general constrained decoding method, is comparable to the state-of-the-art method NeuroLogic designed specifically for dealing with lexical constraints. In particular, Cold achieves a higher coverage of given keywords, at the expense of generating slightly less fluent language. Cold is also substantially better than lexically constrained decoding method TSMH in terms of both coverage and fluency.
4 Additional Analysis
We ablate two important ingredients of our approach, namely the constraints and the top- filtering. Due to space limit, we report the results of constraints and defer the results of top- filtering to the appendix. Table 5 shows the human evaluation results for ablations of the constraints used on the abductive reasoning task (Eq. 7). The -gram similarity constraint provides the largest contribution to the overall coherence. The reverse LM fluency constraint also to some extent helps with the right-side coherence by conditioning on the right-side content . Removing the future-token prediction constraint similarly causes inferior scores in terms of right-side and overall coherence, as expected. Removing the individual constraints leads to better grammaticality due to less competition among different constraints, at the cost of coherence. Our uniform treatment of all constraints as energy terms makes it straightforward to balance the different constraints by controlling the constraint weights.
We report the average runtime of generating one sample on the Counterfactual Story Rewriting data. The table below shows the results (on an NVIDIA Quadro GV100 GPU, batch size=32). We compare with Mix-and-Match , a recent energy-based decoding method with discrete MCMC sampling (Metropolis-Hastings, in particular). Cold, which uses gradient-based sampling, is faster than the gradient-free Mix-and-Match: Cold is 30% faster with base LMs of similar sizes (GPT2-M and BERTLarge), and has roughly the same time cost when using a much larger LM (GPT2-XL).
Related Work
Previous works proposed beam search variants for lexically constrained decoding which enforce constraints in a discrete space. Recent works consider constraint satisfaction by adjusting vocabulary distributions using an additional discriminator or LM . Differing from those approaches that determine the generation token by token auto-regressively, Qin et al. optimize the whole (soft) token sequence via gradient propagation, which facilitates sequence-level semantic constraints (e.g., right-coherence, minimal-edits). Cold also samples complete sequences, while offering a principled and unified formulation based on energy-based modeling. Kumar et al. extend by imposing constraints with a Lagrangian method and optimizing for a single output with gradient descent. In contrast, our approach based on energy-based sampling (§3.1) allows for generating samples for other utilities (e.g., rank-and-select §3.4, estimating expectations). We also introduce components for more fluent generations such as the novel discretization procedure. Also, on the empirical side, we explore a different class of problems and tackle them in the absence of labeled data. The recent CGMH and TSMH , followed by , perform constrained decoding with extended Gibbs sampling or Metropolis-Hastings sampling in the discrete text space. Our energy-based formulation with gradient-based Langevin dynamics sampling produces substantially better results than the discrete TSMH (§4.3). Sha uses gradient information to guide generation, which, however, is specifically designed for lexically constrained generation.
Energy-based models (EBMs) have been used for incorporating additional information to train text generation models . In contrast, we focus on the constrained decoding (inference) that can be directly applied to pretrained LMs without fine-tuning. Langevin dynamics is widely used on EBMs of modalities with continuous values, like images , 3D shapes , latent features , and audio sequences . To our knowledge, we are the first to apply Langevin dynamics for (constrained) discrete text generation (with a continuous approximation) for efficient sampling.
Conclusion
We introduce Cold decoding, an energy-based constrained text generation framework that can express various soft/hard constraints through an energy function, and sample using Langevin dynamics. Cold can be applied directly to off-the-shelf LMs without task-specific fine-tuning. We showcase its flexibility and strong performance on three distinct applications of constrained text generation.
Acknowledgements
This work was funded in part by the Natural Sciences and Engineering Research Council of Canada (NSERC) (funding reference number 401233309), DARPA MCS program through NIWC Pacific (N66001-19-2-4031), the Allen Institute for AI, and Microsoft Research PhD Fellowship. We thank the XLab research group, and our anonymous reviewers for their feedback on this work. We also acknowledge the Beaker team (https://beaker.org) for their support with experiments.
References
Checklist
Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]
Did you describe the limitations of your work? [Yes]
Did you discuss any potential negative societal impacts of your work? [Yes] , see the Appendix.
Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
If you are including theoretical results…
Did you state the full set of assumptions of all theoretical results? [N/A]
Did you include complete proofs of all theoretical results? [N/A]
Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes]
Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] .
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No] Due to computational resource constraints, we didn’t run multiple cross-validation splits, or with enough random seeds to form stable confidence intervals. However, we do a thorough set of ablations across model configurations.
Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Details are in appendix.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…
If your work uses existing assets, did you cite the creators? [Yes]
Did you mention the license of the assets? [No] : we don’t introduce new datasets, and refer readers to the original releases in case license information for those works changes.
Did you include any new assets either in the supplemental material or as a URL? [No]
Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [Yes] All data we experiment with is public.
Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No] We aren’t releasing new data, and existing corpora, to our knowledge and in our experience, do not contain personally identifying information.
If you used crowdsourcing or conducted research with human subjects…
Did you include the full text of instructions given to participants and screenshots, if applicable? [Yes] See Appendix.
Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] Crowdworking studies involving no personal disclosures of standard NLP corpora are not required by our IRB to be reviewed by them. Specifically:
We do not collect personal information. Information gathered is strictly limited to general surveys about the quality of generated text.
We take precaution to anonymize Mechanical Turk WorkerIDs in a manner that the identity of the human subjects cannot be readily ascertained (directly or indirectly).
We do not record or include any interpersonal communication or contact between investigation and subject.
Additional locality-specific details of our IRB withheld to preserve anonymity, but will be made public upon de-anonymization.
Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [Yes] , our pay is always over $15 USD per hour on average (and sometimes more, see Appendix.
Appendix A Ethical Considerations
Automatic text generation, though powerful in generating fluent human-like language, could be potentially used for malicious purposes, such as generating toxic, biased, offensive, or fake information. We hope that our research, as a method to control language model generations by plugging in constraints, can provide a way for steering and harnessing the LMs to alleviate those ethical issues.
Appendix B Experimental Configurations
Configurations of Abductive Reasoning. For the energy function in Eq.(7), we select the constraint weights on the dev set. The overall weight of the fluency constraints is set to 0.5, wherein the and constraints are balanced with a 6:4 ratio, leading to and . The remaining weight 0.5 is assigned to the constraints (b) and (c), with a ratio of 1:0.05, leading to and . Throughout the experiments, we set the number of Langevin dynamics steps to , with a step size (Eq. 2). The text decoded by Cold is set to have length 10 and is completed by the base LM as described in §3.4. We set the for top- filtering. For each (), we generate 16 samples and pick the best one by first ranking by the perplexity of the joint sequence for overall coherence, and then from the top 5 candidates selecting the best one in terms of the perplexity of for enhanced coherence with the right context.
Configurations of Counterfactual Story Rewriting. The constraint weights in the energy function in Eq. (8) are selected on the dev set. The weights of the constraints (a) and (b) are set to and , respectively. For the LM and reverse LM fluency constraints in (a), we use a ratio of 8:2, leading to and . We largely follow the algorithm configurations in §4.1 except that the text length is set to 20, for top- filtering, and we generate 32 samples for each test example and pick the best one ranked by the perplexity of .
Configurations of Lexically Constrained Decoding. The weights of the constraints in energy function Eq. (9) are the same as those in the abductive reasoning task (§4.1) except for the ratio of the n-gram similarity constraint, which is increased to 1:0.1 between constraints (b) and (c), leading to and . We set the for top- filtering. All other configurations are the same as those in §4.1.
Right-to-left language model. The right-to-left LM is publicly released by West et al. . Specifically, the LM was trained following GPT-2 using the OpenWebText training corpus (see section 2.4 in West et al. ).
Computing. All experiments were conducted using a server with 8 NVIDIA V100 GPUs.
Appendix C Human Evaluation Details
We conduct human evaluation for 3 tasks: 1)Lexically Constrained Generation 2)Abductive Reasoning 3)Counterfactual Reasoning. We sampled 200 prompts randomly from the corpus for each human evaluation. We shuffle HITs to eliminate systematic bias of rater availability by time. Figures show the screenshot of instructions for our human evaluation.
C.2 Human Evaluation Payment
Mean hourly pay was determined using a javascript timing tool to be $15/hr.
Appendix D Ablation Study: Top-k Filtering
We investigate the role of top- filtering mechanism (§3.3). Specifically, we investigate its effect on the output performance for different values in Table 6. We can see that the grammar score tends to decrease as increases. This is expected since a larger indicates more flexibility for the generation to satisfy other constraints, often at the expense of fluency. The left coherence shows a similar relationship with the value since it is also enforced by the left-to-right LM through the soft fluency constraint (Eq.3). In contrast, the right and overall coherence generally benefits from a larger due to the increased flexibility for choosing the right words. Interestingly, with a large value (50, 100), the right/overall coherence no longer improves, probably due to the inferior fluency that has affected the meaning and coherence of the generation.
Appendix E Generated Samples
Tables 7, 8, and 9 show generated samples for the abductive reasoning, counterfactual reasoning, and lexically constrained decoding tasks, respectively.