Stabilize LLM reasoning training by replacing mean-based advantage baselines with K-quantile baselines, preventing both entropy collapse and explosion while improving performance on mathematical benchmarks through response-level gating and asymmetric sample weighting.
Stabilize LLM reasoning training by replacing mean-based advantage baselines with K-quantile baselines, preventing both entropy collapse and explosion while improving performance on mathematical benchmarks through response-level gating and asymmetric sample weighting.
Entropy-Safe Advantage Estimation Through Quantile Baselines
Outcome
Achieve stable, efficient reinforcement learning training for language models on reasoning tasks by automatically regulating exploration-exploitation balance through quantile-based advantage estimation. Prevent training instability (entropy oscillation) while improving reasoning accuracy on mathematical benchmarks.
Problem Context
Reinforcement Learning with Verifiable Rewards (RLVR) enables direct optimization of LLMs toward measurable outcomes (correctness on math problems). However, existing value-free RL methods like GRPO and DAPO suffer from an "entropy dilemma": the mean-based advantage baselines they use trigger unstable oscillations between entropy collapse (premature convergence to single outputs) and entropy explosion (excessive stochasticity that degrades learning).
Prior work addressed only entropy collapse with ad-hoc token-level fixes. This leaves entropy explosion unaddressed, causing training to destabilize around convergence plateaus. The core issue: mean baselines treat all samples symmetrically, failing to distinguish between hard queries (needing diversity) and easy queries (needing precision).
Core Concept
Quantile Advantage Estimation (QAE) replaces the empirical mean baseline with a group-wise K-quantile baseline, creating a "two-regime gate" that adapts update behavior to query difficulty:
Hard queries (success rate p ≤ 1−K): Preserve all sampled successes, zero out failures. Reinforces rare wins through positive advantages.
Easy queries (success rate p > 1−K): Zero out successes, preserve failures. Targets corrections through negative advantages.
This asymmetric masking provides provable bounds on entropy changes. Under first-order softmax policy updates, QAE minimizes entropy growth on hard queries (explosion-proof) while maximizing it on easy queries (collapse-proof). A single hyperparameter K controls the quantile threshold, eliminating need for complex token-level tuning.
The method requires only a one-line code change in existing implementations and sparsifies the update signal—approximately 80% of sampled responses receive zero advantage, concentrating gradient flow on informative minorities.
Architecture Overview
Baseline Computation: Replace mean with K-quantile of sample success indicators across a batch.
Response-Level Gating: Mask advantages at the full-response level (not token-by-token), simplifying the update rule.
Asymmetric Weighting: Hard-query samples use positive advantages; easy-query samples use negative advantages. Weights no longer symmetric √p(1−p).
Entropy Safety: Provable one-step bounds on entropy change in both directions via Proposition 4.2.
Composability: Orthogonal to existing token-level methods (Clip-Cov, KL-Cov) and sequence-level methods (GSPO).
Implementation
Step 1: Replace Baseline Computation in Training Loop
In your RLVR training code, swap the mean-based baseline for quantile calculation. Most frameworks compute advantage as (sample_reward - baseline). With QAE, the baseline becomes the K-quantile of rewards rather than the mean.
Python pseudocode for quantile baseline
import numpy as np
from torch.nn.functional import softmax
defcompute_quantile_baseline(rewards, K=0.4):
"""
Compute K-quantile baseline for a batch of rewards.
Args:
rewards: array of shape (batch_size,), scalar rewards per sample
K: quantile parameter (0.4 is default, range [0.0, 1.0])
Returns:
baseline: scalar quantile value
"""return np.quantile(rewards, 1 - K)
defcompute_qae_advantages(rewards, K=0.4):
"""
Compute advantages using quantile-based baseline.
Masks responses at the response level (not token level).
"""
baseline = compute_quantile_baseline(rewards, K)
raw_advantages = rewards - baseline
# Response-level masking: zero out non-conforming updates
advantages = np.where(raw_advantages > 0, raw_advantages, 0)
advantages = np.where(raw_advantages < 0, raw_advantages, 0)
# Net effect: positive advantages on hard queries, negative on easy queriesreturn advantages
Step 2: Integrate with Policy Gradient Update
Modify your policy gradient computation to use QAE advantages. If using DAPO or GRPO, the update step changes only in the advantage computation—the rest of the pipeline remains identical.
Integration pattern for GRPO-style training
defgrpo_update_step_with_qae(
model,
prompts,
sampled_responses,
verifier,
optimizer,
K=0.4,
temperature=0.7):
"""
Single GRPO update with QAE baseline.
Args:
model: language model with forward and logprobs methods
prompts: input prompts for batch
sampled_responses: completions for each prompt (32 per prompt typical)
verifier: deterministic verifier (pass/fail checker)
optimizer: torch optimizer
K: quantile parameter
temperature: sampling temperature for logprobs
Returns:
loss: scalar loss value
"""# Score all samples
rewards = torch.tensor([
float(verifier.verify(response))
for response in sampled_responses
])
# Compute QAE advantages (one line change from mean baseline)
baseline = torch.quantile(rewards, 1 - K)
advantages = rewards - baseline
# Mask at response level
hard_query_mask = (rewards.mean() <= 1 - K)
if hard_query_mask:
advantages = torch.where(advantages > 0, advantages,
torch.zeros_like(advantages))
else:
advantages = torch.where(advantages < 0, advantages,
torch.zeros_like(advantages))
# Compute log probabilitieswith torch.no_grad():
logprobs = model.compute_logprobs(
prompts, sampled_responses, temperature=temperature
)
# Policy gradient loss (standard form)
loss = -(logprobs * advantages).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.item()
Step 3: Hyperparameter Tuning and Monitoring
Set K based on your task's success rate distribution and monitor entropy metrics.
Hyperparameter configuration example
classQAEConfig:
"""Configuration for Quantile Advantage Estimation."""def__init__(
self,
K=0.4,
temperature=0.7,
entropy_warmup_steps=500,
log_interval=50):
self.K = K # quantile thresholdself.temperature = temperature # sampling temperatureself.entropy_warmup_steps = entropy_warmup_steps
self.log_interval = log_interval
@staticmethoddeffrom_baseline_stats(baseline_success_rate):
"""Auto-tune K based on observed success rate."""if baseline_success_rate < 0.2:
return QAEConfig(K=0.3)
elif baseline_success_rate < 0.5:
return QAEConfig(K=0.4)
else:
return QAEConfig(K=0.5)
deftraining_loop_with_monitoring(
model,
train_data,
verifier,
config,
num_steps=5000):
"""Training loop with entropy monitoring."""
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
entropy_log = []
accuracy_log = []
for step inrange(num_steps):
# Sample batch
batch_prompts = [d['prompt'] for d in train_data[:32]]
all_responses = sample_responses(model, batch_prompts, 32)
# Update
loss = grpo_update_step_with_qae(
model, batch_prompts, all_responses, verifier,
optimizer, K=config.K, temperature=config.temperature
)
# Monitor entropy (compute on validation set)if step % config.log_interval == 0:
val_entropy = compute_policy_entropy(model, train_data[:16])
val_accuracy = evaluate_accuracy(model, verifier, train_data[:16])
entropy_log.append(val_entropy)
accuracy_log.append(val_accuracy)
return entropy_log, accuracy_log
Practical Guidance
Hyperparameter Reference
Parameter
Default
Range
When to Adjust
K
0.4
[0.1, 0.7]
Lower K for harder tasks, higher K for easier tasks
temperature
0.7
[0.5, 1.5]
Increase if entropy collapse, decrease if explosion
entropy_warmup_steps
500
[100, 2000]
More steps needed for larger models (14B+)
batch_size
32
[16, 64]
Use larger batches for stable quantile estimates
When to Use QAE
Multi-step reasoning tasks with verifiable rewards (math, code, logic puzzles)
Models showing entropy oscillation in RLVR training (spikes and drops in policy entropy)
Mixed-difficulty datasets where some queries are easy (high success rate) and others are hard
Resource-constrained settings where token-level masking adds computational overhead
Scaling beyond 8B parameters where entropy control becomes critical
When NOT to Use
Single-shot classification tasks without reasoning where entropy control is unnecessary
Reward signals that are noisy or non-binary (continuous rewards violate the discrete success/failure assumption)
Very small batch sizes (<16) where quantile estimates become unstable
Fully solved tasks where success rate is >95% (QAE provides diminishing returns)
Tasks requiring exploration-heavy training where all samples should influence learning equally
Common Pitfalls
Quantile instability with small batches: Ensure batch_size ≥ 16 for stable K-quantile estimates. With smaller batches, the quantile computation becomes jittery.
K-value overfit to initialization: Tune K on a held-out validation set, not the same task you report results on. Different task domains require different K values.
Ignoring success rate distribution: Compute baseline success rate first. If most queries already succeed (>80%), QAE offers less benefit. If most fail (<10%), use K ≈ 0.2.
Skipping entropy monitoring: Track policy entropy every 50-100 steps. If entropy dips below 0.5 nats early in training, increase K. If it spikes above 3 nats, decrease K.
Mixing with incompatible baselines: QAE is incompatible with other baseline-modification schemes. Use only one of: mean baseline (standard), QAE quantile baseline, or actor-critic baselines.
Reference
Paper: Quantile Advantage Estimation: Stabilizing RLVR for LLM Reasoning
Authors: Junkang Wu, Kexin Huang, Jiancan Wu, An Zhang, Xiang Wang, Xiangnan He
arXiv: https://arxiv.org/abs/2509.22611Submission: September 26, 2025 | Last revised: February 28, 2026
Key Results (Qwen3-8B on mathematical reasoning):
AIME'25: +6.7% absolute improvement with Clip-Higher
AIME'24: +21.5% absolute improvement with Clip-Higher
Consistent 4-8% gains across DAPO, GRPO methods without hyperparameter conflicts
