Accelerating Blockwise Parallel Language Models with Draft Refinement
Authors: Taehyeon Kim
Ananda Theertha Suresh
Kishore A Papineni
Michael Riley
Sanjiv Kumar
Adrian Benton
Venue: NeurIPS 2024 poster (2024)
Keywords: Blockwise parallel decoding
Language model
speculative decoding
Primary area: natural_language_processing
TL;DR: We improve blockwise parallel decoding in language models by analyzing token distributions and developing refinement algorithms for block drafts, significantly achieving higher block efficiency across various tasks.

Abstract: Autoregressive language models have achieved remarkable advancements, yet their potential is often limited by the slow inference speeds associated with sequential token generation. Blockwise parallel decoding (BPD) was proposed by Stern et al. [42] as a method to improve inference speed of language models by simultaneously predicting multiple future tokens, termed block drafts, which are subsequently verified by the autoregressive model. This paper advances the understanding and improvement of block drafts in two ways. First, we analyze token distributions generated across multiple prediction heads. Second, leveraging these insights, we propose algorithms to improve BPD inference speed by refining the block drafts using task-independent \ngram and neural language models as lightweight rescorers. Experiments demonstrate that by refining block drafts of open-sourced Vicuna and Medusa LLMs, the mean accepted token length are increased by 5-25% relative. This results in over a 3x speedup in wall clock time compared to standard autoregressive decoding in open-source 7B and 13B LLMs.

Accelerating Blockwise Parallel Language Models with Draft Refinement
Taehyeon Kim Ananda Theertha Suresh² Kishore Papineni
Michael Riley?
Sanjiv Kumar?
Adrian Benton
KAIST AI
Google Research
Abstract
Autoregressive language models have achieved remarkable advancements, yet their potential is often limited by the slow inference speeds associated with sequential token generation. Blockwise parallel decoding (BPD) was proposed by Stern al. [42] as a method to improve inference speed of language models by simultane- ously predicting multiple future tokens, termed block drafts, which are subsequently verified by the autoregressive model. This paper advances the understanding and improvement of block drafts in two ways. First, we analyze token distributions generated across multiple prediction heads. Second, leveraging these insights, we propose algorithms to improve BPD inference speed by refining the block drafts using task-independent n-gram and neural language models as lightweight rescorers, Experiments demonstrate that by refining block drafts of open-sourced Vicuna and Medusa LLMs, the mean accepted token length increased by 5-25% relative. This results in over a 3x speedup in wall clock time compared to standard autoregressive decoding in open-source 7B and 13B LLMs.
1 Introduction
The landscape of natural language processing has been profoundly reshaped by recent advances in autoregressive language models [3, 48, 34, 37, 47]. These models have shown remarkable proficiency across range of text generation tasks, including applications like question answering [38] and summarization [17]. However, a significant obstacle to their wider application is high inference latency, particularly for extremely deep models with hundreds of billions of parameters [18, 35, 7]. This latency, intrinsic decoding with autoregressive language models (LMs), imposes considerable computational burdens and limits real-time deployment.
In response to these challenges, the field has seen shift towards decoding methods aimed reducing the inference latency in large language models (LLMs). One promising development is the concept of blockwise parallel decoding (BPD) [42, 31, 4]. Unlike autoregressive decoding, which generates one token at a time, blockwise parallel LMS are outfitted with a set of prediction heads that propose and verify draft, a block of subsequent tokens, in parallel. While BPD offers one solution to accelerated text generation, it also poses challenge in ensuring that the proposed drafts are fluent and natural.
BPD inference speed depends both the time it takes generate block draft and verification of the draft's agreement with the original LM's output (referred to as base LM from here on) (Figure la). Unlike standard autoregressive LMs that generate tokens sequentially ensuring consistency with all preceding tokens (e.g.. 'Messi' following 'Lionel') BPD employs a non-autoregressive drafting strategy. Here, blockwise parallel LMs simultancously predict multiple token drafts (c.g., 'Lionel' and 'Ronaldo') )) each position produced independently. The primary challenge in BPD drafting is ensuring that these enocurrnnly-generaatee tokens are consistent with each other. An effective block drafter should prefer coherent sequences, such as 'Lionel Messi" over less coherent combinations like
"Work done while at Google Research as student researcher. Corresponding authors: <kimtae- hycon61 10 @aamalll@mmml..mmmm. cmaml.com> and <adhenton@e @ google.com>.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

24 Meria
Loghts
Refined block draft
best refined block draits
Lord Mosi tie best best
Lionel Merd is the best sin
Lionel Ronaild is the best who the
Ussrll Mand the bust who the
Local
Giobal
Small Neural
N with best drafts
Autoregresiive Language Model
Blockwise Parallel Language Model
Block Draft
Neural Rescoring
Poliect Algrom Rescoring
Audiiaad
Adtleal
Seat
Dp-
2 Chiss
2 Menoa
2
2:
Prompt
Who the best soccer player the world Ar
(a) Example of block drafts
(b) Output of our proposed rescoring algorithms
Figure 1: (a) Illustration of two tokens that are decoded by autoregressive decoding vs. two tokens drafted by BPD. (b) Outputs from our proposed algorithms,, where the top-k token-level predictions are refined using local neural or global n-gram rescoring, which selects the p most probable sequences by dynamic programming, for batched verification.
'Lionel Ronaldo', which would be improbable under reasonable L.M. The focus of this paper is on improving the quality of block drafts without altering the underlying model parameters.
2 Our contributions
This paper first investigates properties of the drafts from blockwise parallel LMs across seven tasks. These analyses are based on modest, 1.5 billion (B) parameter LMs. Given our observations, we propose lattice rescoring algorithms to produce higher quality block drafts. Finally, we apply these lattice rescoring algorithms to improve the drafts from large (7B/13B parameter) open-source LLMs, reducing mean per- per-token latency relative to both standard BPD and Medusa decoding across tasks.
2.1 Observations on block drafts
Consecutive repetitions All heads within a block make predictions independently in blockwise parallel LM. Unsurprisingly, we observe that this leads to block drafts with significant token repetition across heads. Consecutive repetition is pervasive across tasks, ranging from 20% to 75% of all neighboring draft tokens, depending on the task (Section 5.1).
Confidence of different heads We analyze the distribution of probabilities within each block head. Our empirical analysis reveals an interesting property of BPD: the block drafter tends to be more confident with initial tokens, and becomes progressively less confident for subsequent tokens. We find that the confidence of block heads correlates strongly with the quality of the block drafter (Section 5.2).
Oracle top-k block efficiency the standard BPD algorithm (Algorithm 1), the most likely token at each head is generated as the draft. As mentioned above, this is prone to two issues: (1) this sequence might contain unnatural, consecutive repetitions and (2) the model might not be confident of the prediction at some of the heads. We use block efficiency, the average number of draft tokens accepted during decoding, to measure the quality of given drafter [28, 46]. ask whether the block efficiency can be improved by considering the top-k most likely tokens at each head. To measure the potential benefit of considering top-k tokens, we define the block efficiency of the oracle path through this top-k lattice oracle top-k block efficiency, and show that there is significant headroom for improvement across tasks (Section 5.3).
2.2 New block draft algorithms with lightweight rescoring
Based on these observations, we propose two algorithms to leverage the top-k predictions at each head and improve average latency for open-source LLMs (Figure 1b). We show that these algorithms can also reduce the average latency in Medusa decoding [4], a recent popular extension of BPD (Section 7). Neither of these algorithms requires changes to the underlying blockwise parallel LMs.
Local rescoring via neural LMs Given the top-k predictions at each head, we refine the block draft by using a small neural, autoregressive LM to greedily rescore these local predictions (Section 6.1). While the block prediction scores are produced independent of each other, neural rescoring should favor sequences that fluent, encouraging coherence between the predictions at each head.
Global rescoring via n-gram LMSs with multi-drafts If the blockwise parallel LM has h heads and we consider the top-k tokens from each head, then there are h candidate drafts of length h that

---

