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

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

Problem Statement
Recent advancements in Large Language Models (LLMs) have highlighted Reinforcement Learning (RL) as a potent technique for enhancing reasoning capabilities, particularly through inference-time scaling. However, the effectiveness of RL is heavily contingent on the priors established in the model prior to RL training. Existing literature suggests that RL with verifiable rewards (RLVR) often performs "distributional sharpening," selectively amplifying behaviors already present in the base model rather than inducing genuinely new capabilities. When a model lacks exposure to diverse problem-solving approaches for a single task, subsequent RL training may struggle to improve performance or may even degrade it under certain evaluation metrics (e.g., pass@k). The core challenge addressed is how to effectively equip a base model with a rich set of diverse reasoning strategies before applying RL, thereby enabling the RL process to compose and refine these strategies rather than merely reinforcing a narrow set of existing behaviors.

Methodology
The authors propose a mid-training framework that serves as an intermediate step between pre-training and RL. This phase utilizes diverse self-generated data to expose the model to multiple valid solution trajectories for the same problem.

1. Heuristic-Guided Data Generation:

   • The framework draws inspiration from George Pólya's problem-solving heuristics (e.g., Analogy, Decomposing & Recomposing, Working Backwards).
   • For each question in a seed dataset (GSM8K), the base model is prompted with specific heuristic descriptions and few-shot exemplars to generate multiple candidate responses.
   • This process is bootstrapped: the model generates its own training data without relying on distillation from stronger teacher models.
   • Filtering and Selection: Generated responses are filtered for correctness using rule-based verifiers (e.g., Math-Verify). The remaining candidates are scored by a reward model ($R_\phi$) based on adherence to the specific heuristic. The highest-scoring response for each question-heuristic pair is selected.
   • The resulting dataset ($D_{Pólya}$) contains $n$ distinct solution variants per question, where $n$ represents the number of heuristics applied.

2. Mid-Training (Supervised Fine-Tuning):

   • The base model is fine-tuned on $D_{Pólya}$. Unlike standard SFT which pairs a question with a single response, this approach pairs each question with $n$ diverse trajectories.
   • The objective is to minimize the negative log-likelihood over these multiple trajectories, encouraging the model to assign non-negligible probability to various next-token options at branching points, effectively creating an N-modal next-token distribution.

3. Reinforcement Learning:

   • Following mid-training, the model undergoes RL training using Group Relative Policy Optimization (GRPO) on a separate dataset (DAPO-Math-17k).
   • The study compares models initialized with varying numbers of heuristics ($n \in \{1, 2, \dots, 64\}$) against baselines: Zero-Shot, Vanilla RL (directly on the base model), and STaR+RL (Self-Taught Reasoner followed by RL).

Theoretical Perspective
The paper provides a theoretical analysis of how mid-training influences policy-gradient updates:

• Uni-modal vs. N-modal Regimes: In a uni-modal regime (high confidence in a single token), a policy gradient update results in negligible changes ($\propto \epsilon^2$). In contrast, in an N-modal regime (probability mass split among $N$ dominant modes), the update induces conservative yet appreciable changes ($\propto 1/N$).
• Composition of Strategies: The authors prove that under a negative advantage update, probability mass removed from a sampled token is redistributed predominantly to the remaining $N-1$ dominant modes. This mechanism incentivizes the model to learn to combine different problem-solving approaches within a single response during RL, rather than collapsing to a single strategy.

Key Results
Experiments were conducted on Llama 3.2–3B–Instruct and Qwen2.5–7B–Instruct across six mathematical reasoning benchmarks (Math-500, AIME 2024/2025, AMC 2023, HMMT 2025, OlympiadBench) and out-of-domain tasks (HumanEval, MuSR).

• Performance Gains: Mid-training with diverse self-generated data consistently improves RL performance, particularly at higher $k$ values (pass@64).
  • On Math-500, pass@64 improved from 87.76% (Zero-Shot) to 88.94% ($n=64$).
  • On AIME 2025, pass@64 increased from 12.84% to 18.66%.
  • Models mid-trained with higher reasoning diversity ($n=64$) outperformed both Vanilla RL and STaR+RL baselines across all benchmarks.
• Diversity vs. Quantity: A controlled experiment showed that learning multiple approaches per question (fewer problems, more approaches) yields greater benefits for RL than learning a single approach across a larger number of problems, given a fixed training budget.
• Correctness is Essential: Mid-training with diverse approaches that lead to incorrect answers failed to improve RL performance, indicating that correctness is a prerequisite for the benefits of diversity.
• Superiority over Distillation: The self-generated data exhibited higher diversity (Vendi Score 13.81) compared to data distilled from a stronger teacher model (QwQ-32B, Vendi Score 10.95), leading to better downstream performance.
• Generalization: Despite the math-centric nature of Pólya's heuristics, the learned strategies generalized to code generation (HumanEval) and narrative reasoning (MuSR), with significant gains on multi-step reasoning tasks.

Significance and Claims
The paper claims that the primary contribution is demonstrating that exposure to multiple problem-solving approaches via self-generated data during mid-training fundamentally shapes the priors of the model, enabling subsequent RL to effectively compose these strategies.

• Mechanism of Improvement: The study posits that RL does not necessarily create "emergent" reasoning capabilities from scratch but rather composes existing behaviors. Mid-training ensures these behaviors are diverse and available for composition.
• Practical Utility: The approach offers a simple, effective, and scalable method to improve RL without requiring human annotation or stronger teacher models.
• Theoretical Insight: The work bridges the gap between the empirical success of RL and theoretical understanding, showing how N-modal distributions prevent mode collapse and facilitate the learning of composite reasoning chains.

The authors remain modest regarding single-sample accuracy (pass@1), noting that their focus is on understanding how diverse exposure shapes RL behavior and improving performance in scenarios requiring multiple sampling attempts (pass@k). They acknowledge that while the heuristics are math-centric, the resulting generalization to other domains suggests the procedure surfaces and consolidates latent capabilities within the base model.