GRIFFIN: Effective Token Alignment for Faster Speculative Decoding

1 Introduction

• Training uses target features and ground-truth tokens, whereas decoding uses draft-generated features and tokens, creating feature and token misalignment that degrades acceptance and speedup.The passage relates these discrepancies to exposure bias.Source: §1 · pp. 1–2
• 20% higher average acceptance length than EAGLE2 and 8% than HASS, with speedup ratios 18% and 7% higher, respectively, on the listed models.The comparison covers LLaMA2-7/13B, LLaMA3-8B, and Vicuna-7B.Source: Figure 1 · p. 2, §1 · p. 2
• EAGLE and HASS address feature misalignment, but token mismatch remains; EAGLE2 reaches 48% token misalignment and HASS reaches 37% in later training steps.The passage notes HASS becomes ineffective for deeper harmonized training at step n ≥3.Source: §1 · p. 2
• Token-alignable training masks loss for tokens whose ground-truth tokens fall within the draft model’s top-k predictions, rather than requiring exact-target matches.This aligns training with decoding’s use of top-k draft trees.Source: §1 · p. 2
• Token-Guided Fusion incorporates input tokens twice to refine draft features, while the draft model and training strategy reinforce one another to reduce misalignment.GRIFFIN maintains lower token misalignment across multiple forward steps, supporting longer accepted sequences and greater speedups in deeper harmonized training.Source: §1 · p. 2

2 Related Work

• Retrieval-based drafters depend on available relevant context, tree-based verification incurs high computational overhead, and resource-reuse methods depend on external-resource quality and access.Chimera and Glide integrate the target model into the draft model at extra computational cost.Source: §2 · p. 3
• Lightweight approaches include Medusa’s MLPs, Hydra and Recurrent Drafter’s RNNs, and EAGLE and EAGLE2’s transformer decoders over feature sequences.FSPAD tailors input sequences and training for lightweight draft-model prediction.Source: §2 · p. 3
• HASS addresses feature misalignment but not token-level mismatch; GRIFFIN adds token-alignable training and a token-alignable draft model to target that gap.The passage characterizes GRIFFIN as complementary to existing methods.Source: §2 · p. 3

3 Motivation: Token Misalignment

• Speculative decoding drafts multiple tokens with a lightweight model, then verifies and accepts a subset in one target-model forward pass.Source: §3 · p. 3
• EAGLE predicts draft tokens and hidden features from ground-truth tokens and target features during training, but must use its own generated tokens and features during decoding.This creates both feature misalignment and token misalignment.Source: §3 · p. 3
• 48% token misalignment for EAGLE2 when generating five draft tokens per cycle and 37% for HASS; early token errors compound across passes.Source: §3 · pp. 3–4
• When HASS uses more than three forward passes, acceptance length plateaus despite continued training, limiting decoding efficiency.Source: §3 · p. 4
• Replacing ground-truth training tokens with draft tokens fails because stored target features create inconsistent input-feature pairs, degrading performance and acceptance length.Source: §3 · p. 4

4.1 Token-Alignable Training

• Training progresses across multiple forward passes, increasingly conditioning on draft-generated tokens and features rather than ground-truth inputs.This makes training inputs more like those encountered during decoding.Source: §4.1 · p. 4
• Top-k predictability masks identify aligned tokens, and a cumulative alignment mask tracks whether earlier unpredictable tokens make later positions misaligned.The strategy avoids penalizing positions following tokens that decoding would reject.Source: §4.1 · p. 5

4.2 Token-Alignable Draft Model

• The draft model places Token-Guided Fusion before autoregression and uses a Token-Enhanced Head afterward for token prediction and feature generation.The head separates the token-prediction and feature-generation objectives.Source: §4.2 · p. 5
• Token-Guided Fusion uses token embeddings alongside draft features to address persistent feature mismatch with the target model.It normalizes and expands the fused representation, then applies nonlinear refinement and a residual connection.Source: §4.2 · pp. 5–6

