Efficient Inference for Large Language Model-based Generative Recommendation

Dear Reviewer Up12,

Thanks for your comments. Your review is very detailed and thorough. we greatly appreciate it! We have provided detailed clarification, explanations, intuitions behind method design, and step-by-step derivation to address each of your concerns. We have also updated our manuscript accordingly (marked in blue). If we have any misunderstanding, please feel free to leave further comments. We eagerly anticipate our discussion with you!

Q1: In Section 3.1, paragraph "Acceptance Rate", the condition given for having $\beta = 1$, namely that $p(\mathbf{y}) \geq p(\mathbf{y}_K)$ for all $\mathbf{y}$ in $\mathcal{Y}_q$, doesn't seem to be fulfillable. By definition of $\mathbf{y}_K$, there exists only $K-1$ sequences with greater probability according to $p$, yet the condition requires that all $N$ ($\geq K$) sequences in $\mathcal{Y}_q$ have a probability greater than $\mathbf{y}_K$. The only possibility for this condition to be realized is if $\mathcal{Y}_q$ contains duplicates, which should not happen in the case of beam search.

Reply: Thank you for pointing out this problem. The acceptance rate in Section 3.1 indeed would be unfulfillable if $N>K$. The only possibility to achieve this condition is when $N=K$ and the top-$K$ ($N=K$) token sequences from the draft model are also the top-$K$ token sequences from the target model. To correct this, we revise the condition as:

$\beta=1$ if $\exists \mathcal{Y}'\_{q}\in\mathcal{Y}\_{q}$ such that $p(\mathbf{y})\ge p({\mathbf{y}}\_k)$ for $\forall \mathbf{y}\in\mathcal{Y}'\_{q}$, where $|\mathcal{Y}'_{q}| = K$, $K$ is the target LLM beam size and $\mathbf{y}_k$ is the sequence that has the $K$-th highest probability in $p$.

Based on this condition, the alignment objective as in Eq.(2) is unaffected, which aims to encourage the top-$N$ generated sequences from the target model to have high probabilities in the target LLM distribution (a high $\frac{p(\mathbf{y})}{p(\mathbf{y}_K)}$). Precisely, if the drafted sequence fails to be in the top-$K$ sequences according to $p$ (i.e., $p(\mathbf{y})<p(\mathbf{y}_K)$), we still encourage it to have a high probability closer to $\mathbf{y}_K$ (i.e., a high $\frac{p(\mathbf{y})}{p(\mathbf{y}_K)}$), aiming to push the top-$N$ distribution of draft model to cover the top-$K$ distribution of the target LLM.

Q2: Eq 3 introduces $\mathcal{D'}$ in which $\mathcal{Y}$ is the top-$K$ of a mixture of $q$ and $p$. Following Gu et al (2024) is mentioned as the motivation for this, but it would be appreciated to have more intuitions on this choice.

Reply: Thanks for your valuable suggestions. For the choice of the mixture of $q$ and $p$ in Eq.(3), we have two main considerations:

• 1) Higher training efficiency. The original $\mathcal{D}$ should include the sequences sampled from the draft model (i.e., $\mathbf{y}\in\mathcal{Y}_q$ in Eq.(2)). However, training over draft model-generated sequences requires the online learning process, which will lead to high computational costs and time costs for training. Because we need to sample the sequences continuously from the draft models for every epoch or even every batch during the alignment training process. Therefore, to alleviate the reliance on the frequent sampling from the draft model $q$, we consider using the mixture of draft model $q$ and target LLM $p$. In practice, the mixture of $p$ and $q$ is achieved by alternating the sequences sampled from $p$ and $q$, where the alternating sessions are controlled by the mixture coefficient $\lambda$. Since the target LLM-generated sequences can be pre-stored, we can improve the training efficiency by significantly reducing the number of sampling sequences from draft model.

• 2) Mitigation of low-quality training data issue. During alignment training, the draft model might generate low-quality sequence (e.g., repeated phrases). However, such low-quality sequences will be rejected by the target LLM during inference since they are invalid identifiers. As such, pushing $q(\mathbf{y})$ closer to $p(\mathbf{y})$ over low-quality sequences will lead to unnecessary and suboptimal alignment. Therefore, we consider utilizing LLM-generated data to ensure high-quality training data for alignment training.

