Know Thyself, Teach Thyself: Internal Information Flow for Selective Self-Distillation

Rui Wang; Ruijie Wang; Bo Chen; Jiangxuan Long; Yingyu Liang

cs.LGSep 29, 2026

Know Thyself, Teach Thyself: Internal Information Flow for Selective Self-Distillation

Authors: Rui Wang, Ruijie Wang, Bo Chen, Jiangxuan Long, Yingyu Liang

Organizations: Peking University · University of Oxford · The University of Hong Kong

Abstract

Self-distillation turns knowledge distillation into a closed learning loop and offers a path toward recursive self-improvement. Without an external teacher, however, the model must determine both what information can improve its supervision and which induced changes should be learned. Existing methods typically improve teacher-generated data or select training examples in isolation, leaving the information transferred between these stages unmeasured. We introduce InFlow, a retrieval-guided on-policy self-distillation framework that models this process as potential-to-realized information flow. InFlow first retrieves potentially informative sources using certainty-calibrated hidden-state trajectories, then measures their realized effect through the Jensen--Shannon divergence between the teacher's initial and retrieval-conditioned answer beliefs. Examples with larger belief shifts are selected for on-policy distillation. Our analysis formalizes the information optimized by retrieval and selection and relates the answer-level shift to the teacher--student distillation gap. Across four open-weight language models and three knowledge domains, InFlow achieves the strongest cross-model average among the compared selection methods, with ablations supporting both stages of the framework. Our code is available at https://github.com/1240148048/INFLOW.

Figures & tables

Figure 1: Overview of InFlow . For each query, Stage I uses internal signals to estimate how much information retrieval can channel from external evidence into the teacher. In Stage II, the teacher reflects through self-comparison and selects the realized information to provide to the student. Stage III distills this new information into the student, where it is expected to shift the original semantic belief without retrieval context.

Model	Domain	Original	Full	Redundancy	Consistency	NEURON	Mutual	InFlow
Qwen3-4B	Math	$45.3_{\pm 0.47}$	$49.0_{\pm 0.82}$	$49.7_{\pm 0.94}$	$48.0_{\pm 1.41}$	$46.7_{\pm 0.47}$	$49.3_{\pm 1.70}$	$\bm{51.3_{\pm 0.94}}$
	Natural Sci.	$78.5_{\pm 0.82}$	$78.8_{\pm 0.24}$	$78.8_{\pm 0.62}$	$78.7_{\pm 0.24}$	$\bm{79.8_{\pm 0.85}}$	$79.5_{\pm 1.08}$	$79.2_{\pm 0.47}$
	Hum. & Soc.	$45.0_{\pm 0.71}$	$43.8_{\pm 0.85}$	$45.7_{\pm 1.03}$	$45.2_{\pm 1.03}$	$45.5_{\pm 0.71}$	$45.7_{\pm 0.62}$	$\bm{46.2_{\pm 0.47}}$
	Overall	$58.5_{\pm 0.09}$	$58.9_{\pm 0.34}$	$59.7_{\pm 0.41}$	$59.1_{\pm 0.52}$	$59.5_{\pm 0.66}$	$59.9_{\pm 0.90}$	$\bm{60.4_{\pm 0.33}}$
Qwen3.5-9B	Math	$49.3_{\pm 0.47}$	$51.3_{\pm 2.49}$	$51.3_{\pm 1.70}$	$49.0_{\pm 1.41}$	$51.0_{\pm 1.41}$	$51.3_{\pm 0.47}$	$\bm{53.3_{\pm 2.05}}$
	Natural Sci.	$80.2_{\pm 1.25}$	$80.8_{\pm 1.25}$	$80.7_{\pm 1.25}$	$79.7_{\pm 0.47}$	$\bm{81.7_{\pm 0.62}}$	$80.8_{\pm 0.24}$	$\bm{81.7_{\pm 0.62}}$

Table 1: Average rollout accuracy at 250 optimizer updates and a 20% selection ratio. Each entry is Avg@3

\pm

SD (%). The final row reports mean runtime per run. Bold and underlined entries denote the best and second-best results, respectively. Maj@3 results are reported in Appendix C.1 .

