PaperlayerView exact revision on arXiv

Source-linked AI summary

Your Language Model is Its Own Critic: Reinforcement Learning with Value Estimation from Actor's Internal States

Yunho Choi, Jongwon Lim, Woojin Ahn, Minjae Oh, Jeonghoon Shim, Yohan Jo

arXiv:2605.07579v2cs.LGcs.AIcs.CL

TL;DR

RLVR baseline estimation is costly because PPO uses a policy-scale critic and GRPO needs multiple rollouts per prompt, limiting prompt diversity and increasing compute. POISE predicts expected reward from the policy’s internal states through a lightweight probe and cross-rollout construction, preserving gradient unbiasedness. Across math reasoning benchmarks, it matches DAPO at lower computational cost while its estimator performs comparably to a separate value model.

  • Problem

    PPO requires an LLM-scale critic and GRPO requires multiple rollouts per prompt, making reliable RLVR baselines computationally expensive.

  • Method

    POISE trains a lightweight probe on prompt and trajectory internal signals, using an independent cross-rollout to keep the baseline conditionally independent of the optimized action.

  • Results

    POISE achieves performance comparable to DAPO on mathematical reasoning benchmarks at lower computational cost and performs as well as a separate policy-scale value model.

  • Takeaways & Limitations

    Policy internal representations can serve as practical optimization signals for stable and scalable RLVR without group-relative baselines or a separate critic model.

  • Takeaways & Limitations

    Experiments use a fixed compute budget, and behavior under substantially longer training horizons remains uncharacterized.

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards (RLVR) for Large Reasoning Models hinges on baseline estimation for variance reduction, but existing approaches pay a heavy price: PPO requires a policy-model scale critic, while GRPO needs multiple rollouts per prompt to keep its empirical group mean stable. We introduce Policy Optimization with Internal State Value Estimation), which obtains a baseline at negligible cost by using the policy model's internal signals already computed during the policy forward pass. A lightweight probe predicts the expected verifiable reward from the hidden states of the prompt and generated trajectory, as well as token-entropy statistics, and is trained online alongside the policy. To preserve gradient unbiasedness despite using trajectory-conditioned features, we introduce a cross-rollout construction that predicts each rollout's value from an independent rollout's internal states. Because POISE estimates prompt value using only a single rollout, it enables higher prompt diversity for a fixed compute budget during training. This reduces gradient variance for more stable learning and also eliminates the compute overhead of sampling costs for detecting zero-advantage prompts. On Qwen3-4B and DeepSeek-R1-Distill-Qwen-1.5B across math reasoning benchmarks, POISE matches DAPO while requiring less compute. Moreover, its value estimator shows similar performance to a separate LLM-scale value model and generalizes to various verifiable tasks. By leveraging the model's own internal representations, POISE enables more stable and efficient policy optimization.

1 Introduction

RLVR baseline estimation improves training stability but imposes substantial compute costs in PPO and GRPO. POISE uses the policy’s internal representations with a cross-rollout probe to provide a lightweight baseline and matches strong RL performance with lower overhead.

  • Motivation: PPO trains an LLM-scale critic, while GRPO requires multiple rollouts per prompt, reducing prompt diversity and increasing gradient-estimation variance under fixed compute.Both approaches also spend compute on prompts whose rollouts receive identical rewards and therefore have zero advantage.
  • Method: POISE trains a lightweight probe on prompt-level and trajectory-level internal signals to predict the policy’s expected reward.The trajectory-level features include terminal hidden states and token-level entropy.
  • Method: POISE uses an independent cross-rollout prediction so trajectory-conditioned baseline features remain conditionally independent of the rollout action.This construction preserves unbiased gradient estimation while training the probe to recover the policy’s expected reward.
  • Advantages: Compared with PPO and GRPO, POISE uses a lightweight estimator and only a rollout pair, redirecting saved compute toward more distinct prompts.Its continuous baseline also avoids extra sampling to detect zero-advantage prompt groups.
  • Results: POISE matches DAPO with less compute and shows that internal representations can serve as practical optimization signals rather than only diagnostic tools.The paper also analyzes estimator behavior during training and evaluates its value prediction against a separate critic model.

2 Preliminaries