Algorithm 1: Blockwise parallel decoding (BPD)
input Blockwise parallel LM M2, initial prompt sequence 7 and target sequence length T.
1: Initialize & {-]
2: while & < ' do
3: /* Stage 1: Predict 8/
4: zi + M-,,, (gt), vi h.
5: Øc+1: 1e-1-Ve2,,...... Dr+h arg yav [v]. arg yax =?[v)..... arg max 2
6:
1* Stage 2: Verify $/
7:
for j t O.....,
....,, h in parallel do E04s +MMo Use 11+1, 22-2: , end for
8:
9:
10:
1 Stage 3: Accept $/
n f max(n bo+s arg max yey yev 244--1ly1,1 1 n}
t + + Bo+s + Ø145, S n and yr+++1 arg max Jev 2i4n [y]
13: end while
can be formed. We propose to use æ-gram model to efficiently rescore all paths, via dynamic programming. and generate the p most probable rescored paths as a batch draft candidates. These p drafts can then be verified in parallel by the blockwise parallel L.M (Section 6.2).
There are two critical distinctions between the proposed algorithms: the amount of context/expressive power available to each class of rescoring model, and fundamental limitations of decoding with each class. While neural rescoring models are potentially more expressive and can leverage unbounded context, n- n-gram LMs can be used to efficiently find the globally most likely rescored drafts from the exponentially-sizee sized set of possible draft candidates. Detailed algorithms are given in Section 6.1.
3 Preliminaries
Autoregressive decoding Let MM be autoregressive LM parameterized by 9. The objective is to generate an output sequence y<r = (yh..... yr) conditioned on an input sequence T. 2t Mo(-[ū, Vs*) is a vector of logits, 2t E TR VI, where V is the vocabulary over tokens. Let 2 [v] denote the logit of symbol y. These logits define a conditional probability distribution at each time c.t.sl
step po(y T., ygg))= =
Evecy e which by the chain rule yields po po(yr((0) = II Po(yel J Ber)).
Sequences are generated autoregressivell, either through ancestral sampling from some form of the conditional next token distribution [19], or by beam search through the space of possible sequences to return a probable sequence. For simplicity, in this paper we focus on greedy decoding, where at each step the next token is predicted as Vt+1 arg max Po(y| Ugt). The goal of BPD is to predict the same tokens as the base model, albeit efficiently.
Blockwise parallel decoding Let Mb be blockwise parallel LM with block size h and let 2f Ma (-1², Ve&) be the vector of logits corresponding to the 7th block given context ã. Ust- This model employs h distinct feedforward neural (FFN) layers, each with single hidden layer, atop the base LM's final hidden layer. The output of each FFN is followed by softmax layer over the vocabulary to predict each of the / subsequent tokens in the block. In our initial analyses, the parameters of the FFNs are learned jointly with the base LM during training, and the weights of all softmax layers are tied to the input embedding table. Similar to [42], the first head is the same as the base LM, i... 2 = Mbh (-1æ, v5() Ma (, Wee) = 2N and the hope is that for subsequent heads i 2 2, 6, (-1z, ver) R Ma (-) 1))
Algorithm describes the BPD greedy decoding procedure. We outline the algorithm below and refer readers to [42] for additional details.
1. Predict: Mb generates draft of h token predictions 91-2. ... Uc conditioned on the prompt, H. and existing generated text, Usa (i.e., Wi+: arg max 2 [[y] Vi < h). Since the first head is same as the base LM, ftt1 is identical to yt+1, the output of the base LM with greedy decoding.
2. Verify: In order to verify the predicted drafts, the base LM greedily generates next- token logits ( } conditioned on the existing prefix and block draft i.e., âtt = Ma (z, ysu. 1e 21+2, ..., Bets) for i E (0.1,...,h, h). Verification amounts to check-
3

---

Table 2: Sample outputs from blockwise parallel LMs finetuned per task. Black indicates standard decoded output, blue indicates accepted draft tokens, and brown is the prompt.
Table I: Per-task test performance of each finetuned model and block efficiency over language modeling (LM), extractive ques- tion answering (QA), and both long and short summarization (L-Sum & S-Sum).
LAMBADA
it's nothing more than a faceless, formless brown blob to me, but I take his word for the resemblance to our genctic makeup. ... [Skip...
Block Efficiency
Task
Dataset
Performance
SQuAD VI
LM
LAMBADA [33]
7.88
3.12
Question: Who was unnouncedd as the LEM contractor in November 1962?
QA
SQuAD VI [38]
57.60
2.08
context: Wiesner kept up the pressure, even making disagreement public (Skip)...
S-SUM
CNN/Daily [17]
39.85
1.74
Answer: Grumman
SAMSUM [14]
37,66
1.27
XSUM
MultiNews (12)
23.08
1.10
L-SUM
XSUM [32]
52.15
1.13
Summarize: /SKip)...
NewsRoom [15]
39.85
1.08
Millions of small businesses will benefit from a reduction of business rate from the Budget Osborne, Chancellor George Oshome has announced.
ing which block draft tokens match the autoregressive greedy decode from the base LM: (arg yax f [v]) == Dr+i+1. Note that the verification of all positions can be performed in parallel under the assumption that the base LM is a decoder-only transformer.
3. Accept:
Finally, the length of the longest contiguous prefix n where draft tokens match the base LM's greedy decode is identified. Since the first head is the same as the base LM, the first token it++1 is always accepted. After accepting the tokens, one free token can be obtained as the conditional probability of the base LM based on accepted tokens have already been calculated. Thus, the decoded sequence is extended by n + tokens and we iterate. Typically, not all h tokens are accepted, with some draft tokens discarded. As the block generation has minimal overhead compared to the base LM's forward pass, even modest gains in accepted prefix length justify the cost of block draft generation.
4 Analysis setup
We train a 1.5 billion (B) parameter decoder-only transformer LM with 9 heads, and investigate the drafts produced by this modest blockwise parallel LM. The 1.5B model and all auxiliary LMs were pretrained on (English) C4 [36] with the causal next token prediction objective tokenized with the GPT3 subword vocabulary [3]. For the 1.5B blockwise parallel LM, all heads were trained jointly to predict the following h tokens at each iteration. During pretraining, we use batches of 2048 subword sequences, each 5 512 tokens in length, amounting to z 200B input tokens in total. Model training/inference was run on TPUv3/TPUv4 [20], and implemented in Jax [2].
We evaluate the potential latency improvement of block drafts by block efficiency [28, 46]. In this context, block efficiency represents the theoretical speedup compared to standard greedy decoding. is defined as the average number of tokens decoded per serial call to the blockwise parallel LM. The formula for block efficiency is given by B= Total number of scrial calls to M
Total number of decoded tokens
In this definition, the total number of decoded tokens is the sum of the number of accepted tokens across decoding steps, not necessarily all h predicted tokens in each block. Only the tokens that pass the "Verify' stage and align with the base LM's predictions are accepted and integrated into the final sequence. This ensures that generated text is identical to the base LMM, while achieving speedup. The total number of serial calls to Ma is the number of times the model processes a block of tokens. A block efficiency of 1 means that one is achieving no speedup relative to standard decoding.
We investigate the drafts produced by this 1.5B blockwise parallel LM on LAMBADA [33] (language modeling). SQUAD V1 [38] (extractive QA), along with five summarization tasks: XSUM [32], MultiNews [12], SAMSum [14], NewsRoom [15] and CNN/DailyMail [17]. For each task other than language modeling, finetune the blockwise parallel LM for that task. 3
"This model follows the original blockwise parallel LM, with modification: we decoder-only models instead of the T5 encoder-decoder architecture. Other than this, the LM architecture consistent with that proposed in Stern et al. [42].
Details are given in Appendix D.

---

Ermma cerbutin fo exch head
3.0
Linear Ragression
Block Efficiency
2.5
SQuaD
2.0
ChMpuy
SAMBUM
MattNews 1a0
1.0
Erenpy
Mmax
(a) Entropy distributions across block draft heads (b) Correlation between block efficiency and Imax
Figure 2: (a) Entropy distributions across block draft heads LAMBADA [33]. The density plots illustrate the entropy distribution for each head in the model. (b) Correlation between block efficiency and hmax, the head until which the average entropy task increases nearly monotonically.
Table shows that block efficiency varies dramatically across task.* Language modeling, most closely matching the pretraining objective, achieves the highest block efficiency followed by the context-constrained task of extractive question answering. Table 2 sketches how BPD acts three examples from each class of tasks.
LM: BPD excels generating common multi-word expressions in single step. For example, (no) "thing more than', and (take) 'his word for the" are each drafted and accepted in single step.
QA: BPD also attains high block efficiency in extractive QA, where correctly drafts multi- token entities copied from the input sequence. In SQuAD VI, it accurately completes the answer 'Grumman' from 'Gru' by adding 'mman', highlighting its ability to process multiple tokens at once and quickly extend answers.
SUM: BPD's effectiveness in SUM tasks varies by dataset. For formulaic summaries like Mall it performs well. reflecting its alignment with LM and QA tasks. However, in marrative-drive datasets like SAMSum and XSUM, where concise summaries are required, the block efficiency of BPD little better than standard decoding.
Exploration of block drafts
5.1 Consecutive repetition
We observe that vanilla block drafts prone to sig- nificant token repetition. This is due to the fact that each head's prediction is independent of the others, and is a limitation shared with con-autoregressive generation in general [16]. Table 3 shows the pro- portion of consecutive tokens in block drafts that identical each other, along with the average maximum length of repeated sequences in block drafts across all decode time steps. We compare these statistics before and after rescoring with a 2-gram LM a trivial rescorer, but one that can encourage local consistency between consecutive draft tokens. Strings of repeated tokens unnat- ural, and unlikely to be generated by a strong base language model. Rescoring the top-k lattice with even a simple language model eliminates signif- icant amount of repetition, reducing the percentage of consecutive repeated tokens from between 9.9% 24.5%. %, depending on the task.
Table 3: Consecutive token repetition in block drafts before and after C4-trained 2-gram rescor- ing the top- lattice. "% Consec" the per- centage of consecutive identical draft tokens out of all pairs of consecutive tokens. "Max run" is the average maximum repeated subsequence length in tokens (upper bound of 9, the number of block draft heads). Higher values correspond to more egregious repetition in drafts.
R. Consec
Max
Task
Dateset
Vanilla
2-gram
Vanilla
2 gram
LM
LAMBADA
20.0
10.7
2.2
188
QA
SQUAD VI
75.5
67.6
6.6
S-SUM
SAMSUM
20.9
21.9
20.0
3.1
2.5
MuliiNews
14.7
3.1
2.1
1-SUM
XSUM
24.0
9.4
2.6
1.7
NewsRuom
47.2
32.1
4.1
33
5.2 Confidence across multiple heads
Intuitively, predicting the identity of the 7th future token becomes harder as i increases. better understand this phenomenon, we measure the confidence of the predictions by the entropy of the
'The performance metric for LM is perplexity, for QA is exact match, and for the remaining summarization tasks, the metric is ROUGE-L.
5

---

2001
2
HIBOLO
JAPILI
Figure An example of top-5 sausage lattice on NewsRoom example. Edge weights correspond to logits. Edges at each time step are ordered in descending weight and green, bolded edges correspond to candidates matching the greedy decode over the next nine tokens: ".. desktop computers with new Intel Corp processors that it.... The initial node in this graph is state 0 and the final node is 9.
token-level probability distribution for each head. In Figure 2a, we plot the normalized histogram of entropy of each head on the LAMBADA dataset. From the normalized histogram, it is clear that the entropy increases token prediction becomes more difficult as i increases.
move from first head to the last head, which agrees with our intuition that
However, we observed that the head entropy does not increase monotonically for all tasks H[] be the average entropy of head i on particular corpus, and let hmax = max(k: Vi < k, H < Hii + 1]}, be the index of the largest head such that the average entropy
function of i. Let
of each head increases monotonically to that point. We observed a strong correlation between mmax and block efficiency (Figure 2b). Heads with lower entropy (indicating more confident predictions) intuitively contribute more to efficiency. A linear regression confirms this with an R-value of 0. 0.77. This analysis suggests that the entropies of block heads could be used as a proxy for block efficiency, and thus inference latency.
5.3 Oracle top-k block efficiency
Oracle efficiency The concept of oracle block ef- ficiency serves as a theoretical benchmark, illus- trating the headroom available from improving the quality of the block draft. To compute oracle block efficiency, we consider the top-k most probable tokens at each head, and form a "sausage" lattice from these. This data structure is weighted di- rected graph, which succinctly represents all pos- sible drafts (and their score under the blockwise parallel LM) that could be formed from selecting one of k tokens from each of the h heads (Figure 3). In the automatic speech recognition and machine translation communities, is known as a "confusion network" [26, 41].
Question: Who the best soccer player in the world? Answer:
Linnel
Ronaido
Messi
hast
beus
Christiono
Messi
Messi
the
the
world
in
Michael
Riche
who
who
Hie
wewrle
Accepted
Refected
Figure 4: Illustration of the output through ora- cle selection. For a given top k: tokens of 3, if we can choose the oracle path successfully, the block efficiency can be improved from to 5.
Given the top-k lattice each decoding step, identify an oracle path that represents the path through the lattice that maximizes the length of the accepted prefix. This exercise, as shown in Figure 4, gives insight into how much headroom exists in improving block drafts.
Potential headroom from oracle selection Oracle drafting is not practical, but rather a reference point. Analyzing the gap between actual BPD performance and the oracle upper bound (Figure 5) helps us to understand the limitations the original block drafts and potential areas for improvement. Additionally, exploring oracle efficiency as a function of the k in the top-k lattice, demonstrates how "close" the block draft was to producing a stronger draft.
6 Lattice rescoring with lightweight rescorers
Having explored the properties of block draft predictions, we propose two drafting algorithms to improve block efficiency through rescoring of the top-k lattice with lightweight auxiliary LMs. This section presents techniques for rescoring the top-k lattice along with empirical results.
Each of these algorithms is a modification of the block drafted in Stage in Algorithm 1. Instcad of using the most likely token each head the prediction, we construct the top-k sausage lattice of likely drafts from each head, where the set of top-k tokens is denoted as S; for head i. This approach allows any token within S to be chosen for position i, yielding total possible combinations of: |St] x |Szl x ... |SAl = hh!
'The number of combinations can be reduced to h-1 by using the fact that the first head is the same as the base LM and hence can set [S1]

---

Haodo
Hoedo Helid)
Deci O Hoado
GISC-12 Haado
Orece la Haado
(a) LAMBADA
(b) SQuAD VI
(c) CNN/Daily
Otedt 13 Housnn
2.40
Ovadt Houes
Oruda Hescil
Crada HHNUE)
60
10
(d) SAMSUM
(c) MultiNews
(1) XSUM
Figure 5: Oracle block efficiency over the top-k lattice function k. Each plot (a-f) represents
different task, demonstrating the relative improvement in block efficiency the oracle draft with respect to the standard block draft as a function of the number of block draft heads used.
In this lattice, any path from the start to final state represents viable draft. Two algorithms are proposed select small number of h-length drafts from this lattice, which are then passed the verification step. The first algorithm employs neural autoregressive transformers (Section 6.1), while the second utilizes n- n-gram language models (Section 6.2).
6.1 Local rescoring via neural models
simple approach uses a small neural rescorer LM, interpolating between the logits of the rescorer LM and vanilla block draft logits with an interpolation weight (Algorithm 2). Recall that z/ is the vector of logits corresponding the juh block. Let Sj denote the set of symbols with top-k values in the logits vector 2/. The rescored prediction for head j is given by:
Algorithm 2: Local rescoring via neural models
input Blockwise parallel LM Mý. top-k indices selection function ToP-k(-), rescoring model Ma, interpolation weight >> 0.
1: 2 MB, Uge), <
2: S; Top-k(z)). Vi S
3: D14+ + arg max sev = vi Si
4: 1a Local lattice rescoring */
5: for 2,..., in parallel do
< =|[s,] + [St)]S,),
6: Tt+j -Ma (- USL Ie. 0+-g-11
7:
[[] 2/[S,] + Petj]S]]
where a is the weight placed on the rescorer's 's prediction and rt4. are the corresponding log- its predicted by the small neural rescoring
8:
2 #[[s]] <- x
9: end for
model, when conditioned on the sequence ys*,, Ut+1. 1....: 0142, j- We also set logits for symbols outside set S, (S5) to be negative infinity, which corresponds zero probability. Note that we do not rescore the first head as is the same as the base LM. We then run Algorithm 1, where instead of using logits directly from the BPD model, We experiment with decoder-only transformers having 32, 61, and 94 million (M) weight parameters (Appendix D).
the rescored logits to generate the draft.
6.2 Global z-gram rescoring
We also evaluate the quality of drafts generated by rescoring with n-gram LM. Recall that blockwise parallel LMs can used to compute lattice representing h possible sequences. We rescore all of these sequences except the first position token with an и-gram model, select the top p most likely sequences and pass them to the verification stage. When p = 1, we refer to this as n-gram rescoring and when 1, refer this p-best n-gram BPD.
While global rescoring typically yields better results compared to local rescoring, rescoring 1 sequences with neural LM and selecting the most likely sequence would take time OA*)), which computationally prohibitive most cases. Hence, we take advantage n-gram LMSs which are unique in that can efficiently select the most likely rescored sequence in time poly(k, h), using

