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

Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models
Oct 5, 2026Heng Liang, Xinwen Zhang, Hongchang Gao
Department of Computer and Information Sciences Temple University
Disentangling Self-Distillation: Measuring and Modeling Acquisition and Retention
Sep 30, 2026Luis Zuin, Alexis Huet, Dario Rossi +1Retention Acquisition
Huawei Technologies Co., Ltd., Paris, France
DemoPSD: Disagreement-Modulated Policy Self-Distillation
Jul 2, 2026Yunhe Li, Hao Shi, Wenhao Liu +5Unsupervised On-Policy Self-Distillation Teacher-Student Distillation
1City University of Hong Kong · 2Tsinghua University · 3Shenzhen University of Advanced Technology +1

Oct 5, 2026, cs.LGJ/K move · Enter open · S save

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

Heng Liang, Xinwen Zhang, Hongchang Gao

Department of Computer and Information Sciences Temple University

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Disentangling Self-Distillation: Measuring and Modeling Acquisition and Retention

DemoPSD: Disagreement-Modulated Policy Self-Distillation