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

Technical Summary: On-Policy Self-Distillation with Sampled Demonstrations Reduces Output Diversity

Problem Statement
On-policy self-distillation (SDSD) has emerged as a powerful technique for improving Large Language Model (LLM) performance, particularly in achieving high pass@1 accuracy. In this paradigm, a single model acts as both teacher and student; the teacher is conditioned on a "privileged" correct demonstration (often a correct rollout generated by the student itself) to provide dense, token-level feedback. While recent methods like SDPO, SDFT, and OPSD have demonstrated strong average performance, this paper investigates a hidden cost: a significant reduction in rollout diversity. The authors observe that SDSD models often exhibit flattened pass@k curves, meaning that generating additional rollouts fails to solve new problems or cover diverse solution strategies. This contrasts with on-policy Reinforcement Learning (RL) methods (e.g., GRPO), where increasing the number of samples typically yields steep improvements in problem coverage. The core problem is that SDSD may trade functional and semantic diversity for average accuracy, leading to poor generalization in out-of-distribution (OOD) settings that require diverse strategies.

Methodology and Theoretical Analysis
The paper formalizes the mechanism behind this diversity collapse through a theoretical derivation of the optimal self-distillation policy.

1. Theoretical Framework:

   • Standard RL: The optimal policy for KL-regularized RL is a tilted version of the base distribution, modulated by the reward function. Crucially, if two rollouts have the same base probability and the same reward, their probability ratio remains unchanged in the optimal policy.
   • SDSD Optimal Policy: The authors derive that the optimal policy for self-distillation with sampled demonstrations (SDSD-KL) tilts the base distribution not by reward, but by the expected Pointwise Conditional Mutual Information (PCMI) between the student's rollout and the sampled demonstration.
   • The Sharpening Mechanism: The PCMI measures how much more likely the teacher finds a specific rollout when conditioned on a demonstration. If a rollout is already more probable under the base policy, it is more likely to be sampled as a demonstration. Consequently, the teacher (which shares the base model's biases) assigns higher probability to rollouts similar to that demonstration. This creates a "rich-get-richer" feedback loop: the PCMI term amplifies existing probability gaps, making likely rollouts even more likely and suppressing less common but equally correct solutions. Unlike standard RL, which preserves ratios among equally correct rollouts, SDSD can sharpen the distribution, leading to mode collapse.

2. Experimental Setup:

   • Controlled Graph Path-Finding: The authors introduce a synthetic task involving graph path-finding with concept chains (e.g., birds, fruits). The task is designed with multiple valid paths (semantic diversity) but biased training data favoring shorter, easier paths. This allows for precise measurement of "semantic diversity" (number of distinct concept categories explored) and "functional diversity" (slope of the pass@k curve).
   • Science QA Benchmarks: Experiments are conducted on verifiable reasoning tasks from SciKnowEval (Biology, Chemistry, Materials, Physics) using Qwen3 and Olmo models.
   • Baselines: The study compares SDSD (conditioned on sampled correct rollouts) against standard GRPO and GRPO with an explicit diversity reward. Variants include using external diverse demonstrations to test if the source of the demonstration mitigates the issue.

Key Contributions

1. Theoretical Characterization: The paper proves that the optimal SDSD policy tilts the base distribution by expected PCMI rather than reward. It demonstrates that this mechanism inherently amplifies probability imbalances among equally correct rollouts, a property absent in standard RL objectives.
2. New Diversity Metrics: The authors argue that token-level entropy is insufficient to capture diversity. They introduce and utilize functional diversity (measured by the slope of pass@k curves) and semantic diversity (measured by the variety of high-level strategies or concept categories used) to diagnose the collapse.
3. Empirical Evidence of Collapse: Through experiments on graph path-finding and science QA, the paper shows that while SDSD matches or exceeds RL in pass@1 accuracy, it exhibits substantially lower functional and semantic diversity. SDSD models fail on OOD tasks (e.g., harder graphs requiring diverse strategies) that RL-based models can solve.
4. Robustness to Demonstration Source: The study demonstrates that diversity collapse persists even when using diverse external demonstrations. The self-distillation mechanism itself, rather than the source of the demonstration, drives the loss of diversity.

Results

• Graph Path-Finding: SDSD models achieved high pass@1 on in-distribution data but showed nearly flat pass@k curves, indicating that additional rollouts did not solve new problems. In contrast, GRPO showed steep pass@k improvement. SDSD models failed completely on "Harder Graph" OOD tasks requiring diverse exploration, whereas GRPO succeeded. Semantic diversity scores for SDSD were the lowest among all methods.
• Science QA: SDSD variants achieved higher pass@1 than GRPO but significantly lower pass@16. The pass@k curves for SDSD flattened quickly, confirming low functional diversity.
• Entropy Analysis: Token-level entropy was found to be an unreliable metric. In some cases (QA tasks), SDSD had higher token-level entropy than GRPO despite having lower functional diversity and worse pass@k performance.

Significance and Claims
The paper claims that on-policy self-distillation with sampled demonstrations introduces compounding biases that progressively reduce rollout diversity. The significance lies in identifying that high average accuracy (pass@1) can mask a critical failure in exploration and generalization. The authors argue that the "hidden cost" of SDSD is a flattened pass@k curve, rendering models brittle in scenarios requiring diverse strategies.

The authors conclude that average accuracy alone is insufficient for evaluating self-distillation methods. They assert that functional and semantic diversity are at risk of collapse and must be explicitly monitored during training and deployment. The work suggests that while SDSD is effective for refining known solutions, it may be inherently unsuited for tasks requiring the discovery of novel solution strategies unless diversity is explicitly enforced through external mechanisms (like diversity rewards in RL), which SDSD lacks natively. The paper does not propose a new algorithm to fix this but rather provides a theoretical and empirical diagnosis of the phenomenon.