---

Table 4: Block efficiency of rescoring methods the top-16 lattice. '16-best O-gram BPD' indicates performance of 16-best best draft verification over the original lattice without n- n-gram rescoring. Relative percent improvement over BPD (Baseline) is indicated in parentheses. Green circles improvement over the Baseline, while red circles denote no improvement.
indicate
Baseline BPD
Local rescoring
Glohal rescorinng 16-best O-gram BPD
Oracle
neural 61M BPD
4-grum BPD
16-best 4-gram BPD
LM
LAMBADA
3.12
(-1.28%)
3.05 (-2.24%)
3.23 (43.5356)
3.29 (+5.45%)
3.67
QA
SQuAD V1
2.08
2.10 2.10 (+0.96%)
2.07 (-0.48%)
2.18 (+4.85%)
2.22 (+6.87%)
2.45
CNN/Daill SAMSUM
1.74
1.73 (-0.57%)
1.73 (-0.57%)
1.82 (+1.666%)
1.83 (+5.41%)
2.26
S-SUM
1.27
1.39 (+9.45%)
1.29 (+1.57%
1.37 (+7.87%)
1,45 (+14.17%)
1.95
MattNews
1.10
1.25 (+13.64%)
1.12 (+1.8%)
1.13 (+2.739%
1.22 (+10.91%)
1.43
L-SUM
XSUM
1.13
(5%)
1.66 (+2.65%
1.18 (+4.42%)
(+11.50%)
NewsRoom
1.08
1.29 (+19.449)
1.18 (+9.26%)
1.11 (+2.789)
1.31 (+21.30%%
1.50
Bock Bockerddro Berde
Book eerd
1.40
4.rmm
p Sequeences
Top Sequencea
Toe p Sequences (c) CNN/Daily
(a) LAMBADA
(b) SQuAD VI
Lgan gaa
Block Effchencc
Block Efficency
Block Emitency
2.15
1.0
The p Seauences (d) SAMSUM
The p Soquences
The Soquences
(e) MultiNews
(I) XSUM
Figure Block efficiency of p-best n-gram BPD methods as a function of the number of top p sequences verified in parallel. The block efficiency of the methods is evaluated with the same number of paths extracted from the top-16 lattice.
dynamic programming. We use the OpenFST library [1] to represent each n-gram LM as a weighted finite state automaton and apply finite state composition with the top-k lattice followed by extraction of the most likely draft sequences. Training details for n-gram LMs are in Appendix D.3.
6.3 Empirical evaluation
Block efficiency Table 4 and Figure demonstrate the impact of lattice rescoring on block efficiency across various tasks. Autoregressive neural, n-gram LM, and p-best n-gram BPD rescoring all demonstrate improvements in block efficiency, although gains are task-dependent.
High initial block efficiency (LAMBADA, CNN/Daily): Both rescoring methods show little to improvement, suggesting that vanilla BPD already produces high quality drafts.
Low initial block efficiency (SQuAD V1, SAMSUM, XSUM, NewsRoom): Both neural and n-gram augmentatiaons lead to block efficiency gains, particularly with neural LMs achieving the best performance in some cases.
Repairing repetitions In Section 5.1, we note that vanilla block drafts are prone token-level repetition and that rescoring with a simple language model reduces the incidence of this. Although rescoring reduces repetition overall in drafts, is this driving improvements in block efficiency? To answer this, we com- pared the drafts generated by greedy rescor- ing with the 61M parameter neural rescorer against vanilla drafts. Time step instances were considered wins/ties/ossee based on the
Table 5; Wins, ties, and losses of 61M neural- rescored and vanilla drafts. "% Repair" corresponds to instances where the rescored draft eliminates rep- etition and "% Regress" corresponds to instances where the rescored draft introduces repetition.
Tie
Lass
LAMBADAA
601.5K
22.35
0.55
SQUAD V1
13.01
12.1%
7241
15.0
1116
15.77
330M
263.02
3000
34.7
633
4.82
410
10.0
Novillws
95.25
070

---

3.3
2.02
SP3
(PC-mecen
Hedua
Med mo-
2.75
1.3
252
1.02
Calogotes in MT Datanet
Categerie Nffdecch Datse
Block efficiency
(b) Speedup ratio
Figure 7: Block efficiency and speedup ratio relative to the standard autoregressive decoding on sub-categories of MT-Bench dataset [53] when greedily decoding with Vicuna 13B.
accepted prefix length of the rescored draft vs. vanilla draft. Table 5 displays the win frequency across tasks along with the percentage of wins/losses attributed to introducig/Climinaain repetition.
Note that in the tasks where rescoring improves block efficiency the most, NewsRoom and MultiNews, high percentage of those repaired instances driven by fixing erroneously repeated tokens. In fact, for MultiNews, 66.23% of block drafts are improved through repetition repair. We also evaluated the performance of rescoring with in no-ommii trained rescoring LMs, but found that they tended to perform no better than C4-trained LMs (Appendix E...
Lattice rescoring on open-source blockwise parallel LLMs
Medusa decoding [4] extends BPD by verifying a set of plausible candidates in parallel. Verification performed efficiently through a tree-attentt ate-attenioon mechanism, requiring only a single forward pass. Other aspects not explicitly mentioned remain the same as described in Algorithm 1. While Medusa employs tree-attention during decoding to efficiently verify subset of likely drafts, our approach focuses on rescoring these draft candidates, making them potentially complementary techniques. We explore this synergy by integrating neural rescoring method into Medusa decoding. In this section, we apply rescoring to large open-source L.Mss, using Vicuna 7B-v1.3 and Vicuna 13B-v1.3 as base models. We report both block efficiency and speedup ratio achieved relative to standard autoregressive decoding using the SpecBench benchmark [49]. To ensure rigorous verification, expand our experiments to include a wider range of datasets. We use existing pretrained Medusa heads as the block drafter Although these base LMs were not trained jointly with the block drafter, this corresponds to the Medusa-1 configuration, which has been shown to result in comparable speedups jointly trained Medusa models [4]. For lattice neural rescoring, we set & to be the full vocabulary size, using 5 heads with the next- woxd-word--gg wwrd--rediceeeeee prediction LM head as one of the heads, following Algorithm 2. All timings
evaluated on single NVIDIA A A100 80GB GPU with batch size 1.
Figure demonstrates the block efficiency and speedup ratio on MT-Bench [53], comparing greedy BPD and Medusa with and without local rescoring for Vicuna 13B models by setting the interpolation weight a to 1.0. The same analysis on Vicuna 7B is described in the Figure 9, which is detailed in Appendix G. A key observation is that even after increasing model size from 7B to 13B, a relatively small neural model (68M) can effectively serve as the rescoring drafter, showcasing the robustness of our approach. The rescoring model used in these experiments is decoder-only LM trained on the C4 and ShareGPT datasets?. Furthermore, we observe consistent performance improvements across both the original BPD and its extension, Medusa, further validating the efficacy of our local rescoring method. While the speedup gains might not always directly correlate with the increase in block efficiency, we consistently observe performance improvements across all categories. This difference suggests that block efficiency does not always translate into equivalent speedup, likely due to system-level factors. However, there remains potential for further acceleration through additional system-level optimizations.
Table 6 further presents speedup ratios across diverse datasets for Vicuna 7B and 13B models, respec- tively. We evaluate not only under greedy decoding (Temperature=0.0) but also under temperature sampling (Temperature=0.7, 1.0), employing typical acceptance for verification [4]. Both BPD and Medusa, enhanced with our local rescoring, consistently yield speedup improvements across all

---

Table Speedup ratio relative to the standard autoregressive decoding for Vicuna models (7B and
13B) on various datasets: MT-bench [53]. S-Sum (CNN/Daily), QA [27]. GSM8K [8], and RAG [21].
Alodd
Matbod
Tonperataon(0)
Tespentured:
Torgeralazc=1.00
AiTibacch
5-Smm
CSMBX
8.4G
S-Serr
GSMM
RANG
MT.hrach
5.8m
Q44
CSMRR
R:4
HPP
1.300
1.430
1,806
1.419
1,012
1.522
1.813
Vieco 7B
rescoding
1
1,561
1.412
1.54
1,422
Medraa [4]
2-430
2
3,055
2.317
1835
2,425
2.054
2113
2.563
1.908
2.511
3,000
2.534
3,000
2,010
* rescoving
2.482 9
2.006
2.144
2.377
20000
2,097
3,238 6
3,379
2440
2.139
2.657 0
2.281
2,366
2
2.173
1:345
1.530
1-085
1.291
1.165
1.851
1,535
1.510
1.875
1.554
2.045
1.891
1.675
31774
1866
Viceaa 138
#Lecal
1.819
1.522
1.519
1889
1.501
1.990
1443
1.761
1.99
1478
2.188
1
1
2,3
1.779
Medra [i]
2:781
2.001
1996
1995
2.000
2.322
2:300
2008
2.256
Table 7: Speedup ratio of efficient LLM inference methods during greedy decoding.
Vicuna 7B
Vicuna 13B
Method
MT-Bench
S-Sum
QA
GSMBK
RAG
M'T-Bench
S-Sum
QA
GSM8K
RAG
Sps [5]
1.432
1.394
1.417
1.364
1.568
1.417
1.424
1.362
1.448
1.606
Lookahead [13]
1.818
1.645
1.503
1.865
1.475
1.118
1.007
1.011
1.324
0.963
PLD [39]
1.676
2.707
1.162
1.605
1.909
1.528
2.384
1,050
1.646
1.876
BPD
1,780
1.509
1.489
1.696
1.409
1.745
1,530
1.488
1.794
1.483
+ Local rescoring
1.922
1.534
1.555
1,780
1.501
1.819
1.522
1.519
1.819
1.501
Medusa
2,430
2,002
2.045
2.317
1.833
2.383
2,000
1,986
2.507
1.945
Local rescoring
2,482
2.076
2,114
2.357
2.000
2.467
2.136
2.154
2.519
2.068
settings. Green circles indicate further improvements from local rescoring, while red circles denote no improvement. Notably, even with larger models, our method delivers consistent gains in latency. Table compares the speedup ratio of various efficient LLM inference methods. While other methods offer speedup in certain scenarios, their performance is inconsistent across task and decoding setting. Overall, we find that local neural rescoring consistently provides additional speedup over both BPD and Medusa decoding.
n-gram rescoring While C4 g-gram lattice rescoring yielded 1-best candidates with improved block efficiency for many tasks, the gains were not as stark locally rescoring with the Vicuna-68m model (Figure 8). The discrepancy is partly due to domain mismatch, since the C4 data used to train the n-gram rescorer differs from the data used to train the base Vicuna LLMs. Unsurprisingly, we fail to see a large improvement in block efficiency for specialized tasks such as math reasoning (GSM8K), but significant gains for tasks where generic English grammaticality is important (QA and summarization). The neural rescorer may also benefit from access to increased context.
3um
Wenbench
5.m
Vicuna
138
Figure& 8: Block efficiency for greedy BPD with n-gram top- 10 lattice rescoring. An interpolation weight of 0.2 was placed on the n-gram LM before interpolating with blockwise parallel logits.
8
Conclusion
This paper presents comprehensive analysis of BPD, highlighting its predictive dynamics and proposing methods to refine the generation of block drafts. Our study offers insights into BPD's behavior, particularly the tendency for drafts to contain consecutive repetitions and its heads to exhibit varying confidence levels in predictions. Two algorithms are proposed for generating higher quality drafts: one for local rescoring with small neural models (i.c., neural BPD) and another for global rescoring with an n- n-gram LM and generating multiple drafts (i.c., p-best n-gram BPD). These algorithms leverage the strengths of both blockwise parallel LMs and small rescoring models to reduce average decoding latency, pushing the boundaries of efficient text generation with BPD. We show that BPD lattice rescoring even complements Medusa decoding, recent extension of BPD, demonstrating further latency reduction for open-source LLMs. We believe this work points to the value in incorporating smaller LMs in improving LLM decoding speed.
10

---

Acknowledgments
We would like to thank Cyril Allauzen for discussions on finite state lattice rescoring throughout the development of this work.
References
[1] Cyril Allauzen, Michael Riley, Johan Schalkwyk, Wojciech Skut, and Mehryar Mohri. Openfst: A general and efficient weighted finite-state transducer library: (extended abstract of an invited talk). In Implementation and Application of Automata: 12th International Conference, CIAA 2007, Praque. Czech Republic, July 16-18, 2007, Revised Selected Papers 12, pages 11-23. Springer, 2007.
[2] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[4] Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10777 2024.
[5] Charlic Chen, Scbastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.013318 2023.
[6] Ethan A Chi, Julian Salazar, and Katrin Kirchhoff. Align-refine: Non-autoregressive speech recognition via iterative realignment. arXiv preprint arXiv:2010.142333 2020.
[7] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmano, et al. Palm: Scaling language modeling with pathways. arXiv preprint 2tt3i::22200021111 2022.
[8] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 10,14688, 2021.
[9] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359 2022.
[10] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3, int8 (): 8-bit matrix multiplication for transformers scale. Advances Neural Information Processing Systems, 35:30318 1--3832,, 2022.
[11] Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. arXiv preprint arXiv: 2019.
[12] Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news: large-scale multi-document summarization dataset and abstractive hierarchical model. arXiv preprint arXiv: 90..9006.7799 2019.
[13] Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of Ilm inference using lookahead decoding. arXiv preprint arXiv:2402.02055 2024.
[14] Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum cor- pus: A human-annotated dialogue dataset for abstractive summarization. arXiv preprint arXiv:1911.12237 2019.

---

[15] Max Grusky, Mor Naaman, and Yoav Artzi. Newsroom: A dataset of 1.3 million summaries with diverse extractive strategies. In Marilyn Walker, Heng Ji. and Amanda Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long Papers), pages 708-719, New Orleans, Louisiana, June 201 8. Association for Computational Linguistics.
[16] Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non- autoregressive neural machine translation. arXiv preprint arXiv: 171100221,,, 2017.
[17] Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines read and comprehend. Advances in neural information processing systems, 28, 2015.
[18] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems, 35:30016-30030, 2022.
[19] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint TTt:::99440995.11 2019.
[20] Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[21] Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprini arXiv:2004.088 049906 2020.
[22] Slava Katz. Estimation of probabilities from sparse data for the language model component of speech recognizer. IEEE transactions acoustics, speech, and signal processing, 35(3):400- 401, 1987.
[23] Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[24] Tachyeon Kim, Joonkee Kim, Gihun Lee, and Se- Young Yun. Instructive decoding: Instruction- tuned large language models are self-refiner from noisy instructions. The Twelfth International Conference on Learning Representations, 2024.
[25] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04455 2020.
[26] Philipp Koehn, Hieu Hoang, Alexandra Birch, Chris Callison-Burch, Marcello Federico, Nicola Bertoldi, Brooke Cowan, Wade Shen, Christine Moran, Richard Zens, et al. Moses: Open source toolkit for statistical machine translation. In Proceedings of the 45th annual meeting of the association for computational linguistics companion volume proceedings of the demo and poster sessions, pages 177-180. Association for Computational Linguistics, 2007.
[27] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
[28) Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. International Conference on Machine Learning, pages 19274-19286. PMLR, 2023.
[29] Xiang Lisa Li, Ari Holtzman, Daniel Fried, Percy Liang, Jason Eisner, Tatsunori Hashimoto, Luke Zettlemoyer, and Mike Lewis. Contrastive decoding: Open-ended text generation as optimization. arXiv preprint arXiv:2210. 2022.
12

---

[30] Xinyin Ma, Gongfan Fang, and Xinchao Wang. Amm-rreer: On the structural pruning of large language models. arXiv preprint arXiv:2305.1162, 2023.
[31] Giovanni Monea, Armand Joulin, and Edouard Grave. Pass: Parallel speculative sampling. arXiv preprint urXiv:2311. 2023.
[32] Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization, arXiv preprint
2018.
[33] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Ruffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint 1606 06031, 2016.
[34] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1,8)9, 2019.
[35] Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.114 11446, 2021.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
[37] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 5551,, 2020.
[38] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprini arXiv: 1606.05500000 2016.
[39] Apoorv Saxena. Prompt lookup decoding, November 2023.
[40] Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456- 17472, 2022.
[41] Wade Shen, Richard Zens, Nicola Bertoldi, and Marcello Federico. The jhu workshop 2006 iwslt system. In Proceedings of the Third International Workshop on Spoken Language Translation: Evaluation Campaign, 2006.
[42] Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018.
[43] Andreas Stolcke. Entropy-based pruning of backoff language models. arXiv preprint cs/0006025, 2000.
[44] Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint mrXiv:2366.11699 2023.
[45] Ziteng Sun, Jae Hun Ro, Ahmad Beirami, and Ananda Theertha Suresh. Optimal block-level draft verification for accelerating speculative decoding. arXiv preprint arXiv:2403. 2024.
[46] Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami. Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. arXiv preprint mxXiv22310.15141 2023.
[47] Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, -aptiSee Alayrac Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805. 2023.
13

---

[48] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, al. Emergent abilities of large language models. arXiv preprint arXiv:2206.06682 2022.
[49] Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: comprehensive survey of speculative decoding. arXiv preprint arXiv:2401 0000110011, 2024.
[50] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087-38099, PMLR, 2023.
[51] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances Neural Information Processing Systems, 35:27168-221 071118227111--11 2022.
[52] Euiin Yi, Taehyeon Kim, Hongseok Jeung, Du-Seong Chang, and Se- Young Yun. Towards fast multilingual IIm inference: Speculative decoding and specialized drafters, 2024.
[53] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36, 2024.
[54] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances neural information processing systems, 34:17723-17736, 2021.
14

---

A Broader impact
Our work on BPD for language models has potential applications in latency-sensitive scenarios. Furthermore, this work suggests that LLMs may benefit from the incorporation of faster, lightweight language models, either to reduce latency or potentially to improve the quality of generated text.
B Limitation and future work
B.1
Limitation
Our current drafting heads closely follow the original design of BPD but leave room for architectural improvements. The structure of the block drafter is essential for optimizing gains from rescoring, and advanced training methods may enable the model to understand the block context effectively, bringing better alignment into the target prediction.
B.2
Future work
Our work proposes to augment a small model to improve the quality the drafts. Possible future directions include (a) combining our lattice rescoring method with alternative sampling strategies (b) scaling the blockwise parallel LM for compatibility with larger-scale LLMs (c) improving training methods for drafting heads (d) using the sequential entropy ordering of heads (Figure 2b) possible halting condition during block draft head training, or to inform how a rescoring LM should be interpolated with the block lattice weights.
C Related work
C.1 Efficient transformer inference
Works on improving transformer efficiency encompass both optimization of existing of model weights, or fundamental change to the model architecture. Examples of the former include tech- niques such as quantization [50, 51, 10] and model pruning [44, 30]. In parallel, neural architecture search has played crucial role in identifying network structures that balance performance with efficiency [25, 54]. Relatedly, Elbayad et [11] propose early-exiting at intermediate layers for faster inference, while Schuster et al. [40] explore confidence thresholding for balancing speed and accuracy. These methods offer insights into optimizing decoding under resource constraints.
One important line of work has focused on modifying the decoding method in LMs. The adoption of non-autoregressive (parallel) decoding strategies [42, 16] marks pivotal shift in this domain, addressing inference latency by simultaneously generating multiple tokens. Subsequent innovations have sought to refine this approach by incorporating additional context [6], iterative refinement [23], and tree- based attention mechanism [4]. However, these refinements often require complex training or additional inference data.
C.2 Efficient autoregressive decoding
There are several recent works that improve the speed of LLM decoding, including pioneering works like BPD and speculative decoding. Speculative decoding leverages a smaller 'draft' model to anticipate the outputs of larger target model, improving average decode latency without loss in generation quality [28, 5, 23, 45, 52]. The draft model is typically trained on the same corpus as the LLM, thus autoregressively generates similar drafts as the target model with reduced latency. Speculative decoding is most successful when a long sequence of speculated tokens accepted by the target LM during verification, avoiding multiple serial calls to the target LM to generate the same sequence.
the surface, contrastive decoding algorithms share some similarities with our proposed draft rescoring approach, insofar weaker model is used to modify the predictions of the target LM [29, 24). However, in this work, we refine block drafts solely to improve latency. Like speculative decoding, our proposals have effeect the quality of the target LM's generated text.
15

---

D Experiment details
D.I Training objective for blockwise parallel LMs
minimized the following loss function train blockwise parallel LMs:
LBPD
An Lr,
h=1
where is the number of heads, A is a non-negative scalar that weights the loss from head h, and Ca denotes the loss for each individual head:
Ch
Σ
EL....BEIA
where 21... the token sequence up position i, Me+h the ground truth token at position + h, and p( is the probability of observing token yi+h given the sequence #1... under the blockwise parallel L.M. We trained models in this work with λ We leave tuning these hyperparameters, improving the block efficiency and quality of the blockwise parallel LM, as future work.
D.2 Neural model details
Table 8: Architecture hyperparameters for each of the transformer- based neural language models.
Type
Model
Layers
Embedding Dim
Hidden Dim
Blockwise Parallel Decoder
1.5B
1,536
12,288
32M
384
1,536
Autoregressive Decoder
61M
12
384
1,536
94M
768
3,072
Each neural rescoring LM is a decoder-only transformer with learned absolute positional embeddings and twelve self-attention heads at each layer. The key architecture hyperparameters given in Table 8. Aside from scale, the only difference between the blockwise parallel [.M and neural rescoring models is the addition of the feedforward neural networks and eight additional block prediction heads. Note that the number of parameters for each of these models also includes the embedding table.
Each model was pretrained on the English C4 corpus for 200K iterations with batch size of 200 x 1M tokens per batch. Dropout was not applied. For the blockwise parallel LM, all heads were trained jointly. The pretraining for the blockwise parallel LMs took about 47 hours 128 TPUv3 units.
For downstream tasks, models were finetuned for maximum 100K iterations with batch size of two examples with maximum sequence length of 2048. Maximum learning rate was fixed to 10-1 for all runs, with a cosine learning rate schedule. Checkpoints were selected based on heldout set model performance. Interpolation weight for all rescoring models was tuned for block efficiency on 100 randomly selected examples from the evaluation set for each task, and performance reported on remainder the evaluation set.
D.3 n-gram details
All it-gram LMs in this work are Katz backoff n n-gram LMs [22] fit on the train split of the GPT3 heoword-tokenizer English C4 corpus with n-gram order [2, 4). We apply entropy pruning [43] to reduce model size to maximum of 100 million n-grams per model, and ensure that each trigram is observed at least twice and each 4-gram is observed least four times. Preprocessing of the text identical to that used to train neural LMs.
16

---

D.4 Datasets
LAMBADA (LAnguage Modeling Broadened Account for Discourse Aspects): A collection of narrative passages designed to test the understanding of long-range dependen- cies in language models, where the task involves predicting the last word of a passage based on the full context [33].
SQuAD V1 (Stanford Question Answering Dataset): reading comprehension dataset that features questions based Wikipedia articles, with answers located within the text [38].
summaries, mainly used to evaluate the summarization capabilities of language models, particularly in abstractive summarization [17].
CNN/Daill Mail:
This dataset includes news articles paired with
SAMSum (Semi-Automatie Machine Summarization). Focuses abstractive summa- rization using news articles and machine -eeerreee summaries, testing models' abilities to refine and improve existing summaries [14].
MultiNews: Comprises news articles from diverse sources for abstractive summarization tasks, evaluating models on handling different writing styles and topics [12].
XSUM: Contains scientific documents and summaries, challenging language models to
process complex scientific information and language [32].
NewsRoom: dataset of news articles aimed at assessing the factual accuracy and informa- tion extraction capabilities of models in generating summaries [15].
All datasets were tokenized using the 50,257 GPT3 subword vocabulary [3].
Templates
We used the following prompts during model finetuning and inference.
SQuAD: "question: [question] context: [context]"
CNN/DailyMaill: "summarize: [text]"
SAMSum: "Here is dialogue: [text] ( Write a short summary!"
MultiNews: "Write summary based this article: [text]"
XSUM: "Summarize: [text]"
NewsRoom: "Please write a short summary for the following article: [title] [text]"
E
Rescoring with in-domain language models
Table 9: Block efficiency from rescoring with in-domain trained rescoring models for 2-gram and
61M parameter neural rescorer.
2-gram
neural-61M
Dataset
C4
In-domain
C4
In-domain
SQuAD V1
2.09
2.04
2.10
2.06
CNN/Daily
1.73
1.73
1.73
1.72
SAMSUM
1.31
1.22
1.39
1.24
MultiNews
1.13
1.14
1.25
1.16
XSUM
1.17
1.18
1.23
1.14
NewsRoom
1.20
1.22
1.29
1.11
We found that in-domain rescorers performed no better than rescorers only trained C4. We suspect this is due to lack of sufficient finetuning data and that the main benefit of rescoring comes from discouraging unnatural artifacts such repetition from the original BPD draft. Table 9 shows block efficiency after rescoring using in-domain models for all tasks besides language modeling.
Neural rescorers were finetuned from C4-pretrained checkpoints. n-gram models were trained from scratch, and unseen vocabulary was added as unigram ares with trivial weight (negative log probability of 1000.0). This was done to ensure that all paths through the lattice were assigned
17

---

non-zero probability by the n-gram model. We also tried interpolating the in-domain n- n-gram model with a unigram model trained on C4, and observed similar performance as simply adding unseen unigrams.
F
Interpolation weights tuned per task
We tuned the interpolation weight, for the 94M parameter neural and 4-gram LM rescorers, and then used this weight to rescore with all other models of that same class. 100 examples from each task's heldout were set aside for tuning, to maximize block efficiency, The remainder of examples were used for evaluation. We swept over a E {0.1, 0.5, 0.75, 0.9, 1.0, 1.1, 1.5, 2.0, 5.0, 10.0}.
Note that for tasks where lattice rescoring much higher weight on the block draft logits (Table 10). This signal that the rescorer does not provide additional information over the original block draft heads.
unhelpful, the interpolation weight, a is tuned to place
Table 10: Tuned interpolation weight per task for neural and n-gram rescoring.
Dataset
Neural
n-gram
LAMBADA
0.1
0.1
SQuAD VI
1.0
0.75
SAMSum
5.0
1.5
CNN/Daily
0.1
0.1
MultiNews
5.0
2.0
XSUM
1.5
1.1
NewsRoom
5.0
2.0
G Local rescoring impact on block efficiency
Table 11 reveals the impact of different rescoring methods on the block efficiency of the block lattice, offering insights into their effectiveness across diverse tasks and models, supporting the investigations in Section 6.1.
Limited improvement for high baselines: For tasks with already high initial block effi- ciency (LAMBADA, CNN/DailyMail,,, rescoring offers minimal or even negative changes in block efficiency compared the baseline BPD system. This suggests that for tasks where standard BPD already achieves significant speed improvements, there is limited room for further gains through rescoring.
Efficacy for poor baselines: In tasks with lower initial block efficiency (SQuAD V1, XSUM, NewsRoom), rescoring using both n-gram and neural language models results in increased block efficiency. Notably, neural rescoring with larger models (61M and 94M parameters) achieves the highest efficiency gains in these tasks, reaching up to 19.44% improvement in NewsRoom. These results highlight the potential of rescoring refine predictions and enhance efficiency for models exhibiting calibration issues.
Task-specific effectiveness: The level of improvement from rescoring varies across different summarization tasks (MultiNews, XSUM, NewsRoom). While all show positive responses, NewsRoom exhibits the largest gains, suggesting that the effectiveness of rescoring is task-dependent.
Comparison with oracle efficiency: The 'Oracle' columns present the upper bound achiev- able only the most likely token at each step chosen with perfect hindsight (k=2 and k=16). While significant gaps remain between current results and the oracle, the observed improvements from rescoring demonstrate progress towards closing this efficiency gap.
Overall, these findings suggest that local rescoring methods can valuable tool for enhancing BPD efficiency, particularly for models with less calibrated predictions. Further exploration of advanced rescoring strategies, especially conjunction with larger neural language models, holds promise for achieving even cooscceto-oracle efficiency levels.
18

---

Table 11: Block efficiency after rescoring of the block lattice. Green circles indicate improvement over the Baseline (BPD), with the percentage changes in block efficiency shown brackets relative to the Baseline. Red circles
denote no improvement.
Thazur
Boadite
Z-gman UP12
Listr sescotnn
J-ymen
4-man BMP
noma--22 BPD
Local reeaenng
ncurdl
nomai/900 BPD
Gracke (kale))
2.5622240
(2.24%)
(-12990)
Q4
SQUAD VI
208
2:00 (e0.68%))
243 (0007%)
2.01 (-048%0)
3:00 (40009)
2.00 (-8.9%%
200 (4.4890
2.66
2443
1:21--8-7700
MuliNeno
L10
1.15 (1.779)
L13 7393
(122)
(3,66))
125 (13
(1299)
1.3
143
1-SLIM
XSUM
1.15
1-17 (e1-5850
1.77
1.8
NewoRoon
108
(110111))
(19.2)))
1.35 (19.2))
(1.4444)))
L17
1,1
1.0
3.3
3F3
BrD-
Muduse
Nodzaertion
3.3
2.0
Gdugatis in Mr-Bashh Doinst
Calagrria Middrrch culast
Block efficiency
(b) Speedup ratio
Figure 9: Block efficiency and speedup ratio relative to the standard autoregressive decoding on sub-categories of MT-Bench dataset [53] when greedily decoding with Vicuna 7B.
For the evaluation of inference time on open-sourced Vicuan 7B model, we provide Figure 9 for the block efficiency and speedup ratio MT-Bench dataset which is parallel to Figure 7 in Section 7.
H Ablation on the number of heads in the blockwise parallel LM
Table 12 summarizes the block efficiency for different head configurations across various language tasks with the same settings discussed in Figure 1.
Both performance and block efficiency tend to increase with the number of heads, up to point. This suggests that using more heads allows the model to capture richer contextual information and make more accurate predictions.
General trend:
Efficiency trade-off: While increasing heads generally improves block efficiency, it also increases the memory for verification stages. Therefore, the optimal number of heads depends on the balance between desired block efficiency and available resources.
Table 12: Test performance task. Test performance of each finetuned model and block efficiency are shown as a function of heads (h E 3, 6, 9). Tasks inclue Language Modeling (LM), extractive Question Answering (QA), and both Long and Short Summarization (L-Sum & S-Sum). The metric for LM is perplexity, for QA is exact match, and for all the remaining (summarization) tasks, the metric is ROUGE-L.
# Heads (f)
Task
Datasct
Performance
LM
LAMBADA
7.88
1.79
2.84
3.12
QA
SQuAD
57.60
1.53
2.03
2.08
CNN/Daily
39.85
1.60
1.71
1.74
S-SUM
SAMSUM
37.66
1.18
1.25
1.27
MultiNews
23.08
1.10
L-SUM
XSUM
52.15
1.11
1.12
1.13
NewsRoom
39.85
1.07
1.08
1.08
19

---

I Practical efficiency of rescoring block drafts
To enhance our understanding of block rescoring within the realm of contemporary deep learning hardware environments, we present an in-depth examination focused on TPU/GPU utilization and the overhead incurred by n-gram rescoring. This analysis is divided into two parts: (1) an analysis of block rescoring through the lens of TPU/GPU utilization, and (2) empirical benchmarks of n- n-gram lattice rescoring. The major takeaways are as follows.
Memory bandwidth (HBM + SRAM) A critical factor in the performance of deep learning applications is the efficient management of memory bandwidth between High Bandwidth Memory (HBM) and Static Random Access Memory (SRAM) [9]. Increasing the block efficiency via the block lattice rescoring reduces the average per token parameter and key-value cache I/O that needs to be communicated from HBM to SRAM.
Overhead in n-gram rescoring æt-gram rescoring actually quite efficient. For the size of lattices we consider in this work, moving the lattice from HBM to DRAM, performing n-best n--gram rescoring, and moving the n-best paths back to HBM requires no more than 2 ms per lattice.
L1 Hardware utilization
We compare our approach against traditional Autoregressive LMs across several metrics (Table 13).
Table 13: Comparative analysis of per decoded token efficiency metries across block rescoring meth-
ods and the standard autoregressive LM (batch size=1). 1). This table shows the average block efficiency, parameter I/O, key-value (KV) cache I/O at varying sequence lengths, and FLOPS evaluated per- per-token basis with batch size 1.
Component
Aulonegeasive
4-gram Neural-61M
16-best O-gram
4-gram
Avg. Block Elliciency
1,000
1.646
1.657
1.724
1.717
1.797
Parameter IO (GB)
3,000
1.823
1.511
1.811
1.747
1.669
KY Cache U/ (GB) ke 128
0,113
0.074
0.073
0.076
0,140
0.134
Cache (GB)
0,453
0,290
KY Cache LOO (GB) kcn 1024
0.906
0.555
0.552
0.574
0.602
0.575
KV Cadhe 1o (CB)-Scu) S5q. 2048
1.812
1.106
1.098
1.144
1.129
1.079
FLOPS (T)
0.931
0.57
0.567
0.635
0.621
0.593
Memory bandwidth and compute efficiency The block rescoring variants achieve significant reductions Parameter I/O and KV Cache I/O compared to autoregressive decoding, suggesting BPD) methods' ability reducing inference times by mitigating the primary latency bottleneck- memory bandwidth.
Comparative latency impact A consistent decrease in memory bandwidth utilization across block- wise parallel LMs, including those leveraging LM rescoring and parallel processing strategies, illustrates our approach's contribution to accelerating inference speed. This underscores the practi- cality and applicability of our enhancements in promoting more efficient language model inference within state-of-the-an computational frameworks.
L.2 Overhead of n-gram rescoring
While the majority computational efforts in block rescoring are dedicated to TPU/GPU utiliza- tion, the implementation of n- n-gram rescoring introduces additional overheads. These primarily attributed to CPU computations and the data transfer between the CPU and HBM. This section provides comprehensive examination of these overheads, drawing benchmarks from rescoring experiments with a 4-gram C4 L.M.
Benchmarks for 4-gram C4 LM rescoring We conducted benchmarks rescoring lattices with a 4-gram C4 LM of x 100M n n-grams. The average latency observed across 10 runs for different numbers of the shortest paths is summarized in Table 14.
Notably, rescoring with a large 4-gram I.M averages less than milliseconds for extracting up to 16 beet paths, despite the lattice containing approximately 4.29 billion possible paths. In
20

