Rethinking Critic Learning in PPO: Understanding and Mitigating Value Flattening

AI summary

Overview

• Research area: Reinforcement learning for large language models (LLM post-training), specifically critic (value function) learning inside Proximal Policy Optimization (PPO).
• Technical level: Advanced. The paper combines empirical LLM training runs, a controlled FrozenLake study, and a formal decomposition of the critic loss and actor gradient error.
• Scope in one sentence: The paper diagnoses a systematic failure mode it names "Value Flattening" — PPO critics produce nearly flat value predictions within a response while Monte Carlo estimates of state value change sharply — explains two causes, and proposes a sparse critic-supervision method, SParse Proximal Policy Optimization (SP³O), that mitigates it and improves policy performance on Qwen3-4B-Base and Qwen3-8B-Base.

What This Paper Is About

PPO trains a critic to estimate the expected return from each intermediate state so that token-level advantages can be built. The authors find that in practice the critic does not track meaningful changes in expected success across steps within a single response: value predictions stay comparatively flat even where Monte Carlo (MC) value estimates from repeated continuations jump sharply. The paper identifies this as Value Flattening, traces it to two structural properties of standard critic training, and shows that supervising the critic at only a few well-separated states mitigates the problem and improves the learned policy.

Key Contributions

1. Identification of Value Flattening. The authors characterize a systematic mismatch in which policy-conditioned state values estimated from multiple Monte Carlo continuations change sharply within individual responses, while PPO critic predictions remain comparatively flat — across training checkpoints, and for both correct and incorrect responses.

2. Two causal explanations. They relate Value Flattening to (a) an implicit variance penalty in the critic's mean squared error loss, which — under terminal-only rewards with shared response-level targets — directly penalizes prediction variation across token positions, and (b) redundant updates from temporally correlated states, since LLM states differ by a single token and produce aligned representations and gradients.

3. A sparse-supervision fix (SP³O). They introduce SParse Proximal Policy Optimization, which applies the value loss only at a small set of well-separated states within each response while the critic still produces values at every generation state. Experiments use three supervised states per response, at response-relative positions 0.3, 0.6, and 0.9, plus a 0.95 state for responses of at least 6144 tokens.

4. Evidence that the fix improves policies, not just diagnostics. SP³O mitigates Value Flattening and consistently improves actor performance across Qwen3-4B-Base and Qwen3-8B-Base on mathematical and out-of-distribution reasoning suites.

Main Findings

• Value Flattening is observable in LLM reasoning. In a Qwen3-4B-Base PPO run on DAPO-Math-17k, MC values often exhibit sharp local transitions while critic predictions remain flat, and in some cases the critic moves in the opposite direction from the MC values. Adjacent-change analysis shows MC value changes span a broad range while critic changes stay concentrated near zero rather than following the diagonal.

• Value Flattening worsens as the state space grows. In a stochastic FrozenLake environment with the training configuration held fixed, critic predictions become progressively smoother and less accurate as the maze size n increases, with weaker local contrast and weaker agreement with ground truth.

• The critic loss contains an implicit variance penalty. Under terminal-only rewards with γ = λ = 1, every position in a response shares the same sampled terminal return, so the per-response MSE decomposes into a term fitting the mean prediction to the outcome plus the empirical variance of the predictions — the latter directly penalizing within-response value variation.

• Dense supervision produces redundant updates. Neighboring LLM states overlap in nearly their entire input. Hidden-representation alignment, gradient alignment, and update energy remain high across PPO training stages, and gradient cosine similarity decreases as the relative separation between token positions increases, so dense token-level supervision aggregates many similar updates.

• SP³O improves critic value resolution. Compared with PPO, SP³O more closely captures the direction and magnitude of MC value changes, reducing profile MSE by 36%, 11%, and 21% at 30%, 60%, and 90% response progress, respectively.

• SP³O improves critic optimization without losing response-level discrimination. In the response AUC versus within-response variance comparison across early (E), middle (M), and late (L) stages, PPO gains discrimination while losing within-response variation, whereas SP³O improves both. SP³O also raises the median effective rank of response-centered hidden states from 4.33 to 5.63, and shows lower RMS discrepancy between value-head gradients induced by terminal-return targets and by MC values.

• Sparse supervision stabilizes actor optimization. In online runs on Qwen3-4B-Base, SP³O shows smaller within-iteration update changes over most of training, higher validation accuracy and rollout reward than PPO after the early training stage, and longer responses.

• In-domain mathematical results (avg@32, Table 1). Qwen3-4B-Base averages: Base 17.95, PPO 37.60, GRPO 39.26, SP³O 45.57. Qwen3-8B-Base averages: Base 23.58, PPO 48.50, GRPO 47.91, SP³O 50.51. The reported gains over PPO reach 7.97 percentage points on in-domain mathematical reasoning.

• Out-of-distribution results (avg@4, Table 2). Qwen3-4B-Base averages: Base 19.26, PPO 51.95, GRPO 56.44, SP³O 59.28. Qwen3-8B-Base averages: Base 38.21, PPO 64.38, GRPO 64.91, SP³O 66.37. The reported gains over PPO reach 7.33 percentage points on out-of-distribution reasoning.

• Supervision density has an inverse relationship with benefit. Varying the number of supervised anchors K with Qwen3-4B-Base, configurations with K in {3, 4, 8} achieve higher mean training rewards than denser configurations and than standard PPO with token-level supervision. Performance is highest at K = 3 and stable through K = 8, but drops substantially at K = 16 and K = 64, approaching the dense PPO baseline. The authors note run-to-run variance is non-negligible.

