Branching On-Policy Distillation

Motivation

• A single trajectory follows only one continuation at each decision, leaving unchosen multi-step histories unsupervised; independent rollouts can regenerate near-identical prefixes.This limits coverage in mathematics, where early steps constrain derivations, and agent tasks, where actions change later observations.Source: §1 · p. 1

BranchOPD

• BranchOPD selects fork points using student uncertainty and teacher–student disagreement, then resamples student-policy continuations for token-level teacher supervision.It selects token positions in mathematics and turns followed by positions within them in agent tasks; shared prefixes enter the loss once.Source: §1 · p. 2, Figure 1 · p. 2

Protocol-boundary repair

• Protocol-boundary repair adds an auxiliary loss to reinforce beneficial and suppress harmful protocol events, reducing repetition and wasted generation.The baselines also use PBR so comparisons isolate the effect of branching.Source: §1 · p. 2

Experiments and analysis

• BranchOPD outperforms Vanilla OPD and Multi-Rollout OPD on all evaluated benchmarks, using about 19% and 31% fewer student tokens than seven independent trajectories on ALFWorld and WebShop.Ablations and diagnostics show that fork-point choice matters.Source: §1 · p. 2

On-policy distillation

• On-policy distillation methods train on student samples and vary the divergence, supervision weighting, or trajectory scheduling; SPOT instead uses teacher-proposed tokens and outcome-verifier scoring.BranchOPD samples branches from the student and supervises them token by token without task rewards.Source: §2 · pp. 2–3

Quality of distillation supervision

• Other approaches improve supervision through trajectory refinement, expert corrections, or teacher continuations from student anchors; BranchOPD's continuations come from the student.Its token-local objective complements work separating token-level updates from sequence-level credit assignment, while PBR supervises specific output events.Source: §2 · p. 3

Tree search and reinforcement learning

• Tree-search and reinforcement-learning methods reuse prefixes to explore alternatives, typically comparing branches through task rewards; BranchOPD uses shared-prefix exploration to gather tokens for teacher supervision.Related approaches also branch at high-entropy tokens, adapt rollout budgets, or roll back environments and reflect.Source: §2 · p. 3

3.1 STUDENT TRAJECTORIES

• In mathematics, a student trajectory is one response ending at a stop token or length limit; in agent tasks, it includes responses, actions, and resulting observations, but observations are not supervised.Mathematics is the special case with one turn and no observations.Source: §3.1 · p. 3
• Resuming an agent trajectory from a prefix requires restoring the environment to its state there so the continuation has the same distribution as an uninterrupted rollout.Restoration can use a snapshot or action replay.Source: §3.1 · p. 3

3.2 ON-POLICY DISTILLATION

• On-policy distillation samples trajectories from a frozen copy of the student and trains on its visited prefixes using teacher scores for sampled tokens.The token-level signal is the difference between teacher and student log probabilities for each sampled token; Vanilla OPD uses one trajectory, while Multi-Rollout OPD uses multiple independent trajectories.Source: §3.2 · pp. 3–4

4 METHOD

• BranchOPD selects fork points on a student trajectory, resamples continuations from the student policy, and applies OPD loss and PBR to generated tokens.Shared prefixes are counted once.Source: §4 · p. 4

4.1 SELECTING FORK POINTS

• Fork selection combines student entropy with teacher–student disagreement aggregated over a 128-token window, then rank-normalizes and mixes the scores.Mathematics selects high-scoring eligible positions; agent tasks select turns and place each fork among that turn's three highest-scoring eligible positions, by default.Source: §4.1 · p. 5

4.2 CONDITIONAL RESAMPLING AND COVERAGE

• Given correctly restored contexts and branches sampled from the student continuation distribution, the branch estimate is conditionally unbiased; independent branches reduce its variance by a factor of k.The probability of seeing a first token different from the main trajectory's increases with k, with diminishing increments that motivate spreading the budget across forks.Source: §4.2 · p. 5

4.3 TRAINING OBJECTIVE

• The training objective masks observations and padding, applies a clipped teacher–student log-probability advantage, and counts each shared prefix once.Because the loss averages over valid tokens, longer suffixes contribute more terms.Source: §4.3 · p. 5

4.4 PROTOCOL-BOUNDARY REPAIR

• Protocol-boundary repair detects constraints at prefixes specifying whether the next token should belong to a nonempty token set, then applies an auxiliary loss to main and branch trajectories.Positive constraints raise the probability of beneficial events and negative constraints lower that of harmful ones; the auxiliary term is zero when no constraints are detected.Source: §4.4 · p. 6

5.1 EXPERIMENTAL SETUP