---

Table 14: Average latency for N-best rescoring an 8-time step lattice with 16 ares per time step. N, the number of shortest paths, is varied from to 16. # Shortest Paths
N-best Rescoring Latency (ms)
1.630
2
1.751
1.878
1.871
16
1.983
initial experiments, increasing the size of the n-gram LM had little effect on n-best rescoring latency. indicating that improvements to rescoring LM quality will incur little additional latency, provided that the rescoring LM fits within DRAM.
Latency predominantly influenced by lattice size, particularly the number of top-k tokens per time step and the number of time steps, depicted in Table 15.
Table 15: I-best rescoring latency by the 4-gram C4 LM for varying lattice sizes.
Number of time steps
Top-k per time step
1-best rescoring latency (ms)
1.038
4
1.050
1.130
4
16
1.237
8
1.061
8
1.144
8
1.234
8
16
1.630
16
1.102
16
1.206
16
8
1.558
16
16
2.174
The benchmarks highlight the fact that the additional overhead introduced by n-gram rescoring, though present, should significantly impact overall latency.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims are provided in Section and Section 2.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are described in Appendix B.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
22

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Ys]
Appendix describes the experimental details both the blockwise parallel LMs used in this paper and proposed rescoring methods.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
23

---

Answer: [No]
Justification: We have not provided open access to the source code used in our experiments. However, we have detailed the data access, architecture, and training processes in the supplemental material to enable replication of our results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. co/
www.iissssssoccoooooooooooosllll for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips.cc/ wttpcccccppppcccececcceccccccccc hwwpbligguuide/////eee/eee htte://ooeSuummieeee wwriinsssoooooooooo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Appendix D provides clear experimental details. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Our paper does not report error bars or statistical tests, as focuses on qualita- tive evaluations and proof-of-conced demonstrations, We provide detailed descriptions of our experiments to support the claims and emphasize practical observations.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources are detailed Section 4 and Section 6.1, including TPU types (TPUv3/TPUv4), training time, batch sizes and use of JAX.
Guidelines
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http:///eeur..peeeeeee wtcppp/lc////eeeeeeeeeeeeeeeeeet Ceallli//Ehhistttteee eehiideueeellmm ines?
Answer: [Yes]
Justification: Our research does not involve human subjects and primarily presents no direct ethical concerns. The datasets with CC-BY 4.0 are used for evaluation. We discuss potential societal impact in Appendix A.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer:
[Ye]]
Justification:
We discuss societal impacts of our research in Appendix A.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
25

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We believe models do not present more risk than the public models compared them to, as they show similar effectiveness but at a lower cost of operation.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The datasets having CC-BY 4.0 license are used for training and evaluation.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
26

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com con/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector:
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper presents no potential risks for IRB approvals or equivalent for research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
27

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "Y1BJqvcgev", "reviewer": "Reviewer_FLJ6", "summary": "This paper provides an improved solution for block-drafting, which is a potential useful way to improve the inference efficiency of LLMs. The work begins with observations of the problems of current block drafting, reveals that the consecutive repetition and drafting confidence are related to the quality of the draft. Rescoring methods are employed to improve the drafting process accordingly.", "questions": "See the weakness part.", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The observations are persuasive and the solution is intuitive.\n\nThe experiments show impressive improvement (up to 20%) of block efficiency, which is potentially useful.", "weaknesses": "The rescoring phase uses yet another model to scoring the candidate. I am wondering whether using different models affect the final performance. Because different models have different token generation distributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback. We understand your concern regarding the use of different models for the rescoring phase and its potential impact on final performance due to varying token generation distributions.\n\n- **Note:** Before reading further, **we kindly ask you to check the `Author Rebuttal by Authors` and `the attached PDF`** for detailed explanations and additional results.\n\n# W1. Additional experiments on different models\n\nTo address this, we conducted additional experiments demonstrating that local neural rescoring methods are robust and consistently improve performance across different target models (Instruction-tuned openLLM Vicuna `7B` and `13B` with `4 block draft heads`). Specifically, we have shown that:\n\n1. Our approach significantly speeds up BPD even further for existing, open-sourced, instruction-tuned LLMs.\n\n2. Local rescoring also accelerates decoding  for different model architectures, including the very recent Medusa model [1] presented at ICML 2024, which employs tree-attention to enhance performance in blockwise parallel LMs, **achieving an additional speedup of over `~20%` relative to BPD/Medusa and `~300% speed-up` relative to vanilla decoding**.\n\nIn the `attached PDF`, we provide detailed results showing consistent performance improvements when our model is applied on top of Medusa decoding. These results highlight the versatility and robustness of our rescoring approach in the field regarding multi-token predictions, confirming that it effectively enhances performance regardless of the underlying model used for token generation.\n\n\nWe hope this addresses your concern and illustrates the robustness and effectiveness of our proposed methods.\n\n[1] Tianle Cai, et al. \"Medusa: Simple LLM inference acceleration framework with multiple decoding heads\", ICML, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yyINAUmzbd", "reviewer": "Reviewer_rNoJ", "summary": "This paper proposes new ways to improve blockwise parallel decoding (BPD), a method to reduce inference latency in large language models. It first analyzes the token distributions produced by multiple prediction heads and then leverages this analysis to develop algorithms to improve BPD inference speed by refining the block drafts using n-gram and language models.", "questions": "Please address the questions listed in the weakness section.", "limitations": "The paper has discussed the limitations sufficiently.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The paper thoroughly studies BPD's behavior, including issues like consecutive token repetition and varying confidence levels across prediction heads, providing new insights for efficiency improvement. It further introduces the oracle top-k block efficiency as a useful metric for understanding the potential headroom for improvement in block drafts.\n\n+ The proposed refinement algorithms (local neural rescoring and global n-gram rescoring) demonstrate improvements in block efficiency across multiple tasks, with gains of up to 21.30% in some cases.\n\n+ The evaluation considers a variety of tasks (language modeling, question answering, and summarization) and datasets, \n\n+ The paper is well-structured and easy to follow, with helpful illustrations and examples.", "weaknesses": "+ The evaluation mainly compares the proposed improvements over the existing BPD baselines but doesn't compare them with other approaches for reducing inference latency, such as quantization or model pruning. It's suggested to make a more thorough comparison.\n\n+ The evaluation is conducted mainly on a 1.5B parameter model. It's unclear how well these findings and improvements would generalize to larger, state-of-the-art models. Further, while the paper focuses on improving block efficiency, there's limited discussion on the additional computational cost of the rescoring methods. It's unclear whether the efficiency gains outweigh any increased computational requirements.\n\n+ The improvements in block efficiency vary significantly across tasks, with some showing little to no improvement. A deeper analysis of why certain tasks benefit more than others would be valuable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback in helping refine our work.\n\n- **Note:** Before reading further, **we kindly ask you to check the `Author Rebuttal by Authors` and `the attached PDF`** for detailed explanations and additional results.\n\n# W1. Comparison with Other Latency Reduction Approaches\n\nWe appreciate the suggestion to include comparisons with other latency reduction approaches, such as quantization or model pruning. Our paper's primary focus is \"Exploring and Improving Multi-Token Prediction (Block Draft)\", which aligns more closely with speculative inference methods. While speculative inference methods like ours are orthogonal to pruning and quantization, we recognize the importance of providing a thorough comparison for a comprehensive understanding.\n\nRecent works such as Medusa [1,2] (presented @ ICML 2024, last month; extended version of BPD)  are examples of speculative inference techniques similar to ours. Although our paper was submitted before many of these studies were published or open-sourced, we have now included a comparative analysis in the `Table E` (see `Author Rebuttal by Authors`) and the `attached pdf`. This analysis demonstrates that **our method consistently enhances the efficiency of blockwise parallel LMs (BPD/Medusa), outperforming other speculative inference approaches (achieving an additional speedup of over `~20%` and `~300% speed-up` relative to vanilla decoding).** This inclusion aims to provide a more comprehensive understanding of our method's effectiveness relative to other approaches in the field.\n\n[1] Tianle Cai, et al. \"Medusa: Simple LLM inference acceleration framework with multiple decoding heads.\" ICML 2024.\n\n[2] Ankner, Zachary, et al. \"Hydra: Sequentially-dependent draft heads for medusa decoding.\" arXiv (2024). \n\n# W2. Generalization to Larger Models and Computational Cost\n\nWe understand the need to evaluate our methods on larger, state-of-the-art models and to discuss the computational costs associated with our rescoring methods. To address this, we have conducted additional experiments with larger models, specifically Vicuna `7B` and `13B` with `4 block draft heads`, as detailed in the `attached PDF`. These results demonstrate that our approach scales effectively, maintaining its performance gains across different model sizes.\n\n**Regarding computational costs, we would like to emphasize that we've already covered a detailed discussion in `Appendix H`**. `Appendix H` provides a detailed discussion on LLM inference speed improvements, covering low-level aspects such as `KV-Cache` and `Parameter I/O` `memory bandwidth` in TPU/GPU experiments. Our findings show that the efficiency gains from our rescoring methods outweigh the increased computational requirements. The `attached PDF` supports this statement in terms of the wall clock speedups relative to the standard autoregressive decoding. \n\n# W3. Variability in Block Efficiency Improvements\n\nWe have extended our experiments to include a wider range of datasets and tasks, as detailed in the `attached PDF`. While most tasks show significant improvements, we acknowledge that not all tasks benefit equally.\n\nAchieving consistent improvements across all tasks and datasets is a challenge faced by many LLM and LLM studies, **not just our own**. Various factors, including the characteristics of specific datasets, can influence the results. **This is not necessarily a limitation of our method but rather a common challenge in the field** [3].\n\n**However, our results demonstrate that our method consistently enhances performance, even as model sizes increase (e.g., `7B` and `13B` parameters).** The improvements are consistent for larger models, as highlighted in the `attached PDF`.\n\n[3] Xia, Heming, et al. \"Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding.\" ACL Findings (2024).\n\nWe hope this mitigates your concern and demonstrates the robustness and effectiveness of our proposed methods.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ghii7j9z7A", "reviewer": "Reviewer_vMaK", "summary": "The paper analyzes the block drafts generated by multiple independent prediction heads of blockwise parallel language models and observes three key features: consecutive repetitions, confidence of different heads, and efficiency gap with oracle top-k block. To address these issues, the paper proposes two algorithms to leverage the top-k predictions at each head: local rescoring via (small) neural LMs and global rescoring via n-gram LMs with multi-drafts. Experimental results show that the proposed two algorithms can improve block efficiency.", "questions": "1. It is not clear how the observations in Section 6.2 and 6.3 contribute to the design of the algorithms.\n2. The neural draft is generated based on the logits of the original prediction, not the rescored ones. If the rescored top-1 token is different the original top-1 token, the following rescoring will be incorrect.", "limitations": "The authors have adequately addressed the limitations", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper identifies a weakness in the existing blockwise parallel decoding algorithms: the predictions are made independently. The paper proposes two algorithms: local rescoring via neural models and global n-gram rescoring to address the weakness.\n2. The paper analyzes the block drafts and gives several observations. The observation of strong correlation between the index of the largest head such that the average entropy of each head increases monotonically to that point and block efficiency is especially interesting.", "weaknesses": "1. The experiments are only conducted on a 1.5B LM pretrained on 200B tokens of C4, without the alignment stage. This is far from the common practice in current LLMs. The authors should consider adding results on more open LLMs with different sizes.\n2. The medusa paper [1] already proposes using top-k predictions for different heads. The contribution of the paper mainly focuses on the two rescoring algorithms. However, the algorithm of local rescoring with a small LM is very similar to speculative decoding with a small LM [2]. The contribution of the global rescoring algorithm with n-gram models is not sufficient for a NeurIPS paper.\n\n[1] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. http://arxiv.org/abs/2401.10774\n[2] Fast Inference from Transformers via Speculative Decoding. https://arxiv.org/abs/2211.17192", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback in helping refine our work.\n- **Note:** Before reading further, **we kindly ask you to check the `Author Rebuttal by Authors` and `the attached PDF`** for detailed explanations and additional results.\n\n# W1. Limited Experimental Scope\nWe acknowledge the limitation of our initial experiments. To address this, we include additional results on more open LLMs with different sizes and more datasets. Specifically, we conducted experiments with Vicuna `7B` and `13B` with `4 block draft heads` in the `attached PDF`.  **These new results also show that our approach scales effectively and consistently improves performance across different model sizes, achieving an additional speedup of over `~20%` and `~300% speedup` relative to vanilla decoding**. Here, the blockwise parallel LMs extended from Vicuna `7B` and `13B` are trained with alignment stages.\n\n# W2. Difference with Existing Work\nWe believe our contributions, including rescoring methods as well as findings, are both valuable and innovative compared to existing works. Here’s why:\n\n### **Medusa**\n- **Similarity:** Both approaches explore potential gains from TopK predictions of blockwise parallel LM heads.\n- **Technical Differences:**\n    - **Medusa** searches for the best candidate in block drafts without altering the original logits\n    - **Our** method involves rescoring the block drafts to obtain the best candidate.\n- **Exploratory Differences**\n    - **Medusa** does not explicitly address the reasons behind its decoding method. \n    - **Our work** is the first to explicitly address issues such as (1) consecutive repetition, (2) confidence of different heads, and (3) top-k oracle block efficiency.\n- **Our contributions are orthogonal and different from Medusa, and our method integrated with Medusa shows even better results, detailed in the `attached PDF`**.\n\n### **Speculative Decoding**\nFirst of all, BPD (presented in 2018) is a predecessor and an instance of speculative decoding.\n- **Similarity:** Both approaches use small language models for efficient LLM inference.\n- **Technical Differences**\n    - **Speculative decoding** uses a small drafter to predict the next multiple tokens\n    - **Our method** uses the small drafter (for both local and global rescoring) to rescore the logits of block drafts with a top-k mask, which are then used for speculative inference.\n\n- **TL;DR**\n    - **Speculative decoding** uses an independent module for drafting\n    - **Our method** corrects the drafts from a dependent speculative module (e.g., BPD and Medusa), **which is totally different from the use of the small LM in speculative decoding.**\n\n### **W2-2. Regarding the Contribution of the Global Rescoring with N-gram Models**\n**We believe the global rescoring with n-gram models is valuable and innovative. Here’s why:**\n\n1. **Efficiency**: The n-gram rescoring is highly efficient, taking only `1.6 ms per lattice` via OpenFST, making it practical for real-time applications (`Table12` & `Table13` in `Appendix H`).\n2. **Effectiveness**: Despite being a classical approach, n-gram models remain effective. Our results show significant improvements in block efficiency and speedup when integrated with BPD.\n\nThe assertion that n-gram models lack novelty is not supported by our evidence. Instead, our results clearly demonstrate their practical value and effectiveness.\n\n# Q1. Clarity on How Observations Contributes to Algorithm Design\n- **Section 6.2 (Confidence across multiple heads):** This section examines the confidence levels across different prediction heads in BPD, providing insights into how well the parallel heads are trained. These observations, while not directly informing the current design of our rescoring algorithm, are included to align with our paper's goal of \"Exploring and Improving Multi-Token Prediction (block draft).\" They suggest potential future improvements, such as using variable interpolation weights based on head confidence levels. We will refine the manuscript to clarify this point.\n\n- **Section 6.3 (Top-k oracle block efficiency):** The concept of top-k oracle block efficiency serves as a theoretical upper bound for the potential improvements achievable with rescoring. By quantifying the maximum possible efficiency, we gain a benchmark against which we can measure the actual performance of our algorithms. This observation led to the realization that current methods, such as BPD and Medusa, have significant room for improvement. As a result, we developed our rescoring methods to approach this upper bound more closely, thereby optimizing block acceptance rates and reducing redundancy.\n\n# Q2. Potential Rescoring Issues\nTo clarify, the original prediction is used for the next-token prediction, but rescoring happens after that (from the 2nd position future token onward) (You can find it in `Section 7` `Algorithm2`. We will make it clearer in the camera-ready version). **The rescored top-1 token can differ from the original top-1 token, which is intentional. The original top-1 token for future positions often fails to be accepted for speculative inference, but our findings show that rescoring improves this. Evidence is provided by improvements in `block efficiency` and `speed up`.**\n\nAdditionally, BPD follows a draft-verify-accept structure. Incorrect predictions in the draft phase do not pass verification, ensuring that BPD and Medusa's predictions are identical to vanilla decoding. We focused on the number of tokens accepted, using block efficiency as a hardware-agnostic metric. Based on reviewer feedback, we have also included latency experiments in the `attached PDF`. `Appendix H` provides a detailed discussion on LLM inference speed improvements, covering low-level aspects such as `KV-Cache` and `Parameter I/O` `memory bandwidth` in TPU/GPU experiments.\n\nWe believe these responses and additional results in the `attached PDF` address your concerns comprehensively and highlight the novelty of our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We extend our gratitude to all the reviewers for providing comprehensive and thoughtful feedback on our manuscript. We appreciate your valuable insights into the strengths and areas for improvement of our work. 

