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

The paper concerns inference-time approaches for controlled generation with language models. Building on [Lew et al 2023], the authors develop a method based on sequential Monte-Carlo. The method proceeds segment-by-segment, first extending the segment (e.g. generating a next-token) subject to a token-level score (e.g., arising from a grammar constraint), then reweighting the candidates based on a partial sequence score (e.g., whether the code-so-far has a runtime error), then determining the next set of candidates by resampling candidates based on their weights.

The authors evaluate the method using prompted Llama 3 (base or instruct, varied by task) on 4 tasks. In each task, they construct task-specific token-level and partial-sequence level potentials, and provide an ablation of each of their three proposed components. The authors study the divergence between the target distribution and the distribution induced by running the algorithm, finding that each component leads to a lower KL divergence. Additionally, they provide a nice visualization of the sampling distributions and target distributions for their molecular generation task, and show that samples from their method improve along various dimensions.

In general the paper is written quite precisely, and several intuitive ideas (e.g., token masking or filtering out a sequence if its code doesn't run) are placed into a probabilistic framework.

The paper introduces a Sequential Monte Carlo (SMC) approach for constrained generation in language models to address semantic and syntactic constraints that traditional methods fail to handle effectively. The novel elements include integrating domain-specific constraints via potential functions, weight correction to mitigate the bias from locally greedy decoding, and adaptive resampling that focuses computational effort on promising sequences. The method was tested on four challenging tasks: Python code generation, text-to-SQL translation, goal inference, and molecular synthesis. The experiments compared the SMC-based approach to several baselines and showed that incorporating weight correction and semantic potentials significantly improved performance, while adaptive resampling further enhanced results by reallocating resources to better particles. The total SMC method outperformed the ablated SMC variants on the selected task.

The authors apply sequential Monte Carlo to tackle semantic parsing problems involving global constraints. They introduce an enhanced SMC approach, incorporating efficient stochastic approximations of full token-masking distributions and semantic potential to boost performance across five diverse datasets. Additionally, they estimate the KL divergence between each method’s output distribution and the global product-of-experts distribution, demonstrating that the latter is well-calibrated.

Controlling LLM generation to follow logical, syntactical or semantical (soft) constraints is a challenging task. The authors propose to unify controllable generation as sampling from the un-normalized distribution p_{LM}(x) \phi(x) where \phi is an energy function specifying the constraints. The authors propose to leverage sequential Monte Carlo to sample from the desired un-normalized conditional distribution. The authors conducted extensive evaluation of their approach on various downstream tasks with different combinations of constraints and have demonstrated significant improvement compared to the LLM baseline. One important ablation study has suggested the positive correlation between approximation accuracy and generation quality, motivating for further research on improving the sampling algorithm or the proposal distribution.

Our work is complementary to both Lew et al. 2023 and Zhao et al. 2024.

Lew et al. 2023 is a workshop paper that proposes, but does not systematically evaluate, the use of SMC to perform inference in a very broad class of probabilistic sequence models. We will revise our paper to better contextualize our work with respect to theirs; here, we summarize several key points:

• One contribution of Lew et al. 2023 is a without-replacement resampling algorithm that they speculate may be useful in the language modeling setting, making SMC more similar to beam search without compromising its unbiasedness. This resampling algorithm is completely compatible with our contributions: we can replace multinomial resampling in our SMC implementation with their without-replacement resampling. We conducted preliminary experiments testing this in the Text-to-SQL domain, and found that without-replacement resampling slightly hurt performance. See Appendix, Table 6.

• Lew et al. 2023 do not give a general recipe for defining intermediate target distributions or proposals; instead, they describe three example instantiations of their SMC framework (for infilling, prompt intersection, and natural language generation subject to word-length constraints). Our work specializes and extends sequential Monte Carlo for code generation under syntactic and semantic constraints, with concrete recipes for defining proposals and intermediate targets in terms of these constraints. Note that our proposal distributions and weight calculations are not quite a special case of Lew et al. 2023’s formalization, because our incremental weights are not deterministic functions of the generated string (see Appendix B). Note also that although Lew et al. 2023 use the word “potential” in their technical exposition, it means something different for them; their “potential” is our “incremental weight.” They have no direct analogue of the potentials from our work.

• Although all examples in Lew et al. 2023 deal with natural language generation and not code generation, the “word-length constraints” example can be straightforwardly adapted to use grammar constraints, given an efficient incremental parser. The resulting algorithm could be seen as an ablation of our method that removes the semantic potentials, but keeps the weight correction and the resampling steps [though it does not use our character-level proposal optimization in Appendix B]. This particular ablation was missing in our initial submission but we have now added it; see the portion of this response labeled “Comparison to other methods.”

• Finally, one contribution of Lew et al. 2023 was to propose language model probabilistic programs as a medium for specifying inference problems (as well as proposals and twist functions); given such a program, SMC can be automated. Conceptually, our method can be understood as a way to compile a set of user-provided potentials into such a probabilistic program. In order to practically implement our approach this way, the Python library of Lew et al. 2023 would have to be extended to support certain forms of CPU parallelism (which we use to efficiently run our incremental parser on multiple particles in parallel) and unbiased importance weight estimation (which we use to compute faster proposal distributions, see Appendix B).

Zhao et al. 2024 presents and thoroughly compares various methods for learning twist functions for natural-language generation tasks such as infilling. (To translate to our setting: if users define a set of potentials on full/complete strings, Zhao et al. 2024’s method can be understood as learning an extension of the potentials to cover partial strings.) Zhao et al. 2024 use these learned twist functions the same way that we use our efficient potentials $\phi_{eff}$: to inform both the proposal distributions and the intermediate targets in SMC.

However, learning twist functions can be very expensive. For each new constrained generation problem, twist functions need to be re-learned, a process that can take hours or days even for small models. (Note that twist learning is more like reinforcement learning than supervised fine-tuning, and may require many cycles of sampling roll-outs of the current SMC algorithm, then optimizing twists.) For this reason, all of Zhao et al. 2024’s experiments are run on either the 33M-parameter TinyStories model or, in one case, GPT-2 Medium (355M parameters). By contrast, our method requires no additional training, and was cheap enough that we could evaluate our method using 1B, 8B, and 70B-parameter variants of the Llama models. Indeed, a key insight of our approach is that although they may not be optimal twists, in structured generation domains we often have access to off-the-shelf tools that can serve as reasonable programmable twists, ranging from prefix parsers (for grammatical constraints) to more complex partial evaluators that incrementally capture more of the generated code’s semantics.

Comparison to other methods: We agree with the reviewers that the previous version of the paper did not make clear how our ablation studies related to other methods in the literature. To address this, we added a new baseline method (see table 1), and rewrote the introduction to section 3 and created new shorthands for the models to make the link more explicit. We summarize the above here:

1. LM + grammar constraints [Locally-constrained decoding] corresponds to the most common approach to grammar-constrained code generation with LMs (Shin et al., 2021; Scholak et al., 2021; Poesia et al., 2022; Willard & Louf, 2023; Moskal et al., 2024; Ugare et al., 2024). This approach applies the grammar constraint locally via token masking or token biasing at each step of generation.
2. LM + grammar constraints + semantic potentials [Sample-rerank] corresponds to a common approach for incorporating an external signal into an LM’s generations post-hoc, for instance a reward model (e.g. [1]), or a boolean encoding validity (e.g. [2]). This approach samples some number of full sequences from the locally-constrained model, then weights these using the semantic potential at the end.
3. LM + grammar constraints + weight correction + resampling: We add this new baseline, which we refer to as “Grammar-only SMC” as a shorthand, and note that it corresponds to a straightforward application of Lew et al. 2023 to incrementally correct locally-constrained decoding to the true Bayesian posterior via resampling; it is related to the method in Park et al. 2024, which also attempts to correct for the greediness of locally-constrained decoding. We find that, in general, the improvement of SMC without semantic potentials over Importance Sampling without semantic potentials is modest; the improvement is large only in the Goal Inference domain, where locally-constrained decoding seems to suffer the most from myopic sampling — however adding semantic potentials leads to large gains even here and still outperforms other methods, consistent with our previous findings.

[1] Nakano, Reiichiro, et al. "Webgpt: Browser-assisted question-answering with human feedback." arXiv preprint arXiv:2112.09332 (2021).

[2] Olausson, Theo X., et al. "LINC: A neurosymbolic approach for logical reasoning by combining language models with first-order logic provers." arXiv preprint arXiv:2310.15164 (2023).

Additional models: In addition to our original experiments on a model with 8 billion parameters (Llama 3.1), we ran additional experiments using both a larger 70 billion parameter LM (Llama 3) and a smaller 1 billion parameter LM (LLama 3.2). Due to computational resource constraints, we could only run the 70B experiment for one domain: we chose to do it for our most challenging domain, DS-1000. We find that:

• As expected, performance across the board is higher for larger LMs.
• The overall ranking of different approaches is preserved across the 1B, 8B, and 70B models: weight correction, resampling, and semantic potentials still generally improve performance, with the highest performing approach still being the combination of all three, that is SMC.
• Most interestingly, we find that, in general, the relative gains in accuracy provided by our method are more pronounced for smaller models: in fact, we find that in three out of our 4 domains (with the exception of text-to-SQL), our approach using the smaller model outperforms local decoding using a model 8 times larger, suggesting that our approach can dramatically improve the performance of smaller LMs. We also find, in DS-1000, that our approach allows an open source model (LLama 3 70B) to outperform a larger, closed-source, fine-tuned model (GPT-3 Codex 002, as reported by [1])

[1] Lai, Yuhang, et al. "DS-1000: A natural and reliable benchmark for data science code generation." International Conference on Machine Learning. PMLR, 2023.

We ran an extensive study of how the number of particles affect performance. As mentioned in our New main takeaways, our SMC method scales better with sample size than Importance Sampling, outperforming the latter even with 5 times (Goal Inference) or 10 times (text-to-SQL, Molecular Synthesis, Data Science) fewer particles. In general, we find that variations in performance across the number of particles depends on the domain. In both the Molecular Synthesis and the text-to-SQL domain, we do not find significant differences in performance across 5, 10 and 50 particles. For the Goal Inference and Data Science domains, on the other hand, we find that increasing the number of particles in general greatly improves performance across different ablations but especially for SMC. Interestingly, the performance gain from larger sample sizes occurs in the most challenging domains, whose semantic potentials run the complex procedures of executing a pddl plan or a python program.

We thank reviewers for their insightful comments: in response to common points brought up, we ran a new suite of experiments varying the base LM and the number of particles — effectively quadrupling the number of LM experiments in the paper. We have edited the manuscript, highlighting changes in blue: the comprehensive new experiments are in Appendix A—we plan to incorporate them into the main text, but wanted to share them soon for discussion. We underscore key takeaways below:

• Our method allows smaller LMs to outperform LMs that are over 8 times larger, and allows an open source model (LLama 3 70B) to outperform a larger, closed-source, fine-tuned model (GPT-3 Codex 002)
• Our SMC method scales better with sample size than Importance Sampling, outperforming the latter even with 5 times (Goal Inference) or 10 times (text-to-SQL, Molecular Synthesis, Data Science) fewer particles.

We post individual general responses in separate comments. Reviewer specific responses are as replies to specific reviews.

We thank the reviewers for their insightful comments. There was a great deal of agreement between reviewers on the points that most needed improvement: we addressed their feedback and effectively quadrupled the number of LM experiments in the paper. We believe this has significantly strengthened the paper and are pleased that the three review responses we’ve received raised their scores by 2, 2, and 3 in light of our revisions.

Below, we summarize the key improvements from the review period:

• Experiments with larger and smaller LMs (in response to points raised by reviewers MMNX, 8wBr, and 4hyJ): In addition to our experiments with an 8-billion parameter LM (Llama 3.1), we added experiments with a larger 70-billion-parameter LM (Llama 3) and a smaller 1-billion-parameter LM (Llama 3.2). We found that our method enables smaller LMs to outperform models over eight times larger and allows an open-source model (Llama 3 70B) to outperform a larger, closed-source, fine-tuned model (GPT-3 Codex 002).
• Particle experiments (in response to points raised by all reviewers): We conducted an extensive analysis of how performance scales with the number of particles. The results show that our SMC method increases in performance faster with more particles than basic importance sampling.
• Literature contextualization and comparison:(in response to points raised by reviewers gr6N, 8wBr, and 4hyJ): To better link our ablation studies with existing literature, we added a new baseline method, rewrote the model descriptions, and introduced clearer shorthand notations. We also clarified our relationship with two earlier papers using SMC in related ways: Lew et al., (2023) and Zhao et al., (2024).
• Evaluation of computational cost (in response to points raised by reviewer 4hyJ): We included a discussion and small experiment analyzing the computational cost of our approach.

This paper presents a novel inference-time method for controlled generation with large language models (LLMs), leveraging Sequential Monte Carlo (SMC) techniques. By extending the probabilistic framework introduced in Lew et al. (2023) and incorporating semantic potentials and resampling, the authors show how SMC-based methods can more faithfully approximate distributions over syntactically or semantically constrained text sequences. They apply their approach to four challenging tasks—Python code generation (DS-1000), text-to-SQL, goal inference, and molecule synthesis—and demonstrate consistent improvements over a variety of baselines and ablations. These gains hold across different model sizes, including smaller (1B) and larger (70B) LLMs, and across tasks that demand global semantic constraints beyond local token-level restrictions.

Accept (Oral)