GRIFFIN: Effective Token Alignment for Faster Speculative Decoding

TL;DR

GRIFFIN tackles token misalignment in speculative decoding for large language models by introducing a token-alignable training strategy and a token-alignable draft model. It employs a dynamic loss masking scheme based on top-$k$ predictions and a Token-Guided Fusion architecture to align draft and target features across multi-pass decoding. Across multiple backbones (LLaMA2/3, Vicuna, Qwen, Mixtral) and tasks (MT-Bench, HumanEval, GSM8K), GRIFFIN delivers higher acceptance lengths and faster speedups than state-of-the-art methods, with ablations validating the necessity of its components. The approach is compatible with existing training workflows and is openly released, enabling practical acceleration for real-world LLM inference.

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.

Figures (3)

• Figure 1: Comparison between our GRIFFIN, EAGLE2, and HASS. (a) Speedup ratio comparison. (b) Acceptance length under different training steps, in which "Step $n$" denotes aligning draft model for $n$ times in training. (c) Misaligned token rate under different forward passes in each drafting-verification cycle, where "Forward $n$" denotes forwarding $n$ passes to generate $n$ draft tokens.
• Figure 2: Token and feature misalignment in EAGLE.
• Figure 3: Structure of GRIFFIN's darft model. (a) Token-Alignable Draft Model. (b) TGF module. The diagram depicts the shared architecture used in both training and decoding phases—arrows indicating token flow correspond to valid data dependencies in both regimes.