# Core Contributions of Our Work
- **Novel Findings:** This work explicitly addresses key issues in blockwise parallel LMs, such as consecutive repetitions and the confidence of different heads, and introduces the concept of oracle top-k block efficiency (potential headrooms for speedup).
- **Novel Technical Contributions:** Building on BPD, our methods effectively remove repetitions and improve block efficiency, achieving up to a `21.30%` improvement. In most cases from `rebuttal experiments`, our approach also results in (a) `~15% speed-up` relative to BPD and Medusa decoding and (b) `~300% speed-up` relative to vanilla decoding.
- **Objective and Scalability:** Our main objective is “Exploring and Improving Multi-token Prediction (Block Draft)”. The framework is easy to plug in and scale up because any type of rescoring LM can be used regardless of the size of target blockwise parallel LM.
- **Oracle Top-k Block Efficiency:** To our best knowledge, this metric has not been measured in the field, but it provides a valuable upper bound indicating potential room for improvement.
- **Instruction-Tuned LLMs:** As shown in the `attached PDF`, this framework performs well even with instruction-tuned LLMs, including models with `7B` and `13B` parameters.

# Summary of Strengths Cited by Reviewers
- **Impact:** We appreciate `Reviewer FLJ6` and `Reviewer rNoJ` for noting the important motivation and impact of our work. The improvements in block efficiency (up to `21.30%`) are significant and have practical applications.
- **Soundness:** `Reviewers FLJ6` and `Reviewer rNoJ` acknowledged the technical soundness of our approach and the thorough study of BPD behavior, including issues like consecutive token repetition and varying confidence levels across prediction heads.
- **Integration:** `All reviewers` highlighted the well-structured and integrated nature of our paper, with a clear presentation and helpful illustrations.
- **Experiments:** `All reviewers` appreciated our comprehensive experiments across a variety of tasks (language modeling, question answering, and summarization) and datasets, demonstrating the robustness and effectiveness of our methods.

