GRIFFIN: Effective Token Alignment for Faster Speculative Decoding

Thank you for your thoughtful feedback. Below, we provide a point-by-point response to address your concerns.

1. Goal of Token Alignment. There are some misunderstandings. We would like to clarify the distinctions and objectives of our method.

Feature Alignment vs Token Alignment. As illustrated in Fig. 2, feature misalignment occurs because during decoding, the draft model uses draft features $\bar{\mathbf{F}} _ t$ to generate $\bar{\mathbf{F}} _ {t+1}$, while during training it uses target features $\mathbf{F} _ t$. Token misalignment arises because during decoding, the draft model uses draft tokens $\bar{\mathbf{x}} _ t$ to generate $\bar{\mathbf{F}} _ {t+1}$, while during training it uses ground truth tokens $\mathbf{x} _ t$.

Feature alignment addresses feature misalignment by replacing $\mathbf{F} _t$ with $\bar{\mathbf{F}} _t$ during training. However, token misalignment cannot be directly resolved by simply replacing $\mathbf{x} _t$ with $\bar{\mathbf{x}} _t$ because we follow the fixed-dataset setting adopted by EAGLE and HASS. We provide the detailed explanations for this constraint in Section 3, lines 124-131. Therefore, GRIFFIN simultaneously 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.

Why Token Alignment ≠ "Dropping Hard Predictions". The goal of our masking is not to discard tokens whose ground-truth falls outside the top-$k$ of the draft model. Here we take one example for illustration. Assume the draft model have generated five draft tokens $\bar{\mathbf x} _{t+1:t+5}$. If $\mathbf x _{t+1}$ and $\mathbf x _{t+2}$ are in top-$k(\bar{\mathbf{x}} _{t+1})$ and top-$k(\bar{\mathbf{x}} _{t+2})$ respectively, but $\mathbf{x} _{t+3}$ is not, we still compute the loss for all three positions $\bar{\mathbf x} _{t+1:t+3}$. Tokens after the first mismatch ($\bar{\mathbf x} _{t+4:t+5}$) are guaranteed to be rejected in verification; training on them provides no useful signal. These positions are therefore masked, not because they are “hard,” but because they are inevitably discarded during speculative decoding.

Corrections and Clarifications. We also checked the whole work to find the reason which leads to your misunderstanding. Firstly, in line 158, we use "aligned" to indicate when the ground-truth token $\mathbf{x} _t$ appears within top-$k$($\bar{\mathbf{x}} _t$). We agree that "predictable" would be more accurate than "aligned" in this context, since this "aligned" does not have the same meaning with "token alignment" which refers to the consistency between training and inference in terms of token level. Secondly, Eqn (3) has an incorrect superscript, and should be $\mathbf{m}_t = \prod _{i=t-n+1}^{t-1}\bar{\mathbf{m}} _i$. This means we do not directly mask out the loss for draft tokens where $\bar{\mathbf{m}} _{t}=0$. We will correct them to enhance clarity and readability.

2. How "gradually" and "progressively" work? Assume the draft model aims to predict $N$ draft tokens $\bar{\mathbf x} _{t+1:t+N}$ by performing forward pass $N$ times:

1. Forward pass 1 (predict $\bar{\mathbf x} _{t+1}$)

   • Inputs: ground-truth tokens $\mathbf{x} _{1:t}$ and target-model features $\mathbf{F} _{1:t}$.

   • Outputs: token $\bar{\mathbf x} _{t+1}$ and feature $\bar{\mathbf F} _{t+1}$.

2. Forward pass $n$ (predict $\bar{\mathbf x} _{t+n}$), $2 ≤ n ≤ N$

   • Inputs:

   – ground-truth tokens $\mathbf{x} _{1:t+n-1}$

   – target-model features $\mathbf{F} _{1:t}$ (unchanged)

   – self-generated draft features $\bar{\mathbf F} _{t+1:t+n-1}$ from the previous $(n−1)$ passes

   • Outputs: token $\bar{\mathbf x} _{t+n}$ and feature $\bar{\mathbf F} _{t+n}$.

Thus, with each additional drafted token, the proportion of self-generated draft features fed back into the draft model increases. This is what we mean by gradually/progressively shifting the model toward using its own outputs.

3. Step & Pass, and why acceptance lengths can be longer than step. We thank the reviewer for raising this important question regarding the usage of the terms step and pass, and for pointing out the potential confusion related to Table 4 and Line 141.

In Line 141, the variable n refers to the index of the training step in GRIFFIN’s progressive training schedule, rather than the total number of forward passes. Specifically, each training step n is designed to train the draft model to perform n consecutive forward passes, thereby predicting n future tokens and their corresponding features. This approach helps the draft model learn to operate under conditions that closely simulate those it will encounter during speculative decoding, where it must rely on its own previously generated tokens and features. For example, at training step 3, the draft model is trained to generate three tokens in a sequence, conditioning on the outputs of its own previous passes.

