Source-linked AI summary

Outcome-based Exploration for LLM Reasoning

Yuda Song, Julia Kempe, Remi Munos

arXiv:2509.06941v1cs.LGcs.CL

TL;DR

LLM reasoning RL can improve correctness while causing diversity collapse, including on unsolved training questions, which matters for test-time scaling. The paper analyzes this transfer, exploits the tractable outcome space with historical and batch exploration, and reports improved accuracy–diversity tradeoffs. It also formalizes outcome-based exploration through outcome-based bandits, with limitations in verifiable, tractable-outcome, single-turn settings.

  • Problem

    Outcome-based RL improves final-answer correctness but systematically reduces generation diversity, limiting a property valuable for test-time scaling.

  • Method

    The paper analyzes RL as sampling and adapts exploration bonuses to final outcomes through historical UCB-style exploration, refined positive and negative signals, and batch repetition penalties.

  • Results

    The proposed methods improve test performance across pass@k metrics and achieve a better accuracy–diversity tradeoff at test time.

  • Takeaways & Limitations

    Outcome-based exploration offers a practical path toward RL reasoning methods that improve accuracy while mitigating diversity collapse.

  • Takeaways & Limitations

    The algorithms currently apply only to verifiable domains with tractable outcome spaces, and evaluation is limited to single-turn benchmarks.

Abstract

from arXiv · show

Reinforcement learning (RL) has emerged as a powerful method for improving the reasoning abilities of large language models (LLMs). Outcome-based RL, which rewards policies solely for the correctness of the final answer, yields substantial accuracy gains but also induces a systematic loss in generation diversity. This collapse undermines real-world performance, where diversity is critical for test-time scaling. We analyze this phenomenon by viewing RL post-training as a sampling process and show that, strikingly, RL can reduce effective diversity even on the training set relative to the base model. Our study highlights two central findings: (i) a transfer of diversity degradation, where reduced diversity on solved problems propagates to unsolved ones, and (ii) the tractability of the outcome space, since reasoning tasks admit only a limited set of distinct answers. Motivated by these insights, we propose outcome-based exploration, which assigns exploration bonuses according to final outcomes. We introduce two complementary algorithms: historical exploration, which encourages rarely observed answers via UCB-style bonuses, and batch exploration, which penalizes within-batch repetition to promote test-time diversity. Experiments on standard competition math with Llama and Qwen models demonstrate that both methods improve accuracy while mitigating diversity collapse. On the theoretical side, we formalize the benefit of outcome-based exploration through a new model of outcome-based bandits. Together, these contributions chart a practical path toward RL methods that enhance reasoning without sacrificing the diversity essential for scalable deployment.

1 Introduction

Outcome-based RL improves LLM reasoning but systematically reduces generation diversity, including on unsolved training questions. The paper analyzes this transfer and proposes historical and batch outcome-based exploration to improve the accuracy–diversity tradeoff.

  • 1 Introduction: Outcome-based RL rewards final-answer correctness effectively but causes systematic loss of generation diversity.Lower diversity matters because diverse samples can amplify performance through test-time scaling.
  • 1 Introduction: Diversity degradation transfers from solved to unsolved questions, reducing the set of candidate answers during training.The paper frames this as concentration on solved questions propagating reduced diversity to unsolved ones.
  • 1 Introduction: Outcome-based exploration adapts UCB-style bonuses to final answers, avoiding intractable token-level exploration over exponentially many sequences.The approach can incorporate both positive and negative outcome signals.
  • 1 Introduction: Historical exploration diversifies training visits, whereas batch exploration penalizes within-batch repetition to promote test-time diversity.The two forms address different objectives: historical exploration improves pass@1, while batch exploration improves pass@k.
  • 1 Introduction: The proposed methods are designed to enhance reasoning without sacrificing diversity needed for scalable deployment.The contributions include refined outcome-based exploration, theoretical analysis, and batch exploration.

2 Diversity Degradation: RL as Sampling

