PaperlayerView exact revision on arXiv

Source-linked AI summary

B-STaR: Monitoring and Balancing Exploration and Exploitation in Self-Taught Reasoners

Weihao Zeng, Yuzhen Huang, Lulu Zhao, Yijun Wang, Zifei Shan, Junxian He

arXiv:2412.17256v2cs.AIcs.CLcs.LG

TL;DR

Self-improvement can reduce reliance on extensive human-curated data, but its drivers and bottlenecks remain poorly understood. B-STAR monitors exploration and exploitation and adapts configurations across iterations. It significantly surpasses other self-improvement methods across mathematical reasoning, coding, and commonsense reasoning.

  • Problem

    Self-improvement’s key drivers and bottlenecks remain poorly understood, despite its potential to reduce reliance on extensive human-curated reasoning data.

    Source: §1 · pp. 1–2

  • Method

    B-STAR monitors exploration and reward exploitation, then adjusts configurations such as sampling temperature and reward thresholds to maximize average balance scores.

    Source: §1 · p. 2

  • Results

    Across mathematical reasoning, coding, and commonsense reasoning, B-STAR significantly surpasses other self-improvement methods and Pass@32 improves without notable degradation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper concludes that balancing exploration and exploitation is important for self-improving training and reports superior performance on reasoning and coding tasks.

    Source: §5 · p. 10

    The analysis focuses on problem-solving domains where response correctness can be verified on labeled datasets.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

In the absence of extensive human-annotated data for complex reasoning tasks, self-improvement -- where models are trained on their own outputs -- has emerged as a primary method for enhancing performance. However, the critical factors underlying the mechanism of these iterative self-improving methods remain poorly understood, such as under what conditions self-improvement is effective, and what are the bottlenecks in the current iterations. In this work, we identify and propose methods to monitor two pivotal factors in this iterative process: (1) the model's ability to generate sufficiently diverse responses (exploration); and (2) the effectiveness of external rewards in distinguishing high-quality candidates from lower-quality ones (exploitation). Using mathematical reasoning as a case study, we begin with a quantitative analysis to track the dynamics of exploration and exploitation, discovering that a model's exploratory capabilities rapidly deteriorate over iterations, and the effectiveness of exploiting external rewards diminishes as well. Motivated by these findings, we introduce B-STaR, a Self-Taught Reasoning framework that autonomously adjusts configurations across iterations to Balance exploration and exploitation, thereby optimizing the self-improving effectiveness based on the current policy model and available rewards. Our experiments on mathematical reasoning, coding, and commonsense reasoning demonstrate that B-STaR not only enhances the model's exploratory capabilities throughout training but also achieves a more effective balance between exploration and exploitation, leading to superior performance.

1 INTRODUCTION

Self-improvement offers a more scalable alternative to scarce human-curated reasoning data, but its drivers and bottlenecks remain poorly understood. B-STAR monitors exploration and exploitation to adapt training configurations, and reports stronger performance across reasoning and coding tasks.

1 INTRODUCTION

  • Scarcity of extensive, high-quality human-curated datasets motivates self-improvement as a more cost-effective and scalable approach to training.Self-improvement iteratively selects model-generated responses to create data for further refinement.Source: §1 · pp. 1–2
  • Self-improvement performance saturates after 3 to 5 iterations, while its key drivers and internal optimization mechanisms remain poorly understood.The paper asks which factors determine self-improvement and how they can guide training.Source: §1 · p. 2
  • Exploration is generating correct, diverse candidates, while exploitation is how effectively external rewards select high-quality solutions; both can stagnate or decline.The paper reports that imbalances between these capabilities can hinder continued improvement.Source: §1 · p. 2
  • B-STAR adaptively adjusts configurations such as sampling temperature and reward thresholds to maximize average balance scores across iterations.Each iteration selects configurations on a small query subset, then applies them to the full training set before updating the model.Source: §1 · p. 2, Figure 2 · p. 3
  • Across mathematical reasoning, coding, and commonsense reasoning, B-STAR significantly surpasses other self-improvement methods; Figure 1 tracks Pass@1 across four benchmarks.The caption identifies GSM8K, MATH, APPS, and ARC-Challenge.Source: §1 · p. 2, Figure 1 · p. 2

2 MONITORING EXPLORATION AND EXPLOITATION IN SELF-IMPROVEMENT

