Tractable Control for Autoregressive Language Generation

2023/08/22

tags: RL Group meeting

Outline

• Introduction
• Related work
• Guiding Autoregressive Generation with Tractable Probabilistic Models
• Efficient Probabilistic Reasoning with Hidden Markov Models
• Experiments
• Conclusion

Introduction

• Autoregressive large language models remain a major challenge to generate text that satisfies complex constraints:

  • Sampling from the conditional distribution
    $\mathrm{Pr}(\mathrm{text}|\alpha )$ is intractable for even the simplest lexical constraints
    $\alpha$.

• We propose to use tractable probabilistic models (TPMs) to impose lexical constraints in autoregressive text generation models, which we refer to as GeLaTo (Generating Language with Tractable Constraints).

• Our goal is to generate text effectively following the conditional distribution

  ${\mathrm{Pr}}_{\mathrm{LM}}(x_{1:n}|\alpha )$ for arbitrary lexical constraints α.
  • TPMs can efficiently compute the joint probability distribution over the input sequence and the constraints, which allows for more precise control over the generation process.
  • Pre-trained LMs only model the next token distribution given some prefix, and conditioning on constraints can be intractable even for simple constraints.

• We use distilled hidden Markov models

  1. We can efficiently compute
     $\mathrm{Pr}(\mathrm{text}|\alpha )$, to guide autoregressive generation from GPT2.
  2. We propose a dynamic programming algorithm that efficiently computes conditional probabilities
     ${\mathrm{Pr}}_{\mathrm{HMM}}(·|\alpha )$

• Our study demonstrates the potential of TPMs in controlling large language models and motivates the development of more expressive TPMs.

1. We train a TPM
   ${\mathrm{Pr}}_{\mathrm{TPM}}$ via maximum likelihood estimation (MLE) on samples drawn from
   ${\mathrm{Pr}}_{\mathrm{LM}}$, which is equivalent to minimizing the KL-divergence between
   ${\mathrm{Pr}}_{\mathrm{TPM}}$ and
   ${\mathrm{Pr}}_{\mathrm{LM}}$;
2. At generation time, we compute
   ${\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}|x_{1:t},\alpha )$ efficiently and combine it with
   ${\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}|x_{1:t})$ to approximate
   ${\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}|x_{1:t},\alpha )$ for reliable control.

Related work

Tractable probabilistic models

• A class of queries
  $\mathbf{Q}$ is tractable on a family of probabilistic models
  $\mathcal{M}$ iff any query
  $q\in \mathbf{Q}$ on a model
  $m\in \mathcal{M}$ can be computed in time
  $\mathcal{O}($ poly
  $(|m|))$.
• We also say that
  $\mathcal{M}$ is a tractable model for
  $\mathbf{Q}$.
• Tractable probabilistic models support efficient probabilistic inference.
• Probabilistic circuits (PCs) is a unified framework for a large family of tractable probabilistic models:
  • hidden Markov models
  • bounded tree-width graphical models
  • sum-product networks (SPNs)

Controllable Autoregressive Language Generation

• One line of research on constrained text generation focuses on modifying the decoding algorithm to inject constraints into the beam search process
  • Search-based
    • constrained beam search
    • NeuroLogic Decoding
    • A*esque NeuroLogic Decoding
  • Token-level
    • NADO
    • FUDGE
  • Insertion-based

Guiding Autoregressive Generation with Tractable Probabilistic Models

• Our goal is to generate from the following conditional distribution:
  
  $${\mathrm{Pr}}_{\mathrm{LM}}(x_{1:n}\mid \alpha )=\prod _t{\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}\mid x_{1:t},\alpha )$$
  • ${\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}|x_{1:t},\alpha )$ is intractable
  • We can assume that
    ${\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}|x_{1:t},\alpha )$ can be efficiently computed.
• We train the TPM model via MLE:
  
  $${\mathbb{E}}_{x_{1:n}\sim {\mathrm{Pr}}_{\mathrm{LM}}}\log {\mathrm{Pr}}_{\mathrm{TPM}}\left(x_{1:n}\right)$$
• Which effectively minimizes their KL-divergence:
  
  $$\begin{array}{rl}& D_{\mathrm{KL}}({\mathrm{Pr}}_{\mathrm{LM}}\Vert {\mathrm{Pr}}_{\mathrm{TPM}})\\ & ={\mathbb{E}}_{x_{1:n}\sim {\mathrm{Pr}}_{\mathrm{LM}}}\log {\mathrm{Pr}}_{\mathrm{LM}}\left(x_{1:n}\right)-{\mathbb{E}}_{x_{1:n}\sim {\mathrm{Pr}}_{\mathrm{LM}}}\log {\mathrm{Pr}}_{\mathrm{TPM}}\left(x_{1:n}\right)\end{array}$$
