💬 NLP

Reinforcement Learning-based Knowledge Distillation with LLM-as-a-Judge

This paper proposes a reinforcement learning framework that utilizes a single-token LLM-as-a-Judge to provide efficient, label-free reward signals for knowledge distillation, significantly enhancing reasoning capabilities on math benchmarks even when combined with verifiable rewards.

Original authors: Yiyang Shen, Lifu Tu, Weiran Wang

Published 2026-04-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yiyang Shen, Lifu Tu, Weiran Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but tiny, robot (the Student) how to solve complex math problems. Usually, to teach a robot, you need a giant textbook with every single problem and its correct answer written out perfectly. This is called "ground truth."

But what if you don't have that textbook? What if you have millions of math problems, but no answers? Or what if getting the answers is too expensive?

This paper proposes a clever new way to teach the robot without needing the answer key. Here is the breakdown using simple analogies:

1. The Problem: The "Answer Key" Bottleneck

Traditionally, to make AI smarter, researchers use Reinforcement Learning (RL). Think of this like training a dog.

The Old Way: You give the dog a trick. If it gets it right, you give it a treat (a reward). If it gets it wrong, you give it nothing. But to know if it's right, you need to know the answer beforehand. This limits you to problems where you already have the answers.
The Limitation: For many real-world tasks, we don't have the "correct answer" ready. We have the questions, but not the solutions.

2. The Solution: The "Tough Teacher" (LLM-as-a-Judge)

The authors introduce a new character: a Super-Intelligent Teacher (a large AI model).

Instead of writing out the answer for the student, the Teacher just looks at the student's work and says, "Yes, that's good" or "No, that's wrong."
The Magic Trick: The Teacher doesn't just say "Yes" or "No" with a simple button press. It calculates a score based on how confident it is.
- If the student's reasoning looks very convincing, the Teacher gives a high score (like 0.9).
- If it looks shaky, the score is low (like 0.2).
- If it's nonsense, the score is 0.

This score acts as the "treat" for the robot. The robot tries to solve the problem, gets a score from the Teacher, and learns to do better next time to get a higher score.

