Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Syntactic and Semantic Control of Large Language Models via Sequential Monte Carlo

João Loula, Benjamin LeBrun, Li Du, Ben Lipkin, Clemente Pasti, Gabriel Grand, Tianyu Liu, Yahya Emara, Marjorie Freedman, Jason Eisner, Ryan Cotterell, Vikash Mansinghka, Alexander K. Lew, Tim Vieira, Timothy J. O'Donnell

TL;DR

The paper addresses constrained text generation by modeling a global posterior $g(\boldsymbol{x}) = \frac{1}{Z} p(\boldsymbol{x}) \Phi(\boldsymbol{x})$ and solving it with an adaptive sequential Monte Carlo framework that can incorporate cheap per-token constraints and expensive task-specific signals. By decomposing signals into efficient and expensive potentials and using incremental Extend/Reweight/Resample steps, the method adaptively allocates computation to promising partial sequences. Across four challenging domains, the Full SMC method with a relatively small LM matches or outperforms larger, fine-tuned models, with downstream performance tracking the quality of the posterior approximation via KL-based diagnostics. The work demonstrates practical gains in synthesis and parsing tasks and provides a programmable interface for applying SMC to controlled generation problems.

Abstract

A wide range of LM applications require generating text that conforms to syntactic or semantic constraints. Imposing such constraints can be naturally framed as probabilistic conditioning, but exact generation from the resulting distribution -- which can differ substantially from the LM's base distribution -- is generally intractable. In this work, we develop an architecture for controlled LM generation based on sequential Monte Carlo (SMC). Our SMC framework allows us to flexibly incorporate domain- and problem-specific constraints at inference time, and efficiently reallocate computational resources in light of new information during the course of generation. By comparing to a number of alternatives and ablations on four challenging domains -- Python code generation for data science, text-to-SQL, goal inference, and molecule synthesis -- we demonstrate that, with little overhead, our approach allows small open-source language models to outperform models over 8x larger, as well as closed-source, fine-tuned ones. In support of the probabilistic perspective, we show that these performance improvements are driven by better approximation to the posterior distribution. Our system builds on the framework of Lew et al. (2023) and integrates with its language model probabilistic programming language, giving users a simple, programmable way to apply SMC to a broad variety of controlled generation problems.

Syntactic and Semantic Control of Large Language Models via Sequential Monte Carlo

TL;DR

The paper addresses constrained text generation by modeling a global posterior and solving it with an adaptive sequential Monte Carlo framework that can incorporate cheap per-token constraints and expensive task-specific signals. By decomposing signals into efficient and expensive potentials and using incremental Extend/Reweight/Resample steps, the method adaptively allocates computation to promising partial sequences. Across four challenging domains, the Full SMC method with a relatively small LM matches or outperforms larger, fine-tuned models, with downstream performance tracking the quality of the posterior approximation via KL-based diagnostics. The work demonstrates practical gains in synthesis and parsing tasks and provides a programmable interface for applying SMC to controlled generation problems.

Abstract

A wide range of LM applications require generating text that conforms to syntactic or semantic constraints. Imposing such constraints can be naturally framed as probabilistic conditioning, but exact generation from the resulting distribution -- which can differ substantially from the LM's base distribution -- is generally intractable. In this work, we develop an architecture for controlled LM generation based on sequential Monte Carlo (SMC). Our SMC framework allows us to flexibly incorporate domain- and problem-specific constraints at inference time, and efficiently reallocate computational resources in light of new information during the course of generation. By comparing to a number of alternatives and ablations on four challenging domains -- Python code generation for data science, text-to-SQL, goal inference, and molecule synthesis -- we demonstrate that, with little overhead, our approach allows small open-source language models to outperform models over 8x larger, as well as closed-source, fine-tuned ones. In support of the probabilistic perspective, we show that these performance improvements are driven by better approximation to the posterior distribution. Our system builds on the framework of Lew et al. (2023) and integrates with its language model probabilistic programming language, giving users a simple, programmable way to apply SMC to a broad variety of controlled generation problems.

Paper Structure

This paper contains 47 sections, 2 theorems, 19 equations, 4 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Monte Carlo Inference for Constrained Semantic Parsing
  3. Notation.
  4. Language models.
  5. Potential functions.
  6. Target distribution.
  7. Locally constrained decoding.
  8. Importance sampling.
  9. Sequential Monte Carlo.
  10. Further extensions.
  11. Experiments
  12. Domains
  13. Evaluation of Downstream Performance
  14. Grammar constraints.
  15. Expensive potentials.
  16. ...and 32 more sections

Key Result

Proposition 1

The set-based proposal speedup's distribution $q$ (def:set-based-proposal-speedup) is a properly weighted proposal for ${\color{SET}\widetilde{\ell}}$.

Figures (4)

  • Figure 1: Controlled generation from LMs via sequential Monte Carlo.Left: We use sequential Monte Carlo to sample from high-quality approximations to posteriors over LM outputs. Partial sequences are repeatedly extended via grammar-constrained generation. We then apply weight corrections to mitigate the greediness of locally constrained decoding, as well as expensive potentials to encode rich information that cannot be included in logit masks. Finally, resampling focuses computation on promising particles. Right: Accuracy gains from these innovations on challenging data science, text-to-SQL, goal inference, and molecule synthesis benchmarks.
  • Figure 2: Left: Performance on the Data Science task (DS-1000) for different models and methods. Codex-002 performance as reported in lai2023ds. Right: Performance across all tasks for Full IS and Full SMC with 5, 10, and 50 particles. Error bars: bootstrapped 95% confidence intervals.
  • Figure 3: Estimated KL between the algorithm and the global product of experts for a representative problem instance in each domain. Values closer to 0 indicate that the algorithm is better at approximating $g$. Significant differences are indicated with ** for $p < 0.01$ and *** for $p < 0.001$ (t-test). Algorithms use $N = 10$ particles.
  • Figure 4: Distributional properties of compounds generated by different methods. Middle: Distribution of drug-likeness as measured by QED score bickerton2012quantifying. Right: Means for other properties of interest such as diversity and de novo similarity (details in \ref{['app:gbd_17']}).

Theorems & Definitions (8)

  • Definition 1
  • Definition 2
  • Proposition 1
  • proof
  • Definition 3
  • Definition 4
  • Proposition 2: Estimating $\log Z_{\sigma} - \mathrm{KL}(q_\text{IS}^* \mathrel\| \sigma)$
  • proof