BASED: Branch-Aware Soft Embedding Distillation for Reasoning LLMs

On-policy distillation (OPD) provides dense teacher supervision along student-generated reasoning trajectories, but each discrete rollout follows only one continuation at a reasoning branch, leaving alternative plausible branches unsupervised. Soft embeddings offer a natural way to represent multiple candidate branches within a continuous trajectory, yet we find that naive soft-trajectory distillation loses branch-specific information and may lead to branch collapse. To explain this limitation, we introduce the Hard-Branch KL objective and derive an exact decomposition into a Mixture KL term and a Branch-Posterior KL term. We show that soft-trajectory KL approximates the Mixture KL while omitting branch-posterior supervision. This reveals that soft distillation preserves mixture-level supervision but discards the branch identities needed to distinguish alternative reasoning paths. Motivated by this analysis, we propose Branch-Aware Soft Embedding Distillation (BASED), which combines soft KL with a local branch-posterior correction to efficiently recover the missing branch-specific supervision. Experiments across eight reasoning benchmarks show that BASED improves Pass@1 over vanilla OPD across both student scales and achieves higher average Pass@8, while ablations and trajectory-diversity analyses provide further evidence that BASED helps preserve diverse reasoning branches.