Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Flatter Tokens are More Valuable for Speculative Draft Model Training

Jiaming Fan, Daming Cao, Xiangzhong Luo, Jiale Fu, Chonghan Liu, Xu Yang

TL;DR

The paper tackles the data inefficiency of training draft models for Speculative Decoding by showing that tokens whose target distributions are flatter (i.e., more uncertain) provide larger gains in acceptance rate. It introduces flatness, a cosine-to-uniform measure, and develops Sample-level-flatness-based Dataset Distillation (SFDD) to filter data offline, preserving SD speedups while reducing training cost. Theoretical analysis under a budgeted KD framework links token-level flatness to improvements in the L1 distance that governs acceptance, and empirical validation demonstrates that SFDD consistently outperforms entropy-based and other heuristics across multiple tasks and datasets. On the EAGLE-2/LLaMA3-8B-Instruct setup with ShareGPT data, SFDD achieves over 2× training speedup at 50% data retention and keeps inference speedups within 4% of the full-data baseline, highlighting a practical, data-centric path to more efficient SD training.

Abstract

Speculative Decoding (SD) is a key technique for accelerating Large Language Model (LLM) inference, but it typically requires training a draft model on a large dataset. We approach this problem from a data-centric perspective, finding that not all training samples contribute equally to the SD acceptance rate. Specifically, our theoretical analysis and empirical validation reveals that tokens inducing flatter predictive distributions from the target model are more valuable than those yielding sharply peaked distributions. Based on this insight, we propose flatness, a new metric to quantify this property, and develop the Sample-level-flatness-based Dataset Distillation (SFDD) approach, which filters the training data to retain only the most valuable samples. Experiments on the EAGLE framework demonstrate that SFDD can achieve over 2$\times$ training speedup using only 50% of the data, while keeping the final model's inference speedup within 4% of the full-dataset baseline. This work introduces an effective, data-centric approach that substantially improves the training efficiency for Speculative Decoding. Our code is available at https://anonymous.4open.science/r/Flatness.

Flatter Tokens are More Valuable for Speculative Draft Model Training

TL;DR

The paper tackles the data inefficiency of training draft models for Speculative Decoding by showing that tokens whose target distributions are flatter (i.e., more uncertain) provide larger gains in acceptance rate. It introduces flatness, a cosine-to-uniform measure, and develops Sample-level-flatness-based Dataset Distillation (SFDD) to filter data offline, preserving SD speedups while reducing training cost. Theoretical analysis under a budgeted KD framework links token-level flatness to improvements in the L1 distance that governs acceptance, and empirical validation demonstrates that SFDD consistently outperforms entropy-based and other heuristics across multiple tasks and datasets. On the EAGLE-2/LLaMA3-8B-Instruct setup with ShareGPT data, SFDD achieves over 2× training speedup at 50% data retention and keeps inference speedups within 4% of the full-data baseline, highlighting a practical, data-centric path to more efficient SD training.

Abstract

Speculative Decoding (SD) is a key technique for accelerating Large Language Model (LLM) inference, but it typically requires training a draft model on a large dataset. We approach this problem from a data-centric perspective, finding that not all training samples contribute equally to the SD acceptance rate. Specifically, our theoretical analysis and empirical validation reveals that tokens inducing flatter predictive distributions from the target model are more valuable than those yielding sharply peaked distributions. Based on this insight, we propose flatness, a new metric to quantify this property, and develop the Sample-level-flatness-based Dataset Distillation (SFDD) approach, which filters the training data to retain only the most valuable samples. Experiments on the EAGLE framework demonstrate that SFDD can achieve over 2 training speedup using only 50% of the data, while keeping the final model's inference speedup within 4% of the full-dataset baseline. This work introduces an effective, data-centric approach that substantially improves the training efficiency for Speculative Decoding. Our code is available at https://anonymous.4open.science/r/Flatness.
Paper Structure (48 sections, 2 theorems, 43 equations, 7 figures, 11 tables)

This paper contains 48 sections, 2 theorems, 43 equations, 7 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Preliminaries
  4. Speculative decoding
  5. Theoretical analysis of KD in SD
  6. The proposed approach
  7. Empirical validation of the theoretical analysis
  8. From token-level insight to sample-level data selection
  9. Sample-level-flatness-based dataset distillation
  10. Experiments
  11. Experimental setup
  12. Models and baselines.
  13. Dataset and tasks.
  14. Evaluation metrics.
  15. Main results
  16. ...and 33 more sections

Key Result

Theorem A.2

There exists a one-parameter family of Gaussian candidates and a unique parameter $\tau^* \in [0,1]$ such that the optimal solution $r^*$ to the problem above is given by For any $\tau\in[0,1]$, we have In particular, $\mu_r^* = \mu_r(\tau^*)$ and $\sigma_r^{2*} = \sigma_r^2(\tau^*)$. Let $\Delta_\mu := \mu_p - \mu_q$ and define Then: Moreover, if $p\neq q$, the function $g(\tau)$ is strictly d

Figures (7)

  • Figure 1: Simulation of importance metrics. We use Gaussian distributions in our simulation because they are analytically tractable and effectively capture key distributional properties like variance. For visualization clarity, all plotted quantities are min–max normalized to $[0,1]$. (a) Across small, medium, and large separations between the means of $p$ and $q$, $\Delta_{L_1}$ consistently increases with the target variance ($\sigma_p^2$). (b) Cosine similarity to uniform increases monotonically with the standard deviation ($\sigma$), validating its use as a practical proxy that directly tracks changes in variance.
  • Figure 2: Target-sorted flatness view. Tokens sorted by the target’s statistic (low$\rightarrow$high flatness). (a) flatness values; (b) one-epoch change in flatness; (c) the one-epoch reduction in the $L_1$ discrepancy, $\Delta_{L_1}$. Curve coloring distinguishes token groups by target flatness: the blue segment represents tokens with low flatness, while the red segment represents those with high flatness. Panels (b,c) additionally annotate Minimal change (left; indicating that the vast majority of points in this segment exhibit small changes) and Greater change (right; indicating that the vast majority of points in this segment exhibit larger changes). (d) Entropy-vs-flatness filtering gap: for each metric, we rank tokens by that metric, take the bottom 35% tokens, and compute their one-epoch average $|\Delta L_1|$; bars plot the difference between the entropy-based and flatness-based bottom-35% averages under different numbers of selected examples. Flatness is as defined in Equation \ref{['eq:token-flatness-def']}.
  • Figure 3: The SFDD workflow: This approach calculates $\textit{flatness}_{\mathrm{sample}}$ by averaging token flatness within each sample, and then uses quantile-derived threshold to select the top-$k$% and filter the dataset for training. The figure illustrates this with a concrete example using 70% retain ratio.
  • Figure 4: Training time as a function of the data retain ratio (including data-selection time). Each point is annotated with the absolute wall-clock time and the corresponding training speedup.
  • Figure 5: Target-sorted entropy view. Tokens sorted by the target's entropy (low$\rightarrow$high entropy). (a) target entropy values; (b) one-epoch change in draft entropy; (c) the one-epoch reduction in the $L_1$ discrepancy, $\Delta_{L_1}$. Similar to the flatness view, we observe that tokens with higher entropy (indicating flatter distributions) contribute more to the training dynamics ($\Delta_{L_1}$) and exhibit larger changes in the draft model.
  • ...and 2 more figures

Theorems & Definitions (6)

  • Definition A.1: Problem Formulation
  • Theorem A.2: Optimal Gaussian Parameters
  • proof
  • Definition B.1: Discretized Gaussian
  • Theorem B.2: Asymptotic Cosine Similarity
  • proof