PaperlayerView exact revision on arXiv

Source-linked AI summary

Mid-Training with Self-Generated Data Improves Reinforcement Learning in Language Models

Aswin RRV, Jacob Dineen, Divij Handa, Mihir Parmar, Ben Zhou, Swaroop Mishra, Chitta Baral

arXiv:2605.08472v1cs.AI

TL;DR

RL may struggle when training data exposes only a narrow range of valid reasoning approaches. This paper mid-trains models on diverse self-generated solutions guided by Pólya’s heuristics before RL, improving subsequent RL performance across mathematical reasoning benchmarks, especially at higher pass@k values.

  • Problem

    Limited exposure to diverse valid reasoning approaches may constrain how effectively subsequent reinforcement learning improves language-model reasoning.

  • Method

    The authors fine-tune models on multiple correct, self-generated solution trajectories per question, guided by Pólya’s problem-solving heuristics, before reinforcement learning.

  • Results

    Across benchmarks, subsequent RL performs better with diverse mid-training, reaching 48.09% average pass@64 versus 44.21% for vanilla RL.

  • Takeaways & Limitations

    Learning multiple problem-solving approaches through self-generated data improves subsequent RL performance, particularly when evaluated at higher pass@k values.

  • Takeaways & Limitations

    The study leaves unresolved whether the observed behaviors are genuinely emergent or primarily compositions of behaviors learned during pre-training.

Abstract

from arXiv · show

The effectiveness of Reinforcement Learning (RL) in Large Language Models (LLMs) depends on the nature and diversity of the data used before and during RL. In particular, reasoning problems can often be approached in multiple ways that rely on different forms of reasoning, and exposure to only a limited range of such approaches in the training data may limit the effectiveness of RL. Motivated by this, we investigate using diverse self-generated data during mid-training as an intermediate step before RL training. Specifically, we adopt a bootstrapped data-generation framework guided by George Polya's problem-solving approaches for generating multiple variants of correct answers for each question in the training data, and then perform fine-tuning. We first provide a theoretical perspective on how mid-training on such data improves RL and explain how policy-gradient updates can incentivize combining multiple approaches. We then empirically demonstrate that RL-trained models initialized with our mid-training data achieve consistent improvements across various mathematical reasoning benchmarks and other OOD tasks like code generation and narrative reasoning. Overall, our investigative study shows that a language model learning multiple problem-solving approaches, through self-generated data helps subsequent RL.

1 Introduction

The introduction argues that RL benefits from diverse reasoning data and investigates self-generated mid-training as a way to improve subsequent RL. The study combines theoretical analysis with experiments showing improved reasoning performance and compositional strategy use.

  • RL improves LLM performance on complex reasoning tasks by encouraging longer, inference-time-scaled reasoning chains with multi-step reasoning, self-verification, and self-correction.
  • RL may fail to elicit self-reflective behaviors or improve beyond the base model without strong priors and suitable task and training-data conditions.
  • The method generates multiple diverse, correct solutions per question from the same base model and fine-tunes on them as a self-improvement mid-training stage.
  • GRPO-based RL consistently improves mid-trained models across mathematical reasoning benchmarks, with larger gains as mid-training introduces more diverse reasoning approaches.
  • At pass@64, mid-training improves vanilla RL by 2.85% on MATH-500 and 5.7% on AIME 2024.
  • Theoretical analysis explains how multiple approaches help RL incentivize combining reasoning strategies within a response, a compositional behavior observed empirically.

2 Related Works

Prior work studies LLM-generated synthetic data for instruction tuning and reasoning, as well as reinforcement-learning methods for preference alignment and skill composition. This work differs by mid-training on diverse self-generated responses using Pólya-style heuristics and analyzing how subsequent RL combines these approaches.

  • Synthetic Data Generation: Synthetic-data research uses LLMs to generate datasets without human annotation, including instruction-response bootstrapping and distillation for smaller models.A pre-trained model can iteratively generate instruction-response pairs to improve instruction following, while other work distills synthetic data from larger models to improve smaller-model reasoning.
  • Reinforcement Learning for LLMs: RLHF aligns LLMs with human preferences through reward modeling and policy optimization, while DPO reframes preference alignment as supervised learning.Prior work also shows that reinforcement learning can enable models to learn novel compositions of atomic skills not explicitly taught during training.
  • Novelty: This work differs by mid-training on diverse self-generated responses produced with Pólya-style problem-solving heuristics, then theoretically analyzing how subsequent RL combines these approaches.The distinction centers on using self-generated data before RL rather than only applying post-training alignment methods.

3 Preliminaries