Panel of Figure 2 — Figure 2: Accuracy over optimizer updates. Shaded bands denote the standard deviation across three rollout-index accuracies.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Domain	Original	Full	Redundancy	Consistency	NEURON	Mutual	InFlow
Qwen3-4B	Math	47.0	50.0	51.0	50.0	46.0	50.0	52.0
	Natural Sci.	78.5	79.5	78.5	78.5	80.0	79.5	79.0
	Hum. & Soc.	44.5	44.0	45.0	46.0	45.0	45.5	46.0
	Overall	58.6	59.4	59.6	59.8	59.2	60.0	60.4
Qwen3.5-9B	Math	53.0	53.0	54.0	49.0	54.0	50.0	54.0
	Natural Sci.	80.5	80.0	81.0	82.0	82.5	81.0	82.5

Appendix

Table 2: 3-rollout Majority-vote accuracy (Maj@3, %) at 250 optimizer updates and a 20% selection ratio. Bold and underlined values are the best and second best in each row. In the event of a tie, the voting outcome is determined by the softmax of the average answer token logit.

JS rank	Fix	Break	W $\rightarrow$ W	R $\rightarrow$ R	Fix $-$ Break	Initial Acc.	Teacher Acc.	Mean JS
0–10%	36.0	30.0	34.0	0.0	6.0	30.0	36.0	0.6931
10–20%	36.0	16.0	48.0	0.0	20.0	16.0	36.0	0.5898
20–30%	12.0	20.0	50.0	18.0	$-8.0$	38.0	30.0	0.3242
30–40%	0.0	0.0	54.0	46.0	0.0	46.0	46.0	0.0182
40–50%	0.0	0.0	26.0	74.0	0.0	74.0	74.0	0.0000
50–60%	0.0	0.0	22.0	78.0	0.0	78.0	78.0	0.0000

Appendix

Table 3: Direction of retrieval-induced belief shifts on Qwen3-4B. Targets are ranked by JS belief shift from largest to smallest. Transition frequencies and accuracies are percentages within each interval; Fix

-

Break is measured in percentage points.

JS rank	Fix	Break	W $\rightarrow$ W	R $\rightarrow$ R	Fix $-$ Break	Initial Acc.	Teacher Acc.	Mean JS
0–10%	26.0	34.0	40.0	0.0	$-8.0$	34.0	26.0	0.6931
10–20%	26.0	20.0	54.0	0.0	6.0	20.0	26.0	0.6414
20–30%	26.0	16.0	56.0	2.0	10.0	18.0	28.0	0.4223
30–40%	4.0	4.0	68.0	24.0	0.0	28.0	28.0	0.1815
40–50%	0.0	0.0	62.0	38.0	0.0	38.0	38.0	0.0461
50–60%	0.0	0.0	38.0	62.0	0.0	62.0	62.0	0.0000

Appendix

Table 4: Direction of retrieval-induced belief shifts on Llama-3.1-8B, using the same protocol and notation as Table 3 .

Figure 5: Qwen3-4B-Instruct ablation over the number of retrieved demonstrations per target. Other generation and optimization budgets are fixed.

Figure 6: Knowledge-retrieval graphs on Qwen3-4B-Instruct. Each panel uses the same target queries, candidate pool, and retrieval depth $k=3$ . Background colors indicate mathematics, natural sciences, and humanities and social sciences. Outlined circles are target queries, squares are retrieved contexts, and connecting edges show the resulting retrieval relations. PC1 and PC2 explain $11.1\%$ and $10.2\%$ of the representation variance, respectively.

Method	Retrieval rule	Teacher-data selection rule
Original	None.	No adaptation.
Full	Retrieve the $k$ sources with the largest cosine similarity between final-layer answer representations.	Retain all generated teacher data.
Redundancy	Use the same final-layer cosine retrieval as Full .	Iteratively remove the example most similar to another retained teacher representation until $\alpha N$ examples remain.
Consistency	No retrieved demonstrations; instead, we obtain more confident training data through multiple sampling, consistency checks, and self-feedback.	Retain the $\alpha N$ targets with the largest modal-answer agreement among $M$ generations.
NEURON	Retrieve the $k$ sources with the largest Jaccard overlap between sets of positively contributing neurons.	Retain the $\alpha N$ targets with the smallest mean number of activated contributing neurons across teacher generations.
Mutual	Retrieve the $k$ sources that provide more conditional mutual information for the target data using a greedy algorithm.	Retain the $\alpha N$ targets with the largest aggregated retrieval-information score.

Appendix

Table 5: Baseline instantiations in the unified comparison framework. Here

k

is the retrieval depth,

M

is the number of teacher generations, and

\alpha

is the selection ratio.