Q3: Moreover, the $\mathcal{D'}$ used later in the relaxed alignment objective (Eq 8) is different. What is the rationale to have different $\mathcal{D'}$ in the strict and relaxed objectives?

Reply: Thanks for your insightful question. The different choices of $\mathcal{D}’$ essentially result from their different expressions of acceptance rate and different alignment objectives. Specifically,

• For the strict verification, since we hope every drafted sequence can achieve high probability in target LLM distribution, the alignment objective is maximizing $\frac{p(\mathbf{y})}{p(\mathbf{y}_K)}$ over the output distribution of draft model. Based on the consideration for higher training efficiency and mitigation of the low-quality training data issue as mentioned in the previous response, we adopt the mixed sampling for AtSpeed-S.

• For the relaxed verification, we are inspired by Eq.(1) in [1] to maximize the sequence acceptance rate over the output distribution of target LLM. This is because, with each approximated with-replacement sequence sampling (Eq.(6)), the expectation over the drafted sequence is essentially the expectation over the model vocabulary given the prefix. While the vocabulary is the same for both the draft model and the target LLM, the prefix of $\mathbf{y}$ in Eq.(6) is accepted by the target LLM in inference and thus follows the with-replacement sampling distribution of target LLM. Therefore, the $\mathbf{y}$ in Eq.(7) of our paper is considered to be sampled over the target LLM distribution for AtSpeed-R.

[1] Yongchao Zhou, et al. DistillSpec: Improving Speculative Decoding via Knowledge Distillation. ICLR 2024.

Q4:What loss is used in practice for LRec? It is only mentioned to be a recommendation loss but additional details would be appreciated (at least in the appendix)

Reply: Thanks for your valuable suggestions. The recommendation loss used in our work is defined as

$$\mathcal{L}\_\text{Rec} = - \frac{1}{N} \sum\_{(\mathbf{x}, \mathbf{y}) \sim \mathcal{D}} \sum\_{t=1}^{|\mathbf{y}|} \log \mathcal{M}\_p({y}\_t|\mathbf{y}\_{<t}, \mathbf{x}),$$

where $\mathbf{x}$ is user’s historical interactions, $\mathbf{y}$ is the user’s next interacted item identifier, and $D=\\{(\mathbf{x}, \mathbf{y})\\}$ denotes the original recommendation dataset. We have also added this to our manuscript in Appendix.

Q5: The relaxation sampling verification strategy misses some intuition: why is $p(\mathbf{y}) \geq q(\mathbf{y})$ the right criterion for accepting a sequence $\mathbf{y}$? Why are such $\mathbf{y}$'s good candidates? Moreover, in the case of rejection, $\mathbf{y}$ is drawn from $p' = norm(max(0, p(\mathbf{y}) - q(\mathbf{y})))$ but it is unclear whether this is an actual distribution and what the norm operator exactly consists of. Is $p'$ denoting the uniform distribution over the $\mathbf{y}$'s such that $p(\mathbf{y}) \geq q(\mathbf{y})$? If so, how does one sample from this in practice? The definition of this distribution would deserve more clarifications and details.

Reply: Thanks for your valuable questions and we appreciate your careful review.

• The intuition behind accepting sequence $\mathbf{y}$ if $p(\mathbf{y}) \ge p(\mathbf{y})$ is that, given a candidate sequence $\mathbf{y}$ with a high probability according to the draft model, if the target LLM is even more confident of generating the sequence (i.e., $p(\mathbf{y}) \ge p(\mathbf{y})$), the candidate sequence is likely to be generated by the LLMs and should be accepted. We have added the intuition explanation in our updated manuscript.

• In the case of rejection, the normalized probability distribution is adjusted as follows:

$$p’ = \text{norm}(\max (0, p(\mathbf{y})- q(\mathbf{y}))) = \frac{max(0, p(\mathbf{y})-q(\mathbf{y}))}{\sum_{\mathbf{y}}{max(0, p(\mathbf{y})-q(\mathbf{y}))}}$$

We then sample a sequence from this normalized distribution to replace the rejected sequence. The clarification of normalized distribution is added in our updated manuscript.

