Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs

HedgeSpec is an adaptive drafter selection framework for speculative decoding in LLMs. Expert drafters could perform strong in-domain, they often degrade sharply outside their expertise, making effective drafter selection challenging (left). HedgeSpec leverages the structure of spec decoding and introduces an evaluation phase to extract feedback for all candidates without extra target calls, enabling lightweight full-information learning and achieving provably no-regret performance (top right). In practice, this leads to the high end-to-end speedups (bottom right). More details can be found in our ICLR'2026 paper.

Enviroment Setup

The repo is based on EAGLE repository. The base environment to run ProxSparse is stored in requirement.txt. You can run the below command to install them.

conda create -n hedgespec python==3.10
conda activate hedgespec
pip install -r req.txt
pip install fschat==0.2.31

Set up the root directory environment variable:

git clone https://github.com/aladinggit/hedgespec.git
cd hedgespec
export EAGLE_PATH=$(pwd)/eagle

Download drafters and crafter curation

We built 21 drafters (for 3 base models) upon the official EAGLE-3 checkpoints (Llama-3.1-8B-IT: yuhuili/EAGLE3-LLaMA3.1-Instruct-8B; Qwen-3-8B: Tengyunw/qwen3_8b_eagle3; Qwen-3-32B: AngelSlim/Qwen3-32B_eagle3) across 7 different domains (datasets: python, math, biology, chemistry, MedQA, CNN_DM and SQL)

Target rollout training datasets

The training of the drafter requires the rollout data from the target model (not directly using the datasets). This can be handled with vllm/sglang framework. We provide below example workflow for handling the rollout procedure.

(On another terminal)
cd hedgespec/drafter_datagen_rollout
conda create -n rollout python==3.12
conda activate rollout
pip install vllm
# launch the server, modify the target model name you want to rollout
bash vllm_launch.sh

# convert gsm8k data into ShareGPT format
python convert_sharegpt_gsm8k_example.py
# put the output ShareGPT file name into the bash file first
bash gen_gsm8k.sh

Drafte training

We use SpecForge for crafting those expert drafters. The drafter curated are trained with rollout datasets from the 7 domains for different base models.

Downloading drafter for different domains

The 21 drafters used in this paper can be accessed in huggingface repo (aladinggit/hedgespec_eagle_drafters). To download those drafters, run:

python script/download_llama3_1_8B_drafter.py
python script/download_qwen3_8B_drafter.py
python script/download_qwen3_32B_drafter.py

So that those domain expert EAGLE3 drafters will be downloaded into current directory (./[lm318, qw38, qw332]/model/[category]).

Evaluation datasets

We provide a list of evaluation datasets in ./data directory. Those datasets are in the ShareGPT format already. To convert more data into this format, please refer to the above converter script under rollout session.

Running HedgeSpec

Below illustrates the step to run HedgeSpec as well as other variants/baselines (i.e. Bandit based.). This repo is mostly inherited from official EAGLE github. The most significant modification we made are three files under ./eagle/model directory: ea_model_back.py, cnets_back.py and util_back.py, where we load multiple drafters and conduct additional evaluation step for online drafter selections. More patching explanation can be found in ./eagle/model/doc.txt.

To run hedgespec, set the base enviroment variable first:

cd hedgespec
export EAGLE_PATH=$(pwd)/eagle

Configuring the drafters to use

To customize the number/categories of the drafter you want to use, write a json drafter specification under script/drafters and pass it as a variable in the scripts. There are currently three json files there in the directory (each with 7 drafters in pool), put the drafter path/name (i.e. from huggingface) and add/delete in the json files.

Running HedgeSpec

We provided two example script for running HedgeSpec, you can modify the corresponding arguments for running different datas/drafters/configurations combinations. The entry point for running the experiments is in script/test_final_multi_speed_all.py. Here if you prefer single GPU setting, please set HEDGE_SINGLE_GPU environmental variable to 1 as well as setting your CUDA_VISIBLE_DEVICES variable. This repo currently supports running of Llama-3.1-8B, Qwen-3-8B as well as Qwen-3-32B, and the repo has been tested on Nvidia H100 GPU cluster. Arguments explanation can be found in test_final_multi_speed_all.py and we listed some below:

base-model: the target model to run

data: The datasets for evaluation

drafter_pool: the path to the json configuration specifying the drafters involved during inference in pool

hedge: configurations to invoking different online learning mechanisms

• hedge = 13: running HedgeSpec with expected length as loss
• hedge = -1: bandit based method with UCBSpec
• hedge = -2: bandit based method with EXP3Spec
• HedgeSpec variants: hedge = 14: running HedgeSpec with adanormalHedge (NormalHedge by default); hedge = 5: running HedgeSpec with token acceptance rate as loss

Below are some example scripts for running HedgeSpec on python/openmath datasets:

bash script/run_hedgespec_llama.sh ## running HedgeSpec on Llama
bash script/run_hedgespec_qwen.sh ## running HedgeSpec on Qwen

To run bandit based method:

bash script/run_bandit.sh 

To play with adanormalHedge or token acceptance rate based loss:

bash script/run_variants.sh 

running single drafter eagle:

Below are script to test the single drafter EAGLE models. The entry point of the program is in script/test_final_all_standard_speed.py.

base-model: the target model to run

data: The datasets for evaluation

eagle: The EAGLE drafter to run, i.e. yuhuili/EAGLE3-LLaMA3.1-Instruct-8B or $EAGLE_PATH/../lm318/lm318/model/math.

To test the single drafter framework:

bash run_vanilla_eagle.sh

running naive target generation:

Below are script to test the naive target generation. The entry point of the program is in script/test_final_naive_speed.py.

bash run_naive_gen.sh

License

This project is licensed under the Apache 2.0 license.

Citation

If you might find our work useful, please cite:

@article{liu2025not, title={Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs}, author={Liu, Hongyi and Huang, Jiaji and Jia, Zhen and Park, Youngsuk and Wang, Yu-Xiang}, journal={arXiv preprint arXiv:2510.20064}, year={2025} }