• Evaluation covers mathematics and three agent benchmarks; agent evaluation uses one trajectory per task, with no branching at evaluation time.The main baselines include Vanilla OPD, seven-trajectory Multi-Rollout OPD, and single-trajectory Entropy-Weighted OPD.Source: §5.1 · p. 6

5.2 MAIN RESULTS

• BranchOPD outperforms KD, GRPO, EOPD, and SPOT by 11.27, 8.81, 6.40, and 5.94 mathematics Avg points, respectively.Its gains over SPOT are 1.36 points on MATH-500 and 8.54 and 7.91 points on AIME24 and AIME25.Source: §5.2 · p. 6
• 18.55-point mathematics Avg and 13.83-point mean ALFWorld success gains over Vanilla OPD; BranchOPD also leads across the three agent benchmarks.With PBR on every method, BranchOPD reaches 44.37 mathematics Avg, 5.43 and 5.28 points above Vanilla OPD and Multi-Rollout OPD (n = 7).Source: §5.2 · pp. 6–7, Table 1 · p. 7, Table 2 · p. 7
• Across the three agent benchmarks, BranchOPD reaches 46.39 SR, while mean ALFWorld success rises from 49.27% for Vanilla OPD (+PBR) to 56.54%.It also raises WebShop SR from 30.47% to 33.59% over Vanilla OPD (+PBR), and gains 0.47 points on Search-QA.Source: §5.2 · p. 7, Table 2 · p. 7

5.3 ABLATIONS: WHERE AND HOW MUCH TO BRANCH

• Mixing student entropy and teacher disagreement beats random forks, raising mathematics Avg from 40.11 to 44.37 and ALFWorld ID/OOD SR by 7.14/8.96 points.Turn-initial branching trails the default by 2.14/1.49 ALFWorld points, so fork position within a turn also matters.Source: §5.3 · p. 7, Table 3 · p. 8
• With B = 3, increasing continuations per fork from k = 1 to k = 2 improves all three metrics, while k = 3 adds no further gain.At k = 2, increasing forks from two to three adds 0.43 mathematics points and 0.73 ALFWorld points.Source: §5.3 · p. 8, Table 3 · p. 8

5.4 BEHAVIOR COVERAGE AND SAMPLING EFFICIENCY

• At selected forks, 74.2% of ALFWorld and 47.7% of WebShop branches change action sequence; shared-prefix reuse saves 19.1% and 31.2% of student tokens, respectively.When the main trajectory fails, at least one branch succeeds in 5.0% of WebShop to 42.7% of mathematics cases.Source: §5.4 · p. 8

5.5 PROTOCOL-BOUNDARY REPAIR

• PBR raises Vanilla OPD mathematics Avg from 25.82 to 38.94 and ALFWorld ID/OOD SR by 7.14/5.97 points, outperforming generic repetition remedies.Repetition penalties achieve 28.70–30.40 Avg, filtering overlong samples reaches 33.20, and early stopping reaches 22.05, versus 38.94 for PBR.Source: §5.5 · pp. 8–9

6 CONCLUSION

• BranchOPD outperforms Vanilla OPD and Multi-Rollout OPD on every evaluated benchmark while generating fewer student tokens than independent rollouts.The conclusion emphasizes that what OPD teaches depends on which of the student's alternative histories it sees.Source: §6 · p. 9

A.1 OPD GRADIENT AND TOKEN SURROGATE

• The OPD token surrogate expresses reverse-KL gradients as old-policy expectations using an importance ratio and detached teacher-student signal, with PPO and signal clipping defining training variants.The derivation assumes teacher support wherever the student has support and student-policy support under the rollout policy at supervised prefixes.Source: §A.1 · p. 18, A THEORETICAL DETAILS · p. 18

A.2 CONDITIONAL RESAMPLING

• Fork selection may depend on the main trajectory without biasing the conditional branch estimate, because the expectation is over fresh branch draws.The result concerns the average of a fixed suffix statistic.Source: §A.2 · p. 18

A.3 COVERAGE AND DIMINISHING RETURNS

• The expected marginal increase in distinct first-token coverage is nonnegative but decreases as more branches are sampled; branches sharing a first token can still diverge later.The paper therefore measures both first-token alternatives and later behavioral differences.Source: §A.3 · p. 19

A.4 TOKEN AGGREGATION AND PROTOCOL REPAIR

• Token-loss aggregation counts generated suffix tokens, includes shared prefixes only once, and uses the mask for normalization; PBR raises beneficial-event probabilities and lowers harmful-event probabilities.PBR detection is fixed during differentiation, and a batch without constraints contributes zero.Source: §A.4 · p. 19

A.5 WEIGHTED FORCED-TOKEN COMPARATOR

