Do Student LLMs Inherit OOD Robustness? Invariance-Weighted Distillation for Reliable Knowledge Transfer

Dileesha Kannangara; Sanghamitra Dutta

cs.CLSep 18, 2026

Do Student LLMs Inherit OOD Robustness? Invariance-Weighted Distillation for Reliable Knowledge Transfer

Authors: Dileesha Kannangara, Sanghamitra Dutta

Organizations: University of Maryland, College Park

Abstract

Knowledge distillation (KD) aims to compress high-performance teacher LLMs into lightweight students. However, distilled students often exhibit substantial performance degradation in out-of-distribution (OOD) settings, a critical gap that remains underexplored. We identify two compounding mechanisms causing OOD performance degradation: (1) data spuriousness: students can learn spurious correlations in the distillation dataset over genuine causal relationships; and (2) teacher capability: standard KD treats all samples uniformly, ignoring whether the teacher is guided by causal features or misled by spurious shortcuts on a given sample. To address these challenges, we propose Invariance-Weighted Distillation (IWD), a theoretically grounded framework that dynamically reweights training samples using an estimate of the teacher's causal reliance derived from prediction invariance across multiple synthetic environments. IWD perturbs spurious cues while preserving core semantics, assigning higher distillation weights to samples whose teacher predictions remain invariant, indicating greater reliance on causal rather than spurious features. We theoretically show that IWD reduces the student's Spurious-to-Causal (S2C) gradient ratio compared to standard uniformly weighted KD, driving the student toward more invariant representations. Experiments on four NLP benchmarks (MNLI, SQuAD-v2, CoNLL-2003 NER, and SST-2) across two model families (DeBERTa-v3 and Qwen-2.5) demonstrate that IWD consistently outperforms strong KD baselines on OOD evaluations while maintaining competitive in-distribution (ID) performance. Specifically, IWD achieves the highest accuracy in 15 out of 16 OOD benchmarks and improves average OOD performance over standard KD by 4.34 percentage points on NLI and 14.94 percentage points on QA.

Figures & tables

Figure 1: Overview of IWD. The pipeline operates in four stages: (1) Counterfactual Environment Generation: Generating semantically equivalent paraphrase variants; (2) Teacher Invariance Distance: Measuring prediction consistency across variants using the teacher predictions; (3) Weight Calculation: Assigning higher reliability weights to samples where teacher predictions remain invariant; and (4) Weighted Knowledge Distillation: Training the student model using the weighted loss objective.

DeBERTa-v3 Family (Encoder-Only)
Task / Dataset	Vanilla KD	RevKL	DKD	LWD	AugKD	IWD (Ours)
NLI (Accuracy)
MNLI (ID)	$82.02\pm 0.29$	$82.08\pm 0.21$	$79.91\pm 0.47$	$82.56\pm 0.12$	$\mathbf{85.31\pm 0.18}$	$\underline{84.66\pm 0.27}$
HANS (OOD-1)	$52.34\pm 0.57$	$51.45\pm 0.81$	$51.33\pm 0.79$	$52.99\pm 1.10$	$\underline{60.95\pm 1.30}$	$\mathbf{61.45\pm 1.01}$
SNLI (OOD-2)	$78.26\pm 0.58$	$77.49\pm 0.26$	$75.19\pm 0.85$	$78.97\pm 0.92$	$\underline{80.21\pm 0.24}$	$\mathbf{80.91\pm 0.50}$
Extractive QA (F1)

Table 1: Main performance comparison of distilled student models across four NLP tasks for both DeBERTa-v3 and Qwen-2.5 families, evaluated on one ID and two OOD benchmarks per task. Results report mean

\pm

standard deviation across multiple random seeds. The best performance per benchmark is in bold, and the second-best is underlined.

DeBERTa-v3
Data Size	MNLI-m	MNLI-mm	HANS	SNLI
( $N$ )	(ID)	(ID)	(OOD1)	(OOD2)
5k	83.27	83.79	51.98	79.73
8k	84.29	84.54	52.40	79.79
10k	84.84	85.11	60.83	80.56
30k	86.39	86.42	62.64	82.47

Table 2: Data scaling analysis of IWD on MNLI across varying sample sizes (

N

) for DeBERTa-v3/Qwen-2.5.

