ArXiv

On-Policy Self-Distillation with Sampled Demonstrations Reduces Output Diversity

2026-06-24 · 17:59 UTC ·Andrei Liviu Nicolicioiu, Mohammad Pezeshki, Aaron Courville ·1 min read
Authors
Andrei Liviu Nicolicioiu, Mohammad Pezeshki, Aaron Courville
Categories
cs.LG, cs.AI
arXiv
https://arxiv.org/abs/2606.26091v1
PDF
https://arxiv.org/pdf/2606.26091v1

Brief

On-policy self-distillation uses one model as teacher and student, with the teacher conditioned on a sampled correct demonstration to give dense token-level feedback. The authors show this boosts pass@1 but decreases rollout diversity and flattens pass@k. They prove the method tilts probabilities by a pointwise conditional mutual information score, amplifying preexisting mode gaps; experiments on a controlled graph path‑finding task and science QA find similar or better average accuracy than RL but much lower diversity and degraded OOD robustness.

Why it matters

On-policy self-distillation (a single model acting as both teacher and student, with the teacher conditioned on a sampled correct demonstration) raises pass@1 accuracy but reduces rollout diversity and flattens pass@k curves so that generating more rollouts often fails to improve accuracy.

Key details

  • The paper proves the optimal self-distillation policy tilts the base distribution by a pointwise conditional mutual information score between a student's rollout and the correct demonstration; unlike ideal on-policy RL (which preserves probability ratios among equally correct rollouts), self-distillation amplifies existing probability gaps and concentrates mass on dominant modes—empirically matching or exceeding RL on average but producing substantially lower functional/semantic diversity and worse OOD performance (graph path‑finding and science QA).
Source evidence

Abstract

On-policy self-distillation achieves strong pass@1 accuracy by using a single model as both teacher and student, with the teacher conditioned on a correct demonstration to provide dense token-level feedback. We show that this could come at a hidden cost: rollout diversity decreases and pass@k curves flatten (i.e., generating more rollouts fails to improve accuracy). We trace this to compounding biases in the design of self-distillation with sampled demonstrations. The teacher scores each student rollout while conditioned on a sampled correct rollout, channeling its feedback through the model's own biases. We theoretically analyze the optimal self-distillation policy and show that it tilts the base distribution by a pointwise conditional mutual information score between the student's rollout and the correct rollout used as context. Unlike the ideal optimal on-policy reinforcement learning (RL), which preserves probability ratios among equally correct rollouts, self-distillation can amplify existing probability gaps, concentrating mass on already-dominant modes. On a controlled graph path-finding task and science question-answering benchmarks, self-distilled models match or exceed RL on average performance but exhibit substantially lower functional and semantic diversity, failing on out-of-distribution settings that require diverse strategies.