• We assume that there exists some “quality” constraint
  $\beta$ such that
  ${\mathrm{Pr}}_{\mathrm{TPM}}(|\beta )$ is even closer to
  ${\mathrm{Pr}}_{\mathrm{LM}}$.
  
  $${\mathrm{Pr}}_{\mathrm{TPM}}(x_{1:n}\mid \alpha ,\beta )=\prod _t{\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}\mid x_{1:t},\alpha ,\beta )$$
• We assume the key independence assumption:
  
  $$\begin{array}{rl}& {\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}\mid x_{1:t},\alpha ,\beta )\\ & \propto {\mathrm{Pr}}_{\mathrm{TPM}}(\alpha \mid x_{1:t+1},\beta )\cdot {\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}\mid x_{1:t},\beta )\\ & \propto {\mathrm{Pr}}_{\mathrm{TPM}}(\alpha \mid x_{1:t+1})\cdot {\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}\mid x_{1:t}).\end{array}$$

  • Unsupervised setting

    • Assume that the base pre-trained LM is not fine-tuned given task-specific supervision.
    • It may still be adapted to generate text in a specific domain or context.
      
      $$p(x_{t+1}\mid x_{1:t},\alpha )\propto {\mathrm{Pr}}_{\mathrm{TPM}}(\alpha \mid x_{1:t+1})\cdot {\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}\mid x_{1:t}).$$

  • Supervised setting

    • Assume that
      ${\mathrm{Pr}}_{\mathrm{LM}}$ is fine-tuned in a sequence-tosequence manner.
    • We adopt an alternative formulation by viewing
      ${\mathrm{Pr}}_{\mathrm{TPM}}(x_{t+1}|x_{1:t},\alpha )$ and
      ${\mathrm{Pr}}_{\mathrm{LM}}(x_{t+1}|x_{1:t})$ as classifiers trained for the same task yet with different biases.
      
      $$\begin{array}{rl}& p(x_{t+1}\mid x_{1:t},\alpha )\\ & \propto {\mathrm{Pr}}_{\mathrm{TPM}}{(x_{t+1}\mid x_{1:t},\alpha )}^w\cdot {\mathrm{Pr}}_{\mathrm{LM}}{(x_{t+1}\mid x_{1:t})}^{1-w}\end{array}$$
• To summarize, GeLaTo consists of two major steps:
  • Distillation - We train a TPM on samples drawn from the pretrained LM via MLE to effectively minimize the KL divergence between
    ${\mathrm{Pr}}_{\mathrm{LM}}$ and
    ${\mathrm{Pr}}_{\mathrm{TPM}}$.
  • Probabilistic reasoning: for each step of autoregressive generation, we compute
    ${\mathrm{Pr}}_{\mathrm{TPM}}(·|\alpha )$ and generate from the conditional next-token distribution
    $p(x_{t+1}|x_{1:t},\alpha )$ defined above.
• Two advantages:
  • The sentences generated following
    $p(x_{t+1}|x_{1:t},\alpha )$ are guaranteed to satisfy the lexical constraint α.
  • The TPM training is independent of the lexical constraint α, which is only enforced at inference time.
    • No need to re-train the TPM model no matter how α changes.

Efficient Probabilistic Reasoning with Hidden Markov Models(HMMs)

• We need to compute

  ${\mathrm{Pr}}_{\mathrm{TPM}}(x_{1:t},\alpha )$:
  • unsupervised setting:
    $\mathrm{Pr}(\alpha |x_{1:t+1})$ =
    $\mathrm{Pr}(x_{1:t+1},\alpha )/\mathrm{Pr}(x_{1:t+1})$
  • supervised setting:
    $\mathrm{Pr}(x_{t+1}|x_{1:t},\alpha )\propto \mathrm{Pr}(x_{1:t+1},\alpha )$

• We describe a dynamic programming algorithm that computes

  $\mathrm{Pr}(x_{1:t},\alpha )$ for HMMs, where α is some lexical constraint encoded in a conjunctive normal form (CNF):
  
  $$(I\left(w_{1,1}\right)\vee \cdots \vee I\left(w_{1,d_1}\right))\wedge \cdots \wedge (I\left(w_{m,1}\right)\vee \cdots \vee I\left(w_{m,d_m}\right))$$
  • $w_{i,j}$ is a string of tokens.
  • $I(w_{ij})$ is the indicator variable that represents whether
    $w_{ij}$ appears in the generated text.

Hidden Markov Models

• The joint probability

  $\mathrm{Pr}(x_{1:n},z_{1:n})$ is defined as:
  
  $$\mathrm{Pr}(x_1\mid z_1)\mathrm{Pr}\left(z_1\right)\prod _{2\le t\le n}\mathrm{Pr}(x_t\mid z_t)\mathrm{Pr}(z_t\mid z_{t-1})$$

• The parameters of HMM are given by the initial probability

  $\mathrm{Pr}(z_1)$, emission matrix
  $\mathrm{Pr}(x_t|z_t)$ and the transition matrix
  $\mathrm{Pr}(z_{t+1}|z_t)$, which stay the same across different positions t.
  
  $$\mathrm{Pr}(x_{t:n}\mid z_t,x_{1:t-1})=\mathrm{Pr}(x_{t:n}\mid z_t).$$

