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

arXiv: https://arxiv.org/abs/2408.08146 · HTML: https://mlsanity.com/paper/2408.08146/

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

## Placed in

- Efficient LLM inference and serving / [Speculative decoding](https://mlsanity.com/inference/speculative-decoding/)

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [PaLM 2 Technical Report](https://mlsanity.com/paper/2305.10403/) (2023)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [Accelerating Large Language Model Decoding with Speculative Sampling](https://mlsanity.com/paper/2302.01318/) (2023)
- [FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance](https://mlsanity.com/paper/2305.05176/) (2023)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://mlsanity.com/paper/2401.15077/) (2024)
- [EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees](https://mlsanity.com/paper/2406.16858/) (2024)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [Adapting Language Models to Compress Contexts](https://mlsanity.com/paper/2305.14788/) (2023)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Blockwise Parallel Decoding for Deep Autoregressive Models](https://mlsanity.com/paper/1811.03115/) (2018)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding](https://mlsanity.com/paper/2401.07851/) (2024)
- [Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time](https://mlsanity.com/paper/2310.17157/) (2023)
- [Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation](https://mlsanity.com/paper/2203.16487/) (2022)
- [DistillSpec: Improving Speculative Decoding via Knowledge Distillation](https://mlsanity.com/paper/2310.08461/) (2023)
- [Accelerating LLM Inference with Staged Speculative Decoding](https://mlsanity.com/paper/2308.04623/) (2023)
- [SpecTr: Fast Speculative Decoding via Optimal Transport](https://mlsanity.com/paper/2310.15141/) (2023)
- [Accelerating Transformer Inference for Translation via Parallel Decoding](https://mlsanity.com/paper/2305.10427/) (2023)
- [Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding](https://mlsanity.com/paper/2402.05109/) (2024)
- [Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decoding](https://mlsanity.com/paper/2307.05908/) (2023)

## Built on by (most cited first, 6 of 6)

- [Add-One-In: Incremental Sample Selection for Large Language Models via a Choice-Based Greedy Paradigm](https://mlsanity.com/paper/2503.02359/) (2025)
- [Advancing Decoding Strategies: Enhancements in Locally Typical Sampling for LLMs](https://mlsanity.com/paper/2506.05387/) (2025)
- [Hawk: Leveraging Spatial Context for Faster Autoregressive Text-to-Image Generation](https://mlsanity.com/paper/2510.25739/) (2025)
- [Pre-Training Curriculum for Multi-Token Prediction in Language Models](https://mlsanity.com/paper/2505.22757/) (2025)
- [FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration](https://mlsanity.com/paper/2605.20022/) (2026)
- [ResiSpec: Enhancing Multi-Candidate Speculative Sampling via Residual Distribution Shaping](https://mlsanity.com/paper/2608.24411/) (2026)

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