Training LLM Judges from Language Feedback via Position-Selective Self-Distillation

AI summary

Overview

Research area: Natural Language Processing — specifically training LLM judges / reward models using reinforcement learning and self-distillation, with a focus on subjective evaluation tasks.

Technical level: Advanced. The paper assumes familiarity with on-policy reinforcement learning for LLMs (GRPO and its variants), knowledge distillation, KL divergence objectives, and token-level entropy analysis.

Scope in one sentence: The paper proposes selecting which token positions to distill when training an LLM judge from natural language feedback, using a per-position entropy shift between a student and a feedback-conditioned teacher to mask out the positions that encourage memorization rather than semantic understanding.

What This Paper Is About

The dominant way to train an LLM judge is outcome-supervised reinforcement learning (e.g., GRPO, Dr. GRPO, DAPO), which rewards an entire rollout based only on whether the final verdict is correct. This gives no separate credit to the tokens where the model actually chooses and weighs which evaluation criteria to invoke — a distinction that matters most on subjective tasks, where picking the wrong criterion produces a wrong verdict even when the criterion is applied rigorously. The paper's goal is to exploit the natural language feedback (preference rationales) that already accompanies preference labels, using self-distillation to give dense per-position supervision, and then to figure out which of those positions are actually worth supervising.

Key Contributions

1. A characterization of two distillation regimes via per-position entropy shift. Defining the entropy shift as the entropy of the student's next-token distribution minus that of the feedback-conditioned teacher's, the authors identify context sharpening (large positive shift: the teacher concentrates probability on a particular feedback-aligned criterion expression) and context spreading (large negative shift: the teacher distributes probability across multiple feedback-aligned alternatives). They interpret sharpening as encouraging memorization of one criterion expression and spreading as promoting semantic understanding.

2. Position masking based on entropy shift. A per-generation mask that drops the top-ρ fraction of positions ranked by entropy shift and keeps the lower tail, implemented as a variant of the on-policy reverse-KL self-distillation loss. Setting ρ = 0 recovers naive self-distillation.

3. Evidence that criterion-choice tokens carry disproportionate distillation loss. On 102 HelpSteer3-Preference validation rollouts, criterion-selection spans are 16.7% of response tokens but 44.4% of total reverse-KL mass (~4.0× the per-token KL of other positions), and criterion-name positions are 0.4% of tokens but 11.1% of reverse-KL mass.

4. Empirical gains over outcome-supervised RL and naive self-distillation, including out-of-distribution judge benchmarks, a criterion-diversity analysis, and a downstream pairwise-selection tournament.

Main Findings

• Both entropy-shift tails carry most of the distillation loss. Across 104,046 response positions from 102 HelpSteer3-Preference validation rollouts generated by Qwen3-30B-A3B-Instruct-2507 (34 examples each from code, general, and STEM), per-position reverse KL has a U-shaped relationship with the entropy shift: partitioning each rollout into bottom 30% / middle 40% / top 30%, both tails carry approximately 20–28× the per-position KL of the neutral middle.

• Criterion-name positions concentrate in the tails. Criterion-name positions show 5.7× larger mean |entropy shift| than other response positions, and 94.8% of them fall within the two extreme 30% tails, which together contain 60% of response positions.

• Sharpening aligns the teacher's top-1 criterion with the rationale. At criterion-name positions with |entropy shift| ≥ 0.50 (87 sharpening positions from 56 rollouts), the teacher's top-1 criterion name has mean similarity 0.713 to the decisive criteria under the matched rationale versus 0.647 under a control that swaps in another example's rationale — a paired difference of +0.066 with 95% CI [0.033, 0.100]. Example: the phrase "correctly includes" drives the teacher to Correctness with top-1 probability 0.97; "difficult to read" drives it to Readability with top-1 probability 1.00.

• Spreading reopens criterion choices across aligned alternatives. At 71 spreading positions from 51 rollouts, mean similarity over the teacher's top-10 criterion names is 0.646 under the matched rationale versus 0.610 under the control — a paired difference of +0.037, 95% CI [0.022, 0.051]. Example: a student committed to Originality (top-1 probability 0.83) while a rationale emphasizing "factually verifiable" details shifted the teacher toward Correctness, Error, and Fidelity.

