Skip to main content Inicio Creadores adu2021 skillxiv adaptive-speculative-decoding
adaptive-speculative-decoding Accelerate LLM inference 1.5-2x using a universal draft model that adapts to user data in real-time, handling different target models and tokenizers via online n-gram cache and hybrid distillation.
Ir a la instalación Skills Marketplace Descubre y explora habilidades de IA creadas por la comunidad.
Instalar con Codex o Claude Copia este prompt, pégalo en Codex, Claude u otro asistente, y deja que revise la página de la skill y la instale por ti.
Mostrar detalles del prompt Un comando directo omite el prompt de revisión. Revisa el origen antes de ejecutarlo.
npx skills add https://github.com/ADu2021/skillXiv --skill adaptive-speculative-decodingEl comando permanece en una sola línea. Desplázate horizontalmente para revisarlo antes de copiarlo.
¿Prefieres una copia local? Descarga los archivos que SkillsMP tiene disponibles ahora.
Descargar Zip Descargando... Más de este repositorio meaningful-kebab-case-name Convert arXiv papers into ready-to-use agent skills using category-aware extraction. First classifies the paper into one or more of 11 research categories, then applies a specialized extraction pipeline for each category — because different types of papers produce different types of usable knowledge. A single paper can yield multiple skills if it spans categories. Use this skill whenever the user wants to turn a paper into a skill, extract practical techniques from research, build a skill library from papers, convert arXiv papers into reusable agent instructions, or batch-process multiple papers into skills. Also trigger when someone asks about extracting actionable knowledge from papers, making research practical for LLM agents, or systematically converting academic contributions into structured agent capabilities.
action-quantization-behavior-cloning Establish regret bounds for behavior cloning with discretized actions combining statistical error and quantization error terms. Prove smoothness requirements for safe quantizer design, show that learning-based quantizers fail these requirements, and propose model-based augmentation to reduce error dependence from H² to H.
adaptive-lora-personalized-ranks Dynamically allocate LoRA ranks per-layer during fine-tuning instead of using fixed uniform ranks. Learn optimal rank for each layer and subject via variational framework with discretized exponential distribution, reducing memory footprint while maintaining fidelity and text-alignment.
Ocupaciones relacionadas SOC
Basado en la clasificación ocupacional SOC
name adaptive-speculative-decoding title OmniDraft: A Cross-vocabulary, Online Adaptive Drafter for On-device Speculative Decoding version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2507.02659 keywords ["Speculative Decoding","On-device Inference","Online Adaptation","Vocabulary Mismatch","Knowledge Distillation"] description Accelerate LLM inference 1.5-2x using a universal draft model that adapts to user data in real-time, handling different target models and tokenizers via online n-gram cache and hybrid distillation.
OmniDraft: Universal and Adaptive Speculative Decoding
Speculative decoding accelerates LLM inference by using a small draft model to generate multiple tokens, which are verified by the target model in parallel. However, deploying this approach on-device faces challenges: (1) draft models trained for one target model may not work with another, and (2) user-specific data causes distribution shift over time. OmniDraft addresses both through a single universal draft model (Llama-68M) that works with any target model via cross-vocabulary token mapping, and online knowledge distillation that continuously adapts the drafter to user data without explicit retraining.
The key innovations are an online n-gram cache that translates between different tokenizer vocabularies, and a hybrid loss combining token-level and distribution-level distillation that updates the drafter as it processes user data. This enables 1.5-2x speedup across reasoning, coding, and text generation tasks on consumer hardware.
Core Concept
OmniDraft operates on two principles: (1) vocabulary-aware adaptation through n-gram caches that learn token mappings between drafters and targets, and (2) online learning that improves the drafter continuously as it encounters user data. Rather than training separate drafters for each target model or retraining when deployment contexts change, a single universal drafter becomes increasingly specialized to its deployment context through lightweight online updates.
The n-gram cache maintains learned mappings showing which draft tokens typically correspond to which target tokens, enabling effective translation even with completely different vocabularies. Online distillation ensures the drafter learns patterns specific to the user's domain without catastrophic forgetting of general knowledge.
Architecture Overview
The system comprises several interconnected components:
Universal Draft Model : Lightweight Llama-68M serving as single drafter for multiple targetsN-gram Cache : Learned bijection between draft and target token sequences, enabling vocabulary translationHybrid Distillation Head : Combines KL divergence (for directly mapped tokens) and NLL (for n-gram mapped tokens)Acceptance Predictor : Lightweight head predicting token acceptance probability, enabling adaptive proposal lengthsOnline Learning Loop : Continuously updates draft model and cache during inference with user data
Implementation
Start with the vocabulary mapping layer:
import torch
import torch.nn as nn
from typing import
Dict
List
Tuple
from
import
class
NGramCache
"""
Learn mappings between draft and target model vocabularies.
Maintains an n-gram cache storing which target token sequences
typically follow draft token sequences, enabling cross-vocabulary adaptation.
"""
def
__init__
self, n: int = 4 , cache_size: int = 50000
self
self
self
lambda
int
self
0
def
record_mapping
self, draft_sequence: List [int ], target_token: int
"""
Record mapping from draft token sequence to target token.
Updates n-gram statistics: "when we see these draft tokens,
the target often outputs this token next."
"""
if
len
self
tuple
self
1
else
tuple
self
1
self
1
if
len
self
self
self
def
get_mapping
self, draft_sequence: List [int ],
top_k: int = 5
List
Tuple
int
float
"""
Get predicted target tokens given draft sequence.
Returns list of (token_id, probability) pairs ranked by likelihood.
"""
if
len
self
tuple
self
1
else
tuple
if
not
in
self
return
self
sum
for
in
sorted
lambda
1
True
return
def
_evict_cold_entries
self
"""Remove least-accessed n-gram patterns to stay within cache budget."""
for
in
self
sum
len
self
10
sorted
lambda
1
for
in
del
self
class
VocabularyAdapter
"""
Adapt draft model vocabulary to target model vocabulary.
Uses n-gram cache to translate draft tokens to target tokens,
handling vocabulary mismatch transparently.
"""
def
__init__
self, draft_vocab_size: int , target_vocab_size: int ,
hidden_dim: int = 768
super
self
self
self
self
4
def
translate_tokens
self, draft_logits: torch.Tensor,
draft_history: List [int ],
target_tokenizer: object
"""
Convert draft model logits to target vocabulary predictions.
Uses learned projection + n-gram cache for accurate translation.
"""
self
float
'-inf'
for
in
range
self
10
for
in
if
self
for
in
range
1e-8
max
return
Implement the online distillation mechanism:
class HybridDistillationLoss (nn.Module):
"""
Combine token-level and distribution-level distillation.
Token-level loss guides direct-mapped tokens; distribution loss
guides sequences found via n-gram cache. Enables effective online learning.
"""
def __init__ (self ):
super ().__init__()
self .kl_loss = nn.KLDivLoss(reduction='batchmean' )
self .nll_loss = nn.NLLLoss()
def forward (self, draft_logits: torch.Tensor,
target_logits: torch.Tensor,
direct_mapped_mask: torch.Tensor,
ngram_mapped_mask: torch.Tensor,
temperature: float = 3.0 ) -> torch.Tensor:
"""
Compute hybrid loss for online distillation.
Args:
draft_logits: logits from draft model
target_logits: logits from target model
direct_mapped_mask: which tokens have direct vocabulary mappings
ngram_mapped_mask: which tokens have n-gram cache mappings
temperature: for soft targets
Returns:
loss: weighted combination of KL and NLL losses
"""
draft_log_probs = torch.log_softmax(draft_logits / temperature, dim=-1 )
target_probs = torch.softmax(target_logits / temperature, dim=-1 )
kl_loss = self .kl_loss(draft_log_probs * direct_mapped_mask.unsqueeze(-1 ),
target_probs * direct_mapped_mask.unsqueeze(-1 ))
target_tokens = torch.argmax(target_logits, dim=-1 )
nll_loss = self .nll_loss(
draft_log_probs[ngram_mapped_mask],
target_tokens[ngram_mapped_mask]
)
total_loss = 0.7 * kl_loss + 0.3 * nll_loss
return total_loss
Implement the adaptive draft model with online learning:
class AdaptiveDraftModel (nn.Module):
"""
Draft model that learns online from user data.
Updates parameters during inference based on target model outputs,
improving specialization to user distribution without explicit retraining.
"""
def __init__ (self, model_name: str = "meta-llama/Llama-2-7b" ,
draft_model_name: str = "gpt2-medium" ,
learning_rate: float = 2e-5 ):
super ().__init__()
from transformers import AutoTokenizer, AutoModelForCausalLM
self .tokenizer = AutoTokenizer.from_pretrained(draft_model_name)
self .draft_model = AutoModelForCausalLM.from_pretrained(draft_model_name)
self .lora_optimizer = torch.optim.Adam(
self .draft_model.parameters(), lr=learning_rate
)
self .vocab_adapter = VocabularyAdapter(
self .tokenizer.vocab_size,
32000 ,
)
self .distillation_loss = HybridDistillationLoss()
self .acceptance_head = nn.Sequential(
nn.Linear(768 , 256 ),
nn.ReLU(),
nn.Linear(256 , 1 ),
nn.Sigmoid()
)
def generate_candidates (self, input_ids: torch.Tensor,
num_candidates: int = 5 ) -> torch.Tensor:
"""
Generate candidate tokens using draft model.
Returns multiple tokens for parallel verification by target model.
"""
outputs = self .draft_model(input_ids)
logits = outputs.logits[:, -1 , :]
top_k = min (num_candidates, logits.shape[-1 ])
_, top_indices = torch.topk(logits, k=top_k, dim=-1 )
return top_indices
def online_update (self, input_ids: torch.Tensor,
target_logits: torch.Tensor,
target_tokens: torch.Tensor ):
"""
Update draft model based on target model outputs.
Lightweight online learning: single gradient step on recent data.
"""
draft_outputs = self .draft_model(input_ids)
draft_logits = draft_outputs.logits
loss = self .distillation_loss(
draft_logits,
target_logits,
direct_mapped_mask=torch.ones_like(draft_logits[:, :, 0 ], dtype=torch.bool ),
ngram_mapped_mask=torch.zeros_like(draft_logits[:, :, 0 ], dtype=torch.bool )
)
self .lora_optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self .draft_model.parameters(), 1.0 )
self .lora_optimizer.step()
return loss.item()
def predict_acceptance_rate (self, draft_features: torch.Tensor ) -> torch.Tensor:
"""
Predict probability that draft tokens will be accepted.
Enables adaptive proposal length: lower acceptance rate = shorter proposals.
"""
return self .acceptance_head(draft_features)
Implement the full speculative decoding pipeline:
class OmniDraftDecoder :
"""
Full on-device speculative decoding with online adaptation.
Generates tokens using draft model, verifies with target model,
and adapts draft model to user data in real-time.
"""
def __init__ (self, target_model_name: str = "meta-llama/Llama-2-7b-chat" ,
draft_model_name: str = "gpt2-medium" ):
from transformers import AutoTokenizer, AutoModelForCausalLM
self .target_tokenizer = AutoTokenizer.from_pretrained(target_model_name)
self .target_model = AutoModelForCausalLM.from_pretrained(target_model_name)
self .draft = AdaptiveDraftModel(target_model_name, draft_model_name)
self .total_tokens = 0
self .accepted_tokens = 0
def decode (self, prompt: str , max_new_tokens: int = 128 ,
num_draft_candidates: int = 5 ) -> str :
"""
Decode with speculative generation and online adaptation.
Uses draft model to propose tokens, target model to verify,
and online learning to improve draft model during generation.
"""
input_ids = self .target_tokenizer.encode(prompt, return_tensors='pt' )
generated = input_ids.clone()
for _ in range (max_new_tokens):
candidates = self .draft.generate_candidates(
input_ids, num_draft_candidates
)
target_outputs = self .target_model(input_ids)
target_logits = target_outputs.logits[:, -1 , :]
target_token = torch.argmax(target_logits, dim=-1 )
accepted = (candidates[0 , 0 ] == target_token).item()
self .total_tokens += 1
if accepted:
self .accepted_tokens += 1
if _ % 10 == 0 :
self .draft.online_update(input_ids, target_logits, target_token)
next_token = target_token.unsqueeze(0 )
generated = torch.cat([generated, next_token], dim=1 )
input_ids = generated[:, -128 :]
result = self .target_tokenizer.decode(generated[0 ], skip_special_tokens=True )
return result
def get_speedup_stats (self ) -> Dict :
"""Return statistics about speedup achieved."""
acceptance_rate = self .accepted_tokens / max (self .total_tokens, 1 )
return {
'total_tokens' : self .total_tokens,
'accepted_tokens' : self .accepted_tokens,
'acceptance_rate' : acceptance_rate,
'estimated_speedup' : 1 + acceptance_rate
}
Practical Guidance Parameter Default Range Notes Draft model size 68M 30M-250M Larger = better proposals but slower verification N-gram cache size 50K 10K-200K Larger = more accurate mappings; more memory N-gram length 4 2-6 Longer = more specific but sparse Online learning rate 2e-5 1e-6 to 1e-4 Very conservative; single-step updates Update frequency 10 steps 1-50 How often to update on new data Distillation temperature 3.0 1.0-10.0 Higher = softer targets; better transfer Num candidates 5 2-10 More candidates = more parallelism; more rejects
You need to deploy LLMs on consumer hardware with limited resources
You want 1.5-2x speedup without changing the target model
You plan to use same draft model across multiple target models
Your deployment has time to learn from user data (chat/interactive)
You need to handle different data distributions per user
You have unlimited compute (just use target model directly)
You need deterministic generation (sampling adds variance)
Your application requires sub-10ms latency (overhead may hurt)
You have very short input contexts (overhead dominates)
You need guaranteed token sequences (rejection sampling changes output)
Vocabulary mismatch underestimation : If draft and target vocabularies differ significantly, n-gram cache becomes sparse. Test on your specific models.Online learning instability : Large learning rates cause draft model to overfit to recent data. Use very conservative rates and sample diverse batches.Cache thrashing : If n-gram patterns change rapidly, cache eviction hurts. Monitor cache hit rates and increase size if needed.Acceptance rate drift : If draft model's adaptation lags behind target distribution, acceptance rate drops over time. Increase update frequency or learning rate.Memory overhead : N-gram cache + online optimizer parameters add overhead. Profile actual memory usage before deployment.Token mismatch issues : If draft tokenizer produces different tokens than target, candidate verification fails silently. Validate tokenization equivalence.
Reference Authors (2025). OmniDraft: A Cross-vocabulary, Online Adaptive Drafter for On-device Speculative Decoding. arXiv preprint arXiv:2507.02659. https://arxiv.org/abs/2507.02659
