COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics

In plain words

Language models write one word at a time, left to right, so it is hard to make them meet requirements that depend on words still to come, such as "use these keywords" or "end the story in this way". This paper's method, COLD, does not train anything. It keeps a whole draft as a grid of soft scores, one row per word position, instead of a fixed list of words.

Each requirement is written as a score for how well a draft fits it. The scores are added up into one "energy" where lower is better. The draft is then nudged downhill, with a little random jitter that shrinks over time, so that it settles into drafts that satisfy the requirements while still reading naturally. At the end, each position is turned into a real word chosen from the model's own likely next words.

Because the requirements are just added terms, the same machinery covers hard rules (include certain words), soft context (fit the text between a start and an end) and minimal-change rewriting. The authors apply it to three tasks with an off-the-shelf GPT-2 model.

TL;DR

• Decoding as sampling: an energy over soft token logits, sampled with Langevin dynamics, then mapped to text with top-k filtering using the same LM.
• Constraints are plug-in terms: LM fluency in both directions, a future-token prediction term, and a differentiable n-gram match.
• Off-the-shelf GPT2-XL, no fine-tuning, three tasks: abductive reasoning, counterfactual story rewriting, lexically constrained generation.
• Author-reported gains over DeLorean and NeuroLogic; the runtime is 33.6 s per sample on counterfactual rewriting, versus 33.5 s for the gradient-free Mix-and-Match baseline tab. 5.
• The released code departs from the equations in several places (optimizer, noise schedule, loss scaling, candidate sets, ranking). Read "Paper vs code" before implementing.

Problem

Left-to-right LMs cannot easily enforce constraints that depend on later words. Decoding methods that enforce hard lexical constraints (beam variants) or sequence-level semantics usually need either task-specific training data or a discrete search that can be slow or get stuck §1 §2 §4.2. The authors frame this as sampling from an energy-based model over text, which is intractable to normalise and slow with Gibbs-style sampling. Text is discrete, so the gradient of the energy is undefined §2.

Method

Soft sequence. Each position t holds a vector of logits over the vocabulary \mathcal{V}, not a token. A softmax with temperature \tau turns it into a distribution that approaches one-hot as \tau \to 0 §3.1. The paper does not give a numeric \tau.

Energy-based target. The target distribution over text is p(\mathbf{y}) \propto \exp\{\sum_i \lambda_i f_i(\mathbf{y})\} with weights \lambda_i \ge 0 and constraint scores f_i (higher means better), i.e. energy E = -\sum_i \lambda_i f_i eq. 1 §3.1. The sample is drawn on the soft sequence \tilde{\mathbf{y}}, so the gradient is defined.

Langevin update. Starting from an initial soft sequence, iterate \tilde{\mathbf{y}}^{(n+1)} = \tilde{\mathbf{y}}^{(n)} - \eta \nabla E(\tilde{\mathbf{y}}^{(n)}) + \epsilon^{(n)} with \epsilon^{(n)} \sim \mathcal{N}(0, \sigma) eq. 2 §3.1. The paper argues that with annealed step size and vanishing noise the chain converges to the target distribution, in the limit only §3.1.

Constraint terms §3.2:

• Soft fluency f_{\text{LM}}^{\to}: the cross-entropy-style match between each soft position and the LM's next-token distribution given the preceding soft tokens. The preceding tokens are fed as softmax-weighted embeddings eq. 3 §3.2.
• Reverse fluency f_{\text{LM}}^{\leftarrow}: the same with a right-to-left LM, conditioned on the tokens to the right §3.2.
• Future-token prediction f_{\text{pred}}: the likelihood the LM assigns to fixed right-hand tokens \mathbf{x}_r after the soft sequence eq. 4 §3.2.
• N-gram similarity f_{\text{sim}}: a differentiable n-gram match to a reference \mathbf{y}_*, described as an approximation of BLEU-n eq. 5 §3.2.

