Small Language Models as Judges for Rubric-Based Reinforcement Learning

AI summary

Overview

Research area: Natural Language Processing / LLM evaluation and reinforcement learning (rubric-based RL, LLM-as-a-judge, small language models).

Technical level: Intermediate. The paper assumes familiarity with reinforcement learning from verifiable rewards, reward models, and linear probing over hidden states, but the core arguments are stated in accessible terms.

Scope: This paper asks whether language models as small as 0.6B–1.7B parameters can replace large generative LLM judges when computing rubric rewards for reinforcement learning, and it introduces two benchmarks plus a reward-model study to answer that question empirically.

What This Paper Is About

Reinforcement learning with verifiable rewards works well when answers can be checked exactly, but many open-ended tasks (for example long-form research reports) can only be graded against instance-specific rubrics covering coverage, factuality, source use, and constraints. Scoring responses against such rubrics is expensive, because every RL step may require a generative LLM judge with 7B parameters or more, often behind a proprietary API. The paper asks whether much smaller language models, used carefully, can act as reliable and efficient criterion-level rubric judges instead.

Key Contributions

1. Two new pointwise rubric-evaluation datasets. PointRubric, adapted from OpenRubrics, and RaR-Science-Static, built from RaR-Science, both provide explicit rubric criteria, multiple responses per prompt, per-criterion satisfaction labels from GPT-4o, and weighted rubric scores.
2. A three-way comparison of readout methods for small judges. Generative verdicts, Yes/No Logprob scoring, and hidden-state Probes are compared on the same Qwen3 backbones at 0.6B, 1.7B, 4B, and 8B parameters, in both base and SFT (supervised fine-tuning) forms.
3. A demonstration that the Qwen3-1.7B Probe judge works as a GRPO reward model, training a Qwen3-4B-Base policy from 0.232 to 0.643 on RaR-Science rubric score, versus 0.594 for an 8B Generative judge baseline, while the baseline requires 10.7× more cumulative reward-judge time.
4. Transfer experiments showing the improved policy also gains on GPQA-Diamond, and that a RaR-Science-trained Probe retains criterion-level agreement on RaR-Medicine without target-domain fitting.

Main Findings

• Probe judges dominate on the controlled benchmark (PointRubric). The Qwen3-1.7B base Probe reaches 0.875 weighted criterion accuracy versus 0.518 for Generative and 0.766 for Logprob at the same scale, and 0.802 at 0.6B versus 0.370 (Generative) and 0.582 (Logprob). SFT helps the 1.7B Generative judge reach 0.902, but SFT Probe stays the strongest readout at 0.6B (0.820) and remains competitive at other scales (1.7B: 0.845; 4B: 0.906; 8B: 0.912).

• Readout choice matters more than model size on RaR-Science-Static. At 1.7B, Generative and Logprob reach only 0.443 and 0.449 macro-F1, while Probe reaches 0.835. Across model sizes, the base Probe goes 0.793 (0.6B), 0.835 (1.7B), 0.851 (4B), 0.864 (8B), versus 0.413/0.433, 0.443/0.449, 0.496/0.738, and 0.609/0.756 for Generative/Logprob respectively.

• The signal is hidden, not absent. Among 4,518 held-out criterion decisions, 33.0% are correct only under Probe, whereas 7.8% are correct only under Generative.

• Probe judgments survive a change of response generator. On an extended 1,200-response bank using Mistral-7B-Instruct-v0.3 and OLMo-2-1124-7B-Instruct responses, the unchanged 1.7B Probe reaches 0.741 macro-F1, versus 0.394 for Generative and 0.424 for Logprob.

• Probe scores induce strong pairwise rankings. The Qwen3-1.7B Probe reaches pairwise accuracies of 0.924 on PointRubric and 0.888 on RaR-Science-Static, while Rubric-RM-4B and Rubric-RM-8B top out at 0.322 and 0.487 when parse failures count as incorrect, or 0.575 and 0.670 on parseable outputs only.

