Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models

Ayan Sengupta; Vaibhav Seth; Tanmoy Chakraborty

cs.CLSep 29, 2026

Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models

Authors: Ayan Sengupta, Vaibhav Seth, Tanmoy Chakraborty

Organizations: Department of Electrical Engineering, Indian Institute of Technology Delhi · Yardi School of Artificial Intelligence, Indian Institute of Technology Delhi

Abstract

Knowledge Distillation (KD) trains a smaller-capacity student model to imitate a larger-capacity teacher model by matching output distributions, implicitly assuming the teacher to be a reliable oracle. In large language models (LLMs), this assumption often fails: teacher predictions can exhibit high entropy and hallucinations, causing standard KD to degrade well-calibrated student priors. We propose CaRE-KD, a confidence-gated distillation framework that replaces static objectives with uncertainty-adaptive optimization. CaRE-KD has two components: a token-level loss (CaRE-Divergence) that adaptively switches between Forward and Reverse KL divergence based on teacher--student confidence, and a batch-level epistemic rejection mechanism (Revival) that suppresses updates when the teacher is more uncertain than the student. We provide a gradient-level analysis showing how this dual-granularity design induces a conditional calibration mechanism that prior static divergences cannot reproduce. Empirically, across eight teacher--student pairs and eleven benchmarks spanning instruction following, chat alignment, code generation, and mathematical reasoning, CaRE-KD delivers consistent gains over strong baselines (Skewed-KL, $α$ -- $β$ divergence). Highlights include up to $+3.2$ average ROUGE-L on instruction-following tasks, $+2.1$ pass@1 on MBPP, $+1.7$ accuracy on GSM8k, and $+1.8$ accuracy on CollegeMath over the strongest baseline, with consistent gains in LLM-as-a-judge factuality (up to $+2.5$ per task over Skewed-RKL). Revival further acts as a principled, loss-agnostic plug-in that systematically strengthens existing distillation objectives by filtering epistemically unreliable teacher supervision.

Figures & tables

Panel of Figure 1 — Figure 1: OpenLLaMA-3B student distilled from a 7B teacher on Dolly-15k with Skewed-RKL [ 6 ] . (a) Distillation fails to reduce student epistemic uncertainty when the teacher is uncertain (KS distance changes by only $0.4\%$ ). (b) Teacher supervision helps when the teacher is confident (median $+4\%$ ) but barely helps when the teacher is uncertain ( $+0.5\%$ ). (c) Static objectives enforce agreement even with unreliable teachers.

Figure 2: Token-level geometric adaptation induced by CaRE-Divergence . When the teacher is confident, the objective behaves as Forward KL (mean-seeking, recall-oriented); when the teacher is uncertain, it transitions to Reverse KL (mode-seeking, precision-oriented), suppressing tail noise.

Teacher $\rightarrow$ Student	Loss	Benchmarks				Avg
Teacher $\rightarrow$ Student	Loss	Dolly	Self-Inst	Super-Nat	Vicuna	Avg
	FKL	23.4/21.2	10.7/15.6	19.0/16.4	15.0/13.5	17.0/16.7
	RKL	24.7/20.7	11.6/16.4	19.2/16.2	16.6/14.6	18.0/17.0
	SRKL	24.9/23.1	12.1/17.6	23.6 /21.9	16.5/15.3	19.3/19.5
	$\alpha$ - $\beta$ Div	23.9/23.1	11.5/16.4	21.9/18.3	16.1/15.0	18.4/18.2
GPT2-XL (1.5B) $\rightarrow$ GPT2-base (0.1B)	CaRE-Div (Ours)	25.6/24.1	12.9/20.1	22.8/ 24.4	17.3/15.3	19.6/21.0

Table 1: Main instruction-following results under SGO training. Cells report ROUGE-L / LLM-as-a-judge factuality on Dolly, Self-Instruct, Super-Natural Instructions, and Vicuna. Across five teacher–student pairs, CaRE-Divergence achieves the best average ROUGE-L on three pairs and the best factuality on three pairs, with statistically significant gains over the second-best baseline (

t=3.93

p=0.03

). Bold denotes the best score per cell. Results are averaged over five seeds; detailed comparisons appear in Appendix 11 .

Student	Loss	Benchmarks				Avg
Student	Loss	Dolly	Self-Inst	Super-Nat	Vicuna	Avg
	FKL	0.22	0.14	-0.64	0.09	-0.05
	RKL	0.32	0.61	3.51	0.12	1.14
	SRKL	0.41	0.54	1.64	0.30	0.72
	$\alpha$ - $\beta$ Div	1.11	0.80	0.94	0.67	0.88
GPT2-base	CaRE	0.35	0.94	1.26	1.04	0.90

Table 2: Impact of Revival on different distillation losses. Values are the per-task delta