• Table 8 uses scaled Horvitz–Thompson weights in a pooled token loss, averaging the main-slot weight across forks; self-normalized RB is a different estimator.The stated HT identity requires the natural draw and continuations to follow the fixed-prefix sampling distribution.Source: §A.5 · p. 19

B.1 STUDENT SIZE AND MODEL FAMILY

• With Qwen3-1.7B, BranchOPD improves over Vanilla OPD (+PBR) by 6.56 points on ALFWorld and 4.68 points on mathematics; with Qwen2.5-3B, the ALFWorld gain is 8.76 points.Table 4 covers Qwen3-4B to Qwen3-1.7B and reports task SR or mathematics mean avg@16; Table 5 covers Qwen2.5-7B-RL to Qwen2.5-3B on ALFWorld.Source: §B.1 · p. 20, Table 4 · p. 20, Table 5 · p. 20

B.2 SELECTION AND BRANCHING BUDGETS

• The entropy-plus-disagreement selector is best among tested rules; disabling mathematics-aware eligibility lowers mathematics Avg to 42.03, while branching only from failed main trajectories lowers it to 36.82.Branches from successful trajectories are also useful.Source: Table 6 · p. 20, §B.2 · p. 21
• At k = 2, increasing the fork budget B from one to three raises mathematics Avg from 42.55 to 44.37, while k = 3 adds nothing over k = 2.The five-trajectory setting (B = 2, k = 2) retains most of the gain at lower generation cost.Source: §B.2 · p. 21, Table 7 · p. 21

B.3 SAMPLING AND PERTURBATION CONTROLS

• Resampling gives the highest ALFWorld, mathematics, and WebShop scores, while RB gives the highest Search-QA SR; resampling also yields fewer low-likelihood suffixes, 5.8% versus 19.2%.Noise-token branching gives smaller ALFWorld gains of 1.09 and 0.73 points over Vanilla OPD (+PBR), depending on fork placement; informative branching's larger gains indicate that fork selection matters.Source: §B.3 · pp. 21–22, Table 8 · p. 21, Table 9 · p. 21

B.4 GENERATION AND ENVIRONMENT COSTS

• Relative to seven independent rollouts, BranchOPD saves 19.1% and 31.2% of student generation on ALFWorld and WebShop while improving SR; estimates give savings of 18.9% on mathematics and 39.4% on Search-QA.With two added continuations, ALFWorld uses 2.55× the generation and 3.02× the environment steps of Vanilla OPD because state replay adds environment cost beyond suffix generation.Source: §B.4 · p. 22, Table 11 · p. 22, Table 12 · p. 22

B.5 PROTOCOL REPAIR AND REPETITION

• 38.94 mathematics Avg with PBR exceeds 33.20 for filtering and 28.70–30.40 for repetition penalties; without PBR, responses reach the length cap around step 30.After step 50, PBR's median response length is about 8,600 tokens and its median truncation rate is 0.28.Source: §B.5 · p. 23, Table 13 · p. 23
• PBR reduces excess leading </think> markers on ALFWorld from Vanilla OPD's 1.94/1.91 ID/OOD averages to zero; trajectory-level protocol violation drops to zero by step 25 and stays near zero.A single response can contain several excess markers.Source: §B.5 · p. 23, Figure 5 · p. 23

C.1 PROBABILITY AVAILABLE AT SELECTED DECISIONS

• At selected forks, alternative-token mass is about 0.63 on ALFWorld and 0.69 on WebShop, while 80.4% and 59.1% of ordinary supervised positions have alternative mass below 0.05.Figure 6 defines alternative mass as the probability mass left unrealized by the main draw; Figure 7 presents its empirical CDF over supervised positions.Source: §C.1 · p. 24, Figure 6 · p. 24, Figure 7 · p. 24

C.2 ALTERNATIVE ACTIONS ACROSS TRAINING

• At step 100, action disagreement (EBR) is 74.2% on ALFWorld and 47.7% on WebShop, while new-action coverage (NSR) is 55.5% and 37.5%, respectively.Coverage persists through training: it stabilizes after an initial drop on WebShop and stays high on ALFWorld.Source: §C.2 · p. 25, Table 14 · p. 25

C.3 RECOVERY OPPORTUNITIES

• The rescue rate ranges from 5.0% on WebShop to 42.7% on mathematics, with successful alternatives continuing to appear throughout training.Rescue means some branch succeeds when the main trajectory fails; harm means all branches fail when the main succeeds, and the successful main remains in training.Source: §C.3 · p. 26, Table 15 · p. 26

C.4 WHERE BEHAVIOR BEGINS TO DIVERGE

• The first changed action follows the fork by roughly 12–18 generated tokens on ALFWorld, 1.5–2 on WebShop, and 30 on Search-QA.Fork locations vary by task and often fall inside a turn; token-level forks can therefore change later actions while reusing preceding context.Source: §C.4 · p. 27, Figure 10 · p. 27

