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

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.

• The paper identifies that SDSD biases reinforcement by emphasizing demonstration similarity, leading to mode collapse in output diversity.
• It establishes that token-level entropy is insufficient for assessing functional and semantic diversity, necessitating explicit diversity metrics.
• Empirical findings on graph tasks and Science QA reveal high pass@1 accuracy but flat pass@k curves, highlighting limitations in generalization.

On-Policy Self-Distillation with Sampled Demonstrations: Analysis of Diversity Collapse

Technical Background

On-policy self-distillation with sampled demonstrations (SDSD) is a paradigm in LLM post-training where the teacher and student are the same model; the teacher, when conditioned on a correct demonstration, provides token-level guidance to the student. This setup eliminates the need for an external teacher, yielding efficient learning via dense feedback and strong pass@1 accuracy. However, the mechanics of how correct rollouts receive reinforcement diverges fundamentally from on-policy RL, which employs binary reward signals. RL rewards all correct solutions equally, preserving the ratio between rollout probabilities, whereas SDSD biases the reinforcement signal by the similarity to sampled demonstrations. This inherent mode-seeking behavior is further amplified by SDSD, leading to a sharpening of the probability mass towards a subset of already-common solution modes.

Figure 1: RL and SDSD treat correct rollouts differently, with SDSD biasing reinforcement towards common rollouts and flattening pass@k curves; diversity collapse emerges prominently as more rollouts fail to improve coverage.

Theoretical Characterization

The paper rigorously derives the optimal policy for SDSD under a KL-regularized reverse KL objective. Unlike standard RL, which exponentiates reward to tilt the base policy, SDSD tilts by the expected pointwise conditional mutual information (PCMI) between the candidate rollout and the demonstration:

$$\pi^*_{\text{SDSD-KL}}(y \mid x) \propto \pi_0(y \mid x) \exp \left( \frac{1}{1+\beta} \mathbb{E}_{y^{corr} \sim p_{corr}} \left[ i(y; y^{corr} \mid x) \right] \right)$$

where $i(y; y^{corr} \mid x)$ denotes the log-ratio of base policy probability with and without conditioning on a demonstration. This formulation reveals that SDSD amplifies existing distributional imbalances, increasing the likelihood for rollouts that are more congruent with sampled demonstrations, irrespective of reward. The resulting policy is mathematically sharper; probability ratios between equally correct solutions are not preserved unless they are equally supported by the expected PCMI. This can yield mode collapse even in environments with multiple equivalent optimal solution regions.

Figure 2: Mode collapse in SDSD: probability mass concentrates on one of two equally-rewarded modes, unaffected by KL regularization—whereas RL maintains coverage of both.

Empirical Findings

Diversity Metrics

To diagnose diversity collapse, the authors introduce two distinct metrics:

• Functional diversity: rate at which additional samples solve new problems (measured by the slope of pass@k curves).
• Semantic diversity: diversity in high-level strategies or solution paths (not token-level entropy).

Token-level entropy is shown to be an inadequate proxy for either functional or semantic diversity; models with similar entropy can exhibit drastically different diversity profiles.

Controlled Graph Path-Finding Task

On a synthetic graph path-finding task with precisely measurable semantic diversity, SDSD models achieve competitive pass@1, but exhibit extremely flat pass@k curves, signifying low functional diversity. Models trained with GRPO (on-policy RL) or GRPO with explicit diversity rewards cover more concept paths and generalize to OOD tasks requiring diverse strategies. SDSD fails on out-of-distribution settings demanding exploration of long or rare paths.

Figure 3: SDSD is performant in pass@1 but its pass@k curve and semantic diversity are notably inferior, demonstrating collapse to a single solution mode.

Science QA Benchmarks

Across multiple scientific QA tasks, SDSD variants (including those using ensembles of demonstrations) consistently achieve higher pass@1, but their pass@k curves flatten rapidly. GRPO-trained models, in contrast, demonstrate much steeper scaling, with significant gains in coverage when multiple rollouts are generated.

Figure 4: Pass@k curves for Science QA reveal that SDSD, though strong on average (pass@1), fails to leverage additional rollouts to improve coverage.

External Demonstrations

The diversity collapse persists even when demonstrations for SDSD come from external models, including datasets with high semantic diversity. The student invariably converges to low diversity, regardless of demonstration variability, signifying a fundamental limitation of the SDSD mechanism.

Figure 5: Training SDSD with diverse demonstrations leads to strong in-distribution accuracy, but semantic diversity remains low—diversity collapse is unaffected by demonstration richness.

Explicit Diversity Measurements

Comparative analysis across multiple model sizes substantiates that SDSD consistently underperforms in both pass@k scaling and explicit concept diversity. Diversity collapse is robust to architectural scaling and teacher variations.

Figure 6: Across model scales, SDSD exhibits flatter pass@k curves, indicating poor coverage in both in-distribution and harder graph settings.

Figure 7: SDSD models have persistently lower semantic diversity over training compared to RL-based baselines.

Token-Level Entropy: Insufficiency for Diversity Assessment

Token-level entropy does not correlate reliably with either functional or semantic diversity. SDSD models can have both higher or lower entropy than baselines, yet exhibit diversity collapse according to other metrics.

Figure 8: Concept Graph instance illustrates semantic diversity in solution paths.

Implications and Future Directions

Theoretical and empirical findings highlight a fundamental limitation of SDSD: a self-reinforcing loop of alignment bias, which narrows the solution distribution. This diminishes model robustness, OOD generalization, and ability to discover diverse strategies in multi-modal reasoning settings. The practical implication is that average accuracy (pass@1) is an insufficient metric; accurate yet non-diverse models can fail catastrophically when deployed on tasks requiring exploration or adaptation.

Potential mitigations include diversity rewards (as in GRPO+diversity), best-of-N training objectives, and explicit balancing between quality and diversity, as explored in recent RLHF and preference optimization literature. There is also scope for further theoretical exploration of privileged information distillation (e.g., using environmental feedback rather than static demonstrations), and for designing distillation objectives that better preserve distributional ratios among correct outputs.

Conclusion

SDSD is effective at maximizing average accuracy but exhibits strong diversity collapse, as evidenced by multiple metrics and benchmarks. The PCMI-based bias in teacher conditioning leads to distributional sharpening and mode collapse, distinct from standard RL. As the field advances towards deploying LLMs in diverse, complex domains, explicit monitoring and optimization of diversity—not just accuracy—will be crucial. Theoretical insights presented here provide a foundation for future algorithmic innovation in distillation protocols and post-training strategies.