The paper analyzes RL post-training as a sampling process and finds that diversity degradation emerges during training, not only in final-model test behavior. RL initially solves questions faster, but later underperforms base-model sampling as reduced diversity transfers to unsolved questions; final answers provide a tractable diversity proxy.

  • 2.1 Preliminaries: The setup models an LLM as a policy generating reasoning traces and final answers, with verifiable rewards determined by final-answer correctness.The experiments use fully on-policy GRPO as the vanilla RL baseline and define reasoning-trace and answer supports for each question.
  • 2.2 Diversity Degradation during RL training: Experiments compare Llama-3.1-8B-Instruct and Qwen-2.5-7B-Base on easy and medium math datasets using solved-question counts and distinct-answer counts.Training epochs are aligned with base-model sampling through k = nt; Figure 2 uses n = 16 for pass@k and n = 8 for diff@k comparisons.
  • 2.2 Diversity Degradation during RL training: RL initially solves questions faster than the base model but eventually solves fewer questions with the same number of samples.As training continues, the RL question-solving rate decreases faster than the base model’s rate.
  • 2.2 Diversity Degradation during RL training: Reduced diversity on solved questions transfers to unsolved questions, causing RL to produce fewer distinct answers there than the base model.The paper attributes this transfer to probability concentration on previously solved questions and tracks cumulative distinct answers across all and unsolved questions.
  • 2.2 Diversity Degradation during RL training: Final answers make diversity tractable in verifiable domains, with the base model averaging fewer than 50 possible answers per question under a large sample budget.Token-space diversity is intractable because it entails exponentially many candidates, whereas final answers provide an operational proxy.

3 Outcome-based Exploration

Outcome-based exploration adapts exploration bonuses to the bounded space of final answers, while distinguishing historical exploration during training from batch exploration at test time. UCB variants improve training and test behavior to different degrees, with baseline-based methods offering stronger generalization and batch exploration preserving more large-k diversity.

  • Historical Exploration via UCB: Outcome-based exploration searches over distinct final answers while retaining the correctness reward, using answer visitation counts to construct UCB bonuses.The bonus depends on how often answer a has been sampled for question x; masking the final answer during policy updates prevents bonus hacking.
  • Historical Exploration via UCB: Naive UCB consistently improves training performance but does not reliably improve test performance across models and datasets.A significant test improvement is observed only on the easy dataset with Llama-3.1 8B.
  • UCB with a Baseline: UCB-Mean uses the batch mean of UCB bonuses as a baseline, rewarding underrepresented answers and penalizing answers that appear more often within the batch.This produces exploration signals that can be positive or negative relative to current-batch frequency.
  • UCB with a Baseline: UCB-Con replaces the batch mean with a tunable constant baseline, controlling the balance between positive and negative exploration signals.With b0 = 0.5, answers visited fewer than 4 times receive positive exploration signals, while more frequently visited answers receive negative signals.
  • UCB with a Baseline: Global exploration is not expected to maintain high pass@k for large k near training completion, because exploration can favor repeatedly generating a least-visited answer.In some settings, final pass@8, pass@16, and pass@32 remain similar to vanilla RL.
  • Batch Exploration: Batch exploration directly encourages diverse outputs within each batch, achieving similar peak pass@k to UCB-Con while generally providing better large-k diversity at the end of training.Its training metrics can be worse because the objective targets test-time accuracy-diversity tradeoffs rather than training exploration metrics.

4 Additional Analysis

The analysis distinguishes historical from batch exploration and examines their diversity dynamics, then develops an outcome-based bandit theory explaining when outcome-level exploration can reduce complexity.

  • Historical and batch exploration: Historical exploration solves more questions and accumulates more diverse answers over time than batch exploration.These metrics are inherently historical, so historical exploration is superior under this comparison.
  • Historical and batch exploration: Batch consistently produces more distinct answers per batch than UCB-Con because it directly optimizes batch diversity.Table 2 measures distinct answers in batches of size 8, including separate statistics for solved and unsolved questions.
  • Historical and batch exploration: Historical and batch exploration can exhibit distinct pathologies, so neither is guaranteed to substitute for the other.Historical exploration may become nearly deterministic, while batch exploration may cycle through a small answer set without improving training dynamics.
  • Outcome-based bandits: The bandit model contains K arms mapped to m much smaller outcomes, but outcome partitioning alone cannot guarantee regret depending only on m.Without generalization across arms sharing an outcome, some instances retain regret lower bounds of Ω(min{T, K}).
  • Outcome-based bandits: With generalization across arms yielding the same outcome, UCB-style exploration over outcomes recovers an m-armed-bandit regret dependence.This assumption supplies the missing bridge from trace-level arms to the smaller outcome space.

