adu2021 PosS: Position Specialist Generates Better Draft for Speculative Decoding
Core Concept
Speculative decoding accelerates LLM inference by using a smaller draft model to predict multiple tokens, then verifying them with the larger target model. PosS identifies a critical limitation: draft models suffer from degraded token quality at later positions due to error accumulation. By deploying multiple position-specialized layers—each handling a narrow, predictable range of feature deviation—PosS improves acceptance length and speed-up ratios by up to 5.7%.
Architecture Overview
- Position-Specific Feature Deviation: Each position in a draft has characteristic feature deviation from the target model. Later positions accumulate errors multiplicatively.
- Position Specialists: Separate transformer layers trained to handle specific positions with their expected deviation levels
- Position-Wise Acceptance Rate (Pos-Acc): New metric analyzing draft quality across different positions in the prediction sequence
- Chain Rule Decomposition: Overall acceptance depends on multiplying individual position acceptance rates; fixing weak positions improves end-to-end throughput
- Training with Simulated Error: Specialists learn from features generated by predecessors, simulating inference conditions
Implementation
Step 1: Analyze Position-Specific Acceptance Degradation
import torch
import numpy as np
from typing import List, Dict
from collections import defaultdict
class PositionAcceptanceAnalyzer:
"""Diagnose where draft models fail in speculative decoding"""
def __init__(self, draft_model, target_model):
self.draft_model = draft_model
self.target_model = target_model
def compute_position_wise_acceptance(self,
prompts: List[torch.Tensor],
num_draft_tokens: int = 4) -> Dict:
"""
For each position k in the draft, measure acceptance rate.
Key finding: pos-acc rapidly deteriorates beyond k=1.
"""
position_stats = defaultdict(list)
for prompt in prompts:
# Generate multiple draft tokens
draft_tokens = self.draft_model.generate_tokens(
prompt, num_tokens=num_draft_tokens
)
# For each position, check if target model accepts
for pos in range(num_draft_tokens):
draft_token_at_pos = draft_tokens[pos]
# Target model's prediction at this position
target_logits = self.target_model.forward(
torch.cat([prompt, draft_tokens[:pos]], dim=0)
)
target_top_token = torch.argmax(target_logits[-1, :])
# Check acceptance
accepted = (draft_token_at_pos == target_top_token)
position_stats[pos].append(float(accepted))
# Aggregate statistics
position_acceptance_rates = {}
for pos in range(num_draft_tokens):
acc_rate = np.mean(position_stats[pos])
position_acceptance_rates[pos] = acc_rate
print(f"Position {pos}: {acc_rate:.1%} acceptance rate")
# Compute cumulative impact via chain rule
overall_acceptance = np.prod(list(position_acceptance_rates.values()))
print(f"\nOverall acceptance (pos[0] × pos[1] × ...): {overall_acceptance:.1%}")
return position_acceptance_rates
def analyze_feature_deviation(self, prompts: List[torch.Tensor]) -> Dict:
"""
Measure feature deviation at each position.
Later positions have higher feature deviation due to error accumulation.
"""
position_deviations = defaultdict(list)
for prompt in prompts:
# Run draft and target forward passes
draft_hidden = self.draft_model.extract_hidden_states(prompt)
target_hidden = self.target_model.extract_hidden_states(prompt)
# Measure feature deviation at each generation step
for pos in range(len(draft_hidden) - 1):
deviation = torch.norm(
draft_hidden[pos] - target_hidden[pos]
)
position_deviations[pos].append(deviation.item())
# Average deviation per position
avg_deviations = {}
for pos in range(len(draft_hidden) - 1):
avg_dev = np.mean(position_deviations[pos])
avg_deviations[pos] = avg_dev
print(f"Position {pos}: avg feature deviation = {avg_dev:.4f}")
return avg_deviations
Step 2: Design Position-Specialist Architecture
class PositionSpecialist(torch.nn.Module):
"""Transformer layer specialized for a specific position"""
def __init__(self, hidden_dim: int, num_layers: int,
position_id: int, expected_deviation: float):
super().__init__()
self.position_id = position_id
self.expected_deviation = expected_deviation
# Adapt model size based on expected deviation
# Higher deviation = need more capacity
adapted_layers = max(2, int(num_layers * expected_deviation / 0.5))
self.specialist_layers = torch.nn.ModuleList([
self.TransformerBlock(hidden_dim)
for _ in range(adapted_layers)
])
# Loss components for training
self.token_loss_weight = 0.7
self.feature_loss_weight = 0.2
self.topk_loss_weight = 0.1
class TransformerBlock(torch.nn.Module):
def __init__(self, hidden_dim):
super().__init__()
self.self_attn = torch.nn.MultiheadAttention(
embed_dim=hidden_dim, num_heads=8
)
self.ffn = torch.nn.Sequential(
torch.nn.Linear(hidden_dim, 4 * hidden_dim),
torch.nn.ReLU(),
torch.nn.Linear(4 * hidden_dim, hidden_dim),
)
self.norm1 = torch.nn.LayerNorm(hidden_dim)
self.norm2 = torch.nn.LayerNorm(hidden_dim)
def forward(self, x):
x = x + self.self_attn(x, x, x)[0]
x = self.norm1(x)
x = x + self.ffn(x)
x = self.norm2(x)
return x
def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
"""Process hidden state through specialist layers"""
x = hidden_state
for layer in self.specialist_layers:
x = layer(x)
return x
class PositionSpecialistEnsemble(torch.nn.Module):
"""Multiple position specialists covering all draft positions"""
def __init__(self, base_model, num_specialists: int = 4):
super().__init__()
self.num_specialists = num_specialists
self.specialists = torch.nn.ModuleList()
# Create specialist for each position
position_deviations = self.estimate_position_deviations(base_model)
for pos in range(num_specialists):
specialist = PositionSpecialist(
hidden_dim=base_model.hidden_dim,
num_layers=base_model.num_layers,
position_id=pos,
expected_deviation=position_deviations.get(pos, 0.5)
)
self.specialists.append(specialist)
def estimate_position_deviations(self, base_model) -> Dict:
"""Estimate expected feature deviation per position"""
deviations = {}
# Empirical: deviation grows with position
for pos in range(self.num_specialists):
# Approximation: exponential growth
deviations[pos] = 0.1 * (1.5 ** pos)
return deviations
def forward(self, hidden_states: torch.Tensor,
position: int) -> torch.Tensor:
"""
Route hidden state through appropriate specialist.
If position >= num_specialists, use last specialist.
"""
specialist_idx = min(position, self.num_specialists - 1)
specialist = self.specialists[specialist_idx]
output = specialist(hidden_states)
return output
Step 3: Training Position Specialists
class PositionSpecialistTrainer:
"""Train specialists using simulated inference conditions"""
def __init__(self, target_model, specialist_ensemble: PositionSpecialistEnsemble):
self.target_model = target_model
self.specialists = specialist_ensemble
self.optimizer = torch.optim.AdamW(
self.specialists.parameters(), lr=2e-4
)
def train_on_batch(self, prompts: torch.Tensor,
completions: torch.Tensor) -> Dict[str, float]:
"""
Train specialists on data with three loss components:
1. Token-level: predict correct next token
2. Feature-level: match target model features
3. Top-K: preserve top-K token distribution
"""
losses = {}
# Extract target model's hidden states for supervision
target_hidden = self.target_model.extract_hidden_states(
torch.cat([prompts, completions], dim=1)
)
target_logits = self.target_model.forward(
torch.cat([prompts, completions], dim=1)
)
total_loss = 0
# Train each position specialist
for pos in range(self.specialists.num_specialists):
# Simulate inference at this position:
# Use specialist for previous positions, then current position
if pos == 0:
current_hidden = target_hidden[-1, :].unsqueeze(0)
else:
# Use specialist output from previous position
prev_output = self.specialists.forward(
target_hidden[-2, :].unsqueeze(0),
position=pos - 1
)
current_hidden = prev_output
# Run through specialist
specialist_output = self.specialists.forward(
current_hidden,
position=pos
)
# Compute three-part loss
token_loss = self.compute_token_loss(
specialist_output,
target_logits[pos],
)
feature_loss = self.compute_feature_loss(
specialist_output,
target_hidden[pos],
)
topk_loss = self.compute_topk_loss(
specialist_output,
target_logits[pos],
k=5
)
position_loss = (
0.7 * token_loss +
0.2 * feature_loss +
0.1 * topk_loss
)
total_loss += position_loss
losses[f'pos_{pos}'] = position_loss.item()
# Backward pass
self.optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self.specialists.parameters(), 1.0)
self.optimizer.step()
losses['total'] = total_loss.item()
return losses
def compute_token_loss(self, specialist_logits: torch.Tensor,
target_logits: torch.Tensor) -> torch.Tensor:
"""Cross-entropy loss on predicted token"""
return torch.nn.functional.cross_entropy(
specialist_logits.view(-1, specialist_logits.size(-1)),
target_logits.argmax(-1)
)
def compute_feature_loss(self, specialist_hidden: torch.Tensor,
target_hidden: torch.Tensor) -> torch.Tensor:
"""L2 loss on hidden state features"""
return torch.nn.functional.mse_loss(
specialist_hidden,
target_hidden.detach()
)
def compute_topk_loss(self, specialist_logits: torch.Tensor,
target_logits: torch.Tensor,
k: int = 5) -> torch.Tensor:
"""KL divergence on top-K token probabilities"""
specialist_probs = torch.nn.functional.softmax(specialist_logits, dim=-1)
target_probs = torch.nn.functional.softmax(target_logits, dim=-1)
# Keep only top-K for divergence
_, topk_indices = torch.topk(target_probs, k)
specialist_topk = specialist_probs[topk_indices]
target_topk = target_probs[topk_indices]
return torch.nn.functional.kl_div(
torch.log(specialist_topk + 1e-8),
target_topk.detach(),
reduction='batchmean'
)
Step 4: Integration with Speculative Decoding
class SpeculativeDecodingWithPosS:
"""Enhanced speculative decoding using position specialists"""
def __init__(self, target_model, draft_model,
position_specialists: PositionSpecialistEnsemble):
self.target_model = target_model
self.draft_model = draft_model
self.specialists = position_specialists
def generate_with_verification(self, prompt: torch.Tensor,
max_length: int = 100,
num_draft_tokens: int = 4) -> torch.Tensor:
"""
Speculative decoding with position-specialized draft:
1. Use specialist-enhanced draft to predict k tokens
2. Verify with target model
3. Accept accepted tokens, resample rejected ones
"""
generated = prompt.clone()
total_draft_tokens = 0
total_verified_tokens = 0
while generated.shape[0] < max_length:
# Generate draft tokens with specialist enhancement
draft_tokens = []
for pos in range(num_draft_tokens):
# Get base draft prediction
draft_logits = self.draft_model.forward(generated)
# Enhance with position specialist
draft_hidden = self.draft_model.extract_hidden_state(generated)
enhanced_hidden = self.specialists.forward(
draft_hidden, position=pos
)
# Blend enhanced and base predictions
enhanced_logits = self.draft_model.decode(enhanced_hidden)
blended_logits = 0.6 * enhanced_logits + 0.4 * draft_logits
draft_token = torch.argmax(blended_logits)
draft_tokens.append(draft_token)
total_draft_tokens += 1
# Verify with target model
draft_sequence = torch.cat([generated, torch.stack(draft_tokens)])
target_logits = self.target_model.forward(draft_sequence)
# Check each drafted token
accepted_count = 0
for pos, draft_token in enumerate(draft_tokens):
target_token = torch.argmax(target_logits[-(len(draft_tokens)-pos)])
if draft_token == target_token:
generated = torch.cat([generated, draft_token.unsqueeze(0)])
accepted_count += 1
total_verified_tokens += 1
else:
# Rejection: sample from target distribution
target_probs = torch.softmax(target_logits[-(len(draft_tokens)-pos)], dim=-1)
new_token = torch.multinomial(target_probs, 1)
generated = torch.cat([generated, new_token])
total_verified_tokens += 1
break
# Compute metrics
acceptance_length = total_verified_tokens / total_draft_tokens
speedup_ratio = total_verified_tokens / (total_verified_tokens / num_draft_tokens + total_verified_tokens)
print(f"Acceptance length: {acceptance_length:.2f}")
print(f"Speedup ratio: {speedup_ratio:.2f}×")
return generated
Practical Guidance
Measure Pos-Acc First: Profile your draft model with position-wise acceptance analysis. Identify which positions have lowest acceptance rates—these are bottlenecks.
Specialist Count: Start with 4-8 specialists covering 4-8 draft positions. More specialists provide finer granularity but increase computational overhead.
Feature Deviation Estimation: Position deviation grows roughly exponentially. Allocate more parameters to later positions which have higher deviation.
Training Strategy: Train specialists with simulated inference conditions using previous specialist outputs, not clean ground truth.
Three-Part Loss: Use token-level (70%), feature-level (20%), and top-K (10%) losses. This balances correctness with feature alignment.
Integration: Enhance draft model predictions by routing through specialists, then blend enhanced logits with base draft logits (60-40 weighting works well).
Reference
- Paper: PosS (2506.03566)
- Key Metric: Position-wise acceptance rate (pos-acc) analysis
- Improvements: 4.5% on acceptance length, 5.7% on speed-up ratio
- Architecture: Multiple position-specialized transformer layers trained with simulated error accumulation