\Delta=\text{Score}_{\text{with-{Revival}}}-\text{Score}_{\text{no-{Revival}}}

in ROUGE-L. CaRE-Divergence shows the most consistent positive response, with a statistically significant average improvement across model pairs (one-sample

t

-test:

t=3.98

p<0.001

). Static baselines such as FKL (

t=-2.55

p=0.06

) often degrade with Revival , while RKL (

t=1.24

p=0.25

) yields insignificant gains, indicating that our confidence-gated geometry is uniquely synergistic with epistemic filtering. Bold marks the largest positive delta per (model, task) cell.

Method	Qwen2-7B-Inst $\rightarrow$ Qwen2-1.5B			DS-Coder-6.9B-Inst $\rightarrow$ DS-Coder-1.3B		Qwen2.5-Math-7B-Inst $\rightarrow$ Qwen2.5-Math-1.5B
	Alpaca	Evol	Ultra	HEval	MBPP	GSM8k	College
	Rouge-L/LLM	Rouge-L/LLM	Rouge-L/LLM	pass@1	pass@1	pass@1	pass@1
Student	8.81/27.3	14.9/31.2	10.8/30.5	32.3	58.5	69.9	37.1
Distillm2	16.3/43.3	26.5/52.1	22.1/ 49.3	43.3	60.3	72.2	44.2
ABKD	16.4 / 44.8	26.7/51.9	22.3/48.4	42.7	59.8	71.0	44.3
GKD	–	–	–	43.1	60.2	71.4	43.9

Table 3: Generalization to specialized domains. Cells report ROUGE-L / LLM-as-a-judge factuality. We evaluate three model families: chat alignment (Qwen2), code generation (DeepSeek-Coder), and math reasoning (Qwen2.5-Math). CaRE-Divergence consistently improves precision-sensitive tasks, achieving +2.1 pp on MBPP, +1.7 pp on GSM8k, and +1.8 pp on CollegeMath over the strongest baseline. Gains are statistically significant (

t=3.05

p=0.018

). Bold denotes the best score per cell. GKD and Distillm are additionally reported on the code and math tasks, where CaRE-KD remains best on MBPP, GSM8K and CollegeMath and tied-best on HumanEval; chat-alignment entries for these two baselines are still running and omitted (“–”).

Figure 3: Ablation study of CaRE-KD with GPT2-base. (a) gating strategy (

g_{\mathrm{hard}}

vs.

g_{\mathrm{soft}}

); (b) MC-dropout sample count

N

; (c) sampling temperature

T

; (d) target skip percentage

\tau

; (e) Revival scheduling strategy (constant, increasing, decreasing); and (f) margin

m

in soft gating.

Figure 4: BALD uncertainty distributions across four instruction-following benchmarks. Teacher models show heavy high-uncertainty tails (red), whereas CaRE-KD students (blue) suppress these regions and maintain concentrated low-BALD profiles. Reported KS distances indicate strong distributional separation, exceeding

0.5

on Vicuna and Dolly.

Panel of Figure 5 — Figure 5: Runtime on Dolly-15k with OpenLLaMA-3B for one epoch. (a) Wall-clock time scales near-linearly with the number of MC dropout samples $N$ for BALD estimation. (b) Higher skip rates $\tau$ reduce BALD overhead by skipping backward passes on rejected batches.

Teacher	$N$	Mean BALD	BALD range	Non-zero BALD	Mean $\Delta$ top-1 prob
GPT2-XL	3	0.049	0.013–0.107	500/500 (100%)	0.074
OPT-2.7B	3	0.152	0.057–0.239	500/500 (100%)	0.079

