Mixture of Attentions For Speculative Decoding
Authors: Matthieu Zimmer
Milan Gritta
Gerasimos Lampouras
Haitham Bou Ammar
Jun Wang
Venue: ICLR 2025 Poster (2025)
Keywords: large language models
speculative decoding
EAGLE
Primary area: generative models
TL;DR: Using dynamical system view of LLM generation for better speculative decoding architecture

Abstract: The growth in the number of parameters of Large Language Models (LLMs) has led to a significant surge in computational requirements, making them challenging and costly to deploy.
Speculative decoding (SD) leverages smaller models to efficiently propose future tokens, which are then verified by the LLM in parallel.
Small models that utilise activations from the LLM currently achieve the fastest decoding speeds.
However, we identify several limitations of SD models including the lack of on-policyness during training and partial observability. 
To address these shortcomings, we propose a more grounded architecture for small models by introducing a Mixture of Attentions for SD.
Our novel architecture can be applied in two scenarios: a conventional single device deployment and a novel client-server deployment where the small model is hosted on a consumer device and the LLM on a server.
In a single-device scenario, we demonstrate state-of-the-art speedups improving EAGLE-2 by 9.5% and its acceptance length by 25%.
In a client-server setting, our experiments demonstrate: 1) state-of-the-art latencies with minimal calls to the server for different network conditions, and 2) in the event of a complete disconnection, our approach can maintain higher accuracy compared to other SD methods and demonstrates advantages over API calls to LLMs, which would otherwise be unable to continue the generation process.

