Tractable Control for Autoregressive Language Generation

A hidden Markov model (HMM) represents a joint probability distribution over $n$ observed variables $X_{1:n}$ and $n$ latent variables $Z_{1:n}$. Specifically, for language modeling, $X_{t}$ represents the token at position $t$ and $Z_{t}$ represents the corresponding latent state; $Z_{t}$ takes values in $\{1,2,\dots,h\}$, where $h$ is the number of latent states. Given observed token sequence $x_{1:n}$ and latent state sequence $z_{1:n}$, the joint probability $\Pr(x_{1:n},z_{1:n})$ is defined as:

in particular, the parameters of HMM are given by the initial probability $\Pr(z_{1})$, emission matrix $\Pr(x_{t}\>|\>z_{t})$ and the transition matrix $\Pr(z_{t+1}\>|\>z_{t})$, which stay the same across different positions $t$. HMMs can also be represented as Bayesian networks (Pearl, 1985); see Figure 2 for an example. To perform probabilistic inference on HMMs efficiently, we leverage the following Markov property:

For example, we can compute the probability of any prefix $\Pr(x_{1:t})=\sum_{z_{t}}\Pr(x_{1:t},z_{t})$, which can be efficiently computed by the following recurrence relation, which is referred to as the forward algorithm (Rabiner & Juang, 1986):

Modeling Variable-length Texts with HMMs. HMMs model distributions over a fixed number of random variables $X_{1:n}$. To model texts with variable lengths, we first determine a maximum sequence length $n$ and pad training texts of length $<n$ with the special EOS (“endoftext”) token to the maximum length. We also construct our HMM in a special way such that an EOS token can only be followed by EOS tokens; that is, sequences that do not satisify this constraint have $0$ probability. Hence, ${\Pr}_{\text{HMM}}(x_{1:n})$ effectively defines a distribution over all texts with length $\leq n$.

We first illustrate the dynamic programming algorithm with a toy example. As shown in Figure 2, assume that we have generated the first $t-1$ tokens “Kids … like” and we are given the constraint:

here we assume “like”, “working”, “eating” and “soccer” are single tokens and $\oplus$ denotes string concatenation. To compute $\Pr(x_{1:t-1},\alpha_{1:n})$, we marginalize out $x_{t}$ in $\Pr(x_{1:t-1},x_{t},\alpha_{1:n})$; in particular, we sum over three possible cases for the next token $x_{t}$: “eating”, “working” or neither, and for each case we reduce $\Pr(x_{1:t-1},x_{t},\alpha_{1:n})$ to $\Pr(x_{1:t-1},x_{t},\alpha^{\prime}_{t:n})$ for some $\alpha^{\prime}$; here $\alpha^{\prime}$ is some CNF formula obtained by removing from the original $\alpha$ the clauses that are already satisfied. Then, we leverage the Markov property of HMMs (see Equation 3) to break down the joint probability $\Pr(x_{1:t-1},x_{t},\alpha^{\prime}_{t:n})$ into sub-problems.

Before we describe how to compute $\Pr(x_{1:t},\alpha_{1:n})$, we establish a recurrence relation for computing terms of the form $\Pr(x_{l:r},\psi_{l:n}\>|\>z_{l})$, where $x_{l:r}$ is either the empty string or a suffix for some keystring in $\alpha$, $\psi$ is a CNF consisting of a subset of clauses in $\alpha$ and $z_{l}$ is a latent state for $Z_{l}$.

Assumptions & Notations. For simplicity, we make the following non-overlapping assumption: for the set of keystrings appearing in $\alpha$, denoted as $\{w_{ij}\}$, the prefix of $w_{ij}$ cannot be a suffix for $w_{pq}$ for all $ij\neq pq$. We also define the following set of strings:

which contains all strings that can be appended to $x$ to form some keystrings in $\alpha$. For the example in Figure 2, for $x=\text{``Kids ... like''}$, $S(x,\alpha)$ is given by $\{\text{``eating''},\text{``working''}\}$. We write $s_{i:j}$ as a shorthand for $X_{i:j}=s$.

Recurrence Relation. $\Pr(x_{l:r},\psi_{l:n}\>|\>z_{l})$ follows the following recurrence relation:

Case 1. $x_{l:r}\neq\emptyset$; then,

here $\oplus$ denotes string concatenation and $\alpha\setminus x_{l:r}\oplus s$ represents the CNF obtained by removing the clauses with any keywords appearing in $x_{l:r}\oplus s$.

Case 2. $x_{l:r}=\emptyset$; we reduce the problem to Case 1 by enumerating $x_{l}$ over the vocabulary:

The recurrence relation presented above gives us a dynamic programming algorithm for computing terms of the form $\Pr(x_{l:r},\psi_{l:n}\>|\>z_{l})$; see appendix for derivations. Note that the boxed terms are the sub-problems and the underlined terms are either HMM parameters or can be pre-computed via the forward algorithm and then cached for later use.

