On-Policy Self-Distillation with Sampled Demonstrations: The Hidden Cost of Reduced Output Diversity

Paper: 2606.26091 Authors: Andrei Liviu Nicolicioiu, Mohammad Pezeshki, Aaron Courville Categories: cs.LG, cs.AI

The Gap

Existing self-distillation methods (e.g., Conditional Training, Self-Consistency, Expert Iteration) focus narrowly on boosting pass@1 accuracy. They treat the teacher’s conditioning on a correct demonstration as a harmless source of dense feedback. The implicit assumption is that this feedback is unbiased — it simply rewards correct patterns. Prior work had not investigated whether this conditioning introduces systematic biases that affect the distribution of generated rollouts.

This paper identifies a concrete blind spot: on-policy self-distillation with sampled demonstrations reduces output diversity, flattening pass@k curves and degrading out-of-distribution performance. The logical path from gap to conclusion is:

+--- Problem: pass@k flattens under self-distillation ---+
|                                                        |
v                                                        v
+--- Assumption: teacher feedback is unbiased -----------+
|                                                        |
v                                                        v
+--- Method: derive optimal self-distillation policy ---+
|          as CMI-weighted tilt of base distribution     |
|                                                        |
v                                                        v
+--- Evidence: graph path-finding + science QA ----------+
|   (GSM8K, MMLU, ARC) — compare RL vs self-distill     |
|                                                        |
v                                                        v
+--- Conclusion: self-distillation amplifies existing    |
|   probability gaps; RL preserves ratio among correct   |
|   rollouts; diversity loss is structural, not random   |
+--------------------------------------------------------+

The Increment

One sentence: Before this paper, practitioners used self-distillation to boost pass@1 without realizing that it systematically reduces output diversity. After this paper, we understand that the degradation arises from a conditional mutual information (CMI) weighting that amplifies existing biases in the model’s own distribution.

Core Mechanism

The method is primarily an analytical diagnosis, supported by controlled experiments. The core insight is that self-distillation with a sampled demonstration as teacher context changes the learning objective in a subtle but critical way compared to on-policy RL.

In standard on-policy RL (e.g., policy gradient), the model generates rollouts, receives a reward based on whether each rollout leads to a correct final answer, and updates to increase the probability of correct rollouts. Crucially, among equally correct rollouts, the policy is free to maintain or even spread probability mass — RL does not penalize diversity as long as all correct actions are rewarded. The optimal policy under RL simply redistributes mass uniformly among rewarding actions (in the simplest case).

In on-policy self-distillation, the teacher is the same model, but it is conditioned on a correct demonstration sampled from the model’s own distribution. The teacher then scores each student rollout based on its similarity to that demonstration (implicitly, through the token-level feedback). This conditioning introduces a bias: the teacher prefers rollouts that resemble the demonstration, even if they differ only in style or strategy. The paper shows analytically that the optimal policy under this objective tilts the base distribution by a term proportional to the conditional mutual information between the student’s rollout and the demonstration.

Data flow in self-distillation:

+----------+               +--------------------------+
|  Model   | -> rollouts ->| Teacher (conditioned on  |
| (student)|               |  sampled correct demo)   |
+----------+               +--------------------------+
     ^                           |
     | (same model,              | (scores: high if similar to demo)
     |  different roles)         v
     +--- correct demo ---+   +----------+
                          +-->| Update   |
                              +----------+

Data flow in on-policy RL:

+----------+               +--------------------------+
|  Model   | -> rollouts ->| Reward (outcome-based:   |
|          |               |  1 if correct, 0 if not) |
+----------+               +--------------------------+
                                |
                                v
                           +----------+
                           | Update   |
                           +----------+

Structural metaphor: Think of a writing class. The student writes several essays (rollouts). In the RL world, the teacher grades purely on factual correctness (e.g., “does the conclusion match the data?”). All correct essays are rewarded equally, so the student can freely explore different arguments. In the self-distillation world, the teacher picks one correct essay from the student’s own previous work (the demonstration), and then grades new essays based on how similar they are to that chosen essay — rewarding similar phrasing, structure, and examples. The student quickly learns to mimic that one essay style, abandoning all other equally correct ways of writing. The bias comes from the teacher using a reference written by the student himself: it reinforces the student’s existing stylistic preferences, creating a positive feedback loop that collapses diversity.

Key Concepts