• Masking higher-entropy-shift positions improves overall accuracy at both scales. Total average goes from 77.70 to 79.04 at Qwen3-4B-Instruct and from 80.17 to 82.50 at Qwen3-30B-A3B-Instruct when comparing naive SD to SD+mask. The largest per-subcategory gains include 30B RM-Bench Code (76.56 to 80.07), 30B Chat (76.74 to 80.53), and 30B RewardBench v2 Precise IF (40.62 to 48.75).

• Self-distillation beats outcome-supervised RL on subjective subcategories by 2–9 percentage points, but not on objective ones. On objective subcategories Dr. GRPO stays competitive: at 30B it matches the SD variants on RM-Bench Math (95.59 for Dr. GRPO versus 94.45 for SD and 95.63 for SD+mask) and beats them on RewardBench v2 Math (87.43 versus 81.42 and 84.15). On subjective ones the gap flips: 30B Chat 74.07 (Dr. GRPO) versus 76.74 (SD) and 80.53 (SD+mask); 30B Factuality 71.79 versus 78.53 and 79.16; 30B Focus 80.00 versus 86.26 and 85.66; 4B Chat 68.91 versus 75.71 and 77.95; 4B Factuality 62.74 versus 68.42 and 70.74.

• The 30B SD+mask judge outperforms leading trained judge baselines on overall accuracy and is comparable to Claude-Sonnet-4. It exceeds RM-R1-DeepSeek-Distilled-Qwen-32B, Llama-3.3-Nemotron-Super-49B-GenRM, and RationaleRM-30B on overall benchmark accuracy. In training curves on RM-Bench (debiased EMA with β = 0.6, evaluated every 20 training steps), SD+mask leads naive SD and Dr. GRPO for both Qwen3-4B-Instruct and Qwen3-30B-A3B-Instruct, and surpasses DeepSeek-R1 and Claude-Sonnet-4 at the 30B scale.

• SD+mask invokes a broader repertoire of evaluation criteria. Clustering criteria extracted from 300 samples each of RM-Bench, RewardBench v2, and HelpSteer3-Preference validation with complete-link agglomerative clustering over Qwen3-Embedding-8B vectors, SD+mask produces more distinct criterion clusters than naive SD at every threshold, with the gap widening from 4 clusters at cosine threshold t = 0.65 to 26 at t = 0.85.

• Gains transfer to a downstream selection task. In a multi-round best-of-eight tournament on Arena-Hard-v2 creative-writing prompts (GLM-4.5-Air generating eight responses at temperature 1.0, top-p 1.0, max 8K new tokens), Claude Sonnet 5 scored the selected responses against the reference responses at a win rate of 0.556 for a Dr. GRPO-trained selector, 0.612 for naive SD, and 0.632 for SD+mask (ρ = 0.7) — a 7.6 percentage-point improvement over Dr. GRPO and 2.0 over naive SD.

• The masking direction matters. In a selector ablation at matched ρ = 0.7 on Qwen3-30B-A3B-Instruct, total average accuracy is 82.50 for the proposed top-entropy-shift mask, versus 79.23 for masking the bottom fraction (direction flip), 79.73 for masking by |entropy shift| (both tails), 79.55 for random masking, 79.57 for student-entropy masking, and 78.61 for teacher-entropy masking.

• Mask fraction sweep supports aggressive masking. Sweeping ρ ∈ {0, 0.3, 0.5, 0.7} on Qwen3-30B-A3B-Instruct, total average rises 80.17 → 80.63 → 80.68 → 82.50, with ρ = 0.7 achieving the highest overall accuracy among tested fractions. The same pattern holds at 4B.

• Results with an alternative feedback format are reported but not in the main text. The authors also generated a per-example rubric for every HelpSteer3-Preference sample using Claude Opus 4.8 (average 2,062 characters, approximately 7.7 times longer than the original preference rationales); those results appear in Appendix B.

Methodology in Plain English

The setup is a pairwise judge: given a prompt and two candidate responses, the model produces an intermediate reasoning trace (which criteria it invokes and how it applies them) and then a final verdict, A or B. Each training example also carries a one- or two-sentence human-written rationale that explains why one response was preferred — feedback that outcome-only RL never uses.

The researchers use the same model in two roles. As the student, it sees only the prompt and generates a rollout. As the teacher, it additionally sees the preference rationale — a piece of context the student never gets at training or deployment time. Training minimizes a reverse KL divergence between the student's next-token distribution and the teacher's at every position in the rollout, so the teacher's distributional signal at criterion-naming tokens directly reflects the rationale.

