AdaDecode: Accelerating LLM Decoding with Adaptive Layer Parallelism
Authors: Zhepei Wei
Wei-Lin Chen
Xinyu Zhu
Yu Meng
Venue: ICML 2025 poster (2025)
Keywords: Autoregressive model
efficient decoding
parallel token processing
Primary area: deep_learning->large_language_models
TL;DR: Built on the proposed adaptive layer parallelism, AdaDecode significantly accelerates autoregressive decoding while ensuring output parity, without the need for auxiliary models or any changes to original model parameters.

Abstract: Large language models (LLMs) are increasingly used for long-content generation (e.g., long Chain-of-Thought reasoning) where decoding efficiency becomes a critical bottleneck: Autoregressive decoding is inherently limited by its sequential token generation process, where each token must be generated before the next can be processed. This sequential dependency restricts the ability to fully leverage modern hardware’s parallel processing capabilities. Existing methods like speculative decoding and layer skipping offer potential speedups but have notable drawbacks: speculative decoding relies on an auxiliary “drafter” model, which can be challenging to acquire and increases memory overhead, while layer skipping may introduce discrepancies in the generated outputs due to the missing key-value cache at skipped layers. In this work, we propose AdaDecode, which accelerates LLM decoding without requiring auxiliary models or changes to the original model parameters, while ensuring output consistency. AdaDecode leverages the insight that many tokens—particularly simple or highly-predictable ones—can accurately be generated at intermediate layers, as further layers often do not significantly alter predictions once the model reaches a certain confidence. By adaptively generating tokens at intermediate layers when confidence is high, AdaDecode enables the next token’s computation to begin immediately. The remaining layer computations for early-predicted tokens are deferred and executed in parallel with subsequent tokens when needed, maximizing hardware utilization and reducing decoding latency. A final verification step ensures that early predictions match the results of standard autoregressive decoding, preserving output parity. Experiments across diverse generation tasks shows that AdaDecode consistently achieves superior decoding throughput compared to baselines with up to 1.73$\times$ speedup, while guaranteeing output parity with standard autoregressive decoding.

