🤖 machine learning

Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

The paper introduces REOPOLD, a novel framework that interprets on-policy distillation as policy optimization and stabilizes it through relaxed constraints and dynamic reward mechanisms, achieving significantly higher sample efficiency and enabling smaller models to match larger teachers' reasoning performance across diverse tasks.

Original authors: Jongwoo Ko, Sara Abdali, Young Jin Kim, Tianyi Chen, Pashmina Cameron

Published 2026-03-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jongwoo Ko, Sara Abdali, Young Jin Kim, Tianyi Chen, Pashmina Cameron

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 young, energetic apprentice (a small AI model) how to solve complex puzzles, like advanced math or visual riddles. You have a wise, master teacher (a massive AI model) who is incredibly smart but takes a long time to think and is very expensive to run.

The goal is to make the apprentice as smart as the master, but faster and cheaper.

The Problem: The "Parrot" vs. The "Thinker"

Traditionally, when we teach the apprentice, we use a method called On-Policy Distillation. Think of this as the apprentice trying to perfectly mimic the master's every move in real-time.

The Issue: If the master makes a tiny mistake or takes a weird path, the apprentice tries to copy it exactly. If the master says, "Don't do that!" with a very harsh tone, the apprentice panics and forgets everything it knew before.
The Result: The apprentice gets confused, stops learning, or starts hallucinating (making things up). It's like a student who is so afraid of getting a question wrong that they stop thinking creatively and just freeze.

The Solution: REOPOLD (The "Relaxed" Mentor)

The authors of this paper introduced a new method called REOPOLD. Instead of forcing the apprentice to copy the master perfectly, they act more like a wise coach who knows when to push and when to let the student figure things out.

Here is how REOPOLD works, using simple analogies:

1. The "Safety Net" (Reward Clipping)

In the old method, if the master said, "That answer is terrible!" (a huge negative score), the apprentice would get a massive shock and crash.

REOPOLD's Fix: It puts a "safety net" under the coach's voice. If the coach gets too angry or harsh, REOPOLD says, "Okay, that's a bad idea, but let's not freak out." It caps the negative feedback so the apprentice doesn't get scared into stopping. It filters out the "noise" and only listens to the useful criticism.

2. The "Highlighter" (Dynamic Sampling)

Imagine the apprentice is reading a long textbook. Most of the pages are boring and obvious (like "1+1=2"). The master and the apprentice already agree on these. But the real learning happens on the few pages with tricky, confusing diagrams.

REOPOLD's Fix: It uses a highlighter to ignore the boring, easy pages where the apprentice and master agree. It focuses 100% of the energy on the confusing, high-uncertainty moments. This is like a tutor who skips the easy homework and only helps with the hard problems, making learning much faster.

3. The "Two-Phase Training" (Exploration then Refinement)

Phase 1 (The Playground): At the start, the coach lets the apprentice try many different crazy ideas. Even if they are wrong, the coach doesn't punish them too hard. This encourages the apprentice to be creative and explore different ways to solve the problem.
Phase 2 (The Drill): Once the apprentice has explored, the coach switches modes. Now, they focus on sharpening the best ideas and cutting out the bad ones. This turns the "creative mess" into a "polished skill."

Why This Matters (The Results)

The paper shows that this new "Relaxed" approach is a game-changer:

Super Efficient: The apprentice learns the same amount of knowledge using 6 to 12 times less data than before. It's like learning a language in a month instead of a year.
Small Models, Big Brains: A small 7-billion-parameter model (the apprentice) trained with REOPOLD can solve visual puzzles almost as well as a massive 32-billion-parameter model (the master), but it does it 3 times faster.
No More Crashes: The training is stable. The apprentice doesn't get confused or give up; it steadily gets smarter.

The Bottom Line

REOPOLD is about changing the relationship between the teacher and the student. Instead of a strict dictator demanding perfect imitation, it acts like a flexible mentor who:

Softens the blow of harsh criticism.
Focuses only on the hard parts.
Encourages exploration before demanding perfection.

This allows small, fast AI models to become reasoning geniuses without needing the massive computing power of their giant teachers.

1. Problem Statement

While Reinforcement Learning (RL) has successfully enhanced reasoning capabilities in large language models (LLMs) like DeepSeek-R1 and OpenAI-o1, replicating this success in Small Language Models (SLMs) is challenging. SLMs often suffer from instability and negative transfer when subjected to direct RL due to limited representational capacity and sparse rewards.

On-policy distillation has emerged as a promising alternative, where a student model learns from its own trajectories guided by a high-capacity teacher. However, vanilla on-policy distillation (specifically using Reverse KL divergence, or RKL) faces three critical optimization bottlenecks:

Instability from Heavy-Tailed Negative Rewards: When the student generates tokens the teacher deems improbable, the log-likelihood ratio approaches $-\infty$ , creating extreme negative rewards that destabilize training.
Inefficiency from Near-Zero Rewards: For tokens where student and teacher distributions align well, the reward signal is near zero. These tokens consume computation but provide negligible learning signals (gradient dilution).
Entropy Collapse: The policy tends to collapse prematurely into a narrow set of outputs, losing the diversity required for effective reasoning exploration.