It is also important to clarify why the acceptance length reported in Table 4 can exceed the value of the training step. The step index defines how the model is trained, not a hard constraint on how many tokens it can generate or how many the target model may accept during decoding. At inference, the draft model is allowed to generate a sequence of any length (e.g., >N), and the target model verifies them all at once. The acceptance length refers to the average number of draft tokens accepted by the target model per draft-verify cycle and is influenced by the quality of the generated tokens. Accordingly, acceptance length can exceed the "step" value used during training.

In our revised manuscript, we will explicitly clarify that "step" refers to the training schedule rather than a decoding limit, and we will add explanatory notes to Table 4 to prevent this misunderstanding. We appreciate the reviewer’s thoughtful feedback, which helps us improve the clarity and precision of our presentation.

4. Loss and Weight in Eqn 2. We follow exactly the same dual-objective loss introduced in EAGLE. The complete definition of loss $\ell$ in Equation 2 is:

$$\ell(\bar{\mathbf{x}}_t, \mathbf{x}_t, \bar{\mathbf{F}}_t, \mathbf{F}_t) = \ell_1(\bar{\mathbf{F}}_t, \mathbf{F}_t) + \omega * \ell_2(\bar{\mathbf{x}}_t, \mathbf{x}_t)$$

where $\ell_1$ is the Smooth L1 distance and $\ell_2$ is the standard cross-entropy loss. For fair comparison, we follow EAGLE and HASS and set the weighted factor $\omega$ to 0.1.

Because we did not tune $\omega$ and simply inherited the canonical setting from prior work, it was omitted from the hyper-parameter by oversight. We will provide full implementation details of the loss function in Appendix C.2 of the revised manuscript.

5. Speedup Ratio and hardware setup. To clarify, the Speedup Ratio reflects the entire speculative decoding process, including both drafting and verification. All evaluation experiments were conducted on a single NVIDIA A100 80G GPU, except for LLaMA3-70B, which required two GPUs due to memory constraints. Both drafting and verification are performed on the same GPU, consistent with prior work (e.g., EAGLE, HASS), ensuring a fair comparison.

Regarding the concern that speculative decoding might be slower than autoregressive decoding on a single GPU: this is a misconception. Although speculative decoding adds some overhead from drafting, it significantly reduces latency by verifying multiple tokens in parallel via a single forward pass of the target LLM model using batched causal masking. In contrast, autoregressive decoding requires one forward pass per token.

This efficiency has been consistently validated in prior work (e.g., EAGLE, HASS), and our results confirm it: GRIFFIN achieves up to 18% speedup over EAGLE-2 and 7% over HASS, using the same GPU configuration. We will clarify this in the final version for clarity.

We greatly appreciate your thoughtful follow-up and your willingness to raise your score. Thank you for your careful reading and constructive engagement throughout the review process. Below, we provide a point-by-point response to address your concerns.

1. KV cache management. Our implementation of KV cache for autoregressive decoding, drafting, and verification strictly follows the standard approaches used in prior works such as EAGLE-2 and HASS. Concretely, every decoding request maintains two independent caches:

• Target-model cache: Persistently stores the key–value pairs of the target model throughout the conversation.
• Draft-model cache: Holds keys and values generated by draft model; this cache is discarded immediately after the corresponding verification step and therefore does not accumulate over the entire sequence.

Detailed implementation specifics can be found in our anonymous code repository (model/kv_cache.py).

2. FLOPs analysis. Your intuition is correct: speculative decoding entails more total FLOPs than autoregressive decoding, since it invokes an extra draft model. However, the draft model is typically less than 5% of the target model’s size, so the incremental FLOPs is minor.

3. Why speculative decoding achieves acceleration? Your understanding of the reasons behind acceleration is correct. The primary reason speculative decoding can accelerate inference is that LLMs are not compute-bound during autoregressive decoding, especially when batch size are relatively small; during this stage, GPU utilization is relatively low. Speculative decoding improves GPU utilization by allowing the target model to verify multiple draft tokens in parallel within a single forward pass. Although the number of FLOPs per forward pass for the target model increases, the time required to verify and accept multiple tokens is comparable to the time needed to generate a single token in conventional autoregressive decoding.

4.GRIFFIN's performance at larger batch sizes. To evaluate GRIFFIN’s performance under batch sizes greater than 1, we integrated GRIFFIN to vLLM following the same speculative decoding interface as EAGLE. We then measured throughput on Llama-3-8B-Instruct using the MT-Bench dataset (temperature $=0$) for batch sizes {$2,4,8,16$}. Because the current vLLM interface does not yet support tree-based drafting, we evaluated all methods in a sequential setting with maximum chain length $=2$. Results are summarised below (baseline $=1.00\times$ is vanilla vLLM decoding):

