メニュー
Accelerate speculative decoding by pruning drafter vocabulary to high-frequency tokens. Achieves 16% speedup in memory-bound settings by eliminating unused vocabulary entries without retraining.
デフォルトでは、最初にソースを確認する Prompt が選択されています。直接コマンドに切り替えるか、ローカルコピーをダウンロードすることもできます。
インストールを決める前に、SKILL.md と SkillsMP に表示されている付属ファイルをお読みください。
SOC 職業分類に基づく
Codex または Claude でインストール この Prompt をコピーして Codex、Claude、または他のアシスタントに貼り付けると、Skill ページを確認してインストールできます。
直接コマンドでは確認用 Prompt が省略されます。実行前にソースを確認してください。
npx skills add https://github.com/ADu2021/skillXiv --skill vocabtrim-speculative-decodingコマンドは1行のまま表示されます。コピー前に横へスクロールして全体を確認してください。
ローカルで確認しますか？SkillsMP が現在取得できるファイルをダウンロードできます。
SKILL.md を表示中
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.
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.
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.
| name | vocabtrim-speculative-decoding |
| title | VocabTrim: Vocabulary Pruning for Efficient Speculative Decoding in LLMs |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2506.22694 |
| keywords | ["Speculative Decoding","Vocabulary Pruning","Inference Optimization","Memory Efficiency","Token Prediction"] |
| description | Accelerate speculative decoding by pruning drafter vocabulary to high-frequency tokens. Achieves 16% speedup in memory-bound settings by eliminating unused vocabulary entries without retraining. |
Speculative decoding uses a small drafter model to propose multiple tokens per inference step, which the target verifier model accepts or rejects. This approach can 2-3× speedup inference when the drafter is fast. However, the drafter's language modeling head (the final layer outputting logits over all vocabulary tokens) becomes a memory bottleneck. For Llama-3 with 128K vocabulary tokens, computing logits over all 128K tokens at every step wastes memory and computation even though the drafter only samples from a tiny subset of frequently-occurring tokens.
VocabTrim solves this by reconstructing the drafter's vocabulary to contain only the high-frequency tokens it actually samples during inference. The insight is that drafters are biased toward "easy-to-predict" tokens (common words, punctuation) and rarely sample rare tokens. By trimming the vocabulary to the most frequent 25-50K tokens, you eliminate 60-75% of the LM head computation with negligible impact on acceptance rates.
VocabTrim works on a simple principle: replace the drafter's full vocabulary with a smaller one containing only tokens it frequently generates. The key insights are:
The result is a smaller drafter that produces fewer logits per step (reduced memory transfers) and faster matrix multiplication (fewer output dimensions), with only minimal increase in rejection rate (1-3%).
VocabTrim modifies only the drafter's language modeling head:
Step 1: Measure token frequencies on calibration data
Sample from the drafter on representative data and count which tokens appear most frequently in the output distribution.
def measure_drafter_token_frequencies(drafter, tokenizer, calibration_data,
num_samples=100000, sample_length=256):
"""
Run the drafter on calibration data and measure which tokens
are sampled most frequently. This identifies which tokens matter for inference speed.
"""
token_counts = {}
total_tokens_sampled = 0
for batch_idx, batch in enumerate(calibration_data):
input_ids = tokenizer.encode(batch, return_tensors='pt').to(device)
# Generate from drafter, collecting all sampled tokens
with torch.no_grad():
for step in range(sample_length):
outputs = drafter(input_ids)
logits = outputs.logits[:, -1, :] # Last position logits
# Sample from the distribution (how the drafter would generate in practice)
probs = torch.softmax(logits, dim=-1)
sampled_tokens = torch.multinomial(probs, num_samples=1)
# Count token occurrences
for token_id in sampled_tokens.flatten().tolist():
token_counts[token_id] = token_counts.get(token_id, 0) + 1
total_tokens_sampled += 1
input_ids = torch.cat([input_ids, sampled_tokens], dim=1)
if (batch_idx + 1) * sample_length >= num_samples:
break
# Sort by frequency
sorted_tokens = sorted(token_counts.items(),
key=lambda x: x[1], reverse=True)
return dict(sorted_tokens)
Step 2: Select vocabulary and create remapping
Choose the top-K tokens and create an index mapping from small vocabulary to full vocabulary.
def create_vocab_mapping(token_frequencies, target_vocab_size=32000):
"""
Select the most frequent tokens and create a mapping.
This tells us which full vocabulary indices to keep.
"""
# Select top-K tokens by frequency
top_tokens = sorted(token_frequencies.items(),
key=lambda x: x[1], reverse=True)[:target_vocab_size]
# Create two mappings:
# 1. full_to_trimmed: maps full vocabulary indices to trimmed indices
# 2. trimmed_to_full: maps trimmed indices back to full vocabulary
full_to_trimmed = {}
trimmed_to_full = {}
for trimmed_idx, (full_token_id, count) in enumerate(top_tokens):
full_to_trimmed[full_token_id] = trimmed_idx
trimmed_to_full[trimmed_idx] = full_token_id
# Calculate coverage: what percentage of real samples are covered
total_count = sum(token_frequencies.values())
selected_count = sum(count for _, count in top_tokens)
coverage = selected_count / total_count
return {
'full_to_trimmed': full_to_trimmed,
'trimmed_to_full': trimmed_to_full,
'vocab_size': target_vocab_size,
'coverage': coverage
}
Step 3: Reconstruct the drafter's LM head
Create a new, smaller language modeling head by extracting only the necessary weight rows.
def reconstruct_lm_head(drafter, vocab_mapping):
"""
Extract the drafter's LM head weights for selected tokens only.
This creates a new head with fewer output dimensions.
"""
original_head = drafter.lm_head
original_weight = original_head.weight # Shape: [vocab_size, hidden_dim]
# Select only rows corresponding to kept tokens
selected_indices = torch.tensor(
[vocab_mapping['trimmed_to_full'][i]
for i in range(vocab_mapping['vocab_size'])],
device=original_weight.device
)
# Extract selected rows
trimmed_weight = original_weight[selected_indices, :]
# Create new head with smaller output vocabulary
trimmed_head = torch.nn.Linear(
original_head.in_features,
vocab_mapping['vocab_size'],
bias=(original_head.bias is not None)
)
# Copy weights
with torch.no_grad():
trimmed_head.weight.copy_(trimmed_weight)
if original_head.bias is not None:
original_bias = original_head.bias
selected_bias = original_bias[selected_indices]
trimmed_head.bias.copy_(selected_bias)
return trimmed_head
Step 4: Implement token remapping in the speculative decoding loop
During generation, convert drafter indices (from trimmed vocabulary) back to full vocabulary indices.
def speculative_decode_with_trimmed_vocab(target_model, drafter,
input_ids, vocab_mapping,
num_steps=256, draft_length=4):
"""
Run speculative decoding with a trimmed-vocabulary drafter.
The drafter outputs indices in the small vocabulary,
which we remap to the full vocabulary before verification.
"""
full_to_trimmed = vocab_mapping['full_to_trimmed']
trimmed_to_full = vocab_mapping['trimmed_to_full']
current_ids = input_ids
for step in range(num_steps):
# Drafter generates draft tokens
draft_tokens = []
drafter_ids = current_ids
for draft_step in range(draft_length):
with torch.no_grad():
# Forward pass through drafter with TRIMMED vocabulary
drafter_outputs = drafter(drafter_ids)
drafter_logits = drafter_outputs.logits[:, -1, :]
# Sample from trimmed vocabulary
trimmed_probs = torch.softmax(drafter_logits, dim=-1)
trimmed_sample = torch.multinomial(trimmed_probs, num_samples=1)
# REMAP: convert trimmed index to full vocabulary index
full_sample = torch.tensor(
[[trimmed_to_full[int(trimmed_sample[0, 0])]]],
device=trimmed_sample.device
)
draft_tokens.append(full_sample)
drafter_ids = torch.cat([drafter_ids, full_sample], dim=1)
# Concatenate draft tokens for verification
draft_ids = torch.cat(draft_tokens, dim=1)
candidate_ids = torch.cat([current_ids, draft_ids], dim=1)
# Target model verifies candidate tokens (with FULL vocabulary)
with torch.no_grad():
target_outputs = target_model(candidate_ids)
target_logits = target_outputs.logits
# Verification: accept tokens while probabilities remain high
for accept_idx in range(draft_length):
position = current_ids.shape[1] + accept_idx
candidate_token = draft_ids[0, accept_idx].item()
# Get target model's probability for this token
position_logits = target_logits[0, position - 1, :]
target_probs = torch.softmax(position_logits, dim=-1)
target_prob = target_probs[candidate_token].item()
# Accept or reject
if torch.rand(1).item() < target_prob:
# Accept: token is added permanently
current_ids = torch.cat([current_ids, draft_ids[:, accept_idx:accept_idx+1]], dim=1)
else:
# Reject: resample from target model
resampled = torch.multinomial(target_probs, num_samples=1)
current_ids = torch.cat([current_ids, resampled], dim=1)
break
return current_ids
Step 5: Evaluate speedup and coverage trade-off
Measure the actual inference speedup and how often the drafter misses tokens outside its vocabulary.
def evaluate_vocabtrim_performance(target_model, drafter, vocab_mapping,
eval_prompts, vocab_mapping_baseline):
"""
Compare trimmed drafter (VocabTrim) against baseline speculative decoding.
Key metrics: speedup, acceptance rate, coverage.
"""
results = {
'trimmed': [],
'baseline': []
}
for prompt in eval_prompts:
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
# Baseline: full vocabulary drafter
start_time = time.time()
output_baseline = speculative_decode_baseline(target_model, drafter,
input_ids, max_length=256)
time_baseline = time.time() - start_time
# Trimmed: VocabTrim drafter
start_time = time.time()
output_trimmed = speculative_decode_with_trimmed_vocab(
target_model, drafter, input_ids, vocab_mapping, num_steps=256
)
time_trimmed = time.time() - start_time
# Compute acceptance rate
# (tokens where drafter was right and target accepted them)
acceptance_baseline = count_accepted_tokens(output_baseline)
acceptance_trimmed = count_accepted_tokens(output_trimmed)
results['baseline'].append({
'time': time_baseline,
'acceptance': acceptance_baseline
})
results['trimmed'].append({
'time': time_trimmed,
'acceptance': acceptance_trimmed
})
speedup = np.mean([r['time'] for r in results['baseline']]) / \
np.mean([r['time'] for r in results['trimmed']])
acceptance_degradation = (np.mean([r['acceptance'] for r in results['baseline']]) -
np.mean([r['acceptance'] for r in results['trimmed']])) / \
np.mean([r['acceptance'] for r in results['baseline']])
return {
'speedup': speedup,
'acceptance_degradation': acceptance_degradation,
'vocab_coverage': vocab_mapping['coverage']
}
| Hyperparameter | Recommended Value | Notes |
|---|---|---|
| Target vocab size | 25K-50K | Higher = more speedup but diminishing returns |
| Coverage threshold | 85-95% | Ensure you're covering most sampled tokens |
| Calibration data size | 10K-50K examples | Measure frequencies on representative data |
| Acceptance rate tolerance | 97-99% | Can afford 1-3% degradation for 16% speedup |
| Remapping overhead | <1% | Negligible compared to inference speedup |
When to use VocabTrim:
When NOT to use VocabTrim:
Common pitfalls:
VocabTrim: Vocabulary Pruning for Efficient Speculative Decoding in LLMs https://arxiv.org/abs/2506.22694