GRIFFIN: Effective Token Alignment for Faster Speculative Decoding

TL;DR

GRIFFIN introduces a token alignment strategy for speculative decoding in large language models, significantly improving inference speed and draft token acceptance rates by addressing token misalignment issues.

Contribution

It proposes a novel token-alignable training and draft model framework that reduces misalignment, enhancing the efficiency of speculative decoding in LLMs.

Findings

Over 8% increase in acceptance length

More than 7% speedup ratio

Outperforms existing speculative decoding methods

Abstract

Speculative decoding accelerates inference in large language models (LLMs) by generating multiple draft tokens simultaneously. However, existing methods often struggle with token misalignment between the training and decoding phases, limiting their performance. To address this, we propose GRIFFIN, a novel framework that incorporates a token-alignable training strategy and a token-alignable draft model to mitigate misalignment. The training strategy employs a loss masking mechanism to exclude highly misaligned tokens during training, preventing them from negatively impacting the draft model's optimization. The token-alignable draft model introduces input tokens to correct inconsistencies in generated features. Experiments on LLaMA, Vicuna, Qwen and Mixtral models demonstrate that GRIFFIN achieves an average acceptance length improvement of over 8% and a speedup ratio exceeding 7%, outperforming current speculative decoding state-of-the-art methods. Our code and GRIFFIN's draft models are released publicly in https://github.com/hsj576/GRIFFIN.