MNLI Benchmark
Environments	ID	OOD1	OOD2	OOD3	OOD4
Datasets	MNLI	HANS	SNLI	–	–
$E=3$	84.34	56.27	78.98	–	–
$E=4$	84.54	57.02	79.68	–	–
$E=5$	84.78	58.41	80.10	–	–
$E=6$	84.80	56.28	80.60	–	–

Table 3: Impact of varying the number of environments (

K

) in IWD across MNLI/SST-2 benchmarks on accuracy.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Training	ID Eval	OOD Sets	Metric
NLI (MNLI)	MNLI train	MNLI-m dev	HANS, SNLI	Accuracy
Sentiment (SST-2)	SST-2 train	SST-2 dev	Yelp-Polarity, CR	Accuracy
QA (SQuAD v2)	SQuAD v2 train	SQuAD v2 dev	NewsQA, Natural Questions	F1
NER (CoNLL-2003)	CoNLL-2003 train	CoNLL-2003 test	WNUT-17, OntoNotes 5.0	Entity-F1

Appendix

Table 4: Evaluation datasets. ID sets share the training distribution; OOD sets target specific distribution shifts. Metrics follow standard task practice.

Task	Environment	Text
MNLI	Original (label = entailment)	P: “The children were playing soccer in the field.” H: “Kids were engaged in a sport.”
	negation	P: “…in the field, and no adult was refereeing .”
	low_overlap	H: “Young people participated in athletic activity.”
	syntactic_shortcut_break	H: “It was a sport in which the kids were engaged.”
SST-2	Original (label = positive)	“The film was absolutely wonderful.”
	negation_invariant	“The film was not disappointing at all.”

Appendix

Table 5: Illustrative examples of counterfactual environment generation across three tasks. Each row shows one environment’s rewrite, with the modified span emphasised. Bold text in MNLI/SST-2 shows the injected transformation; NER entities are shown in brackets with their type.

NLI (DeBERTa-v3, 5k)
Task / Dataset	$\tau=0.1$	$\tau=0.3$	$\tau=1.0$	$\tau=3.0$	$\tau=10.0$	$\tau=30.0$
MNLI (ID)	$81.32\pm 0.27$	$82.30\pm 0.33$	$82.58\pm 0.31$	$\mathbf{82.63\pm 0.42}$	$82.57\pm 0.48$	$82.55\pm 0.50$
HANS (OOD-1)	$52.93\pm 1.74$	$53.56\pm 0.60$	$53.36\pm 1.35$	$\mathbf{54.03\pm 2.28}$	$53.98\pm 1.78$	$53.91\pm 1.61$
SNLI (OOD-2)	$76.39\pm 0.87$	$78.21\pm 0.51$	$78.61\pm 0.29$	$78.52\pm 0.55$	$78.76\pm 0.08$	$\mathbf{78.84\pm 0.20}$
Sentiment Analysis (DeBERTa-v3, 5k)
SST-2 (ID)	$92.20\pm 0.23$	$92.13\pm 0.35$	$\mathbf{92.28\pm 0.24}$	$92.05\pm 0.33$	$92.09\pm 0.23$	$92.13\pm 0.07$

Appendix

Table 6: Sensitivity analysis evaluating the impact of the distillation temperature parameter

\tau

on IWD across four task families using DeBERTa-v3 (trained on 5k subsamples).

MNLI Benchmark
Distance Metric	ID	OOD 1	OOD 2
	MNLI-m	HANS	SNLI
Jensen–Shannon ( $D_{\text{JS}}$ )	82.60 $\pm$ 0.28	53.35 $\pm$ 1.38	78.68 $\pm$ 0.34
Symmetric KL ( $D_{\text{SymKL}}$ )	81.38 $\pm$ 0.40	53.01 $\pm$ 2.12	77.43 $\pm$ 0.74
$L_{2}$ Distance	82.38 $\pm$ 0.27	52.76 $\pm$ 1.03	78.58 $\pm$ 0.33
SST-2 Sentiment Benchmark

Appendix