• forward algorithm:
  

  $$\begin{array}{rlr}& \mathrm{Pr}(x_{1:t},z_t)& =\sum _{1\le z_{t-1}\le h}\mathrm{Pr}(x_t\mid z_t)\mathrm{Pr}(z_t\mid z_{t-1})\mathrm{Pr}(x_{t-1},z_{t-1})\end{array}$$

• ${\mathrm{Pr}}_{\mathrm{HMM}}(x_{1:n})$ effectively defines a distribution over all texts with length ≤ n.

An Efficient Dynamic Programming Algorithm

• $\alpha \prime$ is some CNF formula obtained by removing from the original
  $\alpha$ the clauses that are already satisfied.
• $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$.
• $$S(x,\alpha ):=\{s:\mathrm{\exists }{x}^{\prime }\text{ a suffix of }x\text{ s.t. }{x}^{\prime }\oplus s\text{ lies in }\alpha \}$$
• Case 1.
  $x_{l:r}$
  $\ne \varnothing$; then
  
  $$\begin{array}{rl}& \mathrm{Pr}(x_{l:r},{\alpha }_{l:n}\mid z_l)\\ & =\sum _{z_{r+1}}\underset{―}{\mathrm{Pr}(x_{l:r},z_{r+1}\mid z_l)}(\mathrm{Pr}({\alpha }_{r+1:n}\mid z_{r+1})\\ & +\sum _{s\in S(x_{l:r},\alpha )}\mathrm{Pr}(s_{r+1:r+|s|},{(\alpha \mathrm{\setminus }{x}_{l:r}\oplus s)}_{r+1:n}\mid z_{r+1})\\ & -\sum _{s\in S(x_{l:r},\alpha )}\mathrm{Pr}(s_{r+1:r+|s|},{\alpha }_{r+1:n}\mid z_{r+1}));\end{array}$$
• Case 2.
  $x_{l:r}$
  $=\varnothing$; we reduce the problem to Case 1 by enumerating
  $x_l$ over the vocabulary:
  
  $$\mathrm{Pr}({\alpha }_{l:n}\mid z_l)=\sum _{x_l\in \text{ vocabulary }}\mathrm{Pr}(x_l,{\alpha }_{l:n}\mid z_l)$$
• At step t by computing
  $\mathrm{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:
  
  $$\mathrm{Pr}(x_{1:t},{\alpha }_{1:n})=\sum _{z_1}\mathrm{Pr}\left(z_1\right)\mathrm{Pr}(x_{1:t},{\alpha }_{1:n}\mid z_1)$$
• the time complexity of GeLaTo is O(2|α|nm)
  • |α| is the number of clauses in α
  • n is the maximum sequence length
  • m is the number of different suffixes for all keystrings in α.

Experiments

1. Fine-tuning GPT2-large

   • domain adaptation
   • sequence-to-sequence

2. Training HMMs

   • To enforce lexical constraint in autoregressive. generation

3. Constraint Formulation
   

   $$\begin{array}{rl}& [I(\text{ catch })\vee I(\text{ caught })\vee \dots ]\\ \wedge & [I(\text{ fr }\oplus \text{ is }\oplus \text{ bee })\vee I(\text{ fr }\oplus \text{ is }\oplus \text{ bees })\vee \dots ]\\ \wedge & [I(\text{ snow })\vee I(\text{ snow }\oplus \text{ ing })\vee I(\text{ snow }\oplus \text{ ed })\vee \dots ]\end{array}$$

4. Decoding

   • We adopt beam search to greedily search for
     $x_{1:n}$ that maximizes
     $p(x_{1:n}|\alpha )$.

5. Metrics

   • ROUGE
   • BLEU
   • CIDEr
   • SPICE

Conclusion

• We propose GeLaTo, where we use tractable probabilistic models (TPMs) to impose complex lexical constraints (denoted α) in autoregressive language generation from large language models.
• With hidden Markov model as a running example:
  • We present an efficient dynamic programming algorithm for conditioning HMMs on complex lexical constraints.
  • We demonstrate the effectiveness of GeLaTo on various constrained generation benchmarks.

Appendix

Autoregressive model

• An autoregressive language model is a type of Machine Learning model that uses autoregressive techniques to predict the next word in a sequence of words based on the words that have come before it.
  • $$y(t)=c+w\mathrm{\_}1y(t-1)+w\mathrm{\_}2y(t-2)+\dots +w\mathrm{\_}py(t-p)+e(t)$$

HMMs

• A Hidden Markov Model (HMM) is a statistical model used to describe a sequence of observable events or symbols in terms of an underlying sequence of hidden states.
• Given a sequence of observations, the goal of HMMs is to find the most likely sequence of hidden states that generated those observations.