• Pass@1 vs. Pass@k: Pass@1 measures the accuracy of a single sampled generation. Pass@k measures the probability that at least one of k independent generations contains a correct answer. High pass@k requires diversity: if all k rollouts are identical, pass@k = pass@1. This paper shows that self-distillation improves pass@1 but flattens pass@k, meaning the model becomes a one-trick pony — it generates the same correct answer over and over, wasting the capacity of multiple samples.

• Conditional Mutual Information (CMI): Intuitively, CMI measures how much information a rollout carries about the correct demonstration, beyond what is already known from the input. The paper proves that the optimal self-distillation policy reweights the base distribution by a factor proportional to the pointwise CMI between the rollout and the demonstration. This means rollouts that are statistically similar to the demonstration get boosted, while equally correct but dissimilar rollouts are suppressed. Concretely, if two reasoning paths lead to the same answer, but one shares more tokens with the demonstration, the teacher favors that one, gradually collapsing variety.

• Out-of-Distribution (OOD) Generalization: When the model faces a new problem that requires a strategy it has suppressed during self-distillation (e.g., a different path in the graph, or a different method for a science question), it fails because that strategy no longer has enough probability mass. This is the practical cost of reduced diversity.

Framework Shift

Below is a napkin sketch comparing the mainstream view before this paper with the updated understanding after.

Before (mainstream assumption):        After (this paper's insight):
+----------------------------+         +----------------------------+
| Self-distillation          |         | Self-distillation          |
|   = RL + dense feedback    |         |   = RL + biased feedback   |
|   = better learning        |         |   = better pass@1         |
|                             |         |   but worse pass@k        |
| Assumed effect on diversity:|         | Real effect on diversity: |
|   neutral or slightly +    |         |   negative (amplifies     |
|   (more feedback -> more   |         |   existing modes, kills   |
|    exploration?)           |         |   alternative strategies) |
+----------------------------+         +----------------------------+

Key structural difference:
Before: Teacher uses external correctness signal => unbiased.
After: Teacher uses a demonstration sampled from the student itself => self-referential bias.

One sentence: From treating self-distillation as a harmless upgrade to RL, the core shift is recognizing that conditioning on a sampled correct demo turns the teacher into a self-referential scoring function that systematically reduces diversity.

Expert Assessment

Problem choice: This is a real and timely gap. As LLM-based reasoning systems increasingly use self-consistency and self-distillation, the hidden cost of diversity loss has practical consequences for reliability and OOD performance. The field had been rushing to improve pass@1 without monitoring pass@k. This paper sits at a critical juncture — it’s not a manufacturing gap.

Method maturity: The theoretical derivation (CMI tilt) is a clean analytic result that explains the observed behavior. The controlled graph-path task is a well-chosen “minimal model” to isolate the effect. The science QA benchmarks are standard. If anything, the paper could have explored more diverse transformer architectures, but the core insight is robust. Not brute force — clever and parsimonious.

Experimental integrity: Baselines are fair: they compare self-distillation against on-policy RL with the same model, same number of rollouts, same compute. The metrics include pass@1, pass@k, distinct n-grams, and OOD accuracy. No obvious red flags. One could argue for more diverse OOD settings, but the evidence is convincing.

Writing quality: The paper is well-structured, but the theory section could be more intuition-friendly. The paper uses mutual information without a walkthrough example, which may lose readers. If the authors added a concrete calculation of CMI for a simple symbolic sequence, the paper would be much more accessible. The experimental section is solid.

Verdict: Weak accept — the discovery is important and the evidence is sufficient, but the contribution is more diagnostic than prescriptive (no new algorithm is proposed). Nonetheless, this paper will change how practitioners evaluate self-distillation.

Takeaways

• Monitoring diversity is cheap and mandatory: Whenever you use self-distillation, track pass@k for several k values and compare against a pure RL baseline. Also compute distinct-n-gram ratio or self-BLEU to quantify collapse.
• For problems requiring diverse strategies (code generation, open-ended QA, multi-step reasoning), prefer on-policy RL over self-distillation, or add a diversity regularizer (e.g., KL penalty against a uniform prior over correct rollouts).
• If you must use self-distillation, consider replacing the sampled demonstration with a fixed external set of correct demonstrations (e.g., from a human-written corpus) to break the self-referential bias. This paper’s CMI analysis suggests that the bias comes from sampling the demo from the model itself.
• The theoretical framework (CMI-weighted policy tilt) can be applied to other self-training paradigms — for example, self-imitation learning in offline RL, or bootstrapped DQN in reinforcement learning. The general lesson: conditioning on a self-generated “expert” trajectory can collapse exploration.