5.1 Comparison with SoTAs

• 5–6 tokens per drafting-verification cycle: GRIFFIN has the highest acceptance length and speedup ratio across tested datasets and LLMs.This is roughly three times the amount of standard speculative sampling and 1.5 times the amount of EAGLE.Source: Table 1 · p. 7, §5.1 · p. 7

5.2 Ablation Study

• Removing TAT lowers acceptance length by 0.26/0.28 and speedup ratio by 0.18/0.24 at T=0/1; removing both components produces the steepest decline.Removing TAD lowers acceptance length by 0.19/0.23 and speedup ratio by 0.12/0.14 at T=0/1.Source: Table 2 · p. 8, §5.2 · p. 9

5.1 Comparison with SoTAs

• 13.8% higher speedup ratio than HASS on HumanEval with Vicuna 7B at T=1; GRIFFIN also outperforms EAGLE-3 across LLaMA3 8B benchmarks and temperatures.On Qwen2 7B, the speedup improvement over HASS is 5.8%; on Mixtral-8x7B it is more than 6.6%.Source: §5.1 · p. 8

5.2 Ablation Study

• Acceptance length and speedup ratio peak at k=3; aligning top-k tokens from 1 to 10 improves both metrics over no token alignment.Aligning only the top-1 token is less effective than including additional candidate tokens.Source: Table 3 · p. 8, §5.2 · p. 9
• Increasing TAT steps improves acceptance length through step 5, but gains taper at steps 4–5 as aligned tokens decline; experiments use N=3 for fair HASS comparison.Source: §5.2 · p. 9, Table 4 · p. 9
• 0.83 acceptance-length and 0.48 speedup-ratio reductions follow replacing token embeddings x with raw features F; replacing them with fused h yields reductions of 0.26 and 0.2.Source: Table 5 · p. 9, §5.2 · p. 10
• At Wu=4,096, acceptance length and speedup ratio decline; at Wu=22,016, acceptance length improves slightly while speedup ratio falls with added computation.Source: §5.2 · p. 10

• GRIFFIN combines training that excludes misaligned tokens from loss computation with a draft model that reduces misalignment, and evaluations report the highest speedup ratios and acceptance lengths.Source: §6 · p. 10
• Multi-step token-alignable training adds overhead compared with EAGLE, while the draft model is trained once and inference is described as the primary bottleneck.The paper reports comparable overall training overhead to HASS and over 7% higher speedup ratio than HASS under the same training cost.Source: §6 · p. 10
• The paper says faster decoding could make real-time LLM applications more feasible, including education, healthcare assistants, accessibility tools, and scientific research.Source: §6 · p. 10

NeurIPS Paper Checklist

• The checklist says limitations and positive societal impacts are discussed, the research conforms to the NeurIPS Code of Ethics, and the work poses no model-release risks.It marks theoretical results, human-subject research, and core-method LLM usage as not applicable, and says existing assets are credited with license compliance.Source: NeurIPS Paper Checklist · pp. 13, 16, 17, 18
• The checklist says the authors document datasets, compute resources, training methods, and hyperparameters, and provide code and reproduction guidelines; it also says the new draft models are documented.The code and draft models are made accessible through the project repository.Source: NeurIPS Paper Checklist · pp. 14, 15, 18
• The checklist reports that benchmarking omits error bars, citing deterministic inference workloads with minimal variance and asserting that speedup ratio and acceptance length are statistically significant.The authors say they follow prior speculative-decoding benchmarking practice and run the same inference workloads multiple times per dataset.Source: NeurIPS Paper Checklist · pp. 15–16

A.1 Motivation Behind TGF

