💻 computer science

On-Policy Distillation with Best-of-N Teacher Rollout Selection

This paper introduces BRTS, a Best-of-N Rollout Teacher Selection framework for on-policy distillation that improves reasoning performance by sampling multiple teacher trajectories and selecting the most correct and student-aligned one to provide reliable supervision, thereby overcoming the high-variance limitations of standard methods.

Original authors: Ke Zhang, Yunjie Tian, Dongdi Zhao, Yijiang Li, Yuanye Liu, Vishal M Patel, Di Fu

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ke Zhang, Yunjie Tian, Dongdi Zhao, Yijiang Li, Yuanye Liu, Vishal M Patel, Di Fu

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 apprentice (the Student) how to solve complex math puzzles. You have a master mathematician (the Teacher) who is incredibly smart but sometimes gets distracted, makes silly mistakes, or explains things in a way the apprentice just can't follow.

In the world of AI, this is called On-Policy Distillation. Usually, the apprentice tries to solve a problem, and the master watches what the apprentice does, correcting them step-by-step. The problem is: if the apprentice starts down a confusing path, the master might get confused too, or the master might just give a random, unhelpful answer because they are also "rolling the dice" on how to solve it.

This paper introduces a new method called BRTS (Best-of-N Rollout Teacher Selection). Think of it as a "Smart Coach" system that fixes the flaws in the old teaching method. Here is how it works, using simple analogies:

1. The Problem: The "Bad Roll of the Dice"

In the old method, when the apprentice asks, "How do I solve this?", the master flips a coin and gives one answer.

The Risk: Sometimes the master is having an "off day" and gives a wrong answer. Sometimes the master gives a correct answer, but explains it in a way that is totally different from how the apprentice thinks.
The Result: The apprentice learns from a bad example or a confusing one, which makes them worse at solving problems.

2. The Solution: The "Try a Few, Pick the Best" Strategy

BRTS changes the game. Instead of asking the master for just one answer, the system asks the master to generate several different attempts (a small pool of answers) at once.

Then, a smart filter (the Selector) looks at these attempts and picks the one to show the apprentice based on two simple rules:

Rule #1: Is it Correct? First, the system checks: "Did the master actually get the right answer?" If an attempt is wrong, it gets tossed in the trash.
Rule #2: Does it Match the Student? If the master got the right answer in three different ways, the system picks the one that looks most similar to how the apprentice usually thinks. This ensures the lesson is easy for the apprentice to understand.

3. The "Emergency Backup" (Tier-2 Recovery)

What if the master tries three times and all of them are wrong? (This happens with really hard puzzles).

The Old Way: The system would give up or force the apprentice to learn from a wrong answer.
The BRTS Way: The system has a secret cheat sheet (the Ground Truth). It whispers the correct answer to the master silently and says, "Okay, now that you know the answer, please write out a perfect, natural explanation of how you got there."
The master then produces a high-quality, correct explanation that the apprentice can learn from. This is like a coach stepping in to say, "Here is the right path, now watch how I walk it."

4. The Result: Faster and Smarter Learning

The paper tested this on difficult math competitions (like AIME and AMC).

The Finding: By using this "Pick the Best" method, the apprentice learned much faster and solved more problems correctly than when using the old "random single answer" method.
Why it works: It stops the apprentice from learning from the master's mistakes or confusing explanations. It ensures the apprentice only sees the best, most relevant examples of how to think.

Summary Analogy

Imagine you are learning to cook from a famous chef.

Old Method: You ask the chef, "How do I make this soup?" The chef flips a coin. Heads, they give you a recipe with salt. Tails, they give you a recipe with sugar. You try to learn from whichever one they picked, even if it's the sugar one.
BRTS Method: You ask the chef to write down five different soup recipes. You check them: "Okay, this one has sugar (bad), this one is missing water (bad). This one is perfect, and this one is also perfect but uses a technique you already know." You pick the perfect, familiar one and learn from that. If the chef can't think of a good one on their own, you secretly show them the "correct" recipe card and ask them to explain it again.

The paper claims this simple change—checking multiple options and picking the best one—makes AI models much better at reasoning without needing expensive new training techniques.

Technical Summary: On-Policy Distillation with Best-of-N Teacher Rollout Selection (BRTS)

1. Problem Statement

On-policy distillation (OPD) has emerged as a data-efficient post-training method for Large Language Models (LLMs), where a student model is supervised on its own sampled trajectories using the teacher's per-token log-probabilities. While OPD mitigates exposure bias by aligning supervision with the states the student actually visits, standard implementations suffer from high-variance supervision signals.

The core limitation identified is twofold:

Noisy Student Contexts: Supervision is computed on prefixes generated by an imperfect student, which may drift into noisy reasoning states where local teacher feedback is less informative.
Stochastic Teacher Rollouts: Standard OPD typically relies on a single teacher rollout per prompt. Since the teacher is stochastic, a single sampled trajectory may be incorrect, uninformative, or poorly aligned with the student's current reasoning behavior. This results in unreliable guidance, particularly on difficult prompts where the teacher's sampled outputs vary significantly in correctness and style.

Consequently, standard OPD risks propagating erroneous teacher signals or failing to expose the student to complete, correct reasoning paths, limiting its effectiveness on challenging reasoning benchmarks.

