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

Distillation Traps and Guards: A Calibration Knob for LLM Distillability
Apr 21, 2026Weixiao Zhan, Yongcheng Jing, Leszek Rutkowski +1Language Model Calibration Language Model Distillation
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
WASD: Wasserstein-based Knowledge Distillation for Large Language Models
Oct 6, 2026Byeonghu Na, Donghyeok Shin, Yeongmin Kim +2Language Model Distillation Knowledge Distillation
KAIST
Reinforcement-aware Knowledge Distillation for LLM Reasoning
Feb 26, 2026Zhaoyang Zhang, Shuli Jiang, Yantao Shen +6Knowledge Distillation RL for Knowledge Distillation
AWS Agentic AI

Apr 21, 2026J/K move · Enter open · S save

Distillation Traps and Guards: A Calibration Knob for LLM Distillability

Weixiao Zhan, Yongcheng Jing, Leszek Rutkowski, Dacheng Tao

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

Open paper AskPDF arXiv

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.

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