f_{\text{LM}}^{\to}(\tilde{\mathbf{y}}) = \sum_{t=1}^{T} \sum_{v \in \mathcal{V}} p_{\text{LM}}^{\to}(v \mid \tilde{\mathbf{y}}_{<t}) \log \mathrm{softmax}(\tilde{\mathbf{y}}_t(v)) eq. 3

f_{\text{pred}}(\tilde{\mathbf{y}}; \mathbf{x}_r) = \sum_{k=1}^{K} \log p_{\text{LM}}^{\to}(x_{r,k} \mid \tilde{\mathbf{y}}, \mathbf{x}_{r,<k}) eq. 4

f_{\text{sim}}(\tilde{\mathbf{y}}; \mathbf{y}_*) = \texttt{ngram-match}(\tilde{\mathbf{y}}, \mathbf{y}_*) eq. 5

Task energies §4.1.1 §4.2 §4.3:

• Abductive reasoning: two fluency terms, future-token prediction on the right context, and a 1-gram similarity to the keywords of the right context that are not in the left context eq. 7.
• Counterfactual rewriting: a fluency term conditioned on the new context, a reverse fluency term, and an n-gram similarity (n = 2, 3) to the original ending eq. 8.
• Lexical constraints: two fluency terms, a 1-gram similarity to the given words, and future-token prediction on the words concatenated in arbitrary order eq. 9.

E(\tilde{\mathbf{y}}) = \lambda_a^{lr} f_{\text{LM}}^{\to}(\tilde{\mathbf{y}}; \mathbf{x}_l) + \lambda_a^{rl} f_{\text{LM}}^{\leftarrow}(\tilde{\mathbf{y}}; \mathbf{x}_r) + \lambda_b f_{\text{pred}}(\tilde{\mathbf{y}}; \mathbf{x}_r) + \lambda_c f_{\text{sim}}(\tilde{\mathbf{y}}; \text{kw}(\mathbf{x}_r) - \text{kw}(\mathbf{x}_l)) eq. 7 §4.1.1