Explore similar work

Oct 5, 2026cs.LG

Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models

Large language models have shown strong reasoning capabilities, but their high inference costs make knowledge distillation an important approach for transferring such capabilities to compact models in resource-constrained scenarios. On-policy self-distillation further reduces the reliance on external large teacher models while improving the reasoning ability of compact language models. However, existing methods typically either distill all token positions uniformly or select tokens using fixed heuristic criteria, assigning the same distillation strength to the selected positions rather than adaptively learning which tokens are most beneficial for distillation. To address these limitations, we propose BiToK-SD (Bilevel Top-K Token Selection for Self-Distillation), a bilevel-optimization-based token selection method that learns where distillation should be applied during on-policy self-distillation. Specifically, BiToK-SD is formulated as a bilevel optimization problem, where the lower-level problem models Top-K token selection as a differentiable threshold-based relaxation, allowing the selected positions to adapt as the student policy evolves, while the upper-level problem performs knowledge distillation on the selected positions. Experiments on mathematical reasoning benchmarks show that BiToK-SD achieves the best average performance among all compared methods while requiring only lightweight additional computation.

Heng Liang, Xinwen Zhang, Hongchang Gao

Department of Computer and Information Sciences Temple University

Topic pending

Ask PDF

Sep 30, 2026cs.AI

Disentangling Self-Distillation: Measuring and Modeling Acquisition and Retention

Self-distillation with privileged context adapts a language model from demonstrations by letting the model, once conditioned on a reference response, teach its context-free copy token by token. Our taxonomy reveals existing methods differ along three entangled axes: (i) the rollout source (student or teacher), (ii) the teacher coupling (frozen, or an exponential moving average of the student at some coupling rate) and (iii) the KL direction (reverse or forward), yet these axes are usually studied in fixed combinations and have led to conflicting conclusions. We formalize a unifying framework to encompass all self-distillation methods vs classic supervised fine-tuning: we train every combination of the three axes, on Qwen2.5-7B and Ministral-3-3B across ordinary and contradictory tasks, totaling 1,200 adaptation runs, to systematically investigate the impact of the above axes. We propose a controlled model of the same objective to explain the resulting acquisition-retention trade-offs. We find that (i) the rollout source matters mostly where the task contradicts the pretrained behavior: there teacher rollouts raise acquisition well above what student rollouts achieve, with almost no change in retention; (ii) the teacher coupling changes acquisition most, on every task: acquisition rises with the coupling rate, then falls past a task-specific rate; (iii) switching the KL direction costs retention in one model but not the other so which axis to tune first depends on the model. The controlled model reproduces the three trends.

Luis Zuin, Alexis Huet, Dario Rossi +1

Huawei Technologies Co., Ltd., Paris, France

Retention Acquisition Measures

Ask PDF

Jul 2, 2026cs.LG

DemoPSD: Disagreement-Modulated Policy Self-Distillation

On-policy self-distillation (OPSD) has emerged as a practical method for training large language models (LLMs) to reason, where a single model acts as both the teacher and the student with different levels of information access. However, recent studies have found that the teacher's dense token-level supervision, conditioned on privileged information, can lead to overfitting to in-domain patterns, suppress exploration, and hurt cross-domain generalization, while also introducing a more fundamental issue: privileged information leakage, where the student encodes answer-dependent shortcuts that are unavailable at test time. We introduce DemoPSD, a novel framework that resolves such problems through the idea of selective adoption of teacher guidance. Instead of fitting the full teacher distribution, DemoPSD steers the student toward a reverse-KL barycenter target, a weighted geometric combination of the teacher and student distributions, that naturally balances learning from the teacher with preserving the student's own reasoning capacity. We measure the difference between their distributions and use such a discrepancy to adaptively control the blending at each token position. We provably show that DemoPSD achieves (1) leakage attenuation, i.e., effective mitigation of privileged information leakage; and (2) exploration preservation, i.e., preservation of exploration capacity under dense token-level distillation. Extensive experiments on SciKnowEval across four scientific fields show that DemoPSD outperforms both GRPO and SDPO while maintaining higher training entropy and robustly generalizing to out-of-distribution GPQA benchmarks.

Yunhe Li, Hao Shi, Wenhao Liu +5

1City University of Hong Kong · 2Tsinghua University · 3Shenzhen University of Advanced Technology +1

Unsupervised On-Policy Self-Distillation Teacher-Student Distillation Data Leakage

Ask PDF