Correlation-Aware Exploration for Long-Trace Reasoning

AI summary

Overview

Research area: Machine learning / reinforcement learning for large language model reasoning (RLVR, GRPO, exploration and entropy regularization).

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods (PPO/GRPO), policy entropy, and information theory (mutual information, conditional entropy).

Scope: The paper diagnoses why per-token entropy bonuses destabilize long-chain-of-thought RL training and proposes Info-GRPO, a latent-seed-conditioned mutual-information objective that relocates exploration from the token distribution to the trajectory distribution.

What This Paper Is About

Reasoning models trained with Group Relative Policy Optimization (GRPO) produce chains of thought thousands of tokens long, but such training tends to collapse onto a narrow set of strategies. The standard fix, an entropy bonus, pushes probability mass toward the vast low-probability tail of a ~152K-token vocabulary at every single generated token, which fights the entropy decay that convergence requires and makes its coefficient hard to tune. The authors propose Info-GRPO, which instead conditions the policy on a random latent seed during training and maximizes a generalized mutual information between the seed and the resulting reasoning trace, so that diversity lives between whole trajectories while each trajectory stays confident.

Key Contributions

1. A structural diagnosis of entropy regularization in long traces. The paper shows that the entropy gradient attaches a scalar factor of magnitude at least |1 + log δ| to each of at least V − 1/δ tokens in the low-probability tail, yielding a "Total Tail Instability" (TTI) lower bound of (V − 1/δ)|1 + log δ|, which is applied once per token and thus accumulates as O(T · TTI) over a trace of length T.
2. The Info-GRPO objective. A correlation-aware training paradigm that augments prompts with latent seeds and maximizes α·H(π(𝒯|s₀)) − H(π(𝒯|s₀,z)), a generalized mutual-information regularizer in which the consolidation term drives each seed-conditioned policy toward confidence and a relaxation coefficient α ∈ [0,1] retains a controlled amount of marginal-entropy pressure.
3. A practical, deployment-neutral implementation. Training splits each prompt's G = 16 rollouts into two equal sub-groups (unseeded and seed-augmented), with advantages estimated over the merged set; the latent seed is used only during training, so inference runs on the original prompt unchanged.
4. Broad empirical validation. Consistent gains over GRPO and entropy-regularized baselines across four backbones and six benchmarks, plus a controlled tail-masking ablation that isolates the vocabulary tail as the locus of degradation.

Main Findings