This section defines supervised fine-tuning as likelihood-based learning from question–response pairs and reinforcement learning as autoregressive policy optimization over generated responses. It also identifies GRPO as the policy-gradient method used in the experiments.

  • Supervised Fine-Tuning (SFT): SFT trains a pretrained language model on question–response pairs by minimizing the responses’ negative log-likelihood.The dataset is denoted D_SFT = {(x, y)}, with each response represented as a token sequence.
  • RL via Policy Gradients: In RL, the language model acts as a policy that autoregressively samples response tokens conditioned on the question and prior tokens.Each state is (x, y_<t) and the action is the next token y_t.
  • RL via Policy Gradients: The advantage function measures a response’s relative quality under the reward signal, and the experiments use GRPO because of its simplicity.GRPO is selected among policy-gradient approaches including REINFORCE and PPO.

4 Mid-Training (MT)

The method uses heuristic-guided self-generated trajectories to expose models to multiple correct reasoning approaches during mid-training before RL. This produces multi-modal next-token distributions that support conservative probability updates and help prevent mode collapse during subsequent policy-gradient training.

  • Data construction: Mid-training constructs data associating each question with multiple correct solution trajectories reflecting distinct reasoning behaviors.The approach targets diverse self-generated data as an intermediate phase between pre-training and post-training.
  • Heuristic-guided generation: Pólya-inspired heuristics guide generation by prompting the base model with textual descriptions and few-shot exemplars for each question–heuristic pair.Example strategies include working backward, introducing auxiliary quantities, and decomposing problems into simpler subproblems.
  • Trajectory selection: Sampled trajectories are filtered for correct final answers, scored for heuristic adherence, and selected by highest reward-model score.Rule-based verifiers such as Math-Verify perform correctness filtering before reward-model ranking.
  • Mid-training objective: Training on multiple trajectories encourages several high-probability next-token modes at prefixes shared across different solution strategies.This contrasts with training on a single solution per question, which does not encourage the same multimodal branching behavior.
  • Policy-gradient implications: Under a uni-modal regime, a policy-gradient update changes a highly confident sampled token’s probability only at scale ϵ^2, whereas N-modal distributions yield conservative yet appreciable changes.As N increases, updates become more conservative and reduce commitment to a single mode, helping prevent mode collapse.
  • Policy-gradient implications: With negative advantage, probability removed from a sampled token is redistributed mainly among the remaining N −1 dominant modes.The redistribution provides a conceptual explanation for why diverse solution approaches can benefit subsequent RL.

5 Experiments

Experiments evaluate Pólya-guided mid-training with varying numbers of self-generated solution variants, first directly and then as initialization for GRPO-based RL. Mid-training improves pass@64 substantially and generally outperforms vanilla RL at pass@2 and pass@64 across mathematical reasoning benchmarks.

  • Training Setup: Mid-training uses 7,112 filtered GSM8K questions, each paired with up to 64 heuristic-specific correct-solution variants, with separate models trained for n ∈{1, 2, 4, . . . , 64}.The original GSM8K training set contains 7,473 question–answer pairs and serves as Dseed.
  • Mid-Training Results: Average pass@1 rises from 9.25% at n = 1 to 11.50% at n = 64, while STaR reaches 13.02%, or +1.94% over zero-shot.The reported mid-training gains are steady but modest as the number of heuristics increases.
  • Mid-Training Results: Average pass@64 improves from 46.30% zero-shot and 46.32% with STaR to 48.17% at n = 64, with larger gains on harder benchmarks.AIME 2025 improves from 12.84% to 18.66%, AMC 2023 from 83.49% to 85.18%, and OlympiadBench from 42.13% to 43.57%.
  • Results with RL: At pass@2 on MATH-500, vanilla RL reaches 56.48% and STaR+RL 58.13%, while mid-trained models reach 59–60%, including 59.68% at n = 32.Mid-trained models generally outperform vanilla RL, although performance varies non-monotonically with the number of heuristics.
  • Results with RL: At pass@64, vanilla RL averages 44.21% and STaR+RL 45.69%, versus 48.09% for mid-training at n = 16 and 47.62% at n = 64.On AIME 2025, performance increases from 16.91% to 23.34%; on AMC 2023, it increases from 78.18% to 84.52%.

6 Analysis and Discussion