The observed trend—declining speedup as batch size increases—is theoretically expected. As batch size grows, the target model benefits from better GPU utilization, reducing computational redundancy and diminishing the marginal gain from speculative decoding. Furthermore, large batches begin to approach memory and compute bottlenecks, making the overhead from additional draft model calls more significant relative to the baseline. Nonetheless, GRIFFIN preserves a consistent 6–11% advantage over EAGLE and 4–8% over HASS across all tested batch sizes, demonstrating robust benefits even in high-throughput scenarios.

Thank you for your thoughtful feedback. Below, we provide a point-by-point response to address your concerns.

1. Mask. Thanks. The training mask $\mathbf{m} _t$ was unfortunately misrepresented in our original text due to a typesetting oversight. The correct definition of Equation (3) should be $\mathbf{m} _t = \prod _{i=t-n+1}^{t-1} \bar{\mathbf{m}} _i$. This formulation means that the loss mask at time step $t$ is determined by the prediction quality of the previous $n-1$ draft tokens—not by the current prediction at $t$. Specifically, $\bar{\mathbf{m}} _i = 1$ if the ground truth token at position $i$ is within the top-$k$ predictions of the draft model, and 0 otherwise. Therefore, $\mathbf{m} _t = 1$ only if all previous draft tokens are well-aligned; otherwise, the loss at position $t$ is masked out.

This design is motivated by the behavior of speculative decoding: once a misaligned token is encountered, all subsequent draft tokens are rejected. By mirroring this decoding logic during training, we ensure the model is only updated on token sequences that would have been accepted in actual decoding scenarios. In this way, the loss focuses on segments of the draft that contribute valid outputs, thereby improving the alignment between training and inference distributions.

We will revise the equation and accompanying explanation in our manuscript for clarity.

2. Breakdown of Decoding Latency. Thank you for the thoughtful question. The key reason the speedup ratio is lower than the acceptance length in Table 1 is due to the non-negligible overhead of draft model inference.

To elaborate, let $N$ be the total number of tokens generated. Then the total latency for autoregressive decoding is $\mathbf{T}_a = N \cdot \mathbf{t}$, where $\mathbf{t}$ is the latency per forward pass of the target model. By comparison, the total latency for speculative decoding is $\mathbf{T}_s = \frac{N}{\tau} \cdot \left( \mathbf{t} + \mathbf{d} \cdot \bar{\mathbf{t}} \right)$, where $\tau$ is the average number of tokens accepted per cycle, $\bar{\mathbf{t}}$ is the draft model's per-pass latency, and $\mathbf{d}$ is the draft tree depth. Thus, the speedup ratio becomes:

$$\text{Speedup Ratio (SR)} = \frac{\mathbf{T}_a}{\mathbf{T}_s} = \frac{\mathbf{t}}{\mathbf{t} + \mathbf{d} \cdot \bar{\mathbf{t}}} \cdot \tau$$

Using LLaMA3-8B-Instruct on A100-80G as a representative example, the latency per forward pass of target model is approximately 20-30ms, for draft model is 1-2ms. So, if we set $\mathbf{t} = 25ms$, $\bar{\mathbf{t}} = 1.5ms$, $\tau = 5$ and $\mathbf{d} = 6$, we could roughly calculate the speedup ratio $SR = 3.68$ , which is closely matches our reported experimental results, confirming that the drafting overhead is the primary factor limiting speedup below the acceptance length. We will include a breakdown of draft vs. target model latency in the Appendix H in the revision.

3. Training time comparision. GRIFFIN’s progressive training introduces some overhead due to multiple forward passes per training example. To quantify this, we compare the total training cost of GRIFFIN against HASS, a strong state-of-the-art baseline that also uses a multi-stage training schedule:

While the training cost increases modestly, this overhead is amortized over large-scale deployment. Once trained, the draft model serves millions of inference requests. GRIFFIN yields up to 7% higher decoding speed over HASS at inference (Table 1), making it significantly more efficient in production.

For example, in a system serving ~10,000 requests/day, let alone the billion requests/day of ChatGPTs and other popular LLMs, the additional training cost would be offset by inference gains within one week of deployment. We discuss this trade-off in detail in Appendix F, and believe the training overhead is well-justified by the lasting efficiency gains in real-world usage.

4. Discussion with EAGLE-3. Thank you for opening this discussion about EAGLE-3’s decision to drop feature prediction loss in the draft model. To understand the practical impact of that choice and the possible interaction with our own techniques, we ran both (i) an ablation on GRIFFIN that removes feature-level loss and (ii) a head-to-head comparison with a faithfully re-implemented EAGLE-3 draft model on LLaMA3-8B-Instruct.

