Branching On-Policy Distillation
cs.AI
Submitted: 2026-10-09
Updated: 2026-10-09
Code: https://github.com/chendy25/BranchOPD
License: http://creativecommons.org/licenses/by/4.0/
The gist: On-policy distillation (OPD) provides teacher supervision on student-generated trajectories, but a single trajectory realizes only one continuation at each decision, limiting supervision coverage
Terminology
Abstract
On-policy distillation (OPD) provides teacher supervision on student-generated trajectories, but a single trajectory realizes only one continuation at each decision, limiting supervision coverage under a finite sampling budget. We propose BranchOPD (Branching On-Policy Distillation), which selects informative decision points using student uncertainty and teacher-student disagreement, resamples student-policy continuations from shared prefixes, and applies token-level teacher supervision to both main trajectories and branches. BranchOPD operates at token-level decisions in mathematics and within-turn decisions in agent tasks. Shared-prefix reuse reduces student-generated token costs by approximately 19% on ALFWorld and 31% on WebShop relative to seven independent trajectories. We also introduce protocol-boundary repair (PBR), which uses targeted supervision under protocol constraints to mitigate repetition and wasted generation. BranchOPD outperforms Vanilla OPD and Multi-Rollout OPD (n=7) across all evaluated benchmarks, with gains over Vanilla OPD of 18.55 percentage points in mean mathematics accuracy over 16 sampled answers (avg@16) and 13.8 points in mean ALFWorld ID/OOD success. Mechanistic analyses and selection ablations show that branching at uncertain decisions expands subsequent behavior coverage and outperforms random selection and turn-initial branching.
Sources
- On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- Group-in-Group Policy Optimization for LLM Agent Training
- The Llama 3 Herd of Models
- MiniLLM: On-Policy Distillation of Large Language Models
- Branching Policy Optimization: Sandbox-Native Language Agent Reinforcement Learning
- Measuring Mathematical Problem Solving With the MATH Dataset
- Distilling the Knowledge in a Neural Network
- Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes
- Tree Search for LLM Agent Reinforcement Learning
- Trajectory-Refined Distillation
- Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning
- Entropy-Aware On-Policy Distillation of Language Models
- Sequence-Level Knowledge Distillation
- DistiLLM: Towards Streamlined Distillation for Large Language Models
- Efficient Memory Management for Large Language Model Serving with PagedAttention
- Solving Quantitative Reasoning Problems with Language Models
- CURE: Critical-Token-Guided Re-Concatenation for Entropy-Collapse Prevention
- Let's Verify Step by Step
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems