Guaranteed Generation from Large Language Models

Thank you for your review and constructive feedback on our work! We appreciate your recognition of the strengths and value of our paper. We address below the points raised in your review.

Comparison with the heuristic sampler from Appendix A.2

We assume this question is referring to the heuristic algorithm presented in the Enforcement example paragraph in Appendix A.2 (if not, please correct us). While this algorithm has clear expected drawbacks, we agree it would be interesting to verify its performance numerically in the lexical experiment to include “amazing”. We tested this, in the bounded-length setting of Section 4.1 with a 30-token generation length. In the heuristic sampler, we check whether we have already produced the string "amazing" within the first 29 tokens, and if not, force the generation of "amazing" as the last token. With the Gemma-2B tokenizer, the string "amazing" can be generated with the single token [amazing] (along with [ama, zing], etc.) and we can easily design the sampler to produce this token as the 30th token if “amazing” does not appear earlier. For such a sampler $s$, we found that $KL(g||s)=6.06$, which is a significantly larger divergence than $KL(g||g')=0.633$ for the $a'$ based on warm-start DPG.

GUARD itself does not increases the acceptance rate, but $a’$ does

We agree that our formulation in the third contribution could lead to some confusion. We will modify it to clarify that it is indeed the choice of $a’$ that leads to improving the acceptance rate, rather than the GUARD framework itself.

Remarks for the camera-ready version

Thank you for pointing out these mistakes, we will correct them in the updated version of the submission.

Question 1 — What is $V^*$ in line 119?

This is the “Kleene star” notation for the set of all finite sequences based on the token vocabulary $V$.

Question 2 — What is "ancestral sampling"?

The term “ancestral sampling” denotes basic sampling from the model — namely sampling with temperature 1, directly reflecting the probability distribution associated with the model without any distortions or truncation (unlike techniques such as top-$k$ or nucleus sampling).

Thank you for your detailed feedback. We address below each of the points raised in your review.

Minimization of $KL(g || a')$

As you rightfully pointed out, our goal is to minimize $KL(g || a')$ to obtain a 'good' $a'$ (as justified by Theorem 2 which links this quantity to both efficiency and closeness to $g$). In fact, both SFT and DPG seek to minimize $KL(g || a')$ by definition, as detailed below.

The SFT loss corresponds to the cross-entropy of the model $a'$ to be learned, using samples from $g$, i.e., $-\mathbb{E}\_{y \sim g} \log a'(y)$. Minimizing this term with respect to $a’$ is equivalent to minimizing $KL(g || a')$ since $KL(g || a') = \mathbb{E}\_{y \sim g} \log \frac{g(y)}{a'(y)} = \mathbb{E}\_{y \sim g} \log g(y) -\mathbb{E}\_{y \sim g} \log a'(y) = H_g - \mathbb{E}\_{y \sim g} \log a'(y)$ where $H_g$ is the entropy of $g$ which is a constant independent of $a’$.

DPG’s objective is also to minimize the KL divergence (or equivalently, the cross-entropy) between the target distribution $g$ and the learned policy $a' = \pi_{\theta}$ as mentioned in Section 3, paragraph Approximating the target distribution: “This method samples $y$ from a proposal $a'$ initialized with $a$ and updates $a'$ by performing gradient descent on $KL(g||a')$.” This property is further detailed in Section 3.2 of Parshakova et al (2019). The derivation can be summarized as follows: $\nabla_{\theta} KL(g||\pi_{\theta}) = \nabla_{\theta} (H_g - \mathbb{E}\_{y \sim g} \log \pi_{\theta}(y)) = -\mathbb{E}\_{y \sim g}\nabla_{\theta} \log \pi_{\theta}(y) = -\mathbb{E}\_{y \sim \pi_{\theta}} \frac{g(y)}{\pi_{\theta}(y)} \nabla_{\theta} \log \pi_{\theta}(y)$, where the last step is obtained by applying importance sampling using $\pi_{\theta}$ as proposal. Performing gradient descent on $-\mathbb{E}\_{y \sim \pi_{\theta}} \frac{g(y)}{\pi_{\theta}(y)} \nabla_{\theta} \log \pi_{\theta}(y)$ leads to the DPG update rule shown at Line 21 of Algorithm 2.

We will update the paper to clarify the $KL(g || a')$ minimization property in SFT and DPG.

Rejection sampling in FUDGE

While rejection sampling is indeed briefly discussed in Yang and Klein (2021) as a possible way to improve FUDGE, it is not explored in that paper and only mentioned “in passing” as a potential extension. In contrast, we provide in our paper an in-depth theoretical and empirical investigation of the use of rejection sampling for guaranteed generation.

Design of $b$

We agree that the choice of $b$ is important (as pointed out in our Limitations paragraph in the Conclusion section). However, we still believe that the investigation on the design of $b$ is orthogonal to the focus of our paper — namely, the GUARD approach — as it implies a modification of the core target $g$. In other words, modifying $b$ would affect the target distribution $g$, while our paper's primary objective is to find a guaranteed sampler that approximates $g$ for a given $b$. Studying the choice of $b$ may then warrant a paper of its own, and any improvement of $b$ will directly translate into improving GUARD, as reviewer vMCq also noted.

Consistency in ARM abbreviation usage

Thank you for pointing this out, we will correct this.

Computation of $KL(g || a')$, $KL(g || g')$ and $AR_{a'}$

According to Theorem 1, it is indeed typically impossible to find an autoregressive model that matches distribution $g$ exactly, but we can still obtain gold samples from $g$ by first sampling from $a$ and then filtering with $b$ (though this may be highly inefficient in case of low $AR_a$). These samples are used to estimate $KL(g || g')$ as detailed in Equation 10 in Appendix C. The estimate of $KL(g || a')$ is obtained in the same way. Similarly, $AR_{a'}$ which is equal to $\mathbb{E}\_{y \sim a'} [b(y)]$ can be estimated by drawing samples from $a’$. The consistency of these estimates can also be verified using Theorem 2 which links these three quantities.

Quality of generated answers and capabilities of $g'$ wrt $a$

There are two types of deviations of $g'$ relative to $a$. The first deviation is unavoidable and corresponds to an intrinsic distortion, which is the fact that $g$ itself has to deviate from $a$ to some extent to enforce the constraint $b$ (although this deviation is minimal as explained in Appendix A.1). The other deviation is related to the approximation error, namely the value $KL(g || g')$ that we are trying to minimize.

Note that in some cases the constraint $b$ makes it impossible to maintain the capabilities of the original model (e.g., if $b$ corresponds to a strict notion of harmlessness, it may not preserve the ability of the model to be helpful in certain limit cases), while in some other cases (e.g., if $b$ enforces a polite tone) the capabilities of $a$ are mostly preserved in $g$ (and, in turn, $g'$). Because of these variations, that depend on the nature of $b$, we concentrate on the approximation error $KL(g || g')$ rather than on the intrinsic distortion $KL(g || a)$ in our experiments.

We also provide examples of generated text with DPG-based $g'$ and CAP-based $g'$ in Tables 5 and 9 in the Appendix.

Efficiency of the different approaches

The efficiency at inference time is measured by the acceptance rate of $a’$. The acceptance rate improvement of GUARD over applying rejection sampling on $a$ is up to 180x for the lexical constraint experiment and 60x for the sentiment reversal experiment, as detailed in Sections 4.1 and 4.2, respectively. In more intuitive terms, using $a$ for rejection sampling leads to accepting 1 sample out of 435 drawn samples for lexical constraints and 1 sample out of 200 drawn samples for the sentiment reversal. In comparison, using GUARD, this goes up to accepting approximately 1 sample out of 2 drawn samples and 1 sample out of 3 drawn samples, respectively. We will update the paper to include these numbers, which we agree are easier to interpret.

Thank you for the detailed response!

• I appreciate the clarification on rejection sampling and the fact that SFT and DPG minimize your objective. Could you also explain how CAP aligns with this?
• I still do not fully agree with the statement that investigating the design of $b$ is orthogonal work. If it is impossible to construct a meaningful $b$ for real-world applications, finding a guaranteed sampler that approximates $g$ for a given $b$ adds no practical value.
• Regarding the quality of the answers, I am mainly concerned that $g'$ might underperform on general benchmarks compared to $a$ (e.g., solving math problems when constrained to avoid generating political statements). This concern arises because fine-tuning alters the original model, and the CAP prompts themselves appear to be quite ambiguous (e..g. generating a sentence that contains "amazing" involves a prompt like "Write sentences with the given words. diagnosis: Assessment of microscopical and clinical parameters in the diagnosis of diabetes mellitus. pandas: Column headings differ in spreadsheet that is merged with pandas data. change: How to change the decimal separator in MS Word? amazing:")
• Regarding efficiency, does this imply that 435 samples are required to obtain one additional training sample for SFT and DPG? How many samples did you generate in total for training? Similarly, for computing $KL(\cdot || \cdot)$, how many samples were generated to approximate the divergences between $g$, $g'$ and $a'$?

Q3. Does the fine-tuning or prompting of $g'$ lead to a loss of general performance in comparison to $a$?

First, we agree that using CAP with the few-shot prompt mentioned in your question would lead to a loss of general performance in the resulting $g'$. This is one of the main downsides of prompting, which tends to bias and restrict the scope of the generated output, as pointed out in our answer to Q1.

For a $g'$ obtained through fine-tuning, using DPG or SFT, the loss of general performance is related to the intrinsic distortion and to the approximation error that we discussed in our previous response. Depending on the constraint $b$ of interest, certain capacities from $a$ will not be present anymore in $g$ — which relates to the intrinsic distortion. To follow the example given in your question, using a $b$ that forbids the mention of some political parties X, Y, and Z would make it impossible for $g$ to solve a math problem such as “In a small town, there are three political parties, the X, the Y and the Z, each party has a certain number of supporters, …”. This would then naturally be reflected in the performance of the DPG/SFT-based $g'$ on this specific math problem.

However, on math problems which do not mention parties X, Y, and Z, the performance of $g$ would be equivalent to that of $a$. Then, if the approximation error between $g$ and $g'$ is low (as it is the case for DPG/SFT-based $g'$ in the lexical constraints and sentiment reversal settings), the performance of $g$ (and thus $a$) will be preserved in $g'$.

Q4. Details about the number of samples used for training and evaluation.

You are right that obtaining one additional sample for SFT does require 435 samples from $a$, which is also our argument against using SFT due to its inefficient nature. However, for DPG, this is only true at the beginning of the training of DPG, (when $\pi_\theta = a$ in Algorithm 2, implying few effective gradient updates on line 21), because later on, $\pi_\theta$ gets closer to $g$ and therefore produces more actual updates — this is what makes DPG an adaptive algorithm. When warm-starting DPG with CAP, the samples from $\pi_\theta$ are actually useful from the start, as we mentioned in the answer to your first question. The number of samples used in training (sampling from $a$) is defined as the "sampling budget" in Figures 3 and 5 from the revised PDF, and we experimented with up to 800,000 and 1,500,000 samples, respectively.

For the evaluation, we sampled $2^{20}$ (i.e., approximately 1,000,000) new samples from $a$ to perform the KL estimation — while this is costly, we typically only perform these estimations for scientific analysis purposes. Then, in the lexical constraints scenario (where $AR_a$ = 0.0023), this leads to approximately 2,300 samples from $g$. For the sentiment reversal experiment (where $AR_a$ = 0.005), this yields 5,000 samples from $g$. We have added these details in Appendix D in the revised submission.

We hope that our responses helped clarify the concerns you raised, and we thank you again for your valuable engagement in this exchange.

Thank you once again for taking the time to follow up on our previous responses.

Q1. To me, it remains unclear how to obtain such a prompt. Moreover, this does not provide an argument for guaranteeing that $KL(g||a')$ is minimized.

To clarify any misunderstanding, we did not say that the $a'$ resulting from prompting guarantees a low $KL(g||a')$, but only that it tends to satisfy the constraint much more often than the original model (which is very different) without any training cost, namely that $AR_{a'}$ tends to be better than $AR_{a}$. In fact, we explicitly said the opposite: “CAP tends to have a worse $KL(g||a')$ than those $a'$ that are obtained by either SFT or DPG […]. This implies that when CAP is used, the constraint is enforced by strongly biasing the generation, ie, at the cost of a high $KL(g||a')$”. Experimentally, looking at Fig. 4, we do see that every CAP-based $a'$ have $KL(g||a')$ larger than 10, even worse than $KL(g||a)$ which is around 6. Please note that in order to read $KL(g||a')$ from this figure, one needs to look at the intersection between the dotted line and the x-axis. For instance, the leftmost pale-blue CAP point lies on the dotted line that intersects this axis for $x \approx 10$. By comparison, all the fine-tuned versions of $a'$, either through SFT or through one of the two DPG variants, has a $KL(g||a')$ lower than 4. The same trend can be observed in Fig. 7, for the sentiment-reversal experiment.

On the other hand, what we indicated in our previous response is that a CAP-based $a'$ is a good starting point for the warm-start version of DPG (namely DPG initialized with $a(\cdot|CAP)$). Warm-start DPG yielded better results for $KL(g||a')$ than DPG both in Fig. 4 and more so in Fig. 7, and with a quicker convergence than DPG (see Fig. 3). In both experiments, the best $KL(g||a')$ result is obtained for warm-start DPG, at around 1.4 in Fig. 4 and 1.9 in Fig. 7.

Q2. Your arguments are convincing. They should be discussed in more detail in the main paper.

Thank you! We will update our paper with these arguments as soon as we are allowed again to update the paper.

Q3. Claiming such statements (”the performance of $g$ would be equivalent to that of $a$”) without supporting evidence is not convincing. It is well known that fine-tuning and instruction-prompting can significantly impact a model's performance.

First, we wish to draw your attention to the fact that $g$ itself is not the result of fine-tuning or instruction-prompting from $a$. It is the model obtained as $g(y) = \frac{a(y)b(y)}{Z}$, with $Z$ a normalization constant. Given this formulation, we note that on the support of $g$, corresponding to the set of sequences satisfying constraint $b$ (i.e., the $y$'s such that $b(y) = 1$), we have $g(y) = a(y) / Z$. Intuitively, this means that $g$ is proportional to $a$ on the support of $g$.

Now, let us consider again a math problem which excludes political parties A, B, and C, as discussed in our previous response. If we have a set of numerical answers for this math problem, these answers will be ranked in the same way by $a$ and $g$ due to the proportionality relationship presented above. In other words, if $a$ is able to identify the correct answer among the set of possible answers, then $g$ will also identify the same correct answer. This proves that the performance of $g$ and $a$ are equivalent on the support of $g$ for this example, and this can be easily extended to more general cases. We will clarify this property in the paper when updating is allowed again.

Q4. While I appreciate the insights, the need to generate such a large number of samples for a relatively simple constraint amplifies my concerns about the practicality and scalability of the approach.

In our previous answer, we mentioned that we experimented with up to 800K and 1.5M samples from $a$ for our two settings, which indeed may appear to be large numbers. However, we wish to remind that in practice the warm-start DPG variant is able to get a close-to-optimal $KL(g || a')$ for as little as 100K samples in the lexical constraint scenario and 200K samples on sentiment reversal (see Figs. 3 and 5, respectively). We believe that these numbers of samples remain in a reasonable range.

Moreover, our experiments have been conducted in settings with particularly low acceptance rates (0.0023 for the lexical constraint scenario and 0.005 for sentiment reversal) to better understand the characteristics of guaranteed generation. There is however a large range of concrete applications where constraints such as safety, politeness, or cultural sensitivity would typically lead to a higher acceptance rate. GUARD would show great practical potential for such cases and would not suffer from scalability issues.

Considering the additional elements we provided in response to your latest concerns, do you still believe that a score of 5 is the most appropriate for our work?

Many thanks for your very positive feedback and for the useful suggestions for improvement! We address below the points raised in your review.

Suggestion 1. Algorithm 1 is too short. It is frustrating to see a 3-line algorithm in a paper. Switching its position with that of Algorithm 2 would be much better.

Thank you for the suggestion. We agree that the way Algorithm 1 is presented may give a bad first impression to the reader, and we will update the submission to present it in a more subdued way inside the text. As for Algorithm 2, it would take a lot of space in the main text, is focused on training, not on sampling, and its exact form is actually less essential to the core message of the paper than what the extremely simple Algorithm 1 says, so we prefer to keep it in the Appendix.

Suggestion 2. The content of Theorem 2 is too simple to be a theorem. It's OK to put it as an equation.

It is true that the formulation of this theorem is simple, as well as its proof (see Appendix A.6 for two proofs, including one simple direct derivation). However, its proof is not the most important part, but rather its interpretation, which is core to our work. While we have considered calling it a “Lemma” or a “Fact”, we prefer not to give it such a secondary status in relation with Theorem 1.

Suggestion 3. For Theorem 1 in main text, I suggest to replace it with the complete version in appendix. The concept of PTC is central and should be highlighted in the main text.

We totally agree with this suggestion and will follow it, thanks a lot!

Suggestion 4. In section 2.3, It may not be a good idea to discuss the non-zero probability under softmax as an evidence for limits of autoregressive models, since it is not the key point and such shortcoming can be easily avoid by rejection sampling. I suggest to discuss the PTC property of common models instead, which will serve for Theorem 1.

We will discuss the PTC property, as mentioned above in relation to Theorem 1, but also find the discussion of more mundane properties of autoregressive models useful, and they also connect, here and in the Appendix that expands on them, with some remarks by other reviewers.

Suggestion 5. In experiments, I did not find the exact numbers of AR of DPG. I notice that they are reported in the figures by coordinates, but exact numbers are also necessary.

Thank you for pointing out what we missed. We will explicitly provide the AR values before and after training.

Suggestion 6. Text in figures and tables are too small to be viewed on A4 paper, please consider rearranging the layout.

We are sorry that the paper may be difficult to view on A4 paper, and we do agree that formatting needs improvement, but this is unfortunately difficult to do while remaining within the 10 pages limit. We will do our best to maximize the use of the free space that remains to address feedback from reviewers and to improve the layout.

Question 1. How do you get the gold samples in Fig. 4&7?

By applying Algorithm 1 with $a’=a$ (see second paragraph of Section 3) it is possible to obtain samples from $g$ by first sampling from $a$ and then filtering with $b$ (though this may be highly inefficient in case of low $AR_a$).

Question 2. I notice that you discussed the relation to I-projection in line 136&242. How do such discussions help? I found that removing these contents does not affect understanding.

We agree with you that, concerning line 242, the reference to I-projection is not absolutely necessary, because it is possible to give a simple self-contained derivation (we do so in Appendix A.6) and we will reformulate the text to explicitly mention that derivation. However, we feel that mentioning I-projection in the context of Theorem 2 could be useful for extensions of the work presented here (I-projection is an important information-theoretic concept, that can also be used for soft constraints of the type used in Khalifa et al (2021)).

On the other hand, in the context of line 136, we feel that this notion is really useful, because it provides another motivation for the definition of $g$, namely as the distribution that minimally distorts $a$ while respecting the constraint. So, overall, despite the introduction of not absolutely necessary terminology, we prefer to keep it.

Thank you for your feedback, and in particular for calling our attention to some important references that we had missed at the time of submission.

We feel that [4] is the one most directly related to our work, and we will start by discussing it, stressing common points as well as fundamental differences, before moving to other aspects of your feedback. We will also update our submission (as soon as possible before the rebuttal deadline) with expanded related work reflecting [4] and other references, and will also correct mistaken statements implying that our work is the first to address 100% constraint satisfaction while attempting to keep close to the original LLM.

Relation of our work to [4]

Similar to us, [4] is able to produce outputs that always satisfy the constraint, while maintaining proximity to the original model. It does so by approximating the LLM with an HMM, which, in contrast to the LLM, supports a dynamic programming approach to complying with lexical constraints over the output, and therefore the ability to weigh the long-term consequences of local next-token choices relative to these constraints. When the HMM approximation to the LLM is good enough (which depends in particular on the number of states of the HMM), the same HMM can be used at inference time with different lexical constraints without need of retraining.

We acknowledge the interest of this approach, but note several fundamental differences with what we do.

• We handle arbitrary logical constraints (that is, binary predicates), not necessarily of a lexical nature, which would be difficult to handle through HMMs or similar finite-state mechanisms.
• We give a central status to the “gold” constrained distribution $g$, which is the distribution minimally deviating from the original distribution while still fully satisfying the constraint, and we evaluate the quality of our sampler in terms of distributional distance to this distribution $g$. [4] does not focus on the evaluation of a sampler relative to a reference distribution, but rather on the evaluation of a decoder in terms of downstream tasks which have a looser relation with the constraint.
• We study the trade-off between the quality of the sampler and its efficiency in the context of a simple rejection sampling mechanism based on an autoregressive approximation $a’$ of $g$ and show that quality and efficiency are both directly controlled by the divergence $KL(g||a’)$. This in turns motivates our interest for an approximation technique that focuses on minimizing this divergence, such as DPG.

Note: Apart from [4], we recently came across the follow-up work [7], which generalizes [4] by considering lexical constraints defined through DFA’s (Deterministic Finite Automata). As a side note, on p. 13 of our original appendix (“Solvable instances of the problem”), we mentioned briefly that in case the base autoregressive model and the lexical constraint are based on weighted finite state automata, the intersection of these automata could be constructed and sampled from. Such weighted automata are very much related to HMMs, of course, and we will also update the appendix to refer to [4,7], which had escaped our attention at that time.

[7] Zhang, H., Kung, P., Yoshida, M., Van den Broeck, G., & Peng, N. (2024. August). Adaptable Logical Control for Large Language Models. arXiv preprint arXiv:2406.13892.

Relation to other works mentioned

Thank you for the other references you mentioned. There is indeed a very large literature on controlled generation and we cited a survey from 2023. Apart from [4], the only such reference concerned with strict guarantees is [5], also in the context of lexical constraints, but here the focus is simply on efficient pruning of next tokens that do not satisfy a finite-state constraint, without any global reweighing. Although we were not aware of this work at the time of submission, we did discuss a similar approach for enforcing or avoiding the inclusion of specific keywords in Appendix A.2. We will update this section to add a reference to [5].

Concerning the other references, note that we did actually cite [1], which is indeed relevant to our work, and we will also cite [3], which does have some significant relation to our work. We will mention [2] as well, although it is strongly focused on deterministic decoding rather than sampling.

Concerning the CommonGen challenge [6], the constraints correspond to lists of keywords, and the task is to formulate sentences that contain the keywords and follow a certain common-sense logic. In particular, the dataset includes a few specific target sentences for each list of keywords, which are provided as references. The goal for evaluated models is then to generate sentences as close as possible to these few references. This substantially differs from our distributional objective where we seek to generate sentences that are distributed in a similar way as samples from $g$. For that reason, it does not seem straightforward to evaluate GUARD on CommonGen.

Using NADO distribution [3] as proposal in GUARD

The NADO distribution used by [3] is trained with an objective similar to ours, namely it tries to approximate a distribution comparable to our gold filtered distribution $g$. As a training method to obtain an autoregressive distribution similar to our $a’$, it has some analogies with (i) SFT, by training on samples from the filtered distribution, and (ii) DPG, by also performing a kind of distributional matching, but with more emphasis on local (i.e., token-level) decisions. Contrary to DPG, NADO does not directly have the objective of minimizing the divergence $KL(g||a’)$, which is the determining factor for the success of GUARD. Using NADO for $a’$ would nonetheless be interesting as a follow-up work, to see whether the value of its divergence is competitive with the cases we have explored, leading to improved performance for GUARD.

Dear Reviewers,

Thank you again for the time invested in our paper and for your valuable feedback!

We have uploaded an updated version of our submission, in which we have attempted to address your detailed questions and suggestions. The changes are highlighted in blue for your convenience. We hope that these revisions will help clarify any concern you may have had about our paper and that they do improve the overall quality of our work.

If you feel that our responses and revisions addressed your concerns, we would be most grateful for this to be reflected in score adjustments; otherwise, we would be happy to further discuss any remaining doubts.