Self-improvement iteratively generates responses, selects them with rewards, and updates the model; its effectiveness depends on balancing exploration and exploitation as they change during training. In the MATH case study, exploration metrics decline after initial gains while Answer + PRM retains diversity, even as Reward@K-S continues improving.

2.1 BACKGROUND: SELF-IMPROVEMENT

  • Self-improvement generates K candidates, scores and selects high-quality responses, then uses the selected data to update the policy model.SFT is commonly used for the improving step in reasoning tasks; with SFT, reward-based filtering is also called rejection fine-tuning.Source: §2.1 · p. 3

2.2 THE CRITICAL FACTORS – EXPLORATION AND EXPLOITATION

  • Self-improvement requires diverse high-quality candidates and a reward that distinguishes them; overfitting can reduce diversity, while excessive diversity can impair reward discrimination.These capabilities move with the policy during training, motivating a dynamic balance between exploration and exploitation.Source: §2.2 · p. 4
  • Exploration is tracked with Pass@K-S, the number of unique correct responses among K samples, and diversity among correct generations; Reward@K-S tracks whether the reward-ranked top S are all correct.Pass@K-S is presented as a more stable exploration proxy than Pass@K.Source: §2.2 · p. 4

2.3 DYNAMICS OF EXPLORATION AND EXPLOITATION – A CASE STUDY IN MATHEMATICAL PROBLEM SOLVING

  • Pass@K-S initially increases then declines, with Pass@K-1 nearing the SFT baseline; Answer + PRM overcomes the declining diversity trend.The authors suggest answer-only filtering can produce homogeneous reasoning paths, while fine-grained PRM rewards encourage high-quality paths and preserve diversity.Source: Figure 3 · p. 5, §2.3 · pp. 5–6
  • The MATH case study uses online RFT with Mistral-7B and compares answer-matching rewards with answer matching plus PRM, using τ = 0.Training ran 4,500 steps in nine 500-step iterations, sampling 32 candidates per query at temperature 1.0.Source: §2.3 · p. 5
  • Reward@K-S continues improving on MATH, potentially because the training data is MATH; the authors hypothesize that stalled exploration bottlenecks self-improvement.The reward remains fixed during training, and exploitation performance is related to the policy model’s changing exploration.Source: Figure 3 · p. 5, §2.3 · p. 6

3 B-STAR – BALANCED SELF-TAUGHT REASONERS

B-STAR seeks to balance exploration and exploitation by maximizing a per-query balance score, then adapts sampling temperature and reward threshold as optimal settings change during training.

3 B-STAR – BALANCED SELF-TAUGHT REASONERS

  • The method’s objective is to capture how exploration and exploitation interact and to adjust configurations that control them as training evolves.The section presents the balance score as a measure of the overall contribution of synthetic data under exploration- and exploitation-related configurations.Source: §3 · p. 6, §3.1 · p. 6

3.1 BALANCE SCORE

  • The balance score reflects both the number of unique correct responses and their proportion among selected responses, so neither quantity alone captures synthetic-data contribution.It discounts queries with fewer than n⋆ correct responses and caps that discount at 1; n⋆ roughly represents the target correct responses per query.Source: §3.1 · pp. 6–7

3.2 DYNAMIC OR STATIC – ON THE CONFIGURATIONS OF EXPLORATION AND EXPLOITATION

  • As training progresses, the balance-maximizing temperature shifts from lower to higher settings, while the preferred reward threshold shifts from higher to lower settings.The temperature pattern is attributed to a shift from limited correctness early to preserving solution diversity later; the threshold is relaxed as the model strengthens.Source: §3.2 · pp. 7–8, Figure 4 · p. 7

3.3 B-STAR

  • B-STAR dynamically adjusts sampling temperature and reward threshold at every iteration to maximize average balance score, using a small query subset and fixing sample size to the budget maximum.The authors report that evaluating 600 MATH queries adds negligible cost; larger sample sizes generally help in preliminary trials.Source: §3.3 · p. 8

4 MAIN EXPERIMENTS

B-STAR is evaluated against self-improvement baselines on mathematical reasoning, coding, and commonsense reasoning. It leads in Pass@1 across the reported scenarios, while its dynamic configuration adjustments support exploration and exploitation through training.

4.1 SETUP

  • The evaluation compares B-STAR with STaR/ResT-EM, Iterative RFT, and Online RFT, using reward settings with and without a reward model.The tasks are mathematical problem-solving, coding, and commonsense reasoning.Source: §4.1 · p. 8