• $SR$: Speedup Ratio; $\tau$ : Acceptance Length

EAGLE-3 Implementation Details: For EAGLE-3, we used their official open-source codebase to train a draft model on the ShareGPT dataset. We made every effort to follow their protocol:

• Training protocol and hyperparameters were kept fully consistent with those reported in the EAGLE-3 paper (e.g., learning rate schedule, optimizer, batch size, model configuration).
• The only difference is the scale of training data: due to computational constraints, we used only ShareGPT (rather than the much larger UltraChat-200K + ShareGPT combination), as training on the full protocol would require approximately 2,400 A100-80G GPU hours—far beyond our resources for the rebuttal period. Both GRIFFIN and EAGLE-3 used identical ShareGPT training data, ensuring fair comparison.

Key observations:

• Removing feature-level loss from GRIFFIN causes a sizable drop (10%-15%) in both Speedup Ratio and Acceptance Length, confirming that feature-level loss is still important for our token-alignment objective.
• Under identical ShareGPT training data and settings, GRIFFIN still outperforms EAGLE-3. This suggests that EAGLE-3’s “remove feature-level loss” design does NOT provide an advantage at this data scale.

Applicability of GRIFFIN techniques to EAGLE-3:

• EAGLE-3 does not explicitly address token misalignment, whereas GRIFFIN’s Token-Alignable Draft Model (TAD) and Token-Alignable Training (TAT) tackle this issue directly.
• Both TAD and TAT can in principle be integrated with EAGLE-3-style draft models without altering their external interface.

We will include a detailed discussion of these experimental results and their implications in Appendix I of our revised manuscript.

Thank you for your thoughtful feedback. Below, we provide a point-by-point response to address your concerns.

1. Batch sizes > 1. To evaluate GRIFFIN’s performance under batch sizes greater than 1, we integrated GRIFFIN into vLLM, following the same speculative decoding interface as EAGLE. We then measured throughput on LLaMA3-8B-Instruct using the MT-bench dataset with temperature set to 0, across a range of batch sizes. As shown below, GRIFFIN consistently outperforms both EAGLE and HASS at all batch sizes:

Here, the baseline (1.00×) refers to standard vLLM decoding without any speculative decoding methods. It is important to note that these results were obtained under specific constraints imposed by vLLM’s current speculative decoding implementation. Specifically, the setup does not support tree-based drafting, so all evaluations used sequential speculation with a maximum chain length of 2. As a result, these results are not directly comparable to those reported in the main body of our paper, which uses a different decoding backend and configuration.

The observed trend—declining speedup as batch size increases—is theoretically expected. As batch size grows, the target model benefits from better GPU utilization, reducing computational redundancy and diminishing the marginal gain from speculative decoding. Furthermore, large batches begin to approach memory and compute bottlenecks, making the overhead from additional draft model calls more significant relative to the baseline.

Despite these challenges, GRIFFIN shows robust performance, delivering 6–11% higher throughput than EAGLE and 4–8% higher than HASS across all batch sizes. These results suggest that GRIFFIN’s advantages are preserved even under high-throughput, batch-inference settings, which are common in real-world deployments.

We will include this analysis in Appendix G of the revised manuscript, along with: (1) full implementation details of the vLLM integration; (2) head-to-head comparisons with other speculative decoding methods in batch settings; and (3) a discussion of potential optimization strategies for improving speculative decoding under large-batch conditions.

2. Comparison with EAGLE-3. We compared GRIFFIN with EAGLE-3 on LLaMA3-8B-Instruct across three benchmarks (MT-bench, HumanEval, GSM8K) at two temperature settings:

• $SR$: Speedup Ratio; $\tau$ : Acceptance Length

Since EAGLE-3 doesn't provide pre-trained draft models for LLaMA3-8B, we used their official code to train on the ShareGPT dataset, maintaining all other hyperparameters consistent with their paper. Training EAGLE-3 on ShareGPT alone required over 300 A100-80G GPU hours. Following their paper's full protocol (UltraChat-200K + ShareGPT) would require approximately 2,400 GPU hours, which exceeded our rebuttal period resources. However, both methods used identical ShareGPT training data, ensuring fair comparison.

GRIFFIN outperforms EAGLE-3 across all benchmarks and temperature settings. At $t=0$, GRIFFIN achieves 3.7% higher speedup ratio and 4.5% higher acceptance length compared with EAGLE-3. At $t=1$, GRIFFIN achieves 3.8% higher speedup ratio and 3.9% higher acceptance length compared with EAGLE-3.  Interestingly, EAGLE-3's performance is roughly comparable to HASS levels, which aligns with your expectation. This suggests that while EAGLE-3 introduces architectural improvements, the performance gains are modest when training data is limited.