5 Related Work

Related work connects this paper to RL post-training, diversity-preserving exploration in preference and reasoning settings, and classical bandit exploration methods.

  • RL post-training: RL post-training improves LLM objectives such as human preference and reasoning, while final-answer rewards can improve reasoning accuracy without intermediate rewards.Related work also reports diversity losses after post-training in non-reasoning domains.
  • Exploration in RL: Preference fine-tuning work explores base-model likelihood bonuses, diversity-based data ranking, and theoretical guarantees under linear assumptions.These approaches address exploration during preference optimization rather than the outcome-space formulation developed here.
  • Exploration in RL: Reasoning-domain exploration methods include Random Network Distillation, entropy objectives, and pass@k-based training objectives.These methods encourage exploration through trace novelty, token entropy, or direct test-time diversity objectives.
  • Bandit exploration: Classical count-based and Thompson-sampling methods provide exploration guarantees across tabular and linear settings, while broader function-approximation methods can lack computational efficiency.The paper’s historical exploration shares the bonus-based exploration principle used in much of this literature.

6 Conclusion and Discussion

The paper analyzes diversity degradation in reasoning post-training and proposes outcome-based historical and batch exploration to improve accuracy–diversity tradeoffs.

  • Contributions: The study identifies transfer of diversity degradation and tractability of the outcome space as central phenomena in verifiable reasoning tasks.It analyzes RL post-training as sampling and uses final outcomes as the exploration target.
  • Contributions: Outcome-based UCB exploration, refined with positive and negative signals, improves pass@k test performance across k.The approach adapts classical exploration to the outcome space rather than the exponentially large trace space.
  • Contributions: Batch exploration improves the test-time accuracy–diversity tradeoff by encouraging variation among outputs within a batch.The paper also analyzes how historical and batch exploration interact and concludes they are not mutually exclusive.
  • Limitations: The algorithms apply only to verifiable domains with tractable outcome spaces and are evaluated only on single-turn benchmarks.Extending them to more general settings and studying multi-turn settings remain future directions.

A Theoretical Results

The theoretical model represents reasoning traces as many arms grouped by a smaller set of final-answer outcomes, with rewards determined by outcomes.

  • Arms, outcomes, and partition: The bandit has K arms and m much smaller outcomes, with an unknown mapping assigning each arm to one outcome.The outcome partitions are mutually exclusive and may be imbalanced or balanced.
  • Reward: Rewards depend only on the selected outcome and are stochastic, with outcome means determining the optimal outcome.The model assumes rewards lie in [0, 1] and are 1-sub-Gaussian around their outcome means.
  • Interaction protocol: At each round, a policy selects an arm from the history, after which the environment reveals its outcome and draws a reward with that outcome’s mean.The history records previously selected arms, outcomes, and rewards.
  • Performance metrics: The framework measures pseudo-regret over T rounds and defines outcome discovery times to quantify delays before first encountering outcomes.Discovery times capture delays under the no-generalization constraint.
  • LLM reasoning interpretation: Each reasoning trace corresponds to an arm, its final answer to an outcome, and its verifiable correctness to the outcome-dependent reward.This interpretation explains why the outcome space is smaller than the trace space.

A.2 Lower Bound

The lower-bound analysis shows that outcome-based bandits can incur regret polynomial in the number of arms when the optimal outcome class is small. The proof uses random placement of optimal arms and lower-bounds the delay before discovering one.

  • Lower-bound result: Polynomial regret in K is unavoidable in the worst case without additional assumptions.The theorem considers a partitioned arm set with a unique optimal outcome and establishes a lower bound for any algorithm.
  • Hard instance: The hard instance has an optimal outcome with s⋆ arms, while every other outcome has mean reward 1/2 and the optimal outcome has mean reward 1/2 + ∆.The construction places the s⋆ optimal arms uniformly among K arm indices.
  • Discovery delay: Before the first pull from the optimal class, every observed reward is ∆-suboptimal in expectation.The proof denotes the first optimal-class pull by τ⋆ and charges regret for the preceding rounds.
  • Constant factor: A universal constant c0 > 0 appears in the lower bound, with one example given as c0 = e^-1/2/4.The resulting fixed instance is obtained from the random-placement argument.

