Learning Harmonized Representations for Speculative Sampling
This work addresses a specific bottleneck in speculative sampling for LLM decoding, offering incremental improvements in speed for practical applications.
The paper tackles the problem of inconsistent context and objectives in speculative sampling for accelerating LLM decoding, proposing HASS which achieves 2.81x-4.05x wall-clock time speedup, surpassing EAGLE-2 by 8%-20%.
Speculative sampling is a promising approach to accelerate the decoding stage for Large Language Models (LLMs). Recent advancements that leverage target LLM's contextual information, such as hidden states and KV cache, have shown significant practical improvements. However, these approaches suffer from inconsistent context between training and decoding. We also observe another discrepancy between the training and decoding objectives in existing speculative sampling methods. In this work, we propose a solution named HArmonized Speculative Sampling (HASS) that learns harmonized representations to address these issues. HASS accelerates the decoding stage without adding inference overhead through harmonized objective distillation and harmonized context alignment. Experiments on four LLaMA models demonstrate that HASS achieves 2.81x-4.05x wall-clock time speedup ratio averaging across three datasets, surpassing EAGLE-2 by 8%-20%. The code is available at https://github.com/HArmonizedSS/HASS.