• TGF uses token embeddings to guide fusion and better align draft features with target representations, addressing misalignment that feature-level loss alone cannot eliminate.The authors say ablation results in Table 5 show reduced feature inconsistency and empirical gains from TGF.Source: §A.1 · p. 20
• TGF separately normalizes fused features and token embeddings, then expands feature dimensionality with an Up Projector to improve training stability and representation expressiveness.These enhancements are identified in Fig. 3 (b-ii).Source: §A.1 · p. 20
• SiLU adds nonlinear capacity to recover complex feature interactions, while the Down Projector compresses enriched representations into stable fused features for downstream processing.Source: §A.1 · p. 20
• TGF improves acceptance lengths and speedup ratio, indicating a more accurate approximation of the target model’s output space.The passage points to Table 1 for these measurable improvements but provides no numerical values.Source: §A.1 · p. 20

A.2 Ablation Study on the TGF Architecture

• Removing the initially fused feature h reduced acceptance length by 2.63 (T = 0) and 2.49 (T = 1), and speedup ratio by 1.50 (T = 0) and 1.39 (T = 1).The authors identify this as the most severe degradation.Source: Table 6 · p. 20, §A.2 · p. 21
• Removing the Up Projector and SiLU activation reduced acceptance length by 0.23 (T = 0) and 0.27 (T = 1), and speedup ratio by 0.17 (T = 0) and 0.21 (T = 1).The ablation removes both components simultaneously.Source: Table 6 · p. 20, §A.2 · p. 21
• Removing token embeddings x from secondary fusion reduced acceptance length by 0.10 (T = 0) and 0.13 (T = 1), and speedup ratio by 0.07 (T = 0) and 0.08 (T = 1).The authors connect this loss to reduced injection of token-specific information.Source: Table 6 · p. 20, §A.2 · p. 21

A.3 Ablation on the Expansion Dimension of TGF

• At expansion dimension 4,096, acceptance length fell by 0.12 (T = 0) and 0.13 (T = 1), while speedup ratio fell by 0.06 at both temperatures.GRIFFIN’s chosen dimension is 11,008, matching the target model’s FFN intermediate size.Source: Table 7 · p. 21, §A.3 · p. 21
• At expansion dimension 22,016, acceptance length rose by 0.06 (T = 0) and 0.02 (T = 1), but speedup ratio fell by 0.18 (T = 0) and 0.25 (T = 1).The larger setting adds 0.14B parameters; the authors favor matching the target FFN intermediate size for efficiency.Source: Table 7 · p. 21, §A.3 · p. 21

B Effectiveness of Token-Alignable Draft Model (TAD) Components

• Removing TGF reduced acceptance length by 0.15 and speedup ratio by 0.09 at T = 0 and 0.11 at T = 1.Removing TEH reduced average acceptance length by 0.06 and speedup ratio by 0.04; the reported TGF degradation is greater.Source: Table 8 · p. 22, B Effectiveness of Token-Alignable… · p. 22

C.1 Loss function

• The per-step loss combines weighted token-prediction loss and feature-alignment loss; the latter minimizes L1 distance between draft and target hidden features.The feature term regularizes the draft model toward the target model’s latent space for stable multi-pass alignment.Source: §C.1 · p. 22

C.2 Draft Tree Structure

• Experiments use a dynamic draft tree with 60 draft tokens and depth 6.The configuration follows the optimal setting established in EAGLE-2 and HASS.Source: §C.2 · p. 22

C.3 Training Configuration

• Draft-model training uses AdamW, batch size 4 per GPU, 20 epochs, and 800,000 total steps.The settings also include 2,000 warmup steps, gradient clipping at 0.5, and a maximum sequence length of 2,048 tokens.Source: §C.3 · p. 23

D Clarification of Baseline Methods

• The authors used released draft parameters or official code for several baselines, training HASS and FSPAD models themselves where parameters were unavailable.For LLaMA3-8B, EAGLE-3 was trained on ShareGPT alone, requiring over 300 A100-80G GPU hours; the full protocol was estimated at approximately 2,400 GPU hours.Source: D Clarification of Baseline Methods · p. 23

