Accelerate test-time scaling for diffusion language models by identifying inconsistent tokens, selectively remask and regenerate only uncertain tokens, and aggregate across samples via voting. Achieve 5.5-22× speedup over standard iterative sampling with 6-8% accuracy gains on reasoning tasks.
["Diffusion Language Models","Test-Time Scaling","Voting","Inference Optimization","Ensemble Methods"]
description
Accelerate test-time scaling for diffusion language models by identifying inconsistent tokens, selectively remask and regenerate only uncertain tokens, and aggregate across samples via voting. Achieve 5.5-22× speedup over standard iterative sampling with 6-8% accuracy gains on reasoning tasks.
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
"""
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
"""
len
for
in
range
for
in
max
if
self
return
Step 2: Selective remask and regenerate
classSelectiveRemaskingStrategy:
"""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
defremask_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
defregenerate_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
defiterative_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 inrange(num_iterations):
masked = self.remask_inconsistent_tokens(current, inconsistent_pos)
current = self.regenerate_masked_region(masked)
return current
Step 3: Aggregate predictions via voting
classTokenVotingAggregator:
"""Aggregate multiple samples via majority voting.""" @staticmethoddefmajority_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 isNone:
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
@staticmethoddefconfidence_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 isNone:
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 inenumerate(samples):
token = sample[pos].item()
weight = log_probs[sample_idx][pos].exp().item()
if token notin 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