Table 7: Sensitivity of IWD to distance metric formulations across tasks on DeBERTa-v3 (

N=5\text{k}

JS Distance (Primary Metric)
Variant Strategy	MNLI-M	MNLI-MM	HANS	SNLI
& Metric	(ID)	(ID)	(OOD1)	(OOD2)
All Variants	84.238	84.947	57.323	80.542
Original Variant	82.812	83.411	56.206	77.433
Random Variant	82.873	83.441	53.103	78.185
Sym KL Divergence

Appendix

Table 8: Ablation study evaluating environment variant selection strategies during student distillation across different distance metrics using DeBERTa-v3 on ID and OOD benchmarks.

Backbone	Norm	ID	HANS (OOD1)	SNLI (OOD2)
Task: MNLI
DeBERTa-v3	Batch-Mean	83.25	51.98	79.66
	None	83.22	51.80	79.91
	$\Delta$	-0.03	-0.18	+0.25
Qwen-2.5	Batch-Mean	81.87	53.19	78.39
	None	81.70	53.20	78.23

Appendix

Table 9: Ablation study comparing batch-mean normalization against no weight normalization on MNLI and SST-2 across DeBERTa-v3 and Qwen-2.5 backbones.

Task	$\boldsymbol{w_{i}}$	% $\boldsymbol{w_{i}}$	% $\boldsymbol{w_{i}}$	Effective
	(Mean $\pm$ Std)	$\mathbf{\geq 0.9}$	$\mathbf{<0.5}$	$\boldsymbol{N}$
MNLI	$0.87\pm 0.10$	40.8%	0.0%	4,931 (98.6%)
SST-2	$0.95\pm 0.08$	78.0%	0.0%	4,966 (99.3%)
SQuAD v2	$0.86\pm 0.20$	63.5%	4.4%	4,748 (95.0%)
NER	$0.48\pm 0.06$	0.0%	58.7%	4,931 (98.6%)

Appendix

Table 10: Distribution of per-sample IWD weights across training data at

\tau=1.0

and batch-mean normalization (

N=5{,}000

per task). Effective sample size

(\sum_{i}w_{i})^{2}/\sum_{i}w_{i}^{2}

measures information retention relative to uniform weighting, where values near

N

indicate gentle re-weighting rather than aggressive filtering.

RL for Knowledge Distillation

See all 12 papers

Jun 3, 2026

Invariant Gradient Alignment for Robust Reasoning Distillation

Large language models (LLMs) suffer from shortcut learning: they systematically fail on out-of-distribution (OOD) inputs whose semantic surface differs from training data, even when the logical structure is identical. This undermines knowledge distillation pipelines that transfer chain-of-thought reasoning to smaller students. We introduce Invariant Gradient Alignment (IGA), a training framework that aligns gradient updates across semantically diverse but logically isomorphic examples via three innovations: (i) Logical Isomer Sets, groups of problems sharing identical logical structure across distinct semantic domains (mathematics, medicine, law, science); (ii) a differentiable Continuous Gradient Conflict Mask, that suppresses parameter dimensions with high cross-domain gradient variance while preserving invariant directions; and (iii) a truncated SVD projection of the masked gradient back onto the LoRA low-rank manifold, maintaining parameter efficiency throughout. Theoretically, IGA yields tighter OOD generalization bounds than ERM, scaling with the number of isomer domains, and converges at the standard SGD rate under mild regularity. Empirically, IGA outperforms eight baselines across four benchmarks with accuracy gains up to 14.3 pp over ERM-SFT and a Logical Consistency Score of 0.031 versus 0.142 -- a fourfold improvement in representational invariance.

Zehua Cheng, Wei Dai, Jiahao Sun

University of Oxford, Oxford, UK

CoT Distillation Domain Generalization OOD Generalization Knowledge Distillation

Ask PDF

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

Sep 29, 2026

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

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.

Ayan Sengupta, Vaibhav Seth, Tanmoy Chakraborty

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

Language Model Distillation Selective Knowledge Distillation Uncertainty-Aware Knowledge Distillation LLM Reliability Language Model Calibration

Ask PDF

Abstract

RL for Knowledge Distillation

Explore similar work

Invariant Gradient Alignment for Robust Reasoning Distillation

Distillation Traps and Guards: A Calibration Knob for LLM Distillability

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