2. Methodology: REOPOLD

The authors propose REOPOLD (Relaxed On-Policy Distillation), a framework that interprets on-policy distillation as a form of policy optimization where the teacher-student log-likelihood ratio acts as a token-level reward. To stabilize this process, REOPOLD introduces three core mechanisms:

A. Theoretical Foundation: Stop-Gradient as Control Variate

The authors establish that on-policy distillation is theoretically equivalent to a policy gradient method. By applying a stop-gradient operator to the reward term (the log-likelihood ratio), they treat the reward as a fixed extrinsic signal. This removes the variance caused by the gradient of the reward term itself, acting as a control variate to stabilize optimization.

B. Mixture-Based Reward Clipping

To address the heavy-tailed negative rewards (where $\pi_T \to 0$ ), REOPOLD introduces a principled clipping threshold derived from the convex mixture of teacher and student distributions.

Instead of allowing rewards to diverge to $-\infty$ , the reward is clipped at a lower bound: $\log \frac{\lambda}{1-\lambda}$ .
This prevents gradient explosion from extreme negative signals while preserving the "mode-seeking" nature of the distillation.

C. Entropy-Guided Token-Level Dynamic Sampling

To solve the inefficiency of near-zero rewards, REOPOLD filters tokens based on entropy.

Observation: Low-entropy tokens (deterministic steps) usually have near-zero rewards and offer little learning value. High-entropy tokens (uncertain branching points) contain the most valuable learning signals.
Mechanism: The method dynamically masks out low-entropy tokens, focusing gradient updates only on high-entropy tokens where the student and teacher distributions diverge significantly. This creates a "dense" learning signal.

D. Exploration-to-Refinement Multi-Stage Training

REOPOLD employs a two-phase training strategy to balance exploration and exploitation:

Exploration Phase (Early Stages): Uses reward clipping to suppress excessive negative penalties. This mimics Supervised Fine-Tuning (SFT), allowing the model to explore diverse solution paths without collapsing entropy prematurely.
Refinement Phase (Later Stages): Switches to entropy-guided sampling. This allows negative feedback on critical, uncertain tokens to refine the policy and converge toward high-quality reasoning paths.

3. Key Contributions

Theoretical Insight: Proved that on-policy distillation is equivalent to policy optimization with a fixed reward, identifying the lack of stop-gradient and unclipped rewards as primary causes of instability.
Algorithmic Innovation: Introduced REOPOLD, which integrates reward clipping, token-level dynamic sampling, and multi-stage training to stabilize distillation for SLMs.
Efficiency & Scalability: Demonstrated that relaxing strict imitation constraints allows smaller models to scale effectively, achieving performance comparable to much larger teachers with significantly fewer resources.

4. Experimental Results

The authors evaluated REOPOLD across mathematical, visual, and agentic reasoning tasks.

Mathematical Reasoning (AIME, MATH, etc.):
- Sample Efficiency: REOPOLD achieved 6.7x to 12x greater sample efficiency compared to standard RL (ProRL) and vanilla RKL baselines.
- Stability: While vanilla RKL caused performance drops in 7B models (instability), REOPOLD enabled stable, continuous improvement.
- Robustness: Outperformed baselines regardless of the teacher model size (7B vs. 32B), whereas RKL performance varied significantly based on the teacher.
Visual Reasoning (Geometry3K, MathVerse, etc.):
- Performance: A 7B student model trained with REOPOLD matched the performance of a 32B teacher on visual reasoning tasks.
- Inference Speedup: Due to the smaller model size, REOPOLD achieved a ~3.32x inference speedup compared to the 32B teacher while maintaining similar accuracy (Pass@K).
- Test-Time Scaling: REOPOLD demonstrated superior scaling curves, maintaining high accuracy with increased sampling (Pass@K) compared to RKL.
Agentic Reasoning (Tool Use):
- In tasks requiring visual tool use (e.g., zooming, selecting frames), REOPOLD outperformed both GRPO and RKL, showing better generalization and sample efficiency without complex reward engineering.

5. Significance

Democratizing Reasoning: REOPOLD provides a viable path for deploying high-level reasoning capabilities in compact models (e.g., 1.5B–7B), making advanced AI reasoning accessible on consumer hardware.
Stabilizing Distillation: It resolves the fundamental instability of on-policy distillation, allowing the use of much larger teachers (32B+) to train smaller students without the risk of negative transfer or entropy collapse.
Efficiency: By filtering out uninformative tokens and optimizing the training signal, REOPOLD drastically reduces the computational cost (GPU hours) required to achieve state-of-the-art reasoning performance.

In conclusion, the paper argues that relaxing strict imitation through intelligent signal filtering and multi-stage optimization is essential for scaling reasoning capabilities in resource-constrained models.