• Probes are data-efficient and simple. A linear last-token probe reaches 0.805 macro-F1 with 100 training questions and 0.834 with 1,000, while MLP heads (MLP-32: 0.834; MLP-64: 0.840) and last-4-layer averaging (0.839) add little; mean pooling is much worse at 0.761.

• The 1.7B Probe is the best RL reward model in the matched comparison. With the same Qwen3-4B-Base actor and GRPO configuration, the 1.7B Probe lifts RaR-Science rubric score from 0.232 to 0.643 (+0.411), ahead of 0.562 (0.6B Probe), 0.588 (4B Probe), 0.506 (8B Probe), and 0.594 (8B Generative).

• Static accuracy does not predict downstream reward quality. The 4B and 8B Probes score slightly higher on the fixed bank, but on-policy diagnostics suggest their rewards are more saturated on actor rollouts, reducing within-group reward variation for GRPO.

• Human audits support the comparison. The probe-reward policy is preferred over the base actor on 72 of 100 examples by human annotators and 80 of 100 by GPT-4o; after removing ties, human and GPT-4o agree on 69 of 72 cases (95.8%). Separately, a manual audit of 100 held-out PointRubric responses (600 decisions, 50 paired comparisons) achieves 90.9% weighted criterion agreement and 96.0% pairwise preference agreement with GPT-4o.

• Policy transfer is positive but modest. On the 198-question GPQA-Diamond split across four answer-order runs, the probe-reward policy improves accuracy from 0.335 ± 0.023 to 0.388 ± 0.037, a gain of +0.053.

• Judge transfer across domains is asymmetric. A RaR-Science-trained Probe reaches 0.718 macro-F1 on RaR-Medicine without Medicine labels for fitting, below the Medicine-trained in-domain probe (0.782); Medicine→Science gives 0.702 versus 0.754 for Science→Science.

• Efficiency. By the comparison checkpoint, Generative uses 89,912.1 seconds of cumulative judge time versus 8,389.9 seconds for Probe (10.7×), with validation judge time of 492.0 s versus 31.1 s (15.8×); elapsed RL time is 11h 11m 56s versus 8h 51m 47s (1.26×). The reward Probe is trained once on 500 GPT-4o-labeled responses (3,766 criterion labels, approximately $2.43) and can then be reused.

Methodology in Plain English

The researchers start by defining the task: a judge is given one prompt, one candidate response, and one rubric criterion, and must decide whether the response satisfies that criterion. Those item-level decisions are then combined using the rubric's weights into a single score that can serve as an RL reward.

Because existing datasets do not supply that full structure, they build two. PointRubric takes 3,000 OpenRubrics questions with hard-rule constraints and rewrites each into six criteria (two hard, four soft). For each prompt they generate four candidate responses — a full-score response, a low-score response, a rubric-guided partial response, and an unguided response — then use GPT-4o to label every response–criterion pair. After filtering, PointRubric holds 1,042 questions, 4,168 responses, and 25,008 response–criterion labels, split by question into 417 training, 104 development, and 521 held-out questions. RaR-Science-Static contains 1,500 randomly sampled science questions, each keeping the RaR-Science reference answer plus one greedy-decoded answer from Qwen3-4B, giving 3,000 responses; rubrics have between 6 and 11 criteria (7.52 on average), split into 1,000 training, 200 development, and 300 test questions.

They then compare three ways to read a judgment out of a frozen Qwen3 backbone. Generative judges emit a binary verdict in text. Logprob judges score the next-token probabilities of "Yes" and "No" and use the log-probability margin. Probe judges keep every language-model parameter frozen and train only a lightweight linear classifier on the final non-padding-token hidden state, using binary cross-entropy against GPT-4o labels, with layer selection and thresholding done on the development split. They also test supervised fine-tuning of the backbone and reapply all three readouts afterwards.