Finally, as discussed at the beginning of this section, we guide autoregressive generation from language models at step $t$ by computing $\Pr(x_{1:t-1},x_{t},\alpha_{1:n})$, where $x_{1:t-1}$ denotes the first $t-1$ tokens that have been generated:

here $\Pr(z_{1})$ is the initial probability of the HMM and $\Pr(x_{1:t},\alpha_{1:n}\>|\>z_{1})$ can be computed by the formula in Case 1 (setting $l\!=\!1$), given that all boxed terms are pre-computed by the dynamic programming algorithm.

As an example, Algorithm 1 summarizes how to perform constrained generation with GeLaTo by sampling from the autoregressive distribution $p(x_{t}\>|\>x_{1:t-1},\alpha)$, as defined in Section 2 (unsupervised setting). We can easily adapt Algorithm 1 for other decoding procedures like beam search.

For a rough analysis of the time complexity of Algorithm 1, we treat both the number of latent states $h$ and the vocabulary size as constants; in practice, we can avoid enumerating all tokens in the vocabulary and all latent states of HMM via GPU parallelization (see appendix & code for details). It follows that the time complexity of GeLaTo is $O(2^{|\alpha|}nm)$, where $|\alpha|$ is the number of clauses in $\alpha$, $n$ is the maximum sequence length and $m$ is the number of different suffixes for all keystrings in $\alpha$. We show that GeLaTo scales well in practice in Section 4.3.

CommonGen (Lin et al., 2020) is a benchmark for constrained generation with lexical constraints: the input of each example consists of three to five concepts (keywords) and the goal is to generate a natural sentence using all concepts; in particular, the given keywords can appear in any order or in any form of inflections in the generated sentences. For example, given “car snow drive” as concepts, both “a man drives a car on a snow covered road” and “the car drove through the snow” are considered acceptable. We also evaluate GeLaTo on the Yelp!Review (Cho et al., 2019) and the News (Zhang et al., 2020) datasets. Compared to CommonGen, both Yelp!Review and News share similar formats, except that they require all keywords to be generated in the forms as given (i.e. no inflections allowed) and to follow specific orders.

We compare GeLaTo against constrained generation approaches belonging to different families:

InsNet (Lu et al., 2022a) is a class of insertion-based language models (Susanto et al., 2020) that generate text by repeatedly inserting new tokens into the sequence. InsNet guarantees that the keywords appear in the generated sentence by initializing the token sequence as the keywords, arranged in some order.

NeuroLogic (A*esque) Decoding (Lu et al., 2021, 2022b) are search-based decoding algorithms; they are inference-time algorithms like beam search and do not use any auxiliary models. Leveraging look-ahead heuristics, NeuroLogic A*esque decoding not only optimizes the probability of the generated sentence but also steers the generation towards satisfying the lexical constraints.

NADO (Meng et al., 2022) trains an auxiliary neural model approximating the conditional distribution $\Pr(\alpha|x_{1:t},x_{t+1})$ to guide constrained generation of the base model. As mentioned in Section 2, NADO needs to re-train the auxiliary model for different types of $\alpha$ (e.g., ten keywords) while GeLaTo does not need re-training.

Following the experiment setup of Lu et al. (2021) and Meng et al. (2022), we evaluate GeLaTo under both unsupervised and supervised settings, as described in Section 2.

Main evaluation results are presented in Table 1. GeLaTo outperforms all baselines in both unsupervised and supervised settings by a large margin, achieving not only significantly higher BLEU and ROUGE scores but also 100% constraint satisfaction. The unsupervised setting is more challenging given that the base model is never trained with task-specific supervision; despite this, GeLaTo achieves 30.3 BLEU score in the unsupervised setting, while NADO (the best performing baseline) obtains 30.8 BLEU score in the supervised setting. To provide more insight into GeLaTo, we also conduct the following ablation studies.

Generation Quality vs. Approximation Performance. As discussed in Section 2, GeLaTo assumes that distilled HMMs are good enough approximations for base models; our hypothesis is that the better the HMM approximates the base model, the better the generation quality. With GeLaTo, we generate from different HMM checkpoints from the distillation procedure, and report the average log-likelihoods and BLEU scores (without re-ranking the beams). As shown in Figure 3, as the training proceeds, both log-likelihood and BLEU score improves, exhibiting a clear positive correlation. This finding motivates the development of better tractable probabilistic models for language modeling.

Robustness of Hyperparameter $\boldsymbol{w}$. As described in Section 2, for the supervised setting, the formulation of GeLaTo involves a hyperparameter $0\!\leq\!w\!\leq\!1$ that decides how much the TPM or the base model contributes to generation. For our experiments, $w$ is set to 0.3 based on cross-validation results on the training set. Figure 4 shows the BLEU score (after re-ranking) on the validation set of CommonGen given different values of $w$. The performance of GeLaTo is very robust with respect to different choices of $w$, achieving SoTA BLEU scores for $0.1\!\leq\!w\!\leq\!0.8$.

Effect of Beam Size. GeLaTo uses beam search for generation and we study how its performance is affected by the choice of beam size. Figure 5 shows that for both unsupervised and supervised settings, the performance of GeLaTo improves monotonically as the beam size increases.

Run-time Comparison. We conduct an empirical evaluation of the run-time (in seconds) of GeLaTo on CommonGen, in comparison to NeuroLogic A*esque and vanilla GPT2-large; all methods are evaluated on a single NVIDIA A100 GPU with 40 GB memory; the run-time is measured on 100 randomly sampled examples for each # of concepts.

GeLaTo achieves its best performance with beam-size=128; yet we also report the run-time for beam-size=16, where it achieves performance better than all baselines. For the unsupervised setting, GeLaTo is much faster than NeuroLogic A*esque, which suffers from an unconstrained search space. For the supervised setting, GeLaTo is slower than A*esque but the run-time for beam-size = 16 is still comparable.

Human Evaluation. We conduct human evaluation for sentences generated on CommonGen (dev), following the setup of prior works (Lu et al., 2022b; Meng et al., 2022). Specifically, we mix sentences generated by different methods, and each sentence is presented to one human annotator to be evaluated on four aspects: concepts, plausibility, quality, and overall rating. The results are shown in Table 3. To test statistical significance, we conduct the Wilcoxon signed rank two-sided test with $p$-value $<$ 0.05 and GeLaTo performs best in all metrics compared to prior SoTA. We refer readers to Appendix B for details of human evaluation.

Fixing Order of Keywords. Following prior works (Meng et al., 2022; Lu et al., 2022a), we evaluate GeLaTo on Yelp!Review and News datasets. They are more challenging in that they require keywords to appear in specific orders; besides, the average sequence lengths for both datasets are approximately 64 tokens, twice of that of CommonGen. With a minor modification to Algorithm 1, GeLaTo is easily adapted to generate text with ordered keywords. For both datasets, the training examples do not provide keywords thus there is no immediate way to finetune the base models in a supervised way. Yet, as shown in Table 4, the unsupervised GeLaTo alone achieves SoTA BLEU scores.

Tractable probabilistic models support efficient probabilistic inference (e.g., marginal probability), thus they have been widely used in inference-demanding tasks, including enforcing algorithmic fairness (Choi et al., 2020a, 2021), and making predictions under missing data (Khosravi et al., 2019; Correia et al., 2020; Li et al., 2021; Dang et al., 2022b).

Probabilistic circuits (PCs) is a unified framework for a large family of tractable probabilistic models including hidden Markov models (Rabiner & Juang, 1986), bounded tree-width graphical models (Meila & Jordan, 2000) and sum-product networks (SPNs) (Poon & Domingos, 2011). Recent progress in learning probabilistic circuits for generative modeling (Dang et al., 2022c; Liu et al., 2023) and their efficient implementation (Molina et al., 2019; Peharz et al., 2020; Dang et al., 2021) have been pushing the limits of PC’s expressive power.

The capacity of deep generative models is continuously increasing, while their probabilistic and logic querying ability is restricted. A variety of methods have been developed. Boyd et al. (2022) introduce a general inference typology on autoregressive sequence models that can develop query estimation methods based on beam search and importance sampling. Ahmed et al. (2022) use PCs as a replacement for the SoftMax layer in neural networks such that their outputs are guaranteed to satisfy the constraint.

One line of research on constrained text generation focuses on modifying the decoding algorithm to inject constraints into the beam search process, such as constrained beam search (Post & Vilar, 2018), NeuroLogic Decoding (Lu et al., 2021) and A*esque NeuroLogic Decoding (Lu et al., 2022b). Though they can be easily applied to various language models without training, these search-based methods can be inefficient as they suffer from large search spaces. Recent works like NADO (Meng et al., 2022) and FUDGE (Yang & Klein, 2021) train auxiliary neural models to provide token-level guidance for autoregressive generation. Another family of approaches that enforce keyword-type constraints are insertion-based language models (Lu et al., 2022a; Susanto et al., 2020), where the initial sequences only consist of the desired keywords and the transition phrases are repeatedly inserted to complete the sentences.

This work was funded in part by the DARPA Perceptually-enabled Task Guidance (PTG) Program under contract number HR00112220005, NSF grants #IIS-1943641, #IIS-1956441, #CCF-1837129, an SRA from Meta, a research gift from Amazon Alexa AI, and a gift from RelationalAI. GVdB discloses a financial interest in RelationalAI.