RLVR maximizes expected verifiable reward with policy gradients, using a baseline to reduce variance without biasing the gradient. The preliminaries motivate POISE by showing that broader prompt allocation can reduce estimator noise under fixed compute.

  • 2.1 Policy Gradient and Baseline Estimation: RLVR models prompt–response pairs as a contextual bandit and maximizes expected verifiable reward R(x, y) over sampled responses.The policy samples y from πθ conditioned on prompt x.
  • 2.1 Policy Gradient and Baseline Estimation: The policy-gradient estimator subtracts a baseline b(x) from reward to form the advantage A(x, y) = R(x, y) − b(x), reducing variance.The baseline does not change the target objective when the required independence condition holds.
  • 2.1 Policy Gradient and Baseline Estimation: The value function is the standard near-optimal variance-reduction baseline, but it is unknown and must be estimated.PPO learns a critic, whereas GRPO estimates the prompt baseline from the mean reward of multiple same-prompt responses.
  • 2.2 Baseline Independence: Baseline subtraction preserves unbiasedness when the baseline is conditionally independent of the sampled response given the prompt.The cross-rollout construction computes one response’s baseline from another independent response.
  • 2.3 Gradient Variance and Number of Prompts in the Batch: Proposition 1 decomposes gradient variability into within-prompt and between-prompt covariance components that are fixed by the data distribution, policy, and reward.The decomposition formalizes how rollout allocation affects the batch estimator.
  • 2.3 Gradient Variance and Number of Prompts in the Batch: Under fixed completion budget B = n · m, gradient-estimator variance is minimized at m = 1 and n = B, favoring as many distinct prompts as possible.GRPO’s repeated sampling conflicts with this allocation because it needs multiple responses to estimate a reliable group baseline.

3 Policy Optimization with Internal State Value Estimation (POISE)

POISE uses a lightweight probe over the policy model’s internal states and token-entropy statistics to estimate prompt value, then integrates that estimator into policy optimization with cross-rollout baselines. The method trains online from paired rollouts while preserving conditional independence between each rollout’s update and its baseline.

  • Value estimation: POISE trains a lightweight probe to predict expected verifier reward from prompt-level states, trajectory-level states, and token-entropy statistics.The probe is trained jointly with the policy and uses signals already produced during the policy forward pass.
  • Value estimation: The probe’s target is the prompt-level value, even though its input includes generated reasoning features.Training uses expected reward derived from other responses rather than asking the estimator to verify the same reasoning trajectory.
  • Value estimation: Internal-state probes achieve better held-out value prediction than a separate policy-scale critic while adding only a lightweight regression head.The preliminary benchmark uses reward-labeled rollouts from the DAPO-Math dataset.
  • Online optimization: Training uses two independent rollouts per prompt, evaluates both verifiable rewards, and updates the probe with recent examples plus a trajectory buffer.The buffer stabilizes training under policy drift, while the lightweight probe adds negligible update cost.
  • Cross-rollout baselines: POISE constructs each rollout’s baseline from the other independently sampled rollout’s internal signals, satisfying the conditional-independence requirement.This cross-rollout construction prevents the target from leaking the reward of the rollout whose features are used by the probe.
  • Online optimization: POISE uses a PPO-style clipped surrogate objective to optimize the policy over multiple inner epochs per batch.The cross-rollout baselines provide the advantages used in this policy update.
  • Empirical evaluation: POISE achieves competitive Avg@32 accuracy with baseline models on olympiad-level mathematical reasoning benchmarks.Table 1 reports Avg@32 accuracy across multiple datasets.

4 Experiments

POISE achieves math-reasoning performance comparable to DAPO while reducing training compute and providing a lower-noise, online-adaptive baseline.

  • Main results: POISE achieves Avg@32 scores of 0.500 versus DAPO’s 0.508 on Qwen3-4B, while improving Deepseek-Distill-Qwen-1.5B from 0.296 to 0.303.It outperforms DAPO on several individual benchmarks across both model scales.
  • Training efficiency: POISE’s continuous internal-state baseline avoids the degenerate-group sampling overhead required by DAPO when all sampled rewards are identical.This allows usable advantage signals without repeatedly sampling additional groups.
  • Diagnostics: Figure 3 compares wall-clock time per step and gradient norm between POISE and DAPO, while Figure 4 reports online MAE and Figure 5 compares the estimator with a critic.The figures jointly assess efficiency, optimization stability, calibration, and critic agreement.
  • Training efficiency: POISE requires about 18 versus 24 wall-clock hours on DeepSeek-R1-Distill-Qwen-1.5B and about 36 versus 49 hours on Qwen3-4B versus DAPO.The savings come from avoiding large rollout groups and extra sampling for degenerate prompts.
  • Estimator dynamics: The online estimator remains calibrated as the policy changes and reduces reward variance by roughly one third after the initial training phase.The reported variance reduction ratio remains around 30%.

5 Analysis of the Value Estimator

The value estimator uses multiple internal signals and remains competitive with a separately trained critic across domains and models, while simple probe designs and mid-later layers are effective.

  • Generalizability: The estimator is competitive with and often more accurate than a critic across mathematical reasoning, coding, tool-calling, and instruction-following tasks.This evaluation spans Qwen3-4B and DeepSeek-R1-Distill-Qwen-1.5B/7B, using MAE and correlation against target-policy Avg@8 scores.
  • Input features: Trajectory-level features, including reasoning hidden states and mean entropy, are among the estimator’s ablated input components.The ablation compares prompt hidden states, reasoning hidden states, vocabulary entropy, and response length.
  • Extraction hyperparameters: Mid-later hidden-state layers are optimal, and probe performance is not sensitive to token length.These findings concern hidden-state extraction hyperparameters.
  • Probe architecture: Linear regression is as effective as, and sometimes better than, heavier MLP probes.The paper attributes this efficiency to semantic features encoded as linear directions in Transformer representations.