Finally, they plug the best judge — the Qwen3-1.7B linear last-token Probe, fitted on 450 training and 50 development responses — into GRPO as the reward model, holding the actor, optimizer, schedule, response budget, aggregation rule, and GPT-4o evaluator fixed so that only the reward judge varies. They then check transfer on GPQA-Diamond and across science/medicine rubric splits, and measure judge time, validation judge time, and wall-clock RL time.

Why This Matters

Impact on research. The paper reframes reward-model design for rubric-based RL: instead of assuming a bigger generative judge is better, it shows that a frozen small backbone with a trained linear probe can be the stronger and far cheaper reward signal. It also supplies two reproducible benchmarks that isolate the criterion-level structure rubric RL actually needs, and it documents a case where static judge accuracy and downstream RL utility diverge — a caution for how reward models are selected.

Real-world applications.

• Training agents for long-form research reports, where coverage, factuality, and source use must be graded against explicit criteria rather than exact answers.
• Grading and improving expert-domain responses (the paper demonstrates science and medicine rubrics) at a fraction of the cost of repeated large-judge calls.
• Reducing reliance on proprietary judge APIs, since a frozen local probe avoids dependence on a fixed provider snapshot and can be reused across RL runs.
• Lower-cost experimentation in resource-constrained labs, given the approximately $2.43 one-time labeling cost reported for the RL reward probe.

Industry relevance. Repeated rubric grading is a dominant cost in rubric-based RL pipelines. The reported 10.7× cumulative judge-time and 15.8× validation judge-time reductions, combined with a higher-scoring policy (0.643 versus 0.594), give a concrete quality-per-cost argument for deploying small local probes instead of large generative judges.

Future Directions

• Broaden the human reference beyond GPT-4o. The authors note their supervised target is GPT-4o criterion scoring rather than human ground truth, and that broader human evaluation is needed to characterize agreement across tasks and domains.
• Test harder domains. The paper's strongest evidence is on RaR-Science, with additional checks on GPQA-Diamond and RaR-Medicine; dialogue, long-form instruction following, and safety-critical expert settings may need separate calibration and validation.
• Generalize beyond one policy family and one rubric-RL setup. The matched RL experiments fix the actor, aggregation rule, and training configuration, so the reward-overoptimization checks remain setting-specific.
• Explain and control reward saturation. The finding that larger Probes assign more saturated rewards and thereby hurt GRPO is offered as a possible explanation, pointing to open work on reward-shaping and judge selection driven by downstream policy improvement rather than static accuracy alone.

Target Audience

Researchers and engineers working on RLHF/RLVR and rubric-based RL, reward-model and LLM-as-a-judge practitioners, and evaluation scientists who need cheaper, reproducible criterion-level graders. It is also useful for readers interested in interpretability-adjacent techniques, since the central result is that hidden states carry evaluative signal that generation and verbalizer scoring do not reliably surface.

Authors’ abstract

Rubric-based reinforcement learning extends RL beyond tasks with exact answers or rule-based verifiers by scoring responses against instance-specific criteria. However, this makes reward computation expensive: training requires repeated rubric judging, often with proprietary APIs or local generative LLM judges with 7B parameters or more. We study whether smaller language models can serve as efficient and reliable rubric-based judges. To make this question measurable, we construct PointRubric and RaR-Science-Static, two pointwise rubric-based evaluation datasets with instance-specific criteria and itemwise satisfaction labels. We compare three ways of extracting criterion-level judgments from small models: Generative verdicts, Yes/No Logprob margins, and Probe judges. Across both datasets, the Qwen3-1.7B Probe judge achieves the strongest criterion-level agreement among these methods, outperforming Generative and Logprob judges. Used as a GRPO reward model, it trains a policy from 0.232 to 0.643 on RaR-Science rubric score, compared with 0.594 for an 8B Generative judge baseline, while the baseline requires 10.7$\times$ more reward-judge time. Task and domain transfer experiments further suggest that Probe judges preserve criterion-level reward structure across settings.