Q6: The proof of Lemma 1 in App A.2, which corresponds to Lemma 3.3 from Leviathan et al (2023), is missing the step with the term $1 - \sum_{\mathbf{y}} \frac{p(\mathbf{y})+q(\mathbf{y})-|p(\mathbf{y})-q(\mathbf{y})|}{2}$.

Reply: Thank you for pointing out this typo. Your careful review is greatly appreciated. We have revised it in our updated manuscript.

Step3. Relate $K\_i$ to $N$ and $p\_i$.

We then relate $K\_i$ to $N$ and $p\_i$. Since $K\_i=Np\_i$, we have $\ln \frac{K\_i}{k\_i}=\ln \frac{Np\_i}{k\_i}$. We also have $n=\sum\_{i=1}^{r}k\_i$. Then, we can express the logarithm of hypergeometric distribution in terms of $p\_i$ and $k\_i$ as

$$\begin{aligned} \ln P\_\text{hyper} & \approx \sum\_{i=1}^{r} [k\_i\ln \frac{K\_i}{k\_i}+ \frac{k\_i^2}{K\_i}] - [n\ln \frac{N}{n} + \frac{n^2}{N}] \\\\ & = \sum\_{i=1}^{r} [k\_i\ln \frac{Np\_i}{k\_i} + \frac{k\_i^2}{Np\_i}] - [n \ln \frac{N}{n}+\frac{n^2}{N}] \\\\ & =\sum\_{i=1}^{r} [k\_i (\ln N + \ln p\_i - \ln k\_i)+ \frac{k\_i^2}{Np\_i}] - [n\ln \frac{N}{n} + \frac{n^2}{N}] \\\\ & = \sum\_{i=1}^{r} k\_i\ln N + \sum\_{i=1}^{r}[k\_i(\ln p\_i - \ln k\_i)+ \frac{k\_i^2}{Np\_i}] - n\ln \frac{N}{n} - \frac{n^2}{N} \\\\ & = n\ln N - n\ln N + n\ln n - \frac{n^2}{N} + \sum\_{i=1}^{r} [k\_i(\ln p\_i - \ln k\_i) + \frac{k\_i^2}{Np\_i}] \quad (\text{we have } \sum\_{i=1}^{r}k\_i=n \text{ in last expression})\\\\ & = \sum\_{i=1}^{r} [k\_i(\ln p\_i - \ln k\_i)+ \frac{k\_i^2}{Np\_i}] + n\ln n - \frac{n^2}{N} \\\\ &= n\ln n - \sum\_{i=1}^{r}k\_i \ln k\_i + \sum\_{i=1}^{r} k\_i \ln p\_i - \frac{n^2}{N} + \sum\_{i=1}^{r}\frac{k\_i^2}{Np\_i}. \end{aligned}$$

Now the approximation of the logarithm of multivariate hypergeometric distribution has been finished.

Step4. Compare with multinomial distribution.

Similarly, we approximate the multinomial distribution with Stirling's approximation. We expand the logarithm of multinomial distribution as

$$\begin{aligned} \ln P\_\text{multi}(k\_1, k\_2, \dots, k\_r) & = \ln \frac{n!}{\prod\_{i=1}^{r}k\_i!}\prod\_{i=1}^{r}p\_i^{k\_i} \\\\ & = \ln n! - \sum\_{i=1}^{r} \ln k\_i! + \sum\_{i=1}^{r}k\_i \ln p\_i. \end{aligned}$$

Using Stirling's approximation, we have $\ln n ! \approx n\ln n - n$ and $\ln k\_i! \approx k\_i \ln k\_i -k\_i$. We then substitute $\ln n!$ and $\ln k\_i!$ with the approximation and obtain

$$\begin{aligned} \ln P\_\text{multi} &\approx n\ln n - n -\sum\_{i=1}^{r} (k\_i \ln k\_i - k\_i) + \sum\_{i=1}^{r} k\_i \ln p\_i \\\\ & = n\ln n - n - \sum\_{i=1}^{r}k\_i\ln k\_i + \sum\_{i=1}^{r}k\_i + \sum\_{i=1}^{r}k\_i \ln p\_i. \end{aligned}$$

Since we have $\sum\_{i=1}^{r}k\_i = n$, we have

$$\begin{aligned} \ln P\_\text{multi} &\approx n\ln n - n - \sum\_{i=1}^{r}k\_i\ln k\_i + n + \sum\_{i=1}^{r}k\_i \ln p\_i \\\\ & = n\ln n - \sum\_{i=1}^{r}k\_i \ln k\_i + \sum\_{i=1}^{r}k\_i\ln p\_i. \end{aligned}$$

Now, comparing the approximated logarithm of multinomial distribution with the approximated logarithm of multivariate hypergeometric distribution, we have

$$\begin{aligned} \ln P\_\text{multi} \approx \ln P\_\text{hyper} + \frac{n^2}{N} - \sum\_{i=1}^{r}\frac{k\_i^2}{Np\_i}. \end{aligned}$$

Note that the term $\frac{n^2}{N} - \sum\_{i=1}^{r}\frac{k\_i^2}{Np\_i}$ is negalectable when $N$ is large and $k\_i$ and $n$ are small compared to N. Therefore, we show that when the population size $N$ is large (e.g., all possible sequences for sampling) and $k\_i, n$ are small (e.g., $k\_i=1$ or $0$ since each sequence represents a category and $n$ is usually less than 20 in LLM-based recommendation), the multivariate hypergeometric distribution can be approximated to the multinomial distribution. That is, sampling without replacement is approximately equivalent to sampling with replacement.

Theorem 1. When population size $N$ is large and sample size $n$ is small compared to $N$ (i.e., $n\ll N$), the multivariate hypergeometric distribution approximates the multinomial distribution:

$$P\_\text{hyper}(k\_1, k\_2, \dots, k\_r) \approx P\_\text{multi}(k\_1,k\_2, \dots, k\_r).$$

Proof.

Step1. Apply Stirling's approximation on logarithm of multivariate hypergeometric distribution.

We can expand the logarithm of hypergeometric probability in factorials as follows:

$$\begin{aligned} &\ln P\_\text{hyper}(k\_1,k\_2,\dots,k\_r) =\ln \frac{\prod\_{i=1}^{r}\frac{K\_i!}{k\_i!(K\_i-k\_i)!}}{\frac{N!}{n!(N-n)!}} \\\\ =& \ln \prod\_{i=1}^{r}\frac{K\_i !}{k\_i!(K\_i-k\_i)!} - \ln \frac{N!}{n!(N-n)!} \\\\ =& \sum\_{i=1}^{r}\ln \frac{K\_i!}{k\_i!(K\_i-k\_i)!} - \ln \frac{N!}{n!(N-n)!} \\\\ =& \sum\_{i=1}^{r} [\ln K\_i! - \ln k\_i! - \ln (K\_i-k\_i)!] - [\ln N! - \ln n! - \ln (N-n)!]. \end{aligned}$$

Using the Stirling's approximation, we have:

