mlsanityData to 14 Oct 2026

KOALA: Enhancing Speculative Decoding for LLM via Multi-Layer Draft Heads with Adversarial Learning

arXiv 2408.08146 ↗First posted 15 Aug 2024cs.CLcited by 6

On the map: Efficient LLM inference and serving → Speculative decoding

Abstract (arXiv)

Large Language Models (LLMs) exhibit high inference latency due to their autoregressive decoding nature. While the draft head in speculative decoding mitigates this issue, its full potential remains unexplored. In this paper, we introduce KOALA (K-layer Optimized Adversarial Learning Architecture), an orthogonal approach to the draft head. By transforming the conventional single-layer draft head into a multi-layer architecture and incorporating adversarial learning into the traditional supervised training, KOALA significantly improves the accuracy of the draft head in predicting subsequent tokens, thus more closely mirroring the functionality of LLMs. Although this improvement comes at the cost of slightly increased drafting overhead, KOALA substantially unlocks the draft head's potential, greatly enhancing speculative decoding. We conducted comprehensive evaluations of KOALA, including both autoregressive and non-autoregressive draft heads across various tasks, demonstrating a latency speedup ratio improvement of 0.24x-0.41x, which is 10.57%-14.09% faster than the original draft heads.

Key facts. Cited by 6 papers in the map (data through 14 Oct 2026). Placed in Speculative decoding, one of 12 directions in Efficient LLM inference and serving. Builds on 31 of its 38 references that are in the map. First posted on arXiv 15 Aug 2024.

Builds on

31 of 38 references.

Closest ancestors · 13

Same direction, newest first

  1. 2024EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees24 Jun
  2. 2024Clover: Regressive Lightweight Speculative Decoding with Sequential Knowledge1 May
  3. 2024Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding7 Feb
  4. 2024EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty26 Jan
  5. 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads19 Jan
  6. 2024Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding15 Jan
  7. 2023SpecTr: Fast Speculative Decoding via Optimal Transport23 Oct
  8. 2023DistillSpec: Improving Speculative Decoding via Knowledge Distillation12 Oct
  9. 2023Accelerating LLM Inference with Staged Speculative Decoding8 Aug
  10. 2023Accelerating Transformer Inference for Translation via Parallel Decoding17 May
  11. 2023Accelerating Large Language Model Decoding with Speculative Sampling2 Feb
  12. 2022Fast Inference from Transformers via Speculative Decoding30 Nov
1 more
  1. 2022Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation30 Mar

Other references · 18

  1. 2017Attention Is All You Needcited by 60,361
  2. 2023GPT-4 Technical Reportcited by 18,396
  3. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  4. 2023Judging LLM-as-a-Judge with MT-Bench and Chatbot Arenacited by 8,626
  5. 2023QLoRA: Efficient Finetuning of Quantized LLMscited by 3,418
  6. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,346
  7. 2023PaLM 2 Technical Reportcited by 1,162
  8. 2023SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shotcited by 1,121
  9. 2023FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performancecited by 853
  10. 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 832
  11. 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 717
  12. 2024Better & Faster Large Language Models via Multi-token Predictioncited by 360
6 more
  1. 2023Adapting Language Models to Compress Contextscited by 356
  2. 2020DeeBERT: Dynamic Early Exiting for Accelerating BERT Inferencecited by 332
  3. 2018Blockwise Parallel Decoding for Deep Autoregressive Modelscited by 323
  4. 2020BERT Loses Patience: Fast and Robust Inference with Early Exitcited by 273
  5. 2023Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Timecited by 252
  6. 2023Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decodingcited by 40

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 4; 2026: 2.

6 papers.

  1. 2025Add-One-In: Incremental Sample Selection for Large Language Models via a Choice-Based Greedy Paradigmcited by 6
  2. 2025Advancing Decoding Strategies: Enhancements in Locally Typical Sampling for LLMscited by 4
  3. 2025Hawk: Leveraging Spatial Context for Faster Autoregressive Text-to-Image Generationcited by 3
  4. 2025Pre-Training Curriculum for Multi-Token Prediction in Language Modelscited by 2
  5. 2026FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibrationcited by 2
  6. 2026ResiSpec: Enhancing Multi-Candidate Speculative Sampling via Residual Distribution Shapingcited by 1

Similar, not cited

Close in wording, no citation either way

  1. 2024Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters
  2. 2024Closer Look at Efficient Inference Methods: A Survey of Speculative Decoding
  3. 2024LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
  4. 2025KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization
  5. 2025Scaling Laws for Speculative Decoding
  6. 2026LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding
  7. 2024Direct Alignment of Draft Model for Speculative Decoding with Chat-Fine-Tuned LLMs
  8. 2024Decoding Speculative Decoding