6 Related Work

Prior LLM RL methods estimate variance-reducing baselines with critics or rollout groups, while hidden-state research has mainly treated internal signals as diagnostics or test-time controls.

  • Value estimation in RL: Explicit value models can estimate baselines but incur the training, calibration, and deployment costs of an additional LLM-scale model.This is one major axis along which recent LLM reasoning RL work extends baseline estimation.
  • Outcome-relevant information: Hidden states have been shown to encode outcome-relevant information, including factuality, truthfulness, confidence, and answer correctness.These findings motivate using internal representations as value-related signals.
  • Positioning: POISE differs by incorporating hidden-state signals directly into RL training through an online estimator, avoiding both an auxiliary critic and large rollout groups.The paper positions this as a shift from diagnostic or test-time use toward optimization.

7 Conclusion

POISE uses cross-rollout internal-state value estimation to obtain a cheap baseline, matching DAPO at lower cost while supporting stable training and broader verifiable-task generalization.

  • Conclusion: POISE predicts rollout value from policy internal states rather than a group-mean baseline or separate critic, using cross-rollout construction to preserve unbiasedness.The method combines lightweight value estimation with a design intended to maintain the gradient estimator’s validity.
  • Conclusion: POISE matches DAPO on mathematical reasoning benchmarks at lower computational cost and achieves more stable training.The conclusion also reports comparable performance with a separate policy-scale value model.
  • Conclusion: The value estimator performs as well as a separate policy-scale value model and generalizes to other verifiable tasks.This extends the claimed utility beyond the main mathematical-reasoning setting.

8 Limitations and Future Work

The method is evaluated under a fixed compute budget, with longer-horizon behavior left for future study and extensions proposed beyond sequence-level mathematical reasoning.

  • Limitations: Longer training horizons remain uncharacterized despite consistent trends across backbones and benchmarks.The authors leave this study to future work requiring greater compute resources.
  • Future Work: Future work includes token-level credit assignment, preference-learning applications, and RL training for agentic reasoning and instruction following.These extensions would move beyond sequence-level value prediction and mathematical reasoning.
  • Method: POISE uses cross-rollout baselines with two rollouts per prompt and updates the policy with PPO while training the value estimator online.The algorithm maintains a value buffer and computes advantages from cross-rollout predictions.
  • Experimental Setup: DAPO is used as the baseline RL algorithm, with dynamic sampling and other techniques retained while the original length penalty is removed.The experiments use a binary correctness reward instead of DAPO’s length-aware penalty.
  • Method: Hidden-state features are collected during an already-required teacher-forced forward pass, avoiding additional computation for estimator training.Prompt and reasoning states are pooled separately, while paired-rollout rewards provide supervised targets.
  • Experimental Setup: Experiments use Qwen3-4B and DeepSeek-R1-Distill-Qwen-1.5B on an English-only filtered DAPO-Math-17K corpus.The maximum response length is capped at 8,192 tokens, with batch sizes of 1,024 and 512 prompts respectively.

B.5 Evaluation Protocol

Evaluation covers multiple mathematical competition benchmarks using standardized sampling, binary correctness rewards, and avg@k metrics with k fixed at 32.

  • Benchmarks: AMC23 and AMC24 contribute 80 problems from two American Mathematics Competition editions.Each edition contains 40 problems.
  • Benchmarks: AIME24, AIME25, and AIME26 contribute 90 challenging problems from three examination editions.Each edition contains 30 problems.
  • Benchmarks: HMMT25 and BRUMO25 each contribute 30 problems from high-school mathematics competitions.The datasets represent demanding tournament and olympiad-level problem sets.
  • Evaluation Protocol: Each test problem is evaluated with 32 independently sampled completions using temperature 0.6, top-p 0.95, and an 8,192-token limit.Inference runs on two NVIDIA B200 GPUs in a single node.
  • Metrics: Rewards are binary correctness indicators, and the paper reports avg@k metrics with k = 32.The same reward function is used during reinforcement learning and evaluation.
  • Prompt Format: All mathematical tasks use a single-turn prompt requiring step-by-step solutions and a final line formatted as Answer: $Answer.The answer must appear on its own line after the specified marker.

C Training Dynamics

