Technical Summary: How You Begin is How You Reason: Driving Exploration in RLVR via Prefix-Tuned Priors

1. Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a powerful paradigm for enhancing Large Language Model (LLM) reasoning capabilities in tasks such as mathematical problem solving and theorem proving. However, the authors identify a critical limitation in current RLVR approaches: entropy collapse.

While RLVR often improves single-rollout accuracy (Pass@1), it frequently fails to expand the coverage of successful reasoning trajectories. This manifests as a reduction in output diversity, where the model converges to a narrow set of reasoning paths. Existing exploration techniques, such as entropy regularization or intrinsic rewards, often treat the base model's distribution as fixed and attempt to modify the RL objective or sampling procedure. The paper argues that these methods are insufficient because RL is strongly shaped by the base model's prior distribution. Furthermore, passive exploration techniques often degrade generation quality, resulting in noisy rollouts that do not contribute to effective reasoning.

The core challenge is to encourage effective exploration—discovering diverse, task-relevant, and high-quality reasoning behaviors—without collapsing the policy into redundant or low-quality trajectories, and without the computational cost of updating the full model parameters.

2. Methodology: IMAX Framework

The authors propose Information-Maximizing Augmented eXploration (IMAX), a lightweight framework that shifts the optimization target from the LLM parameters to soft prompt tuning. Instead of modifying the RL process on top of a fixed prior, IMAX learns a pool of soft prefixes that actively reshape the model's prior over reasoning trajectories.

2.1 Prefix-Conditioned Reasoning

The core mechanism involves a pool of $C$ soft prefixes (continuous virtual tokens) $\mathcal{E} = \{e_1, \dots, e_C\}$. For a given input prompt $x$, a prefix $e_z$ is prepended to the input embedding sequence. The backbone LLM remains frozen, and the policy $\pi_E(y|x, z)$ is conditioned on both the prompt and the selected prefix identity $z$.

• Initialization: Prefixes are initialized from a pool of semantically meaningful text prompts (e.g., "Introduce variables explicitly," "Partition into cases") generated by an external LLM. This ensures the prefixes start with distinct reasoning strategies.
• Control: Each prefix acts as a "control knob," inducing a distinct rollout distribution and steering the model toward different regions of its reasoning space (e.g., algebraic vs. diagrammatic vs. code-like solutions).

2.2 Information-Maximizing (InfoMax) Reward

To prevent the learned prefix pool from collapsing into redundant controls (where different prefixes induce identical trajectories), the authors introduce an InfoMax reward. The objective is to maximize the mutual information (MI) between the prefix identity $Z$ and the generated trajectory $Y$, conditioned on the prompt $X$:
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
Intuitively, this encourages different prefixes to produce distinguishable, valid reasoning behaviors for the same question.

2.3 Variational Lower Bound and Training

Directly computing the mutual information is intractable. The authors derive a variational lower bound using a variational posterior $q_\phi(z | x, y)$, which is a lightweight classifier trained to predict which prefix generated a specific response $y$ given prompt $x$.

• Reverse Form: They utilize the reverse decomposition of MI: $I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$.
• Reward Formulation: The InfoMax reward is defined as $\log q_\phi(z | x, y)$. The total objective for the prefix pool becomes:
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  where $R(x, y)$ is the standard verifiable reward (correctness) and $\beta$ is a hyperparameter.

2.4 Alternating Optimization

The training procedure follows an alternating optimization scheme (Coordinate Ascent Variational Inference):

1. Update Posterior ($q_\phi$): With the prefix pool fixed, update the classifier $q_\phi$ to predict the prefix identity from the prompt-response pair (generated without the soft prefix to prevent information leakage).
2. Update Prefix Pool ($E$): With $q_\phi$ fixed, update the soft prefixes using a Group Relative Policy Optimization (GRPO) algorithm augmented with the InfoMax reward.
3. Stratified Sampling: During training, each prompt is paired with distinct prefix identities to ensure uniform coverage of the prefix space within each GRPO group.

3. Key Contributions

1. Lightweight RLVR Paradigm: The paper introduces a method to improve reasoning by optimizing a pool of soft prompts rather than full model parameters. This reshapes the conditional generation distribution of a frozen model without the computational overhead of full fine-tuning.
2. InfoMax Reward for Diversity: The authors propose a response-level InfoMax reward that explicitly encourages different prefixes to induce distinguishable, task-relevant reasoning trajectories, preventing the collapse of the prefix pool into redundant controls.
3. Empirical Validation: Extensive experiments on mathematical reasoning benchmarks (GSM8K, MATH-500, Minerva, AMC) across three model scales (Qwen2.5-1.5B, Qwen3-4B, Qwen3-8B) demonstrate that IMAX improves both reasoning quality and diversity.

4. Experimental Results

The authors evaluated IMAX against several baselines, including full-policy GRPO, text-prompt initialization, and other diversity-oriented RLVR methods (DIVER, MERCI).

• Performance Gains: IMAX consistently outperformed standard prefix-tuned RLVR baselines.
  • On Qwen3-4B, IMAX improved Pass@4 by up to 11.60% and Avg@4 by 10.57% compared to standard baselines on specific benchmarks.
  • On Qwen3-8B, DAPO+IMAX showed significant improvements over DAPO-prefix on GSM8K, MATH-500, and Minerva.
• Diversity and Coverage:
  • t-SNE Analysis: Visualizations of response embeddings showed that IMAX-induced prefixes resulted in more separated clusters compared to other methods, indicating stronger intra-group diversity.
  • Reasoning Components: Heatmaps revealed that different prefixes successfully induced distinct reasoning patterns (e.g., Prefix 1 favoring algebraic manipulation, Prefix 2 favoring case partitioning and verification), aligning with their semantic initialization.
  • Oracle Performance: The combined performance of multiple prefixes (Oracle) exceeded the best single prefix, confirming that the prefixes cover complementary successful regions of the problem space.
• Efficiency: IMAX requires significantly less training cost (3.0x–6.0x less GPU-hours) compared to full-parameter GRPO, as it only updates the lightweight prefix embeddings and the classifier head.

5. Significance and Claims

The paper claims that how a model begins a reasoning process (the prior) fundamentally dictates how it reasons. By treating soft prompts as reusable control variables, IMAX offers a scalable solution to the exploration problem in RLVR.

• Algorithm Agnostic: The framework is designed to be seamlessly integrated into existing RLVR pipelines (e.g., GRPO, DAPO).
• Structured Diversity: Unlike methods that simply increase randomness, IMAX encourages structured reasoning diversity, ensuring that the exploration leads to semantically distinct and high-quality paths.
• Modest Limitations: The authors acknowledge that the method is primarily designed for reasoning tasks, showing limited gains on general instruction-following tasks. Additionally, the effectiveness depends on the quality of the initial text prompts, which may require manual design or retuning across domains. The learned prefixes also introduce a persistent, albeit small, overhead during inference.

In summary, IMAX demonstrates that optimizing the "beginning" of the generation process via soft prefixes is a highly effective strategy for driving exploration and improving reasoning capabilities in LLMs, offering a computationally efficient alternative to full-policy reinforcement learning.