Technical Summary: Privileged Value Functions for LLM Reinforcement Learning

Problem Statement

Reinforcement Learning (RL) for Large Language Models (LLMs) is currently dominated by critic-free, group-relative methods like GRPO (Group Relative Policy Optimization). While these methods reduce gradient variance by sampling multiple responses per prompt and comparing them, they suffer from two primary limitations:

1. Sequence-Level Credit Assignment: They provide advantages only at the sequence level, failing to offer token-level granularity which is theoretically superior for temporal credit assignment.
2. Straggler Effects and Off-Policyness: Training is blocked by the slowest rollout in a group, reducing throughput. In asynchronous settings, this increases off-policyness as the policy drifts while waiting for slow responses.

Learned value functions theoretically address both issues by providing token-level advantages and operating without large groups. However, their adoption has been hindered by:

• Infrastructure Complexity: Maintaining a separate value model adds engineering overhead.
• Fitting Sensitivity: If the critic is poorly fitted, value baselines can introduce noise or bias, often performing worse than simple group means.
• Lack of Justification: The practical success of critic-free methods has made it difficult to justify the added cost of value functions.

Methodology

The paper proposes two complementary strategies to improve value-function-based RL, aiming to justify their inclusion by enhancing performance and robustness.

1. Privileged Value Functions (PVF)

The core insight is that value functions can condition on privileged information—data available during training but hidden from the policy during inference—without violating the unbiasedness of the policy gradient.

• Mechanism: A Privileged Value Function $V_\phi(h_{i,t}, z_{i,t})$ conditions the critic on the standard policy history $h_{i,t}$ and additional context $z_{i,t}$.
• Admissibility: To ensure the policy gradient remains unbiased, the privileged information $z_{i,t}$ must be conditionally independent of the current token $y_{i,t}$ given the history $h_{i,t}$. It cannot contain future tokens or rewards generated by the current trajectory.
• Forms of Privileged Information:
  • Reference Solutions: Oracle answers, proof sketches, or gold patches (e.g., in Sudoku or code repair).
  • Leave-One-Out (LOO) Group Info: Other sampled responses and their rewards from the same prompt group. This treats the group as an in-context learning prompt for the critic.
  • Verifier Rubrics: Detailed task specifications unavailable to the policy.
• Theoretical Benefit: Conditioning on richer information reduces the variance of the value prediction ($E[(R - E[R|H, Z])^2] \le E[(R - E[R|H])^2]$), thereby reducing the variance of the policy gradient.

2. TETHER (Adaptive Group-Value Baseline)

To address the risk of poorly fitted critics, the authors introduce TETHER, a baseline that adaptively interpolates between group-relative and value-based baselines.

• Mechanism: TETHER computes a mixed baseline $b^{TETHER}_{i,t} = (1-\rho)b^{LOO}_i + \rho V_{i,t}$, where $b^{LOO}_i$ is the leave-one-out group mean and $V_{i,t}$ is the learned token value.
• Adaptive Coefficient ($\rho$): The mixing ratio $\rho$ is not a fixed hyperparameter. It is fitted via least squares to minimize the prediction error of the Monte Carlo returns ($R_i$) on completed groups.
  • $\rho = 0$: Recovers the group baseline (reliable, no token-level credit).
  • $\rho = 1$: Recovers the pure value baseline (token-level credit, sensitive to critic fit).
  • Intermediate values: Smoothly trade off between the two.
• Training Dynamics: The coefficient $\rho$ is updated using an Exponential Moving Average (EMA) on batches after the policy update for that batch is computed, ensuring the baseline for a batch does not depend on its own returns (preserving unbiasedness).

Key Contributions

1. Privileged Value Functions (PVF): A formal mechanism to inject task-relevant, token-level signals into the value function using information hidden from the policy. This improves value estimation and reduces gradient variance without altering the policy objective.
2. TETHER: An adaptive baseline that dynamically balances the reliability of group means with the granularity of value functions based on the critic's current accuracy, removing the need for manual hyperparameter tuning to manage the transition.
3. Empirical Validation: Demonstration that both strategies consistently outperform standard value baselines and are competitive with or superior to mean-baseline GRPO across multiple reasoning-heavy tasks.

Experimental Results

The authors evaluated their methods on four tasks: Reasoning Gym (procedural reasoning), CodeIO (program reasoning), Sudoku (multi-turn grid filling), and MiniF2F (formal mathematics).

• PVF Performance:

  • PVF consistently outperformed both the standard Value Function (VF) and the Group Mean (MEAN) baselines across all settings.
  • Sudoku: Showed the largest gains. Access to the solved grid (privileged info) simplified the critic's task of validating intermediate moves, which otherwise requires implicit solving of the puzzle.
  • CodeIO: PVF utilized LOO group responses and rewards as privileged info, outperforming both VF and MEAN, demonstrating that group information can be effectively leveraged even without oracle answers.
  • Explained Variance: PVF consistently explained more return variance than VF, correlating with higher final rewards.

• TETHER Performance:

  • TETHER consistently outperformed the standard VF baseline across all tasks.
  • It successfully mitigated the degradation seen in VF when the critic was poorly fitted early in training.
  • In Reasoning Gym and MiniF2F, TETHER outperformed the MEAN baseline. In Sudoku, it narrowed the gap with MEAN, and in CodeIO, it matched MEAN.
  • Coefficient Dynamics: The adaptive $\rho$ started near 0 (relying on group means) and increased as the critic improved, confirming the method's ability to self-regulate based on critic accuracy.

Significance and Claims

The paper argues that value functions should be reconsidered as a practical component of LLM RL pipelines. The authors claim:

• Justification of Cost: By leveraging privileged information (PVF) and adaptive mixing (TETHER), the performance gains of value functions can outweigh their infrastructure costs, even in scenarios where critic-free methods are currently dominant.
• Robustness: TETHER provides a reliable pathway to introduce value functions into existing GRPO pipelines without the risk of training instability caused by a poorly fitted critic.
• Future Directions: The work suggests that extensive value pretraining on diverse policies and systematic tuning of the $\lambda$ parameter (for bias-variance trade-off) could further enhance the scalability and utility of value-based RL for LLMs.

The authors emphasize that their approach maintains the theoretical guarantees of unbiased policy gradients while offering practical improvements in training efficiency and final performance.