Table 4: Verification that MC-dropout BALD is non-degenerate on the two teachers (500 Dolly sequences,

N=3

, dropout forced to

p=0.1

Teacher $\rightarrow$ Student	Loss	Benchmarks				Avg
Teacher $\rightarrow$ Student	Loss	Dolly	Self-Inst	Super-Nat	Vicuna	Avg
	FKL	23.0/20.4	9.7/14.2	16.8/14.9	14.9/13.0	16.1/15.6
	RKL	23.0/20.1	10.1/14.9	17.7/15.3	15.5/13.7	16.6/16.0
	SRKL	23.5/21.7	11.1/16.2	21.5/18.8	14.8/14.2	17.7/17.7
	$\alpha$ - $\beta$ Div	23.0/21.6	10.3/15.3	19.7/17.0	14.3/13.8	16.8/16.9
GPT2-XL (1.5B) $\rightarrow$ GPT2-base (0.1B)	CaRE-Div	25.2/23.5	12.5/18.6	23.7/21.8	16.3/14.8	19.5/19.7

Table 5: Comparison of divergence measures without student-generated outputs (non-SGO) [ 6 ] . Each cell reports ROUGE-L / LLM-as-a-judge factuality scores. Even without exploratory rollouts, CaRE-Divergence is competitive with the strongest baseline on every model pair, achieving the highest average ROUGE-L on three of five teacher–student pairs and matching SRKL on the remaining two within standard error. Improvements over FKL/RKL are statistically significant in a paired

t

-test (

t=12.84

p<0.005

). Bold marks the best score per (model, task, metric) cell.

Method	Dolly	Self-Inst	Vicuna	Super-Nat	Avg.
SFT	23.33 (0.22)	10.56 (0.54)	15.12 (0.47)	17.08 (0.29)	16.52
KD [ 13 ]	23.52 (0.19)	11.23 (0.41)	15.92 (0.37)	20.68 (0.14)	17.84
SeqKD [ 1 ]	23.38 (0.37)	10.18 (0.18)	15.01 (0.28)	15.08 (0.12)	15.91
ImitKD [ 62 ]	21.63 (0.51)	10.85 (0.38)	14.70 (0.29)	17.94 (0.10)	16.28
MiniLLM [ 5 ]	23.84 (0.26)	12.44 (0.28)	18.29 (0.36)	22.62 (0.26)	19.30
GKD [ 2 ]	23.75 (0.15)	12.73 (0.24)	16.64 (0.24)	23.05 (0.23)	19.04

Table 6: Baseline comparison for GPT2-XL (1.5B)

\rightarrow

GPT2-base (0.1B). ROUGE-L on four held-out instruction benchmarks; values in parentheses are standard deviations across 5 seeds. CaRE-KD achieves the best average ROUGE-L. Bold marks per-task best.

Method	Dolly	Self-Inst	Vicuna	Super-Nat	Avg.
SFT	21.78 (0.19)	8.09 (0.39)	14.40 (0.17)	13.45 (0.20)	14.43
KD [ 13 ]	20.54 (0.38)	9.16 (0.29)	14.65 (0.47)	15.79 (0.26)	15.04
SeqKD [ 1 ]	20.72 (0.59)	8.94 (0.37)	13.56 (0.33)	16.80 (0.36)	15.01
ImitKD [ 62 ]	20.16 (0.19)	8.95 (0.50)	15.05 (0.52)	14.72 (0.21)	14.72
MiniLLM [ 5 ]	22.24 (0.32)	9.92 (0.47)	16.97 (0.49)	16.58 (0.19)	16.43
GKD [ 2 ]	22.46 (0.27)	10.59 (0.36)	16.25 (0.65)	19.33 (0.20)	17.16

Table 7: Baseline comparison for OPT-2.7B

\rightarrow

OPT-125M. ROUGE-L on four held-out instruction benchmarks; values in parentheses are standard deviations across 5 seeds. CaRE-KD achieves the best average ROUGE-L. Bold marks per-task best.

Method	Dolly	Self-Inst	Vicuna	Super-Nat	Avg.
SFT	25.11 (0.58)	16.52 (0.56)	16.33 (0.39)	29.28 (0.45)	21.81
KD [ 13 ]	20.95 (0.49)	16.12 (0.80)	15.39 (0.39)	27.93 (0.26)	20.10
SeqKD [ 1 ]	24.67 (0.47)	15.83 (0.62)	17.06 (0.47)	29.06 (0.28)	21.66
ImitKD [ 62 ]	24.53 (0.26)	17.80 (0.80)	17.36 (0.22)	31.50 (0.12)	22.80
MiniLLM [ 5 ]	27.88 (0.28)	19.94 (0.58)	20.50 (0.41)	36.91 (0.28)	26.31
GKD [ 2 ]	26.30 (0.31)	19.56 (1.03)	18.66 (0.34)	35.71 (0.20)	25.06

Table 8: Baseline comparison for OpenLLaMA2-7B

\rightarrow

OpenLLaMA2-3B. ROUGE-L on four held-out instruction benchmarks; values in parentheses are standard deviations across 5 seeds. CaRE-KD is competitive with the strongest baseline (Distillm) on average and achieves the best score on Self-Instruct and Super-Natural Instructions. Bold marks per-task best.

Panel of Figure 6 — Figure 6: Cross-entropy loss on the Dolly-15k development set.

Panel of Figure 7 — Figure 7: Exact match (pass@1) accuracy on the Dolly-15k development set.

Figure 8: Ablation results with OPT-125M student.

Panel of Figure 9 — Figure 9: Epistemic uncertainty (BALD) distributions of the OpenLLaMA-3B student under Skewed RKL, with and without Revival .

Panel of Figure 10 — Figure 10: Epistemic uncertainty (BALD) distributions of the OpenLLaMA-3B student under $\alpha$ – $\beta$ divergence, with and without Revival .

Panel of Figure 11 — Figure 11: Proportion of batches triggering $M_{\mathrm{Revival}}$ over training iterations for different skip targets and scheduling strategies.

Panel of Figure 12 — Figure 12: Proportion of $M_{\mathrm{Revival}}$ activations for different numbers of MC dropout samples.

Panel of Figure 13 — Figure 13: Runtime analysis for different BALD sample configurations for specialized tasks.

Objective	Avg. confidence	ECE
Forward KL	0.67	0.22
CaRE-KD (softest configuration)	0.68	0.28
$\alpha$ – $\beta$ divergence	0.75	0.28
Reverse KL	0.77	0.31
Skewed-RKL	0.79	0.31
CaRE-KD (default temperature)	0.79	0.34

Table 9: ECE on Dolly for the GPT2 students, ordered by average confidence. With a single-reference target, ECE tracks confidence almost one-for-one and is largely a proxy for sharpness.

Measure (LLM-judged)	Region where Revival fires	Opposite region	Mann–Whitney $p$
Question hardness	2.87	2.98	0.37
Response diversity	3.96	3.99	0.63

Table 10: Selection-bias check on OpenLLaMA-Dolly. LLM-judged question hardness and response diversity in the region where Revival fires vs. the complementary region. Neither difference is significant.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

Apr 21, 2026

Distillation Traps and Guards: A Calibration Knob for LLM Distillability

Knowledge distillation (KD) transfers capabilities from large language models (LLMs) to smaller students, yet it can fail unpredictably and also underpins model leakage risks. Our analysis revealed several distillation traps: tail noise, off-policy instability, and, most fundamentally, the teacher-student gap, that distort training signals. These traps manifest as overconfident hallucinations, self-correction collapse, and local decoding degradation, causing distillation to fail. Motivated by these findings, we propose a post-hoc calibration method that, to the best of our knowledge, for the first time enables control over a teacher's distillability via reinforcement fine-tuning (RFT). Our objective combines task utility, KL anchor, and across-tokenizer calibration reward. This makes distillability a practical safety lever for foundation models, connecting robust teacher-student transfer with deployment-aware model protection. Experiments across math, knowledge QA, and instruction-following tasks show that students distilled from distillable calibrated teachers outperform SFT and KD baselines, while undistillable calibrated teachers retain their task performance but cause distilled students to collapse, offering a practical knob for both better KD and model IP protection.

Weixiao Zhan, Yongcheng Jing, Leszek Rutkowski +1

Generative AI Lab, College of Computing and Data Science Nanyang Technological University, Singapore 639798 · Systems Research Institute of the Polish Academy of Sciences, AGH University of Krakow, 30-059 Kraków, and the SAN University, 90-113, Łód´z, Poland

Language Model Calibration Language Model Distillation Teacher-Student Learning Knowledge Distillation

Ask PDF

Oct 6, 2026

WASD: Wasserstein-based Knowledge Distillation for Large Language Models

Autoregressive large language models (LLMs) have rapidly advanced in capability, but their increasing scale comes with substantial computational and memory costs at inference time. Knowledge distillation (KD) offers a practical solution by transferring knowledge from a large teacher model to a smaller student model via alignment of discrete probability distributions. However, existing KD methods for LLMs primarily rely on divergences that evaluate discrepancies through probability values at each vocabulary index, without explicitly leveraging token-level semantic information. We propose Wasserstein-based knowledge distillation (WASD) for LLMs, which incorporates token-level semantic information via the Wasserstein-based distance with a cost matrix derived from token embeddings. To ensure computational tractability, we adopt the Sinkhorn divergence and derive a gradient-equivalent objective that can be efficiently optimized without introducing additional networks. Experiments across multiple LLM families and scales show that WASD consistently improves distillation performance on diverse tasks, including instruction following, mathematical reasoning, and code generation. Our results highlight the importance of semantic information encoded in the token space for effective distribution alignment in LLM distillation. The implementation is publicly available at https://github.com/aailab-kaist/WASD .

Byeonghu Na, Donghyeok Shin, Yeongmin Kim +2

KAIST

Language Model Distillation Knowledge Distillation

Ask PDF

Feb 26, 2026

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

Zhaoyang Zhang, Shuli Jiang, Yantao Shen +6

AWS Agentic AI

Knowledge Distillation RL for Knowledge Distillation On-Policy Distillation Selective Knowledge Distillation RL for Language Model Reasoning

Ask PDF

Abstract

Language Model Distillation

Explore similar work

Distillation Traps and Guards: A Calibration Knob for LLM Distillability

WASD: Wasserstein-based Knowledge Distillation for Large Language Models

Reinforcement-aware Knowledge Distillation for LLM Reasoning