Training dynamics show that POISE’s predicted values track improving rewards while its online error remains controlled and its learned baseline generally reduces advantage variance.

  • Qwen3-4B: On Qwen3-4B, reward rises from roughly 0.35 to above 0.65 before stabilizing around 0.70, while predicted value reaches 0.70-0.75.The matching trends indicate tracking of the changing reward scale.
  • Qwen3-4B: On Qwen3-4B, online target MAE stabilizes after early training and the advantage variance ratio stays below one for most updates.This indicates variance reduction despite a changing policy target.
  • DeepSeek-R1-Distill-Qwen-1.5B: On DeepSeek-R1-Distill-Qwen-1.5B, reward increases from roughly 0.15-0.20 to around 0.45-0.50 despite noisier dynamics.The estimated value also rises after an initial drop.
  • DeepSeek-R1-Distill-Qwen-1.5B: On the smaller model, online target MAE remains bounded and the advantage variance ratio stays below one for most training after early updates.The baseline continues reducing variance under a lower and more variable reward distribution.
  • Online Evaluation: The online target is an empirical Avg@8 score from the current actor checkpoint, and estimator quality is assessed with Pearson correlation and MAE.Checkpoints are evaluated every 10 training steps using eight sampled responses per prompt.
  • Critic Comparison: The internal-state estimator closely tracks a separately trained critic across both model scales while using hidden-state and entropy features already produced by the policy.The critic can be slightly more accurate because it is an LLM-scale model trained on accumulated rollout data.

D.3 Comparisons on Multiple Domains and Models

Across verifiable-reward domains and policy backbones, the internal-state estimator is often competitive with or better than a separately trained critic, with especially strong gains on Qwen3-4B outside mathematics.

  • Scope: The evaluation spans five domains—math, coding, tool use, and instruction following—and three policy backbones.Datasets include DAPO-Math, DeepScaleR, AceCoder, ToolDial, and IF-RLVR; backbones include Qwen3-4B, DeepSeek-R1-Distill-Qwen-1.5B, and DeepSeek-R1-Distill-Qwen-7B.
  • Qwen3-4B: On Qwen3-4B, the internal-state estimator outperforms the critic across all five domains in MAE and Pearson correlation.Both estimators are trained on the same reward-labeled rollout data and evaluated against held-out empirical Avg@8 values.
  • Qwen3-4B: On AceCoder, Pearson correlation increases from 0.056 to 0.612 for the internal-state estimator.This is one of the largest reported gains outside the original mathematical setting.
  • Qwen3-4B: On IF-RLVR, Pearson correlation increases from 0.150 to 0.642 for the internal-state estimator.The result supports value prediction beyond the DAPO-Math training domain.
  • DeepSeek-R1-Distill-Qwen-1.5B: On DeepSeek-R1-Distill-Qwen-1.5B, results are mixed: the estimator wins on DeepScaleR and ToolDial, while the critic wins on AceCoder and slightly on IF-RLVR MAE.The estimator remains competitive in correlation in the weaker cases.
  • DeepSeek-R1-Distill-Qwen-7B: On completed DeepSeek-R1-Distill-Qwen-7B settings, the internal-state estimator improves over the critic on both DAPO-Math and the reported metrics.These results support applicability beyond the mathematical-reasoning training experiments.

E.1 Ablations of Hyperparameters During Hidden State Extraction

The ablations show that POISE is robust to hidden-state extraction choices, with layer and pooling decisions affecting performance modestly and a linear probe offering the strongest correlation. The main experiments therefore use a shared layer-19, last-10-token configuration and a linear ridge estimator.

  • Layer index: Layer 19 gives Qwen3-4B the best Pearson correlation, while layer 33 gives the lowest MAE; layer 19 is selected for strong correlation and near-optimal MAE.The default prioritizes correlation with verifier value while retaining competitive absolute error.
  • Layer index: For DeepSeek-R1-Distill-Qwen-1.5B, the earliest layer performs best, but absolute MAE differences are small and layer 19 remains competitive.Layer 19 is retained for both backbones to avoid model-specific tuning and maintain implementation consistency.
  • Pooling window: Last-10 and last-15 pooling perform almost identically on Qwen3-4B, while last-5 pooling is weaker; last-10 is chosen for the best Pearson correlation with a shorter window.For DeepSeek-R1-Distill-Qwen-1.5B, the three windows are close, with last-10 achieving the best MAE and Pearson correlation.
  • Overall choice: The ablations support layer 19 with last-10-token mean pooling as a robust shared extraction setting and a linear probe as POISE’s default estimator.Value-relevant information is largely linearly accessible, although larger nonlinear probes can improve calibration in some cases.
  • Probe architecture: The best 3-layer, width-1024 MLP lowers MAE from 0.124 to 0.117, but the linear ridge probe achieves the highest Pearson correlation at 0.834.The linear estimator is selected because it is cheaper to fit online, has fewer hyperparameters, and better supports stable advantage formation.
Loading 2605.07579v2…