A.3 Balanced Partitions

The balanced-partition analysis shows that outcome-level exploration can avoid dependence on the total number of arms when outcome classes are balanced, while generalization determines how efficiently outcomes are discovered. The resulting algorithms combine discovery with UCB over observed outcome representatives.

  • Balanced partitions: Balanced partitions allow an algorithm whose regret is independent of K.Under the balanced condition |A_o| = K/m, the analysis separates discovery from post-discovery bandit learning.
  • Balanced Outcome UCB: Balanced Outcome UCB uniformly probes unseen arms until all outcomes are represented, then applies outcome-level UCB to representative arms.The algorithm maintains discovered outcomes, representative arms, and a pool of unseen arms.
  • Regret decomposition: After discovery, the problem reduces exactly to an m-armed stochastic bandit because rewards depend only on outcomes.The total regret is decomposed into discovery and post-discovery contributions, with the latter bounded using standard UCB analysis.
  • Strong generalization: Strong generalization lets the learner identify and exclude an entire outcome class after observing one outcome, making discovery take at most m fresh probes.The corresponding algorithm stores one representative per outcome and removes its class from the fresh-probe pool.
  • Soft generalization: Soft generalization permits excluding only a fraction ρ_o of an observed outcome class, interpolating between strong and no generalization.Algorithm 3 uses soft exclusion while retaining outcome-level UCB after discovery.
  • Soft-generalization analysis: The soft-generalization bound depends on the aggregate exclusion parameter ρ̄, which summarizes exclusion across non-optimal outcomes.The analysis first bounds the time to discover the optimal outcome under this partial exclusion model.

B Additional Experiment on Hard Dataset

On the hard dataset, outcome-based exploration improves both training accuracy and the number of unique solutions explored relative to the baseline. The experiments compare UCB-Con, Batch, UCB variants, and GRPO using training and pass@k test metrics.

  • Training performance: The training comparisons plot solved-question fraction and the number of different answers sampled over gradient updates.The hard-dataset comparison includes UCB-Con, Batch, and GRPO with three random seeds and mean performance.
  • Training performance: Outcome-based exploration improves the number of questions solved and unique solutions explored on the hard dataset.For Llama-3.1-8B-Instruct, the hard dataset contains 996 questions, and the improvement is reported as larger than on the base model.
  • UCB variants: Additional training figures compare UCB variants with GRPO on medium and easy datasets using solved-question fraction and answer diversity.For unsolved questions, the figures count different answers after historically sampling one correct answer.
  • Test performance: Test comparisons report pass@k for k ∈ {1, 2, 4, 8, 16, 32} using 32 evaluation samples per question.The comparisons use three random seeds and plot mean performance at regular training checkpoints.

D Quantitative Results

The quantitative tables compare baseline methods at the best and final checkpoints using pass@1 and pass@32 across three random seeds. UCB-Con generally has the strongest peak performance, while Batch generally has the strongest final pass@32 performance.

  • Best checkpoint: UCB-Con generally achieves the best peak performance at the best checkpoint.Table 3 reports pass@1 and pass@32 means and standard deviations over three random seeds.
  • Final checkpoint: Batch generally achieves the best final pass@32 performance at the final checkpoint.Table 4 compares pass@1 and pass@32 means and standard deviations over three random seeds.

E Implementation Details

The experiments use fixed exploration hyperparameters and compare strategies by the number of different answers sampled in a batch.

  • Llama experiments use bonus coefficient c = 0.1, while Qwen experiments use c = 0.2.
  • UCB-Con uses b0 = 1 for the easy dataset and b0 = 0.5 for the medium dataset.
  • Table 5 compares exploration strategies by the number of different answers sampled in a batch.Results are reported as means with standard deviations across two random seeds.
Loading 2509.06941v1…