E Parameter Sizes of GRIFFIN’s Draft Models

• GRIFFIN draft models range from 0.41B to 2.07B parameters across 7B–70B targets and are 0.45B for Mixtral-8x7B; their average speedup improvement exceeds 8%.The draft models contain 0.17B–1.08B more parameters than EAGLE-2 and HASS, while remaining similar in size to FSPAD.Source: E Parameter Sizes of GRIFFIN’s Draft… · pp. 23–24

• All compared methods use ShareGPT for draft-model training and the same number of training tokens.Source: F Training Overhead of GRIFFIN · p. 24
• 150, 250, and 600 NVIDIA A100 80G GPU hours are reported for GRIFFIN’s 7B, 13B, and 70B models, versus 130, 220, and 500 hours for HASS.Source: F Training Overhead of GRIFFIN · p. 24
• Roughly 8% improvement in speculative decoding speed over HASS is reported, with the paper stating that inference efficiency justifies the slightly higher training overhead.The draft model is trained once and used extensively during inference.Source: F Training Overhead of GRIFFIN · p. 24

G Throughput of GRIFFIN

• 6–11% higher throughput than EAGLE and 4–8% higher than HASS is reported across tested batch sizes.Source: Table 9 · p. 24, G Throughput of GRIFFIN · p. 24
• The vLLM evaluation uses LLaMA3-8B-Instruct on MT-Bench at temperature 0, with sequential speculation capped at chain length 2 rather than tree-based drafting.Its throughput values are not directly comparable to main-text results using native tree-structured speculation.Source: G Throughput of GRIFFIN · p. 24, Table 9 · p. 24
• As batch size increases, target-model GPU utilization improves while draft-model memory and compute overhead grows, narrowing relative speedup.Source: G Throughput of GRIFFIN · p. 24

H Breakdown of Decoding Latency

• The analysis attributes the gap between acceptance length and speedup ratio primarily to the non-negligible latency of draft-model inference.Source: H Breakdown of Decoding Latency · p. 25
• The latency breakdown models speculative decoding as target verification of τ candidate tokens plus draft-model passes over rollout depth d, with average draft-pass latency t̄.Source: H Breakdown of Decoding Latency · p. 25
• With t = 25 ms, t̄ = 1.5 ms, τ = 5, and d = 6 on LLaMA3-8B-Instruct/A100-80G, the estimate closely matches measurements and identifies draft latency as the primary speedup limit.Source: H Breakdown of Decoding Latency · p. 25
• Draft-model latency accumulates across rollout steps, particularly in deep or large-d configurations; proposed directions include distillation, asynchronous drafting, and kernel fusion or caching.Source: H Breakdown of Decoding Latency · p. 25

• The comparison uses LLaMA3-8B-Instruct on MT-Bench, HumanEval, and GSM8K at temperatures 0 and 1, with EAGLE-3 trained using matching optimizer, schedule, and architecture settings.Training used ShareGPT only, excluding the additional UltraChat-200K corpus.Source: I Discussion with EAGLE-3 · p. 25
• Removing GRIFFIN’s feature-prediction term degrades speedup ratio and acceptance length by roughly 10–15% across evaluation settings.Source: I Discussion with EAGLE-3 · p. 26, Table 10 · p. 26
• Under identical training data and inference settings, full GRIFFIN outperforms EAGLE-3 on all metrics reported across the three benchmarks and two temperature settings.Source: I Discussion with EAGLE-3 · p. 26, Table 10 · p. 26
• GRIFFIN’s Token-Alignable Draft architecture and Token-Alignable Training procedure target token misalignment, and could in principle be applied to EAGLE-3-style models without changing their external decoding interface.Source: I Discussion with EAGLE-3 · p. 26
• The study concludes that feature-level supervision remains beneficial with large-scale data and that token-alignment provides complementary improvements beyond EAGLE-3.Source: I Discussion with EAGLE-3 · p. 26