Published as a conference paper at ICLR 2025
MIXTURE OF ATTENTIONS FOR SPECULATIVE DECODING
Matthieu Zimmer', Milan Gritta" & Gerasimos Lampouras
Huawei Noah's Ark Lab, fi name www.aae...ssttttee lastname@huawei com
Haitham Bou Ammar
Huawei Noah's Ark Lab, UCL Centre for Artificial Intelligence
Jun Wang
UCL Centre for Artificial Intelligence
ABSTRACT
The growth in the number of parameters of Large Language Models (LLMs) has led to significant surge in computational requirements, making them challeng- ing and costly to deploy. Speculative decoding (SD) leverages smaller models efficiently propose future tokens, which are then verified by the LLM par- allel. Small models that utilise activations from the L.LM currently achieve the fastest decoding speeds. However, we identify several limitations of SD models including the lack of on-policyness during training and partial observability. To address these shortcomings, we propose a more grounded architecture for small models by introducing Mixture of Attentions for SD. Our novel architecture be applied in two scenarios: a conventional single device deployment and a novel client-server deployment where the small model is hosted on consumer device and the LLM on server. In a single-device scenario, we demonstrate state-of- the-art speedups improving EAGLE-2 by 9.5% and its acceptance length by 25%. In client-server setting, our experiments demonstrate: 1) state-of-the-at laten- cies with minimal calls to the server for different network conditions, and 2) the event of complete disconnection, our approach can maintain higher accuracy compared to other SD methods and demonstrates advantages over API calls LLMs, which would otherwise be unable to continue the generation process.
INTRODUCTION
Auto-regressive inference with LLMs has become quite cost-prohibitive due the increasing pa- rameter count of recent transformer-based LLMs (Vaswani 2017). Different types of (usually or- thogonal) solutions have been proposed address this challenge, e.g. Mixture of Experts (Jacobs) et al. 1991). Flash Attention (Dao et al.] 2022), Model Quantization and Distillation (Polino et al. 2018 Linear/Sparse Self-Attention (Zhang et al. 202I), Tensor Parallelism (Shoeybi et al. 2019) and others. In this work, we focus on a recent LLMM acceleration technique called Speculative De- coding, which leverages efficient models (smaller but less capable) to draft future tokens, which are verified by the LLM (more capable but much less efficient) in parallel (Leviathan et al.) (023).
The most recent state-of-the-art SD methods, like EAGLE (Li et al. 2024b) and MEDUSA (Cai et al. 2024a), leverage activations from the LLM. However, those methods have some architectural limitations including partial observability and the lack of on-policyness. Partial observability occurs when the small (draft) model lacks complete information about the state of the LLM, leading to suboptimal predictions. The lack of on on-oolicssess during training arises because the small model is often trained under ideal conditions, assuming perfect inputs. This does not reflect the real-wor worldworld scenario where the small model generates some inputs. The longer we draft new tokens using only
'These authors contributed equally to this work
1 Corresponning authors: haitham.amma ammar@ hammar@huawei.comm huawei.com, jun.wang@cs.ucl.cc

---

Published as a conference paper at ICLR 2025
the small model, the bigger the distribution shift from the training setting. These limitations can degrade the performance and reliability of speculative decoding.
To address these challenges, propose novel architecture for speculative decoding that enhances the small model's ability to accurately predict future tokens and aligns its training more closely with the inference process. Our architecture introduces several key improvements, including Layer Self- Attention (LSA) to mitigate partial observability, Cross-Attention (CA) to improve on-policyness and training efficiency, and flexible Target Layer Inference (TLI) mechanism balance compu- tational efficiency and prediction accuracy. We evaluate our approach in the standard single-device setting where we demonstrate state-of-the of-t speedups.
Furthermore, the SD paradigm is also ideal in the following scenario a) the model size is limited by some external factor c.g. the computational capabilities of a client device, and b) we can assume access to a larger model e.g. an LLM hosted on server. Under this paradigm, the goal is to minimise server-side inference as well as maintain high accuracy in the event of a total disconnection. This is important consideration because it could pave way for serving LLMs enabling them to generate responses offline while leveraging the capabilities of the large model. To this end, we extend our methodology to a client-server scenario. In this setting, we demonstrate tatte-of-the-ars latency and minimal server calls under various network conditions (4G, 5G). Our method maintains higher accuracy in the event of disconnection, making it preferred choice over independent small models or API calls that would be unable to continue generation.
edge devices,
Contributions We introduce a Mixture of Attentions architecture for Speculative Decoding that addresses current limitations such as partial observability as well as enabling efficient (more on- policy) training while being auto-regressive. We reuse additional activations from the LLM in the small model, enabling a trade-off between drafting speed and response quality. We conduct extensive experiments demonstrate the effectiveness of our approach. Compared to EAGLE-2, 9.5% decoding speedup with a 25% higher acceptance rate in a single-device scenario and a 84% speedup with 53% higher acceptance rate in client-server scenario. Finally, we propose new framework for LLM serving in speculative client-server settings and show its effectiveness.
show
2
BACKGROUND
We first present the background knowledge required for the remainder of the paper, i.e. the decoding mechanisms of LLMs well as the drafting verification techniques that ensure correct generation.
2.1 LLM DECODING
Decoding refers to the process by which L.LMs generate tokens in response to input queries. This generation is typically done auto-regressively, where each new token yı is sampled from the LLM's distribution, conditioned on both the query and the preceding tokens yer. We explore decoding from the perspective of dynamic systems, providing foundation for developing
decoding mecha- nisms that combine large and small models (Kong et al., 2024). The internal workings of LLMs can be best understood from dynamic system perspective, which evolves as tokens generated. Given a large model Miarge, we can describe the state transition model of vanilla decoding as:
h 0++1 01+1 = ( wow.oookeeebbben soken_embed(yy)).
M++1 Softmax Coetmax(LM-hhall 01+1)), (1)
where ha represents the key and value tensors of every layer until the current time-step f. Ye is the most recent token and yt+1 is the next token, which is sampled from softmax distribution. Furthermore, token_embed is a lookup table, it assigns an embedding to particular token of the vocabulary V. LM Mead a projection from the embedding size to the vocabulary size [V]. Finally, fLarge(.) is the function aggregating all decoder layers of MLarge and o++1 is the activation of the final decoder layer. With this, the state of the dynamic system is composed of (hgr, Ve the minimal information needed sample the next token from MLange-
2.2 SPECULATIVE DECODING
Some of the earliest work on speculative decoding was introduced by Stern et al. (2018), later ex- tended to non-greedy sampling (Leviathan et al. 2023). These methods are motivated by the pro-
2

---

Published as a conference paper at ICLR 2025
hibitive cost of auto-regressive generation with MLarge that could be alleviated by using draft model Mssall that can more efficiently generate tokens that do not require the full capability of Maarge- These hypotheses, commonly referred to as drafts, can then be verified in parallel with MLarge using rejection sampling (Leviathan et al. 2023), i.e. discard all tokens after the first mis- match. We follow the standard draft and verification cycle that iterates over the following two steps until the "end-of-sequence" token or the maximum sequence length has been reached.
Msmall generates new tokens Ve+1, .......... 3/- where K is the length of the draft sequence, auto-regressively (Xia et al.; 2023 Li et al. 2024b) or in parallel (Ankner et al. 2024).
The number of tokens is typically fixed, which was the standard approach until recently
when dynamic K was proposed (Nair et al. 2024 Mamou al. 2024
Huang al. 2024
2. Verify K drafted tokens in single forward pass with MLagee- Verification uses ei- ther greedy (exact) matching (Xia et al. (023), speculative sampling (Zhou et al. [023] Leviathan et al.,, 2023) "approximate" verification (Stern et al., (0088) which relaxes the acceptance criteria but does not guarantee MLarge'ss 's output distribution. In all cases, after the first rejection, all subsequent/future tokens are typically discarded.
When the drafting of y/+1, Mt+KK is done auto regressegrys token by token, refer to this strategy as chain drafting. Several works (Cai et al. 2024b; [i et al. 2024b) extended this method to free drafting for additional optimisation. In this case, multiple tokens V be proposed for every future position i. Verification is performed with Tree Attention (Miao et al. (024) to efficiently handle multiple branching paths proposed by tree drafting. Consequently, this leads an increase in acceptance lengths and reduces the number of calls to MMLare compared to chain drafting. For small batch sizes, the LLM generation memory- better leverage the spare compute especially with Tree Attention.
this where speculative decoding can
In this paper, we employ tree drafting from EAGLE-2 (Li et al.) 2024b) to construct trees with variable structure. Starting from the root node, we expand the B most probable tokens from the model Msmall (+t+1l+). Then, for a fixed depth D, we recursively perform the following steps: for each existing branch, we compute the joint probability lhe Mssall (yr (tt,,,.,....) of their B child tokens, leading to B expansions
have B branches, each with B children. From the B expansions, we select the top B branches based on their joint probabilities for the next tree layer expansion. Upon reaching the maximum depth, we retain up m tokens from the total B + (D - 1) nodes, selecting those with the highest joint probability for verification.
2.3 ARCHITECTURE OF MSMALL
Speculative decoding architectures broadly fall into two categories, independent and self-drafting
Independent drafters are typically smaller versions of MLarge from the same model family (Li et al. 2024a Zhao et al. 2024 He et al. 2023) while self-drafting methods leverage either subset of Cai et al. 2024a).
MLange and/or newly initialised parameters (Ankner et al. 2024
Our contribution is built on EAGLE (Li et al. 2024b), b, self-drafting architecture which has shown the best results on the Beec-Bench (Xia et al. 2024) leaderboard so far. The drafter reuses the token embed and LM head parameters of Ma (1). It takes as input the ground-truth activations of the last decoder layer of ML 01, ... Of and the tokens of the sequence "1,".. Ut to predict the next activations 0(41, which is passed to the L.M LMMhead to predict the next token distribution:
the
...
6,-1 = M ((O
EAGLE ,0)), token eeneeeembedyyy... - .yt)), Dc+1 Softmax (LM head((,+)) 1))
((o)....
The process repeated by appending 0r41,9c41 1 the inputs to auto-regressively draft tokens Br42-
METHODOLOGY
3.1 MIXTURE OF ATTENTIONS
We begin by defining important properties of Msmall followed by detailed architectural choices.
3

---

Published as a conference paper at ICLR 2025
3.1.1 PARTIAL OBSERVABILITY
In Markov Decision Processes (Kallblnng et al. 1998), partial observability is common challenge where the agent does not have enough information about the true underlying state to take optimal de- cisions. This limitation can significantly degrade the agent's 's performance. Several approaches have been proposed to mitigate this, e.g., adding additional previous observations (Mnih et al. 2015). In drafting, it is important not to suffer from partial observability to draft future tokens more accurately. We extend this notion in our context with the following:
Property 3.1 (Partial observability). Given a ground truth function F X + Z, function f: y + Z and an observation function g : X + y, such that for any E f(gxx)) models F(x), f suffers from partial observability if g non-injective: 3(x,x') X2, x # x' g(z) ) = g(z").
We can observe that the EAGLE drafter suffers from partial observability when F is faares f is MEAGLE and 01. Or = g(hce (y1, -W)).. In other words, 01. , is only a partial ob- servation of the true state (hgr, V1) of the dynamic system ( hindering the capacity of Msmall to predict the right tokens of MLarge.
Layer Self- Attentionn Aiming to alleviate this, our new architecture takes as input the state of the dynamic system (T). However, h&g is large tensor of shape (T, L, 2 where T is the sequence length, L the number of layers in MLagge, and Ex the embedding size of the key and values. There- fore, we introduce Layer Self-Attention (LSA) followed by a mean aggregation operation to reduce its dimension to (T, 2 and extract the most relevant token information from every layer (Fig I Self-attention is performed over the layer dimension and each token is treated independently in this layer. During drafting, we have access the past key values of all the layers, therefore, the attention mask of LSA is bidirectionalffull, see Figure We only perform the LSA computations once at the for a detailed algorithm of the information flow.
start of each drafting phase, see Appendix A.3
h from M Marge
(T, L, 2B)
Layer Self- Attention
(T. L, 2.,..
(T, 2
K,,
Cross- Attention
(T,E)
L+1-N 0+1
tokens y/, " Vt (T.
MLarge embed_ tobee_tokens
self- Attention
Q
(T,
(T,
Figure I: A schematic overview of the mixture of attentions information flow. Layer Self-Attention and mean aggregation are called only once per drafting cycle, i.e after each verification. New tokens drafted ruge-segrysyyvely using
updating only the Cross-Attention layer query.
3.1.2 LACK OF ON-POLICYNESS
Discrepancies between training and testing scenarios arise because, during training, transformer models are typically conditioned on ground-truth sequences, assuming that all previous inputs are correct. If this assumption seems unproblematic for the standard training of transformers, assuming it for training Msmull in speculative-decoding scenarios is much more delicate. is known that some of the previous inputs are generated directly from MSmall, therefore much less accurate. The more tokens we predict with Msmall only, the more error accumulation we can expect. To alleviate this, EAGLE adds uniform noise into its observations (01,-.. ,O) at training time, but this is not ideal.
In order to train Msmall optimally, we need to ensure that its training and inference conditions are closely matched. Specifically. this means training MSmall as if some of the previous tokens were generated by itself. Additionally, we should account for situations where the activations from MLarge are not available, i.e. during the drafting cycle. This approach is called on-policy training- In on-policy training, the data used for training is generated by the same policy (or model) that is currently being trained. For example, when we train a transformer using next-token prediction on static dataset, this is considered off- because the data doesn't change based on the model's decisions. However, if we mix this static dataset with data generated by the model itself during training. we move towards more on- on-policy approach. Similarly, if the model won't have access to certain information, e.g. the activations of Miarge during generation, then always training Msmall with that information would also be considered offlpoy
4

---

Published as a conference paper at ICLR 2025
However, on-policy training is very costly because would need to generate from the model during training. To formalise this limitation, we introduce the concept of T-step boundedness:
Property 3.2 (T-step bounded), A drafter f is said to T-ttep bounded if, in a single forward pass, it can predict up to T future tokens without additional input from MLarge i.e., f y2...... ................................ , (9c+1.
This property is important efficiently train the drafter. For instance, the EAGLE drafter 1-step bounded. If one wanted to perform prediction at time t + 2, two forward passes would be required due to the auto-regressive layer that requires the previous ô.+1 as input, which would be very costly to train on-policy. By contrast, drafter that is Ttstep bounded with T > can predict multiple future tokens within a single forward pass.
Cross-Attention In order to make our drafter partly T-step bounded with > the main com- ponent of our architecture is a Cross-Attention (CA) layer where the query comes from the tokens and the key and values come from Marge activations. More precisely, the key and values come from the output of LSA. Having input queries for time & + 1 to K coming into the CA layer and keys-values from MLarge only up to time effectively means the CA layer is K-step bounded. This allows us to train the CA layer more on-policy efficiently because it simulates what would happen during generation: we only have access to the activations from MLarge up to time t but still have to make prediction for up time / K. Note that it is still not fully on-policy yet the input queries for time + 1 to & + K are not assumed to be generated from MSmalll During training, multiple K are sampled to simulate different lengths of accepted drafts by changing the CA layer mask. For instance, in Figure 2 we have K = 4 followed by K = 3. On the contrary, during generation, we do apply masking as want let Small attend all the currently available activations of MLarge
Self-Attention In order to motivate the introduction of self-attention (SA) layer, start by observing that the cross-attention layer is input-independer (3.3 w.r.t. the input queries, i.e one input query does not influence the results of another query.
Property 3.3 (Input-independence). A layer f imput- inpattindependent for any choice of inputs x [11, ), we have f(x) =
(Ix)).
Therefore, the queries of the CA layer came directly from the embedded tokens yI..... yt. would not have been aware of previously drafted tokens. It would only know the previous token treated by MLagge and the most recent ye- But, in order to make accurate predictions, MS needs to be aware of the previously drafted tokens. Hence, introduce causal self-attention layer the queries to mitigate this problem, shown in Figure I and summarised in Table
Msmall
Table 1: Comparison the properties of our new architecture.
Msmall
Autoregressive
T-step bounded
More on-policy
Observability
Ours
SA layer
variable T for CA layer
CA & LSA layers
LSA-enhanced
EAGLE-2
X
partial
Medusa
X
fixed T
x
partial
3.2 TARGET LAYER INFERENCE
Previous work assumed that the final hidden layer before LMahead was the most appropriate tar- get (activations) Msmall should predict. However, we challenge that assumption by hypothesising that targeting deeper MLarge layer may be more advantageous in terms of draft quality. We thus decompose the dynamic system (T) layer-by-layer by introducing I the (superscript) layer index:
1+1 = token_embed(yrl
h h 1+1 = fdecuder (n's o(.1).
1 Softmax LLMM (111))
1 1
where fleonter is the decoder layer of MLagee at layer l. The state of this new dynamic system L+1
composed of (o h stt1 h 12) We observe that to perfectly predict 0 it is sufficient to
5

---

Published as a conference paper at ICLR 2025
T
7
model layers
prompt response
calls Muag calls Muage
Layer Self-Attention
Self-Attention
Cruss-Attention
Large activations are transposed so that attention computed over the layer dimension order aggregate token activations across layers. Self-Attention: The first 3 tokens represent the prompt, speculative decoding starts at token 4. Cross-Attention: Tokens 4 to only attend to the prompt while tokens 10 attend to the first 7 tokens once MLarge was called for the second time allowing MSmall to use the activations from the newly verified tokens.
Figure 2:
Layer Self-Attention:
perfectly predict of and reuse the faleodeer of MLagg and the already computed KV cache h of the layer L at time t. The same recursive reasoning can be made to predict 0/41 from o etc. We assume (and later show) that predicting t'+1 always easier than predicting of for < & due to o'+1 undergoing fewer layer transformations. Hence, we introduce new hyperparameter TLI to refer to the target layer 14+1-TU that the Msmall should predict. When TLI > 0, the TL.I last layers of MLarge (kept frozen during training) and their KV cache are used to output o L+1,, Henceforth, we use notation (TLI !) where is an integer, to denote the target layer for inference. We can now provide the equation describing our M for a given TLI assuming was the last time we verified Small
with MLarge:
1L+1-IL1 = Muuu (h token embed(yi."". T-1
Mr. Det1, ,0r)),
h 0+1 Saocnder ( (1 sT) o'+1),
1 I- TLI...,,
D1+1 Bootmax(LM L+11
3.3 LOSS
Let MSmall be parameterised by 0, we use a similar training loss as EAGLE, i.e. a forward-KL. loss, with Smooth-LI loss C between the predicted activations of the Mssall 64+1-ILI and the target one obtained from MLarge
arg min X M |M (0) + X1 (6 -0-+1-T1I
(2)
To keep the training lightweight, we do not generate from MMLagg or Msmall during training. This loss is only defined over the response part of the prompt of fixed training dataset.
4 EXPERIMENTS
In all experiments, we use -Lama3-8B-Instree
(Dubey et al. 2024) as MLarge We train all Msmall
on the Ultrachat dataset (Ding et al. 2023) without system prompt and do not assume that we know the system prompt at test time as it was observed that the training dataset can have significant impact on the final performance (Yi al.) (024). Mssall trained with the standard Ilata3-Instruct chat template. Ultrachat is composed of around 200k prompts with around 240M tokens using the I.Lama3 tokenizer. We use multiple test datasets for generation including various tasks such as reasoning, code generation, multi-turn conversation and summarisation. We notably relied on the SpecBench benchmark (Xia et al. 2024) and the following datasets: MT-Bench (Zheng al. 2023), HumanEval (Chen et al. (2021) GSM8K (Cobbe al. 2021). Alpaca (TTori et al.
2023 CNN/Daily Mail (Nallapati et al. 2016) and Natural Questions (Kwiatkowski et al. 2019).
We describe additional hyperparameters and experimental settings in Appendix
We compare our method to EAGLE-2 and an independent distilled M of similar size (denoted "Independent"). In order train the EAGLE model, assume TI.I = 0 in the distillation loss ?

---

Published as a conference paper at ICLR 2025
The independent Msmall leverages the token_embed and LM_head parameters of MLarge with only the decoder layers trained using an identical distillation loss 2 and t = O. We do not compare to Medusa as EAGLE has consistently demonstrated superior speedups on various benchmarks (Xia et al. 2024). We also compare the performance of the official EAGLE-22 weights shared by [Li et al. (2024b). We refer to this as "EAGLE-2 off.". Note that this model was trained on different data and with a fixed system prompt. We take care to match the number of model parameters, i.c. "Ours (N=0)", "EAGLE-2", "EAGLE-2 off.", "Independent 1.3B" and "Glide" all have 1.3B parameters (250M trainable and 1.05B frozen, for the "LM head" and "token embed" layers). We chose 250M trainable parameters to be directly comparable to EAGLE-2 and their official checkpoint. For tree decoding, we use a max breadth of 8, depth of 6 and 62 max tokens to verify. We use f except for the attention softmax weights that are upscaled to float32.
We use standard metries: token-per-sen peree-dd and speedup ratios to measure walltime improvements as well as hardware-independeee metrics: average acceptance length t (the average number of . Mmall tokens accepted by MLarge) and the number of calls to Margge
4.1 SINGLE DEVICE
We now present the main single-device experiments using the SpecBench Xia et al. (2024) bench- mark without a system-prompt to ensure a fair comparison between models.
Table 2: Speedup ratio and acceptance length T on SpecBench using prompts from MT-Bench. HumanEval, GSM8K, Alpaca, Sum and QA datasets. Each model is fine-tuned for 30 epochs and uses EAGLE-2 tree decoding.
Total Trainable|
MT-bench Speedup
HumanEval r Speedup
GSM8K
Alpaca Speedup
CNN/DM
Natural Ques. Speedup
Mean
Mamill
size
size
Speedup
r
Speedup
Speedup
Ours (ILI=3)
L8B
250M
1.74
4.65
2.02
5.41
1.74
4.65
181
4.80
L88
5.04
1.59
4.23
1.79
4.79
Ours (TL==1)
1,55B
250M
183
4.19
2.29
5.30
1.83
4.19
2.02
4.65
2.04
4.74
1.71
3.94
1.95
4.50
Ours (TLI=0)
1.3B
250M
1.80
3.86
2.28
4.98
1.80
3,86
2.03
4.36
2.10
4.55
1.72
3.73
1.95
4.22
EAGLE-2
1.3
250M
1.77
3.55
1.95
3.92
3.36
1.89
3.77
1.84
3.69
3.60
EAGLE-2
1.3B
250M
1.75
3.52
2.06
4.15
1.80
3.60
1.70
3,37
1.60)
3.19
1.38
2.75
1.71
3.43
Independent
1.7B
650M
1.50
3.63
1.91
4.64
1.26
3.01
1.57
3.81
1.56
3.78
1.72
3.94
1.58
3.80
Independent
1.3B
250M
1.23
3.50
1.50
4.36
0.95
2.70
1.33
3.79
1.28
3.59
1.10
3.13
1.23
3.51
Glide
1.3
250M
1.69
3.62
2.06
4.43
1.54
3.27
2
4.27
1.6
3.37
1.59
3.41
1.74
3.72
Looking at Table 2 we can see that our Mixture of Attentions for SD achieves SOTA speedups
when TLI = and TLI = O. Compared to EAGLE-2, we are on average 9.5% faster in terms of tokens-per-xeconl generated. We also increase the acceptance length by 25% when N = 1. More
single device experiments c.g. on the full HvmanEval dataset are shown in Appendix A.4
4.2 CLIENT-SERVER
In this study, we investigate how self-drafting with our method performs in a client-server scenario. To do so, we place Mssall on client device and host MLange on server (see Appendix A.2 for an illustration). The server is performing verification and sends the relevant MLarge activations to the client, which turn is proposing new tokens. The server has times more f tflop than the client. The devices are located in two different cities, separated by ~300 km. The ping between the devices is around 9 ms and the bandwidth 50 Mbits/sec. In order to simulate realistic client-server scenario, we are using 5G and 4G network profiles. In 4G, we assume a maximum of 20 Mbits/sec with a normally distributed delay of 21 ms + 19 ms and a 0.1% chance of dropping packets. In 5G, we assume a normally distributed delay of 10 ms + 10 ms with a 0.1 1 % chance of dropping packets. To do so, we rely on the Linux traffic control subsystem.
In this scenario, the eoken-per-secont performance also depends on the size of the messages. To this end, we analyse the length of the messages sent between the client and the server (see Table()). There is clear distinction between self-drafting methods that need to send/receive activation tensors and independent methods that only exchange text (e.g. token ids). Therefore, we shall analyse whether the improvement in drafting quality offset the increase in message lengths. On the client, encode each node in the draft tree using 3 bytes for the token and byte for its position in the tree. The server answers with the accepted tokens encoded using bytes each plus the associated
7

---

Published as a conference paper at ICLR 2025
activations, if required. For Llama3-8B 8B-Instruct and N I our architecture's payload is less than or equal to EAGLE message lengths. In order to further reduce message sizes, we quantise the E and Exs tensors bits. For both EAGLE and Mixture of Attentions, the initial message sent by the server (before the first token drafted) typically the biggest as it represents the activations of the entire prompt. Therefore, additionally gzip-compress this message after quantisation.
Table 3: Performance HumanEval with EAGLE-2 tree decoding under 5G and 4G profiles.
Msmall
Total size
Trainable size
Tokens per sccond Acceptance length
Calls Maage
5G
4G
Ours (TLI=3)
1.8B
250M
25.0
14.6
4.99
20.8
Ours (TLI=1)
1,55B
250M
30.6
20.3
4.68
22.5
Ours (TLI=0)
1.3B
250M
34.1
25.1
4.30
24.1
EAGLE-2
1.3B
250M
24.3
13.6
2.81
36.4
EAGLE-2 off.
1.3B
250M
28.6
15.0
3.51
29.5
Independent
1.7B
650M
28.5
23.7
3.73
27.1
Independent 1.3B
250M
18.3
16.1
3.16
32.4
In Table 3 we can observe that "Ours (TLI=0)" achieves the fastest decoding speeds. Interestingly, even faster than independent small models that do not exchange any activation tensors. As expected, our Mixture of Attentions is not as fast as in the single device setting, but it can recover the speed of vanilla decoding in a single device setup (33 pokens-per-secome see Appendix[A.4)
However, for this setting to be viable, just recovering the speed of vanilla decoding is not sufficient as it does not provide an advantage over an API call to Maarge Therefore, we show that our model can continue to generate the response by simulating complete disconnection from the server.
Table 4: The success rate (pass @ 1, greedy decoding) on HumanEval in the event of an interrupted connection between the client and the server. EAGLEE-2 tree decoding used.
disconnection occurs after new tokens.
Msmall
Total size
Trainable size
B =
B = 10
B = 25
B 50
B = oo
Ours (TLI=3)
1.8B
250M
2.48 %
11.18 %
18.01 %
31.67 %
45.9 %
Ours (TL[=l)
1,55B
250M
3.10 %
10.55
21.11
30.43 %
45.9
Ours (TLI=0)
1.3B
250M
2.48 %
9.31 %
19.2 %
29.81 %
45.9 %
EAGLE-2
1.3B
250M
0%
8.07 %
16.77 %
27.32 %
45.9 %
EAGLE-2 off.
1.3B
250M
1.24 %
6.83 %
18.01 %
28.57 %
45.9 %
Independent
1.7B
650M
0 %
6.83
18.63 %
29.81 %
45.9
Independent
1.3B
250M
0 %
6.21 %
18.01 %
27.95 %
45.9 %
Generation stops after B new tokens.
Without local model (lower bound)
0%
5.59 %
16.77 %
27.32 % 45.9 %
In Table we can see that indeed, a disconnection occurs, unlike API calls to MLarge we can continue generate the response right the device, i.e. complete additional correct solutions to competitive programming problems HumanEyal. Therefore, with an acceptable speed and the possibility to generate useful responses after disconnection, prove the viability our proposed client-server setting, paving the way for a new framework for serving LLMs with small devices.
4.3 ABLATION STUDY
We now present important ablation results for different components of our Mixture of Attentions architecture. Since multiple models were required to be fine-tuned for this study, have limited each run to 10 epochs. For this ablation, we introduce the "Ours (TLI=I, -LSA)" variant that does not rely LSA and takes input ... O₁ the keys and values of the CA layer. We also include two more EAGLE baselines, one with additional trainable parameters "EAGLE (more params)" and another with additional decoder layers "EAGLE (more layers)" but equal number of trainable
8

---

Published as a conference paper at ICLR 2025
parameters. This is to ensure that the benefit of our architecture does not come from simply adding decoder layers or parameters. this experiment, we use the HumanEval dataset with striet stopping criteria, exiting decoding soon as the model no longer generates source code.
Table 5:
An ablation study of our proposed architecture, tested on HumanEval. Each model is trained on ~2.4B tokens. Chain (not tree) drafting with maximum 4 tokens used for this study. The averages are computed over around 8500 drafting-verificatinn cycles.
Msmull
Total size
Trainable size Tokens per sccond Acceptance length (r)
Ours (TII==)
1.8B
250M
39
2.54
Ours (TI.==1)
1,55B
250M
39
2.25
Ours (TI.=0)
1.3B
250M
40
2.14
Ours (TLI=1, -LSA)
1.55B
250M
1.28
Ours (TLI=1, -LSA, 01, ....... o, inputs)
1,55B
250M
36
2.04
Ours (TLI=0, -SSA on, inputs)
1.3B
250M
38
1.93
EAGLE
1.3B
250M
30
1.45
EAGLE (more params)
1.45B
400M
29
1.28
EAGLE (more layers)
1.3B
250M
1.01
Does the on-policyness (brought with the CA layer) and the T-step bounded property have positive impact on the quality of the drafts? In Table [S] we compare EAGLE with "Ours (TLI=0, -LSA)" for an answer to this question. We can see that these components provide major improvement of 26% in
well as improved acceptance length of 33%.
How does partial observability influence the drafter acceptance rate? In Table pare "Ours (TLI=0, -LSA)" to "Ours (TLI=0)" as well as "Ours (TLI=1, -LSA)" to "Ors (TLI=1)" and report that the per-second performance improves by 6% by introducing LSA, decreasing partial observability. Its impact is less crucial than the on-policyness brought by the CA layer.
we can com-
Does increasing TLI increase the acceptance rate? Finally, by looking at the variation of TLI in Tables 2 and increasing TIII also increases the acceptance length, as we hypothesised. However, this does not always have a positive impact on the tokens-per-secon rate as also increases the computational time of drafting. In the event of a complete disconnection in client-server setting, however, a higher TLI will improve the quality of responses, which is something to consider when deploying Mixture of Attentions for SD on mobile devices.
5 RELATED WORK
Medusa Cai et al. 2024a) is one of the earliest works leveraging the activations of Marge as inputs to Msmall for the purpose of SD. Thanks their work, speculative decoding can be applied any
LLM by distilling an MSmall- It generates K future tokens in parallel by training K new LM heads where each head predicts token at position & K (Gloeckle et al. 2024). It was later extended by Kim et al. (2025) by refining the block drafts using task tesskinndeendeent n- n-gram and neural language models as lightweight rescorers. EAGLE (Li et al., (024c) and Hydra (Ankner al., 2024) are auto-regressive extensions of Medusa. They observe that non-auto-regressive generation limits the acceptance length as MSsmall is not aware of previous tokens. We not compare to Medusa or Hydra as EAGLE is ranked higher on the SpecBench leaderboard.
Tandem Transformers (Nair al., 2024 propose an effective integration of MLarge and Msmall by letting Msmall attend to the down-projected hidden states of MMarge These rich contextualised rep- resentations enable MSmall to draft hypotheses with a higher acceptance rate as the two models are aligned on shared hidden states. We were not able to compare with them because of the lack of open- source implementation, the use closed-source LLMs and an undisclosed amount of data/compute to reproduce the work. Moreover, tandem transformers appear to have high communication over- head between big and small models, making it unrealistic for a client/server setting.
Orthogonal to work, researchers have recently proposed training-free SD methods. Lookahead Decoding (Fu et al. 2024) generates new tokens with single M1 using Jacobi iterations. ex- tended by CLLM Kou al. (2024) and Ouroborous Zhao al. 2024). We evaluated the latter in

---

Published as a conference paper at ICLR 2025
settings, however, it was shown be less efficient than the EAGLE-2 tree decoding strategy, Appendix A.4 For additional related and orthogonal work in the extended SD landscape, we refer the reader to Xia et al. (2024) for detailed and highly informative speculative decoding survey.
Du et al. 2024) previously proposed to leverage the KV-cache of some layers of MLarge. They do not theoretically justify why using the KV-cache instead of the output of each layer, nor how to exactly choose which layer to include as input of MSmall- However, with our dynamical system point of view, showed that the KV-cache of all the layers is part of the state. The introduction of LSA allows to exploit it in its whole with a limited number of layers, whereas Du et al. (2024) would need have the same number of layers in Msmall and Maagge to fully capture resulting in a slow drafting speed.
Although we focused on improving the current SOTA method (EAGLE-2), our observations (partial observability, ne-policyness and target inference layer) are true for many self-drafting methods, for instance, could also be applied Medusa (Cai et al. 2024a). MLP Speculator (Wertheimer al.)
2024 or Gloeckle et al. (2024); Kim et al. (2025).
Regarding non-self-drafting SD, should studied on a case-by-case basis. For instance, target inference layer could potentially be applied to independent small models. Many student-teacher dis- tillation frameworks (Gu et al. [2024] (/hou et al. 2023), already leverage the on-polieyness property by generating directly from the student but mostly 1-step bounded (therefore expensive to train). For SD methods based on lookahead decoding, it would generally not apply. One exception Ouroboros (Zhao al., (222)) that leverages small model with lookahead decoding. Their small model could also benefit from our solutions.
CONCLUSION
We have introduced a Mixture of Attentions architecture for Speculative Decoding to effectively address several limitations of existing state-of-the-art methods. In order enhance drafting accu- racy of MSmalll we proposed a mixture of attention layers: Layer Self-Attention to mitigate partial observability and Self-Attention followed by Cross-Attention to train more on-policy. We have then introduced Target Layer Inference, a novel method that lets MSmall reuse the last N layers of MLarge. enabling trade off between the drafting speed and accuracy. Experimental results show that we achieve state-of-the of-the-art decoding speedups in the standard single-device setup, improving over EAGLE-2 by 9.5% and extending acceptance lengths by up 25%. We have also introduced a client-server paradigm and demonstrated that our self-drafting speculative decoding method is viable alternative to API calls to MLarge Under this paradigm, the client can continue to generate responses with the highest accuracy and speed after a complete disconnection from the network. As a future direction, it would be interesting to investigate whether N could predicted by MM to automatically balance speed and accuracy depending on the current network conditions.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Ragan- Kelley, and William Brandon. Hydra: Sequentially-dependee draft heads for medusa decoding. arXiv preprint arXiv:2402.05100 2024.
Tianle Cai, Yuhong Li. Zhengyang Geng, Hongwu Peng, Jason D Lee. Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401. 100::20411..744, 2024a.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple IIm inference acceleration framework with multiple decoding heads, 2024b.
URL
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray. Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel feeberrtVooss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichrr Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 14168, 2021.
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.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations, 2023.
Cunxiao Du, Jing Jiang, Xu Yuanchen, Jiawei Wu, Sicheng Yu, Yongqi Li, Shenggui Li, Kai Xu, Liqiang Nie, Zhaopeng Tu, et al. Glide with a cape: A low-hassle method to accelerate speculative Ttt::2402.0208822 2024.
decoding. arXiv preprint
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang. Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez. Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song. Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Ander- son, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Ma- hadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong. Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Al- wala, Kartikeya Upasani. Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala- Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo,

---

Published as a conference paper at ICLR 2025
Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Man- nat Singh, Manohar Paluri, Marcin Kardas, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang. Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhar- gava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sum- baly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov. Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang. Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Kar, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vladan Petro- vic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whimey Meers, Xavier Martinet, Xiaodong Wang, Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen. Zoe Papakipos, Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexci Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti. Bran- don Spence, Brani Stojkovic. Brian Gamido. Britt Montalvo, Carl Parker. Carly Burton, Catalina Mejia, Changhan Wang. Changkyu Kim. Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Due Le. Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Ar- caute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Gold- man, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Baantttp Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang. Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen. Lakshya Garg, Lavender A. Leandro Silva, Lee Bell, Lei Zhang. Liangpeng Guo. Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Ke- neally, Michael Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mo- hammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navy- ata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong, Ning Zhang, Norman Cheng, Oleg Chemoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Rohan Maheswari, Russ Howes, Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Sa- tadru Pan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lind- say, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shiva Shankar, Shuqiang Zhang, Shuqiang
12

---

Published as a conference paper at ICLR 2025
Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Tim- othy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vitor Albiero, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Con- stable, Xiaocheng Tang, Xiaofang Wang, Xiaojian Wu, Xiaolan Wang, Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd of models, 2024. URL pettps://aaxiv...o
abs 2407 21 783
Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of Ilm infer- ence using lookahead decoding. arXiv preprint AXXiv:2402.02057, 2024.
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. arXiv preprint arXiv:2404. /9737, 2024.
Yuxian Gu, Hao Zhou, Fandong Meng, Jie Zhou, and Minlie Huang. Miniplm: Knowledge distilla- tion for pre-training language models. arXiv preprint arXiv:2410.17215. 2024.
Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He. Rest: Retrieval-based speculative decoding. arXiv preprini arXiv:2311.08255 2023.
Kaixuan Huang, Xudong Guo, and Mengdi Wang. Specdec++: Boosting speculative decoding via adaptive candidate lengths. arXiv preprint arXiv:2405. /9 2024.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101(1):99- 1998. ISSN 0004-3702. doi: http:///b.ogg///0000000000000010
URL https://wwwt
Taehyeon Kim, Ananda Theertha Suresh, Kishore Papineni, Michael D Riley, Sanjiv Kumar, and Adrian Benton. Accelerating blockwise parallel language models with draft refinement. Advances in Neural Information Processing Systems, 37:34294-34321, 2025.
Lingkai Kong, Haorui Wang, Wenhao Mu, Yuanqi Du Yuchen Zhuang, Yifei Zhou, Yue Song, Rongzhi Zhang, Kai Wang, and Chao Zhang. Aligning large language models with representation editing: control perspective, 2024. WtRLhttts::// /ttp///aaivvvv//v http://awwwv...g/.....//// abs////////// 2406. 05954
Siqi Kou, Lanxiang Hu, Zhezhi He, Zhijie Deng, and Hao Zhang. Clims: Consistency large language models. arXiv preprint arXiv:2403.00885, 2024.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 466, 2019.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023.
Minghan Li, Xilun Chen, Ari Holtzman, Beidi Chen, Jimmy Lin, Wen-tau Yih, and Xi Victoria Lin. Nearest neighbor speculative decoding for Ilm generation and attribution. arXiv preprint arXiv:2405. 19325, 2024a.
13

---

Published as a conference paper at ICLR 2025
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. 2: Faster inference of language
models with dynamic draft trees. arXiv preprint arXiv:2406.16858, 2024b.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires
rethinking feature uncertainty, 2024c.
Jonathan Mamou, Oren Pereg, Daniel Korat, Moshe Berchansky, Nadav Timor, Moshe Wasserblat, and Roy Schwartz. Accelerating speculative decoding using dynamic speculation length. arXiv preprint arXiv:2405.04:: 043044 2024.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang. Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, et al. Specinfer: Accelcrating large lan- guage model serving with tree-based speculative inference and verification. Proceedings of the 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 932-949, 90244022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski. et al. Human-level control through deep reinforcement learning, nature, 51 518(7540):529-333, 2015.
Pranav Ajit Nair, Yashas Samaga, Toby Boyd, Sanjiv Kumar, Prateek Jain, Praneeth Netrapalli, et al. Tandem transformers for inference efficient llms. arXiv preprint arXiv:2402.08644 2024.
Ramesh Nallapati, Bowen Zhou, Caglar Gulcehre, Bing Xiang, et al. Abstractive text summarization using sequence-to-seqqence rnns and beyond. arXiv preprint arXiv:1602.00023,, 2016.
Antonio Polino, Razvan Pascanu, and Dan Alistarh. Model compression via distillation and quanti- zation. arXiv preprint arXiv: 0et..0002.00000 2018.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv:1909.0800 2019.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances in Neural Information Processing Systems, 31, 2018.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An follooii llama model. https://gtt / thub. com/ httr//attuuuute labp st//stanna sttafordd_alppaeeoooo 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Davis Wertheimer, Joshua Rosenkranz, Thomas Parnell, Sahil Suneja, Pavithra Ranganathan, Raghu Ganti, and Mudhakar Srivatsa. Accelerating production Ilms with combined token/embedding speculators. arXiv preprint arXiv:2404.19124 2024.
Heming Xia, Tao Ge. Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. Speculative decod- ing: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 3909-3925, 2023.
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:240..0785/ 2024.
Euiin Yi, Tachyeon Kim, Hongseok Jeung, Du-Scong Chang, and Se-Young Yun. Towards fast multilingual I1m inference: Speculative decoding and specialized drafters. arXiv preprint
arXiv:2406. 24152880..6678. 2024.
Biao Zhang, Ivan Titov, and Rico Sennrich. Sparse attention with linear units. arXiv preprint arXiv:2104.0701. 2021.
Weilin Zhao, Yuxiang Huang, Xu Han, Chaojun Xiao, Zhiyuan Liu, and Maosong Sun. Ouroboros: Speculative decoding with large model enhanced drafting. arXiv preprint 2rXiv:2402.137200 2024.
14

---

Published as a conference paper at ICLR 2025
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 mntcbench and chatbot arena. Advances in Neural Information Processing Systems, 36:46595-46623 2023.
Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, and Rishabh Agarwal. Distillspec: Improving speculative decoding via knowledge distillation, arXiv preprint arXiv:2310.08661, 2023.
15

---

Published as a conference paper at ICLR 2025
A APPENDIX
The source code is publicly available at tree/m ofwre--f-attenti.ooo
A.1 HYPERPARAMETERS
Table 6: List of our hyperparameters.
Distillation
Learning rate for gradient descent Total numbers of transformer updates
3 10-3
186000
Minibatch size
32 yes, float16
Mixed precision training
Mixed-precision
Weight reserve KI. loss (Ao)
Weight of L1 smooth loss (A1)
1.0
L2 gradient clipping
1.0
T Ttstep bounded mask for the CA layer
Uniform between 5 to 15
Architecture
Number of layers of MLange Embedding dimension E of MLarge Embedding dimension of keys and values Exo of MLarge Dropout rate
4096
1024
0.0
Embedding dimension of Layer Self-Attention
2048
Embedding dimension of Self-Attention Embedding dimension of Cross-Attention Size of the MIP projection after Layer Self-Attention Size of the MIP projection after Self-Attention Size of the MLP projection after Cross-Attention Embedding dimension of keys and values of Layer Self-Attention
4096
4096
6144
512
7168
1024
Embedding dimension of keys and values of Self-Attention
512
Embedding dimension of keys and values of Cross-Attention
1024
A.2 CLIENT SERVER DEPLOYMENT
activations
(T, (()
(T.
veri fied tokens
K.V
Server
mean
Miszge embed_tokens tokens
Cross- Attention
L, 2
Layer self- Attention
(T, F)
h<t
MLarge
self- Attention
(T, F)
(T, L, 2
(T, E)
Le_deed
client
drafted tokens
Figure 3: A client-server setting for our mixture of attentions architecture with N = 0.
16

---

Published as a conference paper at ICLR 2025
Table 7: The size of the message (before quantisation) in bytes. M number of nodes in the draft
tree, A = number of accepted tokens, E = hidden size, Eie = hidden size of key and query vectors.
Msmall
Sent by Client
Sent by Server
Ours
4M
3A + 2.A (TLI+ + 1)
EAGLE
4M
3A AE
Independent
4.
3A4
A.3 ALGORITHM
Algorithm Generation algorithm for Monn assuming chain decoding
Require: Input sequence (y1, 32..... V1). draft length K, target layer inference TLI
Obtain fuga activations and V1+1 with forward pass in Maarge given input y
2: y - (y. (1+1)
3: kv + LSA -SS_Jayer later_with with mean (hgt)
4: while stopping criteria is not meet y do to 5: for SA_layerteeen.mmbellll
6:
7:
61.+1-N < CAlayyerqq kv)
8:
NN > then
9:
for I to do
10:
hh' 66.0/+144 flecoder ((' h -1))))))))))))
11:
end for end
12:
13:
Softmax(LM head(ō < (v:
14:
15:
end for
16:
Identify K' verified tokens out of the K latest tokens of y, obtain associated h' and obtain y/ with forward pass in MLarge with inputs y K... lvl and hgg
kv LSA LSyleaeerr with mean (h')
18:
ko + (kv,,
Update by appending the new h' components
Discard previous h
21:
y yI,.. Ju...pl-KKa K+K (keep only the verified tokens)
22:
+ yl
y (y)
24: end while
return y
A.4 ADDITIONAL EXPERIMENTS
We several experiments to assess the quality of the generated responses using greedy decoding. We focused on 3 datasets from SpecBench (HumanEval, GSM8K and CNN/DM) that do not require access proprietary models/APIs for evaluation (Ilm-as-a-judge)
Accuracy of the generated text
17

---

Published as a conference paper at ICLR 2025
Table 8: Quality of the generated text.
Vamilla decoding
HumanEval (pass @ I)
GSMSK (accuracy)
CNN/DM (Rouge-I. (-score)
Mamma-8B-Innsruce
62.5%
80%
0.3071
Speculative Decoding
HumanEval (pass @ I)
GSMSK (accuracy)
CNN/DM (Rouge-L. f-score)
Ours (TLI=3)
62.5%
80%
0.3053
Ours (TI.==I)
62.5%
81.25%
0,3068
Ours (TLI=0)
62.5%
80%
0.3070
EAGLE-2
62.5%
81.25%
0,3062
EAGLE-2 off
62.5%
80%
0,3056
Independent 1.7B
62.5%
80%
0.3067
Independent 1.3B
62.5%
80%
0.3064
We report the results in T The pass@1 @1 on HumanEval is the same across all methods. The accuracy on GSM8K actually improves w.r.t the base model on one question for Ours (TLI=l) and EAGLE-2.. Finally, the ROUGE scores are also extremely similar, leading to conclude that any differences to the base model are negligible and almost certainly appear due to using float
Qwen2.5 3B We trained additional small models on the Ultracaat dataset to accelerate Qwen2.5 3B. EAGLE recommends to one decoder layer of the big LLM define the size of the small LM, which leads to a trainable size 80M parameters. We kept the shared "embed tokens/LM head" layer frozen.
Table 9: Speedup ratio and acceptance length T on SpecBench using prompts from MT-Bench, HvmanEval, GSM8K, Alpaca, Sum and QA datasets with Qwen2.5-3B Instruct.
Total
Trainable
MT-bench
HvmanEEval
GSMSK
Alpaca
CNN/DM
Natural Ques.
Mean
Msmall
size
size
Speedup
r
Speedup
r Speedup
Speedup
Speedup
Speedup
Speedup
Ours (TLI=0) 0.4B
1.71
3,72
2.18
4.76
1.60
3.46
3,9
EAGLE-2
0.48
80M
1.59
3.2
1.84
3,70
1.53
3.06
1.81
3.54
1.60
3.23
1.62
3.17
1.66
3.31
Independent 0.4B
80M
1.59
3.37
2.04
4.38
144
3.03
1.70
3.52
1.54
3.27
1.50
3.12
1.63
3.44
Higher batch size with vLLM We implemented our approach vLLM (Kwon et al. 2023
without tree decoding to support higher batch sizes and continuous batching.
Online inference in vLLM for Llama3-888 BBB-ttruttt (without Tree Attention, HumanEval 80 prompts)
1.8
1.7
1.6
Latency Speed
1.5
1.4
MUA (1L--0)
1.3
Eack 6-2
1.2
11
0.5
Request per second (with Poisson law for arrival time)
Figure 4: vLLM inference with continuous batching.
HumanEval single device To perform this experiment, we reuse the same full HumanEval dataset with strict stopping criteria done in the ablation study in the single device setting.
18

---

Published as a conference paper at ICLR 2025
Table 10:
Test on Human Eval, each model is trained for 30 epochs.
Msmall
Decoding
Total size
Trainable size
Tokens per second
Acceptance length (r)
Ours (TLI=3)
EAGLE-2
1.8B
250M
54
5.02
Ours (TLI=1)
EAGLE-2
1.55B
250M
58
4.70
Ours (TLI=0)
EAGLE-22
1.3B
250M
57
4.30
EAGLE
EAGLE-2
1.3B
250M
43
2.82
EAGLE off.
EAGLE-2
1.3B
250M
52
3.50
Independent
EAGLE-2
1.7B
650M
46
3.72
Independent
EAGLE-2
1.3B
250M
34
3.17
Independent
Ouroboros
1.7B
650M
39
2.37
Baseline
Vanilla
33
From Table 10 we can observe we are 26% faster than EAGLE/EAGLE-2. We are also faster than independent small models and Ouroboros (Zhao al. 2024).
A.5 COMPLEXITY ANALYSIS
Let us analyze the standard decoder-only transformers doing vanilla decoding:
in the first prefill stage, it grows in O(LKE with K imput tokens and an embedding size of E
K)) (K)) given we have L self-attention layers
for the K' new decoded tokens, it grows in O K' T(F² F(KK + +)))
O(LE(ER' +
KK' + K²))).
we assume E and L are fixed, grows in O((K + +')²) overall. For speculative decoding, the first prefill stage is the same. Assuming tokens verified at time, the verification would grow O Kf L(SE² SE(K +))) O(LE(EK' KK' K K))),, leading to the same complexity (E(K+
as vanilla decoding. It dominates the complexity of self-drafting, but we can still analyse it. For EAGLE, decoding new token grows in OEE² EK) as it is a single self-attention layer. For Mixture of Attentions architecture, the Self-Attention and Cross-Attention layers also grow in + EK). The Layer-Self Attention is only called once after every verification stage, so not every decoding step, it grows in OALEE2 + A. L2) if A is the number of accepted tokens in the previous phase. In our experiments, number deecodeed.ooeeee and number _oʃ doecons is 48. Similarly for the second term, AE L is usually smaller than number_of _decoded tokens x EK soon as the request contains more than 24 tokens. Therefore, the time complexity is the same as EAGLE overall.
we look at the first term, ALE? is smaller than E Ex, is times smaller than E, is 32, A is in average 4.5
A.6 PRIVACY APPLICATION
Another advantage of the client-server setup is that we can selectively ensure privacy for the client by only sending the non-sensitive part of the prompt the server. Essentially, the client split their input into consecutive "safe" text and a "private" text. The server processes only the "safe" text, which could be general context non-sensitive information. The client keeps the "private" text, such as confidential data or sensitive instructions, and handles this part locally with M
For instance, the client might send the server some Python code along with a general description. However, any sensitive information, such as the login and password to inject into the code, remains on the client side and is not transmitted to the server. It is only passed to M This approach leverages the activations of MLagee to increase the accuracy of Msmmall for parts of the task while ensuring that sensitive information is never exposed outside the client's environment.
19


Reviews:
[{"review_id": "Brfxb1HWRR", "reviewer": "Reviewer_eFLF", "summary": "The method addresses two key limitations of existing SD approaches, (1) partial observability and (2) lack of on-policyness, by incorporating Layer Self-Attention (LSA) and Cross-Attention (CA).", "questions": "See Weakness.\n\nWill update the score after looking at the results of Weakness.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The client-server framework with the ability to handle disconnections positions the approach as a practical advancement for deploying LLMs.\n\n- The introduction of LSA and CA layers to mitigate partial observability and improve on-policyness makes sense.", "weaknesses": "1. The paper does not thoroughly justify the choice of parameter configurations and its training in its experiments.  As discussed in the Yi et al. (2024), the training dataset and the choice of number of parameters can significantly affect the SD performance, but this paper does not [A].\n\n[A] Yi et al., 2024. Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters, EMNLP 2024-main.\n\n2. Discussions for the memory-bound nature of LLM is required in the paper.\n\n3. The effectiveness of the proposed method on models having 3B~13B parameters is unclear. Current experiments focus on relatively smaller models, and the results may not hold for state-of-the-art LLMs, which typically exhibit different scaling dynamics and memory behavior.\n\n4. Typo? line 466.\n\n5. Parallel to Medusa, actually there are concurrent works regrading non-autoregressive heads for SD. It would be good to put discussions for those areas. \n\n- Gloeckle et al. (2024), Better & Faster Large Language Models via Multi-token Prediction.\n\n- Stern et al. (2024), Blockwise Parallel Decoding for Deep Autoregressive Models\n\n- Kim et al. (2024), Accelerating Blockwise Parallel Language Models with Draft Refinement. (https://openreview.net/forum?id=KT6F5Sw0eg)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "McTz8t67jE", "reviewer": "Reviewer_2FbK", "summary": "This paper proposes a method, Mixture of Attentions for SD, which improves upon the standard SD approach that is used to increase LLM efficiency. The improvements are specifically targeting the partial observability and lack of on-policyness problems of traditional SD. The Mixture of Attentions method is shown to lead to improvements for a single device, but the paper also shows its efficiency/accuracy benefits to a client-server setting.", "questions": "Why did you choose to only experiment with LLama3-8B-Instruct? It is good that there is justification for only comparing to EAGLE (and not Medusa or Tandem Transformers), but there is no justification for your LLama model choice. How do you think your method would work with other architectures?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The organization and flow of the paper is very good. The background section is particularly thorough and helpful.\n\nThe problem is well-explained (e.g. partial observability and lack of on-policyness are both detailed when explaining the methodology) so it is made clear what exactly the Mixture of Attentions method is aiming to solve. Additionally, the related work is well-addressed. It is clear exactly how this work is different from prior solutions.\n\nIt is great that the client-server scenario is tested in a practical setting with different devices (having different resource capacities) and the devices have distance between them. This setting is not only realistic, but it is also well-explained in the text.\n\nThere is good theoretical support in the Methodology section, which provides additional justification for the proposed method being superior to the current SOTA (EAGLE, Medusa).\n\nThe thorough experiment detail (particularly in the appendix) makes the method highly reproducible.", "weaknesses": "The experimentation is very narrow, especially since it only focuses on one model architecture and the improvements over EAGLE seem relatively small and inconsistent. It is therefore not convincing that this method would be effective more generally.\n\nIt is not very clear why this problem/contribution is important. The paper would be stronger if the method was motivated by some real-world example where SD may be used, but would lead to significant problems that Mixture of Attentions would mitigate. I understand that the computational requirements of LLMs is an issue, but the introduction could do a better job of explaining why SD should be focused on as a solution for the computational expense and therefore is important to build improvements for. It is also difficult to understand how this work can be have a broader impact or inspire future work. The future work that is suggested at the end of the conclusion seems very specific and narrow. Essentially, the contribution just seems very narrow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "VOaZSLmlfQ", "reviewer": "Reviewer_Wo7q", "summary": "This paper focuses on an LLM acceleration technique called Speculative Decoding, which leverages efficient models (smaller but less capable) to draft future tokens, which are verified by the LLM (more capable but much less efficient) in parallel. In particular, it addresses the limitations of previous methods by proposing a Mixture of Attentions architecture on top of a prior work EAGLE to improve its performance. They demonstrate the effectiveness of their approach in both a single-device setting and a novel client-server setting, achieving speedups and improved accuracy. They also present a framework for using LLMs on edge devices, allowing for offline text generation with minimal dependence on a server.", "questions": "(also see weaknesses)\n- The experiment on N is arbitrarily set to 0, 1, 3. Why not continuously evaluate 0,1,2.... to the number of layers in the large model?\n- In the client-server mode with N > 0, are the last layers of the large model copied onto the client side?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The work improves above EAGLE-2 and seems to achieve state-of-the-art results.\n- The work provides a good background on speculative decoding\n- The work proposes an interesting client-server setup that fits well with the speculative decoding technique", "weaknesses": "- The work lacks an overall view and clear statements that can improve readability.\n    - Method intuition: the method section only lays out the information of each component but does not provide an overall view of the proposed method as well as motivating intuitions for each design. The necessary intuitive descriptions are also not found in the appendix.\n    - Experiment result: the work only compares to one prior work, EAGLE-2, as a baseline, but did not provide information on how well EAGLE-2 compares with other prior works.\n    - the hyperparameter N is used many times. Giving it a name can help readability.\n- This work overlooks overall performance metrics, concentrating instead on metrics specific to the speculative decoding framework, such as acceptance length. This focus may inadvertently encourage adversarial scenarios where the smaller model aims merely to deceive the larger model into accepting its outputs rather than genuinely enhancing result quality.\n- The work neglects analysis of the time/computational complexity of its method. \n- The ablation study compares different hyperparameters but does not ablate the other components. Thus, the importance and contributions of the designs in Section 3.1 and Section 3.3 is not provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "YDPkR2oHTb", "reviewer": "Reviewer_Sr5b", "summary": "The paper proposes an improved architecture for speculative decoding, which incorporates Layer Self-Attention and Cross-Attention mechanisms to address common limitations like partial observability and lack of on-policy training. Experimental results demonstrate that this approach achieves state-of-the-art speed and accuracy in single-device and client-server deployments, maintaining high accuracy even during disconnections. The limitation analysis of SD in the paper provides insights into this field.", "questions": "Please refer to the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "(1) The limitation analysis of SD in the paper is very interesting, it links the limitation of SD (e.g., the SD may still give different predictions compared to only using the large model even when the previous tokens are all accepted ) with theoretical analysis, i.e., partial observability. This may provide further insights of optimizing SD algorithms.\n\n(2) It is a very good point that the paper design and evaluate the scenario that the small model is on a resource-limited client and the large model is on the server. It is a very realistic scenario to use SD.\n\n(3) The paper uses SOTA models such as Llama-3 for evaluation. Also the benchmarks used (e.g., HumanEval, MT-Bench) are challenging enough for evaluation.", "weaknesses": "(1) The link between the limitation of SD and the proposed attention-based method is not clear enough. Why Layer Self-Attention can solve partial observability? Further intuitive explanations or theoretical analyses are needed.\n\n(2) Authors mention that their work is based on a SOTA SD method named EAGLE. Then my concern is, is the proposed algorithm transferable to other SD algorithms? For example, can the algorithm be applied to non-self-draft SD? The point here is that, if the algorithm can be applied to most of the SD frameworks, there will be more contribution. Otherwise, it's just an optimization of one of the previous works.\n\n(3) No accuracy is displayed in the experiment section. I know that this is common in SD papers. However, when I was trying SD codes, I found that the SD performance is usually not as good as the large model, which is the purpose of SD. Thus, because this paper aims to solve problems such as partial observability, will it also increase the performance of SD, in addition to the efficiency? I believe experiments related to this point will make the paper more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)