• AIME24 gains over GRPO: Info-GRPO improves Avg@8 over GRPO by 3.75%, 1.66%, 4.16%, and 4.17% on the four backbones respectively.
• Average gains scale with backbone capacity: Overall average over GRPO improves by 4.01%, 3.54%, and 2.89% on the three stronger backbones (Qwen2.5-Math-7B: 45.04 → 49.05; DeepSeek-R1-Distill-Qwen-7B: 62.45 → 65.99; DeepSeek-R1-Distill-Llama-8B: 58.73 → 61.62).
• A boundary condition on the weakest backbone: On vanilla Qwen2.5-7B, Info-GRPO (41.42%) is nearly tied with GRPO (41.22%); the authors argue correlation-aware exploration only converts diversity into accuracy if the backbone can already produce valid reasoning paths to diversify over.
• The vocabulary tail is the culprit: With only the entropy scope varied on DeepSeek-R1-Distill-Qwen-7B, the full-vocabulary bonus drops AIME24 below no-entropy GRPO (57.92 → 53.23), whereas restricting the same objective to head tokens restores it (58.13). Seed-augmented MI reaches 62.08.
• Exact mutual information is not the practical setting: At α = 1.0 (Eq. 5 coincides with MI), AIME24 falls to 52.50 and AMC to 75.90, below vanilla GRPO; α = 0.5 is best (62.08 AIME24, 83.13 AMC), while α = 0.05 is too weak (58.33, 81.93).
• Gains concentrate in the long-trace regime: Increasing the trace budget from 2K to 8K improves AIME24 from 56.67 to 62.08 and AIME25 monotonically from 37.08 to 45.83.
• The MI regularizer, not the seed, does the work: Standard GRPO with seed-augmented prompts but no MI term degrades Avg@8 from 45.04% to 42.23% on Qwen2.5-Math-7B; full Info-GRPO reaches 49.05%.
• Neutral seeds beat semantic ones: Digital seeds achieve AIME24 25.00 with diversity 0.300; semantic heuristics achieve 23.33 with diversity 0.307, exposing an "alignment tax" from explicit linguistic constraints.
• Small λ is best: λ = 0.005 gives 62.08 / 45.83 (AIME24 / AIME25), versus 0.05 (58.75 / 40.42), 0.002 (62.08 / 44.17), and 0.5 (57.08 / 39.58).
• Test-time compute efficiency: Under majority voting on AIME24 with DeepSeek-R1-Distill-Qwen-7B, Info-GRPO reaches 80.00 at N = 8, the accuracy GRPO attains only at N = 64 — an 8× smaller budget. On Qwen2.5-Math-7B, GRPO does not reach Info-GRPO's peak of 30.00 even at N = 128.
• Code transfer: On LiveCodeBench with DeepSeek-R1-Distill-Qwen-7B, Info-GRPO improves Pass@1 from 37.63 to 40.86, Pass@8 from 50.18 to 51.97, and Avg@8 from 38.13 to 40.14.
• Diversity profile matches the design goal: Different seeds produce markedly higher inter-group semantic diversity and lower intra-group diversity, and the seed–trajectory correlation appears as a block-diagonal structure; gains stabilize around 10 seeds.
• Faster convergence: Info-GRPO shows faster entropy reduction and steeper training-reward growth than GRPO despite applying a decreasing conditional-entropy term.

Methodology in Plain English

The team starts from a simple observation about why the usual entropy bonus misbehaves. Entropy maximization pushes probability away from the few tokens the model likes and toward everything else. With 152,064 tokens in the vocabulary — of which, in Qwen2.5-7B, fewer than five exceed a 10⁻⁶ probability — almost all of that pressure lands on tokens that carry essentially no mass, and in a long trace the bonus is applied at every single step, so the wasted pressure compounds with both trace length and vocabulary size. They formalize this as a bound on "Total Tail Instability" and then verify it experimentally by simply removing the tail tokens from the entropy sum, which flips a degradation into an improvement.

Their fix keeps the GRPO loop intact but changes what is randomized. Half of each prompt's rollouts are generated normally; the other half get a short prefix telling the model to "choose random seed {z}", with z drawn uniformly from a small digital set such as {1,...,10}. Advantages are computed over the combined pool, exactly as in GRPO. On top of this they add a regularizer with two opposing terms: a mild marginal-entropy bonus on the unseeded rollouts (weighted by α = 0.5, an "inference-aligned anchor" since test time uses unseeded prompts) and a strong negative conditional entropy on the seeded rollouts, which pushes each seed toward its own confident solution style. Because the seeds are only a training-time device, deployment requires no changes at all.

Evaluation spans four backbones (two Qwen2.5 models, two DeepSeek-R1-Distill models) trained on DAPO-Math-17K (17,917 problems, 4,096-token budget) or the 48,371-sample set of He et al. (8,192-token budget), with a learning rate of 1e-6, batch size 128, G = 16 rollouts at temperature 1.0, λ = 0.005, and α = 0.5 on a single node of 8 NVIDIA B200 GPUs. Accuracy is reported as Pass@1, Pass@8, and Avg@8 over 8 trials at temperature 0.6 on AIME24 (30 problems), AIME25 (30), AMC (83), MATH500 (500), Minerva (272), and LiveCodeBench (279); diversity is measured as average pairwise Euclidean distance between JINA-v2 embeddings of reasoning paths.

Why This Matters