2. Methodology: BRTS Framework

The authors propose BRTS (Best-of-N Rollout Teacher Selection), a framework designed to curate reliable teacher trajectories before they are used for supervision within the OPD loop. BRTS augments standard student-context distillation with a teacher-context supervision branch constructed from a selected, high-quality teacher trajectory.

2.1 Correctness- and Alignment-Aware Trajectory Curation

For each prompt, BRTS samples a pool of $N$ unconditioned teacher rollouts ( $\{y_{T,i}\}_{i=1}^N$ ). The selection process follows a priority hierarchy:

Priority 1: Correctness. The system filters for rollouts where the extracted final answer matches the ground truth ( $y^*$ ).
Priority 2: Student Alignment. If multiple correct rollouts exist, BRTS selects the one with the highest token-level overlap with the student's top- $K$ candidate set. This ensures the selected trajectory is compatible with the student's current reasoning distribution.
Ground-Truth-Guided Recovery (Tier-2). If all unconditioned teacher samples fail (i.e., none are correct), BRTS invokes a recovery step. It constructs a modified prompt that silently injects the ground-truth answer as a validation signal while instructing the teacher to produce a natural derivation. A new rollout is sampled; if correct, it is selected. If this also fails, the system falls back to the unconditioned rollout with the highest student overlap.

2.2 Teacher-Context Supervision

The training objective combines two loss components:

Student-Context Loss ( $L_{stu-ctx}$ ): The standard OPD objective, minimizing the reverse KL divergence between the student and teacher distributions conditioned on the student's generated prefix. This corrects the student on the states it actually visits.
Teacher-Context Loss ( $L_{tea-ctx}$ ): An auxiliary loss computed on the selected teacher trajectory ( $y'$ ). Here, both the student and teacher distributions are conditioned on the teacher's prefix. This exposes the student to a coherent, correct reasoning path that may lie outside the student's current high-probability region.

The total loss is defined as:
$L_{total} = L_{stu-ctx} + \lambda L_{tea-ctx}$
where $\lambda$ is a weighting coefficient (empirically set to 10).

2.3 Top-K Candidate Direction

To implement these losses efficiently, BRTS uses top- $K$ aggregation:

For the student-context branch, the candidate token set is derived from the student's top- $K$ tokens under the student prefix.
For the teacher-context branch, the candidate set is derived from the teacher's top- $K$ tokens under the selected teacher prefix. This allows the auxiliary branch to supervise tokens preferred by the teacher that the student might not currently consider, while maintaining a compact candidate set.

3. Key Contributions

The paper summarizes its contributions as follows:

Identification of High-Variance Supervision: The authors highlight that standard student-context supervision in OPD is limited by noisy student prefixes and stochastic teacher rollouts, rendering the teacher signal unreliable on hard prompts.
Introduction of BRTS: A lightweight framework for selecting and distilling reliable teacher trajectories. It employs a "correctness-first, alignment-second" selection rule and includes a ground-truth-guided fallback for failed unconditioned samples.
Empirical Improvement: BRTS significantly outperforms standard OPD on hard reasoning benchmarks. The paper provides diagnostic analyses demonstrating how teacher candidate quality, recovery coverage, and backbone choices influence performance.

4. Experimental Results

Experiments were conducted on AIME 2024, AIME 2025, and AMC 2023 using a 1.5B student model (DeepSeek-1.5B) and various teacher configurations (JustRL-1.5B and DeepSeek-R1-Distill-Qwen-7B).

Performance Gains: BRTS consistently improves over the student-only baseline. The gains are most pronounced on the harder AIME 2025 dataset. For instance, on AIME 2024, using four teacher candidates achieved an average accuracy of 0.4000 compared to 0.3917 for the baseline.
Impact of Tier-2 Recovery: The ground-truth-guided recovery step (Tier-2) provides substantial improvements on difficult prompts where unconditioned sampling fails. On AIME 2025, adding Tier-2 increased the majority-vote accuracy from 0.2748 (Tier-1 only) to 0.3133.
Early Training Peaks: BRTS achieves a higher performance peak in early training stages compared to baselines, suggesting faster convergence and a better compute-performance trade-off.
Robustness: The method remains effective when swapping teacher backbones (e.g., using a 7B teacher for a 1.5B student) and is not brittle to mild prompt perturbations, which actually help decorrelate teacher samples.
Computational Cost: The overhead is modest; increasing the teacher candidate pool from 1 to 2 adds approximately 59 seconds per training step on 8×B200 GPUs, which is considered acceptable given the performance gains.

5. Significance and Claims

The paper claims that BRTS addresses a critical bottleneck in on-policy distillation: the unreliability of stochastic teacher signals on difficult tasks. By structuring the teacher supervision to be both correctness-aware (filtering out errors) and alignment-aware (matching the student's distribution), BRTS transforms stochastic teacher rollouts into structured, high-quality supervision.

The authors position BRTS not as a replacement for OPD, but as a necessary augmentation that allows students to learn from complete, reliable reasoning paths without incurring the catastrophic forgetting often associated with standard supervised fine-tuning or the high variance of standard OPD. The method is presented as a scalable, lightweight extension that improves reasoning capabilities, particularly in domains where ground-truth answers are available for verification.