# Additional Experiments in PDF
1. Local rescoring improves BPD with tree-attention
2. Additional experiments with temperature sampling (`T=0.7`, `1.0`), beyond greedy decoding
3. Instruction-tuned openLLM results with Vicuna `7B` & `13B` target model with `4 block draft heads`
4. Additional tasks
   - MT-Bench (including `writing`, `roleplay`, `reasoning`, `math`, `coding`, `extraction`, `stem`, `humanities`)
   - Natural Questions (QA) [3]
   - GSM8K [5]
   - RAG [2]
6.  Local rescoring vs. vanilla Medusa decoding (extended version of BPD; presented @ ICML24) [1]

# TL;DR for PDF
- We have conducted extensive experiments with additional open-source LLMs and downstream tasks
- **Our results consistently show improved performance from local neural rescoring in terms of block efficiency as well as wall clock time, for LLMs up to `13 billion` parameters.**

### **Additional results for comparing other methods for efficient LLM inference:**
`Table E`. Speedup ratio (i.e. latency) relative to the standard autoregressive decoding.
| NVIDIA A100 GPU   | Vicuna 7B | |       |       |       | Vicuna 13B |       |       |       |       |
|-------------------|-----------|-------|-------|-------|-------|------------|-------|-------|-------|-------|
| Relative Speedup  | MT-Bench  | Sum   | QA    | GSM8K | RAG   | MT-Bench   | Sum   | QA    | GSM8K | RAG   |
|-------------------|-----------|-------|-------|-------|-------|------------|-------|-------|-------|-------|
| Sps  [6]| 1.432     | 1.394 | 1.417 | 1.364 | 1.568 | 1.417| 1.424 | 1.362 | 1.448 | 1.606 |
| Lookahead  [7]| 1.818     | 1.645 | 1.503 | 1.865 | 1.475 | 1.118| 1.007 | 1.011 | 1.324 | 0.963 |
| PLD   [8]   | 1.676     | 2.707 | 1.162 | 1.605 | 1.909 | 1.528      | 2.384 | 1.050 | 1.646 | 1.876 |
|-------------------|-----------|-------|-------|-------|-------|------------|-------|-------|-------|-------|
| BPD| 1.752     | 1.509 | 1.489 | 1.696 | 1.409 | 1.745      | 1.530 | 1.488 | 1.794 | 1.483 |
| + Our local rescoring | 1.843     | 1.534 | 1.555 | 1.780 | 1.501 | 1.819      | 1.522 | 1.519 | 1.819 | 1.501 |
|-------------------|-----------|-------|-------|-------|-------|------------|-------|-------|-------|-------|
| Medusa| 2.254     | 2.002 | 2.045 | 2.317 | 1.833 | 2.232      | 2.000 | 1.986 | 2.507 | 1.945 |
| + Our local rescoring | 2.482     | 2.076 | 2.114 | 2.357 | 2.000 | 2.467      | 2.136 | 2.154 | 2.519 | 2.068 |