4.2 RESULTS

  • B-STAR achieves higher Pass@1 than the compared self-improvement methods across mathematical reasoning, coding, and commonsense reasoning.The comparisons include Rest-EM, Iterative RFT, online RFT, and reward-model variants.Source: §4.2 · p. 8
  • Online RFT outperforms predominantly offline Rest-EM, while B-STAR shows significantly better Pass@K-S values associated with enhanced exploration.Pass@K-S reflects generation of a wider range of high-quality responses.Source: §4.2 · p. 9
  • Other baselines slow and eventually stagnate after early gains, whereas B-STAR maintains substantial growth and consistently outperforms them.The authors suggest balancing exploration and exploitation is crucial for stable and efficient self-improvement.Source: §4.2 · p. 9
  • Early training uses a lower sampling temperature, around 0.5, and stricter reward thresholds; both settings are relaxed as training advances.The authors describe this schedule as supporting cautious early exploration before broader sampling.Source: §4.2 · p. 10

4.3 ABLATION STUDY AND ADDITIONAL RESULTS

  • Performance significantly deteriorates when either temperature or reward-threshold adjustment is fixed, and even the best grid-searched fixed configurations underperform B-STAR.The ablation identifies dynamic adjustment of both configurations as critical.Source: §4.3 · p. 10
  • B-STAR also generalizes to Llama-3.1-8B, evaluated on MATH, GSM8K, APPS, and ARC-Challenge.The reported setup uses smaller sample sizes for these evaluations because of computational resource constraints.Source: §4.3 · p. 10, Table 4 · p. 10

5 DISCUSSION

The paper argues that balancing exploration and exploitation improves self-improving training and proposes B-STAR as an algorithm for that purpose. Its current approach uses hyperparameter configurations, while more flexible control is left for future work.

5 DISCUSSION

  • The authors identify balancing exploration and exploitation as important to self-improving training and propose B-STAR to improve that balance.They report superior performance on reasoning and coding tasks.Source: §5 · p. 10
  • B-STAR currently manipulates exploration and exploitation through hyperparameter configurations, a simple control approach.Source: §5 · p. 10
  • Advanced decoding to control exploration and reward-model updates to improve exploitation are proposed as future ways to make control more flexible.Source: §5 · p. 10

A.1 EXPERIMENT SETUP

The experiments train on MATH and use Mistral-7B for supervised fine-tuning and process-reward modeling. The setup then uses the fine-tuned model for nine self-improvement iterations.

A.1 EXPERIMENT SETUP

  • Training uses 4,000 MATH test-set problems and 7,500 original training problems, with MATH500 reserved for testing and 500 sampled problems for validation.The model's mathematical reasoning is tested on MATH and GSM8K test sets.Source: §A.1 · p. 14
  • Mistral-7B is SFT-trained for 3 epochs, but the first-epoch checkpoint P0 starts self-improvement, which runs for 9 iterations of 500 training steps each (batch size 128).At each iteration’s start, 32 candidate responses per query are sampled at temperature 1.0.Source: §A.1 · p. 14
  • The process reward model is trained on about 270K process annotations, and responses are selected when their final reward score exceeds 0.0.The final reward combines the normalized lowest solution-step score with a sparse reward.Source: §A.1 · p. 14

A.2 RESULTS ON GSM8K

On GSM8K, adding process reward models yields additional improvements, though the gains are less pronounced on MATH. The appendix also describes B-STAR's iterative configuration, sampling, reward, and policy-update procedure.

A.2 RESULTS ON GSM8K

  • Incorporating process reward models brings additional improvements on GSM8K, but the improvements are less pronounced on MATH.The authors hypothesize that MATH difficulty may exceed the discriminative capability of the 7B reward model.Source: §A.2 · p. 14
  • At each iteration, B-STAR selects temperature and reward threshold settings by maximizing the average balance score.The algorithm then draws queries from the training dataset and samples candidate responses from the current policy using the selected temperature.Source: §A.2 · p. 15
  • A reward model scores candidates, responses above the selected threshold form the update dataset, and the policy is updated while optimizer and scheduler state are inherited.Source: §A.2 · p. 15

B EXPERIMENT SETUP FOR MAIN EXPERIMENTS

