Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

Thank you for your strong review and thoughtful questions! We address comments below.

stochastic runtime: runtime is stochastic and could vary significantly across instances, which can be an issue in production.

Thanks for pointing this out. Yes, stochastic runtimes (particularly if the distribution of runtimes has long tails) can lead to issues in production. In Appendix H, we present a modified version of our algorithm that incorporates a user-specified upper bound on the number of attempts AWRS makes to generate a satisfying next token. This requires a slight change to the computation of proper weights for use within SMC, also presented in Appendix H. Note that, using this variant of our algorithm, if a single particle within SMC has found itself in a poor position, it will stop early with a weight of $0$, allowing resampling to cull the particle (and clone a more promising particle) without slowing down the rest of the algorithm. If a constraint is so difficult that the upper bound is hit frequently (across all particles in an SMC algorithm), the user may opt to increase their upper bound, dispatch to a larger model more capable of satisfying the constraint, or increase the number of particles in SMC to mitigate the issue. We will revise the paper to include a clearer discussion of these considerations.

limited performance gap in some domains: when the local constraint is already exact, e.g. Molecular Synthesis and Pattern Matching, improvements are less dramatic.

As noted, in some settings the constraint is already quite exact, so additional corrections based on proper weighting may be unnecessary. We will make this point more clearly in our discussion of the results. In general, we would recommend that users try multiple settings for $M$, the number of particles in AWRS-SMC, and use the lowest setting that meets the functional requirements of their application. If $M=1$ is sufficient, then proper weighting has no effect, and the user can simply run ARS-LCD. However, in some domains, the proper weights really do help, and we see the ability to compute them without summing over the entire token vocabulary as a key benefit of our approach.

computational overhead of multiple loops: while faster than token masking, the algorithm still requires multiple rejection loops to compute the unbiased Z estimator.

It is correct that the algorithm requires multiple loops to compute the unbiased $Z$ estimator. However, the overhead of this is rather minimal. In our experiments, we find that only $2$ total loops are sufficient, and these loops are not independent. Rather, the 2nd loop proceeds immediately after the first, which has already adaptively rejected some number of non-conforming tokens. The 2nd loop does not need to reject these again, and can even sample the same correct token that was already just sampled. Algorithm runtime is thus dominated by the first loop. For example, in the context-sensitive pattern matching domain, the second loop samples only $1$ token in $92$%, $85$%, and $72$% of runs for Llama 70B, 8B, and 1B, respectively. We will revise to discuss this point.

Thank you for the review and for these thoughtful questions and comments! We answer all questions below.

There is no discussion of how this algorithm might be combined with truncation based sampling such as topp or \eta sampling. It would be great if the authors would include some discussion about that.

Thanks for this suggestion! There are a few ways one might combine our method with truncation-based modifications to an LM’s sampling procedure.

The most straightforward way is to think of truncation-based sampling as modifying the base LM distribution, on top of which our constrained decoding approach can then be applied. For example, in top-p sampling with $p=0.9$, we zero out the bottom $10$% of token probabilities and renormalize to obtain a modified next-token probability distribution. We can then apply our algorithm as though the base LM had generated this truncated next-token probability distribution: we adaptively rejection-sample from the truncated distribution, rather than from the full next-token distribution. One caveat is that the user must ensure that truncation does not remove all valid next tokens (according to the constraint) from the support.

Another way to combine our methods with truncation-based sampling is to consider the truncation as part of the constraint: we write a new constraint function that accepts a token only if it passes the original constraint and is in the top $90$% of logits from the base model. If AWRS-SMC were applied with this constraint, the weights would correct not only for myopia from locally constrained decoding, but also for myopia from top-p decoding. As the number of particles grows, the algorithm’s output distribution would converge to the global posterior distribution that arises when the base LM is conditioned on the event that (1) the string satisfies the original constraint, and (2) each token is in the top $90$% of the LM’s predictions given the previous tokens. Note that, as above, the user must be careful to avoid the situation where there are no strings that satisfy this more stringent constraint.

We will revise the paper to include discussion of these points!

Thank you for this review and for your thoughtful questions and suggestions! We address your comments below.

It is not clear to me that evaluating constraints is a major bottleneck in many constrained decoding scenarios, especially when LLM forward pass costs may dwarf constraint costs. That being said I can imagine more expensive constraint scenarios. It would be helpful to have comments in the paper about the relative costs of different common constraints.

Thanks for this suggestion. The cost of constraint evaluation indeed varies dramatically depending on the complexity of the constraint. We will include the following table in the final version of the text, which shows the mean cost (ms/eval) of the constraints we test in this paper.

As can be seen, costs range from $<0.01$ms to $>10$ms, and yet all can still be integrated into the decoding process with ARS, incurring a comparatively low constant-factor overhead (mean $1.22\times$ sec/ex across our benchmarks). Contrast this with the full token-masking approach, which evaluates the constraint on all tokens at each step. This method is so costly that we could only tractably run the full token masking baseline for the context-sensitive pattern matching domain. Even there, the approach incurs $>50\times$ overhead. Without our algorithm, these more complex constraint evaluations can dwarf even LLM forward pass costs.

Alternatively, most existing implementations of token masking are for simple grammar constraints, where highly-engineered parsers in performant languages like Rust or C++ can exploit the limited constraint expressiveness and systems engineering to evaluate the constraint for multiple possible next tokens simultaneously while navigating runtimes on the order of 10 microseconds (Dong et al., 2024). In these cases, your intuition that LLM evals dwarf constraint evaluation does hold, and token masking can be a highly effective technique.

What our work contributes is a new way to integrate arbitrary “black-box” constraints that do more than check adherence to a simple grammar, for which such extensive performance engineering may not be feasible. Indeed, in our experiments, our constraints are written as short Python functions with minimal performance engineering; AWRS lets us use them in the decoding loop without a severe hit to overall decoding speed. An impact of this work is thus both an increase in the complexity of what can be efficiently constrained at decode time as well as a usability benefit for constraint programmers, who can now afford milliseconds per constraint evaluation rather than being restricted to microseconds. We see many opportunities for interesting constraints at this time scale, e.g., incremental type systems, simulators, and so on. Consider the Goal inference domain, for example, where a PDDL planner can be integrated into the decoding loop with ARS, whereas naive token masking would require $11$ms/tok * $256,000$tok > $45$min per timestep for a model with a vocab size of $256$k, such as Gemma.

References:

• Dong, Y., Ruan, C. F., Cai, Y., Lai, R., Xu, Z., Zhao, Y., & Chen, T. (2024). Xgrammar: Flexible and efficient structured generation engine for large language models. arXiv preprint arXiv:2411.15100.

Thank you for this encouraging and thoughtful review! We address your comments below.

not clear whether implementation will be released, since this approach is engineering heavy, the lack of a reference implementation would be a significant bottleneck for future work

Thanks for pointing out this was unclear. To clarify, in the final version of the paper, we plan to link to:

1. a lightweight, performant reference implementation that we have contributed to an existing open-source library for constrained generation.
2. a static snapshot of the code required to reproduce the experiments and analyses presented in the paper.

The reference implementation is, in fact, already available, but we have withheld the link for anonymity during review.