\left\\{ \begin{aligned} \ln K\_i ! &\approx K\_i \ln K\_i - K\_i + \frac{1}{2} \ln 2\pi K\_i ,\\\\ \ln k\_i ! &\approx k\_i \ln k\_i - k\_i + \frac{1}{2} \ln 2\pi k\_i , \\\\ \ln (K\_i-k\_i) ! &\approx (K\_i-k\_i) \ln (K\_i-k\_i) - (K\_i-k\_i) + \frac{1}{2} \ln 2\pi (K\_i-k\_i) ,\\\\ \ln N! &\approx N\ln N -N + \frac{1}{2}\ln2\pi N ,\\\\ \ln n! &\approx n\ln n -n + \frac{1}{2}\ln2\pi n ,\\\\ \ln (N-n)! &\approx (N-n)\ln (N-n) -(N-n) + \frac{1}{2}\ln2\pi (N-n). \end{aligned} \\right.

Then, we can substitute the logarithm of factorials with the approximation as:

$$\begin{aligned} &\ln P\_\text{hyper}(k\_1,k\_2,\dots,k\_r) \\\\ = & \sum\_{i=1}^{r} [\ln K\_i! - \ln k\_i! - \ln (K\_i-k\_i)!] - [\ln N! - \ln n! - \ln (N-n)!] \\\\ = &\sum\_{i=1}^{r} [ K\_i \ln K\_i - K\_i + \frac{1}{2}\ln 2\pi K\_i - (k\_i \ln k\_i - k\_i + \frac{1}{2}\ln 2\pi k\_i) \\\\ & \quad\quad - ((K\_i-k\_i)\ln (K\_i-k\_i) - (K\_i-k\_i) + \frac{1}{2}\ln2\pi (K\_i-k\_i) )] \\\\ & \quad\quad - [N\ln N - N + \frac{1}{2}\ln 2\pi N - (n\ln n - n + \frac{1}{2} \ln 2\pi n ) \\\\ & \quad\quad - ((N-n)\ln (N-n) - (N-n) + \frac{1}{2} \ln 2\pi (N-n) )] \\\\ = &\sum\_{i=1}^{r} [ K\_i \ln K\_i - k\_i \ln k\_i - (K\_i-k\_i) \ln (K\_i-k\_i) + \frac{1}{2} \ln 2\pi K\_i - \frac{1}{2} \ln 2\pi k\_i - \frac{1}{2}\ln 2\pi (K\_i-k\_i) ] \\\\ & \quad\quad - [N\ln N -n\ln n - (N-n) \ln (N-n) + \frac{1}{2} \ln 2\pi N - \frac{1}{2} \ln 2\pi n - \frac{1}{2} \ln 2\pi (N-n) ]. \end{aligned}$$

Since $N$ is a very large number and $n \ll N$, $k\_i \ll K\_i$, we have $\frac{1}{2} \ln 2\pi K\_i - \frac{1}{2} \ln 2\pi k\_i - \frac{1}{2}\ln 2\pi (K\_i-k\_i) \approx 0$ and $\frac{1}{2} \ln 2\pi N - \frac{1}{2} \ln 2\pi n - \frac{1}{2}\ln 2\pi (N-n) \approx 0$. Then, the logarithm of multivariate hypergeometric distribution is approximated as:

$$\ln P\_\text{hyper} \approx \sum\_{i=1}^{r} [K\_i \ln K\_i - k\_i \ln k\_i - (K\_i-k\_i) \ln (K\_i-k\_i)] - [N\ln N -n\ln n - (N-n) \ln (N-n)].$$

Step2. Approximate $\ln (K\_i-k\_i)$ and $\ln (N-n)$ using Taylor expansion.

Since $k\_i$ is small compared to $K\_i$, we can expand $\ln (K\_i-k\_i)$ using Taylor expansion

$$\ln (K\_i-k\_i) = \ln K\_i - \frac{k\_i}{K\_i} - \frac{1}{2}(\frac{k\_i}{K\_i})^{2} + \dots,$$

where we can neglect the high-order terms and obtain

$$\ln (K\_i-k\_i) \approx \ln K\_i - \frac{k\_i}{K\_i}.$$

Similarly, for $\ln (N-n)$, we have $\ln (N-n) \approx \ln N - \frac{n}{N}.$

We can then substitute $\ln (K\_i-k\_i)$ and $\ln (N-n)$ in logarithm of multivariate hypergeometric distribution and obtain

$$\begin{aligned} \ln P\_\text{hyper} &\approx \sum\_{i=1}^{r} [K\_i \ln K\_i - k\_i \ln k\_i - (K\_i-k\_i) \ln (K\_i-k\_i)] - [N\ln N -n\ln n - (N-n) \ln (N-n)] \\\\ & = \sum\_{i=1}^{r} [K\_i\ln K\_i - k\_i \ln k\_i - (K\_i-k\_i) (\ln K\_i + \frac{k\_i}{K\_i})] - [N\ln N - n\ln n - (N-n)(\ln N + \frac{n}{N})] \\\\ & = \sum\_{i=1}^{r} [K\_i\ln K\_i - k\_i\ln k\_i - (K\_i\ln K\_i - k\_i\ln K\_i + k\_i - \frac{k\_i^2}{K\_i}) ] - [N\ln N - n\ln n - (N\ln N - n\ln N + n - \frac{n^2}{N})] \\\\ & = \sum\_{i=1}^{r} [k\_i \ln \frac{K\_i}{k\_i} + \frac{k\_i^2}{K\_i} - k\_i] - [n\ln \frac{N}{n} + \frac{n^2}{N} -n] \\\\ & = \sum\_{i=1}^{r} [k\_i \ln \frac{K\_i}{k\_i} + \frac{k\_i^{2}}{K\_i}] - [n \ln \frac{N}{n}+ \frac{n^2}{N}]. \quad\quad (\text{we have} -\sum\_{i=1}^{r}k\_i+n=0 \text{ in last expression}) \end{aligned}$$

Q13: The "with-replacement sampling approximation" paragraph in App A.2 is difficult to follow so it would benefit from being reworked.

Reply: Thanks for your valuable comments. The with-replacement sampling approximation is mainly based on Stirling’s approximation, and we provide the detailed step-by-step proof to facilitate better understanding. The detailed derivation has also been updated in the Appendix of our updated manuscript.

We aim to prove that the distribution of sampling without replacement is approximately equivalent to that of sampling with replacement. Our proof is mainly based on Stirling's approximation. In the following, we will first clarify notations, and introduce multivariate hypergeometric distribution and the multinomial distribution, which are used to model the sampling without replacement and sampling with replacement, respectively. We then present Stirling's approximation, and show the step-by-step proof.

Notations. To model the sampling, we have the total population size $N$, sample size $n$, the category size $r$, the number of items in category $i$ in the population $K\_i$, and the number of items in category $i$ in the samples $k\_i$. In the case of sequence sampling in LLM decoding, the population includes every possible sequence. Every possible sequence is a unique category, and the sample size $n$ is the beam size, the population size is the total number of all possible sequences at each beam search step. Now we assume that the population sizes go to infinity in such a way that $p\_i=\frac{K\_i}{N}$. And we have $\sum\_{i=1}^{r}k\_i=n$ and $\sum\_{i=1}^{r} K\_i = N$.

Multivariate hypergeometric distribution. Formally, when sampling without replacement, the probability of drawing $k\_1, k\_2, \dots, k\_r$ items from each category is given by the multivariate hypergeometric distribution

$$\begin{aligned} P\_\text{hyper}(k\_1,k\_2,\dots,k\_r) &= \frac{\prod\_{i=1}^{r}\binom{K\_i}{k\_i}}{\binom{N}{n}} \\\\ &=\frac{\prod\_{i=1}^{r}\frac{K\_i!}{k\_i!(K\_i-k\_i)!}}{\frac{N!}{n!(N-n)!}}. \end{aligned}$$

Multinomial distribution. Formally, when sampling with replacement, the probability follows the multinomial distribution as

$$\begin{aligned} P\_\text{multi}(k\_1, k\_2, \dots, k\_r) = \frac{n!}{\prod\_{i=1}^{r}k\_i!}\prod\_{i=1}^{r}p\_i^{k\_i}. \end{aligned}$$

Stirling's approximation. Stirling's approximation gives us the approximation of the logarithm of factorials as:

$$\ln n! \approx n\ln n -n + \frac{1}{2}\ln(2\pi n).$$

Step3. Move expectation over $\mathbf{y}\_{<t}$.

We aim to align the draft model with the target LLM at every beam search step. Notably, at each beam search step $T$, the sequence lengths are fixed and are independent with $t$. Therefore, we can rewrite the objective into:

$$\begin{aligned} & - \sum\_t \mathbb{E}\_{\mathbf{y}\_{<t}\sim q} [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ] \\\\ = &- \mathbb{E}\_{\mathbf{y}\_{\le T}\sim q} \sum\_{t=1}^{T} [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ]. \end{aligned}$$

Step4. Normalize by length and rewrite the objective.

Since we aim to align at every step $T\in\{1,\dots,L\}$, where $L$ is the length of item identifier in LLM-based recommendation, we further normalize the expression by sequence length to prevent different scales on alignment loss across different steps. As such, for every step $T\in\{1,\dots,L\}$, the objective can be rewritten as

$$\begin{aligned} & - \mathbb{E}\_{\mathbf{y}\_{\le T}\sim q} \frac{1}{|y\_T|} \sum\_{t=1}^{|y\_T|} [ \sum\_{y\_t} q(y\_t) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ]. \end{aligned}$$

However, the expectation over the sequence space, (i.e., $\mathbb{E}\_{\mathbf{y}\_{\le T}\sim q}$) is intractable, so we follow previous work (Wen et al., 2023; Kim & Rush, 2016) to approximate it by sampling top-$K$ sequences generated by draft model $\mathcal{M}\_q$. Now we can rewrite our alignment objective for strict top-$K$ verification and obtain Eq.(3) in our paper:

$$\begin{aligned} & \mathop{\arg\max}\_{\theta\in\Theta} - \mathbb{E}\_{(\mathbf{x},\mathcal{Y})\sim D'} \sum\_{\mathbf{y}\in \mathcal{Y}} \frac{1}{|\mathbf{y}|} \sum\_{t=1}^{\mathbf{|y|}} [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ] \\\\ =& \mathop{\arg\min}\_{\theta\in\Theta} \mathbb{E}\_{(\mathbf{x},\mathcal{Y})\sim D'} \sum\_{\mathbf{y}\in \mathcal{Y}} \frac{1}{|\mathbf{y}|} \sum\_{t=1}^{|\mathbf{y}|} [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ], \end{aligned}$$

where $\mathcal{Y}$ denotes the top-$K$ beam sequences generated from the mixture distribution of draft model and target LLM (as stated in the "Alignment Objective" paragraph in Section 3.1, line 216-218).

Q7: In Eq 7, what does $\sum_K$ denote? It seems like there is an index variable missing there.

Reply: Thanks for your question. The $\sum_{K}$ in Eq.(7) is written in the form of summation to denote the derivation from the joint distribution of $K$ sequences, i.e., $K$ in $\log \beta = - \sum_{K} \log \text{TVD}(q, p)$ in the “Alignment Objective” paragraph of Section 3.2. To clarify this, we explained the meaning of $K$ in the context of Eq.(7) in line 281-284 in our submitted manuscript, i.e., line 284-286 in our updated manuscript (“Eq.(7) minimizes TVD with the strength of K for the same sequence. Nevertheless, considering the beam search with sampling would obtain $K$ different sequences, we alternatively leverage the top-$K$ sequences from the target LLM”). This also partially explains our choice of $\mathcal{D}’$ to sample Top-$K$ sequences from the target LLM for the alignment training under relaxed sampling verification.

Q8: The strategy relying on tree-based attention to speed up inference would benefit from being described in more details (at least in Appendix). Section 3.3 only refers to Figure 2(c) to describe the strategy, but this figure is not self-explanatory.

Reply: Thanks for your insightful comments. We provide more details and a concrete example of how tree-based attention is implemented in practice. We have also added them to our updated manuscript.

The main idea of tree-based attention is to eliminate the repeated self-attention calculation for the same prefix of the beam sequences during beam search.

• Detailed explanation of tree-based attention. In tree-based attention strategy, we first compress the N beam sequences into a single flattened sequence, and then construct the sparse tree-based attention mask for efficient target LLM verification. More precisely, given $\gamma N$ drafted beam sequences with different lengths, where $N$ is the draft beam size and $\gamma$ is the number of drafted beam steps,

  • we first flattened the beam sequences by sequentially adding the newly generated tokens from the beam sequences at each step. We denote the length of the flattened sequence as $L_f$.

  • Then, based on the flattened sequence, we construct an attention mask with the shape of $L_f \times L_f$. Specifically, each row in attention mask represents a specific beam sequence, and for each row, we set the corresponding column of the last token and the preceding tokens in the beam sequence as 1, otherwise 0.

• Example. Here, we give a concrete example to illustrate how tree-based attention is implemented to save repeated calculation. We set the beam size of 2 and collect the beam sequences of 3 steps. The collected beam sequences are as follows:

  step 1 beam 1: a1

  step 1 beam 2: a2

  step 2 beam 1: a1 b1

  step 2 beam 2: a2 b2

  step 3 beam 1: a1 b1 c1

  step 3 beam 2: a1 b1 c2

  The flattened sequence will be a1 a2 b1 b2 c1 c2. The constructed tree-based attention is shown in Figure 2(c) of our manuscript, where each row represents a specific beam sequence. For each row, the tickled cell represents the preceding tokens and the last tokens of each beam. And the different colors represent the different steps of beam search. This flattened sequence and sparse tree-based attention enable efficient verification since it saves the repeated calculation of the same prefix across different beam sequences, e.g., a1b1.

Q9: The WS metric represents the walltime speedup, but with respect to which baseline? I assume this is in comparison to directly running the target model without speculative decoding, but it would be helpful for the reader to mention this when defining the metric.

Reply: Thanks for your valuable question and suggestion. Your understanding is correct. The walltime speedup is compared to the execution of original target LLM without speculative decoding. And the walltime speedup is defined as $WS=\frac{T}{T'}$, where T is the time for running target LLM without speculative decoding and T’ is the time for running LLM with speculative decoding with a specific draft model. We have also added the clarification to our updated manuscript.

Step2. Decompose $\sum\_{\mathbf{y}}q(\mathbf{y})$ and obtain step-wise alignment.

Since $\mathbf{y}=(y\_1, y\_2, \dots, y\_n) = (\mathbf{y}\_{<t}, y\_t, \mathbf{y}\_{>t})$ is a sequence, we can decompose $\sum\_{\mathbf{y}} q(\mathbf{y})$ into nested sum $\sum\_{\mathbf{y}\_{<t}} \sum\_{y\_t} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}) = \sum\_{\mathbf{y}\_{<t}} \sum\_{y\_t} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{<t}, y\_t, \mathbf{y}\_{>t})$, where the nested sum over three parts of $\mathbf{y}$, i.e., $\mathbf{y}\_{<t}, \mathbf{y}\_t$, and $\mathbf{y}\_{>t}$ can cover all possible sequence $\mathbf{y}$.