The distinctive step is deciding which positions to include in that loss. For each position they compute the entropy shift: the entropy of the student's next-token distribution minus the entropy of the teacher's. A large positive shift means the rationale made the teacher more confident and concentrated it on one criterion expression; a large negative shift means the rationale made the teacher less confident and spread probability across several alternatives. Since the positive-shift positions look like one specific wording being memorized while negative-shift positions look like multiple valid alternatives being preserved, they mask out the largest-shift positions and keep only the bottom fraction, controlled by a single hyperparameter ρ (the main experiments use ρ = 0.7, keeping the 30% of positions with the lowest entropy shifts). The mask is detached from the computation graph.

To check the interpretation, they annotate criterion-selection spans and criterion-name positions with GPT-5.5, compare the teacher's top candidate criterion names against the rationale's decisive criteria using Qwen3-Embedding-8B cosine similarity (with a swapped-rationale control and 2,000 bootstrap resamples at the rollout level), and measure criterion diversity of the trained judges by clustering extracted criteria.

Training data is a cleaned HelpSteer3-Preference: rows with domain == "multilingual" or overall_preference == 0 (tie) are filtered out, within-split deduplication by sha1(context, response1, response2) keeps the first occurrence (upstream contains roughly 35% byte-identical duplicate rows after the first step), and cross-split deduplication drops validation rows whose hash appears in train (upstream has roughly 918 of roughly 1,553 filtered-deduped validation rows byte-identical to a train row). Evaluation is on RM-Bench and RewardBench v2, both out-of-distribution relative to training.

Why This Matters

Impact on research. The paper reframes token-level supervision in post-training as a selection problem with a directional structure: it is not that high-signal positions are simply the uncertain ones (student-entropy and teacher-entropy masking both underperform here), but that the sign of the teacher–student entropy shift carries meaning. This gives a concrete, computationally cheap criterion for choosing imitation targets in on-policy self-distillation, and it connects the choice/rigor distinction in judging to measurable token statistics.

Real-world applications.

• Reward models for RLHF pipelines, where a judge decides which on-policy outputs get reinforced; the tournament result suggests better selection of preferred responses.
• Standalone LLM-as-a-judge evaluation, particularly for open-ended or subjective quality where the choice of criteria drives the verdict.
• Preference data annotation leverage, since preference rationales are already collected in many datasets and this method puts them to use rather than discarding them.
• Model routing or response selection in multi-candidate generation systems, where repeated pairwise comparisons determine the final output.

Industry relevance. The training data is a public preference dataset, the base models are open instruct models, and the method adds a masking step to a standard distillation loop, so it is implementable without new annotation. The 30B SD+mask judge is reported as outperforming leading trained judge baselines and remaining comparable to Claude-Sonnet-4 on the evaluated benchmarks, which is directly relevant to teams building in-house reward models.

Future Directions

• Full RLHF validation. The authors note that the strongest downstream test would be training separate policies with

Authors’ abstract

We study training LLM judges from natural language feedback, especially for subjective tasks where the verdict depends strongly on which evaluation criteria the judge invokes and how it weighs them. The dominant approach, outcome-supervised RL (e.g., GRPO), credits every token in the rollout with a single scalar determined only by the accuracy of the final verdict, providing no separate credit at the criterion-choice tokens and ignoring the rich language feedback (e.g., preference rationales) that naturally accompanies preference labels. Self-Distillation (SD) is one natural way to use this language feedback: the same model, conditioned on this feedback, acts as a teacher providing dense, position-level supervision. However, not all positions carry equally useful signal. Using the per-position entropy shift between teacher and student, we identify two regimes: context sharpening, where the teacher concentrates probability on a particular feedback-aligned criterion expression, and context spreading, where the teacher distributes probability across multiple feedback-aligned alternatives. We interpret these patterns as follows: sharpening encourages memorization of a particular criterion expression, whereas spreading promotes semantic understanding by preserving these alternatives. Motivated by this asymmetry, we introduce position masking based on the entropy shift that retains the lower tail of the entropy-shift distribution. Experiments show that masking higher-entropy-shift positions improves out-of-distribution generalization over naive SD. The resulting self-distilled judges outperform judges trained with outcome-supervised RL by 2-9 percentage points on the evaluated subjective subcategories, while remaining competitive on objective ones.