Start Classifying: Categorical Critics for LLM Reinforcement Learning
Zhijian Zhou, Long Li, Xuan Zhang, Zongkai Liu, Yulei Qin, Ke Li, Xing Sun, Xiaoyu Tan, Chao Qu, Yuan Qi
COLM 2026
cs.LG
2026-08-03
Replacing PPO's MSE value head with a Gaussian-smoothed categorical classifier (actor unchanged) lifts Qwen2.5 math reasoning pass@256 by 9.6 points and avg@256 by 2.86.
When PPO trains a large model to reason, a value function (the critic) scores each intermediate state of the reasoning trace, estimating how much reward it can still earn from that point. Subtract that score from the actual return and you get the advantage, the signal PPO uses to update the policy. The critic is almost always a single scalar output trained with MSE regression, the obvious default.
The catch is not that scalar regression is wrong. Under RLVR's binary rewards (correct = 1, wrong = 0) with γ = λ = 1, the true value of an intermediate state is just the probability of eventually solving the problem, and scalar MSE does converge to that in principle. The real trouble is optimization and calibration: a small overestimate by the critic gets amplified into a lopsided advantage. If the critic thinks a doomed trajectory still has a 60% chance, that trajectory earns a negative advantage of -0.6 while a successful one earns only +0.4. Penalties outweigh rewards, and the policy update skews. The authors measure a negative-to-positive advantage ratio of about 2.6× for the standard MSE critic.
HL-Gauss PPO touches only the critic's output layer and loss; the actor, rollouts, and sampling are unchanged. Three steps:
Why this design: sharp one-hot targets push predictions toward the corners of the probability simplex, where the Fisher information matrix is nearly singular and optimization is ill-conditioned. Gaussian smoothing keeps targets in the interior and the critic from becoming overconfident. One-hot, two-hot, and a Bernoulli two-bin control together separate the roles of classification itself, a larger output head, and target smoothing. Smoothing is what matters; a bigger head or binary classification alone falls short.
Mathematical reasoning (Qwen2.5-Math-7B):
| Method | avg@256 | pass@256 |
| PPO (MSE critic) | 15.88 | 38.48 |
| DAPO (critic-free) | 16.00 | 42.34 |
| HL-Gauss PPO | 18.74 | 48.06 |
Average score rises 2.86 and pass@256 rises 9.58, with the largest single-dataset gain on AIME25 (avg@256 from 18.60 to 23.80). The trend holds on Qwen3-4B-Base, tool-augmented math, and Search-R1 retrieval QA, where the seven-dataset mean@1 climbs from 44.91 to 47.13.
The calibration numbers are the most direct evidence. On AIME24 reasoning prefixes (oracle success probability V estimated from 256 Monte Carlo rollouts), HL-Gauss cuts the Brier score from 0.203 to 0.164, ECE by 18%, and maximum calibration error by 28%. In advantage terms, the MSE critic hands failed rollouts a mean negative advantage of -0.597 but successful ones only +0.227, a 2.63× ratio; HL-Gauss sits at -0.466 versus +0.458, nearly symmetric. And as training proceeds, the MSE skew worsens while HL-Gauss stays close to one.
For anyone running PPO-style reasoning training, this is a near-zero-cost upgrade. The 101× larger value head is still only about 0.006% of a 7B model, and per-step wall-clock time is indistinguishable from vanilla PPO. A different loss function buys a consistent gain, and the authors have open-sourced the code.
It also surfaces a point the field has underweighted: under RLVR the critic must estimate without systematically penalizing the right trajectories and rewarding the wrong ones. A critic can look fine on mean-squared error yet still bias high on failing states, quietly sabotaging the policy.
The authors concede that every experiment runs on the Qwen family at 4B–7B scale, with no larger models or other architectures tested, and that the calibration and advantage analyses are diagnostic rather than proof that advantage symmetry is the sole cause of the gains.
A few colder reads. First, the gains concentrate in pass@k while avg@k moves modestly (2.86), suggesting the method mostly helps the model occasionally crack hard problems rather than sharpening its single best answer. Second, the Bernoulli two-bin control is mildly awkward: it pushes pass@256 to 45.67, close to HL-Gauss's 48.06, yet its avg@256 dips below MSE, so if you only care about best-of-k a binary head may suffice without 101 bins. Third, the optimal bandwidth σ = 0.75Δ is sharp; widening to σ = Δ drops the average to 17.00, so every task's reward range needs its own support and bandwidth rather than a universal default. Fourth, there is no head-to-head with DisPPO, the distributional LLM PPO; the distinction is drawn by positioning, not by experiment.