AdaDecode: Accelerating LLM Decoding with Adaptive Layer Parallelism
Zhepei Wei Wei-Lin Chen
Xinyu Zhu'
Yu Meng'
Abstract
preserving output parity. Experiments across di- verse generation tasks shows that AdaDecode con- sistently achieves superior decoding throughput compared to baselines with up to 1.73 x.73 speedup, while guaranteeing output parity with standard autoregressive decoding."
Large language models (LLMs) are increasingly used for long-content generation (e.g. long Chain- of-Thought reasoning) where decoding efficiency becomes critical bottleneck: Autoregressive de- coding is inherently limited by its sequential to- ken generation process, where each token must be generated before the next can be processed. This sequential dependency restricts the ability to fully leverage modern hardware's parallel pro- cessing capabilities. Existing methods like specu- lative decoding and layer skipping offer potential speedups but have notable drawbacks: speculative decoding relies on an auxiliary "drafter" model, which can be challenging to acquire and increases memory overhead, while layer skipping may in- troduce discrepancies in the generated outputs due to the missing key-value cache at skipped layers. In this work, we propose AdaDecode, which accelerates LLM decoding without requir- ing auxiliary models or changes to the original model parameters, while ensuring output consis- tency. AdaDecode leverages the insight that many tokens-partinn particullly simple or predictbee ones can accurately be generated at intermediate layers, further layers often do not significantly alter predictions once the model reaches tain confidence. By adaptively generating tokens at intermediate layers when confidence is high, AdaDecode enables the next token's 's computation to begin immediately. The remaining layer com- tokens are deferred and executed in parallel with subsequent tokens when needed, maximizing hardware utilization and reducing decoding latency. A final verifi- cation step ensures that early predictions match the results of standard autoregressive decoding,
1. Introduction
The autoregressive decoding process in large language mod- els (LLLMs) is increasingly becoming a critical efficiency bottleneck for text generation (Khoshnoodi et al., 2024). As each token generation depends previously generated ones, the inherently sequential nature of this process severely lim- its parallelization capabilities on modern hardware (Miao et al., 2023a). This challenge is becoming more pressing due to two recent key trends. Firstly, LLMs continue to grow exponentially in size (Kaplan al., 2020; Hoffmmann et al., 2022; DeepSeek, 2024), with billions or even trillions of parameters (Achiam et al., 2023; Zeng et al., 2023; Dubey et al., 2024; Jiang et al., 2024; Anthropic, 2024), resulting in substantially more time-consuming and resource-intensive computations at each step of token generation. Secondly, model- generraee outputs are hecoming progressively longer, driven by emerging applications such as long-form content creation (Pham et al., 2024; Bai et al., 2024) and long chain- of-thought (CoT) reasoning (OpenAI, 2024; Snell al., 2024; Brown et al., 2024; Guan et al., 2025; Kimi et al., 2025; DeepSeek et al., 2025), which require massive infer- ence steps (Huang et al., 2025). The confluence of these factors leads to significant latency in text generation, ampli- fying the urgent need for more efficient decoding methods.
putations for
To accelerate autoregressive decoding, two primary ap- proaches have emerged: speculative decoding and layer skipping, as illustrated in Figure 1. Speculative decod- ing (Leviathan et al., 2023; Chen et al., 2023; Liu et al., 2023; Miao et al., 2023b; He et al., 2024; Huang et al., 2024; Li et al., 2024a;b; Du et al., 2024) employs a lightweight secondary model, called the "drafter", to generate candidate tokens at lower latency, which are then verified in paral- lel by the larger main model. However, the reliance
University of Virginia. Zhepei Wei <cheeeiiwei@@vigee virginia.edu>> Wei-Lin Chen wwwccen@virgimia.ddu edu>> Xinyu Zhu <xinyuzhu@ @viiginiaiidd,,, Yu Meng <yumeng3 @ virginia.edu>.
Correspondence to:
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Code and artifacts available at com/we i zhepei Decode

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Speculative Decoding
Layer Skipping
AdaDecode (Ours)
Verify paraliel
Skipped layers do not have KV cache
Same-colored layers in parallel
White layers follow normal decoding (no parallel)
f2
f4
f3
f4
Main Model
Venil
Verify
Run paralle!
E1
E2
t
t4
Run nn paralle!
Drafter
Same tokenizer
& vocabulary
13
t2
Generate sutoregressively
Figure 1. (Left) Speculative decoding relies on auxiliary drafter model, leading to increased memory usage and requiring the same tokenizer and vocabulary as the main model. (Middle) Layer skipping bypasses certain layers, which results in missing KV cache those layers and can introduce discrepancies in future token predictions. (Right) AdaDecode (Ours) accelerates decoding by adaptively predicting future tokens early hased on confidence (e.g., 12 and fo are predicted from different intermediate layers), enabling earlier progression to subsequent tokens. When future token steps require KV caches from the skipped layers (due to early predictions), these missing computations are executed in parallel with subsequent token processing (same-colored layers). A final verification step is employed to ensure output consistency with standard autoregressive decoding.
separate drafter model increases memory overhead and can be impractical in many cases, since the drafter must share the same tokenizer and vocabulary as the main model to ensure token compatibility. Layer skipping (Huang et al., 2018; Elbayad et al., 2020; Elhoushi et al., 2024; Del Corro et al,, 2023; Raposo et al., 2024; Geva et al., 2022; Din et al.,, 2023), in contrast, reduces computation cost by selectively bypassing certain layers during token generation. This ap- proach often requires designing new model architectures and intricate training methods (Elhoushi et al., 2024; Ra- poso et al., 2024). Although effective at reducing latency, layer skipping often leads to discrepancies in output qual- ity compared to standard autoregressive decoding (Schuster al., 2022; Liu et al., 2024a). Speeifically, skipped layers do not compute the key-value (KV) cache, which is essential for maintaining consistency in the model's predictions of future tokens. As result, while both speculative decoding and layer skipping provide promising speedups, they come with trade-offs that pose challenges for their widespread adoption in practice.
frozen. Our preliminary studies (Figure 3b) show that when predictions at intermediate layers are sufficiently confident, subsequent layers are unlikely to significantly alter the out- put. Based on this observation, AdaDecode predicts tokens using the hidden state at intermediate layers when confi- dence high, and immediately initiates processing of the next token. The remaining layers' KV cache computations are deferred and performed in parallel future decoding steps. This parallelism mitigates the sequential bottleneck of vanilla autoregressive decoding by processing multiple tokens simultaneously, thus maximizing hardware utiliza- tion and significantly boosting overall decoding throughput. Once all KV cache computations are complete, the standard autoregressive decoding result is obtained, allowing us to verify the correctness of early predictions. Compared to speculative decoding and layer skipping, AdaDecode ac- celerates autoregressive decoding while maintaining output parity, without requiring auxiliary models or modifications to the original model parameters.
Our contributions are as follows: (I) We propose a lightweight intermediate-layer LM head training approach that optimizes token prediction at early layers, enabling high-confidencce early predictions, with the original model parameters unchanged. Our lightweight LM head achieves performance comparable to fully parameterized LM heads from prior works (Stern et al., 2018; Cai et al., 2024), while using 31 x fewer parameters; (2) We introduce adaptive layer parallelism that concurrently processes multiple early- predicted tokens generated by the lightweight LM heads from different layers, significantly improving hardware uti- lization and decoding speed; and (3) AdaDecode demon-
In this work, we propose AdaDecode, a fast and accurate decoding method that accelerates autoregressive decoding through adaptive layer parallelism. AdaDecode builds the insight that many simple and predictable tokens can be accurately generated at intermediate layers, without requir- ing a full pass through all model layers (Schuster et al., 2022; Del Corro et al., 2023). To optimize token generation qual- ity at these layers, we introduce lightweight language model (LM) heads intermediate layers, trained to minimize the KL divergence between their predictions and those of the final layer, while keeping the original model parameters
2

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Table Comparison of AdaDecode with relevant methods. Drafting Strategy represents how the draft tokens are generated: Acceleration refers to the direction of acceleration: horizontal acceleration speedup generation across decoding time steps, while vertical acceleration focuses on speedups within each decoding time step. Early Exiting indicates how the model makes early predictions (e.g. exiting a fixed layer or dynamically exiting at different layers); Verification indicates whether verification step is employed to correct the draft token; Output Guarantee represents whether the output of the method is guaranteed to be consistent with the standard autoregressive decoding-Meese -Medusa violate the guarantec depending on its setup; # of Params refers to the number of parameters that need to be trained for enabling the method, which can impact its efficiency and effectiveness. The analysis of trainable parameters is based on Llama3. 3.1 -------InBttccs which has total of 32 layers with hidden size of 4096. For standard speculative decoding (SpecDec), separate draft model typically needs to be trained. and its number of trainable parameters may vary depending on the model architecture.
Method
Drafting Strategy
Acceleration
Early Exiting
Verification
Output Guarantee
of Params 8B
FREE (Bac et al., 2023)
Layer Skipping
Vertical
Fixed Depth
x
LITE (Varshney et al., 2023)
Layer Skipping
Vertical
Dynamic Depth
x
8B
SpeeDee (Leviathan et al., 2023)
Autoregressive
Standard
N/A
drafi model
LayerSkip (Elhoushi et al., 2024)
Layer Skipping
Vertical
Fixed Depth
X
8B
EESD (Liu et al., 2024b)
Layer Skipping
Vertical
Fixed Depth
0.7B
Medusa (Cai et ul., 2024)
Multi-hop Heads
Horizontal
N/A
/ & X 1.5B
EAGLE (Li al., 2024a)
Autoregressive
Horizontal
N/A
/
0.3B
AdaDecode (Ours)
Layer Skipping
Vertical
Dynamic Depth
/
48M
strates superior decoding throughput across various text generation tasks with up to 1.73x speedup, while ensuring identical outputs to standard autoregressive decoding.
final-layer LM head to intermediate layers results in mostly low predicted probabilities for the tokens generated by the model, making early predictions with standard LMs chal- lenging. Hence, prior research on carly exiting often in- volves designing specific model architectures or fine-tuning existing models enable intermediate- layer predictions (El- houshi al., 2024; Din et al., 2023; Del Corro et al., 2023). As result, these acceleration methods typically fail to produce outputs consistent with those of the original off- shelf LMs due to changes in architecture and parameters.
2. Method: A AdaDecode
In this section, we present our proposed method AdaDecode, as illustrated in Figure 2. The core concept is to start pro- cessing the initial layers of subsequent tokens early in an adaptive manner, while completing the remaining layers of all early-predicted tokens in parallel, thereby enhancing overall decoding throughput via parallel computation. We introduce the techniques for enabling early predictions using intermediate layer representations in Section 2. followed by detailed explanation of our adaptive layer parallelism processing approach in Section 2.2. Note that in this work, we focus on vertical acceleration within each decoding step, which is orthogonal to methods that target horizontal accel- cration across decoding time steps. Moreover, our approach preserves the exact output of standard autoregressive de- coding, with consistency guaranteed by verification step. However, not all acceleration methods that incorporate ver- ification can ensure such output parity. A comprchensive comparison with relevant methods is provided in Table 1.
Training intermediate-layer LM heads with original model parameters frozen. We hypothesize that many intermediate-layer representations may already contain suf- ficient information for predicting the next token, but the original LM head cannot directly harness their potential. To facilitate early predictions using intermediate- repre- sentations without further fine-tuning them, we introduce trainable LM heads g (e,") at each candidate early prediction layer 7(0) They take the hidden represen- tations h ( at layer ((a) as frozen features and predict the next word distribution Pæte: (1 (+), which are trained to approximate the last-layer prediction p² (t(h") (h" is the last- hidden states) by minimizing the following KL divergence loss:
2.1. Lightweight LM Heads Enable Early Predictions
exp exp(e,) (0)
Pars ())
Off-the-shelf LMs struggle with early predictions. Many tokens in natural language, such as stopwords and common phrase completions, can be easily predicted and do not re- quire the full capacity of a model for accurate generation. However, off-the-shelf LMs typically have difficulty uti- lizing intermediate layers for next-token prediction, as the final-layer LM head is not trained work with intermediate- layer representations. As shown in Figure 3a, applying the
Evey exp (et (1) (2)))
((0(0)) = KL (p' (p"(t/h")llpe (tl(())
Since we do not update h (8) the original model parameters remain unchanged, and only the newly added LM heads
are trained. As demonstrated in Figure 3c, training these intermediate-layer LM heads results in good approxima-
3

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
AdaDecode:
AdaDecode vs. Vanilla Decoding
Early predictions speed up decoding Venification ensures accuracy
Run in paralle!
Run in paralto!
AdaDecode (Oura)
Henty
Venly
Decaded tokens
1.3
Layer
Layer
Layer
f/2
Layer
Layer
Layer
Verify fg
Run parallel
11
Layer
Layer
Layer
Verify
Eady pedictioon
Run parallel
Eanly predction
Vanilla Decoding
Layer
Layer
Layer
t1
Layer 1
Layer 2
Layer 3
Token processing across layers
Figure 2. (Left); A simplified example demonstrating how early predictions enable parallelization. Here, t1 triggers the early prediction of to at Layer 2, and f2 triggers the early prediction of i3 at Layer 1. As result, the Layer 2 computations of t2 and followed by parallelized Layer 3 computations for t1, 12, and ta. (Right): Vanilla autoregressive decoding processes tokens strictly in sequence, limiting opportunities for parallelization. In contrast, AdaDecode adaptively starts processing the first few layers of the next token once the model confidently makes early predictions using intermediate LM heads. The remaining layers for all early-predicted tokens then computed in parallel, accelerating the decoding process.
run in parallel,
Who was the first American to win the Nobel Prize?
Who was the first American to win the Nobel Prize?
Layer S
32
7.5
Laver
0.75
0.75
Layer Loss
24
24
5.0
0.50
16-
0.25
2.5
8
The first
Americanco
win
Nobel Prize
The first
Americanii
win
Nobel Prize
10
20
30
40
50
Generated Tokens
Geverated Tokens
#
Epocks
(a) w/ original LM head.
(b) w/ trained LM heads.
(c) KI. loss curve.
Figure 3. Probabilities of model-ggnnerate tokens predicted at the 8th, 16th, 24th, and 32nd (final) layers of the fine-tuned Llama-3.1-88 1-8B- Instruct shown using the original last-layer language model (I.M) head and (b) newly introduced lightweight LM heads. These new LM heads are trained to minimize the KL divergence loss relative to the final layer predictions, while keeping all original model parameters frozen. The LM heads enable close approximation the final predictions, as seen in (b), where many tokens have contident predictions at intermediate layers, and in (c), where the KI. divergence loss relative to the final layer predictions is minimal.
tions of the last-layer outputs, as indicated by the low KL divergence loss at the end of training. This supports our hypothesis that intermediate-layer representations contain ample information for predicting the next token, and simple transformations via new LM heads can effectively extract this information. Consequently, with our trained LM heads, many tokens' predicted probabilities are notably high at intermediate layers, as shown in Figure 3b.
states h (3) to compute the probability over the vocabulary Pat Softmax (h (7) E(1)T ). reduce the parameter cost, we decompose it as E = E'T((() where E is the last- layer LM head weights, and (T()) E BBddd is a learnable transformation matrix. As d < | for most LLMss, learning ( is much more parameter-efficiee
than learning E(9) directly. We defer the proof that learning ( retains full expressiveness of learning E(() to Appendix A. In practice, in Llama3.1 InStrrrruuct model (Dubey al., 2024), each lightweight LM head ( introduces only 16M parameters (d = 4096). while a full hall-parameterized LM head E would require 0.5B parameters (DV) 128K). To further reduce the number of trainable parameters, low-rank adaptation methods (Hu et al., 2022) can be explored, which we leave for future work.
Lightweight LM head implementation. The LM heads g( -el tevv are typically represented by a weight matrix (()) E R|Vixd where d is the model dimension size (Stern et al., 2018; Cai et al., 2024). Given the large vocabulary size of LMs, learning a separate LM head for each carly prediction layer leads to substantial increase in the number of parameters. Based on the observation that the weight matrix O is always applied the hidden

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Algorithm 1: LLM Decoding Acceleration via Adaptive Layer Parallelism (AdaDecode)
Input: User prompt æ = [T,,. [MM]
Parameter: Early prediction threshold ?
Output: Generated output sequence y = (L,, U1, 12,...,,x,) Iv]
initiaaliz KV cache by processing user prompt // initialize output sequence with [BOS] token #1 P tracks the list of tokens be processed in parallel each layer / while 1, / [FOS] // terminate generation upon generating FOS token do
KV cache +- (x+)
i 6 0; 14 f [BOS]: y - [t:]
P P { }: P[[] t []. L
for 1 1 S. L // process t; from the first to the last layer do
If add current token the parallel processing list layer update KV cache layer / by processing all tokens in P[]] in parallel
P[[] + P[[] @ t;
fa+1 ~ Potc (t' (t
// sample from the intermediate LM head at layer I
if Pares [t++] 1h ( > y if probability surpasses the threshold then
y 1- y 1 1i41
|/ append f,+1 to output sequence || add t, to the parallel processing list of all deeper layers If immediately start processing the next token
break
if L the final layer is reached then
Vt E P[[], accept or reject based on Eq. (2)
I/ verify all tokens in P[[] in parallel // empty P[4
P[]] <-
4- 4 44+1
// remove rejected tokens, append 4+11 to output sequence, and resume from the correct token
8 6 i i+11
return y
2.2. Adaptive Layer Parallelism
then verifies the early predictions using a modified rejection sampling scheme (Chen et al., 2023; Leviathan et al., 2023). Specifically, we accept the early predicted token t, sam- pled from the intermediate layer's distribution Pota (t' h(())) with probability:
Early predictions trigger parallel processing. As shown in Figure 3b, when token's 's predicted probability is suf- ficiently high with intermediate LM heads (introduced in Section 2.1,, subsequent layers are unlikely to change the predictions significantly. Based on this observation, we generate the next token at layer ((0) when its probability surpasses a predefined threshold:
min {
D (t(t'(n") Peto (t' h (1)
1.
(2)
If is accepted, we move on to the verification of the next early-predicted token, until all early predictions have been accepted or a token gets rejected. When an early prediction is rejected, we resample a token 1 from the adjusted distribution Normalize(maxx( p= (t'lh") Poc (t' (t'( as a replacement, remove the KV caches for tokens after f, and resume generation from t" onward.
t -~ Poto (t'(he (0) and Pato (t|h (t)
>
where y is a hyperparameter, and any sampling strategy can be employed (e.g., greedy or nucleus sampling (Holtzman et al., 2020)). This early prediction can adaptively happen at different layers (i.e., any layer with an intermediate LM head) once the probability surpasses 7, allowing flexible parallel calculation of KV caches across multiple tokens. Notably, it is necessary to finish processing the subsequent layers of the early predicted token to obtain their KV cache- omitting this step would result in a missing KV cache at deeper layers, which would cause inconsistencies when computing future token representations.
Although rejecting early predictions can lead to some wasted computation, we observe in practice that this happens rarely. For example, as indicated in Figure 5c, with 7 0.85, only about 6% of all early predictions are rejected, which sup- ports our observation that high-confidence predictions from intermediate layers tend to align with the final predictions.
Early prediction verification. Regardless of the thresh- old hyperparameter y set in Equation (1), there is always possibility that the early predicted token from intermediate layers differs from the final prediction. To ensure consis- tency with standard autoregressive decoding, we introduce a verification step that first completes the KV cache of the re- maining layers for all early- early-predicted tokens parallel, and
Overall algorithm of AdaDecode is presented in Algo- rithm 1.
3. Experimental Setups
Backbone models and evaluation tasks. We evaluate our method on a diverse set of text generation tasks, including
5

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Table
Throughput (average tokens per second) and speedup comparison between AdaDecode and haseline decoding methods across three tasks, including text summarization, code generation, and mathematical reasoning. All compared methods ensure generation parity with vanilla autoregressive decoding, and speedup results are computed relative to benchmarks established on the vanilla setup (i.., speedup 1 x for vanilla autoregressive decoding). The best performance is highlighted in bold.
Text Summarization
Code Generation (HumanEval)
Mathematical Reasoning
(XSum)
(GSM8K)
Method
Throughput (Tokens/s)
Speedup (vs. Vanilla)
Throughput (Tokens/s)
Speedup (vs. Vanilla)
Throughput (Tokens/s)
Speedup (vs. Vanilla)
Dlama3. 1-88
Vanilla Decoding
33.31
1.00
32.58
1.00x
33.13
1.00
SpecDecode (Leviathan al., 2023)
w/ drafter Llama 3.2...2.. 2211Bma 11ar
35.64
1.07x
46.26
1.42x
45.38
1.37x
AdaDecode (Ours)
38.09
1.14x
49.21
1.51
49.17
1.48
CodeLlama-13B:n1
Vanilla Decoding
27.80
1.00
27.55
1.00x
28.25
1.00
SpecDecode (Leviathan et al., 2023)
w/ drafter Codellama waa
26.97
0.97x
29.75
1.08
28.53
1.01
Self-SpeeDDecode (Zhang et al., 2024a)
28.63
1.03
31
1.14x
31.64
1.12x
LookAhead (Fu al., 2024)
33.08
1.19x
41.04
1,49x
38.42
1.36x
SWIFT (Xia et al., 2025)
30.02
1.08x
36.64
1.33x
30.51
1.08x
AdaDecode (Ours)
37.99
1.37
46.78
1.69x
44.28
1,57
CodeLlama-34B...
Vanilla Decoding
17.68
1.00
18.91
1.00x
19.16
1.00
SpecDecode (l.eviathan et al., 2023)
w/ drafter CodeLLama- ser
19,09
1,08x
26.66
1,41x
24.14
1.26x
w/ drafter Codel.
:
20.86
1.18x
23.25
1.23x
21.07
1.10x
Self-SpeeDeecode (Zhang et al., 2024a)
18.97
1.07
21.55
1.14x
21.84
1.14x
LookAhead (Fu al., 2024)
20.15
1.14X
26.28
1.39x
27.01
1,41x
SWIFT (Xia et al., 2025)
21.92
1.24x
26.47
1.40x
25.29
1.32x
AdaDecode (Ours)
24.35
1.38x
32.78
1.73x
30.68
1.60
text summarization (i.e., XSum (Narayan et al., 2018)),
will demonstrate later ($ 4.4), they can only lead quite limited speedup even negative speedup results compared to standard decoding unless careful hyperparameter tuning is performed. A more detailed introduction of baselines and implementation details is provided in Appendices and D.
code generation (i.e., HumanEval (Chen al., 2021)), and mathematical reasoning (i.e., GSM8K (Cobbe et al., 2021)), covering broad spectrum of language model capabilities. We employ three instruction-tuned models as the backbone, ranging from 8B to 34B parameters (i.e., LLama3. 1--BB Instruct, Codel Jama- -ama-113--nsstttt Instrutt, and CodeLlama-34B- Instruct), with the default chat template for decoding. Due to the page limit, please refer Appendix B for detailed introduction of the benchmarks.
4. Results
4.1. Main Results
AdaDecode consistently achieves superior inference speedup across all benchmarks. To validate the effective- ness of our method, we compare the proposed AdaDecode with state-of-the-an efficient decoding methods across a wide range of challenging text generation tasks. As pre- sented in Table 2, our method consistently delivers superior speedup compared to all baseline approaches regardless of the backbone model size, achieving up 1.73x speedup compared to standard autoregressive decoding.
Baselines. In our work, we primarily focus on compar- ing with efficient decoding baselines that provide output parity guarantees with standard autoregressive decoding techniques, including speculative decoding (Leviathan et al., 2023; Chen al., 2023), LookAhead (Fu al., 2024), self- speculative decoding (Zhang al., 2024a) and its variant SWIFT (Xia et al., 2025). Despite their conceptual advan- tages, these methods have practical limitations. They often come with inherent constraints regarding model selection or necessitate task-specific model architectures. Such re- quirements significantly compromise their practical utility and present difficulties for real-world application- as we
Speculative decoding (mostly) performs better when as- sisted with smaller models. Table 2 shows that a smaller drafter model (e.g., Codel IB--aammt generally leads
6

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Table Ablation study on AdaDecode. We report the performance and required additional parameters of our method by analyzing the impact of verification, adaptive layer prediction, and ablating the lightweight LM head. We report the speedup and consistency ratio (measured by string match between generation results) compared to vanilla autoregressive decoding of Llama3. 1-B.. -8--Instruct on code generation task (vala
Method
# New Heads
# New Params
Consistency Ratio
Speedup
AdaDecode
48M
0.996
1.51 x
w/o verification
48M
0.652
1.64x
w/ fixed-layer early prediction
1
16M
0.996
1.37
w/ original LM head
0M
0.995
0.84x
w/ mixed-domain LM head
48M
0.998
1.29
w/ full paratttrrreerr LM head
1.5B
0.997
1.49x
1.0
0.97
0:98
0.39
2.0
not yield sufficient speedup, as the performance gains stem primarily from reducing the latency of skipped layers.
Consistency Ratio
1.75
LookAhead requires surplus FLOPs to achieve meaning- ful speedups on larger models. Unlike standard specula- tive decoding methods, LookAhead achieves higher speedup the 13B model than the 34B model. This is be- cause LookAhead employs Jacobi decoding, which gen- erates multiple disjoint n-grams in parallel within a single step. While this approach reduces the number of decoding steps, it comes at the cost of increased FLOPs. In situations where hardware FLOPs become the bottleneck, especially with large models like the 34B, the surplus FLOPs needed for parallel n-gram generation may not be satisfied, thus leading diminished speedup.
Speetuu
0.5
1.5
1.41
1.40
1.14
1.0
SpecDes Self SperDe: LookÅheod
SWIFT
AdaDecode
Figure 4. Consistency ratios and speedups of different methods vs. vanilla decoding on HumanEval with CodeLlama-344 34B-Instruet
to higher speedup for SpecDecode compared to a moderate- size model (e.g., CodeLlama- 13B-IIstruct) This is mainly because smaller drafters have fewer parameters, requiring less computation and inference time to generate draft to- kens, which are then passed to large-scale verifier (e.g., CodeLlama-34BI for parallel verification and correction, significantly reducing the overall time. One notable ex- ception is in the text summarization task, where using the smaller model as the drafter yields a lower speedup than the moderate counterpart. We speculate that this is because smaller models like CodeLlama-77 In-Irrcttt have limited capacity to handle extremely long texts, thus producing low- quality draft tokens, which lead higher rejection ratio, thereby increasing the overall latency.
4.2. Output Parity with Vanilla Decoding
AdaDecode guarantees output parity with standard au- toregressive decoding. In principle, all compared methods in our work are designed to guarantee output consistency with standard autoregressive decoding as a result of the verification steps. However, due to numerical precision in- accuracies and potentially tied probabilities during the com- putation, the generation results might vary in practice and are subject experimental environment and hardware spec- ifications. As shown in Figure 4, empirically evaluate the output consistency ratio of all baseline methods against the vanilla decoding and report their corresponding speedup ratios on the HumanEval benchmark with Coode-lama--344 Instruct as the backbone. Despite slightly deviating from the- oretical guarantees, all of these methods empirically achieve nearly 100% consistency ratio against the standard autore- gressive decoding. One potential reason is the numerical precision inaccuracies during inference, as we use FP16 pre- cision for inference. which may cause the model to select slightly different tokens compared to standard decoding.
SelfoS and its variant SWIFT tend to achieve higher speedups as the model size increases. By skip- ping a larger portion intermediate layers, such methods can generate draft tokens much faster than the full model, thereby achieving significant speedup. However, this ap- proach requires task-specific model architecture configu- ration to achieve meaningful speedups. This is because skipping too many layers negatively impacts the quality of the draft tokens, which will lead to a high rejection rate during verification, and consequently increase the decod- ing latency. On the contrary, skipping too few layers does
Pnttps: http::::////ttt////////ee com/ huggingiace/ transformers iisuussssssssseeeeeee massuus/3041 /0013

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Ours Vinila Devoding
Early Predlicton Rale
60-
Verilicatio Bejection Rale
1.50
Percentage (96)
100
Peronnaage (%6
40-
Speedup
1.25
50-
20
1.00
0.2
0.6 threshold ?
0.4
0.8
0.2
0.4
0.5
0.8
0
0.2
0.4
0.6
0.8
throshold
thrrehhodd
(a) Inference speedup.
(b) Early prediction rate.
(c) Verification rejection rate.
Figure 5. Hyperparamter study of AdaDecode by varying the early prediction threshold These figures present the evaluation results on HumanEval, including (a) the speedup curve, (b) the early prediction rate, and (c) the verification rejection rate.
4.3. Ablation Study
Our lightweight LM head is sufficient for early predic- tion compared to the full-parameterized LM head. We also implement variant of our method by training full- parameterized intermediate LM heads for carly prediction. Despite significantly increasing the number of hyperparam- eters, its speedup does not gain further boost compared to our lightweight head. This echoes our hypothesis that intermediate-layer representations contain enough informa- tion for predicting the next token, and simple transforma- tions via lightweight LM heads can effectively extract it.
As shown in Table 3, we conduct a comprehensive ablation study to understand the effectiveness of each module in our proposed AdaDecode framework.
The verification step is essential for ensuring output con- sistency. The first row shows that removing the verification step from AdaDecode significantly reduces the consistency ratio, even though it achieves a slightly higher speedup. Without verification, there is no control over the generated tokens, which can lead AdaDecode producing entirely different content as early predictions may not always be re- liable. This results in the acceptance of potentially incorrect early predicted tokens and causes significant deviations from standard autoregressive decoding. This finding underscores the importance of the verification step in our method.
4.4. Hyperparameter Sensitivity Study
Despite the encouraging numbers reported in the literature, replicating them is not trivial in practice, they typically require extensive hyperparameter search. To ensure robust- ness, we extensively study the performance of these meth- ods. Below, we present the results of AdaDecode, and the analysis for baselines can be found in Appendix C.
daptive early prediction allows for more flexible layer parallelism. With fixed-layer early prediction, we only need one intermediate LM head make early predictions at fixed layer. While this requires less memory overhead, it negatively impacts the speedup compared to the full-fle full-fledged AdaDecode that allows early prediction at different layers, demonstrating the importance of adaptive early prediction.
AdaDecode is robust to hyperparameter for early pre- dictions. As introduced in Equation (1), the hyperparameter 7 controls early predictions and triggers adaptive layer par- allelism. To study its impact, we conduct experiments with and Figure 5 presents the speedup results, early prediction rate, and verification rejection rate, with varying values of 7 = [0, 0.2, 0.4, 0.6, 0.8, 1]. Specifi- cally, 7 = essentially means no early predictions, as no token can have probability greater than 1, while 7 0 triggers early predictions every inference step. It can be observed that the early prediction rate decreases with increasing 7, while the verification rejection rate also de- creases. The inference speed generally increases with the threshold unless becomes excessively high. This find- ing suggests that AdaDecode works reliably as long as the threshold is within a reasonable range. We attribute this to the fine-tuned lightweight LM head, which enables early predictions with high confidence, resulting in lower verifi- cation rejection rate (e.g., around 5% when ? = 0.85) while maintaining comparatively high early prediction rate, thus achieving meaningful speedups.
Llama3. 1-8B-Innstenn
Off-the-shelf original LM head for early prediction will not lead to speedup. In this study, we replaced the trained lightweight LM head with the original final-layer LM head for early prediction. Although additional parameters are required, it substantially compromises speedup because ap- plying the original LM head to intermediate layers for early prediction only produces very low confidence (as shown in Figure 3a), thus slowing down the decoding process.
Specialist LM head is better than generalist LM head. By default, the intermediate lighter LM heads are trained for specific tasks using the training split in each benchmark. To test their effectiveness, we also train the LM heads on a mixed set of training data from all benchmarks, cover- ing multiple domains. However, the ablation study shows that this significantlly slows down the generation speed, sug- gesting that AdaDecode works more efficiently when the intermediate heads are trained in in-domain
8

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
5. Related Work
allelism with our newly introduced lightweight LM heads. A verification step ensures consistency with standard au- toregressive outputs. Notably, it demonstrates consistent improvements in decoding throughput across various gener- ation tasks compared to baselines, while imposing minimal constraints (e.g., no auxiliary model or fine-tuning of ex- isting model parameters required). Further discussion on limitations and future work is presented in Appendix F.
5.1. Early Exiting
Early exiting enables language models to complete predic- tion at intermediate layers, reducing computational over- head and accelerating generation. Previous approaches achieved this by adding decision branches language mod- eling heads at various depths (Teerapittayanon et al., 2016; Huang et al.. 2018; Elbayad et al., 2020; Schuster et al., 2022; Yang et al., 2024; Raposo et al., 2024; Bae et al., 2023; Varshney et al., 2023). Most of these methods alter both decoding speed and output quality, while our method strictly preserves the exact output of standard autoregres- sive decoding. Notably, similar to method, BESD (Liu et al., 2024b) applies early exiting with verification step to ensure output parity. However, it is limited to fixed-depth exiting and requires an extra decoder layer and full LM head, whereas AdaDecode uses only lightweight heads and allows dynamic-depth exiting, making it more efficient and flexible. Moreover, EESD trains its auxiliary modules us- ing cross-entropy loss and applies Thompson Sampling for stopping decisions. In contrast, AdaDecode optimizes its lightweight heads using KI. divergence and employs sim- ple thrpliilyttrrshoddiig strategy for draft termination. We refer the readers to Khoshnoodi et al. (2024) for a more detailed discussion.
Acknowledgments
The authors would like to thank Tianyu Gao and Mengzhou Xia from Princeton NLP group for their valuable feedback and discussions. This research was supported in part by the NVIDIA Academic Grant. We thank anonymous re- viewers and area chair for their constructive and insightful comments.
Impact Statement
Our work focuses on improving the efficiency of autore- gressive LLM decoding, a well-established objective within the machine learning community. While AdaDecode does not introduce novel ethical concerns beyond those already associated with autoregressive LLMs, it remains crucial to approach the development and deployment of language models responsibly, with careful attention to issues such as fairness, bias, and potential misuse.
5.2. Speculative Decoding
Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) has emerged as an effective approach for accelerating generation. Following the standard speculative decoding approach, most works in this line attempt to draft single chain-based sequence of draft tokens and verify them in parallel (Kim et al., 2023; Hooper et al., 2023; He et al., 2024; Fu et al., 2024; Liu et al., 2024a;b; Qin et al., 2024; Zhang et al., 2024a; Elhoushi et al., 2024; Xia et al., 2024a). Our method also falls into this category and in this work, we only compare with baselines that guarantee output parity with standard decoding. Another line of research tree- based speculative decoding, where multiple draft sequences generated and verified in parallel using tree-based attention mechanism (Miao et al., 2023b; Cai et al., 2024; Ankner et al., 2024; Li et al., 2024a). While orthogonal to our main contribution, we further integrate adaptive layer parallelism to tree-based decoding and share our findings in Appendix E, which we believe merit further investigation.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Ankner, Z., Parthasarathy, R., Nrusimha, A., Rinard, C., Ragan-Kee Kelley, J., and Brandon, W. Hydra: Sequentially- dependent draft heads for medusa decoding. In First Conference on Language Modeling, 2024.
Anthropic.
Anthropie: Introducing claude 3.5 son- net, 2024. URL https:/ //wwwwww.ooooo///oo anthropic.com. news 1llaude-3-5-sonme
Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., et al. Program synthesis with large language models. arXiv 2021.
preprint arXiv:2108.077
6. Conclusion
Bae, S., Ko, J., Song, H., and Yun, S.-Y. Fast and robust carly-cxiting framework for autoregressive language mod- els with synchronized parallel decoding. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp- 5910-5924, 2023.
We introduced AdaDecode, a new approach designed to ac- celerate L.M decoding while preserving output consistency. AdaDecode achieves this by adaptively generating the next tokens at intermediate layers using an adaptive layer par-

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Bai, Y., Zhang, J., Lv, X., Zheng, L., Zhu, S., Hou, L.., Dong, Y., Tang, J., and Li, J. LongWritee Unleashing 10,000+ word generation from long context Ilms, arXiv preprint arXiv:2408.07055 2024.
Din, A. Y., Karidi, T., Choshen, L., and Geva, M. Jump conclusions: Short-cutting transformers with linear transformations, ArXiv, 0tss/2333.09433. 2023.
Du, C., Jiang, J., Xu, Y., Wu, J., Yu, S., Li, Y., Li, S., Xu, K., Nie, L., Tu, Z., and You, Y. Glide with a cape: A low-hassle method to accelerate speculative decoding. In ICML, 2024.
Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Lc, Q. V., Ré, C., and Mirhoseini, A. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint urXiv:2407211787, 2024.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Llama herd of models. arXiv preprint 2024.
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. Medusa: Simple LLM inference acceleration framework with multiple decoding heads. In Forty-first International Conference on Machine Learning, 2024.
Elbayad, M., Gu, J., Grave, E.,, and Auli, M. Depth-adaptive transformer. In International Conference on Learning Representations, 2020.
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318. 2023.
Elhoushi, M., Shrivastava, A., Liskovich, D., Hosmer, B., Wasti, B., Lai, L., Mahmoud, A., Acun, B., Agarwal, S., Roman, A., Aly, A., Chen, B., and Wu, C.-J. Layer- Skip: Enabling early exit inference and self-speculative decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 12622-12642, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
Chen, L., Qiu, R., Yuan, S., Liu, Z., Wei, T., Yoo, H., Zeng. Z., Yang. D., and Tong, H. WAPITI: A water- mark for finetuned open-source LLMs. arXiv preprint arXiv:2410.06467, 2024.
Chen, M., Tworek, J., Jun, H., Yuan, Q., Pinto, H. O.D. Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021.
Fu, Y., Bailis, P., Stoica, I., and Zhang, H. Break the se- quential dependency of LLL inference using lookahead decoding. In Forty-first International Conference on Ma- chine Learning, 2024.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10..11661.411. 2021.
Geva, M., Caciularu, A., Wang, K., and Goldberg, Y. Trans- former feed-forward layers build predictions by promot- ing concepts in the vocabulary space. In EMNLP, 2022.
Guan, X., Zhang, L. L., Liu, Y., Shang, N., Sun, Y., Zhu, Y., Yang, F,, and Yang, M. rStar-Math: Small LLMs can master math reasoning with self-evolved deep thinking. arXiv preprint arXiv:2501 04519,, 2025.
Dao, Flashattention-22 2: Faster attention with better paral- lelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
He, Z., Zhong, Z., Cai, T., Lee, and He, D. REST: Retrieval-based speculative decoding. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (iiiesssVoe (Volume 1: Long Papers), pp. 1582-1595, 1595, 2024.
DeepSeck. DeepSeck LLM: Scaling open-source language models with longtermism. ArXiv, hiss/2001.029544 2024.
DeepSeek, Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., et al. DeepSeek-RI: Incentivizing reasoning capabil- ity in LLMs via reinforcement learning. arXiv preprini ar arXiv:2501.12948 2025.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L. Training compute-optimal large language models. ArXiv, abs/2203. 15556,, 2022.
Del Corro, L., Del Gioro, A., Agarwal, S., Yu, B., Awadal- lah, A., and Mukherjee, S. Skipdecode: Autoregressive skip decoding with batching and caching for efficient LLM inference. arXiv preprint arXiv:2307.02628 2023.
Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. QLoRA: Efficient finetuning of quantized LLMs. In NeurlPS, 2023.
Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. In ICLR, 2020.
10

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Hooper, C.,, Kim, S., Mohammadzadeh, H., Gene, H., Keutzer, K., Gholami, A., and Shao, S. Speed: Spec- ulative pipelined execution for efficient decoding, arXiv preprint arXiv:2310.1200 2023.
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Inter- national Conference Machine Learning, pp. 19274- 19286. PMLR, 2023.
Hu, J. E., Shen, Y., Wallis, P., Allen-Zhu Z., Li, Y., Wang, S., and Chen, W. LoRA: Low-rank adaptation of large language models. In ICLR, 2022.
Li, Y., Wei, F., Zhang, C., and Zhang, H. EAGLE: Specula- tive sampling requires rethinking feature uncertainty. In Forty-first International Conference on Machine Learn- ing. 2024a.
Huang, C., Huang, L.,, Leng, J., Liu, J., and Huang, J. Effi cient test-time scaling via self-calibration. arXiv preprint arXiv:2503.0000 2025.
Li, Y.. Wei, F., Zhang, C., and Zhang, H. EAGLE-2: Faster inference of language models with dynamic draft trees. In Proceedings of the 2024 Conference on Empirical Meth- ods in Natural Language Processing, PP. 7421-7432, Miami, Florida, USA, November 2024b. Association for Computational Linguistics.
Huang, G., Chen, D., Li, T., Wu, F., van der Maaten, L.,, and Weinberger, K. Multi-scale dense networks for resource efficient image classification. In International Conference on Learning Representations. 2018.
Liu, F., Tang. Y., Liu, Z., Ni, Y., Han, K., and Wang, Y. Kangaroo: Lossless saff-speculative decoding via double early exiting, arXiv preprint arXiv:2404. 18911, 2024a.
Huang, K., Guo, X., and Wang, M. Specdec++: Boost- ing speculative decoding via adaptive candidate lengths. arXiv preprint arXiv:2405.1972
2024.
Liu, J., Wang, Q., Wang, J., and Cai, X. Speculative de- coding via early-exiting for faster LLM inference with Thompson sampling control mechanism. In Findings of the Association for Computational Linguistics ACL 2024, pp. 3027-3043, Bangkok, Thailand and virtual meeting, August 2024b. Association for Computational Linguis- tics.
Huang, L.,, Yu, W., Ma, W., Zhong, W., Feng, Z., Wang, H., Chen, Q., Peng, W., Feng, X., Qin, B., and Liu, T. A sur- vey hallucination in large language models: Principles, taxonomy, challenges, and open questions. ArXiv, 2023.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. I., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint ar arXiv:2401.04080 2024.
Liu, X., Hu, L., Bailis, P., Stoica, I., Deng, Z., Cheung, A., and Zhang, H. Online speculative decoding. arXiv preprint arXiv:2310.07177 2023.
Kanai, S., Fujiwara, Y., Yamanaka, Y., and Adachi, S. Sig- softmax: Reanalysis the softmax bottleneck. In Neural Information Processing Systems, 2018
Liu, Z., Amjad, R. A., Adkathimar, R., Wei, T., and Tong, H. SelfElicit: Your language model secretly knows where is the relevant evidence. arXiv preprint arXiv:2502.08767, 2025.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint T0t:22000.0836.,, 2020.
Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Jin, H., Chen, T., and Jia, Z. Towards efficient generative large language model serving: A survey from algorithms systems. arXiv preprint urXiv:2312.152333 2023a.
Khoshnoodi, M., Jain, V., Gao, M., Srikanth, M., and Chadha, A. A comprehensive survey of accelerated gener- ation techniques in large language models. arXiv preprint arXiv:2405. 13000,, 2024.
Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Zhang, 7.,, Wong, R. Y. Y., Zhu, A., Yang, I.,, Shi, X., et al. SpecInfer: Accelerating generative large language model serving with tree- besed speculative inference and verification. arXiv preprint arXiv:2305.0m 330::2335500781. 2023b.
Kim, S., Mangalam, K., Moon. S., Malik, J., Mahoney, M. W., Gholami, A., and Keutzer, K. Speculative decod- ing with big little decoder. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
Narayan, S., Cohen, B., and Lapata, M. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Proceed- ings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 1797-1807, 2018
Kimi, Du, A., Gao, B., Xing, B., Jiang, C., Chen, C., Li, C., Xiao, C., Du, C., Liao, C., et al. Kimi K1.S: Scal- ing reinforcement learning with LLMs. arXiv preprint arXiv:2501. 12599, 2025.
OpenAL.
Introducing OpenAl ol: Learn- ing to reason with large language models, 2024.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1442.0000 142,6690, 2014.
URL https::/ /oppos/// com/ index/ -earninn--o----esoon--ii

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
Pham, C. M., Sun, S., and lyyer, M. Suri: Multi-constraint instruction following for long-form text generation. arXiv preprint arXiv:2406. 0224406617777 19371, 2024.
Yang, S., Lee, G., Cho, J., Papailiopoulos, D., and Lee, K. Predictive pipelined decoding: A compute-latency trade- off for exact LLM decoding. Transactions Machine Learning Research, 2024. ISSN 2835-8856.
Qin, Z., Hu, Z., He, Z. Prakriya, N., Cong, J., and Sun, Y. Optimized multi-token joint decoding with auxiliary model for LLM inference. arXiv preprint MxXiv24007009722. 2024.
Yang. Z., Dai, Z., Salakhutdinov, R., and Cohen, W. W. Breaking the softmax bottleneck: A high-rank RNN lan- guage model. In ICLR, 2018.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Per- formance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Yu, E., Li, J., Liao, M., Wang, S., Zuchen, G., Mi, F. and Hong, L. CoSafe: Evaluating large language model safety in multi-turn dialogue coreference. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 17494 17508, 2024.
Zeng, A., Liu, X.,, Du, Z.. Wang. Z., Lai, H., Ding, M., Yang. Z., Xu, Y., Zheng, W., Xia, X., Tam, W. L., Ma, Z., Xue, Y.. Zhai, J., Chen, W., Liu, Z., Zhang, P., Dong, Y., and Tang, J. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations, 2023.
Raposo, D., Ritter, S., Richards, B., Lillicrap. T., Humphreys, P. C., and Santoro, A. Mixture-of-Deppe Deeth::: Dynamically allocating compute in transformer- based lan- guage models. arXiv preprini arXiv:2404.02258, 2024.
Schuster, T., Fisch, A., Gupta, J.. Dehghami, M., Bahri, D., Tran, V., Tay, Y., and Metzler, D. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456-1747 2022.
Zhang, J., Wang. J., Li, H., Shou, L., Chen, K., Chen, G., and Mehrotra, S. Draft& verify: Lossless large language model acceleration via self-speculative decoding. In Pro- ceedings of the 62nd Anmual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 11263-1 11222, Bangkok, Thailand, August 2024a. Association for Computational Linguistics.
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling LLM test- time compute optimally can be more effective than scal- ing model parameters. arXiv preprint 2.Xiv:24080033144 2024.
Stern, M., Shazeer, N., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018.
Zhang, Z., Lei, L., Wu, L., Sun, R., Huang, Y., Long, C., Liu, X., Lei, X., Tang, J., and Huang, M. BafettBench: Evalu- ating the safety large language models with multiple choice questions. In ACL, 2024b.
Teerapittayanon, S., McDanel, B., and Kung, H.-T. Branchynet: Fast inference via early exiting from deep neural networks. In 2016 23rd international conference on pattern recognition (ICPR), pp. 2464-2469. IEEE, 2016.
Varshney, N., Chatterjee, A., Parmar, M., and Baral, C. Accelerating Llama inference by enabling intermediate layer decoding via instruction tuning with LITE. arXiv preprint ar arXiv:2310.1 (tt.:2330./888. 2023.
Xia, H., Yang, Z., Dong, Q., Wang. P., Li, Y., Ge, T., Liu, T., Li, W., and Sui, Z. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. arXiv preprint arXiv:2401.078: 078511 2024a.
Xia. H., Li, Y.. Zhang, J., Du, C., and Li, W. SWIFT: On-the-fly self-speculative decoding for LLM inference acceleration. In ICLR, 2025.
Xia, M., Gao, T., Zeng, Z., and Chen, D. Sheared LLaMA: Accelerating language model pre-training via structured pruning. In ICLR, 2024b.
12

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
A. Proofs
Table 4. Empirical validation of the last-layer LM head weights E" being full-rank. The rank (ie., number of non-zero singular values) of E" is identical to the model's hidden size, confirming that E" is indeed full-rank. Note that even the smallest singular value is significantly greater than 0, demonstrating that the matrix is full-rank not simply because of randomness or noise.
Model
E Shape
# of Singular Values
# of Non-Zero Singular Values
Smallest Singular Value
L.ama-3.1-800
128256 x 4096
4096
4096
1.45271
CodeLlama-I33 13B
32016 x 5120
5120
5120
1.10411
CodeLlama-344
32000 x 8192
8192
8192
0.83576
Lemma A.1. For any E(() E M there exists a T E Rdddd such that E E*T(O.
Proof. We first prove that E" (the last-layer LM head weights) is full-rank, and then present an explicit form of T() that satisfies ( = E During LLM pretaining, the last-layer LM head weights E are trained together with last layer hidden states H" E IR plxd (iD) the total number of training tokens the corpus) to learn the ground-truth next-token
prediction probability P R PP
Pa Softmax H E"
As the softmax function cannot increase matrix rank (Kanai et al., 2018), we have rank(Softmax H E E) < rank H (OI(HEE")) ) S min (rank(nk(())) E"), rank {H*)] S rank (E"). Thus, to achieve a good approximation of P rank E*) must closely match rank(P")) Given the complexity and diversity of natural language, the empirical distribution P derived from the pretraining data is extremely high-rank (Yang et al., 2018). Therefore, to accurately model this high-rank distribution, E" must be full-rank, as indicated by Lemma in Yang et al. (2018): the language modeling problem can be completely solved (i.e., achieve a O loss) if rank P) d. However, this is practically infeasible due the inherent complexity of natural language, suggesting that rank(P) > d. In addition, we also provide empirical validation in Table 4, which further confirms that E" is indeed full-rank.
Given Ex being full-rank, U = E E" is invertible. We can define (+)) as ( = -U-EE" E(0) then:
==
ETT = EEU-EE E. E
E (E E" E' EA'EE"EE"" E E(i)
Note that P = -") is the projection matrix onto the column space of E (as = Since E' is full rank, (x)) lies in the column space of E Therefore, applying P to ( gives E itself, confirming that E(() can be expressed as E+)())
B. Benchmark Details
We evaluate our method a diverse set of text generation tasks, including text summarization, code generation, and mathematical reasoning, covering broad spectrum of language model capabilities.
Text summarization. For text summarization, we use the widely adopted extreme summarization (XSum) dataset (Narayan et al., 2018), where the models are prompted to produce a single-sentence summary of a news article, testing their ability to identify and precisely summarize the most salient information in a coherent sentence. Following previous works (Zhang et al., 2024a), we randomly sample K instances from the test split for evaluation, and 10K instances from the training split for training the lightweight LM head.
Code generation. For code generation, we evaluate our method on the HumanEval (Chen et al., 2021) benchmark, which assesses Python programming skills through a variety of coding problems, ranging from basic tasks to complex problem-solving challenges. Since the standard HumanEval benchmark does not provide a training set, we use the entire MBPP (Austin et al., 2021) dataset for training, which contains a set of crowd-sourced Python programming problems designed to be solvable by entry-level programmers, covering programming fundamentals and standard library functionality. This results in a total of 974 training samples and 164 test samples for this task.
13

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
thresholi 0.3
thrshhod 0.3
14
XSUM
1.4
threshold
threshokd
HumanEoal
throshokd 0.8
Speedup (x)
threshold 0.8
Speedup (x)
GSMSK
Speedup
1.2-
1.2
1.2
1.0-
1.0
1.0
2
8
14
50
100
150
200
# draft tokens
of draft tokens
#
Bayesian Optimizations Iterations
SpecDec with 7B drafter.
(b) SpecDec with 13B drafter.
(c) Self-SpecDecode BO search.
Figure 6. Hyperparameter study of baseline models with CodeLlama-34B as the backbone model. (a) SpeeDecode with CodeLlama-785 as drafter. (b) SpecDeende with Lama- 13B as drafter. (c) SpecDedeeee optimized with Bayesian optimization (BO) search.
Mathmatical reasoning. We use GSM8K (Cobbe al., 2021) as the benchmark for mathematical reasoning, which contains diverse grade-school math word problems created by human problem writers. The dataset consists of 7.5K training problems and IK test problems. These problems typically require multiple reasoning steps to solve and involve performing sequence of basic arithmetic operations (such as addition and subtraction) to arrive at the final answer. The goal of this task is specifically to evaluate the LLM's ability in multi-step mathematical reasoning.
C. Baselines
By default, we employ three instruction-tuned models as the backbone, ranging from 8B 34B parameters (i.e., LLama3.1- 8B-Instruct, CodeLlama- methods. Since LookAhead is a training-fire method, we directly use the optimized hyperparameter configuration in their released codebase to reproduce the results of backbone models on three benchmarks. For SpecDecode, we consider three configurations: (1) Llama3. 1-8-IInsruct as the main model (i.e., verifier) with Llama3.2- 1-ttma...----nstttiee B-Instruct as the assistant model (i.c., drafter), (2) CodeLlama- 33--Instruct as the verifier with CodeLlama- Instrrutct as the drafter, and (3) -ame-33-m--3 Instruct as the verifier with CodeLlama- Lts--Instruc/CooeeLliiicc 7 I-Intrutct the drafter.
and CodeeLlama34BBInnstee
with the default chat template for all compared
For Self-Spe and its extension SWIFT, adopt two backbone models: (1) Codel Lode- Ins-rutt and (2) CodeLlama- 34ntttrrtttt since their released codebase does not support recent models like the Llama3 newer series. Following Zhang et al. (2024a), we adopt an adaptive confidence threshold strategy and set the initial threshold 70 = 0.6, the max number of draft token K = 12, and the max number of generated token = 512. We run the Bayesian optimization search for 200 iterations to determine skipped layers for configuring the drafter model. We use instances randomly sampled from the training set of each task for the Bayesian optimization search as suggested the original implementation. More implementation details regarding training and inference can be found in Appendix D.
C.1. Hyperparameter Sensitivity Study on Baseline Methods
SpecDecode sensitive to hyper-parameter configurations. Figure 6 presents the hyperparameter study for SpecDecode, indicating that is challenging to achieve consistent speedup without careful tuning. In this study, we use CodeLlama-3444 Instruct the backbone model, and performed comprehensive hyperparameter search the HumanFval benchmark, with
and CodeLlama- 1allmmmmm333 1ata--11-B--- Ins-rust as the assistant models, respectively. We explored various configurations with different max number of draft tokens [2, 8, 14] and confidence threshold = [0.3, 0.5, 0.8). Figure 6a and Figure 6b reveal that significant speedup variance exists across different settings. Moreover, the variance becomes even larger when the confidence threshold shifts. Notably, when n = 14 and the threshold is 0.3, the speedup 13B drafter merely 0.944, which even worse than standard decoding. These findings highlight that careful tuning hyperparameters essential to optimize SpecDec's speedup performance and avoid potential slowdowns.
Self-e and its variant SWIFT requires task-specific architecture for optimized speedup performance. Recent works like Self- Seef-SppeDecode have been proposed as a training free acceleration method to be employed in plug- and-play manner. However, it requires an additional Bayesian optimization (BO) process to first obtain set of layers to skip (as the drafter) before it can be adopted for inference. The time-consuming BO process greatly limits its practicality, moreover, it requires a set examples as validation data to select the desired drafter. To further investigate the effect of this search process, we optimize CodeLlama-3A4 34B-Insstuttt with different numbers of BO iterations, as shown Figure 6c. In general, more number of iterations could lead improved evaluation performance. Yet, on HumanEval, demonstrates a
14

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
decrease in performance in the process. This implies scaling the number of iterations does not monotonically reflect better results, and the iteration process might require careful tuning to select optimized drafter for inference.
D. Implementation Details
Training details. The lightweight LM heads in our method are trained through full-parameter fine-tuning using the hignmenn-handbood repository" with 8xNvidia A 100 GPUs. Specifically, we utilize DeepSpeed ZeRO-3 (Rajbhandari et al., 2020) along with FlashAttention (Dao, 2024) for distributed training, and we enable BF16 mixed precision training to enhance training efficiency. We generate on- on-poliey data to train the lightweight LM heads by prompting the off-the-shelf Llama3. 1-8--Instruc or CodeLlama- -amaa-34---aa Instruct to produce responses using greedy sampling for each benchmark. By default, our models are trained using the Adam optimizer (Kingma & Ba, 2014) for 100 epochs, with a batch size of 128, a learning rate of Se-3, and cosine learning rate schedule with 3% warmup steps.
Inference details. During inference, we adopt the zero-shot evaluation by directly prompting the model to generate responses and apply the corresponding chat templates to format the prompts, as all backbone models used in our work are instruction--uned versions. The framework is implemented using the HuggingFace Transformers library, 4 and we set the sampling temperature zero in all methods for a reproducible comparison. Following Zhang et al. (2024a), the maximum number of new tokens is set to 512. The threshold y in Equation (1) is set to 0.75, and to ensure balance between the early prediction rate and the rejection rate, we limit the maximum number of early predictions to 5.
E. Integration with Tree-based Speculative Decoding
Naive combination of adaptive layer parallelism and tree-based decoding does not yield performance improvement. As discussed in Section 5, tree-based speculative decoding methods (Miao et al., 2023b; Cai et al., 2024) and early exiting techniques are orthogonal each other: the former accelerates decoding horizontally by generating multiple tokens across time steps, while the latter reduces per-token computation through vertical acceleration via adaptive layer parallelism enabled by early predictions. These approaches address different bottlenecks in the decoding process.
To explore the effect of adaptive layer parallelism in the context of tree-based decoding, we augment Llama3.1 1-8-311-8--IInttue by inserting and training two additional lightweight LM heads at the final layer (layer 32) and three lightweight heads at each of two intermediate layers (layers 16 and 24). These extra LM heads support multi-hop next-token predictions (i.e., one-hop, two-hop, and three-hop), and we construct draft trees following predefined structure used in Medusa (Cai et al., 2024). To establish the vanilla tree-based decoding baseline, only the heads at the final layer are used to generate multi-hop predictions, while interneiiate-layer heads are inactive. Like Medusa, the original LM head (for one-hop prediction) at the final layer is used to perform verification and determine the final output tokens.
In contrast, to evaluate the impact of adaptive layer parallelism, we further allow the model to generate draft trees at intermediate layers using the lightweight heads. Draft tokens produced at shallower layers can be verified by subsequent deeper layers and ultimately finalized by the last layer. Although this introduces additional verification costs compared to the vanilla tree-based decoding baseline, our intuition is that if early draft trees are reasonably reliable with high acceptance rates, the extra verification cost could be offset by reduced computation from early exits.
We train and evaluate this model on the HumanEval task, and the results of this exploratory study demonstrate that combining vertical acceleration (via adaptive layer parallelism) with horizontal acceleration (via tree-based speculative decoding) generally leads to greater speedups than using vertical acceleration alone. However, it is unexpected that allowing that model to generate draft trees at intermediate layers does not appear to outperform the vanilla tree-based method that generates all draft tokens only at the final layer. In other words, the naive combination of vertical and horizontal acceleration does not yield improvements over horizontal acceleration alone, suggesting that more thoughtful integration strategy is needed to fully exploit their complementary strengths. We hypothesize several possible reasons for this observation. First, draft tokens generated from intermediate layers are often less reliable, resulting in a higher rejection rate during final-layer verification and ultimately diminishing the potential efficiency gains. Second, different layers may favor different tree shallower layers, due to their lower confidence, may benefit from narrower trees with fewer nodes at deeper hops (i.e., making fewer two-hop or three-hop predictions), while deeper layers can support wider trees with more aggressive
"nttps://github www coam/huggingtace//laee 'https://gi thub. com http///ugginggaacc/ttaaaffaroaoe
15

---

AdaDecode: Accelcrating LLM Decoding with Adaptive Layer Parallelism
exploration. Therefore, naively applying a predefined tree structure across all layers may lead suboptimal results. These findings suggest that, while combining vertical and horizontal acceleration holds promise, further research is needed to address these challenges and fully realize the potential of their integration.
F. Limitations and Future Work
Limitations. Our work focuses on accelerating LM decoding without directly improving the quality of the outputs, so AdaDecode may encounter similar issues general LM decoding. such as hallucinations (Huang et al., 2023; Liu et al., 2025) and generating unsafe content (Zhang et al., 2024b; Yu et al., 2024; Chen et al., 2024). Additionally, while AdaDecode consistently accelerates standard autoregressive decoding, it may occasionally incur higher FLOPs, particularly when an early-predicted token does not match the gold token, necessitating a reversion. However, such cases are rare in practice, and we find that strict parity with standard autoregressive decoding is not always necessary. The early-predicted tokens, even when they differ from the final prediction, are typically still meaningful. Therefore, one might consider relaxing strict consistency requirements to avoid unnecessary computational waste.
Future work. In the future, we plan to integrate AdaDecode with other efficiency-enhanccng techniques like model pruming and quantization. Model pruning (Xia et al., 2024b) offers the potential to reduce model size and parameter space by identifying and removing less critical weights neurons. When combined with AdaDecode, pruning could lead to faster decoding times and lower memory usage without significantly impacting performance. Similarly, AdaDecode can also be seamlessly integrated with quantization (Dettmers et al., 2023), which reduces the precision of model weights and activations, substantially lowering memory and compute requirements. It would also be interesting to apply the layer-dropping training strategy, as employed in LayerSkip (Elhoushi et al., 2024), to our method and investigate the trade-offs between training efficiency and model performance. Additionally, exploring the Thompson sampling strategy used in EESD (Liu et al., 2024b) as an alternative to our current threshold-basee stopping criterion may also provide complementary benefits and enhance overall flexibility. On the other hand, discussed in Appendix E, we aim to explore effective strategies for integrating tree-based methods (Cai et al., 2024; Miao et al., 2023b) with AdaDecode by generating multiple candidate tokens at intermediate layers of the model, which may yield even better speedups. These directions could be particularly valuable for efficient inference on mobile and edge devices.
16


Reviews:
[{"review_id": "bGCCbcz4Po", "reviewer": "Reviewer_2V7R", "summary": "The paper proposes to improve autoregressive decoding process in LLMs by decoding by using intermediate layer outputs when the confidence is high. A lightweight LM head is trained to decode the next token from such intermediate layer output enabling this method to be applied on pretrained models without requiring retraining from scratch. Remaining layer computations are executed in parallel with subsequent tokens as needed. The approach shows speedup over baselines in experiments.\n\n## update after rebuttal\n\nThe authors have clarified my concerns around some of the technical details of the work and therefore I am increasing my score to 3. I still feel that the novelty is a bit limited and my concerns around the rank of the transformation matrix $E^*$ and the cost of domain specific heads have not been fully alleviated, but I would not be opposed to accepting this paper as it does add to the discussion in the literature on adaptive inference.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s insightful feedback.\n\n**[Q1]: I would suggest saying that \"$E^\\star$ is likely to be full rank\", as opposed to \"must be full rank\" in Appendix A and modifying the corresponding claim in Section 2.1 accordingly.**\n\n\n**[A1]**: We provide an empirical validation confirming that $E^*$ is indeed full rank in this [PDF](https://anonymous.4open.science/r/AdaDecode-ICML2025-132E/rebuttal/To_Reviewer_2V7R_and_KXLJ.pdf).\n\nMoreover, as indicated by Lemma 1 of (Yang, 2018), the language modeling problem can be completely solved (i.e., achieve a 0 loss) if $\\text{rank}(P) < d$. However, this is practically infeasible due to the inherent complexity of natural language, suggesting that $\\text{rank}(P) > d$.\n\nWe’ll clarify our theoretical claim in the revision for improved precision.\n\n**[Q2]: The work appears to be a synthesis of ideas from the early exit and speculative decoding areas and so the overall novelty is a bit limited.**\n\n**[A2]**: To better position our contributions, we provide a comparison table (Table 2) in the same PDF. We also clarify our novelty in Q1 to Reviewer u7As.\n\n**[Q3]: Domain-specific heads for intermediate layers will significantly increase the memory/cost requirements for this approach in real LLM-based services where requests come from a mix of domains/users.**\n\n**[A3]**: While we acknowledge the advantage of a generalist head, we would like to highlight the growing trend of developing specialized models for domain-specific use cases such as programming (e.g., Cursor/Github Copilot) and mathematical reasoning (e.g, Qwen-2.5-Math). In such applications, only a specific task type is considered, making a domain-specific LM head more preferable.\n\nThat said, we believe our method can also produce good generalist LM heads. We hypothesize that the lower performance observed in the ablation using a generalist head is primarily due to the relatively small size of our mixed-domain dataset (<20K examples). With more extensive mixed-domain training data, we expect the performance of the generalist head to improve significantly. In this work, we focus on demonstrating that lightweight training (<2 GPU hours) is sufficient to produce high-quality intermediate-layer LM heads and achieve substantial speedups with our method. Exploring more comprehensive training for generalist heads will be interesting future work.\n\n**[Q4]: How can layer 2 of t2 and t3 run in parallel (Fig 2) when the 2nd layer of t3 is dependent on the output (KV cache) from the second layer of t2?**\n\n**[A4]**: While it is true that, in standard autoregressive decoding, the 2nd layer of token $t_3$ depends on the KV cache from the 2nd layer of token $t_2$, our method eliminates this dependency by generating the next token earlier. As illustrated in Fig 2, $t_3$ is produced using only the 1st layer output of $t_2$, enabling the 2nd-layer KV caches of $t_2$ and $t_3$ to be computed in parallel. This is analogous to prefilling, where knowing multiple tokens upfront allows for simultaneous KV cache calculation across those tokens. This parallel computation is the core mechanism behind our method's efficiency.\n\n**[Q5]: How will adding $t_i$ (token generated by the intermediate layer) to the parallel processing list of deeper layers in line 235 help in parallel processing? Don't we need to add the layer output instead to compute the outputs of subsequent layers?**\n\n**[A5]**: Line 235 updates $\\mathcal{P}$, the list of tokens that need to be processed at each layer. For instance, if $t_1$ is an early prediction at layer 8, its KV cache at layers 9-32 remains empty. Then for the next token $t_2$, we have the following adaptive parallelism strategy to calculate their KV caches:\n- Phase 1: The first 8 layers of $t_2$ are processed individually\n- Phase 2: The remaining layers 9-32 of $t_2$ are processed in parallel with $t_1$\n\nSpecifically, at the beginning of Phase 2, we concatenate the hidden representation of $t_1$ at layer 8 with that of $t_2$ at layer 8 to enable parallel processing. These intermediate layer outputs (hidden representations) are managed accordingly when updating $\\mathcal{P}$. Since the KV cache computation automatically requires the layer outputs, such operations on layer output are omitted in the algorithm for simplicity.\n\n**[Q6]: Shouldn't there be a condition after line 238 for resuming processing from any rejected token $t$, as described after equation (2), after the final layer is reached at any position?**\n\n**[A6]**: Indeed, this \"resume processing\" operation is internally handled in Line 240 – if any rejection happens, Line 240 will remove disgarded tokens from $y$ and add the replacement token to $y$ so that the generation resumes from the correct token, as described in Section 2.2. We will update the annotations in Algorithm 1 to clarify this.\n\nIf you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!", "claims_and_evidence": "The authors claim that $E^{(i)}$, the lightweight LM head at layer $i\\ \\forall i$, used to check the exit condition can be represented as a linear transformation of the last layer LM head $E^*$. However, the proof presented in Appendix A requires $E^*$ to be full rank but does not conclusively demonstrate that that is the case. The authors only show that the rank of $E^*$ is lower bounded by the rank of a matrix $P$ that is likely to be high rank. However, this is neither a deterministic nor a probabilistic guarantee on $E^*$ being full rank.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I checked the proof of Lemma A.1 in Appendix A and believe that it isn't entirely correct as described above.", "experimental_designs_or_analyses": "I checked the experiments in the main paper and do not have any issues with them.", "relation_to_broader_scientific_literature": "The paper builds upon prior work in speculative decoding and early exits. Unlike early exit which skips the computation of KV cache in post exit layers, this work performs those computations in parallel with subsequent tokens. Unlike speculative decoding which requires a separate drafter model for generating outputs this work uses the same model for both generation (intermediate layers) and verification (last layer).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n\n1. The observation that the existing LM head does not work well with intermediate layer outputs and that training an LM head for intermediate layers will work enables this process to work with pretrained models which is a very useful feature in the era of foundation models and a clear improvement over prior works that require training from scratch.\n\n2. The efficiency modifications for minimizing the complexity of the new LM head and using adaptive layer parallelism to compute KV cache on the fly can help reduce inference cost.\n\nWeaknesses:\n\n1. The work appears to be a synthesis of ideas from the early exit and speculative decoding areas and so the overall novelty is a bit limited.\n\n2. The proof of Lemma A.1 does not appear to be entirely correct (as mentioned above) and this weakens the argument that representing the intermediate LM heads as a linear projection of the final LM head will never lead to a loss of expressivity.\n\n3. Certain technical details appear to be inconsistent (see questions below) \n\n4. It is mentioned in Section 4.3 that domain specific heads for intermediate layers leads to significantly faster generation than heads trained on a mix of domain. However, this will significantly increase the memory/cost requirements for this approach in real LLM-based services where requests come from a mix of domains/users.", "other_comments_or_suggestions": "I would suggest saying that \"$E^*$ is likely to be full rank\", as opposed to \"$E^*$ must be full rank\" in Appendix A and modifying the corresponding claim in Section 2.1 accordingly.", "questions_for_authors": "1. How can layer 2 of t2 and t3 run in parallel (Fig 2) when the 2nd layer of t3 is depended on the output (KV cache) from the second layer of t2?\n\n2. How will adding $t_i$ (token generated by the intermediate layer) to the parallel processing list of deeper layers in line 235 help in parallel processing? Don't we need to add the layer output instead to compute the outputs of subsequent layers?\n\n3. Shouldn't there be a condition after line 238 for resuming processing from any rejected token $t$, as described after equation (2), after the final layer is reached at any position?", "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": {"ethical_review_concerns": "N/A"}}, {"review_id": "c5hUeklBB6", "reviewer": "Reviewer_u7As", "summary": "The authors propose to use early exiting to accelerate autoregressive decoding in LLMs while leveraging adaptive layer parallelism for efficient hardware deployment. The early exiting framework uses a lightweight head at intermediate layers to enable high-confidence early token predictions. An additional verification step is also added to ensure early-predicted tokens match the results of standard autoregressive decoding. AdaDecode achieves upto 1.73X speed up in token generation on several token generation tasks such as summarization, codegen, and mathematical reasoning.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comments. \n\n**[Q1]: The method offers no novelty compared to prior work and it is unclear what the authors are contributing.**\n\n**[A1]**: We would like to highlight that our work provides an efficient solution to tackle the key limitations of speculative decoding and early exiting through our technical innovations.\n\nSpeculative decoding and early exiting have notable limitations: (1) speculative decoding typically incurs substantial training and memory costs due to the need for an additional draft model, and (2) early exiting is unable to leverage off-the-shelf intermediate layer representations without requiring further training or modifications, which can cause output deviations from standard autoregressive decoding.\n\nTo address these challenges, we propose the following key innovations:\n\n- **Leveraging off-the-shelf intermediate layer representations**: For the first time, we demonstrate that, despite the challenges of using off-the-shelf intermediate layers for next-token prediction (Fig. 3a), learning a simple linear projection on these features yields surprisingly good next-token prediction (Fig. 3b). This finding has significant practical implications: The full fine-tuning of the entire model (e.g., LayerSkip [3]) or the introduction of a full additional layer and full LM head (e.g., EESD [4]) may not be necessary. Learning a simple linear projection based on frozen intermediate-layer features significantly reduces not only training complexity but also facilitates our design objective of guaranteeing output parity to vanilla decoding.\n\n- **Efficient lightweight LM heads without loss of expressiveness**: Through theoretical analysis in Lemma A (which is unknown in prior works), we prove that our lightweight intermediate LM heads (i.e., the transformation matrices) are lossless proxies for full LM heads, which reduce memory costs by 30x. This innovation allows us to serve multiple lightweight LM heads at various intermediate layers (otherwise introducing multiple full LM heads will cause expensive training and memory costs), enabling dynamic early exiting and delivering significant speedups (as shown in the “w/ fixed-layer early prediction” ablation of Table 2).\n\nTo better position our contributions, we present a comparison table that highlights how our contributions differ from recent works. Please refer to Table 1 in this [PDF](https://anonymous.4open.science/r/AdaDecode-ICML2025-132E/rebuttal/To_Reviewer_u7As.pdf).\n\n\n**[Q2]: While the authors have done a good job in comparing their method to approaches other than early exiting, they have not performed the obvious comparison to the numerous publications on early exiting in the literature.**\n\n**[A2]**: To address the reviewer’s concern,  we would like to make the following clarifications.\n\n- **Why not compare with output-inconsistent methods (e.g,. FREE [1], LITE [2], LayerSkip [3])**: Our work is explicitly designed under the constraint of output consistency—producing the same outputs as standard autoregressive decoding. Relaxing this constraint leads to fundamentally different problem settings and design objectives, making direct comparisons with output-inconsistent methods inappropriate.\n\n- **Why not compare with tree-based decoding methods (e.g,. SpeccInfer [8], Medusa [5])**: Tree-based decoding methods are orthogonal to early exiting methods: The former generates multiple tokens simultaneously, accelerating **horizontally (across time steps)**, while the latter reduces per-token computation by parallelizing deep layers after early exit, accelerating **vertically (within each time step's forward pass)**. A direct comparison between these two orthogonal approaches would conflate their respective benefits. As precedent, SWIFT [7], LayerSkip [3], or Self-SpecDecode [6] also didn't compare with tree-based decoding methods for the same reason.\n\nIf you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!\n\n**References**\n\n[1] Bae et al. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. 2023\n\n[2] Varshney et al. Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite. 2023\n\n[3] Elhoushi et al. LayerSkip: Enabling early exit inference and self-speculative decoding. 2024.\n\n[4] Liu et al. Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism. 2024\n\n[5] Cai et al. Medusa: Simple llm inference acceleration framework with multiple decoding heads. 2024\n\n[6] Zhang et al. Draft & verify: Lossless large language model acceleration via self-speculative decoding. 2023\n\n[7] Xia et al. SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration. 2024\n\n[8] Miao et al. Specinfer: Accelerating large language model serving with tree-based speculative inference and verification. 2024", "claims_and_evidence": "The authors provide multiple experiments to support the claims in the paper.", "methods_and_evaluation_criteria": "The benchmarks and speedup comparisons  between the proposed method and the baselines are meaningful.", "theoretical_claims": "The authors provide a proof in Lemma A.1 in the appendix which appears valid. The assumption that $P^*$ is high rank makes sense. Additionally, as indicated, since $E^*$ is full rank, matrix $E^{(i)}$ can be expressed as a linear transformation of $E^*$.", "experimental_designs_or_analyses": "Please refer to the weaknesses section. While the experimental results are valid, they lack comparison to the many prior Early-Exiting works.", "relation_to_broader_scientific_literature": "The paper does not provide any new contributions compared to prior methods on early exiting for autoregressive decoding and fails to mention or compare to them. At its current state, the paper provides no novelty or new contribution to the field.", "essential_references_not_discussed": "The paper seems to miss many of the existing related works in the literature with identical/significant overlap:\n\n[1] Bae, Sangmin, et al. \"Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding.\" arXiv preprint arXiv:2310.05424 (2023).\n\n[2] Elhoushi, Mostafa, et al. \"LayerSkip: Enabling early exit inference and self-speculative decoding.\" arXiv preprint arXiv:2404.16710 (2024).\n\n[3] Liu, Jiahao, et al. \"Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism.\" arXiv preprint arXiv:2406.03853 (2024).\n\n[4] Cai, Tianle, et al. \"Medusa: Simple llm inference acceleration framework with multiple decoding heads.\" arXiv preprint arXiv:2401.10774 (2024).\n\n[5] Varshney, Neeraj, et al. \"Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite.\" arXiv preprint arXiv:2310.18581 (2023).", "other_strengths_and_weaknesses": "While the authors have done a good job in comparing their method to approaches other than early exiting, such as SpecDecode and Swift, they have not performed the obvious comparison to the numerous publications on early exiting in the literature.\n\n- The method offers no novelty compared to prior work and it is unclear what the authors are contributing. Plugging in an intermediate head for early exiting inside a language model is not a new contribution.\n\n- Combining verification / speculative decoding with early exiting is not a new idea and has been published in the references provided above.\n\n- The related work section on early exiting just mentions a few works and refers the reader to a survey. The authors should clarify how their method is differentiated from prior work.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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": "Tuh7QHb6iH", "reviewer": "Reviewer_KXLJ", "summary": "The authors present AdaDecode, a methodology to accelerate decoding without auxiliary models or modification of the model. The proposed approach adaptively predicts tokens from an intermediate layer based on confidence, using a set of additional lightweight LM heads whose predictions are verified using a rejection sampling scheme.\nAdaDecode retains output consistency and achieves a speed-up of 1.73x over vanilla decoding, outperforming four baselines: speculative decoding, self-speculative decoding, LookAhead, and SWIFT.\nMoreover, the authors also include ablation studies and hyperparameter sensitivity experiments to validate design decisions and better understand the impact of the algorithm's different components.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the insightful comments and appreciation of our work.\n\n**[Q1]: While approaches like SpecInfer (Miao et al., 2023b), Medusa (Cai et al., 2024), and its extension HYDRA (Ankner et al., 2024) are orthogonal to AdaDecode, it could be interesting to discuss further to contextualize the contribution better.**\n\n**[A1]**: We agree that these tree-based speculative decoding methods are orthogonal to AdaDecode, but we believe it is possible and interesting to explore how our method could potentially be combined with tree-based decoding techniques. For instance, one potential approach could be generating multiple tokens at intermediate layers using our lightweight LM heads, rather than introducing additional full LM heads as in Medusa and always drafting tokens at the last layer. This strategy could allow for a more efficient draft token generation while maintaining the benefits of tree-based verification, and we believe this will be an interesting direction to explore in future work.\n\n\n**[Q2]: It would be interesting to show an empirical study on the models to show E\\* is full rank.**\n\n\n**[A2]**: Thanks for the great suggestion. Please refer to Table 1 in this [PDF](https://anonymous.4open.science/r/AdaDecode-ICML2025-132E/rebuttal/To_Reviewer_2V7R_and_KXLJ.pdf).\n\n\n**[Q3]: The paper and appendix clearly describe experimental designs and analyses. I tried to run the code, but at the current stage, some needed artifacts are missing (e.g., adadecode/llama_3.1_8b_instruct)**\n\n\n**[A3]**: Please find the datasets and model checkpoints at this anonymous repo: https://huggingface.co/AnonyResearcher\n\n\n**[Q4]: Regarding the passage on deviation from theoretical guarantees, the authors hypothesise a possible cause originating from FP16 precision used at inference time. It would be interesting to see some experiments performing inference tweaking the precision.**\n\n**[A4]**: Following the reviewer’s suggestion, we re-run the experiments in Figure 4 with FP32. While FP32 offers slightly higher numerical consistency than FP16, it still does not reach 100% consistency. \n|Method|SpecDec|Self-SpecDec|LookAhead|SWIFT|AdaDecode|\n|-|-|-|-|-|-|\n|FP16|0.97|0.98|0.98|0.98|0.99|\n|FP32|0.98|0.98|0.99|0.98|0.99|\n\nThis aligns with [the finding in the open-source community](https://github.com/huggingface/transformers/issues/30413) that the output of speculative decoding can differ slightly from standard decoding due to numerical precision inaccuracies and minor variations in token probabilities during computation (please refer to this [discussion](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535) for a detailed study on the impact of different precisions). We would like to note that, during our experiments, we also noticed that different hardware specifications and library versions contribute to variances, demonstrating that half-precision is not the only factor affecting consistency. \n\n**[Q5]: The results that show that generalist heads are significantly slower might suggest a limited applicability of the approach in a \"production\" setting.**\n\n**[A5]**: While we acknowledge the advantage of using a generalist head that can support a mix of domain requests, we would like to highlight the growing trend of developing specialized models for domain-specific use cases. \nIn many real-world production settings, instead of training a single monolithic model for all domains, specialized models offer greater efficiency and improved performance in their respective areas. For example, Cursor and GitHub Copilot are optimized for programming assistance, while models like Qwen-2.5-Math and  DeepSeek-Prover are designed for mathematical reasoning and theorem proving. In such applications, only a specific task type is considered, making a domain-specific LM head more preferable.\nThat said, we believe our method can also produce good generalist LM heads. We hypothesize that the lower performance observed in the ablation using a generalist head is primarily due to the relatively small size of our mixed-domain dataset (<20K examples). With more extensive mixed-domain training data, we expect the performance of the generalist head to improve significantly. In this work, we focus on demonstrating that lightweight training (<2 GPU hours) is sufficient to produce high-quality intermediate-layer LM heads and achieve substantial speedups with our method. Exploring more comprehensive training for generalist heads will be interesting future work.\n\n\nWe hope the reviewer finds our response helpful. We are also happy to incorporate additional suggestions you might have!  Thanks again for your time!", "claims_and_evidence": "The author claims are factual and revolve around algorithm speed-up and output consistency. Both claims are clearly backed up by the presented experiments.", "methods_and_evaluation_criteria": "The methodology used to evaluate the approach makes sense. My only comment is about the limited selection of the baselines.\nThe authors in the related work reasonably exclude methods that do not preserve output consistency but mention approaches like SpecInfer (Miao et al., 2023b), Medusa (Cai et al., 2024), and its extension HYDRA (Ankner et al., 2024). While orthogonal to AdaDecode, it could be interesting to discuss further or report their results to contextualize the contribution better.", "theoretical_claims": "Checked the proof for Lemma A.1.\nThe proof is clear, but I invite the authors to expand on why P* high-rank implies E* is full rank, as this seems speculative. It would be interesting to show an empirical study on the models considered for the experiments in the same appendix.", "experimental_designs_or_analyses": "The paper and appendix clearly describe experimental designs and analyses.\nI tried to run the code, but at the current stage, some needed artifacts are missing (e.g., adadecode/llama_3.1_8b_instruct).", "relation_to_broader_scientific_literature": "The paper's findings are extremely relevant, given the strong need to improve latency in LLM without increasing the already high hardware requirements. The work is timely and builds on existing concepts addressing key limitations.", "essential_references_not_discussed": "I'm not aware of any essential related work not referenced.", "other_strengths_and_weaknesses": "Already highlighted in previous sections.", "other_comments_or_suggestions": "Regarding the passage on deviation from theoretical guarantees, the authors hypothesise a possible cause originating from FP16 precision used at inference time.\nIt would be interesting to see some experiments performing inference tweaking the precision.\n\nThe results that show that generalist heads are significantly slower might suggest a limited applicability of the approach in a \"production\" setting.\nIt would be interesting to expand more on this in the discussion and suggest potential strategies to mitigate this effect.", "questions_for_authors": "No important further questions besides the comments.", "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": "AlQXHYZCfF", "reviewer": "Reviewer_5F7j", "summary": "This paper introduces AdaDecode, a self-speculative decoding method with an early exiting mechanism. Based on empirical findings that many simple and predictable tokens can be accurately generated at intermediate transformer layers, the authors propose three key contributions. First, they introduce a lightweight intermediate-layer LM head training approach that enables high-confidence early predictions without modifying the original model parameters. This lightweight head achieves performance comparable to fully parameterized LM heads from prior works, and is tuned to minimize KL divergence with final layer outputs. Second, they develop adaptive layer parallelism that concurrently processes multiple early-predicted tokens generated by the lightweight LM heads from different layers, significantly improving hardware utilization and decoding speed. Third, by incorporating early exiting with self-speculative decoding, AdaDecode enables KV-cache sharing between draft and verification stages, reducing computational resources while maintaining output consistency.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's insightful feedback.\n\n**[Q1]: Could you present additional experiments demonstrating that AdaDecode's performance is robust to different choices of early exiting layers?** \n\n**[A1]**: Please refer to Table 1 in this [PDF](https://anonymous.4open.science/r/AdaDecode-ICML2025-132E/rebuttal/To_Reviewer_5F7j.pdf). The results show that AdaDecode achieves consistent speedups across different layer configurations.\n\n**[Q2]: The paper is missing a critical citation Early-exiting Speculative Decoding (EESD). Could you elaborate on what specific technical innovations in AdaDecode differentiate it from EESD's approach?**\n\n**[A2]**:  We have cited EESD in our original submission (Line 565). Below is a detailed comparison.\n\n- **Efficiency**: To enable reliable early exiting, EESD requires an extra decoder layer and a full LM head (vocab size × hidden size), adding ~0.7B trainable parameters.  In contrast, AdaDecode introduces only three lightweight LM heads (hidden size × hidden size), totaling just 48M additional parameters, making it significantly more efficient.\n\n- **Flexibility**: While both methods use early exiting for drafting, AdaDecode supports dynamic-depth early exiting that allows adaptive layer parallelism, whereas EESD is restricted to fixed-depth early exiting.\n\n- **Training and stopping strategy**: EESD trains its additional modules using cross-entropy loss and employs Thompson Sampling for stopping decisions. AdaDecode instead optimizes lightweight heads via KL divergence and uses probability-thresholding for termination.\n\n**[Q3]: The paper would benefit from more comprehensive evaluation details such as: (1) detailed hardware specifications for reproducibility, (2) acceptance rates for speculative tokens across different scenarios, (3) memory utilization statistics, and perhaps (4) theoretical computation reduction metrics.**\n\n\n**[A3]**:  Please refer to Appendix D and Figure 5 for a detailed discussion on (1) and (2). For (3), we maximize the memory utilization during training (~80G per GPU), and the memory utilization during inference varies depending on context length and batch size, starting from ~18G with FP16 precision.\n\nFor (4), theoretical computation reduction depends on the number of total generated tokens and number of accurate early predictions. Let $\\alpha$ be the fraction of layers needed for early exiting. Given $N$ tokens and per-token latency $T$, vanilla decoding takes $NT$,  while AdaDecode requires $T (1+ \\alpha (N-1))$. As $N \\to \\infty$, the theoretical speedup is $1/\\alpha$. In practice, the value of  $\\alpha$ can vary based on model size. For instance, in a 48-layer model (e.g., CodeLlama-34B), exiting at layer 12 ideally gives $\\alpha \\sim 0.25$, leading to a potential $4\\times$ speedup.\n\n**[Q4]: In the comparison with baseline methods like Self-SpecDecode, did you conduct experiments using the exact configurations reported in the original papers? If modifications were necessary, could you provide more details about how you ensured fair comparison?** \n\n\n**[A4]**: Yes, we use **exactly the same configuration** as Self-SpecDecode for CodeLlama-13B, as it was also adopted in the original paper. However, since they did not report results for LLaMA-3.1-8B, we used their Bayesian optimization script to search extensively for the best hyperparameters (Appendix C), ensuring a fair comparison.\n\n**[Q5]: Provide the exact scripts used for Bayesian optimization in their comparisons, including random seeds to ensure deterministic results.**\n\n**[A5]**: We used the same Bayesian optimization script from the official repo of Self-SpecDecode: https://github.com/dilab-zju/self-speculative-decoding/blob/main/search.ipynb.\n\n**[Q6]: Context Length Impact on Rejection Rates.**\n\n\n**[A6]**:  Our AdaDecode exhibits consistent speedups across different scenarios with varying context length. Please refer to Table 2 in this [PDF](https://anonymous.4open.science/r/AdaDecode-ICML2025-132E/rebuttal/To_Reviewer_5F7j.pdf).\n\n**[Q7]: Maximum New Token Length Analysis.**\n\n**[A7]**:  Please refer to Table 3 in the same PDF, which confirms our method’s effectiveness with different max new token lengths.  \n \n**[Q8]: Algorithm 1 should explicitly express how KV cache management is handled between draft and verification stages.** \n\n**[A8]**: Indeed, KV cache management is handled internally in Line 238. If any rejection occurs, the KV cache of discarded tokens will be cleared, and the entire KV cache length is truncated to the last verified token. This ensures that generation resumes from the correct token, as described in Section 2.2. We will update the annotations in Algorithm 1 to make this clearer.\n\nIf you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!", "claims_and_evidence": "- The paper's core claims about the effectiveness of lightweight intermediate-layer LM heads and adaptive layer parallelism are generally well-supported by the presented evidence:\n    - The proposed method of using KL divergence between final layer and intermediate layer output distributions is well-motivated and effectively implemented. This approach successfully calibrates log probability between draft and verification stages, as implicitly demonstrated in the experiments about relation between **γ** and rejection rate (Figure 5(c))\n    - The claim of achieving comparable performance with lightweight head is supported by comprehensive experimental results across multiple model architectures (e.g. Llama3.1-8B, CodeLlama-13B/34B-inst) and datasets (e.g. XSum, HumanEval, GSM8K)\n- However, there are concerns about the robustness and justification for the early exiting layer selection policy:\n    - The authors select three intermediate layers for early exiting (8th, 16th, 24th) across different model architectures without providing sufficient analysis justifying this specific configuration. While their experiments with LLaMA3.1-8B, CodeLLama-13B, and 34B show the approach works well, the paper lacks ablation studies or theoretical analysis explaining why selecting three layers works well.\n    - Previous work in layer skipping or early exiting typically incorporates additional fine-tuning or specialized architectures to determine optimal exiting points. The absence of a more sophisticated approach to layer selection raises questions about whether the reported speedups represent the maximum potential of the method.\n    - A more thorough analysis of different layer configurations would strengthen the paper's claims and potentially provide insights for adapting AdaDecode to other model architectures or domains in future work.", "methods_and_evaluation_criteria": "- The paper's proposed methods are technically sound and well-aligned with the goal of improving inference efficiency through speculative decoding:\n    - The lightweight intermediate-layer LM head approach is a practical solution that avoids modifying original model parameters while enabling efficient early predictions.\n    - The adaptive layer parallelism technique represents a meaningful advancement in hardware utilization for speculative decoding systems.\n- Regarding evaluation, the authors employ standard metrics for the field:\n    - The paper uses conventional evaluation metrics for speculative decoding: throughput (Tokens/s) and relative speedup compared to vanilla decoding. These metrics do allow for basic comparisons with prior work in the field.\n    - The evaluation across different model architectures (LLaMA3.1-8B, CodeLLama-13B, 34B) and diverse tasks provides good coverage of practical applications.\n- However, there is a notable limitation in the reproducibility of evaluation results:\n    - While throughput and speedup are standard metrics, they are hardware-dependent, which can make cross-study comparisons challenging.\n    - The paper would benefit from more comprehensive evaluation details such as: detailed hardware specifications for reproducibility, acceptance rates for speculative tokens across different scenarios, memory utilization statistics, and perhaps theoretical computation reduction metrics. These additional measures would provide a more complete picture of the method's efficiency beyond raw throughput numbers, especially for researchers working with different hardware configurations.\n- Also, the comparison methodology with previous state-of-the-art speculative decoding methods presents a significant limitation in the paper's evaluation:\n    - According to Supplement section C, the authors appear to have compared AdaDecode with previous methods like Self-SpecDecode using configurations that differ from those reported in the original papers. This approach undermines fair comparison, which is a cornerstone of reproducible machine learning research\n    - The specific example of comparing with Self-SpecDecode on XSum using LLaMa-3.1-7B is particularly problematic since this configuration wasn't reported in the original Self-SpecDecode paper. This is especially concerning because Self-SpecDecode utilizes Bayesian optimization for hyperparameter tuning, which is highly sensitive to initial conditions and configuration settings.\n    - Hyperparameter tuning significantly impacts model performance and reproducibility. When comparing methods that rely on different tuning approaches, ensuring consistent evaluation conditions becomes critical.\n- For improved reproducibility, the authors should:\n    - Provide the exact scripts used for Bayesian optimization in their comparisons, including random seeds to ensure deterministic results.\n    - Alternatively, conduct experiments using the exact model architectures and datasets reported in the original papers they compare against.", "theoretical_claims": "The paper provides theoretical justification for its lightweight LM head approach in Supplement A. Upon examination, the proof focuses on establishing the mathematical existence of a transformation matrix T that enables parameter-efficient implementation of intermediate layer heads. The mathematical derivations appear correct and thoroughly support the claims made in the main paper.", "experimental_designs_or_analyses": "I have reviewed the experimental designs and analyses in the paper, with particular focus on their methodological soundness:\n\nThe ablation studies examining the robustness of AdaDecode to the hyperparameter **γ** (used as a threshold for the drafting stage) are well-designed and provide valuable insights into the method's stability. The authors appropriately vary this key parameter and demonstrate consistent performance across a reasonable range of values, which strengthens confidence in the practical applicability of the method.\n\nThe overall experimental framework is methodologically sound, with the authors: \n- Testing on multiple model architectures (LLaMA3.1-8B, CodeLLama-13B, 34B) to demonstrate generalizability\n- Evaluating on diverse tasks including summarization, mathematical reasoning, and code generation\n- Providing appropriate baselines for comparison including vanilla decoding and alternative speculative methods\n- Measuring both throughput improvement and output consistency to balance speed and quality\n\nHowever we do have some suggestion that could provide more insight for the future readers of this paper.\n\n1. Context Length Impact on Rejection Rates\n\nWhile the paper demonstrates effectiveness on standard benchmarks, it lacks analysis of how rejection rates scale with context length. This is critical because:\n\n- Draft quality and acceptance rates often degrade with longer contexts due to increased model strain\n- Smaller models like CodeLlama-7B may struggle to maintain draft quality for long-context inputs (as mentioned in paper)\n\nSuggested Additions:\n\n- Experiments measuring acceptance rates across varied context lengths (e.g., 4K → 32K → 128K tokens)\n- Comparison of rejection ratios between AdaDecode and baselines under long-context scenarios\n- Analysis of whether lightweight LM heads mitigate long-context challenges better than conventional drafting approaches\n\n2. Maximum New Token Length Analysis\n\nThe current experiments only test up to 512 new tokens, which may not reveal efficiency patterns for different generation demands:\n\nSuggested Additions:\n\nComparative tests with max_token ∈ {256, 512, 1024} to evaluate:\n- Throughput consistency across generation lengths\n- Rejection rate trajectories over extended sequences\n- Memory utilization patterns during prolonged generation\n\nThese additions would better demonstrate AdaDecode's robustness across practical deployment scenarios while addressing inherent challenges in speculative decoding systems.", "relation_to_broader_scientific_literature": "AdaDecode's lightweight LM head approach builds upon previous early exit methods, but distinguishes itself by offering a more parameter-efficient implementation. While prior works like LayerSkip[1] also explore early exiting with shared LM heads, AdaDecode's specific parameter efficient LM head appears to be orthogonal contribution. As author mentioned future works for using LoRA instead of lightweight LM head is also expected approach.\n\nThe paper's approach of combining multiple early exiting layers with self-speculative decoding connects to concurrent works. AdaDecode's adaptive layer parallelism for concurrent processing of multiple early-predicted tokens represents a potentially orthogonal optimization approach.\n\nThe reuse of KV-cache between drafting and verification stages appears in several recent works including LayerSkip[1], EESD[2]\n\n[1] Elhoushi, Mostafa, et al. \"LayerSkip: Enabling early exit inference and self-speculative decoding.\" arXiv preprint arXiv:2404.16710 (2024).\n[2] Liu, Jiahao, et al. \"Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism.\" arXiv preprint arXiv:2406.03853 (2024).", "essential_references_not_discussed": "The paper is missing a critical citation related to its core contributions. EESD (Early-exiting Speculative Decoding), published in ACL 2024 Findings presents remarkably similar techniques and should be acknowledged. \n\n[1] Liu, Jiahao, et al. \"Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism.\" arXiv preprint arXiv:2406.03853 (2024).", "other_strengths_and_weaknesses": "- Although the paper's contributions are well aligned with its motivation, there is concern that these contributions may overlap significantly with concurrent works. For instance, EESD[1] proposes a self-speculative decoding framework with early exiting layers that bears notable similarities to AdaDecode:\n    - Both methods train early exiting layers while keeping the original model's parameters fixed\n    - EESD[1] uses self-distillation which performs similarly to the KL divergence approach in AdaDecode\n    - Both methods highlight that KV cache created at the draft stage is reusable during the verification stage\n- Given these substantial overlaps in core technical approaches, the authors should more clearly articulate what contributions of AdaDecode are orthogonal to or extend beyond those of EESD[1]. This clarification would help position the paper's contribution.\n\n\n[1] Liu, Jiahao, et al. \"Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism.\" arXiv preprint arXiv:2406.03853 (2024).", "other_comments_or_suggestions": "Algorithm 1 provides a helpful overview of how AdaDecode works, which aids reviewer understanding. However, the verification stage is not adequately presented in the algorithm.\n\nSpecifically, the algorithm should explicitly express how KV cache management is handled between draft and verification stages. There should be clear notation or steps indicating that the KV cache values generated for accepted draft tokens are preserved and reused during the verification stage, while KV cache values for rejected drafts are discarded. This KV cache reuse is a critical aspect of the method's efficiency gains, as it prevents redundant computation for tokens that pass verification.\n\nThis clarification would improve the algorithm's completeness and better highlight one of the key advantages of the AdaDecode approach - the ability to share computation between draft and verification stages through KV cache reuse. Without this explicit indication, readers might miss this important implementation detail that contributes significantly to the method's performance benefits.", "questions_for_authors": "1. The paper presents experiments with pre-selected early exiting layers across different model architectures. Could you provide preliminary experimental results that justified this specific selection? Alternatively, could you present additional experiments demonstrating that AdaDecode's performance is robust to different choices of early exiting layers? This would strengthen the generalizability claim of your approach.\n\n2. In the comparison with baseline methods like Self-SpecDecode, did you conduct experiments using the exact configurations reported in the original papers? If modifications were necessary, could you provide more details about how you ensured fair comparison? This would address concerns about reproducibility and the validity of reported performance improvements.\n\n3. While the benchmark datasets used are reasonably diverse, they may not fully represent challenging scenarios where early exiting might underperform. Specifically, tasks requiring processing of long contexts or generating open-ended responses (e.g. ELI5, NarrativeQA) might present different dynamics for speculative decoding. Have you evaluated AdaDecode on such tasks, and if so, could you share those results?\n\n4. Could you elaborate on how AdaDecode is orthogonal to recently published Early-exiting Speculative Decoding (EESD) in ACL 2024 Findings? EESD similarly uses early-exiting structures and self-distillation (comparable to your KL divergence approach), while employing a Thompson Sampling Control Mechanism. What specific technical innovations in AdaDecode differentiate it from EESD's approach, and how might these differences contribute to performance improvements?\n\nELI5 : https://facebookresearch.github.io/ELI5/\nNarrative QA : https://huggingface.co/datasets/deepmind/narrativeqa_manual", "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:

Meta-review: The paper proposes an approach to speed up LLM inference by combining early-exiting with speculative decoding. The core idea is to allow intermediate layers to predict the next token, with these predictions subsequently being verified in parallel via the speculative decoding mechanism. This guarantees quality-neutrality, while allowing for speed ups owing to a reduction in the amount of sequential processing.

Reviewer opinion on the work was mixed. On the one hand, reviewers found the method to be sound and intuitive, with a convincing set of empirical results and ablations. On the other hand, some reviewers raised concerns around novelty compared to a large number of existing works also touching upon similar ideas. It was also requested that there be empirical comparisons to some of these works.

The author response argued that the the present work _is_ novel compared to existing works, which either forego quality-neutrality by not making use of speculative decoding, or which do not combine speculative decoding with depth-wise adaptivity.

Owing to the disagreement in scores, the AC also reviewed the paper, and some of the prior references suggested by reviewers. Our reading is as follows:
- there are several works which introduce some form of early-exiting, but without explicit quality-neutrality guarantees (e.g., [1, 2]). These are certainly topically relevant, but the present paper's focus on quality-neutrality is a reasonable distinction.
- there are several works which leverage speculative decoding, with innovations in how the drafting is performed (e.g., [3]). These are certainly topically relevant, but the present paper's focus on drafting based on an intermediate layer is a reasonable distinction.
- LayerSkip [4] and EESD [5] both consider a combination of early-exiting and speculative decoding. Both papers _are_ cited in the present work; however, we believe a far **more detailed discussion** of their similarities and differences is needed.
- from our reading of [4], we believe that Section 4.3 discusses a highly similar idea to the present paper, wherein intermediate-layer drafts are verified via the speculative decoding process. However, a distinction is in the mechanism used for _training_ the underlying model: [4] proposes a layer dropout strategy, which allows for re-use of the final unembedding matrix by any intermediate layer. This requires fine-tuning of the _entire_ network, which can result in performance differences to the original network. By contrast, the present work attaches additional trainable parameters at each layer, while the backbone is _frozen_. This allows for quality preservation of the original network performance.
- from our reading of [5], we believe that again the core idea illustrated in Figure 2 is highly similar to the present paper. However, three distinctions are the use of a _fixed_ layer to perform drafting, rather than a dynamic layer selection in the present paper; the use of Thompson sampling to allow for a dynamic number of draft tokens, rather than a fixed selection in the present paper; and the use of a separate Transformer model on top of the early layer embeddings, as opposed to a lightweight linear layer. The last point in particular means that the present proposal involves far fewer parameters than [5].
- from the above, we believe the present paper _can be reasonably argued to offer novelty_ over [4, 5]; and that these novelties can have practical implications (specifically, freezing of the backbone to guarantee quality neutrality, and use of lightweight linear layers). This said, per above, it is really incumbent on the authors to state such arguments clearly; we believe these are not sufficiently spelled out in the original submission, though the response did touch upon these.
- finally, with the above understanding, one question is whether or not an empirical comparison is necessary. For [4], we believe that a comparison to the layer dropping training strategy is of considerable interest; in the settings that the present paper considers, does this actually lead to a significant quality drop compared to the original network? For [5], we believe that a comparison of the linear head versus a Transformer head would be of considerable interest, and could drive home the point of parameter efficiency of the present method (assuming one ends up with similar acceptance rates). We agree that the Thompson sampling strategy is somewhat orthogonal, and so a comparison may not be strictly necessary.

Overall, on balance we believe the paper does have some distinctiveness to the existing literature, albeit it may still be argued to synthesize a few different ideas (which is not necessarily a bad thing). The paper must however make the discussion of these contributions far more detailed. Further, an empirical comparison to the key strategies one of [4, 5] would really help drive home the advantages of the proposed approach.

[1] Bae, Sangmin, et al. "Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding." arXiv preprint arXiv:2310.05424 (2023).

[2] Varshney, Neeraj, et al. "Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite." arXiv preprint arXiv:2310.18581 (2023).

[3] Cai, Tianle, et al. "Medusa: Simple LLM inference acceleration framework with multiple decoding heads." arXiv preprint arXiv:2401.10774 (2024).

[4] Elhoushi, Mostafa, et al. "LayerSkip: Enabling early exit inference and self-speculative decoding." arXiv preprint arXiv:2404.16710 (2024).

[5] Liu, Jiahao, et al. "Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism." arXiv preprint arXiv:2406.03853 (2024).
Decision: Accept (poster)