Approximately Aligned Decoding

Thank you for your review. We would like to respectfully correct a misunderstanding: while our method is inspired by speculative decoding, it is not ASAp plus speculative decoding---there is no small speculative model, and if no errors are encountered during generation, the prefix selection algorithm is never invoked.

The main insight behind AprAD is that instead of re-using the entire prefix (constrained decoding), or re-using none of the prefix (ASAp), it is possible to re-use part of the prefix in an error-free generation task. The speculative decoding algorithm has some properties (1) that make it particularly suited as a prefix selection algorithm between these two extremes, but other prefix selection algorithms could work in its place (see Section 6.2 and Appendix C).

If AprAD were just ASAp with speculative decoding, then we would agree with most of the stated weaknesses. Speculative decoding would not change the sampling distribution of ASAp, and we would have structured our evaluation environment so that only one experiment runs at a time per machine, to prioritize consistent wall-clock times over machine utilization. Speculative decoding usually increases the total amount of computation in terms of FLOPs, and derives a latency improvement by parallelizing the compute and reducing memory reads- ASAp with speculative decoding would have a higher generation ratio compared to the original version of ASAp.

In contrast, AprAD derives its performance improvement by reducing the total amount of computation relative to ASAp; the tradeoff is that our method introduces a distribution shift (2). This leads to a lower generation ratio, but higher KL-divergence, though not as much as constrained decoding.

If desired, traditional speculative decoding could be layered on top of the generation phase of AprAD, in between backtracking iterations, as an orthogonal improvement to improve its latency at the expense of additional computation.

We will update the wording to clarify the distinction between AprAD versus ASAp with speculative decoding.

(1) In particular, that it scales prefix selection likelihood by the amount of probability mass that is "lost" along a given path after incorporating an error. It also usually selects a prefix that ends in a high-entropy token.

(2) The distribution is skewed because a non-violating sequence may either be generated directly, or by re-sampling after a violating sequence occurs. In contrast, with speculative decoding, a sequence is never generated directly---the output is always resampled after any speculative sequence (even non-violating ones). See discussion with reviewer 5S8W for more analysis.

Thank you for your review.

Representing $P^B$ in an implementation

The main structure we use to track token probabilities is a trie. However, note that there are times where the program must use both $P^B$ and $P^{B \cup \\{x\\}}$. To avoid the need to copy the trie, which can grow quite large, our program first collects the probabilities that will later be required from $P^B$ during SpecSample (the token probabilities along the "path" to $x$) before calling AddBadSample to transform the trie into $P^{B \cup \\{x\\}}$.

Additional implementation details are located in Appendix F. We will release our code along with the camera-ready version.

Runtime reporting

Unfortunately, the wall-clock runtimes would not be meaningful as we maximized machine utilization by running several experiments in parallel, leading to varied times depending on exact resource contention at the time of a given experiment. As the runtime is dominated by the LLM forward pass, we used generation ratio as a proxy for the required amount of computation.

Analysis of probability amplification

The recursive/backtracking nature of AprAD meant that we did not have complete theoretical results available at submission time, and omitted a theoretical discussion of the method to focus on the empirical results. If the reviewers feel that this additional analysis would be helpful, we will add it to the Appendix, with a disclaimer:

Intuitively, during a traditional speculative decoding procedure with speculative model $S$ and main model $P$, there are two ways that $y$ may be generated at the end of the process:

1. $y$ is drawn directly from $S$, and then it is kept during the speculative sampling process.
2. Some sample $x \neq y$ is drawn from $S$, and then during the speculative sampling process, the it is "re-sampled" to $y$.

In AprAD we use $P$ as the "speculative" model. $P^\mathcal{B}$ would ideally be the "main" model, though we must approximate it with $P^{\\{x\\}}$ after the first iteration.

If $y \notin \mathcal{B}$ is initially drawn, AprAD always keeps it, in contrast to its possible discarding by the speculative sampling procedure. This increases the overall probability of generating $y$, but only by as much as the probability of sampling $y$ in the first place (when speculative sampling would be likely to discard $y$). The second way of generating $y$ does not increase in probability with AprAD. Therefore, the probability of generating $y$ is at most $2 P^{\\{x\\}}(y)$. As $P^{\\{x\\}}(y) \leq P^{\mathcal{B}}(y)$ for $y \notin \mathcal{B}$, there is a maximum amplification of 2.

Note that this amplification is per-iteration, where an iteration is defined as encountering an error sequence, potentially backtracking, and resampling. In practice, the per-iteration amplification is likely much less, as several of the inequalities involved are very loose. However, there is a cumulative effect as more iterations are required in dense error sets.