The analysis finds that RL combines diverse problem-solving approaches introduced during mid-training, with multiple approaches per problem outperforming broader single-solution coverage and improving out-of-domain performance. It also shows stronger diversity than teacher distillation while leaving the source of apparent emergent behavior unresolved.

  • Trace Analysis: An LLM-based judge classifies reasoning traces using few-shot prompts covering 64 Pólya heuristics, producing one binary decision per heuristic.The analysis compares traces from RL and mid-trained models to identify which problem-solving approaches and combinations they use.
  • Emergent Combination: RL consistently combines problem-solving approaches introduced during mid-training, including heuristic combinations not explicitly demonstrated before RL.This provides empirical evidence that RL composes approaches learned during mid-training within reasoning chains.
  • Limitations: Whether RL creates genuinely emergent behaviors or primarily composes behaviors learned during pre-training remains unresolved and is left for future investigation.The authors note that limited visibility into pre-training data may make existing behaviors appear emergent.
  • Approaches versus Problems: Mid-training with multiple problem-solving approaches per question outperforms training on more distinct problems with one solution across all k, despite equal training-instance counts.The experiment directly varies unique problems and approaches per problem while holding total mid-training instances fixed.
  • Teacher Distillation: The mid-training dataset reaches a Vendi Score of 13.81 versus 10.95 for teacher distillation, while achieving better pass@64 and comparable or better pass@1 after RL.Teacher-distilled RL rollouts are also notably more verbose and repetitive.
  • Out-of-Domain Generalization: On HumanEval, n = 32 and n = 64 models achieve 52.34% and 52.82%, respectively, versus 51.14% for Vanilla RL; on Murder Mystery, they achieve 56.94% and 57.36%, versus 53.15%.The mid-trained models outperform Vanilla RL on both evaluated out-of-domain benchmarks, with larger gains on MuSR.

7 Conclusions

The paper presents a mid-training strategy that uses diverse self-generated solution trajectories before RL. This approach improves subsequent RL across mathematical reasoning benchmarks, especially at higher pass@k values.

  • Method: The method generates multiple correct solution trajectories for each training question using George Pólya’s problem-solving heuristics, then fine-tunes before RL.This mid-training stage is designed to expose the model to diverse reasoning approaches before reinforcement learning.
  • Findings: The proposed mid-training strategy improves the performance of subsequent RL across a range of mathematical reasoning benchmarks.The conclusion attributes these gains to training with diverse self-generated data before RL.
  • Findings: The improvements are especially pronounced at higher pass@k values.This indicates that the method is particularly beneficial when evaluating multiple sampled solutions.

A Appendix · A.1 Limitations

The approach relies on math-centric Pólya heuristics, so a fully domain-general taxonomy remains unresolved. Results on HumanEval and MuSR suggest transfer beyond mathematics, while the procedure may surface and consolidate behaviors already present in pre-training data.

  • A.1 Limitations: The Pólya-based heuristics are inherently math-centric, leaving a fully domain-general taxonomy as an open question.HumanEval and MuSR results suggest the resulting approaches can transfer beyond mathematics.

A.2 Broader Impact

The work examines how pre-RL data shapes the behaviors elicited by reinforcement learning in language models. It argues, with theoretical and empirical evidence, that multiple problem-solving approaches per question can encourage RL to compose them within one reasoning chain.

  • A.2 Broader Impact: The study investigates mechanisms underlying reinforcement learning in language models rather than treating RL as a black-box capability technique.It focuses on how data exposure before RL shapes the behaviors that RL subsequently elicits.
  • A.2 Broader Impact: Exposure to multiple problem-solving approaches per question is presented as a way to encourage RL to compose those approaches within a single reasoning chain.The paper supports this mechanism with both theoretical and empirical evidence.
  • A.2 Broader Impact: The broader contribution is understanding how pre-RL training data influences the behaviors later elicited by RL.This frames the work as part of a broader effort to understand RL training mechanisms in language models.

A.3 Diversity of the Mid-Training Dataset

The mid-training dataset’s reasoning-trace diversity is quantified with the Vendi Score, which increases as more heuristic-guided variants are generated, though overlapping strategies prevent it from reaching the dataset size and downstream gains are not strictly monotonic.

  • Diversity measurement: The Vendi Score measures the effective number of dissimilar reasoning traces from the eigenspectrum of a pairwise similarity matrix, ranging from 1 for identical items to n items.It is calculated at the question level to quantify diversity in the heuristic-guided mid-training dataset.
  • Diversity trends: As n increases, the Vendi Score increases, indicating that heuristic-guided data creation introduces more diverse reasoning approaches.The score’s increase tracks the expansion of generated variants.
  • Diversity trends: The Vendi Score remains below n because some problem-solving strategies overlap, including analogy-based and auxiliary-problem strategies.The metric does not reach the total number of items when strategies are not fully dissimilar.
  • Diversity trends: Downstream gains are not strictly monotonic with respect to the Vendi Score, although the top-performing models tend to have higher scores.This indicates that greater measured diversity does not translate into uniformly increasing downstream performance.

A.4 Training Details

