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

Kaiqi Zhang; Jing Zhao; Rui Chen

Home/Archive/arXiv

arXiv · 2408.08146

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

Kaiqi Zhang·Jing Zhao·Rui Chen

View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

cs.CL

Keep this discovery

Explore connections, maps & timelines

BibTeX RIS

Kaiqi Zhang, Jing Zhao, Rui Chen. 2024-08-15. KOALA: Enhancing Speculative Decoding for LLM via Multi-Layer Draft Heads with Adversarial Learning. https://arxiv.org/abs/2408.08146

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Universal Topological Regularity of Syntactic Structures

Despite their widespread use, the principles governing the organisation of syntactic dependency trees remain poorly understood. I analyse dependency trees from 124 typologically, genetically, and geographically diverse languages. Their topology departs systematically from randomness. Relative to uniformly sampled random trees, dependency trees exhibit greater structural robustness and lower branching heterogeneity. I propose that these universal regularities emerge naturally from incremental grammatical encoding. I model this process using sublinear preferential attachment. The model accurately reproduces the observed topology. More generally, the results demonstrate how a universal statistical property of syntax can emerge from a simple, cognitively motivated generative process. They further illustrate a broader principle of efficiency by construction: communicatively efficient syntactic structures can emerge without direct optimisation for communication.

cs.CL↗

arXiv2026

An Incomplete Loop: Deductive, Inductive, and Abductive Reasoning in Language Models

Modern language models (LMs) can learn to perform new tasks in different ways: in instruction following, the target task is described explicitly in natural language; in few-shot prompting, the task is specified implicitly with a small number of examples; in instruction inference, LMs are presented with in-context examples and are then prompted to generate a natural language task description before making predictions. Each of these procedures may be thought of as invoking a different form of reasoning: instruction following involves deductive reasoning, few-shot prompting involves inductive reasoning, and instruction inference involves abductive reasoning. How do these different capabilities relate? Across four LMs (from the gpt and llama families) and two learning problems (involving arithmetic functions and machine translation) we find a strong dissociation between the different types of reasoning: LMs can sometimes learn effectively from few-shot prompts even when they are unable to explain their own prediction rules; conversely, they sometimes infer useful task descriptions while completely failing to learn from human-generated descriptions of the same task. Our results highlight the non-systematic nature of reasoning even in some of today's largest LMs, and underscore the fact that very different learning mechanisms may be invoked by seemingly similar prompting procedures.

cs.CL↗

arXiv2026

Thinking beyond the anthropomorphic paradigm benefits LLM research

Anthropomorphism, or the attribution of human traits to technology, is an automatic and unconscious response that occurs even in those with advanced technical expertise. In this position paper, we analyze hundreds of thousands of research articles to present empirical evidence of the prevalence and growth of anthropomorphic terminology in research on large language models (LLMs). We argue for challenging the deeper assumptions reflected in this terminology -- which, though often useful, may inadvertently constrain LLM development -- and broadening beyond them to open new pathways for understanding and improving LLMs. Specifically, we identify and examine five anthropomorphic assumptions that shape research across the LLM development lifecycle. For each assumption (e.g., that LLMs must use natural language for reasoning, or that they should be evaluated on benchmarks originally meant for humans), we demonstrate empirical, non-anthropomorphic alternatives that remain under-explored yet offer promising directions for LLM research and development.

cs.CL↗