AdaDecode: Accelerating LLM Decoding with Adaptive Layer Parallelism

We sincerely appreciate the reviewer's insightful feedback.

[Q1]: Could you present additional experiments demonstrating that AdaDecode's performance is robust to different choices of early exiting layers?

[A1]: Please refer to Table 1 in this PDF. The results show that AdaDecode achieves consistent speedups across different layer configurations.

[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?

[A2]: We have cited EESD in our original submission (Line 565). Below is a detailed comparison.

• 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.

• 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.

• 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.

[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.

[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.

For (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.

[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?

[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.

[Q5]: Provide the exact scripts used for Bayesian optimization in their comparisons, including random seeds to ensure deterministic results.

[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.

[Q6]: Context Length Impact on Rejection Rates.

[A6]: Our AdaDecode exhibits consistent speedups across different scenarios with varying context length. Please refer to Table 2 in this PDF.

[Q7]: Maximum New Token Length Analysis.

[A7]: Please refer to Table 3 in the same PDF, which confirms our method’s effectiveness with different max new token lengths.

[Q8]: Algorithm 1 should explicitly express how KV cache management is handled between draft and verification stages.

[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.

If you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!

We sincerely thank the reviewer for the insightful comments and appreciation of our work.

[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.

[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.

[Q2]: It would be interesting to show an empirical study on the models to show E* is full rank.

[A2]: Thanks for the great suggestion. Please refer to Table 1 in this PDF.

[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)

[A3]: Please find the datasets and model checkpoints at this anonymous repo: https://huggingface.co/AnonyResearcher

[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.

[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.

This aligns with the finding in the open-source community 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 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.

[Q5]: The results that show that generalist heads are significantly slower might suggest a limited applicability of the approach in a "production" setting.

[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. In 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. That 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.

We hope the reviewer finds our response helpful. We are also happy to incorporate additional suggestions you might have! Thanks again for your time!

We thank the reviewer for the thoughtful comments.

[Q1]: The method offers no novelty compared to prior work and it is unclear what the authors are contributing.

[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.

Speculative 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.

To address these challenges, we propose the following key innovations:

• 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.

• 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).

To 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.

[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.

[A2]: To address the reviewer’s concern, we would like to make the following clarifications.

• 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.

• 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.

If you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!

References

[1] Bae et al. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. 2023

[2] Varshney et al. Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite. 2023

[3] Elhoushi et al. LayerSkip: Enabling early exit inference and self-speculative decoding. 2024.

[4] Liu et al. Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism. 2024

[5] Cai et al. Medusa: Simple llm inference acceleration framework with multiple decoding heads. 2024

[6] Zhang et al. Draft & verify: Lossless large language model acceleration via self-speculative decoding. 2023

[7] Xia et al. SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration. 2024

[8] Miao et al. Specinfer: Accelerating large language model serving with tree-based speculative inference and verification. 2024

We appreciate the reviewer’s insightful feedback.

[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.

[A1]: We provide an empirical validation confirming that $E^*$ is indeed full rank in this PDF.

Moreover, 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$.

We’ll clarify our theoretical claim in the revision for improved precision.

[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.

[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.

[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.

[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.

That 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.

[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?

[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.

[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?

[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:

• Phase 1: The first 8 layers of $t_2$ are processed individually
• Phase 2: The remaining layers 9-32 of $t_2$ are processed in parallel with $t_1$

Specifically, 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.

[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?

[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.

If you find our response satisfactory, we would be grateful if you could consider raising your score. Thanks again for your time!