The training setup uses H100 GPUs, AdamW with cosine learning-rate scheduling for mid-training, and GRPO-based RL with specified rollout and batch configurations.

  • All experiments run on 4×NVIDIA H100 GPUs.
  • Mid-Training & STaR Setup: Mid-training is implemented with the trl framework and uses AdamW under a cosine learning-rate schedule.
  • Reinforcement Learning: RL training uses the DAPO-Math-17k dataset with GRPO in verl, accelerated by vLLM for rollout generation and evaluation.
  • Reinforcement Learning: Prompts are truncated to 1024 tokens, responses capped at 3072 tokens, and 16 responses sampled per prompt at temperature 1.0.
  • Reinforcement Learning: RL training uses a batch size of 32 and mini-batch size of 8, corresponding to one on-policy update and three off-policy updates.

A.5 Experiments on Qwen2.5–7B–Instruct

On Qwen2.5–7B–Instruct, zero-shot performs strongly, while mid-trained models generally improve RL performance over Vanilla RL and STaR+RL. Among Pólya configurations, n = 8 is usually strongest, aligning with rollout size g = 8.

  • Base-model and mid-training results: Zero-shot achieves the highest average pass@1 (18.47%) and pass@64 (51.39%), indicating substantial reasoning diversity in the base model.STaR underperforms zero-shot on average at 15.44% / 48.97%.
  • Base-model and mid-training results: Pólya n = 8 achieves the strongest average performance, with pass@1 (17.06%) and pass@64 (50.50%), and exceeds zero-shot on AMC 2023 pass@64.The AMC 2023 comparison is 92.88% vs. 92.66%; increasing n beyond 8 yields no further gains.
  • RL results: Mid-trained models generally outperform Vanilla RL and STaR+RL across the four mathematical benchmarks.Among Pólya configurations, n = 8 leads on AIME 24, AMC 23, and HMMT 25, while n = 4 leads on AIME 25 at pass@64 (42.54%).
  • RL results: On HMMT 25, n = 8 dominates at lower k, but at pass@64 Vanilla RL reaches 27.62%, STaR+RL 29.92%, and n = 8 reaches 26.87%.With rollout size g = 8, the best-performing n matches g on three of four benchmarks.

A.6 Proofs and Derivations

This section formally derives Theorem 4.1 and Proposition 4.2, starting with the exact first-order probability-mass shift and specializing it to uni-modal and N-modal regimes.

  • A.6 Proofs and Derivations: The derivations establish the exact analytical form of the first-order probability mass shift before specializing it to uni-modal and N-modal regimes.The section provides formal derivations for Theorem 4.1 and Proposition 4.2.

A.6.1 Preliminaries … A.11 Pólya’s Heuristics

The appendix formalizes how policy-gradient updates alter softmax probabilities, showing that negative-advantage updates in an N-modal policy redistribute mass mainly among alternative dominant modes. It also specifies heuristic-guided data generation, evaluation procedures, compositional analysis, and the Pólya-inspired reasoning approaches used in training and qualitative analysis.

  • A.6.1 Preliminaries: A single policy-gradient update changes token logits through the sampled token’s log-probability gradient under a softmax language-model policy.The setup defines πθ(· | x), token sampling, the learning rate η, scalar advantage A, and the softmax-logit gradient.
  • A.6.2 Derivation of First-Order Probability Change: The appendix derives a general first-order formula for the sampled token’s probability change and analyzes uni-modal and N-modal special cases.The derivation uses a Taylor expansion and the softmax Jacobian before separating the sampled-token term from the remaining tokens.
  • A.6.3 Redistribution of probability mass in the N-Modal Regime: Under a negative advantage A < 0, probability mass removed from the sampled token is redistributed predominantly to the remaining N −1 dominant modes rather than tail tokens.The tail receives only O(ϵ2/(N|V|)) mass per token, negligible relative to dominant-mode changes of O(η|A|/N 2).
  • A.7.1 Data Generation: Heuristic-specific generation uses descriptions and 4 few-shot exemplars, samples 128 responses per question–heuristic pair, and selects the highest-adherence correct response with Skywork-Reward-V2.The reward-model pipeline filters candidates for correctness, ranks them by heuristic adherence, and selects one response for mid-training.
  • A.7.2 Reinforcement Learning: RL prompts pair each question with a chain-of-thought instruction, while correctness is verified programmatically.This is the specified RL training prompt format and verification procedure.
  • A.8 Evaluation Datasets: Evaluation covers open-ended math reasoning benchmarks, with reported results restricted to each dataset’s test set and answers verified with Math-Verify.The listed datasets include Math-500, AIME 2024, AIME 2025, AMC 2023, HMMT 2025, and OlympiadBench.
  • A.9 Human Evaluation: Human evaluation independently labels 50 sampled heuristic-classification examples from the n = 16 mid-trained model on MATH-500, while qualitative examples show RL composing multiple approaches.Annotators receive the heuristic description, canonical example, and full reasoning chain; examples depict restatement, decomposition, analogy, and combinatorial reasoning.
Loading 2605.08472v1…