In contrast, the probability amplification after encountering even a single error sequence in constrained decoding is unbounded, as the best non-error token may have an arbitrarily low probability.

A more in-depth explanation is as follows:

Let $SRS(y|x, S, P)$ (SpeculativeReSample) be the probability of eventually selecting sequence $y$ with speculative decoding, conditional on having drawn $x$ from speculative model $S$. The original version of speculative decoding implicitly relies on the identity $P(y) = \sum_{x \in \Sigma^*} S(x) SRS(y|x, S, P)$.

If $B \subseteq \mathcal{B}$, then $SRS(y|x, S, P^B) \leq SRS(y|x, S, P^\mathcal{B})$ for $y \notin \mathcal{B}$. Intuitively, $P^\mathcal{B}$ excludes more invalid sequences compared to $P^B$, so the probability mass of these sequences during the re-sampling process should be distributed among the non-error sequences, including $y$.

For similar reasons, $P^B(x) \leq P^\mathcal{B}(x)$ for $B \subseteq \mathcal{B}, x \notin \mathcal{B}$, as excluding additional errors $\mathcal{B} \setminus B$ and renormalizing the probability means that all other sequences become more likely.

The probability of sequence $y \notin \mathcal{B}$ being generated after the first iteration of AprAD is equal to the probability that $y$ is generated directly, plus the probability that $y$ is re-sampled after some invalid sequence is generated:

$AprAD(y|P, \mathcal{B}) = P(y) + \sum_{x \in \mathcal{B}} P(x) SRS(y|x, P, P^{\\{x\\}})$

We apply both inequalities described above:

$\leq P^\mathcal{B}(y) + \sum_{x \in \mathcal{B}} P(x) SRS(y|x, P, P^\mathcal{B})$

As $P(x)$ and $SRS(y|x, P, P^\mathcal{B})$ are always nonnegative, we can add additional elements to the sum---expanding it to include all sequences, rather than just error sequences---while maintaining the inequality:

$\leq P^\mathcal{B}(y) + \sum_{x \in \Sigma^*} P(x) SRS(y|x, P, P^\mathcal{B})$

And finally use the speculative decoding identity:

$\leq 2 P^\mathcal{B}(y)$

Thank you for your comments- we will incorporate this feedback into the final version.

The algorithm is general purpose; are there more problem / errors sets you can experiment with?

You are correct that the method is general purpose; however, we focused on only a few evaluation domains to allow for a more in-depth comparison between methods---BigCodeBench as a representative coding task that relies on multiple libraries, and Lipogram as a free-form generation task with a dense but easily understandable error set.

One potential interesting application of AprAD is to enforce syntactically correct tool use in an agentic system- we leave this for future work.

Thank you for your review.

Why were more detailed comparisons to Fudge or Ctrl-G not covered in the paper?

Due to time constraints, we focused our main comparison on training-free methods that work with black-box constraints. Both FUDGE/Ctrl-G require an extra training step to learn a discriminator/distilled HMM. Additionally, Ctrl-G is limited in the types of constraints that are expressible as a DFA, and as such couldn't be used for the hallucination avoidance task (the allowed text depends on variable assignments, imports from the Python environment, etc).

While we could not run a full evaluation on these methods, the discussion in Section 6.1 was informed in part by preliminary experiments on posterior-estimation based methods that did not make it into the paper.

We observe that PˆB and PˆB∪{x} are almost always near-identical distributions, with PˆB∪{x} generally as a “more accurate” distribution because it incorporates an additional error sample

How does this comparison hold with different values of x? For example, when the model is certain about x vs when the model is not certain about it, or, when x is a function word vs when it is not a function word?

Generally, the more likely that $x$ was to be generated by $P^B$, the larger the change in the distribution between $P^B$ and $P^{B\cup \\{x\\}}$- if $x$ was already low probability in $P^B$, then adding it as a violating sample will not change the distribution very much. This is irrespective of if the last token of $x$ is a function word, content word, or if the domain isn't natural language at all.

More concretely, consider a lipogram generation excluding the letter "u"

Prefix: "Attention Is All"

Next token probabilities:

You - 0.5
Important - 0.4
Natural - 0.1

P^B∪{"Attention Is All Natural"} is going to be closer to P^B compared to P^B∪{"Attention Is All You"}, because the latter requires excluding a sequence of much higher probability. However, both are "more accurate" than $P^B$ because they have both excluded a sequence that is incorrect according to the lipogram constraint.

Minor

Thank you, we will make these clarity changes.

Dear Reviewer,

Your active participation in the review process is crucial. Please respond to the authors' response and acknowledge you have done so.

Thanks.

-AC