We will include our experimental results of EAGLE-3 in the revised manuscript.

Thank you for your thoughtful feedback. Below, we provide a point-by-point response to address your concerns.

1. Architectural changes may limit the applicability. We would like to clarify that GRIFFIN’s draft model preserves identical input and output interfaces to those used in EAGLE and HASS, ensuring that our approach does not impose restrictive architectural modifications, at least in terms of input and output, and thus maintains broad applicability which means the same applicability as EAGLE and HASS.

Moreover, as shown in Table 1, GRIFFIN consistently demonstrates strong performance across a diverse set of model families (LLaMA, Vicuna, Qwen, Mistral) and sizes (7B, 8B, 13B, 70B), confirming the robustness and generalizability of our method beyond any specific architecture or scale.

Importantly, the two core components of our approach—Token-Alignable Training (TAT) and the Token-Alignable Draft Model (TAD)—are designed to be modular and independently integrable. TAT can be applied to any draft model training pipeline to effectively resolve input misalignment issues, while TAD’s architectural improvements can be adopted in existing draft model designs with minimal structural changes.

2. Adding conceptual summary for clarity. Thank you for your valuable feedback. To improve clarity, we will add a “Conceptual Overview” section at the start of Section 3 in the revised manuscript. This overview will provide a high-level intuition behind the token misalignment problem, outline the logical flow from problem identification to our proposed solutions, and clearly map technical components to their respective roles.

Additionally, we will carefully define key technical terms such as “top-k” and “forward n” upon their first introduction to ensure readers can follow the methods more easily. We believe these changes will significantly enhance the paper’s readability and accessibility without sacrificing technical rigor.

3. Draft model trained from scratch? GRIFFIN’s draft model does indeed need to be trained from scratch, rather than fine-tuned from existing pretrained draft models. This is primarily due to the architectural modifications introduced in our Token-Alignable Draft Model (TAD), which alter the internal computation flow in a way that is not compatible with the weight initialization of standard pretrained draft models of EAGLE and HASS. However, training the draft model remains substantially less expensive than training full-scale language models, and we believe the efficiency gains during inference justify the added training cost.

4. Training data size for updating the draft model. In our experiments, we followed established practice by using the entire ShareGPT dataset, which aligns with the standard scale adopted in prior speculative decoding work such as EAGLE and HASS.

Our current paper does not include a systematic analysis of the relationship between training data volume and draft model performance. This is primarily due to computational constraints—specifically, limited GPU availability during the rebuttal period—which prevented us from conducting large-scale data scaling experiments.

With additional resources, we would pursue a more thorough investigation on this good direction, as we believe it is a critical step toward making speculative decoding stronger.

We would like to sincerely thank the Area Chair and the reviewers for your thorough and thoughtful feedback. We greatly appreciate the time and effort you have dedicated to evaluating our submission. Below, we provide a point-by-point response to address each of your concerns and clarify aspects of our work.

1. Mask Indexing

Clarification of the Mask–Indexing Scheme (Eq. 2, lines 160–161). We thank the AC and the reviewers for pointing out the imprecision in our exposition of the masking strategy. The essential point is that token misalignment arises only from the second forward pass onward. Accordingly, the binary mask $\mathbf m_t$ must not influence the first–pass loss but becomes indispensable in later passes.

First Forward Pass ($n=1$). Because the draft model processes the identical prefix of ground truth tokens $\mathbf x _{1:t-1}$ at both training and decoding time, all predicted tokens $\bar{\mathbf{x}} _{t}$ are perfectly aligned. The correct loss at step $t$ therefore reads:

$$\mathcal{L} _{\mathcal{M}}^{(1)} = \ell(\bar{\mathbf{x}} _{t}, \mathbf{x} _{t}, \bar{\mathbf{F}} _{t}, \mathbf{F} _{t}),$$

$n$-th Forward Pass ($n \ge 2$). From the second forward pass, during decoding, speculative decoding may reject a draft token $\bar{ \mathbf x} _ {t}$, in which case all later tokens $\bar{\mathbf x} _ {t+1}, \bar{\mathbf x} _ {t+2}, \ldots$ in that draft sequence are also discarded. So, during training, the subsequent draft tokens $\bar{\mathbf x} _ {t+1},\bar{\mathbf x} _ {t+2},\ldots$ in this draft sequence are misaligned tokens. Training on those misaligned tokens provides no useful signal. Therefore, to prevent the model from being penalized for such inevitably rejected positions, we introduce a cumulative binary mask $\mathbf m _ t = \prod _ {i=t-n+1}^{t-1} \bar{\mathbf m} _ i.$ The loss for the $n$-th forward pass is therefore:

$$\mathcal{L} _ {\mathcal{M}}^{(n)} = \frac{1}{\sum _ {t=1}^{l} \mathbf{m} _ {t}} \sum _ {t=1}^{l} \mathbf{m} _ {t} \ell(\bar{\mathbf{x}} _ {t}, \mathbf{x} _ {t}, \bar{\mathbf{F}} _ {t}, \mathbf{F} _ {t}),$$

Manuscript revisions.

• Equation (2) will be corrected to the unmasked first-pass loss shown above.
• A new equation will explicitly present the masked loss for $n\ge 2$.
• Lines 160–161 will clarify that $\mathbf m _t = 1$ for all positions in the first pass and becomes selective only in subsequent passes.

These changes are notational only; the implementation and all reported results remain intact. We appreciate the reviewers’ careful reading and will incorporate the corrected equations and explanatory remarks in the final version.

2. Justification for the top-$k$ masking rule.

We would like to clarify that the top-$k$ filter is applied to the draft model's logits—not the target model’s. Our key motivation is to ensure that the draft model is trained in a manner that is strictly aligned with its actual behavior during decoding.

During speculative decoding, GRIFFIN follows the same draft tree expansion strategy as EAGLE-2. Specifically, at each speculation step, only the top-$k$ tokens from the draft model’s output are retained, while those ranked below the top-$k$ threshold are pruned. These lower-ranked tokens are never examined or evaluated by the target model and therefore have no influence on acceptance or rejection.

To preserve this operational consistency, we adopt the same top-$k$ criterion during training by masking out tokens beyond the top-$k$ set in the loss computation. This ensures that the learning signal is focused exclusively on those draft tokens that could realistically be explored—and potentially accepted—by the target model during decoding. Thus, our masking strategy is deliberately chosen to align the draft model's training distribution with its decoding-time behavior, enhancing both efficiency and accuracy.

6. Draft step in Evaluation

For comparability with HASS and EAGLE-2, we follow their settings and adopt identical hyper-parameters throughout our experiments. Specifically, the total number of draft tokens per draft–verify cycle is fixed to $60$ for 7B/8B models, $50$ for 13B models, and $48$ for 70B models. These tokens are generated across a draft-tree of depth $6$, i.e., the draft model executes six successive draft steps before the verifier is invoked.

7. Calculation of Misalignment Token Rate

Misalignment Token Rate (MTR) is calculated based on the top-$k$ alignment mask $\mathbf{m} _ t$:

$$\text{MTR} = \frac{\sum _ {t=1}^{l}\mathbf{m} _ {t}}{l}$$

where $l$ denotes the sequence length of training data and $\mathbf{m} _ t$ denotes the alignment mask.

In Fig. 1(c) we report $\text{MTR}$ averaged over the entire ShareGPT training corpus for each forward pass. Crucially, the rate for pass 4 is computed over all tokens in the dataset; it is not exclusively over cases where passes 1-3 of the draft tokens were aligned.

8. Clarification of the term “outputs.”

We acknowledge that the manuscript occasionally uses the word outputs without specifying whether it refers to token outputs or feature outputs. In the revised version, Token outputs and feature outputs will be used only when the context already makes the level (token vs. feature) explicit; otherwise we will employ the full expressions above.

9. Clarification of Training Steps

Thanks. To clarify, our use of “training step” follows the terminology introduced in HASS and refers specifically to the lookahead depth (the number of forward passes for draft model).

Each training step is independent in terms of computation. Specifically:

• At training step $n$, the draft model is trained to perform exactly $n$ forward passes to predict $n$ future tokens.
• We do not accumulate passes from previous steps. That is, training for Step 4 requires only 4 forward passes—not 1+2+3+4 =10.

Training proceeds in stages as follows:

• Step 1. Initialize from $\mathcal{M}_0$ and train for $1$-step to obtain $\mathcal{M}_1$: Starting with initialization parameters $\mathcal{M}_0$, we train the draft model on the entire dataset to perform a single forward pass per position, yielding a new model checkpoint $\mathcal{M}_1$.
• Step 2. Train $\mathcal{M}_1$ for $2$-step to obtain $\mathcal{M}_2$: Using $\mathcal{M}_1$ as initialization, we now train the draft model (over the entire dataset) to perform two forward passes per position, resulting in $\mathcal{M}_2$.
• Step $n$. Train $\mathcal{M}_{n-1}$ for $n$-step to obtain $\mathcal{M}_n$: Recursively, at training step $n$, we initialize from $\mathcal{M}_{n-1}$ and train the draft model to conduct exactly $n$ forward passes per position, producing the updated model $\mathcal{M}_n$.

Thus, during the curriculum phase corresponding to training step $n$, the model is run exactly $n$ forward passes per sequence, not the cumulative sum $1{+}\dotsb{+}n$. This design provides a stable, progressively harder training curriculum: each subsequent step starts from a well-converged checkpoint and adapts it to a deeper draft length.