The experiments compare iterative self-improvement baselines and evaluate B-STAR in mathematical problem-solving, coding, and commonsense reasoning. Each domain uses its own dataset, base model, and reward setup.

B.1 BASELINE

  • STaR/ResT-EM retrains from scratch each iteration, Iterative RFT resumes from the previous checkpoint, and Online RFT also carries over optimizer state and the learning-rate schedule.Source: §B.1 · p. 15

B.2 MATHEMATICAL PROBLEM-SOLVING

  • Math experiments use Mistral-7B after three SFT epochs on MATH, with 67,500 samples and 11,500 queries per iteration and 64 responses per query.Temperature and reward threshold vary from 0.5 to 1.2 and -1.0 to 1.0, respectively, in 0.1 increments; Pass@1, Pass@K-S, and Reward@K-S track performance and exploration and exploitation.Source: §B.2 · p. 15

B.3 CODING CHALLENGES

  • Coding experiments use Llama-3-8B on APPS, with 13,500 samples and 2,627 queries per iteration; unit tests provide binary rewards, so B-STAR adjusts temperature only.The training set contains 13K examples sampled at five responses per original question; Mistral-7B was excluded because it performed poorly on coding tasks.Source: §B.3 · pp. 15–16

B.4 COMMONSENSE REASONING

  • Commonsense experiments use Mistral-7B-Instruct on ARC-Challenge without SFT, with ground-truth answers as binary rewards and Pass@1 as the sole reported metric.The multiple-choice response space makes Pass@K and Pass@K-S uninformative for this dataset.Source: §B.4 · p. 16

C DETAILS OF SELF-IMPROVEMENT

This section defines reward functions and rejection-sampling fine-tuning, then compares iterative baselines and reports B-STAR accuracy on GSM8K and MATH.

  • A fixed reward is static, while a trained reward model adapts during training and assigns continuous scores learned from supervision.Binary correctness or passing a unit test exemplifies fixed feedback; output-based and process-based reward models exemplify learned rewards.Source: C DETAILS OF SELF-IMPROVEMENT · p. 16
  • STaR/ResT-EM restarts training from scratch, Iterative RFT inherits the prior checkpoint, and Online RFT additionally inherits optimizer state and the learning-rate schedule.Source: C DETAILS OF SELF-IMPROVEMENT · p. 16
  • B-STAR’s GSM8K accuracy is 53.1%, while its MATH accuracy is 27.8%.Source: Figure 8 · p. 17

D MORE FINE-GRAINED CONFIGURATION ADJUSTMENTS

Finer-grained temperature and reward-threshold adjustments produce more dynamic configuration changes and a higher balance score.

D MORE FINE-GRAINED CONFIGURATION ADJUSTMENTS

E IMPACT OF FIXED CONFIGURATION COMBINATIONS

Fixed online-RFT configurations remain less effective than B-STAR’s dynamic adjustments. The section relates this approach to prior dynamic-hyperparameter work and frames binary-reward self-improvement as data selection and supervised training.

E IMPACT OF FIXED CONFIGURATION COMBINATIONS

  • Grid-searched fixed configurations improve Online RFT somewhat, but remain less effective than B-STAR’s dynamic configuration adjustments.The comparison is used to emphasize the need to balance exploration and exploitation dynamically during training.Source: E IMPACT OF FIXED CONFIGURATION… · p. 17

F RELATED WORK OF DYNAMIC HYPERPARAMETER ADJUSTMENT

  • Related work spans dynamic learning-rate adjustment, hyperparameter optimization, model populations, sequential tuning, meta-learning, and adaptive hyperparameters in reinforcement learning.Source: F RELATED WORK OF DYNAMIC… · p. 18

G THEORETICAL JUSTIFICATION FOR EXPLORATION AND EXPLOITATION

  • The paper frames self-improvement as reinforcement learning, with the current policy sampling responses and a reward function evaluating them.The formulation identifies the dataset, sampled input and ground-truth answer, sampled response, policy at iteration t, and reward.Source: G THEORETICAL JUSTIFICATION FOR… · p. 18
  • With binary rewards, the reinforcement-learning objective becomes data selection and supervised training; exploration samples candidates, while exploitation selects using rewards.The text says insufficient exploration may cause stagnation, while insufficient exploitation can lead to instability and large variance.Source: G THEORETICAL JUSTIFICATION FOR… · p. 18

Related papers

Loading 2412.17256v2…