Impact on research. The paper reframes exploration in RLVR as a correlation problem rather than a confidence problem, which resolves the long-standing tension between entropy bonuses and convergence in a principled way. The TTI diagnostic gives the community a way to reason about entropy regularization as a function of vocabulary size and trace length, and the tail-masking ablation provides a clean causal test that other methods can reuse. Relocating exploration to a latent variable also opens a path to steering and analyzing cognitive strategies rather than merely sampling more.

Real-world applications:

• Competition mathematics and other verifiable-reward domains where majority voting over multiple samples is used at inference — the 8× sample reduction directly lowers serving cost.
• Code generation assistants, where the LiveCodeBench results show the method transfers beyond mathematics.
• Any deployed reasoning model where inference must remain identical to the original prompt — the training-only latent seed means no retooling of the serving stack.
• Analyzing or auditing model behavior, since distinct seeds produce interpretable, distinguishable reasoning strategies that can be inspected separately.

Industry relevance. Both the majority-vote efficiency result and the "no inference-time change" property matter commercially: they reduce test-time compute without altering the deployment surface. The authors are at JD.com and Microsoft, and the method's compatibility with existing GRPO pipelines means it can be dropped into established RLVR training stacks.

Future Directions

• Trace lengths beyond 8K. The longest budget tested is 8,192 tokens (4,096 for the Qwen2.5 series), set by the official backbone configurations. Since the misallocated pressure is predicted to grow as O(T · TTI), the effect should be at least as pronounced at 32K or 128K, but the authors have not verified this due to compute constraints.
• Isolating what the seed contributes. The ablations separate the MI regularizer from prompt augmentation, but not the effect of varying the seed from that of the conditional-entropy term — a matched control holding the seed fixed across all augmented prompts would measure the latter. Seed scheduling over training and per-problem seed selection are also unexplored.
• Richer latent spaces. The current design uses a flat, unstructured discrete latent space; hierarchical or continuous latents allowing compositional control over reasoning are flagged as a substantial standalone problem.
• Broader evaluation and adaptive modulation. Evaluation is restricted to mathematics and code to avoid noisy or hackable reward models, leaving multi-step tool use, varying reward sparsity, and open-ended reasoning untested. Adaptive modulation of the exploration term for backbones that cannot produce valid paths to diversify over (as on Qwen2.5-7B) is left to future work.

Target Audience

Researchers and engineers working on RL post-training for language models — particularly those implementing or tuning GRPO-style pipelines — will get the most from this paper, along with information-theory-inclined readers interested in mutual-information objectives for structured exploration. Practitioners deploying reasoning models with majority-vote sampling will find the efficiency results actionable. Readers without a background in policy gradients, entropy, and conditional entropy will need to work through the background and framework sections carefully, as the argument is developed formally through a proposition, two theorems, and three remarks.

Authors’ abstract

Reasoning models trained with Group Relative Policy Optimization (GRPO) now produce chains of thought thousands of tokens long, and recent studies reveal that such training is prone to policy collapse. Entropy regularization is the standard remedy, but it works by pushing probability mass away from the few tokens that carry it and toward the vast low-probability tail of the vocabulary, and it applies this force once for every token of the reasoning trace. This puts the bonus in direct conflict with the entropy decay that convergence requires, and makes its coefficient difficult to set in the long-trace regime. To resolve this tension, we propose Info-GRPO, an information-theoretic framework that cultivates correlation between the policy and a latent prior. Rather than raising entropy everywhere, Info-GRPO builds on the mutual information between a latent variable and the trajectory it induces: augmenting prompts with latent seeds during training lets the model explore diverse policies correlated with the prior, while the seed-conditioned entropy is guided toward convergence. Diversity is therefore carried across whole reasoning traces rather than by per-token randomness, and exploration no longer fights entropy reduction. Experiments show that Info-GRPO outperforms vanilla GRPO and entropy-regularized GRPO across diverse reasoning benchmarks. On AIME24, it improves $Avg@8$ over GRPO by 3.75%, 1.66%, 4.16%, and 4.17% on four backbones, and it reaches the same majority-vote accuracy with $8\times$ fewer samples. Code and models will be publicly released.