• Anchor placement matters. At K = 3 on Qwen3-4B-Base, PPO baseline accuracy is 37.60, random placement is 36.59, anchors at 0.2/0.5/0.8 give 44.65, and anchors at 0.3/0.6/0.9 give 45.57. Adding the late-tail anchor improves performance from 44.10 to 45.57 and reduces repetition from 18.33 to 1.12.

Methodology in Plain English

The authors first build a diagnostic reference: instead of relying on the single sampled return the critic is trained on, they take a fixed intermediate state and independently sample many continuations from the same policy, averaging the terminal rewards. Under the paper's binary terminal-reward setting, this average is the empirical success rate from that state. They compare these Monte Carlo value estimates against the critic's predictions at the same states, both within single responses and across training checkpoints.

To test whether the pattern is specific to language models, they run a controlled stochastic FrozenLake experiment. There, a grid cell's true value is simply the probability of eventually reaching the goal, so critic predictions can be compared directly against ground truth while the maze size n is increased and everything else is held fixed.

They then analyze the training objective itself. Because rewards are terminal-only and γ = λ = 1, every token position in a response is supervised with the same response-level return. Writing out the per-response mean squared error shows that it splits into a term for the response mean plus the variance of the predictions within the response. Separately, they measure how similar critic hidden representations and gradients are for nearby token positions, since adjacent LLM states share almost all their tokens.

The proposed fix, SP³O, keeps the actor objective, rollout procedure, and return targets exactly as in standard PPO, and only changes which states receive the value loss — a small number of widely spaced states. The default uses three anchors per response (with an additional tail anchor for long responses), and the paper ablates the number of anchors, their placement, and the presence of the tail anchor.

Why This Matters

The paper reframes critic quality in LLM reinforcement learning as a within-response resolution problem rather than only a response-level prediction problem, and shows that a nearly free change to a few lines of the critic loss produces measurable policy gains. It also connects an LLM-specific failure to a phenomenon visible in a classic control environment, suggesting the underlying mechanism is not an artifact of language modeling alone.

Real-world applications that follow from this setting:

• Reasoning-model post-training: Training long chain-of-thought policies where credit must be assigned across thousands of intermediate steps.
• Mathematical and scientific problem solving: The paper's in-domain suite (AIME24, AIME25, AIME26, AMC23, MATH500, Minerva, OlympiadBench) mirrors the workloads of deployed math and STEM assistants.
• General-capability assistants: The out-of-distribution suite (ARC-C, MMLU-Pro, GPQA, AGIEval-English, BigBenchHard, ZebraLogic-Grid) reflects the broad reasoning and knowledge tasks users bring to general assistants.
• Reinforcement-learning infrastructure for LLMs: Any team running critic-based PPO at scale can adopt sparse value supervision without changing actor-side code, rollouts, or advantage computation.

Industry relevance: PPO with a learned critic is a standard post-training tool for large reasoning models, and critic-side supervision is a substantial share of training compute and engineering effort. A method that reduces supervised value-loss positions per response while improving outcomes is directly relevant to teams balancing training cost, stability, and downstream accuracy. The gains over PPO reported here (7.97 points in-domain, 7.33 points out-of-distribution) come alongside a moderate drop at K = 16 and K = 64 back toward the dense PPO baseline, which matters for anyone tuning supervision budgets.

Future Directions

• Understanding the density trade-off. Why do K = 3 to K = 8 anchors work best while K = 16 and K = 64 regress toward the dense PPO baseline? The paper notes the variance across runs is non-negligible, so the shape of this curve is not fully pinned down.
• Principled anchor placement. The paper shows well-spaced fixed anchors beat random placement, but the strongest tested schedule (0.3/0.6/0.9) is a simple uniform choice. Whether placement can be derived from state information rather than response position is left open.
• Extending beyond the terminal binary-reward setting. The analysis exploits γ = λ = 1 and terminal-only rewards. How Value Flattening and sparse supervision behave with per-step shaping, nonzero KL penalties, or denser rewards is not established.
• Generalizing across environments and methods. FrozenLake is used as a controlled demonstration and the LLM results are on two Qwen3-Base sizes with DAPO-Math-17k. Whether the same diagnosis and fix transfer to other task distributions, other model families, and critic-based methods other than standard PPO remains unresolved.

Target Audience

Researchers and engineers working on reinforcement learning for large language models, particularly those training or debugging PPO-style critic-based post-training pipelines. It is also useful for RL theorists interested in the interaction between dense supervision, temporal correlation, and value-function resolution, and for practitioners who need a low-cost, drop-in modification to critic supervision. Readers should be comfortable with PPO, generalized advantage estimation, value-function regression targets, and basic linear algebra such as effective rank.

Authors’ abstract

In reinforcement learning for large language models, Proximal Policy Optimization (PPO) commonly uses a critic to estimate state values and reduce the variance of policy updates. However, we uncover a systematic failure mode in PPO critics, which we call Value Flattening: state values, estimated from multiple Monte Carlo continuations, change sharply across intermediate states while critic predictions remain comparatively flat. We further observe this phenomenon in a controlled FrozenLake environment and find that it becomes more pronounced as the state space grows. Our theoretical and empirical analyses relate Value Flattening to an implicit variance penalty in the critic loss and redundant updates from temporally correlated states with similar gradients. Motivated by these findings, we introduce SParse Proximal Policy Optimization (SP$^3$O), which applies the value loss to only a few well-separated states in each response to mitigate both effects. Experiments on Qwen3-Base show that SP$^3$O with only three states supervised per response can mitigate Value Flattening and consistently improve the learned policy across model sizes and evaluation suites. Together, our results identify Value Flattening as an important yet overlooked failure mode of critic learning in standard PPO and show that a simple sparse supervision strategy can mitigate it.