On-Policy Self-Distillation with Sampled Demonstrations Reduces Output Diversity
arXiv:2606.26091v1 [cs.LG] 24 Jun 2026
Andrei Liviu Nicolicioiu1,2,3,†
Mohammad Pezeshki3,*
Aaron Courville1,2,4,*
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, selfdistillation can amplify existing probability gaps, concentrating mass on alreadydominant modes. On a controlled graph path-finding task and science questionanswering 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.
1
Introduction
Current LLM post-training approaches to instill capabilities in models have different tradeoffs, with supervised fine-tuning (SFT) learning initial behaviors and on-policy RL methods refining and exploring new approaches (Zhang et al., 2025). In between, on-policy distillation (Agarwal et al., 2024; Lu and Thinking Machines, 2025) uses a stronger teacher to guide a student using studentgenerated data. Self-distillation goes further by eliminating the external teacher entirely: the same model, conditioned on privileged information, such as a correct solution or environmental feedback, provides dense token-level feedback on the student’s own generations. Recent methods, including SDPO (Hübotter et al., 2026), SDFT (Shenfeld et al., 2026), OPSD (Zhao et al., 2026; Penaloza et al., 2026), and OPCD (Ye et al., 2026), instantiate this approach, achieving strong performance across several tasks such as scientific question-answering, continual learning tasks, and agentic tasks. In this paper, we investigate a specific case, Self-Distillation with Sampled Demonstrations (SDSD) where student rollouts are guided by a teacher with demonstrations in its context. The demonstrations could come from correct student rollouts (exactly the setup of Hübotter et al. (2026) with student demonstrations) or from external models (Zhao et al., 2026; Penaloza et al., 2026). Here, we find that good accuracy might come at a hidden cost. SDSD models with sampled demonstrations exhibit pass@k curves with small or nearly flat slopes (Figures 3, 4): generating more rollouts fails to solve new problems. By contrast, models trained with on-policy RL (e.g., GRPO) show steep pass@k improvement, where each additional sample meaningfully increases problem coverage. SDSD could thus trade rollout diversity for average accuracy. 3 FAIR at Meta
4 CIFAR AI Chair
† Correspondence: [email protected]
1 Mila
* Equal advising
Preprint.
2 Université de Montréal
Self-Distillation (SDSD)
Question x → Student generates:
Question x → Student generates:
y1 y2 y3 y4 ✓ ✓ ✓ × 1
1
1
0
y1 y2 y3 y4 ✓ ✓ ✓ × y 1 selected for teacher context
Verifier: same reward for all correct scores every correct y i : equal y1 : equal y2 : equal y3 : All correct reinforced equally
i , y1 ) Teacher: p(yti | x, y<t
reinforces every rollout differently: strong y1 : some y2 : y3 :
little
Functional Diversity: pass@[1 → k]
1.0 0.9 0.8 0.7 0.6 0.5 0.4 0.3
Performance
pass@k
Regular RL (e.g., GRPO)
GRPO SDSD 1
2
4
8
k
16 32 64
Figure 1: RL and self-distillation treat correct rollouts differently, with consequences for rollout diversity. Left: A binary verifier gives equal reward to all correct rollouts, so RL reinforces them uniformly. Middle: Self-distillation conditions the teacher on a sampled correct rollout, typically the most probable one, so the teacher’s feedback is strongest for similar rollouts and weakest for those taking a different approach. Right: On a Graph Path finding task, this gap manifests as flatter pass@k curves for Self-Distillation with Sampled Demonstrations (SDSD): generating more rollouts fails to solve new problems, unlike GRPO where each additional sample meaningfully increases coverage.
We argue that SDSD introduces compounding sources of bias that can progressively reduce rollout diversity. Intuitively, a rollout is aligned with a demonstration when the two share more structural or stylistic features, causing the teacher conditioned on that demonstration to assign it higher probability and therefore reinforce it more strongly during training. This creates a bias toward solutions that resemble the sampled demonstrations. In particular, a teacher may struggle to effectively guide a correct but less typical rollout when conditioned on a more standard or canonical demonstration, simply because the two trajectories share fewer common patterns. As a result, distinctive yet valid solution strategies receive weaker learning signals. Over repeated training updates, this preference can compound: rollouts that are more aligned with previously sampled demonstrations become increasingly reinforced, while less aligned, but still correct, solutions are gradually suppressed (See Fig. 1). We hypothesize that this feedback loop contributes to the reduced rollout diversity and flattened pass@k scaling observed in SDSD models. We formalize this effect by deriving the optimal self-distillation policy (Proposition 2). The resulting policy is a tilted version of the base distribution, where the tilt is determined by the expected pointwise conditional mutual information (PCMI) between a student rollout and the sampled demonstration. PCMI measures how much conditioning on a demonstration increases the model’s preference for a particular rollout. Unlike a binary reward, which treats all correct rollouts equally, PCMI distinguishes among equally valid solutions and assigns greater weight to rollouts that are already more compatible with the demonstrations and the base policy. Consequently, self-distillation with sampled demonstrations can amplify existing probability imbalances: likely rollouts become increasingly likely, while less common but correct solutions are progressively suppressed, reducing rollout diversity. To diagnose the reduced diversity, we use two notions of diversity beyond the commonly-used token-level entropy. Functional diversity is the rate at which additional samples solve new problems, reflected in the slope of pass@k curves. Semantic diversity measures whether rollouts differ in their high-level strategy (e.g., different paths through a graph, different proof approaches in math) rather than just surface-level wording. We show that token-level entropy fails to capture either notion (§4.4). Our contributions are: • We prove that the optimal policy of self-distillation with sampled demonstrations tilts the base distribution by expected PCMI rather than reward, and that this can amplify probability gaps among equally correct rollouts, a property absent from standard RL (Prop. 2, Remark 1). • We introduce a graph path-finding task, in which semantic diversity, the number of distinct concept categories explored, is precisely measurable and directly predicts out-of-distribution generalization (§4.1, Fig. 3). 2
• On graph path-finding and science QA tasks (Feng et al., 2024), we show that self-distillation with sampled demonstrations achieves competitive pass@1 but substantially lower functional and semantic diversity than RL, failing on out-of-distribution tasks that require diverse strategies (Fig. 3, 4).
2
Background: Self-Distillation with Sampled Demonstrations
We review the self-distillation with sampled demonstration framework of Hübotter et al. (2026), which uses a correctly verified rollout as privileged context for the teacher. For each question x, the student policy generates a group of N rollouts yn ∼ πθ (· | x), Y(x) = {y1 , . . . , yN }. Let C(x) ⊆ Y(x) denote the subset verified as correct. For each rollout y ∈ Y(x), a correct rollout y corr ∈ C(x) is sampled uniformly from this subset and provided to the teacher as context. The teacher is an exponential moving average (EMA) of the student with parameters θ̄, conditioned on the question x and the correct rollout y corr : πθ̄ (· | x, y corr , y<t ). The training objective minimizes the token-level KL divergence between the student and this context-conditioned teacher: |y| X X 1 Eycorr ∼C(x) KL(πθ (· | x, y<t ) ∥ sg[πθ̄ (· | x, y corr , y<t )]) , LSD (θ; x) = N t=1 y∈Y(x)
where sg[·] denotes stop-gradient. The corresponding token-level gradient (Hübotter et al., 2026) is: |y| X 1 X π (ŷ | x, y ) θ t <t ∇θ LSD (θ; x) = ∇θ log πθ (ŷt | x, y<t ) . Eycorr Eŷt ∼πθ log corr , y ) (ŷ | x, y N π t <t θ̄ t=1 y∈Y(x)
Self-distillation thus replaces a single scalar reward for the full sequence with a dense per-token correction signal derived from a context-conditioned version of the model itself. On-policy RL is mode-seeking, SDSD is even more. Any on-policy method exhibits mode-seeking behavior concentrating its rollouts into a smaller subset (Tajwar et al., 2024; Chen et al., 2025a). SDSD as a reverse KL objective (Bishop, 2006) has the same behavior but even more pronounced due to additional compounding biases. In our setting, self-distillation involves two crucial samplings: that of the student rollout and that of the demonstration (either from the student’s correct rollouts or from an external source). The alignment between these two introduces a bias towards common responses. Consider how two student rollouts with equal reward, one common and one highly novel, are treated. It is more likely to sample a demonstration that resembles the common rollout. Then, the teacher has a bias to give more probability to rollouts that are similar to its context, meaning similar to the demonstration. Together this will lead to the common rollout being upweighted more than the unique one. This leads to a rich-get-richer (likely-get-likelier) behavior stronger than in standard RL, arising from both double sampling and its alignment and the preferences of the teacher. Moreover, while all on-policy methods exhibit mode-seeking behavior due to optimization, selfdistillation uniquely introduces incentives for loss of diversity even at the level of the optimal policy.
3
Optimal Policy of Self-Distillation
We derive the optimal self-distillation policy and characterize how it differs from standard RL. For ease of presentation, we use a sequence level objective first, and refer to the Appendix B for full derivations and token-level presentation section B.3. Let x denote the input prompt, y an output sequence, and π0 (y | x) the base policy. We optimize a student policy π(y | x). All distributions are assumed strictly positive on their support. Proposition 1 (Optimal policy for standard KL-regularized RL). The standard RL objective is: max Ey∼π(· | x) [R(y | x)] − βRL KL(π(· | x) ∥ π0 (· | x)) . (1) π
It is well known (Korbak et al., 2022; Rafailov et al., 2023) that the optimal policy of this objective is the following tilted distribution: 1 ∗ πRL (y | x) ∝ π0 (y | x) exp R(y | x) . (2) βRL 3
The optimal RL policy consists of the base policy modulated by the reward. Thus, two rollouts with the same probability under the base and the same reward will be as likely under the optimal policy. We now derive the analogous result for self-distillation. We take the teacher to be the fixed base policy conditioned on a correct demonstration, and optimize the student to minimize the reverse KL with the teacher, regularized by a KL penalty to the base model. In practice, explicit KL regularization to the base policy is not always used in RL or self-distillation. However, since training starts from the base policy and models are rarely trained to convergence, the policy typically remains close in KL to its starting point, making this a reasonable modeling choice. Let pcorr (· | x) denote a reference distribution over correct demonstrations for input x, from which y corr is sampled. In practice, this can be the empirical distribution over C(x), the student’s own correct rollouts (§2), or a distribution over external demonstrations (§4.3). Let us define the following objective that takes the expectation over demonstrations. Proposition 2 (Optimal policy for SDSD-KL). Let the self-distillation + KL objective: h i min Eycorr ∼pcorr (· | x) KL(π(· | x) ∥ π0 (y | x, y corr )) + β KL(π(· | x) ∥ π0 (· | x)) . π
(3)
The optimal policy of (3) is ∗ πSD-KL (y | x) ∝ π0 (y | x) exp
1 Eycorr ∼pcorr (· | x) i(y; y corr | x) , 1+β where the pointwise conditional mutual information (PCMI) is defined as: i(y; y corr | x) := log
π0 (y | x, y corr ) . π0 (y | x)
(4)
(5)
The RL optimal policy (2) tilts the base policy by the reward. The self-distillation optimum instead tilts by the expected PCMI (5): a log-ratio measuring how much more likely the teacher finds y after conditioning on a demonstration. When y corr is relevant and supports y, the PCMI is positive; when y corr is contradictory to y, it is negative. The base policy is thus tilted not by task reward, but by the teacher’s assessment of how well each rollout aligns with correct demonstrations. Remark 1 (Ratio for two correct sequences under SDSD-KL). Let y1 and y2 be two correct sequences for the same input x, and suppose π0 (y1 | x) = k π0 (y2 | x) for some k ≥ 1. Then ∗ 1 πSD-KL (y1 | x) corr corr corr = k exp i(y ; y | x) − i(y ; y | x) . (6) E 1 2 y ∼pcorr (· | x) ∗ (y2 | x) πSD-KL 1+β When does sharpening occur? SDSD-KL preserves the base-policy ratio k only when the two sequences have the same expected PCMI. If the demonstrations y corr support on average y1 more than y2 , the ratio between the two rollouts becomes even larger under self-distillation, leading to sharpening where likely rollouts become even more likely. This contrasts with the RL optimal policy (2), which maintains the initial ratio of k when both rollouts are equally correct, since the reward tilt cancels out. This shows sharpening occurs when the expected PCMI is higher for the already-probable rollout. Such rollouts are more likely to be aligned to the demonstrations, and the teacher shares the same biases as the student, thus this becomes highly likely. The same derivations and implications carry over to the token-level objective, resulting in a bias for sharpening the distribution of the next-token, leading to loss of diversity of the whole rollout (see Appendix B.3). All on-policy learning methods, like GRPO or self-distillation have mode-seeking behavior due to optimization, but the previous remark shows that self-distillation has an optimal policy that can be sharper than the initial one. 3.1
Illustrative Example: Mode Collapse Under Self-Distillation
We verify the sharpening predicted by our theory in a minimal controlled environment. The action space is D=100 discrete actions split into four quarters: the first and third quarters are rewarded, while the other two are not. An ideally diverse policy would place mass on both rewarded modes. 4
Base model
0.04
Trained =0.1
Avg Reward=0.976 coverage=0.50
Trained =0.5
Avg Reward=0.970 coverage=1.00
Avg Reward=0.877 coverage=1.00
P(y)
GRPO
Trained =0.0
Avg Reward=0.348 coverage=1.00
0.02 0.00
0
50
75
100 0
25
50
75
100 0
25
50
75
100 0
25
50
75
100 0
Avg Reward=0.990 coverage=0.50
25
50
75
100 0
25
50
75
100 0
Avg Reward=0.977 coverage=0.50
25
50
75
100
25
50
75
100
Avg Reward=0.899 coverage=0.50
P(y)
SDSD
0.04
25
Avg Reward=0.348 coverage=1.00
0.02 0.00
0
y
y
y
y
Figure 2: Illustrative example: SDSD collapses to a single high-reward mode, regardless of KL regularization, while GRPO with KL regularization > 0 covers both modes. The environment has two equally-valued reward regions, so an ideal policy would maintain coverage of both.
We parameterize the student policy πθ over a four-bump base distribution. The teacher for SDSD is constructed as πθ̄ (y | y corr ) ∝ πθ (y) · K(y, y corr ), where K is a Gaussian kernel centered on a correct sample y corr drawn from the student’s own correct outputs. This locally upweights the student’s mass near each observed correct sample, exactly the mechanism that produces PCMI sharpening in our theory. Figure 2 shows the result. GRPO recovers both rewarded modes for any β > 0. SDSD collapses to whichever rewarded region the base policy slightly favors and stays there for every β: the more probable region produces more correct samples, those samples become teacher contexts, and the teacher’s kernel-shaped feedback reinforces nearby points, a self-reinforcing loop that no level of KL regularization to the base undoes.
4
Experiments
4.1
Concept-Graph Setting: Loss of Semantic Diversity and OOD Performance
We design a controlled setting, challenging for LLMs, with a precise definition of semantic diversity and a direct link between diversity and downstream performance. Controlled experimental setting. We introduce a graph path-finding task, where we generate multiple graphs and create a query for each one. For each query, we prompt an LLM with a representation of the graph in context and ask it to generate a path between two points. See section B.4 for an example of such query. A graph node represents an instance of a named concept (e.g., specific birds: heron, pigeon; fruits: orange, cherry) as seen in Fig. 6a. The graph has a star structure: a central start node connects to multiple concept chains, each consisting of nodes from the same concept (e.g., all birds or all fruits), each ending at a shape node (e.g., diamond, square). Two of the endpoints have the same name and represent the target; the remaining two are distractors. Multiple valid paths to the targets exist, each passing through a different concept chain. Training graphs are biased: some paths to the target are short (11 nodes, easier), while others are long (15 nodes, harder), creating an incentive for models to exploit easy routes and ignore harder but equally valid alternatives. To check the robustness of the learned models, we evaluate on an in-distribution test set and two out-of-distribution datasets: a larger-graphs dataset, in which all paths to end nodes have a fixed length of 20, and a harder-graphs dataset, in which all paths have fixed length of 11, but one edge is removed from chains leading to one of the two target nodes, giving fewer valid solutions. Baselines. We train Qwen3-1.7B (Yang et al., 2025) models using variants of GRPO and SDSD, on a dataset of 16k graphs for training and 128 for testing. We train for 1000 steps, with a mini-batch 5
Harder Graphs eval
pass@k
pass@k
pass@k
1
2
4
8
k
16
32
64
1
2
4
8
k
16
32
64
1.2
1.4
0.2
0.2
1.1
1.2
0.0 1
2
4
8
k
16
32
64
GRPO + Diversity GRPO SDSD
1.3
1.6
0.4
0.4
0.85
Larger Graphs Diversity 1.4
1.8
0.6
0.6
0.90
In Distribution Diversity
0.8
0.8
Concept Diversity
Larger Graphs eval
1.0
0.95
Concept Diversity
In Distribution eval
1.00
1.0
1.0
0
200
400
600
Training Step
800
1000
0
200
400
600
Training Step
800
1000
Figure 3: In-distribution and Larger Graphs evaluations show SDSD has good pass@1 performance but, toward pass@k, the curve has a small slope highlighting low functional diversity. The third setup requires the model to have learned diverse rollouts during training, and is completely unsolved by self-distillation. Additionally, the last two figures show that the explicitly defined semantic diversity of SDSD is the lowest. All runs train Qwen3-1.7B; mean and min/max runs with 3 seeds are shown. size of 16 queries, 4 rollouts per query, and a maximum generation length of 8,192 tokens. The experiments are trained on a single GPU using the library of Kazemnejad et al. (2025). GRPO. Standard GRPO (Shao et al., 2024) serves as the primary baseline. At each iteration, the policy model generates N=4 rollouts per query. Within each group of N rollouts, a scalar reward is used to compute a scalar advantage for the whole sequence. GRPO+diversity. GRPO+diversity adds a diversity reward to the score reward to encourage the model to explore different concept chains across its N rollouts. For each rollout, its diversity score is the fraction of the other N − 1 rollouts in the same group that used a disjoint set of concepts. A rollout using different concepts than all the other rollouts of the same query, will have a diversity score of 1. This relates to Li et al. (2025), who multiply the score reward by a clustering-based diversity score. SDSD. SDSD implementation following SDPO (Hübotter et al., 2026), where for each student rollout, the teacher is conditioned on another correct student rollout for the same query. All models add to the main loss a KL regularization to the reference model. The results in Fig. 3 show that SDSD achieves good in-distribution pass@1 and the best pass@1 on the Larger Graph setting. However, its pass@k performance increases very slowly, or not at all for the harder settings. This flat pass@k curve indicates low functional diversity: successive samples rarely solve new queries. The Harder Graph setting can only be solved by models that learned diverse rollouts during training. SDSD’s failure there confirms that it relied exclusively on easy routes. Semantic diversity. We seek rollouts with semantic diversity, capturing meaningful variations in their trajectories, such as different high-level strategies or approaches. In mathematical reasoning, for instance, one might want geometric vs. algebraic approaches, or different theorems. In the graph setting, we define the semantic diversity of a set of rollouts as the number of unique concepts present across all of them. This measures whether a model explores fundamentally different strategies (e.g., following animal chains vs. flower chains) rather than mere surface-level token variation. For each query in the in-distribution and larger graphs testsets, we sample 64 rollouts and compute the average number of unique concepts of the nodes present in them, and show these diversity scores across training in Fig. 3. We find that GRPO has relatively low semantic diversity, but adding a diversity reward significantly improves it. On the other hand, SDSD has the lowest semantic diversity scores, which are correlated with the low functional diversity (slope of pass@k) and low scores in the Harder Graphs dataset that requires diversity. 4.2
Science QA: Functional Diversity in Practice
We highlight the interplay of accuracy and diversity of SDSD and GRPO in science QA settings. Setup. We evaluate on four verifiable reasoning tasks spanning scientific knowledge drawn from SciKnowEval (Feng et al., 2024), a benchmark of multiple-choice science questions. Across all tasks, we train on the training split and evaluate on the held-out test split by generating N =16 rollouts per question and reporting mean accuracy (pass@1) as well as pass@k. 6
Pass@k (%)
Qwen3-8B
biology
Pass@k (%)
material
85
84
80
82
75
78
65
76
60 2
4
8
N (number of rollouts)
16
1
biology
2
4
8
N (number of rollouts)
16
chemistry
70 60 50 2
4
8
N (number of rollouts)
16
1
2
4
8
N (number of rollouts)
2
4
8
N (number of rollouts)
16
16
86 84 82 80 78 76 74
75 70 60 1
2
4
8
N (number of rollouts)
16
physics
2
4
8
1
90
85
85
80
80
75
75
70
16
2
4
2
4
8
N (number of rollouts)
16
16
ScienceQA
GRPO SDSD K=1 SDSD K=3
60 1
8
N (number of rollouts)
65
65
N (number of rollouts)
GRPO SDSD K=1 SDSD K=3
65
70 1
ScienceQA
85 80
material
85 80 75 70 65 60 55
80
1
1
physics
95 90 85 80 75 70 65
80
70
1
Olmo-3-7B-Instruct
chemistry
90
80 75 70 65 60 55 50
1
2
4
8
N (number of rollouts)
16
Figure 4: Pass@k curves for Science QA (4 tasks and average across tasks). SDSD achieves better pass@1 but its curves flatten quickly, indicating low functional diversity. mean ± stderr over 3 seeds.
All experiments follow the implementation and configuration of Hübotter et al. (2026), and use Qwen3-8B (Yang et al., 2025) and Olmo-3-7B-Instruct (OLMo et al., 2025) as the base models, trained with AdamW (Loshchilov and Hutter, 2019) for up to 30 epochs bounded by a 5h training time budget. At each training step, the policy generates N =8 rollouts per question for a batch of 32 questions, sampled with temperature 1.0. Each configuration is run with 3 seeds on 4 Nvidia H200 GPUs. We compare the following models: GRPO. Standard GRPO generating N =8 rollouts per question for a batch of 32 questions. The batch is optimized in mini-batches of 8 questions. SDSD (K = 1). SDPO-style self-distillation (Hübotter et al., 2026) where the teacher is the base model conditioned on one correct student rollout. SDSD (K = 3). We introduce a baseline using an ensemble of K = 3 teachers. Everything is the same as above, but we collect three distinct correct demonstrations (if available) and create an independent teacher from each, then average their distillation losses per student rollout. This requires K forward passes through the teacher, though it adds only 5.8% wall-clock time per training step since generating the student rollouts dominates computation. Fig. 4 shows SDSD variants have flatter pass@k curves than GRPO, indicating lower functional diversity. Table 1 reports mean accuracy across 3 seeds, using the best checkpoint per method by average accuracy. As shown in Fig. 9, SDSD variants attain higher pass@1 but lower pass@k throughout training, consistent with reduced diversity. Table 1: Pass@1 and Pass@16 at the best checkpoint per dataset, for Qwen3-8B and Olmo-3-7BInstruct. Overall both SDSD variants have better Pass@1 but worse Pass@16, indicating low diversity between rollouts. Mean ± std across 3 seeds. Bold = best per column. Biology
Method
Chemistry
Material
Physics
Average
Pass@1
Pass@16
Pass@1
Pass@16
Pass@1
Pass@16
Pass@1
Pass@16
Pass@1
Pass@16
57.2±3.9 57.5±1.2 61.8±1.5
69.3±2.9 60.3±3.2 64.8±1.0
76.6±3.6 78.8±0.9 78.5±0.4
87.3±1.7 86.5±0.1 83.9±1.1
79.3±0.4 78.1±2.3 77.7±1.0
82.3±0.7 80.1±2.6 80.1±0.6
74.4±2.9 76.6±2.6 75.8±2.8
95.6±0.5 88.1±2.7 85.4±3.3
71.9±2.0 72.7±1.1 73.4±0.8
83.6±1.1 78.7±1.2 78.5±0.6
80.9±5.8 56.4±2.4 56.4±3.3
59.8±5.7 79.5±0.5 80.8±1.1
85.2±2.6 83.2±0.8 84.3±0.3
74.9±0.6 76.6±1.3 78.6±1.2
79.6±1.0 84.1±1.4 85.1±1.7
63.1±1.3 67.1±2.7 67.4±1.8
90.5±1.4 78.3±1.5 79.5±2.7
62.0±2.4 69.1±0.9 68.7±2.7
84.0±1.4 75.5±0.9 75.3±2.6
Qwen3-8B GRPO SDSD K=1 SDSD K=3
Olmo-3-7B-Instruct GRPO SDSD K=1 SDSD K=3
4.3
50.2±3.2 53.4±1.4 53.2±3.3
Diverse External Demonstrations Still Lead to Diversity Collapse
Previously, we evaluated SDSD on demonstrations coming from the student’s own correct rollouts. Similar to approaches like OPSD (Zhao et al., 2026), we now investigate the case when the demon7
Performance
Student Diversity
Student performance
1.0 0.8 0.6 0.4 0.2 0.0
In-distribution perf Harder graphs perf y=x
0.0
0.2
0.4
0.6
0.8
Demo source performance
1.0
Semantic Diversity
2.0 1.8 1.6 1.4 1.2 1.0 1.0
1.2
1.4
1.6
1.8
Demonstrations Diversity
2.0
Figure 5: We train SDSD models using external demonstrations that are both correct and diverse. We use multiple datasets of demonstrations with increasing levels of diversity. Each demonstration dataset leads to SDSD models with good in-distribution performance (left plot), but low semantic diversity (right plot), regardless of the level of diversity in the demonstration. This shows that, even with external demonstrations, we still have a problem of diversity collapse, regardless of the level of diversity of the demonstrations. strations come from external models. We will see that self-distillation with external demonstrations still suffers from diversity collapse, regardless of the diversity level of the demonstrations. We analyze this in the Concept Graph setup. We create multiple demonstration datasets, with different levels of diversity, where each query has multiple correct solutions. We obtain the demonstrations from different checkpoints of Qwen3-1.7B GRPO+diversity regularizer models. For each dataset, we train SDSD models as in section 4.1 but this time we condition the EMA teacher on these external demonstrations instead of self-generated demonstrations. This way, we can control the diversity of the teacher demonstrations and see their influence on the diversity of the learned student. All resulting SDSD models have high in-distribution performance (pass@8), similar to the models that generated the demonstrations (Fig. 5). Nevertheless, on the Harder Graph setup (that requires diverse exploration during training) pass@8 performance of the students remains low. Moreover, we compare the semantic diversity of the demonstrations and the diversity of the resulting SDSD models for in-distribution questions. We observe that the student models have low semantic diversity, regardless of the level of diversity of the demonstrations. This shows that using diverse demonstrations in the teacher does not fix the diversity problem. 4.4
Token Entropy Is Not a Sufficient Metric of Diversity
In Fig. 6b on Concept Graphs, SDSD has clearly lower average token-level entropy than the GRPO baselines. This correlates well with SDSD’s lack of semantic diversity. On the other hand, token-level entropy cannot distinguish between GRPO and GRPO+diversity, even though GRPO+diversity is clearly more semantically diverse and has better performance on the Harder Graphs task that requires diverse rollouts during training. Conversely, in the QA setting (Fig. 6c), SDSD has higher token-level entropy than GRPO despite lower functional diversity and pass@16. Again, token-level entropy does not correlate well with a meaningful notion of diversity or performance. This suggests that we need more nuanced notions of diversity than entropy at the token level.
5
Related Work
On-policy self-distillation. Distillation transfers knowledge from a teacher model to a student (Hinton et al., 2015; Bucilă et al., 2006), and on-policy distillation gives teacher guidance on studentgenerated data (Agarwal et al., 2024). Moreover, self-distillation methods like SDPO (Hübotter et al., 2026), OPSD (Zhao et al., 2026; Penaloza et al., 2026), SDFT (Shenfeld et al., 2026), RLSD (Yang et al., 2026), and OPCD (Ye et al., 2026) use the same model as the teacher to give guidance, by conditioning it on privileged information, achieving strong results. As privileged information they use their own correct rollouts, demonstrations from stronger models (a setting that we also use), correct answers, or environment feedback on the student rollouts, such as runtime errors. All these approaches use token-wise supervision to give dense feedback to the student, with Zhao et al. (2026) noting that significant gains come from matching the teacher and student over the whole vocabulary. This dense feedback makes self-distillation achieve high performance in small number of steps (Zhao 8
magpie
triangle
hummingbird
anchovy pelican anchovy
goose
sparrow
canary
salmon
halibut
peacock
cod
perch
start pear peach
magnolia
pineapple
lime blackberry
papaya
square
hibiscus
magnolia pear hibiscus
peach
poppy
azalea
biology
GRPO+diversity SDSD GRPO
0.3 0.2 0.1
0.8 0.6 0.4
0.0 0 orchid
square
SDSD K=3 SDSD K=1 GRPO
1.0
sunflower camellia
raspberry
Concept Graph Entropy
0.4
salmon
mackerel
Token Entropy
turkey
bass
Token Entropy
turkey
triangle
200
400
600
Training Step
800
1000
(b) Token Entropy of Concept Graph models
0
50
100
150
Training Step
200
250
(c) Token Entropy of QA models.
(a) Sample Concept Graph.
Figure 6: (a) A Concept Graph instance with chain length 3: four concept chains radiate from start; the two yellow triangle endpoints are valid targets, the two gray squares are distractors. (b, c) Token-level entropy does not tell the whole story. For ConceptGraph, token-level entropy cannot explain the higher semantic diversity (Fig. 3) of GRPO+diversity compared to GRPO. In the QA setting, average token-level entropy does not correlate with functional diversity or low pass@k, since GRPO has the lowest token-level entropy but highest functional diversity and pass@16.
et al., 2026; Yang et al., 2026) before plateauing or decreasing. SDPO (Hübotter et al., 2026) uses two settings: one where the demonstrations are collected from correct student responses and one where coding runtime feedback is used as privileged information. For the second approach, the alignment between feedback and student rollout should be implicitly higher, since the feedback is generated based on the rollout, thus the teacher should have an easier time understanding their relation and conversely provide good guiding signal. RLSD (Yang et al., 2026) points to an irreducible gap in the objective of self-distillation, which results in privileged information leakage when the optimization is done, as usual, in mini-batches. They propose to fix this by using the teacher guidance to change the magnitude of the RL gradient, but keep the direction given by the verifiable reward. Entropy collapse and mode-seeking. Entropy collapse under on-policy training is widely documented (GX-Chen et al., 2026; Yue et al., 2026; Wu et al., 2025) with prior work attributing it primarily to mode-seeking of the on-policy training. Our analysis identifies an additional mechanism specific to self-distillation: unequal alignment between the sampled rollouts and the sampled demonstrations, as defined by PCMI. Nagarajan et al. (2025) analyze the diversity and creativity of LLMs using controlled graph understanding tasks. Prior work on RLHF shows that RL-trained models are less diverse than SFT models (Kirk et al., 2024). This is addressed by Pass@k-aware training objectives (Chen et al., 2025b), and best-of-N fine-tuning (Chow et al., 2025) that directly optimize for output coverage. Li et al. (2025) uses LLM embeddings to partition the rollouts and use them to compute a diversity score. Multiplying the reward with this diversity score improves the diversity of math reasoning and creative writing.
6
Conclusion and Limitations
Conclusion. We analyzed on-policy self-distillation with sampled demonstrations (SDSD) through the lens of rollout diversity. While SDSD models achieve strong average accuracy, their pass@k curves are often flat, indicating collapsed functional diversity. Theoretically, the optimal selfdistillation policy tilts the base distribution by a pointwise conditional mutual information (PCMI) score between the student rollout and the demonstration. Unlike standard RL objectives, which preserve probability ratios among equally correct rollouts, this mechanism amplifies pre-existing imbalances and reinforces solutions that already align with the demonstrations and base policy. Empirically, we observed this diversity collapse on controlled graph path-finding, scientific QA, and synthetic tasks. Our results show that average accuracy alone is insufficient for evaluating self-distillation methods. Functional and semantic diversity are at risk of collapse and should be explicitly monitored during training and deployment. 9
Scope and limitations. We focus specifically on the variant of self-distillation that uses sampled correct rollouts as demonstrations. We do not analyze settings where the teacher is conditioned on richer privileged signals, such as runtime errors in coding (Hübotter et al., 2026), environmental feedback, or external verifiers, where the learning dynamics may differ substantially. Our theoretical analysis assumes a teacher frozen at the base policy, whereas most practical SDSD implementations, including those used in our experiments, employ an EMA teacher derived from the student itself. In addition, our derivation assumes demonstrations are sampled from the base policy, which more closely resembles OPSD (Zhao et al., 2026) and our external-demonstration graph experiments than the fully self-generated setup of Hübotter et al. (2026). In practice, both EMA teachers and self-generated demonstrations introduce additional forms of self-selection bias beyond those captured by our analysis. Finally, our derivation is presented at the sequence level; however, an analogous token-level derivation yields the same PCMI-based tilt at each next-token distribution, causing the effect to compound autoregressively along a trajectory, as discussed in section B.3.
References Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Ramos Garea, S., Geist, M., and Bachem, O. (2024). On-policy distillation of language models: Learning from self-generated mistakes. In The twelfth international conference on learning representations. Bishop, C. M. (2006). Pattern recognition and machine learning, volume 4. Springer. Bucilă, C., Caruana, R., and Niculescu-Mizil, A. (2006). Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 535–541. Chen, H., Razin, N., Narasimhan, K., and Chen, D. (2025a). Retaining by doing: The role of on-policy data in mitigating forgetting. arXiv preprint arXiv:2510.18874. Chen, Z., Qin, X., Wu, Y., Ling, Y., Ye, Q., Zhao, W. X., and Shi, G. (2025b). Pass@ k training for adaptively balancing exploration and exploitation of large reasoning models. arXiv preprint arXiv:2508.10751. Chow, Y., Tennenholtz, G., Gur, I., Zhuang, V., Dai, B., Kumar, A., Agarwal, R., Thiagarajan, S., Boutilier, C., and Faust, A. (2025). Inference-aware fine-tuning for best-of-n sampling in large language models. In The Thirteenth International Conference on Learning Representations. Feng, K., Shen, X., Wang, W., Zhuang, X., Tang, Y., Zhang, Q., and Ding, K. (2024). Sciknoweval: Evaluating multi-level scientific knowledge of large language models. arXiv preprint arXiv:2406.09098. GX-Chen, A., Prakash, J., Guo, J., Fergus, R., and Ranganath, R. (2026). KL-regularized reinforcement learning for generative modelling is designed to mode collapse. In The Fourteenth International Conference on Learning Representations. Hinton, G., Vinyals, O., and Dean, J. (2015). Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Hübotter, J., Lübeck, F., Behric, L., Baumann, A., Bagatella, M., Marta, D., Hakimi, I., Shenfeld, I., Kleine Buening, T., Guestrin, C., and Krause, A. (2026). Reinforcement learning via selfdistillation. arXiv preprint arXiv:2601.20802. Kazemnejad, A., Aghajohari, M., Sordoni, A., Courville, A., and Reddy, S. (2025). Nano aha! moment: Single file "rl for llm" library. https://github.com/McGill-NLP/nano-aha-moment. GitHub repository. Kirk, R., Mediratta, I., Nalmpantis, C., Luketina, J., Hambro, E., Grefenstette, E., and Raileanu, R. (2024). Understanding the effects of RLHF on LLM generalisation and diversity. In The Twelfth International Conference on Learning Representations. Korbak, T., Perez, E., and Buckley, C. (2022). Rl with kl penalties is better viewed as bayesian inference. In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 1083–1091. 10
Li, T., Zhang, Y., Yu, P., Saha, S., Khashabi, D., Weston, J., Lanchantin, J., and Wang, T. (2025). Jointly reinforcing diversity and quality in language model generations. arXiv preprint arXiv:2509.02534. Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations. Lu, K. and Thinking Machines, L. (2025). On-policy distillation. Thinking Machines Lab: Connectionism. https://thinkingmachines.ai/blog/on-policy-distillation. Nagarajan, V., Wu, C. H., Ding, C., and Raghunathan, A. (2025). Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction. In Forty-second International Conference on Machine Learning. OLMo, T., Ettinger, A., Bertsch, A., Kuehl, B., Graham, D., Heineman, D., Groeneveld, D., Brahman, F., Timbers, F., Ivison, H., Morrison, J., Poznanski, J., Lo, K., Soldaini, L., Jordan, M., Chen, M., Noukhovitch, M., Lambert, N., Walsh, P., Dasigi, P., Berry, R., Malik, S., Shah, S., Geng, S., Arora, S., Gupta, S., Anderson, T., Xiao, T., Murray, T., Romero, T., Graf, V., Asai, A., Bhagia, A., Wettig, A., Liu, A., Rangapur, A., Anastasiades, C., Huang, C., Schwenk, D., Trivedi, H., Magnusson, I., Lochner, J., Liu, J., Miranda, L. J. V., Sap, M., Morgan, M., Schmitz, M., Guerquin, M., Wilson, M., Huff, R., Bras, R. L., Xin, R., Shao, R., Skjonsberg, S., Shen, S. Z., Li, S. S., Wilde, T., Pyatkin, V., Merrill, W., Chang, Y., Gu, Y., Zeng, Z., Sabharwal, A., Zettlemoyer, L., Koh, P. W., Farhadi, A., Smith, N. A., and Hajishirzi, H. (2025). Olmo 3. Penaloza, E., Vattikonda, D., Gontier, N., Lacoste, A., Charlin, L., and Caccia, M. (2026). Privileged information distillation for language models. arXiv preprint arXiv:2602.04942. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. (2023). Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36:53728–53741. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. (2024). Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300. Shenfeld, I., Damani, M., Hübotter, J., and Agrawal, P. (2026). Self-distillation enables continual learning. arXiv preprint arXiv:2601.19897. Tajwar, F., Singh, A., Sharma, A., Rafailov, R., Schneider, J., Xie, T., Ermon, S., Finn, C., and Kumar, A. (2024). Preference fine-tuning of llms should leverage suboptimal, on-policy data. In Proceedings of the 41st International Conference on Machine Learning. Wu, F., Xuan, W., Lu, X., Liu, M., Dong, Y., Harchaoui, Z., and Choi, Y. (2025). The invisible leash: Why rlvr may or may not escape its origin. arXiv preprint arXiv:2507.14843. Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Gao, C., Huang, C., Lv, C., et al. (2025). Qwen3 technical report. arXiv preprint arXiv:2505.09388. Yang, C., Qin, C., Si, Q., Chen, M., Gu, N., Yao, D., Lin, Z., Wang, W., Wang, J., and Duan, N. (2026). Self-distilled rlvr. arXiv preprint arXiv:2604.03128. Ye, T., Dong, L., Wu, X., Huang, S., and Wei, F. (2026). On-policy context distillation for language models. arXiv preprint arXiv:2602.12275. Yue, Y., Chen, Z., Lu, R., Zhao, A., Wang, Z., Yue, Y., Song, S., and Huang, G. (2026). Does reinforcement learning really incentivize reasoning capacity in LLMs beyond the base model? In The Thirty-ninth Annual Conference on Neural Information Processing Systems. Zhang, C., Neubig, G., and Yue, X. (2025). On the interplay of pre-training, mid-training, and rl on reasoning language models. arXiv preprint arXiv:2512.07783. Zhao, S., Xie, Z., Liu, M., Huang, J., Pang, G., Chen, F., and Grover, A. (2026). Self-distilled reasoner: On-policy self-distillation for large language models. arXiv preprint arXiv:2601.18734.
11
A
Additional Results
A.1
Concept-Graph
Similar to Fig. 3, we show performance and diversity results in Fig. 7 and Fig. 8. In Distribution eval
0.800
0.750
0.750
0.600
0.500
0.500
0.910
0.400
0.250
0.250
0.880
0.200
0.000
2
4
8
16
k
32
64
In Distribution eval
1
2
4
8
k
16
32
64
Larger 15 Graphs eval
1.000
0.000 1
2
4
8
16
k
32
64
Larger Graphs eval
1.000
1
0.750
0.750
0.800
0.500
0.500
0.910
0.700
0.250
0.250
0.880
0.600
0.000
pass@k
0.900
0.940
2
4
8
16
k
32
64
In Distribution eval
1.000
1
2
4
8
k
16
32
64
Larger 15 Graphs eval
1.000
2
4
8
16
k
32
64
Larger Graphs eval
1.000
1
0.750
0.700
0.500
0.500
0.500
0.550
0.250
0.250
0.250
0.400
0.000
0.000
pass@k
0.750
4
8
16
k
32
64
1
2
4
GRPO
8
k
16
32
64
8
k
16
32
64
2
4
8
k
16
32
64
Harder Graphs eval
1.000
0.750
2
4
0.000 1
0.850
1
2
Harder Graphs eval
1.000
0.970
1
Harder Graphs eval
1.000
0.940
1.000
Qwen3-8B
Larger Graphs eval
1.000
0.970
1
Olmo-3-7B-Instruct
Larger 15 Graphs eval
1.000
pass@k
Qwen3-4B
1.000
0.000 1
2
4
GRPO + Diversity
8
16
k
32
64
1
2
4
SDPO
8
k
16
32
64
Figure 7: Compare GRPO, GRPO + Diversity, SDSD in terms of performance (pass@k curves) for models of different sizes. We observe that SDSD usually has more flat pass@k curves for In-Distribution and Larger Graphs (15 nodes or 20 nodes) showing low functional diversity. The Harder Graphs that require learning diverse trajectories during training remains mainly unsolved by SDSD, showing again a lack of output diversity. Mean and min/max over 3 seeds are shown. In Distribution Diversity
Qwen3-4B
Concept Diversity
2.05
1.8
1.69
1.52
1.53
1.46
1.26
1.27
1.23
1
1
0
Qwen3-8B
Concept Diversity
200
400
600
Training Step
800
1000
In Distribution Diversity
0
200
400
600
Training Step
800
1000
Larger 15 Graphs Diversity
2.07
1
1.8
1.34
1.52
1.53
1.23
1.26
1.27
1.11
1
1
400
600
Training Step
800
1000
In Distribution Diversity
1.7
Concept Diversity
200
0
200
400
600
Training Step
800
1000
Larger 15 Graphs Diversity
1.62
1
1.47
1.23
1.35
1.31
1.15
1.18
1.16
1.08
1
1
200
400
600
Training Step
800
GRPO
1000
0
200
400
600
Training Step
800
1000
GRPO + Diversity
1
200
400
600
Training Step
800
1000
Larger Graphs Diversity
0
200
400
600
Training Step
800
1000
Larger Graphs Diversity
1.3
1.53
0
0
1.45
1.77
0
Larger Graphs Diversity
1.92
1.79
2.03
Olmo-3-7B-Instruct
Larger 15 Graphs Diversity
2.07
0
200
400
600
Training Step
800
1000
SDPO
Figure 8: Compare GRPO, GRPO + Diversity, SDSD in terms of semantic diversity for models of different sizes. We observe that SDSD has the lowest concept diversity (as defined in the main text) across training. This explicit measure of diversity gives additional evidence of the loss of diversity in SDSD models.
12
A.2
Science QA
Accuracy (pass@1) (%)
We show the average performance and diversity scores evolution across training steps in Fig. 9.
biology
80 75 70 65 60 55 50 45 40
60 55 50 45 40 35 30 0
100
200
Training Step
300
chemistry
material
80
physics
75 70 65 60 0
100
200
Training Step
300
0
50
100 150 200 250 300
77.5 75.0 72.5 70.0 67.5 65.0 62.5 60.0 57.5
GRPO SDSD K=1 SDSD K=3 0
Training Step
50
100
150
Training Step
200
Figure 9: Pass@1 scores of SDSD and GRPO across 5h of training time.
B
Detailed Derivations and Additional Results
B.1
Standard KL-Regularized RL
We first consider the standard KL-regularized reinforcement learning objective max Ey∼π(· | x) [R(y | x)] − βRL KL(π(· | x) ∥ π0 (· | x)) . π
Proposition 3 (Optimal policy for standard KL-regularized RL). The optimizer of (7) is 1 ∗ πRL (y | x) ∝ π0 (y | x) exp R(y | x) . βRL
(7)
(8)
Proof. Fix x and suppress it in the notation. The objective is max π
X
π(y)R(y) − βRL
y
X
π(y) log
y
π(y) π0 (y)
subject to
X
π(y) = 1.
y
Its Lagrangian is L(π, λ) =
X
π(y)R(y) − βRL
y
! X π(y) +λ π(y) log π(y) − 1 . π0 (y) y y
X
Differentiating with respect to π(y) gives π(y) + 1 + λ = 0. R(y) − βRL log π0 (y) Rearranging, log π(y) = log π0 (y) +
1 R(y) + c, βRL
where c is a constant independent of y. Exponentiating and normalizing yields 1 ∗ πRL (y) ∝ π0 (y) exp R(y) , βRL which, after reintroducing the x into the notation, proves (8). Remark 2. Equation (8) shows that KL-regularized RL preserves the base policy while exponentially tilting it by the reward. 13
B.2
SDSD-KL: Distillation from a Correct Demonstration
Consider as teacher the base policy conditioned on a fixed correct demonstration. Let y corr denote a correct reference demonstration for the same input x. We define the teacher by π0 (y | x, y corr ).
(9)
The corresponding pointwise conditional mutual information (PCMI) is i(y; y corr | x) := log
π0 (y | x, y corr ) . π0 (y | x)
(10)
This quantity measures how much conditioning on the fixed correct demonstration changes the base policy’s log-probability of the candidate sequence y. Thus, it can be interpreted as how much support y corr brings for y. But the demonstration is not fixed, it is sampled from the correct solutions. Let y corr ∼ pcorr (· | x)
(11)
be a correct demonstration drawn from a reference distribution over correct solutions. For each realized y corr , we define the teacher qycorr (y | x) := π0 (y | x, y corr ).
(12)
The SDSD-KL objective averages the distillation loss over demonstrations: h i min Eycorr KL(π(· | x) ∥ π0 (· | x, y corr )) + β KL(π(· | x) ∥ π0 (· | x)) .
(13)
π
Proposition 4 (Optimal policy for SDSD-KL). The optimizer of (13) is 1 ∗ corr | x) . πSD-KL (y | x) ∝ π0 (y | x) exp Eycorr ∼pcorr (· | x) i(y; y 1+β Proof. Fix x and suppress it in the notation. Expanding (13) gives " # X X π(y) π(y) min Eycorr π(y) log + β π(y) log corr ) π π (y | y π 0 0 (y) y y
subject to
X
(14)
π(y) = 1.
y
Interchanging expectation and summation, X X X min(1 + β) π(y) log π(y) − π(y)Eycorr [log π0 (y | y corr )] − β π(y) log π0 (y). π
y
y
y
The Lagrangian is therefore L(π, λ) = (1+β)
X
! X X X corr π(y) log π(y)− π(y)Eycorr [log π0 (y | y )]−β π(y) log π0 (y)+λ π(y) − 1 .
y
y
y
Differentiating with respect to π(y) gives (1 + β)(log π(y) + 1) − Eycorr [log π0 (y | y corr )] − β log π0 (y) + λ = 0. Hence, log π(y) =
1 β Eycorr [log π0 (y | y corr )] + log π0 (y) + c. 1+β 1+β
Using (10), log π0 (y | y corr ) = log π0 (y) + i(y; y corr ). Therefore, since taking the expectation over y corr does not influence the first term: Eycorr [log π0 (y | y corr )] = log π0 (y) + Eycorr i(y; y corr ) . 14
y
Substituting back, log π(y) = log π0 (y) +
1 Eycorr i(y; y corr ) + c. 1+β
Exponentiating and normalizing yields: π(y) ∝ π0 (y) exp
1 Eycorr i(y; y corr ) . 1+β
(15)
Restoring x to the notation gives (14). Remark 3 (Interpretation of SDSD-KL). SDSD-KL has the same formal structure as KL-regularized RL: it exponentially tilts the base policy. The effective reward is now the expected PCMI, Eycorr ∼pcorr (· | x) i(y; y corr | x) , which measures how strongly, on average over correct demonstrations, the base policy shifts toward the candidate sequence y. Remark 4 (Ratio for two correct sequences under SDSD-KL). Let y1 and y2 be two correct sequences for the same input x, and suppose π0 (y1 | x) = k π0 (y2 | x) for some k > 0. Then, by (14), their ratio under the optimal SDSD-KL policy is ∗ 1 πSD-KL (y1 | x) corr corr corr i(y ; y | x) − i(y ; y | x) = k exp E . 1 2 y ∼pcorr (· | x) ∗ πSD-KL (y2 | x) 1+β
(16)
Equivalently, ∗ πSD-KL (y1 | x) 1 exp(i(y1 ; y corr | x)) = k exp Eycorr ∼pcorr (· | x) log . ∗ πSD-KL (y2 | x) 1+β exp(i(y2 ; y corr | x)) Thus SDSD-KL preserves the base-policy ratio k only when the two sequences have the same expected PCMI. Otherwise, SDSD-KL further reweights them according to the gap in how strongly correct demonstrations support y1 versus y2 . B.3
Optimal Policy for Token-Level SDSD-KL
Previously, we discussed the sequence level objective and its implications. This was done for ease of presentation, but in practice, we use a token-level objective. We will see that exactly the same derivations carry in the token-level case, and we will discuss the implications. Consider the optimization problem at a single generation step t. The policy generates the next token yt from the vocabulary, conditioned on the input x and the generated prefix y<t . The token-level SDSD-KL objective averages the distillation loss over demonstrations at the next-token distribution: min Eycorr ∼pcorr (·|x) [KL(π(·|x, y<t )||π0 (·|x, y<t , y corr ))] + βKL(π(·|x, y<t )||π0 (·|x, y<t )) (17) π
Proposition 5 (Optimal policy for token-level SDSD-KL). The optimizer of the token-level objective is: 1 corr ∗ πSD−KL (yt |x, y<t ) ∝ π0 (yt |x, y<t ) exp Eycorr ∼pcorr (·|x) [i(yt ; y |x, y<t )] (18) 1+β where the token-level pointwise conditional mutual information (PCMI) is defined as: i(yt ; y corr |x, y<t ) := log The proof follows exactly the same steps as before. 15
π0 (yt |x, y<t , y corr ) π0 (yt |x, y<t )
(19)
Implications of the token-level objective. In the case of sequence-level objective, we have seen that there is a bias for common rollouts. That bias comes from the alignment between rollouts and demonstrations and the preference of the teacher. A similar phenomenon is happening at the token-level. Some tokens are more aligned to the context determined by the demonstration and the previous tokens of the student rollout. Again, next-tokens that will move the current rollout into a novel or uncommon direction might be less aligned to the context. Moreover, given the same context, some next-tokens are preferred over others by the teacher. On top of this, differently from the sequence level objective, here the teacher is myopic to the full student rollout. It guides the next token without taking into account the relation to future tokens. Thus, some commonalities between student rollout and demonstrations might only be seen at a sequence level and could be harder to establish at intermediate token positions. Thus, the alignment at intermediate token positions might be low, causing learning to be even more biased. This leads to a bias for common next-tokens, which turns into common entire distributions as we discussed before. B.4
Concept Graph Task: Additional details
In the Concept Graph task, each query is a problem of finding a path given to the LLM. Each query has a different graph with randomly sampled structure and node names. Here is an example: You are given the following graph structure . Nodes : 0 start [ hub ] 1 triangle [ shape ] 2 pigeon [ concept = birds ] 3 parrot [ concept = birds ] 4 sparrow [ concept = birds ] 5 crow [ concept = birds ] 6 heron [ concept = birds ] 7 eagle [ concept = birds ] 8 square [ shape ] 9 tuna [ concept = fish ] ... Edges : 0 ( start ) -- 2 ( pigeon ) 0 ( start ) -- 5 ( crow ) 0 ( start ) -- 9 ( tuna ) 0 ( start ) -- 21 ( plum ) ... 2 ( pigeon ) -- 3 ( parrot ) 3 ( parrot ) -- 4 ( sparrow ) 1 ( triangle ) -- 4 ( sparrow ) ... 21 ( plum ) -- 22 ( mango ) 22 ( mango ) -- 27 ( triangle ) ... Your task is to generate a path from the start node to the target node named triangle . For the output , print the names of the nodes in the path , use the following format : A path from start to triangle is \ boxed { start , node_1_name , ... , triangle } for example \ boxed { start , Jacobi , Hamilton , ... , star }.
A correct response following the birds concept chain would be: A path from start to triangle is \boxed{start, pigeon, parrot, sparrow, triangle}
An equally valid alternative following a different concept (fruits) would be: \boxed{start, plum, mango, triangle}
Both receive the same maximum reward, but a model that only produces bird-paths across different graphs exhibits lower solution diversity than one that explores both bird and fruit paths. 16