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

TL;DR

KOALA tackles high inference latency in LLMs by enhancing speculative decoding through a multi-layer draft head and adversarial training. By expanding single-layer draft heads into a multi-layer architecture and coupling them with a discriminator-driven adversarial objective, KOALA narrows the gap between draft heads and target LLMs, improving next-token prediction accuracy. Empirical results on Medusa and EAGLE across Vicuna models show latency speedups of $0.24x$ to $0.41x$, outperforming original draft heads by up to $14.09\%$ in speed, with higher average acceptance lengths and token efficiency. The approach offers a practical, orthogonal enhancement to speculative decoding that can be integrated with existing draft-heads to accelerate LLM inference with controlled overheads.

Abstract

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.

Figures (7)

• Figure 1: Comparison between the traditional draft head (upper panel) and the KOALA-optimized draft head (lower panel). KOALA expands the conventional single-layer structure to a multi-layer architecture and incorporates adversarial learning into traditional supervised training. While KOALA slightly increases drafting overhead, it substantially enhances speculative decoding efficiency by improving the draft head's accuracy in predicting subsequent tokens.
• Figure 2: Comparison of single-layer and multi-layer draft head structures. For each Medusa Head, KOALA expands the single $\text{ResBlock}$ to $K$ layers. In the EAGLE Head, KOALA extends the single $\text{Decoder Layer}$ to $K$ layers. For simplicity, each draft head predicts only the next two tokens, $\bar{x}_1$ and $\bar{x}_2$, based on the input sequence $x_1, x_2, \cdots, x_n$.
• Figure 3: Training process for multi-layer draft heads which incorporates adversarial learning into supervised training. The target LLM, featuring a snowflake logo, and its parameters remain unupdated throughout the process. The discriminator and draft head are trained adversarially, co-evolving until they reach a Nash equilibrium, whereupon the training terminates.
• Figure 4: Speedup ratios of Medusa, EAGLE, and their KOALA-optimized versions achieving maximum speedup improvement, denoted by superscript $^{\bigstar}$. All configurations achieve maximum speedup at $K$ = 2, except Medusa on Vicuna-33B, which peaks at $K$ = 3.
• Figure 5: Speedup ratios of Medusa and EAGLE with varying layer structures. "M w/ 1" and "E w/ 1" represent the original single-layer Medusa and EAGLE, respectively.