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

2026-06-24

AI Summary

This study investigates the degradation of output diversity in policy self-distillation caused by using the model’s own sampled demonstrations. Through an information-theoretic lens, it reveals for the first time that the optimal self-distillation policy skews the base distribution via conditional mutual information, amplifying existing probability gaps and thereby introducing bias—contrasting with reinforcement learning’s tendency to preserve uniformity among equivalent solutions. Combining theoretical analysis with controlled graph-path tasks and scientific question-answering benchmarks, the work demonstrates that while self-distillation matches or even surpasses reinforcement learning in average performance, it significantly reduces functional and semantic diversity in generated outputs and exhibits weaker generalization under out-of-distribution settings.

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.

Problem

• self-distillation
• output diversity
• on-policy learning
• demonstration bias
• rollout diversity

Innovation

• on-policy self-distillation
• output diversity
• sampled demonstrations
• conditional mutual information
• pass@k degradation