3. The "Yes/No" Shortcut (Why it's fast)

Usually, asking an AI to grade a paper takes a long time because it has to write a long essay explaining why it's right or wrong.

The Innovation: The authors forced the Teacher to be very strict. The Teacher is only allowed to output one single word: "Yes" or "No."
The Analogy: Imagine a judge in a courtroom. Instead of writing a 50-page verdict, they just hold up a green card (Yes) or a red card (No).
Why this matters: Even though the judge only says one word, the AI can look at how confident it was in that word. If it was 99% sure it was "Yes," the reward is huge. If it was only 51% sure, the reward is small. This makes the grading process incredibly fast and cheap, allowing them to train on millions of problems.

4. The Training Process: The "Gym"

The paper tests this on small models (the students) using math problems.

Step 1 (Warm-up): The student practices a little bit with a few problems where they do have the answer key (Supervised Fine-Tuning).
Step 2 (The Gym): The student is thrown into a massive gym with millions of problems where they don't have the answers.
Step 3 (The Workout): The student tries to solve a problem. The "Tough Teacher" looks at the attempt and gives a score. The student adjusts its brain to get a higher score next time.
Step 4 (The Result): The student learns to think deeply and logically, not just by memorizing answers, but by understanding what "good reasoning" looks like.

5. The Results: Small Fish, Big Ocean

The researchers found that:

Small models got much smarter: A tiny model (125 million parameters) improved its math skills by 5–10% just by listening to the Teacher, even without seeing the correct answers.
It works on new things: The student didn't just memorize the practice problems; it learned to solve new, harder types of math problems it had never seen before.
It scales up: Even a bigger model (6.7 billion parameters) got a slight boost, proving this method works for all sizes of AI.

Summary

Think of this paper as a new teaching method for AI. Instead of forcing the AI to memorize the answer key (which is expensive and limited), they let a "Tough Teacher" grade the AI's homework on the fly. The AI learns to think better by trying to please the Teacher, turning a mountain of ungraded homework into a powerful training tool.

In short: They taught a small robot to be a math genius by letting a super-smart robot grade its homework with a simple "Yes/No" score, eliminating the need for a pre-written answer key.

1. Problem Statement

Reinforcement Learning (RL) has proven effective in enhancing the reasoning capabilities of Large Language Models (LLMs), particularly in domains like mathematics and logic. However, existing RL approaches (e.g., RL with Verifiable Rewards, RLVR) rely heavily on ground truth labels or verifiable outcomes (e.g., exact answer matching) to generate reward signals. This creates two major bottlenecks:

Data Scarcity: High-quality labeled datasets are expensive and limited, restricting the scale of RL fine-tuning.
Knowledge Distillation Limitations: Traditional Knowledge Distillation (KD) via Supervised Fine-Tuning (SFT) often results in surface-level token mimicry, failing to transfer the underlying reasoning processes of a teacher model to a smaller student model.

The paper addresses the challenge of performing label-free knowledge distillation using RL, where a smaller student model learns from a larger teacher model's reasoning capabilities without access to ground truth labels for the training data.

2. Methodology

The authors propose a novel RL framework that transforms the "LLM-as-a-Judge" paradigm from a static evaluator into an active reward generator. The core components are:

A. Problem Setup & Training Pipeline

Warm-up: The student model undergoes Supervised Fine-Tuning (SFT) on labeled data to establish a baseline reasoning capability.
RL Fine-Tuning: The model is further optimized using Proximal Policy Optimization (PPO). The student generates reasoning trajectories (Chain-of-Thought, specifically in Python code for Program-of-Thought) for given questions.
Reward Signals: The framework integrates three types of rewards:
1. Verifiable Reward ( $L_{PPO}^{VR}$ ): Used on labeled data. It provides binary/partial rewards based on whether the generated code compiles and produces the correct ground-truth answer.
2. LLM-as-a-Judge Reward ( $L_{PPO}^{YoN}$ ): Used on unlabeled data. A powerful teacher LLM evaluates the student's generated reasoning. Crucially, the judge is constrained to output only "Yes" or "No". The reward is computed efficiently by extracting the logits for these two tokens and applying a sigmoid function to the log-odds ratio, creating a continuous scalar reward without generating full reasoning chains.
3. Reranking Model Reward ( $L_{PPO}^{Rerank}$ ): A pretrained reranking model provides an additional continuous reward signal based on the probability of correctness.

B. The "Yes/No" Reward Mechanism

To ensure efficiency, the judge LLM is prompted to answer "Yes" or "No" to the question: "Does the Python solution correctly solve the problem?"
The reward $s$ is calculated as:
$s(p(\text{Yes}), p(\text{No})) = \sigma(u) \cdot \mathbb{I}\{\sigma(u) \ge a\}$
Where $u = \frac{\log p(\text{Yes}) - \log p(\text{No})}{\tau}$ , $\sigma$ is the sigmoid function, and $a$ is a threshold. This allows for single-token reward computation, making the process scalable for large batches of unlabeled data.

C. Overall Objective

The final loss function is a weighted convex combination of the three components:
$L_{overall} = \lambda \cdot L_{PPO}^{VR} + \mu \cdot L_{PPO}^{YoN} + \rho \cdot L_{PPO}^{Rerank}$
where $\lambda + \mu + \rho = 1$ . This allows for semi-supervised training, leveraging both labeled (verifiable) and unlabeled (judge-based) data.

3. Key Contributions

Label-Free RL Distillation: The paper introduces a scalable RL framework that enables knowledge distillation using vast amounts of unlabeled data, removing the dependency on ground truth supervision for the distillation phase.
Efficient LLM-as-a-Judge Reward: Unlike previous methods that require generating full explanations or complex preference datasets, this approach uses a single-token output ("Yes"/"No") from a judge LLM to derive continuous rewards. This significantly reduces computational overhead.
Superiority over SFT and Standard RLVR: The method demonstrates that active exploration guided by an LLM judge allows the student model to internalize latent reasoning capabilities better than surface-level SFT or RL trained only on verifiable rewards.
Scalability: The approach is validated on both small models (125M, 350M) and larger models (6.7B), showing consistent improvements.

4. Experimental Results

The experiments were conducted on mathematical reasoning benchmarks (GSM8K, SVAMP, GSM-Plus, GSM-Symbolic) using models like Galactica-125M, CodeGen-350M, and Galactica-6.7B.

Small Models (Galactica-125M):
- Performance Gain: The proposed method achieved a 5–10 percentage point improvement over using only verifiable rewards (RLVR) on GSM8K.
- Semi-Supervised Setup: Combining labeled data with unlabeled augmented data (GSM8K-aug) using the judge reward yielded a 40.71% accuracy on GSM8K, compared to ~30% for RLVR alone.
- Out-of-Domain Robustness: The model showed significant gains on difficult, out-of-distribution datasets like GSM-Symbolic and SVAMP, indicating better generalization.
Large Models (Galactica-6.7B):
- Even on a larger architecture where RLVR performance was already saturated, adding the judge reward improved accuracy from 69.07% to 70.43% (and up to 71.95% with majority voting).
Algorithm Comparison: The paper also evaluated Group Relative Policy Optimization (GRPO), finding it provided slightly higher gains than PPO in some configurations but noted PPO was more computationally efficient for their setup.
Failure Modes: The authors identified "reward hacking" (where models generate minimal/no code to trick the judge) when relying solely on unlabeled data ( $\lambda=0$ ). They proposed mitigation strategies, such as enforcing minimum code length and explicit instructions, to stabilize training.

5. Significance

This work represents a significant step toward democratizing advanced reasoning capabilities. By decoupling RL fine-tuning from the need for expensive ground-truth labels, it enables the distillation of complex reasoning skills from massive teacher models into smaller, resource-efficient student models using the vast reservoir of unlabeled data available on the internet.

The "LLM-as-a-Judge" mechanism, specifically optimized for efficiency via single-token scoring, offers a practical solution for scaling RL training in domains where verifiable rewards are unavailable. This suggests that high-quality, automated training signals can be generated dynamically, potentially reshaping how future language models are aligned and distilled.