E(\tilde{\mathbf{y}}) = \lambda_a^{lr} f_{\text{LM}}^{\to}(\tilde{\mathbf{y}}; \mathbf{x}_l') + \lambda_a^{rl} f_{\text{LM}}^{\leftarrow}(\tilde{\mathbf{y}}) + \lambda_b f_{\text{sim}}(\tilde{\mathbf{y}}; \mathbf{x}_r) eq. 8 §4.2

E(\tilde{\mathbf{y}}) = \lambda_a^{lr} f_{\text{LM}}^{\to}(\tilde{\mathbf{y}}) + \lambda_a^{rl} f_{\text{LM}}^{\leftarrow}(\tilde{\mathbf{y}}) + \lambda_b f_{\text{sim}}(\tilde{\mathbf{y}}; \mathcal{W}) + \lambda_c f_{\text{pred}}(\tilde{\mathbf{y}}; c(\mathcal{W})) eq. 9 §4.3

Discretisation §3.3. The soft sample is mapped to tokens one position at a time. The candidate set at position t is the LM's top-k next tokens given the tokens chosen so far. The token is the candidate with the largest soft logit eq. 6. The candidate set is enlarged with constraint tokens in the abductive and lexical tasks §3.3.

y_t = \arg\max_{v \in \mathcal{V}_t^{k}} \tilde{\mathbf{y}}_t(v) eq. 6

Algorithm. The paper's summary of the sampler is alg. 1. The pseudocode below follows the paper's equations, not the released code (see "Paper vs code").

# COLD decoding with Langevin dynamics  [alg. 1] [§3.1]
# Input:  LM, weighted constraints (weight, f), length T, steps N
# Output: token ids of length T

def cold_decode(lm, constraints, length, steps):
    # Initial soft sequence, one logit vector per position  [§3.4]
    soft = initialize_soft_sequence(lm, length)          # (T, V)
    for step in range(steps):
        # Noise std for this step from the annealing schedule  [§3.4]
        sigma = noise_schedule(step)
        # Energy is minus the weighted sum of constraint scores  [eq. 1]
        energy = -sum(weight * f(soft) for weight, f in constraints)
        # Langevin step: gradient descent on the energy plus noise  [eq. 2]
        soft = soft - eta * grad(energy, soft) + normal(0, sigma)
    # Discretise left to right with top-k filtering  [eq. 6] [§3.3]
    tokens = []
    for t in range(length):
        # Next-token candidates from the LM given tokens so far  [§3.3]
        candidates = lm_top_k(lm, tokens, k)
        # Keep the candidate with the largest soft logit  [eq. 6]
        tokens.append(argmax_over(soft[t], candidates))
    return tokens

Sample-and-select. Several samples can be drawn; one is picked by a criterion such as an energy term §3.4. The task-specific rules are in the experimental configurations app. B.

Long sequences. The sampled sequence has a fixed length. If it ends mid sentence, the base LM continues it with greedy decoding §3.4.

Hyperparameters

Values are from the paper app. B unless the Where column says code. Values in the code files are the arguments of the released run scripts (commit 0cda54b).

Results

Author-reported. Setup: GPT2-XL, test sets of each benchmark, comparison with the baselines listed in each table. Automatic metrics against human references. The paper gives no test-set sizes for automatic metrics. Table and figure numbers were counted from the LaTeX source, not read from the rendered PDF.

Abductive reasoning (αNLG-style bridge generation) tab. 1 §4.1.2:

Automatic columns are the metrics in §4.1; human columns are 5-point Likert ratings on 200 examples §4.1. The paper says COLD gives the best overall coherence; DeLorean is higher on left coherence (4.23 vs 4.00).

Counterfactual story rewriting tab. 2 §4.2:

Mix-and-Match has a much higher Min-Edit overlap but a lower BERTScore coherence, and was excluded from human evaluation §4.2. Human evaluation uses a 3-point scale.

Lexically constrained decoding (CommonGen, canonical setting) tab. 3 §4.3:

COLD reports higher keyword coverage than NeuroLogic (94.50% vs 91.00%), at the cost of higher perplexity (54.98 vs 28.61) and lower human fluency (2.07 vs 2.53) §4.3.

Ablation of constraints on abductive human evaluation (125 examples) tab. 4 §4.4:

Top-k filter ablation (abductive setting) tab. 6 app. D:

Runtime per sample on counterfactual rewriting, batch size 32, one NVIDIA Quadro GV100 tab. 5 §4.4:

The paper reports COLD 30% faster than Mix-and-Match with similar-size LMs, and about equal cost with GPT2-XL §4.4. Appendix B says experiments ran on 8 V100 GPUs app. B; the two hardware statements differ and the paper does not say which applies to which result.

Paper vs code

Code paths are relative to the repository at code_commit. The code is Apache-2.0. Each entry is a difference between the equations or text and the released code, with both citations.

1. Optimizer. The paper's update is a plain Langevin step with step size \eta eq. 2 §3.1. The code optimises an offset epsilon with torch.optim.Adam at learning rate \eta, and a StepLR with ratio 1, so no decay cold_decoding.py cold_decoding.py cold_decoding.py. \eta therefore sets an Adam learning rate, not a gradient scale.
2. Noise placement. Section 3.4 says the noise is added to the gradient §3.4, but Eq. 2 adds it to the iterate eq. 2. The code adds the Gaussian noise to the base logits each step cold_decoding.py.
3. Counterfactual noise schedule. The paper says counterfactual follows the abductive configuration app. B. The counterfactual script uses thresholds 50, 200, 500 with stds 0.5, 0.1, 0.05 counterfactual.sh; the abductive script uses the paper's schedule abductive.sh.
4. Counterfactual fluency split. The paper gives 8:2, i.e. \lambda_a^{lr}=0.64, \lambda_a^{rl}=0.16 app. B. The counterfactual script passes --lr-nll-portion 0.9, which gives 0.72 and 0.08 for the two fluency terms counterfactual.sh counterfactual.sh cold_decoding.py.
5. Weight scale. The code loss is (1-c)[\alpha L_{lr} + (1-\alpha) L_{rl}] + c\,L_{\text{constraint}} with no normalisation cold_decoding.py. For abductive (c=0.5, \alpha=0.6, second constraint weight 0.05) the effective weights are 0.3, 0.2, 0.5 and 0.025, which sum to 1.025. The paper's 0.48 and 0.02 are the normalised shares app. B. The ratio 1:0.05 is kept.
6. Lexical weights. The paper's lexical values (sim 0.05, pred 0.45) sum to 1.0 app. B. The code gives pred 0.5 and sim 0.05 commongen.sh commongen.sh.
7. Labels in the appendix. The abductive paragraph in the appendix names the weights \lambda_c^{lr}=0.48 and \lambda_a^{rl}=0.02, but Eq. 7 has no such symbols; the values appear to belong to \lambda_b (pred) and \lambda_c (sim) app. B eq. 7.
8. Term averaging. Eqs. 3 and 4 sum over positions and tokens eq. 3 eq. 4. The code averages instead: the soft NLL averages over the sequence length util.py, and the future-token cross-entropy takes a mean over the right-context tokens cold_decoding.py. The effective weights therefore depend on sequence length.
9. Reverse LM input. The paper defines the reverse fluency term on soft tokens §3.2. In the counterfactual code the reverse model receives the hard argmax of the soft sequence, so its predicted distribution is a constant; the soft sequence receives gradient only through the log-softmax term cold_decoding.py. For abductive and lexical, the reverse term uses soft inputs with a vocabulary offset of one token util.py cold_decoding.py and a detached repetition penalty that subtracts 1e4 times the next-position softmax cold_decoding.py. The paper mentions neither.
10. Straight-through and temperatures. The paper introduces \tau without a value §3.1. All run scripts pass --straight-through. The forward pass uses logits divided by 0.001, with the gradient passed straight through cold_decoding.py. The future-token term uses 0.3 cold_decoding.py; the backward LM input uses 1e-5 cold_decoding.py.
11. Initialisation. The paper uses greedy-decoding logits §3.4. The code default init_mode is random, which runs a soft rollout: each step feeds the softmax-weighted embedding of the previous logits back in, at temperature 1, with no sampling cold_decoding.py util.py. The original mode initialises from the target ending, which the paper does not describe cold_decoding.py.
12. Last update skipped. The update and noise are applied only for iter < num_iters - 1, so the decoded soft sequence reflects N−1 updates cold_decoding.py cold_decoding.py. The paper's Algorithm 1 has N updates alg. 1.
13. Best-iterate option unused. --min-iters is defined as "record best only after N iterations", but nothing reads it cold_decoding.py. Decoding uses the final iterate cold_decoding.py.
14. Candidate set for counterfactual. The paper enlarges the top-k candidates with constraint tokens for abductive and lexical only §3.3. The counterfactual code also adds the non-stopword tokens of the original ending to every candidate set cold_decoding.py.
15. Abductive keywords. The paper's constraint uses keywords of the right context that are not in the left context eq. 7 §4.1.1. The keyword filter against the left context runs only with --abductive-filterx, which the abductive script does not pass cold_decoding.py abductive.sh. Keyword extraction also keeps proper nouns and drops adverbs cold_decoding.py, which the paper does not describe.
16. Abductive selection not in code. The paper selects among 16 samples: the top 5 by perplexity of the full sequence, then the best of those by perplexity of the right-hand part app. B. The abductive script writes all candidates with no ranking cold_decoding.py, and the README says ranking "will add soon" README.md.
17. Counterfactual selection adds a BLEU rerank. The paper picks the best of 32 samples by perplexity of the counterfactual premise plus output app. B. The code de-duplicates, drops repeated or very short candidates, ranks by perplexity, keeps 1 to 3 candidates depending on the perplexity gaps, and then picks the lowest n-gram loss against the original ending util.py cold_decoding.py.
18. Lexical output not selected. The lexical code writes all 32 × 8 = 256 candidates per concept set with no ranking cold_decoding.py commongen.sh. The paper says the other lexical configurations match the abductive ones app. B, which would imply 16 samples and a ranking. Which candidate Table 3 evaluates is not stated.
19. Post-processing. The paper says the base LM completes a sample with greedy decoding §3.4. The code completes greedily up to a maximum length (40 or 50 tokens) and then truncates to the first sentences with at least length words util.py. The length rule and the maximum are not in the paper.
20. Data and scope. The paper evaluates on test-set data §4. The run scripts read the dev or small files shipped in the repo (small_data.json with 2 lines, dev_data.json, commongen.dev.jsonl) and process examples 0 to 5, i.e. six inputs abductive.sh counterfactual.sh commongen.sh cold_decoding.py.
21. Evaluation not released. The metrics (BLEU, ROUGE, CIDEr, BERTScore, Min-Edit, PPL), the crowdsourcing protocol and the human-evaluation scripts are not in the repository README.md.

Pitfalls

• Adam, not SGD. An implementation that follows Eq. 2 literally will not match the paper's behaviour. Decide whether to follow the code's Adam update or the paper's equation, and say which.
• \tau and the straight-through estimator. The forward pass uses very sharp softmax values, while the gradient passes straight through (see items 10 and cold_decoding.py). Without a numeric \tau, reproduction needs these values from the code.
• Loss scale depends on length. Averaging versus summing (item 8) changes the balance of constraints with the fixed weights, and the sum of weights is not 1 in the code (item 5). Changing the sequence length can silently re-balance the energy.
• Candidate set and discretisation. The top-k candidate set is computed from a soft context. Its mask is taken from the previous iteration cold_decoding.py, so it lags the current soft sequence. The final text is rebuilt in a separate pass over the discrete context util.py.
• Stochastic n-gram loss. The n-gram term in the code uses Gumbel-softmax sampling over the n-gram terms when the n-gram is shorter than the output bleuloss.py. Outputs therefore depend on the random seed and the GPU random state, and the paper does not mention this. The loss also clamps log-probabilities at about −20 before the n-gram match bleuloss.py.
• Hard-coded CUDA. The BLEU loss and the noise sampling call .cuda() or device='cuda' bleuloss.py bleuloss.py cold_decoding.py. The --no-cuda flag does not cover these, so a CPU-only run will fail there.
• Costs the abstract hides. Each sample runs 2000 forward and backward passes through GPT2-XL with a full vocabulary-sized soft sequence, and the paper uses sample-and-select with 16 or 32 samples app. B. The runtime table tab. 5 is per sample, so the total cost of the selection scales with the sample count.
• Figures and tables counted from source. The table and figure numbers in this card were counted from the LaTeX source order. Check them against the PDF before citing a number in another document.
• Unstated values. The paper does not state: the sweep grid for the weights, the value of \tau, the number of samples for lexical generation, the selection rule for lexical and abductive output, the maximum length, or the weights of the two n-gram orders in counterfactual. Use the values in the Hyperparameters table, which come from the paper or the scripts.

Lineage

• Builds on: Langevin dynamics for sampling from energy-based models §2: Welling and Teh (2011). The reverse and backward language model comes from West et al. (2021), Reflective Decoding app. B.
• Compared against: DeLorean, the backprop-based decoding method of Qin et al. (2020) §4.1 §4.2; NeuroLogic, the beam-search method of Lu et al. (2020) §4.3; TSMH, the Monte Carlo tree search method of Zhang et al. (2020) §4.3; Mix-and-Match, the energy-based decoding of Mireshghallah et al. (2022), arXiv:2203.13299 §4.2 tab. 5.
• Uses a differentiable n-gram match from Liu et al. (2021), arXiv:2106.15078 §3.2.
• Superseded by: not stated in this card's sources.