Since $q(\mathbf{y}\_{<t}, y\_t, \mathbf{y}\_{>t})=q(\mathbf{y}\_{<t})q(y\_t|c\_{<t})q(\mathbf{y}\_{>t}|c\_{\le t})$, we can rewrite the nested sum of $q(\mathbf{y}\_{<t}, y\_t, \mathbf{y}\_{>t})$ over $\mathbf{y}$:

$$\begin{aligned} & \sum\_{\mathbf{y}} q(\mathbf{y}) \\\\ =&\sum\_{\mathbf{y}\_{<t}} \sum\_{y\_t} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{<t})q(y\_t|c\_{<t})q(\mathbf{y}\_{>t}|c\_{\le t}) \\\\ = &\sum\_{\mathbf{y}\_{<t}} \sum\_{y\_t} q(\mathbf{y}\_{<t})q(y\_t|c\_{<t}) \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) \\\\ = &\sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}). \end{aligned}$$

We then can substitute $\sum\_{\mathbf{y}} q(\mathbf{y})$ with $\sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|\mathbf{y}\_{<t}) \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t})$ in previous expansion and obtain:

$$\begin{aligned} & \quad - \sum\_t \sum\_{\mathbf{y}} q(\mathbf{y}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} + \sum\_t \sum\_{\mathbf{y}} q(\mathbf{y}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} \\\\ & = - \sum\_t [ \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) ] \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} + \sum\_t [ \sum\_{\mathbf{y}\_{<t}} q(c\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) ] \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} \\\\ & = - \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) + \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}). \end{aligned}$$

Since $\sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t})=1$, we can remove $\sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t})$:

$$\begin{aligned} & - \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) + \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} \sum\_{\mathbf{y}\_{>t}} q(\mathbf{y}\_{>t}|c\_{\le t}) \\\\ =& - \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} + \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} \\\\ =& - \sum\_t \sum\_{\mathbf{y}\_{<t}} q(\mathbf{y}\_{<t}) [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ] \\\\ =& - \sum\_t \mathbb{E}\_{\mathbf{y}\_{<t}\sim q} [ \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p(y\_t|c\_{<t})} - \sum\_{y\_t} q(y\_t|c\_{<t}) \log \frac{q(y\_t|c\_{<t})}{p (\mathbf{y}\_{K,t}|\mathbf{y}\_{K,<t})} ]. \end{aligned}$$

Now we have a step-wise alignment over all sequence from $q$, i.e., $\mathbb{E}\_{\mathbf{y}\_{<t}\sim q}$.