## References

[1] Tianle Cai, et al. "Medusa: Simple LLM inference acceleration framework with multiple decoding heads." ICML, 2024.

[2] Vladimir Karpukhin, et al. "Dense passage retrieval for open-domain question answering." arXiv, 2020.

[3] Tom Kwiatkowski, et al. "Natural questions: a benchmark for question answering research." Transactions of the Association for Computational Linguistics, 7:453–466, 2019.

[4] Lianmin Zheng, et al. "Judging LLM-as-a-judge with MT-Bench and Chatbot Arena." NeurIPS, 2024.

[5] Cobbe, Karl, et al. "Training verifiers to solve math word problems." arXiv  (2021).

[6] Chen, Charlie, et al. "Accelerating large language model decoding with speculative sampling." arXiv (2023).

[7] Fu, Yichao, et al. "Break the sequential dependency of llm inference using lookahead decoding." arXiv (2024).

[8] Apoorv Saxena. 2023. Prompt lookup decoding

We believe these additions and clarifications address the reviewers' concerns comprehensively and strengthen our manuscript. Thank you for your constructive feedback, which has significantly contributed to the improvement of our work. We look forward to your favorable consideration.

Meta-review: The paper addresses block-drafting, an approach to improve the inference efficiency of language models. In particular, the authors proposed a new algorithm involving top-k prediction and local/global rescoring.

In general, reviewers found the approach interesting and effective in their experiments. However, the language model chosen is a little bit small compared with the current practice.
Decision: Accept (poster)