adu2021 dVoting: Fast Voting for dLLMs
Problem Context
Diffusion language models can generate tokens in any order and refine them through masking and redenoising. Standard test-time scaling generates multiple samples and averages predictions. However, most tokens remain consistent across samples—only a small subset vary. dVoting exploits this by identifying variable tokens, regenerating only those selectively, and using voting to aggregate.
Core Concept
dVoting operates in three phases: (1) identify tokens with cross-sample disagreement, (2) selectively remask inconsistent tokens, (3) regenerate and aggregate via voting. This reduces computation by focusing computation on the small disagreement set.
Architecture Overview
- Consistency detection: Compare tokens across multiple samples
- Uncertainty identification: Flag tokens with disagreement
- Selective remask: Only remask uncertain tokens
- Regeneration: Run denoising only on uncertainty regions
- Voting aggregation: Combine candidate answers via vote
Implementation
Step 1: Generate samples and identify inconsistencies
import torch
from typing import List, Dict, Tuple
from collections import Counter
class DiffusionConsistencyAnalyzer:
"""Identify tokens showing cross-sample disagreement."""
def __init__(self, agreement_threshold: float = 0.8):
self.agreement_threshold = agreement_threshold
def detect_inconsistent_tokens(
self,
samples: List[torch.Tensor], # [num_samples, seq_len]
seq_len: int
) -> Tuple[torch.Tensor, List[int]]:
"""
Identify tokens with low agreement across samples.
Args:
samples: List of generated token sequences
seq_len: Sequence length
Returns:
agreement_scores: Agreement ratio per token [seq_len]
inconsistent_positions: Indices of inconsistent tokens
"""
num_samples = len(samples)
agreement_scores = torch.ones(seq_len)
inconsistent_positions = []
for pos in range(seq_len):
tokens_at_pos = [sample[pos].item() for sample in samples]
token_counts = Counter(tokens_at_pos)
max_count = max(token_counts.values())
agreement = max_count / num_samples
agreement_scores[pos] = agreement
if agreement < self.agreement_threshold:
inconsistent_positions.append(pos)
return agreement_scores, inconsistent_positions
Step 2: Selective remask and regenerate
class SelectiveRemaskingStrategy:
"""Remask only inconsistent tokens for efficient regeneration."""
def __init__(self, model, num_denoising_steps: int = 10):
self.model = model
self.num_denoising_steps = num_denoising_steps
def remask_inconsistent_tokens(
self,
sample: torch.Tensor, # [seq_len]
inconsistent_pos: List[int]
) -> torch.Tensor:
"""Create mask for inconsistent tokens."""
masked_sample = sample.clone()
for pos in inconsistent_pos:
masked_sample[pos] = self.model.mask_token_id
return masked_sample
def regenerate_masked_region(
self,
masked_sample: torch.Tensor,
context: torch.Tensor = None
) -> torch.Tensor:
"""
Run diffusion denoising on masked region.
Leverages dLLM's ability to refine specific positions.
"""
refined = self.model.denoise(
masked_sample,
num_steps=self.num_denoising_steps,
context=context
)
return refined
def iterative_refinement(
self,
sample: torch.Tensor,
inconsistent_pos: List[int],
num_iterations: int = 2
) -> torch.Tensor:
"""Multiple passes of selective remask-regenerate."""
current = sample.clone()
for iteration in range(num_iterations):
masked = self.remask_inconsistent_tokens(current, inconsistent_pos)
current = self.regenerate_masked_region(masked)
return current
Step 3: Aggregate predictions via voting
class TokenVotingAggregator:
"""Aggregate multiple samples via majority voting."""
@staticmethod
def majority_vote(
samples: List[torch.Tensor], # [num_samples, seq_len]
positions: List[int] = None
) -> torch.Tensor:
"""
Compute majority vote for each position.
Args:
samples: Generated samples
positions: Positions to vote on (default: all)
Returns:
voted_sequence: Final tokens from voting [seq_len]
"""
num_samples = len(samples)
seq_len = samples[0].shape[0]
if positions is None:
positions = list(range(seq_len))
voted_sequence = samples[0].clone()
for pos in positions:
tokens_at_pos = [sample[pos].item() for sample in samples]
token_counts = Counter(tokens_at_pos)
majority_token = token_counts.most_common(1)[0][0]
voted_sequence[pos] = majority_token
return voted_sequence
@staticmethod
def confidence_weighted_voting(
samples: List[torch.Tensor],
log_probs: List[torch.Tensor], # [num_samples, seq_len]
positions: List[int] = None
) -> torch.Tensor:
"""
Weight voting by model confidence scores.
Higher likelihood tokens weighted more heavily.
"""
num_samples = len(samples)
seq_len = samples[0].shape[0]
if positions is None:
positions = list(range(seq_len))
voted_sequence = samples[0].clone()
for pos in positions:
# Weighted vote by log probability
vote_weights = {}
for sample_idx, sample in enumerate(samples):
token = sample[pos].item()
weight = log_probs[sample_idx][pos].exp().item()
if token not in vote_weights:
vote_weights[token] = 0.0
vote_weights[token] += weight
best_token = max(vote_weights.items(), key=lambda x: x[1])[0]
voted_sequence[pos] = best_token
return voted_sequence
Step 4: Full dVoting pipeline
class DVotingPipeline:
"""Complete dVoting strategy for fast test-time scaling."""
def __init__(
self,
model,
agreement_threshold: float = 0.8,
num_regeneration_steps: int = 10,
num_refinement_passes: int = 1
):
self.model = model
self.analyzer = DiffusionConsistencyAnalyzer(agreement_threshold)
self.remask_strategy = SelectiveRemaskingStrategy(model, num_regeneration_steps)
self.aggregator = TokenVotingAggregator()
self.num_refinement_passes = num_refinement_passes
def dvoting_forward(
self,
prompt: str,
num_samples: int = 8,
compute_log_probs: bool = False
) -> Dict[str, torch.Tensor]:
"""
Execute dVoting for fast ensemble inference.
Args:
prompt: Input prompt
num_samples: Initial samples to generate
compute_log_probs: Return confidence scores
Returns:
final_sequence: Voted final prediction
inconsistent_count: Number of inconsistent tokens
speedup_factor: Estimated speedup vs. standard voting
"""
# Step 1: Generate initial samples (full diffusion)
samples = []
log_probs_list = []
for _ in range(num_samples):
sample, log_probs = self.model.generate_with_logprobs(
prompt, num_steps=50
)
samples.append(sample)
log_probs_list.append(log_probs)
# Step 2: Identify inconsistencies
agreement_scores, inconsistent_pos = self.analyzer.detect_inconsistent_tokens(
samples, samples[0].shape[0]
)
inconsistent_count = len(inconsistent_pos)
total_tokens = samples[0].shape[0]
# Step 3: Selective regeneration
regenerated_samples = []
for sample in samples:
refined = self.remask_strategy.iterative_refinement(
sample, inconsistent_pos,
num_iterations=self.num_refinement_passes
)
regenerated_samples.append(refined)
# Step 4: Vote on inconsistent positions
if compute_log_probs:
final_sequence = self.aggregator.confidence_weighted_voting(
regenerated_samples, log_probs_list, inconsistent_pos
)
else:
# On consistent positions, use original agreement
final_sequence = regenerated_samples[0].clone()
final_sequence = self.aggregator.majority_vote(
regenerated_samples, inconsistent_pos
)
# Estimate speedup
regeneration_fraction = inconsistent_count / total_tokens
speedup_vs_full = 1.0 / max(regeneration_fraction, 0.1)
return {
'final_sequence': final_sequence,
'agreement_scores': agreement_scores,
'inconsistent_positions': inconsistent_pos,
'inconsistent_count': inconsistent_count,
'speedup_factor': speedup_vs_full
}
Step 5: Training and evaluation
def evaluate_dvoting_speedup(
model,
test_prompts: List[str],
verifier,
num_samples: int = 8,
device: str = 'cuda'
) -> Dict[str, float]:
"""
Benchmark dVoting speedup and accuracy.
"""
dvoting = DVotingPipeline(model, num_samples)
total_correct = 0
total_speedup = 0.0
num_tests = 0
for prompt in test_prompts:
result = dvoting.dvoting_forward(prompt, num_samples=num_samples)
# Check correctness
prediction = model.decode(result['final_sequence'])
is_correct = verifier(prediction)
total_correct += is_correct
# Track speedup
total_speedup += result['speedup_factor']
num_tests += 1
return {
'accuracy': total_correct / num_tests,
'avg_speedup': total_speedup / num_tests
}
Practical Guidance
When to use: Diffusion language models needing test-time scaling; ensemble prediction with limited compute budget
Hyperparameters:
- agreement_threshold: 0.75-0.85 (higher = more selective)
- num_regeneration_steps: 5-20 (tradeoff: quality vs speed)
- num_samples: 4-8 (initial sample count)
- num_refinement_passes: 1-2 (iterations on inconsistent tokens)
Key advantages:
- 5-22× speedup vs. standard iterative sampling
- 6-8% accuracy improvements via voting
- Focuses computation on uncertainty
- Works with any diffusion LM
Common pitfalls:
- agreement_threshold too high → misses real disagreements
- Too few refinement passes → suboptimal regeneration
- Not weighting by confidence → loses signal quality
Scaling: Speedup scales with inconsistency sparsity; best on diverse generation tasks.
Reference
Paper: https://arxiv.org/abs/2602.12153 Related work: Diffusion models, test-time scaling, ensemble methods, majority voting Benchmarks: GSM8K, MATH, reasoning tasks