10. Clarfication of Fig. 3.

Figure 3 depicts the architecture of the proposed token-alignable draft model, not a standalone training diagram. Because this architecture is shared between training and decoding, the figure necessarily shows drafted tokens flowing from one step to the next; such arrows are valid in both regimes. To pre-empt confusion, we will revise the caption to state explicitly that the illustration applies to both training and decoding.

5. Difference with HASS

Clarifying “gradually shifting the model toward using its own outputs’’. The quoted phrase refers to a progressive self-conditioning mechanism where the draft model increasingly relies on its own generated features rather than target features during the training process.

For example, at training step $N \ge 2$, the draft model performs exactly $N$ forward passes to predict the sequence of draft tokens $\bar{\mathbf x} _ {t+1:t+N}$.

1. Forward pass 1 (predict $\bar{\mathbf x} _ {t+1}$):

• Inputs: ground-truth tokens $\mathbf{x} _ {1:t}$ and target-model features $\mathbf{F} _ {1:t}$.
• Outputs: token $\bar{\mathbf x} _ {t+1}$ and feature $\bar{\mathbf F} _ {t+1}$.

2. Forward pass $n$ (predict $\bar{\mathbf x} _ {t+n}$), $2 \leq n \leq N$:

• Inputs:
  • ground-truth tokens $\mathbf{x} _ {1:t}$ (unchanged).
  • target-model features $\mathbf{F} _ {1:t}$ (unchanged).
  • self-generated draft features $\bar{\mathbf F} _ {t+1:t+n-1}$ from the previous $(n-1)$ passes.
  • self-generated draft tokens $\bar{\mathbf x} _ {t+1:t+n-1}$ from the previous $(n-1)$ passes.
• Outputs: token $\bar{\mathbf x} _ {t+n}$ and feature $\bar{\mathbf F} _ {t+n}$.

We can observe that along with $n$ increases, the draft model uses less ground-truth tokens $\mathbf{x} _ {1:t}$ and target-model features $\mathbf{F} _ {1:t}$ as the input, but takes more self-generated draft tokens $\bar{\mathbf x} _ {t+1:t+n-1}$ and draft features $\bar{\mathbf F} _ {t+1:t+n-1}$ from the previous $(n-1)$ passes. This is what we mean by gradually shifting the model toward using its own outputs.

How HASS differs. While GRIFFIN and HASS share similar training schemas, the misalignment each addresses is fundamentally different. As illustrated in Fig. 2,

• Feature misalignment (addressed by HASS) occurs because during decoding, the draft model uses draft features $\bar{\mathbf{F}} _ t$ to generate $\bar{\mathbf{F}} _ {t+1}$, while during training it uses target features $\mathbf{F} _ t$.
• Token misalignment (GRIFFIN focus) arises because during decoding, the draft model uses draft tokens $\bar{\mathbf{x}} _ t$ to generate $\bar{\mathbf{F}} _ {t+1}$, while during training it uses ground truth tokens $\mathbf{x} _ t$.

HASS resolves feature misalignment by aligning the inputs at the feature level—specifically, it replaces $\mathbf F _ t$ with $\bar{\mathbf F} _ t$ during training. Token misalignment, however, cannot be fixed by a naive substitution $\mathbf x _ t \rightarrow \bar{\mathbf x} _ t$ because we operate under the fixed-dataset regime established by EAGLE and HASS (Sec. 3, lines 124–131). The core issue is as follows:

Frameworks such as EAGLE and HASS precompute and store $\mathbf{F} _ t$ for all $\mathbf{x} _ t$ before training begins, to avoid the computational cost of regenerating training data on the fly. As a result, naively replacing $\mathbf{x} _ t$ with $\bar{\mathbf{x}} _ t$ produces mismatched input-feature pairs:

• The precomputed target feature $\mathbf{F} _ {t+1}$ was computed on the ground-truth tokens $\mathbf{x} _ {1:t}$ using the target model;
• The corresponding draft feature $\bar{\mathbf{F}} _ {t+1}$, however, would be computed using tokens $\mathbf{x} _ {1:t-1}$ concatenated with the self-generated token $\bar{\mathbf{x}} _ t$.

Consequently, the feature-level loss $\ell(\mathbf{F} _ {t+1}, \bar{\mathbf{F}} _ {t+1})$ becomes ill-defined since the two features are conditioned on different token inputs. As confirmed in Appendix B of HASS, this naive substitution undermines training and ultimately degrades final model performance.

Our solution: token-alignable training and draft model. GRIFFIN not only preserves the strengths of HASS in addressing feature misalignment, but also introduces two complementary mechanisms specifically tailored to token misalignment:

• Token-alignable training. A loss-masking scheme discards highly misaligned positions, preventing their erroneous gradients from corrupting learning.
• Token-alignable draft model. Architectural modifications explicitly incorporate input tokens twice (initially with features and later to refine them) to reconcile discrepancies in the generated draft features.

By jointly employing token-alignable training and a token-alignable draft model, GRIFFIN robustly addresses token misalignment while fully retaining the feature-level alignment advances established by HASS. This comprehensive approach yields both improved training stability and superior model performance, as substantiated in our empirical results in Table 1.

3. Definition of TEH Module

Thanks. In the submission, we did not carefully introduced TEH module but only cited [31] due to two reasons: (1) TEH is not our technical contribution, as it is directly adapted from the design in [31]; and (2) given limited space, we prioritized detailing our novel components and omitted the full description of TEH. We will include the following more detailed explanation of TEH in the appendix of revision for completeness.

To clarify, TEH corresponds to the decoupled projection head proposed in [31] (see Figure 5 in [31]). It is a compact and single-layer MLP that replaces the conventional output projection in the transformer decoder layer of the draft model. Formally, in our context, let $\mathbf{g} \in \mathbb{R}^c$ denote the output of an autoregressive layer in our draft model (see Figure 3(a)), where $c$ is the intermediate size of autoregressive layer. TEH maps $\mathbf{g}$ into two features via:

$$[\bar{\mathbf{F}}^P, \bar{\mathbf{F}}^R] = \mathbf{g} \mathbf{W} _ {\text{TEH}} + \mathbf{b} _ {\text{TEH}}$$

where $\mathbf{W} _ {\text{TEH}} \in \mathbb{R}^{c \times 2d}$ and $\mathbf{b} _ {\text{TEH}} \in \mathbb{R}^{2d}$ are the projection weights and bias, and $d$ is the model’s hidden size. Here $\bar{\mathbf{F}}^P \in \mathbb{R}^d$ is fed to the LM head for token prediction; while $\bar{\mathbf{F}}^R \in \mathbb{R}^d$ serves as the recurrent feature, passed to the next layer for feature propagation.

The core motivation of TEH is to explicitly decouple the two objectives of decoding: accurate token prediction and stable hidden state propagation. A shared representation often leads to a conflict between these objectives, especially in draft models with minimal recurrence. By introducing separate projections, TEH improves the robustness of recurrence without sacrificing prediction quality [31].

4. TGF

Relation between TGF and token alignment. TGF module is explicitly designed to enhance the draft model’s token prediction accuracy, which in turn reduces the number of misaligned tokens during token-alignable training.

TGF module is motivated by the limitations of the standard concat-then-MLP strategy, as adopted in EAGLE, which does not fully capture the complementary information between token embeddings and draft model features. In practice, features generated by the draft model often remain misaligned with the target model’s representations, a discrepancy that cannot be effectively eliminated with feature-level loss minimization alone. TGF addresses this challenge by explicitly leveraging token embeddings to guide the fusion process, aligning feature distributions more closely to those of the target model. As confirmed by ablation results (Table 5), this targeted architectural enhancement significantly reduces feature inconsistency, demonstrating that the modest complexity introduced by TGF provides strong empirical gains.

Empirically, the reduction in misalignment rate increases the proportion of draft tokens that pass the alignment mask, effectively expanding the usable training signal and improving overall learning efficiency. Therefore, TGF plays a direct role in boosting alignment success through improved token quality.

Which components foster token alignment in TGF? In Table 6 of Appendix A.2, we have conducted a comprehensive ablation study, where we selectively removed each component of TGF while maintaining comparable parameter counts and FLOP budgets. The key observations are:

• Removing both the Up-Projector and SiLU lowers acceptance length by $0.23$ ($T{=}0$) and $0.27$ ($T{=}1$), and reduces speed-up by $0.17$ and $0.21$, respectively.

• Omitting the token embedding $\mathbf x$ from the second fusion step degrades acceptance length by $0.10$ ($T{=}0$) and $0.13$ ($T{=}1$), with speed-up dropping by $0.07$ and $0.08$.

• Excluding the recurrent feature $\mathbf h$ yields the most pronounced decline: acceptance length falls by $2.63$ ($T{=}0$) and $2.49$ ($T{=}1$), while speed-up plummets by $1.50$ and $1.39$.

These controlled ablations demonstrate that each component of TGF plays a non-trivial role in improving the quality of draft tokens that ultimately pass verification. In particular, the recurrent integration of the fused feature $\mathbf{h}$ proves essential for capturing high-fidelity, alignable representations. Importantly, we did not identify any alternative architecture with similar parameter/FLOP budgets that matched TGF’s performance, underscoring the necessity and effectiveness of its specific design choices in addressing the token alignment challenge.