adu2021 Low-Probability Tokens Sustain Exploration
Core Concept
During RLVR (RL with Verifiable Rewards) training on reasoning tasks, models suffer from entropy collapse—the policy converges too tightly to common answer patterns, eliminating valuable exploratory reasoning. Low-probability tokens (rare in pretraining) represent genuine thinking diversity but are systematically eliminated during RL. This skill introduces Low-probability Regularization (Lp-Reg): create a proxy distribution filtering noise while protecting "reasoning sparks"—low-probability tokens that enable novel solution paths.
The key insight: Not all entropy is useful exploration; distinguish signal-carrying low-probability tokens from noise, and protect only the signal.
Architecture Overview
- Reasoning Sparks Identification: Detect low-probability tokens that carry reasoning value
- Proxy Distribution Creation: Filter out noise while preserving signal
- Soft Regularization: KL divergence-based entropy maintenance targeted to valuable tokens
- Stable On-Policy RL: Enable continuous scaling without entropy collapse
- Token-Level Analysis: Fine-grained understanding of what enables diverse reasoning
Implementation Steps
1. Identify Reasoning Sparks
Analyze which low-probability tokens contribute to diverse reasoning paths.
class ReasoningSparkDetector:
"""
Identify low-probability tokens that enable novel reasoning.
"""
def __init__(self, model, reference_model):
self.model = model
self.ref_model = reference_model # Pretrained baseline
def identify_reasoning_sparks(self, reasoning_traces, solutions, correctness):
"""
Determine which low-probability tokens lead to correct solutions.
Args:
reasoning_traces: List of reasoning token sequences
solutions: Corresponding final answers
correctness: Boolean list of correctness
Returns:
Set of token IDs that represent reasoning sparks
"""
spark_candidates = {}
for trace, solution, is_correct in zip(reasoning_traces, solutions, correctness):
tokens = tokenize(trace)
for token_id in tokens:
token_prob_pretrain = self.ref_model.get_token_prob(token_id)
# Low-probability tokens: prob < 0.1
if token_prob_pretrain < 0.1:
if token_id not in spark_candidates:
spark_candidates[token_id] = {
"correct_count": 0,
"total_count": 0,
"solutions": []
}
spark_candidates[token_id]["total_count"] += 1
spark_candidates[token_id]["solutions"].append(solution)
if is_correct:
spark_candidates[token_id]["correct_count"] += 1
# Filter: keep tokens with high correctness rate
reasoning_sparks = set()
for token_id, stats in spark_candidates.items():
correctness_rate = stats["correct_count"] / (stats["total_count"] + 1e-8)
# Keep tokens that appear in correct solutions >50% of the time
if correctness_rate > 0.5 and stats["total_count"] > 3:
reasoning_sparks.add(token_id)
return reasoning_sparks
2. Create Proxy Distribution
Build a filtered probability distribution that removes noise while preserving reasoning sparks.
def create_proxy_distribution(logits, reasoning_sparks, noise_threshold=0.01):
"""
Create proxy distribution: filter noise, amplify reasoning sparks.
Args:
logits: Raw model logits (vocab_size,)
reasoning_sparks: Set of token IDs to protect
noise_threshold: Probability below which tokens are filtered
Returns:
Filtered probability distribution
"""
# Convert logits to probabilities
probs = torch.softmax(logits, dim=-1)
# Create mask for tokens to keep
keep_mask = torch.zeros_like(probs)
for token_id in reasoning_sparks:
keep_mask[token_id] = 1.0
# Remove low-probability noise (except protected sparks)
for token_id in range(len(probs)):
if token_id not in reasoning_sparks and probs[token_id] < noise_threshold:
probs[token_id] = 0.0
# Re-normalize to valid probability distribution
proxy_dist = probs / (torch.sum(probs) + 1e-8)
# Amplify reasoning sparks slightly
for token_id in reasoning_sparks:
proxy_dist[token_id] *= 1.2 # 20% boost
# Final normalization
proxy_dist = proxy_dist / (torch.sum(proxy_dist) + 1e-8)
return proxy_dist
3. Implement Low-Probability Regularization (Lp-Reg)
Apply soft KL regularization using the proxy distribution.
class LowProbabilityRegularization:
"""
Maintain entropy using low-probability token protection.
"""
def __init__(self, reference_model):
self.ref_model = reference_model
def compute_lp_reg_loss(
self,
model_dist,
proxy_dist,
reasoning_sparks,
temperature=1.0,
reg_weight=0.1
):
"""
Compute regularization loss protecting reasoning sparks.
Args:
model_dist: Current model's probability distribution
proxy_dist: Filtered proxy distribution
reasoning_sparks: Protected token set
temperature: Distribution sharpness
reg_weight: Regularization strength
Returns:
Scalar regularization loss
"""
# Apply temperature to soften distributions
model_dist_soft = model_dist ** (1 / temperature)
model_dist_soft = model_dist_soft / torch.sum(model_dist_soft)
proxy_dist_soft = proxy_dist ** (1 / temperature)
proxy_dist_soft = proxy_dist_soft / torch.sum(proxy_dist_soft)
# KL divergence: encourage model to match proxy distribution
# But only for tokens related to reasoning sparks
kl_loss = 0
for token_id in reasoning_sparks:
kl = model_dist_soft[token_id] * (
torch.log(model_dist_soft[token_id] + 1e-8) -
torch.log(proxy_dist_soft[token_id] + 1e-8)
)
kl_loss += kl
# Mild regularization: don't fully constrain, allow some flexibility
regularization = reg_weight * kl_loss
return regularization
4. Full RL Loop with Lp-Reg
Integrate regularization into standard RL training.
def train_with_lp_regularization(
model,
reference_model,
dataset,
config
):
"""
RL training with low-probability token protection.
"""
# Phase 1: Identify reasoning sparks
print("Phase 1: Identifying reasoning sparks")
# Collect initial rollouts
rollouts = []
for problem in dataset[:1000]: # Sample for spark detection
trace = model.generate(problem, max_tokens=500)
solution = extract_solution(trace)
is_correct = evaluate_correctness(solution, problem)
rollouts.append({
"trace": trace,
"solution": solution,
"is_correct": is_correct
})
detector = ReasoningSparkDetector(model, reference_model)
reasoning_sparks = detector.identify_reasoning_sparks(
[r["trace"] for r in rollouts],
[r["solution"] for r in rollouts],
[r["is_correct"] for r in rollouts]
)
print(f"Identified {len(reasoning_sparks)} reasoning sparks")
# Phase 2: RL training with Lp-Reg
print("\nPhase 2: RL training with low-probability regularization")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
lp_regularizer = LowProbabilityRegularization(reference_model)
for epoch in range(config.num_epochs):
epoch_loss = 0
for batch_idx, problem in enumerate(dataset):
# Generate reasoning with current policy
trace = model.generate(problem, max_tokens=500)
solution = extract_solution(trace)
# Evaluate reward
is_correct = evaluate_correctness(solution, problem)
reward = 1.0 if is_correct else -1.0
# Get model logits and probabilities
logits = model.get_logits_for_sequence(trace)
model_dist = torch.softmax(logits, dim=-1)
# Create proxy distribution for this sequence
proxy_dist = create_proxy_distribution(
logits,
reasoning_sparks,
noise_threshold=0.01
)
# Policy gradient loss
log_prob = model.compute_log_prob(trace)
pg_loss = -reward * log_prob
# Lp-Reg loss: protect reasoning sparks
reg_loss = lp_regularizer.compute_lp_reg_loss(
model_dist,
proxy_dist,
reasoning_sparks,
temperature=1.0,
reg_weight=0.1
)
# Total loss: RL + regularization
total_loss = pg_loss + reg_loss
# Update
optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
epoch_loss += total_loss.item()
if (batch_idx + 1) % 100 == 0:
avg_loss = epoch_loss / (batch_idx + 1)
print(f"Epoch {epoch}, Batch {batch_idx}: loss={avg_loss:.4f}")
return model
5. Evaluation: Stability and Performance
Assess that model maintains exploration while improving correctness.
def evaluate_lp_reg_training(model, benchmark_dataset, reference_model):
"""
Evaluate exploration stability and performance.
"""
results = {
"accuracy": 0,
"entropy": 0,
"diversity": 0
}
correct = 0
entropies = []
solutions_set = set()
for problem in benchmark_dataset:
# Generate multiple samples to assess diversity
traces = []
for _ in range(5):
trace = model.generate(problem, max_tokens=500, temperature=0.8)
traces.append(trace)
solution = extract_solution(trace)
solutions_set.add(solution)
# Check correctness
if evaluate_correctness(solution, problem):
correct += 1
# Compute entropy of solutions for this problem
solution_entropy = compute_entropy([extract_solution(t) for t in traces])
entropies.append(solution_entropy)
results["accuracy"] = correct / (len(benchmark_dataset) * 5) * 100
results["entropy"] = np.mean(entropies)
results["diversity"] = len(solutions_set)
print(f"Accuracy: {results['accuracy']:.1f}%")
print(f"Entropy: {results['entropy']:.3f}")
print(f"Solution diversity: {results['diversity']}")
return results
Practical Guidance
Hyperparameters:
- Noise threshold: 0.01 (filter very rare tokens)
- Spark correctness threshold: 0.5 (keep tokens in correct solutions >50%)
- Regularization weight: 0.1 (soft constraint)
- Temperature for soft distributions: 1.0-2.0 (sharpness control)
- Spark detection sample size: 1000+ problems
When to Use:
- RL training on reasoning tasks (math, code, logic)
- Experiencing entropy collapse during training
- Want to preserve diverse solution paths
- Tasks have multiple valid approaches
When NOT to Use:
- Single-solution tasks where diversity not valued
- Simple supervised learning (not RL)
- Streaming inference where entropy control not critical
- Tasks where rare tokens are pure noise
Reference
Low-Probability Tokens Sustain Exploration in Reasoning RL — arXiv:2510.03222