D.1 SHARED TRAINING CONFIGURATION

• The default uses B = 3 forks and k = 2 continuations per fork, scoring fork candidates with entropy and teacher–student disagreement mixed at α = 0.5.For agent tasks, selected turns are scored from eligible tokens, while PBR uses an unlikelihood weight λ = 0.01 and a weighted-loss cap of 0.02.Source: §D.1 · p. 27

D.2 ADDITIONAL MATHEMATICS BASELINES

• The mathematics baselines include offline KD on fixed teacher samples and outcome-based GRPO, which updates from binary answer-correctness rewards without a teacher.KD combines token cross-entropy and forward KL; both KD and GRPO are trained without PBR in Table 1.Source: §D.2 · pp. 27–28
• EOPD adds teacher-entropy-triggered forward KL but samples no extra suffixes, whereas SPOT forces teacher-proposed candidates and scores student continuations with an outcome verifier.BranchOPD instead resamples from the student and applies dense teacher supervision along new suffixes.Source: §D.2 · p. 28

D.3 EOS ALIGNMENT FOR THE BASE STUDENT AND INSTRUCTION TEACHER

• The student adopts the teacher’s EOS configuration and stops on either 151645 or 151643, aligning their end-of-response convention for teacher probabilities and terminal supervision.Source: §D.3 · p. 28

• Mathematics evaluation samples 16 answers per problem at temperature 0.6, with top-p = 0.95 and top-k = 20.Prompt and response limits are 2,048 and 16,384 tokens.Source: §D.4 · p. 28
• Agent evaluation samples one trajectory per task at temperature 0.4 with optional thinking disabled; WebShop uses 128 held-out tasks and seed 42.The WebShop environment disables human interaction; Search returns three results per call with a 30-second timeout.Source: §D.4 · p. 29
• Agent Avg. SR is the unweighted mean of ALFWorld aggregate, WebShop SR, and Search-QA SR, while mathematics Avg is the unweighted mean of three benchmark accuracies.ALFWorld aggregate averages its 140 ID and 134 OOD tasks without weighting.Source: §D.4 · p. 29

E.1 MATHEMATICAL REASONING

• Mathematics evaluation covers MATH-500 and AIME 2024 and 2025, with avg@16 averaging correctness across 16 independent solutions per problem and then across problems.Training problems come from DAPO.Source: §E.1 · p. 29

E.2 ALFWORLD

• ALFWorld tests text-based household tasks, reporting success separately on 140 ID and 134 OOD tasks and their unweighted mean.During branch training, the selected interaction state is restored or preceding actions are replayed; observations are context, not supervision.Source: §E.2 · p. 29

E.3 WEBSHOP

• WebShop is a simulated shopping benchmark in which agents search, browse, select options, and purchase items; success is reported on 128 held-out tasks.Source: §E.3 · p. 29

E.4 SEARCH-QA

• Search-QA interleaves search queries, retrieved results, and a final answer, with success measured by exact match to an accepted normalized answer.Each search returns up to three results, with at most four interaction steps.Source: §E.4 · p. 29

E.5 BENCHMARKS DISCUSSED AS BACKGROUND

• The background benchmark list spans grade-school and advanced quantitative reasoning, science and web agents, and question answering that can require multi-hop evidence.Examples include GSM8K, Minerva, ScienceWorld, WebArena, AgentGym, HotpotQA, TriviaQA, and MuSiQue.Source: §E.5 · p. 29

E.6 ADDITIONAL LEARNING AND SYSTEMS CONTEXT

• Related context distinguishes task-reward agent training from the protocol-event loss used here and notes that execution and memory systems complement reduced student-token generation.The experiments use Qwen3 and Qwen2.5 families; rollout policy is fixed within each sampling group in the analysis.Source: §E.6 · p. 30

F TRAINING ALGORITHM

• Each iteration freezes the rollout policy, samples a main trajectory, and selects up to B token forks for mathematics or turns followed by eligible tokens for agent tasks.If no eligible fork is found, training uses the main trajectory alone.Source: F TRAINING ALGORITHM · p. 31
• At each selected prefix, the algorithm independently samples k student-policy continuations, retains them regardless of outcome, and queries teacher log probabilities for retained sampled tokens.Observations and duplicate branch-prefix tokens are excluded from the loss; applicable protocol-boundary events add the PBR loss.Source: F TRAINING ALGORITHM · p. 31

G PROMPT TEMPLATES

• The prompt templates distinguish mathematics answers, structured ALFWorld and WebShop actions, and Search-QA retrieval feedback; their output boundaries are checked by PBR.Source: G PROMPT TEMPLATES · p. 32