Paper deep dive
Do Student LLMs Inherit OOD Robustness? Invariance-Weighted Distillation for Reliable Knowledge Transfer
Dileesha Kannangara, Sanghamitra Dutta
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 9/22/2026, 4:53:52 AM
Summary
The paper introduces Invariance-Weighted Distillation (IWD), a framework designed to improve the out-of-distribution (OOD) robustness of student Large Language Models (LLMs) during knowledge distillation. Standard KD often fails to transfer OOD robustness because students learn spurious correlations and treat all teacher predictions uniformly, even when the teacher is misled by spurious features. IWD addresses this by generating counterfactual variants of training data that perturb spurious cues while preserving semantics, measuring the teacher's prediction invariance across these variants, and using this invariance to weight the distillation loss. Theoretical analysis shows IWD reduces the Spurious-to-Causal (S2C) gradient ratio, and empirical results on NLP benchmarks (MNLI, SQuAD-v2, CoNLL-2003, SST-2) demonstrate significant improvements in OOD performance over standard KD baselines.
Entities (11)
Relation Signals (10)
Invariance-Weighted Distillation → evaluatedon → MNLI
confidence 95% · Experiments on four NLP benchmarks (MNLI...)
Invariance-Weighted Distillation → evaluatedon → SST-2
confidence 95% · Experiments on four NLP benchmarks (...and SST-2)
Invariance-Weighted Distillation → evaluatedon → CoNLL-2003
confidence 95% · Experiments on four NLP benchmarks (...CoNLL-2003 NER...)
Invariance-Weighted Distillation → evaluatedon → SQuAD v2
confidence 95% · Experiments on four NLP benchmarks (...SQuAD-v2...)
Invariance-Weighted Distillation → improves → Out-of-Distribution Robustness
confidence 95% · IWD consistently outperforms strong KD baselines on OOD evaluations while maintaining competitive in-distribution (ID) performance.
Invariance-Weighted Distillation → reduces → Spurious-to-Causal Gradient Ratio
confidence 92% · We theoretically show that IWD reduces the student's Spurious-to-Causal (S2C) gradient ratio compared to standard uniformly weighted KD
Invariance-Weighted Distillation → appliedto → DeBERTa-v3
confidence 90% · across two model families (DeBERTa-v3 and Qwen-2.5)
Invariance-Weighted Distillation → appliedto → QWEN 2.5
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2609.22566v1
- Canonical: https://arxiv.org/abs/2609.22566v1
Trouble viewing inline? Open PDF directly →
Full Text
94,998 characters extracted from source content.
Expand or collapse full text
Do Student LLMs Inherit OOD Robustness? Invariance-Weighted Distillation for Reliable Knowledge Transfer Dileesha Kannangara Sanghamitra Dutta 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. University of Maryland, College Park dileesha@umd.edu, sanghamd@umd.edu Introduction Large Language Models (LLMs) are central to modern applications due to their remarkable ability to solve complex tasks (OpenAI 2023; Yang et al. 2024; DeepSeek-AI 2025). However, high computational and memory demands hinder training, fine-tuning, and deployment in resource-constrained environments (Girija et al. 2025). Knowledge distillation (KD) addresses these challenges by training a compact student model to mimic the predictive behavior of a high-capacity teacher model across a designated distillation dataset (Hinton et al. 2015; Xu et al. 2024). Distilled students deliver competitive performance with drastically reduced resource overhead, enabling efficient real-world deployment. However, student models often demonstrate significant performance degradation on out-of-distribution (OOD) data compared to the in-distribution (ID) data used during training (Du et al. 2023; Stanton et al. 2021). OOD data can be formally defined as data originating from a distribution that differs from that encountered during any stage of training prior to deployment (Shen et al. 2021). Although teacher models themselves are not fully robust to OOD data, we observe that the performance gap between teacher and student models becomes substantially larger under OOD settings than under ID settings. This suggests that current KD techniques do not effectively transfer the teacher model’s existing ability to handle OOD data and directly undermine this capability in the resulting student models (Stacey and Rei 2024). This leads to two central questions: To what extent do student LLMs inherit OOD robustness from their teachers, and how can we transform distillation to enable reliable knowledge transfer under distribution shifts? We identify two main limitations of current KD techniques that cause OOD performance degradation. First, under existing distillation methods, students learn spurious correlations from the distillation dataset instead of genuine causal relationships. Standard KD approaches train the student model to mimic the teacher’s outputs (Hinton et al. 2015; Gu et al. 2024) or intermediate representations (Romero et al. 2015; Sun et al. 2019; Sun et al. 2020; Jiao et al. 2020) by minimizing discrepancies over the distillation dataset. Additionally, distillation minimizes the standard cross-entropy loss over these training samples. In doing so, the student learns dataset-specific spurious correlations, reducing its ability to generalize when encountering OOD data. Second, standard distillation treats all training samples uniformly, ignoring the teacher’s ability to isolate causal signals from spurious correlations across individual samples. In practice, a teacher model may successfully resist spurious shortcuts on samples with strong causal semantics, yet succumb to spurious correlations on others. Standard KD applies a uniform weight across all samples, inadvertently forcing the student to imitate the teacher even when the teacher relies on brittle, spurious associations. For reliable knowledge transfer, distillation should prioritize samples where the teacher exhibits high fidelity to invariant causal features while suppressing spurious influences. 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. To address these limitations, we propose Invariance-Weighted Distillation (IWD), a novel strategy that uses a weighted objective function to prioritize samples where the teacher relies on causal relationships over spurious correlations (see Fig. 1). We first generate multiple synthetic environments for each training sample by altering its spurious features while preserving its original semantic meaning and label. This environment generation step resolves the first limitation by exposing the student to diverse spurious perturbations across counterfactual variants with fixed core semantics. Distilling across these environments prevents the student from learning dataset-specific spurious correlations and forces it to learn invariant representations that generalize across distribution shifts. To address the second limitation, we measure the discrepancy among the teacher’s output distributions across these environments. This distance quantifies the teacher’s prediction invariance, i.e., how strongly the teacher relies on invariant features versus environment-specific cues. We incorporate this distance into the distillation loss as a weight, increasing emphasis on samples where the teacher exhibits high prediction invariance and, consequently, high fidelity to invariant causal features. Contributions Our main contributions are as follows: • Conceptualizing OOD Vulnerabilities in KD. We identify two key failure modes of standard KD under distribution shifts: (1) student models learn spurious correlations from the distillation dataset, and (2) current KD methods treat all training samples equally, failing to recognize the teacher’s ability to distinguish between causally robust and spurious-dominated samples. • Invariance-Weighted Distillation (IWD). We propose IWD, a novel distillation strategy that assigns sample-specific weights based on the teacher’s prediction invariance across synthetically perturbed environments (see Fig. 1). To overcome Limitation 1, IWD varies spurious features while preserving core causal semantics, exposing the student to diverse counterfactual variations of each sample. To address Limitation 2, IWD prioritizes samples where the teacher exhibits high invariance, and thus high fidelity to causal features. Our strategy equips the student with improved robustness to spurious features over standard KD and conventional data augmentation. • Theoretical Guarantee. We provide a theoretical guarantee for IWD by leveraging a metric called the Spurious-to-Causal (S2C) gradient ratio (see Theorem 1). We prove that IWD yields a lower S2C ratio than standard distillation, guaranteeing that the student’s expected gradient aligns more strongly with invariant signals while reducing spurious feature transfer. • Empirical Validation. We evaluate IWD across four distinct NLP tasks (NLI, extractive QA, NER, and sentiment classification) and two model families (DeBERTa-v3 and Qwen-2.5), using several OOD datasets per task under various subsampling regimes. IWD consistently improves OOD robustness while maintaining ID performance. For instance, across model families, IWD improves average OOD performance over standard KD by 4.34 and 14.94 percentage points on NLI and QA, respectively. Related Works Existing KD methods (Hinton et al. 2015; Zhu et al. 2024; Xu et al. 2024; Romero et al. 2015; Sun et al. 2019; Zagoruyko and Komodakis 2017; Tian et al. 2020; Kim and Rush 2016; Hamman et al. 2025a; Dissanayake et al. 2025) focuses on matching teacher and student outputs, intermediate features, or attention patterns. However, the robustness of distilled models under distribution shifts remains underexplored. While ShiftKD (Zhang et al. 2025) shows that distillation suffers substantial performance degradation under covariate and correlation shifts, other approaches attempt to enhance robustness via confidence-guided data augmentation (Popp et al. 2025), adversarial distillation (Goldblum et al. 2019), or distillation from robust foundation models (Zhou et al. 2023). Yet, these methods primarily target conventional deep learning models, leaving OOD robustness in distilled LLMs underexplored. Prior work on OOD robustness in distilled LLMs seeks to improve performance through data augmentation, sample selection, or adaptive distillation strategies.(Stacey and Rei 2024) use domain-targeted data augmentation and sampling strategies to improve NLI distillation under distribution shift. Prior works Chen et al. (2023) and Hamman et al. (2025a) leverage LLMs to generate label-flipping counterfactuals, prompting students to learn causal features over spurious correlations. Other works introduce difficulty-aware distillation using model pruning (Du et al. 2023) or OOD-guided data generation to expose students to more diverse samples (Gholami et al. 2024). Notably, these works do not explicitly enforce invariance across semantically equivalent perturbations and rely on uniform distillation objectives that treat all training samples equally. In contrast, our approach evaluates the teacher’s prediction invariance across synthetically perturbed environments with varying spurious features, leveraging the teacher-invariance signal to weight the distillation loss toward more reliable samples. Preliminaries Knowledge distillation (Xu et al. 2024; Hamman et al. 2025a) trains a smaller student fθSf_ _S to imitate a larger teacher fθTf_ _T. For an input x with target label y, the student’s loss is: ℒKD=(1−λ)ℒCE(y,pS(x))+λℒKL(pT(x)∥pS(x)),L_KD=(1-λ)\,L_CE(y,p_S(x))+λ\,L_KL\,(p_T(x) p_S(x)), (1) Here, ℒCEL_CE is the cross-entropy (CE) loss, ℒKLL_KL is the KL divergence (Kullback and Leibler 1951) at each generation step over the vocabulary, pS(x)=softmax(fθS(x))p_S(x)=softmax(f_ _S(x)) and pT(x)=softmax(fθT(x))p_T(x)=softmax(f_ _T(x)) are the output distributions, and λ∈[0,1]λ∈[0,1] balances hard-label CE and soft target imitation. Out-of-distribution (OOD) refers to the setting in which a deployed model make predictions on inputs that come from a different distribution from the one it was trained on. A model is said to be OOD-robust if its performance does not degrade substantially under such a shift. Formally, let PtrainP_train denote the training distribution and POOD≠PtrainP_OOD≠ P_train denote a shifted evaluation distribution. Standard supervised learning offers no formal guarantee in this regime, and trained models often degrade sharply (Hendrycks and Dietterich 2019; Geirhos et al. 2020), particularly when they exploit features that are spuriously correlated with the label in PtrainP_train but uncorrelated in POODP_OOD. OOD robustness is measured via the metric: ϕ(θ)=(x,y)∼POOD[ 1[fθ(x)=y]],φ(θ)\;=\;E_(x,y) P_OOD\! [\,1[f_θ(x)=y]\, ], (2) where θ denotes the parameters of the model under evaluation, fθ(x)f_θ(x) is its predicted label for input x, and [⋅]1[·] is the indicator function returning 11 when its argument is true and 00 otherwise. The expectation (x,y)∼POOD[⋅]E_(x,y) P_OOD[·] averages over input-label pairs drawn from the shifted distribution, so ϕ(θ)φ(θ) is the expected accuracy of fθf_θ on POODP_OOD. Problem Setup. Our goal is to study how KD in LLMs affects the OOD robustness of the student model. Formally, let =(xi,yi)i=1ND=\(x_i,y_i)\_i=1^N be a labeled training set drawn from an in-distribution (ID) source PtrainP_train, and let POODP_OOD represent a shifted evaluation distribution. Given a teacher model fθTf_ _T pre-trained on PtrainP_train, our objective is to learn a smaller student model fθSf_ _S that maximizes the OOD evaluation metric ϕ(θS)φ( _S). Specifically, our goal is to distill student parameters θS _S that learn genuine causal relationships while suppressing spurious correlations, and achieve higher OOD performance than standard distillation, i.e., ϕ(θS)>ϕ(θSKD)φ( _S)>φ( _S^KD), where θSKD _S^KD denotes parameters learned via standard KD. Main Contributions We introduce Invariance-Weighted Distillation (IWD), a framework enhancing OOD robustness in distilled LLMs. Our strategy operates in four stages: (i) generation of counterfactual environment variants that vary spurious features while preserving the semantic meaning; (i) computation of the teacher-invariance distance that quantifies the teacher’s prediction variation across those environments; (i) mapping distance to sample weights; and (iv) distilling using our proposed weighted distillation objective. Stage 1: Counterfactual Environment Generation. For each training sample (xi,yi)∈(x_i,y_i) , we construct a set of K counterfactual variants: ℰi=xi(0),xi(1),…,xi(K−1),E_i\;=\; \\,x_i^(0),\,x_i^(1),\,…,\,x_i^(K-1)\, \, (3) where xi(0)≜xix_i^(0) x_i is the original sample, and xi(1),…,xi(K−1)x_i^(1),…,x_i^(K-1) are paraphrases generated across different environments e(k)e^(k). These environments are designed to perturb surface-level spurious features such as register, syntactic structure, domain vocabulary, capitalization, and sentence length, while holding the core semantic features determining yiy_i fixed. Formally, each input can be decomposed into an invariant causal component ic_i that determines the label and an environment-specific spurious component i(k)s_i^(k). The generator ensures these variants satisfy counterfactual label invariance: P(yi|i,i(k))=P(yi|i)∀k∈0,1,…,K−1.P (y_i\, |\,c_i,\,s_i^(k) )=P (y_i\, |\,c_i ) ∀\,k∈\0,1,…,K-1\. (4) Equation 4 guarantees that the label distribution conditioned on causal features remains identical across all environments, isolating variation strictly to the spurious components. We construct ℰiE_i offline by prompting an auxiliary language model (GPT-4.1-mini) to generate specific spurious perturbations. To enforce label invariance, any paraphrase that would alter the target label yiy_i is rejected. The union of these environment sets forms the augmented training dataset aug=⋃i=1NℰiD_aug= _i=1^NE_i, which is precomputed once and shared across all experiments. Our semantics-preserving counterfactual environment generation strategy directly addresses the first limitation by preventing the student from learning dataset-specific spurious correlations. Exposing the model to these variants influences both components of the training objective: (i) distillation loss (ℒKLL_KL), which forces the student to match teacher predictions across varied spurious contexts rather than context-dependent noise; and (i) task loss (ℒCEL_CE), which trains the student directly on perturbed inputs to ensure ground-truth predictions rely on invariant semantics rather than surface shortcuts. Stage 2: Teacher Invariance Distance. Using the environment set ℰiE_i, we directly quantify the invariance of the teacher’s predictions across counterfactual perturbations. Definition 1 (Teacher Invariance Distance). Let y^i(k)=fθT(xi(k)) y_i^(k)=f_ _T (x_i^(k) ) denote the teacher’s prediction on the k-th variant in a task-appropriate output space. The teacher invariance distance DiD_i is defined as the mean pairwise disagreement across all environment pairs: Di≜(K2)−1∑0≤u<v≤K−1d(y^i(u),y^i(v)).D_i\; \; K2^-1 _0≤ u<v≤ K-1d ( y_i^(u),\, y_i^(v) ). (5) Here, d(⋅,⋅)d(·,·) is a non-negative, symmetric distance metric. Since causal features ic_i remain fixed across variants, DiD_i directly isolates the teacher’s sensitivity to spurious perturbations i(k)s_i^(k). A small DiD_i indicates that predictions are invariant and grounded in causal semantics, making the sample a reliable source for distillation. Conversely, a large DiD_i reveals that predictions fluctuate with surface cues, signaling the teacher’s reliance on spurious features that will not generalize in OOD settings. The distance metric d(⋅,⋅)d(·,·) is tailored to the output space of the task. For classification (e.g., SST-2, MNLI), d measures distributional divergence over softmax outputs; for structured symbolic outputs like QA text spans or NER entities (SQuAD-v2, CoNLL-2003), d computes per-sample F1F_1 distance (1−F11-F_1) on decoded string predictions. Stage 3: Weight Calculation. To convert the invariance distance DiD_i into a per-sample reliability weight, we apply a monotonically decreasing exponential function: wi=exp(−Di/τ).w_i\;=\; (-D_i/τ). Here, τ>0τ>0 is a temperature hyperparameter. A small τ heavily penalizes teacher variance, while a large τ approaches uniform weighting, recovering environment-augmented distillation as τ→∞τ→∞. Furthermore, to prevent implicit learning rate shifts caused by varying DiD_i distributions across mini-batches, we normalize weights over each mini-batch ℬB: w~i=|ℬ|⋅wi∑j∈ℬwj, w_i\;=\;|B|\;·\; w_i _j w_j, ensuring 1|ℬ|∑i∈ℬw~i= 1. 1|B| _i w_i\;=\;1. This normalization preserves a unit average weight per batch while adjusting relative per-sample importance. Stage 4: Weighted Knowledge Distillation. During training, we sample mini-batches ℬ⊆augB _aug, where each item corresponds to an environment variant xi(k)x_i^(k) derived from parent sample i. Given the batch-normalized weights w~i w_i, the weighted distillation objective over the mini-batch is: ℒIWD(θS)=∑(i,k)∈ℬw~i⋅KL(fθT(xi(k))∥fθS(xi(k))).L_IWD( _S)= _(i,k) w_i·KL (f_ _T(x_i^(k)) f_ _S(x_i^(k)) ). (6) Combining ℒIWDL_IWD with the student task loss across variants, ℒCE=∑(i,k)∈ℬCE(yi,fθS(xi(k)))L_CE= _(i,k) CE (y_i,f_ _S(x_i^(k)) ), yields our total training loss: ℒtotal=(1−λ)ℒCE+λℒIWDL_total=(1-λ)L_CE+ _IWD (Eq. 1). By dynamically scaling the loss with w~i w_i, IWD prioritizes samples where the teacher demonstrates invariant predictions while aggressively down-weighting unstable targets (wi→0w_i→ 0). This suppresses teacher signals on brittle samples, forcing the student to rely on ground-truth supervision instead. The complete training procedure for IWD is detailed in Algorithm 1. Algorithm 1 Invariance Weighted Distillation (IWD) Input: =(xi,yi)i=1ND=\(x_i,y_i)\_i=1^N, teacher fθTf_ _T, student fθSf_ _S Parameters: No. of environments K, temperature τ, learning rate α, distance metric d(⋅,⋅)d(·,·) Output: Trained student fθSf_ _S 1: – Stage 1: Counterfactual Environment Generation – 2: aug←∅D_aug← 3: for each (xi,yi)∈(x_i,y_i) do 4: Set xi(0)←xix_i^(0)← x_i Original sample 5: Generate K−1K-1 paraphrases xi(1),…,xi(K−1) \x_i^(1),…,x_i^(K-1) \ 6: ℰi←xi(0),xi(1),…,xi(K−1)E_i← \x_i^(0),x_i^(1),…,x_i^(K-1) \ 7: aug←aug∪ℰiD_aug _aug _i 8: end for 9: – Stage 2: Teacher Invariance Distance – 10: for each i∈1,…,Ni∈\1,…,N\ do 11: y^i(k)←fθT(xi(k))∀k∈0,…,K−1 y_i^(k)← f_ _T(x_i^(k)) ∀\,k∈\0,…,K-1\ 12: Di←(K2)−1∑0≤u<v≤K−1d(y^i(u),y^i(v))D_i← K2^-1 _0≤ u<v≤ K-1d\! ( y_i^(u),\, y_i^(v) ) 13: end for 14: – Stage 3: Weight Calculation – 15: for each i∈1,…,Ni∈\1,…,N\ do 16: wi←exp(−Di/τ)w_i← (-D_i/τ) Parent weight once per-sample 17: end for 18: – Stage 4: Weighted Knowledge Distillation – 19: while not converged do 20: Sample mini-batch ℬ⊆augB _aug, where each item is indexed by (i,k)(i,k) 21: w~i←|ℬ|⋅wi/∑(j,k′)∈ℬwj∀(i,k)∈ℬ w_i←|B|· w_i\;/\; _(j,k ) w_j ∀\,(i,k) 22: ℒCE←∑(i,k)∈ℬCE(yi,fθS(xi(k)))L_CE← _(i,k) CE (y_i,\,f_ _S(x_i^(k)) ) 23: ℒIWD←∑(i,k)∈ℬw~i⋅KL(fθT(xi(k))∥fθS(xi(k)))L_IWD← _(i,k) w_i·KL (f_ _T(x_i^(k))\,\|\,f_ _S(x_i^(k)) ) 24: ℒtotal←(1−λ)ℒCE+λℒIWDL_total←(1-λ)L_CE+ _IWD 25: θS←θS−α∇θSℒtotal _S← _S-α\, _ _SL_total 26: end while Theoretical Guarantee. We now provide a theoretical guarantee for weighting the distillation objective by the teacher invariance distance. We first introduce the Spurious-to-Causal (S2C) gradient ratio (Definition 2), which measures gradient contamination from spurious feature variations under fixed causal semantics. We show that IWD achieves a strictly lower S2C ratio than uniform distillation, aligning student updates more closely with invariant causal features. Definition 2 (Spurious-to-Causal (S2C) Gradient Ratio). Let pT(x)=softmax(zT(x))p_T(x)=softmax (z_T(x) ) and pS(x)=softmax(zS(x))p_S(x)=softmax (z_S(x) ) denote teacher and student output distributions for logits zT,zSz_T,z_S. Decomposing the teacher’s output distribution into invariant and spurious components, pT(xi(e))=πi+si(e),πi:=e[pT(xi(e))],p_T (x_i^(e) )\;=\; _i+s_i^(e), _i\;:=\;E_e\! [p_T (x_i^(e) ) ], (7) where πi∈ΔC−1 _i∈ ^C-1 is the invariant component, and si(e)s_i^(e) is the spurious component with e[si(e)]=E_e [s_i^(e) ]=0 and variance e[‖si(e)‖2]=σi2E_e [\|s_i^(e)\|^2 ]= _i^2. The S2C gradient ratio under sample weighting W is defined as: ρ2(W)=i∼W[‖giS‖2]/i∼W[‖giC‖2],ρ^2(W)\;=\;E_i W\! [\|g_i^S\|^2 ] /E_i W\! [\|g_i^C\|^2 ], (8) where giC=(πi−pS(xi))⊤Jig_i^C= ( _i-p_S(x_i) )^\! J_i is the causal gradient component, giS=si(ei)⊤Jig_i^S=s_i^(e_i) J_i is the spurious gradient component, and Ji:=∇θSzS(xi)∈ℝC×|θS|J_i:= _ _Sz_S(x_i) ^C×| _S| is the student logit Jacobian. A lower S2C ratio indicates a causally dominated gradient, ensuring the student prioritizes invariant predictions over non transferable surface perturbations. We analyze the distillation objective on dataset D across environments ℰii=1N\E_i\_i=1^N, comparing uniform weighting (Wi=1/NW_i=1/N) against IWD weighting (WiIWD∝exp(−Di/τ)W_i^IWD (-D_i/τ)). As K→∞K→∞, DiD_i converges to a strictly increasing function of spurious variance σi2 _i^2 (Proposition 1 in Appendix A.4 ), allowing DiD_i to directly isolate teacher spurious reliance. See Appendix A for more details. Our analysis relies on three standard assumptions: (A1) Zero-Mean Spurious Noise: Spurious components si(ei)s_i^(e_i) are independent across samples with zero mean and variance σi2 _i^2, representing unstructured prediction variation rather than systematic dataset bias. (A2) Independence of Spurious Variance: Spurious variance is independent of both the causal residual and gradient scale: σi2⟂(ci2,Φi) _i^2 (c_i^2, _i ), where ci2:=‖πi−pS(xi)‖2c_i^2:=\| _i-p_S(x_i)\|^2. This decouples teacher prediction invariance from student task difficulty. (A3) Isotropic Logit Jacobians: Student logit Jacobians satisfy JiJi⊤=ΦiICJ_iJ_i = _i\,I_C with per sample scale Φi⟂σi2 _i _i^2, ensuring parameter updates do not inherently favor spurious directions over causal ones. Theorem 1 (IWD Reduces the S2C Gradient Ratio). Under Assumptions (A1)–(A3), let WiIWD=exp(−σi2/τ)/ZW_i^IWD= (- _i^2/τ)/Z and Wiunif=1/NW_i^unif=1/N. Then, ρ2(WIWD)≤ρ2(Wunif)ρ^2\! (W^IWD )≤ρ^2\! (W^unif ), with equality if and only if σi2 _i^2 is constant across all samples. Experiments We evaluate IWD on four tasks spanning the major NLP output structures—paired classification, single-sentence classification, span extraction, and token-level labeling—across two model families of markedly different scale and architecture. Additional details are given in Appendix B. Task / Dataset Vanilla KD RevKL DKD LWD AugKD IWD (Ours) DeBERTa-v3 Family (Encoder-Only) NLI (Accuracy) MNLI (ID) 82.02±0.2982.02± 0.29 82.08±0.2182.08± 0.21 79.91±0.4779.91± 0.47 82.56±0.1282.56± 0.12 85.31±0.1885.31± 0.18 84.66±0.27¯ 84.66± 0.27 HANS (OOD-1) 52.34±0.5752.34± 0.57 51.45±0.8151.45± 0.81 51.33±0.7951.33± 0.79 52.99±1.1052.99± 1.10 60.95±1.30¯ 60.95± 1.30 61.45±1.0161.45± 1.01 SNLI (OOD-2) 78.26±0.5878.26± 0.58 77.49±0.2677.49± 0.26 75.19±0.8575.19± 0.85 78.97±0.9278.97± 0.92 80.21±0.24¯ 80.21± 0.24 80.91±0.5080.91± 0.50 Extractive QA (F1) SQuAD-v2 (ID) 70.64±0.20¯ 70.64± 0.20 72.40±0.8172.40± 0.81 68.01±1.8168.01± 1.81 70.97±1.3570.97± 1.35 51.50±0.5551.50± 0.55 51.40±0.7751.40± 0.77 NewsQA (OOD-1) 39.47±0.7739.47± 0.77 39.95±0.8739.95± 0.87 44.26±0.6844.26± 0.68 36.10±2.4836.10± 2.48 54.35±0.66¯ 54.35± 0.66 55.11±0.4555.11± 0.45 Natural Questions (OOD-2) 28.23±0.7428.23± 0.74 28.11±2.6728.11± 2.67 30.52±2.3330.52± 2.33 25.02±5.1725.02± 5.17 51.93±0.94¯ 51.93± 0.94 52.06±0.3352.06± 0.33 NER (Entity F1) CoNLL-2003 (ID) 92.18±0.1192.18± 0.11 90.59±0.2190.59± 0.21 90.11±0.2190.11± 0.21 90.71±0.2090.71± 0.20 92.63±0.22¯ 92.63± 0.22 92.64±0.1092.64± 0.10 WNUT-17 (OOD-1) 47.84±0.42¯ 47.84± 0.42 43.68±0.5143.68± 0.51 46.88±0.3546.88± 0.35 46.34±0.5246.34± 0.52 46.78±0.4846.78± 0.48 48.01±0.3148.01± 0.31 OntoNotes 5.0 (OOD-2) 78.99±0.2178.99± 0.21 76.26±0.1976.26± 0.19 77.68±0.3277.68± 0.32 78.12±0.1178.12± 0.11 80.27±0.22¯ 80.27± 0.22 80.38±0.1180.38± 0.11 Sentiment Analysis (Accuracy) SST-2 (ID) 92.73±0.1992.73± 0.19 92.84±0.39¯ 92.84± 0.39 92.71±0.2292.71± 0.22 92.00±0.3092.00± 0.30 93.21±0.1593.21± 0.15 92.78±0.1892.78± 0.18 Yelp-Polarity (OOD-1) 91.39±0.19¯ 91.39± 0.19 91.28±0.2491.28± 0.24 91.36±0.3891.36± 0.38 90.26±0.6590.26± 0.65 90.98±0.5690.98± 0.56 91.42±0.1291.42± 0.12 CR (OOD-2) 87.29±1.17¯ 87.29± 1.17 87.08±0.7587.08± 0.75 87.18±0.4887.18± 0.48 85.81±0.2585.81± 0.25 85.32±0.4885.32± 0.48 87.58±0.9887.58± 0.98 Qwen-2.5 Family (Decoder-Only) NLI (Accuracy) MNLI (ID) 81.88±0.5081.88± 0.50 82.13±0.3082.13± 0.30 79.24±0.5379.24± 0.53 80.27±0.3680.27± 0.36 82.48±0.20¯ 82.48± 0.20 83.21±0.1983.21± 0.19 HANS (OOD-1) 51.69±0.3851.69± 0.38 52.04±0.8552.04± 0.85 51.61±1.0651.61± 1.06 51.20±0.6351.20± 0.63 52.79±1.21¯ 52.79± 1.21 54.95±1.6854.95± 1.68 SNLI (OOD-2) 78.25±0.6378.25± 0.63 79.03±0.96¯ 79.03± 0.96 75.28±1.6875.28± 1.68 76.38±0.8676.38± 0.86 76.99±0.7976.99± 0.79 80.57±0.6780.57± 0.67 Extractive QA (F1) SQuAD-v2 (ID) 75.48±0.3175.48± 0.31 75.02±0.23¯ 75.02± 0.23 74.68±0.3774.68± 0.37 49.86±0.2949.86± 0.29 47.96±0.4347.96± 0.43 48.84±0.2148.84± 0.21 NewsQA (OOD-1) 16.43±0.7316.43± 0.73 16.03±1.3016.03± 1.30 14.88±2.1114.88± 2.11 5.11±1.215.11± 1.21 19.53±0.65¯ 19.53± 0.65 20.34±0.1820.34± 0.18 Natural Questions (OOD-2) 40.89±0.8140.89± 0.81 43.73±2.1943.73± 2.19 35.57±1.0235.57± 1.02 2.37±1.052.37± 1.05 57.12±0.60¯ 57.12± 0.60 57.25±0.5057.25± 0.50 NER (Entity F1) CoNLL-2003 (ID) 68.30±0.2568.30± 0.25 67.69±0.3067.69± 0.30 65.00±0.2265.00± 0.22 35.52±0.3635.52± 0.36 69.83±0.2069.83± 0.20 69.29±0.10¯ 69.29± 0.10 WNUT-17 (OOD-1) 27.48±0.32¯ 27.48± 0.32 24.88±0.4224.88± 0.42 26.13±0.5226.13± 0.52 12.81±1.3012.81± 1.30 25.02±0.4325.02± 0.43 27.77±0.3327.77± 0.33 OntoNotes 5.0 (OOD-2) 54.90±0.40¯ 54.90± 0.40 54.93±1.0054.93± 1.00 53.57±0.6853.57± 0.68 26.58±0.9026.58± 0.90 54.67±0.4254.67± 0.42 55.03±0.6755.03± 0.67 Sentiment Analysis (Accuracy) SST-2 (ID) 93.46±0.5493.46± 0.54 93.17±0.5893.17± 0.58 93.43±0.22¯ 93.43± 0.22 60.17±3.9760.17± 3.97 92.04±0.8992.04± 0.89 90.11±0.6690.11± 0.66 Yelp-Polarity (OOD-1) 84.69±3.1184.69± 3.11 85.99±1.5985.99± 1.59 84.99±2.91¯ 84.99± 2.91 51.17±1.6351.17± 1.63 77.80±0.3177.80± 0.31 83.64±0.8683.64± 0.86 CR (OOD-2) 74.57±10.6674.57± 10.66 76.87±7.02¯ 76.87± 7.02 74.95±9.5074.95± 9.50 62.07±3.3362.07± 3.33 73.00±1.0973.00± 1.09 77.80±2.2977.80± 2.29 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 ± standard deviation across multiple random seeds. The best performance per benchmark is in bold, and the second-best is underlined. Datasets and Models. We train on standard ID splits and evaluate OOD generalization across two robustness benchmarks per task (Table 4). To cover diverse architectures, we distill both encoder-only models (deberta-v3-xsmall, ∼ 71M parameters from -base teachers) and decoder-only models (Qwen2.5-0.5B, ∼ 494M from -1.5B teachers); full setup details are in Appendix B.2. Baselines. We compare against five key KD baselines: (1) Vanilla KD (Hinton et al. 2015) (forward KL divergence); (2) RevKL (Gu et al. 2024) (reverse KL divergence); (3) DKD (Zhao et al. 2022) (decoupled target and non-target KL); (4) LWD (Romero et al. 2015; Sun et al. 2019; Sun et al. 2020; Jiao et al. 2020) (intermediate-layer representation alignment); and (5) AugKD (Vanilla KD on counterfactually augmented data, isolating our invariance weighting scheme from pure data augmentation). Training Setup. We train models using AdamW in mixed precision (T=2.0,τ=1.0,K=6T=2.0,τ=1.0,K=6); full implementation details and hyperparameters are provided in Appendix B.4. Results and Analysis IWD outperforms baselines in OOD settings. Across both architectures, IWD achieves the top performance in 15 of 16 OOD settings, with the largest gains in Extractive QA and NLI (see Table 1). On DeBERTa-v3 Extractive QA, IWD outperforms Vanilla KD by +23.83 points on Natural Questions (28.23%→52.06%28.23\%→ 52.06\%) and +15.64 points on NewsQA (39.47%→55.11%39.47\%→ 55.11\%). On Qwen2.5, Natural Questions F1 rises by +16.36 points (40.89%→57.25%40.89\%→ 57.25\%). On NLI (HANS), IWD outperforms Vanilla KD by +9.12 points on DeBERTa-v3 (61.45%61.45\%) and +3.26 points on Qwen2.5 (54.95%54.95\%). For NER and Sentiment Analysis, IWD consistently matches or exceeds top baselines across all OOD splits (e.g., 80.38%80.38\% on OntoNotes 5.0 and 87.58%87.58\% on CR). IWD is stable across experiments. Non-augmented distillation baselines exhibit high performance variance across tasks and architectures (see Table 1). For instance, LWD achieves competitive DeBERTa-v3 NLI performance (52.99%52.99\% on HANS) but collapses on Qwen2.5 QA (2.37%2.37\% on Natural Questions). Similarly, DKD yields OOD gains on DeBERTa-v3 QA (44.26%44.26\% on NewsQA) yet lags Vanilla KD on NLI and Sentiment Analysis, while RevKL degrades on DeBERTa-v3 QA despite strong Qwen2.5 results. This variance highlights a key advantage of IWD: while standard baselines depend heavily on the task and architecture, IWD remains stable and consistently top-performing across both encoder-only and decoder-only models. This instability of baselines stems from rigid assumptions embedded within these objectives, such as target logit geometry, specific head dynamics, or a particular teacher softmax shape, that hold only for select backbone-task pairs. IWD instead evaluates teacher invariance under counterfactual perturbations. Invariance to spurious shifts is a model-agnostic property; thus, IWD operates seamlessly across both encoder-only and decoder-only students and stays in the top-performing tier on every OOD split we tested. IWD outperforms basic data augmentation. Compared to AugKD, which uses augmented inputs without variance-aware weighting, IWD achieves a 100%100\% win rate across 16 of 16 OOD settings, improving by +1.60 points on average (see Table 1). Advantages are most pronounced on decoder-only architectures like Qwen2.5, where IWD outperforms AugKD on Yelp-Polarity (+4.84 points), CR (+4.80 points), and SNLI (+3.58 points). This confirms that applying invariant weight allocation (wiw_i) to environment variants provides a clear advantage over basic data augmentation alone. Encoder vs. Decoder OOD Performance. Under IWD, encoder-only DeBERTa-v3 achieves higher OOD performance than decoder-only Qwen2.5, widening the gap observed under Vanilla KD (see Table 1). On HANS, DeBERTa’s lead over Qwen expands from +0.65 points under Vanilla KD to +6.50 points under IWD; on NewsQA, this lead reaches +34.77 points. We attribute this gap to how teacher distance is measured. Encoder-only models output probabilities over just a few task labels, making teacher disagreement clean and straightforward to measure. In contrast, decoder-only models generate text token-by-token across a vast vocabulary. Here, minor surface variations, such as alternative word choices or sentence structures, add noise to the distance metric, even when the underlying meaning remains identical. Thus, while generative students still benefit from IWD, fully unlocking their performance may require measuring invariance at the semantic level rather than the token level. IWD Performance on ID. In-distribution (ID), IWD remains fully competitive with standard baselines across tasks, achieving high accuracy on MNLI (84.66%84.66\%), CoNLL-2003 (92.64%92.64\%), and SST-2 (92.78%92.78\%) (see Table 1). However, a limitation emerges in Extractive QA: on SQuAD-v2, both AugKD and IWD exhibit a sharp ID performance drop compared to Vanilla KD (51.40%51.40\% vs. 70.64%70.64\% on DeBERTa-v3, and 48.84%48.84\% vs. 75.48%75.48\% on Qwen2.5). This ID–OOD trade-off likely stems from environment generation shifting the training distribution away from the ID baseline, highlighting a direction for future research. Task Structure and IWD Performance. IWD gains correlate with the baseline OOD degradation and shortcut susceptibility of a task. On DeBERTa-v3 Extractive QA, Vanilla KD drops over 4242 percentage points from the ID SQuAD-v2 benchmark (70.64%70.64\%) to OOD Natural Questions (28.23%28.23\%). IWD dramatically recovers this gap, yielding +23.83+23.83 points on Natural Questions (28.23%→52.06%28.23\%→ 52.06\%) and +15.64+15.64 points on NewsQA (39.47%→55.11%39.47\%→ 55.11\%). Conversely, Sentiment Analysis shows high baseline stability under Vanilla KD—dropping just 1.341.34 points from SST-2 (92.73%92.73\%) to Yelp-Polarity (91.39%91.39\%)—resulting in modest IWD gains (+0.03+0.03 to +0.29+0.29 points). This aligns with task geometry: Extractive QA and NLI heavily exploit lexical and positional shortcuts exposed by counterfactuals, whereas Sentiment Analysis and NER are inherently context-invariant. Thus, while IWD reliably improves shortcut-heavy tasks, stable tasks require precise tuning to maximize OOD generalization. Data Size MNLI-m MNLI-m HANS SNLI (N) (ID) (ID) (OOD1) (OOD2) DeBERTa-v3 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 50k 87.37 87.17 68.19 82.72 Qwen-2.5 5k 81.87 82.69 53.19 78.39 8k 82.70 83.82 54.09 79.34 10k 83.50 84.36 59.05 81.17 30k 84.98 86.04 60.27 82.15 50k 85.56 86.09 66.22 81.74 Table 2: Data scaling analysis of IWD on MNLI across varying sample sizes (N) for DeBERTa-v3/Qwen-2.5. Dataset Scaling Effect. Across datasets, performance improves as data grows from 5k to 50k samples (see Table 2). DeBERTa-v3 exhibits stronger OOD scaling than Qwen2.5: on HANS, its accuracy increases by +16.21 points (51.98%→68.19%51.98\%→ 68.19\%) versus a +13.03-point gain for Qwen2.5 (53.19%→66.22%53.19\%→ 66.22\%). On SNLI, both models scale consistently, each gaining roughly +3.00 points. Environments ID OOD1 OOD2 OOD3 OOD4 MNLI Benchmark Datasets MNLI HANS SNLI – – E=3E=3 84.34 56.27 78.98 – – E=4E=4 84.54 57.02 79.68 – – E=5E=5 84.78 58.41 80.10 – – E=6E=6 84.80 56.28 80.60 – – SST-2 Benchmark Datasets SST-2 SemEval CR Yelp FPB E=2E=2 92.09 88.72 87.23 91.39 90.75 E=3E=3 92.55 89.08 87.50 91.54 90.14 E=4E=4 92.89 89.02 87.23 91.72 90.29 E=5E=5 92.43 89.38 86.44 91.55 90.09 Table 3: Impact of varying the number of environments (K) in IWD across MNLI/SST-2 benchmarks on accuracy. Impact of Environment Count (K). Table 3 evaluates performance across environment variants (K∈3,4,5,6K∈\3,4,5,6\). Increasing K beyond 33 generally improves ID and OOD accuracy by exposing the model to broader shortcut perturbations. On MNLI, scaling from K=3K=3 to K=5K=5 yields a +2.14-point gain on HANS (56.27%→58.41%56.27\%→ 58.41\%), while SNLI accuracy rises to 80.60%80.60\% at K=6K=6. On SST-2, K=5K=5 achieves peak ID performance (92.89%92.89\%) and peak Yelp OOD accuracy (91.72%91.72\%). Overall, K=5K=5 provides the optimal balance between environmental diversity and efficiency. Ablations. We summarize the key takeaways of our ablations (complete setups in Appendix C): (1) Sensitivity to Weight Temperature (τ): We evaluate the weighting temperature τ∈[0.1,30.0]τ∈[0.1,30.0], which controls the sharpness of weighting (Table 6). While classification and NER remain highly robust across τ∈[0.3,3.0]τ∈[0.3,3.0], Extractive QA requires sharp down-weighting (τ=0.1τ=0.1) to suppress high-variance span predictions, yielding substantial OOD gains of up to +9.8+9.8 F1. (2) Sensitivity to Distance Metric Selection: We evaluate how distance metric formulations impact performance across task domains (Table 7). Aligning metric geometry with the output space maximizes performance: bounded Jensen–Shannon divergence (DJSD_JS) stabilizes classification by bounding extreme softmax tails, while structural metrics like Levenshtein distance and F1 (QA answer string drift) and Entity Set F1 (span-level NER) eliminate per-token noise to achieve top OOD performance on NQ (38.1938.19 F1) and WNUT17 (46.0446.04 F1). This proves that choosing the teacher invariance distance metric is a key structural design decision following directly from the output space. (3) Environment Variant Utilization: We compare distilling with all counterfactual variants against restricting the student to the original clean text or a single random variant (Table 8). Distilling across All Variants consistently yields the highest OOD robustness—outperforming single-variant baselines by up to 4.24.2 points on HANS—proving that full exposure to environmental perturbations is necessary for invariant feature learning. (4) Weight Normalization: We evaluate default weight normalization against an unnormalized baseline across experiments (Table 9). Empirical performance between the two is virtually identical (<0.2<0.2-point shift), but batch-mean normalization is essential as it stabilizes training by decoupling learning rates from teacher divergence scales. (5) Weight Distribution Across Tasks: We inspect per-sample weight statistics across benchmarks to analyze how IWD allocates sample importance to tasks of varying complexity (Table 10). Across all datasets, effective sample sizes remain high (≥95.0%≥ 95.0\%), confirming that IWD performs smooth, task-adaptive reweighting rather than aggressive data pruning. Discussion and Limitations In this work, we first ask: do distilled student LLMs inherit OOD robustness under existing methods? We present Invariance-Weighted Distillation (IWD), a framework that enhances OOD robustness in student LLMs without requiring architectural modifications, auxiliary classification heads, or adversarial training. By consistently outperforming KD and data augmentation baselines across OOD tasks, our results demonstrate a clear advantage over current strategies. Limitations include: (i) the computational overhead of synthetic environment generation, which can impact distillation performance; and (i) a reliance on output-level distance measures, leaving fine-grained token-level invariance as a promising unexplored research direction. Future work would look into extensions to instruction tuning (Fu et al. 2026), particularly leveraging token-level measures, as well as, broader paraphrasing techniques (Hamman et al. 2025b). It would also be interesting to study transferability of other kinds of spurious correlations (Halder et al. 2025; Egea et al. 2025) beyond OOD robustness. References Arjovsky et al. (2019) M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz Invariant risk minimization. External Links: 1907.02893, Link Cited by: Appendix A. Bowman et al. (2015) S. R. Bowman, G. Angeli, C. Potts, and C. D. Manning A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, p. 632–642. Cited by: Appendix B. Chen et al. (2023) Z. Chen, Q. Gao, A. Bosselut, A. Sabharwal, and K. Richardson DISCO: distilling counterfactuals with large language models. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Toronto, Canada, p. 5514–5528. External Links: Link, Document Cited by: Related Works. DeepSeek-AI (2025) DeepSeek-AI DeepSeek-r1: incentivizing reasoning capability in llms via reinforcement learning. External Links: 2501.12948, Link Cited by: Introduction. Derczynski et al. (2017) L. Derczynski, E. Nichols, M. van Erp, and N. Limsopatham Results of the WNUT2017 shared task on novel and emerging entity recognition. In Proceedings of the 3rd Workshop on Noisy User-generated Text, p. 140–147. Cited by: Appendix B. Dissanayake et al. (2025) P. Dissanayake, F. Hamman, B. Halder, I. Sucholutsky, Q. Zhang, and S. Dutta Quantifying knowledge distillation using partial information decomposition. In Proceedings of The 28th International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research, Vol. 258, p. 4474–4482. Cited by: Related Works. Du et al. (2023) M. Du, S. Mukherjee, Y. Cheng, M. Shokouhi, X. Hu, and A. H. Awadallah Robustness challenges in model distillation and pruning for natural language understanding. In Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics (EACL), Cited by: Related Works, Introduction. Egea et al. (2025) D. Egea, B. Halder, and S. Dutta Vision: robust and interpretable code vulnerability detection leveraging counterfactual augmentation. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society, Vol. 8, p. 812–823. Cited by: Discussion and Limitations. Fu et al. (2026) Y. Fu, F. Hamman, and S. Dutta T-shirt: token-selective hierarchical data selection for instruction tuning. Advances in Neural Information Processing Systems 38, p. 113932–113958. Cited by: Discussion and Limitations. Geirhos et al. (2020) R. Geirhos, J. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann Shortcut learning in deep neural networks. Nature Machine Intelligence 2 (11), p. 665–673. Cited by: Preliminaries. Gholami et al. (2024) M. Gholami, M. Akbari, C. Hu, V. Masrani, Z. J. Wang, and Y. Zhang GOLD: generalized knowledge distillation via out-of-distribution-guided language data generation. In Findings of the Association for Computational Linguistics: NAACL 2024, Cited by: Related Works. Girija et al. (2025) S. S. Girija, S. Kapoor, L. Arora, D. Pradhan, A. Raj, and A. Shetgaonkar Optimizing llms for resource-constrained environments: a survey of model compression techniques. In Proceedings of the 2025 IEEE 49th Annual Computers, Software, and Applications Conference (COMPSAC), p. 1272–1277. External Links: Document Cited by: Introduction. Goldblum et al. (2019) M. Goldblum, L. Fowl, S. Feizi, and T. Goldstein Adversarially robust distillation. External Links: 1905.09747, Link Cited by: Related Works. Gu et al. (2024) Y. Gu, L. Dong, F. Wei, and M. Huang MiniLLM: knowledge distillation of large language models. In International Conference on Learning Representations (ICLR), Cited by: Appendix B, Introduction, Experiments. Halder et al. (2025) B. Halder, F. Hamman, P. Dissanayake, Q. Zhang, I. Sucholutsky, and S. Dutta Towards formalizing spuriousness of biased datasets using partial information decomposition. Transactions on Machine Learning Research. External Links: ISSN 2835-8856, Link Cited by: Discussion and Limitations. Hamman et al. (2025a) F. Hamman, P. Dissanayake, Y. Fu, and S. Dutta Few-shot knowledge distillation of llms with counterfactual explanations. In Advances in Neural Information Processing Systems, Vol. 38. Cited by: Related Works, Related Works, Preliminaries. Hamman et al. (2025b) F. Hamman, C. Zhu, A. Kumar, X. Peng, S. Dutta, D. Liu, and A. Samuel Improving consistency in retrieval-augmented systems with group similarity rewards. arXiv preprint arXiv:2510.04392. Cited by: Discussion and Limitations. He et al. (2021) P. He, X. Liu, J. Gao, and W. Chen DeBERTa: decoding-enhanced bert with disentangled attention. International Conference on Learning Representations (ICLR). Cited by: Appendix B. Hendrycks and Dietterich (2019) D. Hendrycks and T. Dietterich Benchmarking neural network robustness to common corruptions and perturbations. In Proceedings of the Seventh International Conference on Learning Representations (ICLR), Cited by: Preliminaries. Hinton et al. (2015) G. Hinton, O. Vinyals, and J. Dean Distilling the knowledge in a neural network. External Links: 1503.02531, Link Cited by: Appendix A, Appendix B, Related Works, Introduction, Introduction, Experiments. Hu and Liu (2004) M. Hu and B. Liu Mining and summarizing customer reviews. In Proceedings of the Tenth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, p. 168–177. Cited by: Appendix B. Jiao et al. (2020) X. Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu TinyBERT: Distilling BERT for Natural Language Understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020, Stroudsburg, Penn., p. 4163–4174. Cited by: Introduction, Experiments. Kim and Rush (2016) Y. Kim and A. M. Rush Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing (EMNLP), p. 1317–1327. Cited by: Related Works. Kullback and Leibler (1951) S. Kullback and R. A. Leibler On information and sufficiency. The Annals of Mathematical Statistics 22 (1), p. 79–86. Cited by: Preliminaries. Kwiatkowski et al. (2019) T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, et al. Natural questions: a benchmark for study of corpus-level language understanding. Transactions of the Association for Computational Linguistics 7, p. 453–466. Cited by: Appendix B. Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), External Links: Link Cited by: Appendix B. Malo et al. (2014) P. Malo, A. Sinha, P. Korhonen, J. Wallenius, and P. Takala Good debt or bad debt: detecting semantic orientations in economic texts. Journal of the Association for Information Science and Technology 65 (4), p. 782–796. Cited by: Appendix B. McCoy et al. (2019) R. T. McCoy, E. Pavlick, and T. Linzen Right for the wrong reasons: diagnosing syntactic heuristics in natural language inference. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, p. 3428–3448. Cited by: Appendix B. OpenAI (2023) OpenAI GPT-4 technical report. External Links: 2303.08774, Link Cited by: Introduction. Popp et al. (2025) N. Popp, M. Hein, K. A. Laube, and L. Schott Improving knowledge distillation under unknown covariate shift through confidence-guided data augmentation. External Links: 2506.02294, Link Cited by: Related Works. Pradhan et al. (2013) S. Pradhan, A. Moschitti, N. Xue, H. T. Ng, A. Björkelund, O. Uryupina, Y. Zhang, and Z. Zhong Towards robust linguistic analysis using ontonotes. In Proceedings of the Seventeenth Conference on Computational Natural Language Learning, p. 143–152. Cited by: Appendix B. Rajpurkar et al. (2018) P. Rajpurkar, R. Jia, and P. Liang Know what you don’t know: unanswerable questions for SQuAD. Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics 2, p. 784–789. Cited by: Appendix B. Romero et al. (2015) A. Romero, N. Ballas, S. Ebrahimi Kahou, A. Chassang, C. Gatta, and Y. Bengio FitNets: hints for thin deep nets. In Proceedings of the Third International Conference on Learning Representations (ICLR), San Diego, California, USA. Cited by: Related Works, Introduction, Experiments. Rosenthal et al. (2017) S. Rosenthal, N. Farra, and P. Nakov SemEval-2017 task 4: sentiment analysis in twitter. In Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017), p. 502–518. Cited by: Appendix B. Shen et al. (2021) Z. Shen, J. Liu, Y. He, X. Zhang, R. Xu, H. Yu, and P. Cui Towards out-of-distribution generalization: a survey. IEEE Transactions on Knowledge and Data Engineering. Cited by: Introduction. Socher et al. (2013) R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y. Ng, and C. Potts Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, p. 1631–1642. Cited by: Appendix B. Stacey and Rei (2024) J. Stacey and M. Rei Distilling robustness into natural language inference models with domain-targeted augmentation. In Findings of the Association for Computational Linguistics: ACL 2024, Cited by: Related Works, Introduction. Stanton et al. (2021) S. Stanton, P. Izmailov, P. Kirichenko, A. A. Alemi, and A. G. Wilson Does knowledge distillation really work?. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 34, p. 6906–6919. Cited by: Introduction. Sun et al. (2019) S. Sun, Y. Cheng, Z. Gan, and J. Liu Patient knowledge distillation for bert model compression. In Empirical Methods in Natural Language Processing (EMNLP), p. 4323–4332. Cited by: Related Works, Introduction, Experiments. Sun et al. (2020) Z. Sun, H. Yu, X. Song, R. Liu, Y. Yang, and D. Zhou MobileBERT: A Compact Task-Agnostic BERT for Resource-Limited Devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (ACL 2020), Stroudsburg, Penn., p. 2158–2170. Cited by: Introduction, Experiments. Tian et al. (2020) Y. Tian, D. Krishnan, and P. Isola Contrastive Representation Distillation. In Proceedings of the Eighth International Conference on Learning Representations (ICLR-20), Addis Ababa, Ethiopia. Cited by: Related Works. Tjong Kim Sang and De Meulder (2003) E. F. Tjong Kim Sang and F. De Meulder Introduction to the CoNLL-2003 shared task: language-independent named entity recognition. In Proceedings of the Seventh Conference on Natural Language Learning at HLT-NAACL 2003, p. 142–147. Cited by: Appendix B. Trischler et al. (2017) A. Trischler, T. Wang, X. Yuan, J. Howland, A. Sordoni, P. Harmar, C. Offermann, and K. Suleman NewsQA: a machine comprehension dataset. In Proceedings of the 2nd Workshop on Representation Learning for NLP, p. 191–200. Cited by: Appendix B. Weischedel et al. (2013) R. Weischedel, M. Palmer, M. Marcus, E. Hovy, S. Pradhan, L. Ramshaw, N. Xue, A. Taylor, J. Kaufman, M. Franchini, et al. OntoNotes release 5.0. Linguistic Data Consortium 23. Cited by: Appendix B. Williams et al. (2018) A. Williams, N. Nangia, and S. R. Bowman A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol. 1, p. 1112–1122. Cited by: Appendix B. Xu et al. (2024) X. Xu, M. Li, C. Tao, T. Shen, R. Cheng, J. Li, C. Xu, D. Tao, and T. Zhou A survey on knowledge distillation of large language models. External Links: 2402.13116, Link Cited by: Related Works, Introduction, Preliminaries. Yang et al. (2024) A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, et al. Qwen2.5 technical report. External Links: 2412.15115, Link Cited by: Appendix B, Introduction. Zagoruyko and Komodakis (2017) S. Zagoruyko and N. Komodakis Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer. In Proceedings of the Fifth International Conference on Learning Representations (ICLR-17), Toulon, France. Cited by: Related Works. Zhang et al. (2025) S. Zhang, Y. Luo, Z. Lyu, and X. Chen ShiftKD: benchmarking knowledge distillation under distribution shift. Neural Networks 188, p. 107504. External Links: ISSN 0893-6080, Document Cited by: Related Works. Zhang et al. (2015) X. Zhang, J. Zhao, and Y. LeCun Character-level convolutional networks for text classification. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 28. Cited by: Appendix B. Zhao et al. (2022) B. Zhao, Q. Cui, R. Song, Y. Qiu, and J. Liang Decoupled knowledge distillation. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 11953–11962. Cited by: Appendix B, Experiments. Zhou et al. (2023) A. Zhou, J. Wang, Y. Wang, and H. Wang Distilling out-of-distribution robustness from vision-language foundation models. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), Cited by: Related Works. Zhu et al. (2024) X. Zhu, J. Li, Y. Liu, C. Ma, and W. Wang A survey on model compression for large language models. Transactions of the Association for Computational Linguistics 12, p. 1556–1577. External Links: Document Cited by: Related Works. Appendix Appendix A A. Theoretical Background and Proof of Theorem 1 A.1. Background: Invariance across Environments Standard empirical risk minimization treats all training data as a single pool, allowing models to exploit spurious correlations that change or disappear at test time. To address this, frameworks like Invariant Risk Minimization (IRM) (Arjovsky et al. 2019) conceptually distinguish between causal features (which yield a stable, invariant predictive relationship across environments) and spurious features (whose importance fluctuates depending on the environment). Crucially, standard IRM frameworks operate in a traditional supervised learning setting, enforcing invariance globally at the model level (e.g., via regularisation penalties across macroscopic environment datasets). In contrast, our framework adapts this philosophy specifically for knowledge distillation, operationalising invariance entirely at the per-sample level. Rather than constraining the student’s representations directly, we posit that a teacher model’s prediction on a single example across counterfactual environments can be decomposed into two parts: an environment-invariant base prediction, and an environment-dependent (spurious) residual. This sample-level view of teacher invariance provides the necessary foundation for the gradient decomposition we derive next. A.2.Derivation of Definition 2: Exact Gradient Decomposition under the KL Objective We derive the causal and spurious gradient components giCg_i^C and giSg_i^S of Definition 2 directly from the KL distillation objective used in Algorithm 1 (Line 18). Although the softmax is non-linear in the logits, the gradient of the KL loss with respect to the student’s logits is exactly linear in the teacher’s output probabilities. Since Definition 2 places the IRM decomposition on the teacher’s output distribution (7), the causal/spurious decoupling below is exact: no linearisation or Taylor approximation is required. Distillation loss. For a single training example i observed in environment eie_i, the temperature-scaled KL distillation loss is ℒi=τKD2KL(pT(xi(ei))∥pS(xi)).L_i\;=\; _KD^2\,KL\! (p_T\! (x_i^(e_i) )\, \|\,p_S(x_i) ). (9) The KL gradient is linear in the teacher target. Writing KL(pT∥pS)=∑cpT,clogpT,c−∑cpT,clogpS,cKL(p_T\|p_S)= _cp_T,c p_T,c- _cp_T,c p_S,c, the first (entropy) term is constant in θS _S, so minimising ℒiL_i is equivalent to minimising the cross-entropy H(pT,pS)H (p_T,p_S ). Using the standard softmax–cross-entropy derivative, ∂[−∑c′pT,c′logpS,c′]/∂zS,c=1τKD(pS,c−pT,c)∂ [-\! _c p_T,c p_S,c ]/∂ z_S,c= 1 _KD (p_S,c-p_T,c ), we obtain ∇zSℒi=τKD(pS(xi)−pT(xi(ei))), _z_SL_i\;=\; _KD (p_S(x_i)-p_T\! (x_i^(e_i) ) ), (10) the classical distillation gradient of Hinton et al. (2015); the scalar τKD _KD is absorbed into the learning rate below. The essential observation is that (10) is affine in the teacher’s probability vector pTp_T. The softmax non-linearity acts on the teacher’s logits before the loss is formed; once the IRM decomposition is stated at the level of the teacher’s output distribution, the gradient inherits the decomposition additively. Chain rule. With Ji=∇θSzS(xi)J_i= _ _Sz_S(x_i) the Jacobian of the student’s logits, ∇θSℒi=(pS(xi)−pT(xi(ei)))⊤Ji. _ _SL_i\;=\; (p_S(x_i)-p_T\! (x_i^(e_i) ) )^\! J_i. (11) As in the MSE case, JiJ_i is a Jacobian with respect to the student’s parameters, encoding how each weight affects each output logit. Substituting the teacher decomposition. Substituting pT(xi(ei))=πi+si(ei)p_T (x_i^(e_i) )= _i+s_i^(e_i) from (7) into (11), the error term splits additively: ∇θSℒi=(pS(xi)−πi)⊤Ji⏟causal gradient (pre-flip)−si(ei)⊤Ji⏟spurious gradient (pre-flip). _ _SL_i\;=\; (p_S(x_i)- _i )^\! J_i_causal gradient (pre-flip)\;-\; s_i^(e_i) J_i_spurious gradient (pre-flip). (12) This step is exact and uses only the linearity of (11) in pTp_T. Gradient descent and sign convention. Under gradient descent with step size α, substituting (12) into θS←θS−α∇θSℒi _S← _S-α _ _SL_i and collecting the terms added to θS _S: θS←θS+α(πi−pS(xi))⊤Ji⏟giC+αsi(ei)⊤Ji⏟giS. _S\;←\; _S+α\, ( _i-p_S(x_i) )^\! J_i_g_i^C+α\, s_i^(e_i) J_i_g_i^S. (13) The subtraction in gradient descent flips (pS−πi) (p_S- _i ) to (πi−pS) ( _i-p_S ), so giCg_i^C points the student toward the teacher’s environment-invariant output distribution πi _i; giSg_i^S needs no flip because its minus sign in (12) is cancelled by the descent subtraction. Under (A3), ‖giS‖2=Φ‖si(ei)‖2\|g_i^S\|^2= \,\|s_i^(e_i)\|^2 and ‖giC‖2=Φci2\|g_i^C\|^2= \,c_i^2 with ci2:=‖πi−pS(xi)‖2c_i^2:=\| _i-p_S(x_i)\|^2, exactly as required in Step 1 of the proof of Theorem 1. Per-example spurious reliance σi2 _i^2. Averaging over environments, e[‖si(e)‖2]=σi2≥0E_e [\|s_i^(e)\|^2 ]= _i^2≥ 0 defines the per-example teacher spurious reliance. σi2=0 _i^2=0 implies pT(xi(e))=πip_T (x_i^(e) )= _i for all e: the teacher’s output distribution for example i is completely invariant across environments. Remark 1 (MSE distillation as a special case). If the distillation loss is instead the squared error ‖fθS(xi)−fθT(xi(ei))‖2\|f_ _S(x_i)-f_ _T(x_i^(e_i))\|^2 on raw outputs, the identical derivation applies with pT,pS,zSp_T,p_S,z_S replaced by fθT,fθS,fθSf_ _T,f_ _S,f_ _S and πi _i by μi=e[fθT(xi(e))] _i=E_e[f_ _T(x_i^(e))]: the MSE gradient 2(fθS−μi−si(ei))⊤∇θSfθS2 (f_ _S- _i-s_i^(e_i) )^\! _ _Sf_ _S is likewise affine in the teacher target. Both objectives therefore induce the same causal/spurious gradient structure, and Theorem 1 applies to either. Remark 2 (Probability-space vs. logit-space invariance). Definition 2 states invariance in the space of output distributions: the spurious component is the zero-mean deviation of pT(xi(e))p_T (x_i^(e) ) around its environment average πi _i. Because the softmax is non-linear, zero-mean spurious variation in probability space does not coincide with zero-mean variation in logit space (Jensen’s inequality); the two decompositions define distinct but equally valid notions of invariance. We adopt the probability-space formulation throughout, for two reasons: (i) it is the space in which the KL distillation gradient is linear in the teacher target, making the decomposition in (12) exact; and (i) it is the natural formalisation of prediction invariance — a teacher is invariant on example i precisely when its predictive distribution does not change across environments. A.3.Discussion of Assumptions in Theorem 1 Theorem 1 rests on three assumptions. We discuss their meaning, when they are expected to hold, and what would happen if they were violated. (A1) Independent spurious components with zero mean. We assume si(ei)s_i^(e_i) are independent across training samples, with e[si(e)]=0E_e[s_i^(e)]=0 and e[‖si(e)‖2]=σi2E_e[\|s_i^(e)\|^2]= _i^2. The zero-mean condition is the defining property of spurious features under IRM: any signal that has a consistent direction across all environments would, by Definition 1, be part of the causal component μi _i rather than the spurious component si(e)s_i^(e). Independence across samples holds when counterfactual environments are generated independently per sample, as in our paraphrase-based construction. (A2) Independence of spurious and causal residuals: σi2⟂ci2 _i^2 c_i^2. This is the key structural assumption. σi2 _i^2 measures how sensitive the teacher is to environmental surface variation for sample i. ci2=‖μi−fθS(xi)‖2c_i^2=\| _i-f_ _S(x_i)\|^2 measures how far the student currently is from the teacher’s causal prediction for the same sample. (A2) requires these two quantities to be independent across training samples. The assumption holds when spurious correlations arise from environmental surface confounders — such as lexical register, syntactic style, or demographic markers in NLI — that affect the teacher’s cross-environment invariance but are unrelated to the semantic difficulty of the inference problem (which drives ci2c_i^2). It would be violated if, for instance, semantically harder samples systematically attracted more spurious surface features, creating a positive correlation between σi2 _i^2 and ci2c_i^2. This can be tested empirically by computing the Pearson correlation between DiD_i (our proxy for σi2 _i^2) and the per-sample training loss of the student; a correlation near zero supports (A2). (A3) Homogeneous student gradient norms. We assume ‖∇θSfθS(xi)‖2=Φ\| _ _Sf_ _S(x_i)\|^2= is constant across all training samples. This simplifies the S2C ratio to a ratio of scalar weighted sums (Step 1 of the proof), making the Chebyshev argument clean. The assumption can be relaxed: if gradient norms Φi _i are allowed to vary but are additionally assumed independent of σi2 _i^2 (i.e. Φi⟂σi2 _i _i^2), then the Chebyshev step applies to the products Φiσi2 _i _i^2 and the conclusion of Theorem 1 still holds. We adopt the simpler homogeneous form here for exposition. A.4.Proposition 1 and Proof. In the main text, we substitute the theoretical spurious variance σi2 _i^2 with the empirical distance DiD_i computed over K counterfactual environments. This substitution is mathematically justified under Assumption (A1), as we formalise below. Proposition 1 (Convergence of Empirical Invariance Distance). Let Di=2K(K−1)∑k<k′‖pT(xi(k))−pT(xi(k′))‖2D_i= 2K(K-1) _k<k \|p_T (x_i^(k) )-p_T (x_i^(k ) ) \|^2 be the empirical average pairwise distance between the teacher’s softened output distributions across K independently sampled counterfactual environments. Under Assumption (A1), as K→∞K→∞, DiD_i converges in probability to 2σi22 _i^2. Proof of Proposition 1. Substitute the IRM decomposition pT(xi(k))−pT(xi(k′))=(πi+si(k))−(πi+si(k′))=si(k)−si(k′)p_T(x_i^(k))-p_T(x_i^(k ))=( _i+s_i^(k))-( _i+s_i^(k ))=s_i^(k)-s_i^(k ) into the distance metric. The environment-invariant causal component μi _i cancels out: Expanding the squared L2L_2 norm gives: ‖si(k)−si(k′)‖2=‖si(k)‖2+‖si(k′)‖2−2⟨si(k),si(k′)⟩.\|s_i^(k)-s_i^(k )\|^2=\|s_i^(k)\|^2+\|s_i^(k )\|^2-2 s_i^(k),s_i^(k ) . (14) Taking the expectation over the independent environments k and k′k , and applying Assumption (A1) (which states [si(k)]=0E[s_i^(k)]=0 and [‖si(k)‖2]=σi2E[\|s_i^(k)\|^2]= _i^2): [‖si(k)−si(k′)‖2] [\|s_i^(k)-s_i^(k )\|^2 ] =[‖si(k)‖2]+[‖si(k′)‖2] =E [\|s_i^(k)\|^2 ]+E [\|s_i^(k )\|^2 ] −2⟨[si(k)],[si(k′)]⟩ -2 [s_i^(k) ],E [s_i^(k ) ] =σi2+σi2−2⟨0,0⟩ = _i^2+ _i^2-2 0,0 =2σi2. =2 _i^2. (15) Because DiD_i is a U-statistic estimating this expected pairwise distance across K independent samples, by the Law of Large Numbers for U-statistics, DiD_i converges in probability to its expectation as K→∞K→∞: Di→2σi2.D_i p2 _i^2. (16) Since f(x)=2xf(x)=2x is a strictly increasing function for x≥0x≥ 0, DiD_i serves as a consistent, strictly increasing proxy for the spurious variance σi2 _i^2, justifying its use in the IWD weighting function. ∎ A.5. Complete Proof of Theorem 1 We prove the four steps in full. Proof of Theorem 1. Step 1: Reduce ρ2(W)ρ^2(W) to a ratio of weighted sums. Under (A3), ‖giS‖2 \|g_i^S\|^2 =si(ei)⊤JiJi⊤si(ei)=Φi‖si(ei)‖2, =s_i^(e_i) J_iJ_i s_i^(e_i)= _i\,\|s_i^(e_i)\|^2, ‖giC‖2 \|g_i^C\|^2 =Φici2, = _i\,c_i^2, (17) so that, taking e[‖si(e)‖2]=σi2E_e[\|s_i^(e)\|^2]= _i^2 from (A1), ρ2(W)=∑iWiΦiσi2∑iWiΦici2.ρ^2(W)= _iW_i\, _i\, _i^2 _iW_i\, _i\,c_i^2. (18) Step 2: Bound the numerator via the Chebyshev covariance inequality. Since WiIWDW_i^IWD is a deterministic decreasing function of σi2 _i^2 and Φi⟂σi2 _i _i^2 by (A3), for large N the Law of Large Numbers gives ∑iWiIWDΦiσi2 _iW_i^IWD\, _i\, _i^2 →[Φ]⋅[a(σ2)σ2], \; p\;E [ ]·E [a(σ^2)\,σ^2 ], wherea(σ2) a(σ^2) :=limN→∞NWIWD(σ2). := _N→∞NW^IWD(σ^2). (19) where the factorisation uses Φi⟂σi2 _i _i^2. Because a(⋅)a(·) is strictly decreasing, the Chebyshev (covariance) inequality gives [a(σ2)σ2]≤[a(σ2)][σ2]=[σ2]E[a(σ^2)σ^2] [a(σ^2)]\,E[σ^2]=E[σ^2] (using [a(σ2)]=1E[a(σ^2)]=1 from normalisation). Hence the IWD numerator is asymptotically bounded by the uniform numerator [Φ][σ2]E[ ]\,E[σ^2] (which also uses Φi⟂σi2 _i _i^2). Step 3: Preserve the denominator. By (A2), σi2⟂(ci2,Φi) _i^2 (c_i^2, _i), so the weights are independent of the products Φici2 _ic_i^2 and ∑iWiIWDΦici2→[Φc2]←∑iWiunifΦici2. _iW_i^IWD\, _i\,c_i^2\; p\;E [ \,c^2 ]\; p\; _iW_i^unif\, _i\,c_i^2. (20) 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 Table 4: Evaluation datasets. ID sets share the training distribution; OOD sets target specific distribution shifts. Metrics follow standard task practice. Step 4: Combine to obtain the ratio bound. From Step 3, for large N, the denominators converge to the same constant [Φc2]E[ \,c^2]. Substituting this into the S2C ratio from (18) and applying the asymptotic numerator bound from Step 2: ρ2(WIWD)→[Φ][a(σ2)σ2][Φc2]≤[Φ][σ2][Φc2]←ρ2(Wunif).ρ^2\! (W^IWD ) p E[ ]\,E[a(σ^2)σ^2]E[ \,c^2]\;≤\; E[ ]\,E[σ^2]E[ \,c^2] pρ^2\! (W^unif ). (21) Equality condition. Equality in the Chebyshev covariance inequality holds if and only if σi2 _i^2 is constant across all training samples (i.e. σi2=σ¯2 _i^2= σ^2 for all i), in which case a(σ2)≡1a(σ^2)≡ 1. In that case, both weighting schemes are asymptotically equivalent and ρ2(WIWD)→ρ2(Wunif)ρ^2(W^IWD) pρ^2(W^unif). ∎ Remark (Homogeneous case). If Φi≡Φ _i≡ for all i, everything above reduces to the original finite-N argument (the Chebyshev sum inequality applies, so no Law of Large Numbers is needed for the numerator bound). We present this homogeneous case as a special case in Appendix A.3. Appendix B B. Experimental Details B.1.Dataset Details Natural-language inference (MNLI). We train on the MultiNLI matched split (Williams et al. 2018) (393k premise–hypothesis pairs, three-way labels) and evaluate ID on its matched validation set. For OOD evaluation, we use HANS (McCoy et al. 2019) (30k samples) to test reliance on syntactic shortcuts (lexical overlap, subsequence, and constituent heuristics), and SNLI (Bowman et al. 2015) (570k image caption samples) to evaluate generalization across a linguistic register shift. Sentiment Classification (SST-2). We train on SST-2 training subsets (Socher et al. 2013) and evaluate ID on the development set. OOD evaluation uses four binary sentiment benchmarks targeting distinct domain and register shifts: SemEval-2017 Task 4A (Rosenthal et al. 2017) (social media), Customer Reviews (CR) (Hu and Liu 2004) (consumer products), Yelp Polarity (Zhang et al. 2015) (business reviews), and Financial PhraseBank (FPB) (Malo et al. 2014) (financial news). 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.” social_media_register “tbh the film was absolutely wonderful [music] #mustwatch” domain_transposition “The product was absolutely excellent.” NER Original [Barack Obama]PER_PER visited [Paris]LOC_LOC . lowercase_transform [barack obama]PER_PER visited [paris]LOC_LOC . entity_substitution_same_type [Amina Watanabe]PER_PER visited [Zurich]LOC_LOC . social_media_register omg [barack obama]PER_PER pulled up to [paris]LOC_LOC [music] #travel 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. Extractive Question Answering (SQuAD v2). We train on SQuAD v2 subsets (Rajpurkar et al. 2018) (including unanswerable questions) and evaluate ID on the standard development split. OOD performance is evaluated on NewsQA (Trischler et al. 2017) for a news-domain shift, and Natural Questions (NQ) (Kwiatkowski et al. 2019) to test generalization to real Google search queries paired with Wikipedia articles. Named-entity recognition (CoNLL-2003). Models are trained on English CoNLL-2003 (Tjong Kim Sang and De Meulder 2003) (4 entity types: PER, ORG, LOC, MISC) and evaluated ID on the test set. OOD evaluation uses WNUT-17 (Derczynski et al. 2017) (register and entity-distribution shifts) and OntoNotes 5.0 (Weischedel et al. 2013; Pradhan et al. 2013) (genre and label-schema shifts, projecting its 18-type schema down to PER/ORG/LOC). All setups report entity-level micro-F1. B.2.Model Details DeBERTa family. The student is deberta-v3-xsmall (He et al. 2021) (22M parameters). Teachers are task-specific deberta-v3-small (44M parameters) models fine-tuned on each task’s full training set and selected by ID validation performance. We use the sequence-classification head for MNLI and SST-2, the question-answering head for SQuAD v2, and the token-classification head for NER. Student and teacher share the DeBERTa-v3 tokenizer with a maximum sequence length of 256 tokens for classification and NER, and 384 for SQuAD v2. Qwen family. The student is Qwen2.5-0.5B (Yang et al. 2024) and the teacher is Qwen2.5-1.5B. We use the sequence-classification head for MNLI and SST-2, the question-answering head for SQuAD v2, and the token-classification head for NER. Student and teacher share the Qwen2.5 tokenizer with a maximum sequence length of 256 tokens for classification and NER, and 384 for SQuAD v2. B.3. Baseline Implementation Details All baselines use the same student architecture, tokenizer, training schedule, and hardware as inv_kd. Only the distillation objective differs. Vanilla KD (Hinton et al. 2015). The student minimizes a weighted sum of the ground-truth cross-entropy loss and the forward KL divergence between temperature-scaled teacher and student outputs: ℒKD=(1−α)ℒCE+αT2KL(σ(zT/T)∥σ(zS/T)),L_KD=(1-α)\,L_CE+α\,T^2\,KL\! (σ(z_T/T)\,\|\,σ(z_S/T) ), (22) where zTz_T and zSz_S are teacher and student logits, T is the temperature, and α is the KD weight. For SQuAD v2, the KL is computed separately for start and end logits. Training uses the original (non-augmented) data. RevKL (Gu et al. 2024). We replace the forward KL with reverse KL, KL(pS∥pT)KL(p_S\|p_T), adapting the objective to classification, extractive QA, and token classification. Training uses the original (non-augmented) data. Task / Dataset τ=0.1τ=0.1 τ=0.3τ=0.3 τ=1.0τ=1.0 τ=3.0τ=3.0 τ=10.0τ=10.0 τ=30.0τ=30.0 NLI (DeBERTa-v3, 5k) MNLI (ID) 81.32±0.2781.32± 0.27 82.30±0.3382.30± 0.33 82.58±0.3182.58± 0.31 82.63±0.4282.63± 0.42 82.57±0.4882.57± 0.48 82.55±0.5082.55± 0.50 HANS (OOD-1) 52.93±1.7452.93± 1.74 53.56±0.6053.56± 0.60 53.36±1.3553.36± 1.35 54.03±2.2854.03± 2.28 53.98±1.7853.98± 1.78 53.91±1.6153.91± 1.61 SNLI (OOD-2) 76.39±0.8776.39± 0.87 78.21±0.5178.21± 0.51 78.61±0.2978.61± 0.29 78.52±0.5578.52± 0.55 78.76±0.0878.76± 0.08 78.84±0.2078.84± 0.20 Sentiment Analysis (DeBERTa-v3, 5k) SST-2 (ID) 92.20±0.2392.20± 0.23 92.13±0.3592.13± 0.35 92.28±0.2492.28± 0.24 92.05±0.3392.05± 0.33 92.09±0.2392.09± 0.23 92.13±0.0792.13± 0.07 SemEval (OOD-1) 88.16±0.3688.16± 0.36 88.48±0.2188.48± 0.21 88.45±0.4688.45± 0.46 88.63±0.3288.63± 0.32 88.62±0.4588.62± 0.45 88.61±0.5188.61± 0.51 CR (OOD-2) 88.21±0.9388.21± 0.93 88.30±0.9688.30± 0.96 88.21±0.7788.21± 0.77 87.77±0.9687.77± 0.96 87.85±0.9387.85± 0.93 87.94±1.3487.94± 1.34 Yelp-Polarity (OOD-3) 91.60±0.5691.60± 0.56 91.67±0.4391.67± 0.43 91.70±0.4991.70± 0.49 91.43±0.5391.43± 0.53 91.43±0.5291.43± 0.52 91.51±0.5591.51± 0.55 FPB (OOD-4) 92.65±0.1692.65± 0.16 93.05±0.6693.05± 0.66 92.78±0.6792.78± 0.67 92.68±0.7592.68± 0.75 93.24±0.4393.24± 0.43 92.93±0.3692.93± 0.36 Extractive QA (DeBERTa-v3, 5k) SQuAD-v2 (ID) 47.64±2.7247.64± 2.72 47.14±2.6047.14± 2.60 47.51±1.2747.51± 1.27 46.75±0.7546.75± 0.75 45.80±0.8745.80± 0.87 45.36±0.7445.36± 0.74 NewsQA (OOD-1) 29.20±2.0029.20± 2.00 27.56±1.9727.56± 1.97 24.82±0.4124.82± 0.41 23.00±1.6123.00± 1.61 22.95±0.9122.95± 0.91 21.57±1.0321.57± 1.03 Natural Questions (OOD-2) 45.32±1.4745.32± 1.47 42.99±2.6842.99± 2.68 35.87±0.2935.87± 0.29 33.37±2.1933.37± 2.19 38.32±1.8738.32± 1.87 35.49±2.4735.49± 2.47 NER (DeBERTa-v3, 5k) CoNLL-2003 (ID) 91.42±0.0891.42± 0.08 91.72±0.2391.72± 0.23 91.62±0.2291.62± 0.22 91.50±0.4991.50± 0.49 91.06±0.2491.06± 0.24 91.05±0.2591.05± 0.25 WNUT-17 (OOD-1) 45.54±0.6545.54± 0.65 45.89±0.1845.89± 0.18 46.04±1.1146.04± 1.11 45.26±2.1945.26± 2.19 43.21±0.8443.21± 0.84 43.25±0.7643.25± 0.76 OntoNotes 5.0 (OOD-2) 78.51±0.2578.51± 0.25 78.60±0.2678.60± 0.26 78.56±0.3078.56± 0.30 78.13±0.9378.13± 0.93 77.01±0.6477.01± 0.64 77.00±0.5877.00± 0.58 Table 6: Sensitivity analysis evaluating the impact of the distillation temperature parameter τ on IWD across four task families using DeBERTa-v3 (trained on 5k subsamples). DKD (Zhao et al. 2022). Decoupled KD separates the KD loss into target-class (TCKD) and non-target-class (NCKD) terms: ℒDKD=βtℒTCKD+βntℒNCKD.L_DKD= _t\,L_TCKD+ _nt\,L_NCKD. (23) We use the original settings (βt=1.0 _t=1.0, βnt=1.0 _nt=1.0) and adapt the objective to all tasks. LWD. Layer-Wise Distillation augments output-level KD with MSE matching between corresponding hidden representations. For DeBERTa, we align the first, middle, and last encoder layers; for Qwen, the first, third, and last decoder layers. AugKD. Vanilla KD trained on the same augmented dataset as IWD, treating all K paraphrases as independent training samples with equal weight. This isolates the effect of invariance weighting from data augmentation alone. B.4. Training and Hyperparameter Details Optimiser and schedule. All students use AdamW (Loshchilov and Hutter 2019) (β1=0.9 _1=0.9, β2=0.999 _2=0.999, ϵ=10−8ε=10^-8, weight decay 0.00.0) with 200 warmup steps and linear decay. DeBERTa students train for 3 epochs (5 for NER) with learning rate ×10−57\!×\!10^-5, batch size 32, and fp16 precision. Qwen students train for 3 epochs (5 for NER) with learning rate ×10−52\!×\!10^-5, batch size 8, gradient accumulation of 2 (effective batch size 16), gradient checkpointing, and bf16 precision. KD hyperparameters. All methods use temperature T=2.0T=2.0. The KD weight α is 5 for DeBERTa and 1 for Qwen, selected from 0.5,1,2,5\0.5,1,2,5\ on MNLI ID validation and fixed for all tasks. IWD hyperparameters. We use τ=1.0τ=1.0 with batch-mean weight normalization (∑i∈ℬw~i=|ℬ| _i w_i=|B|). The disagreement metric is Jensen–Shannon divergence for MNLI and SST-2, one-minus SQuAD-F1 for SQuAD v2, and one-minus entity-set F1 for NER. The results in the major result table 1 for NLI, NER and Extractive QA tasks were acheived by these hyperparameters while for better results in sentiment analysis task we had to use τ=0.1τ=0.1 with SymKL as the distance measure of teacher invariance. Per-task disagreement metric DiD_i. DiD_i averages pairwise disagreement over the (K2) K2 variant pairs. For MNLI and SST-2, we compute Jensen–Shannon divergence over teacher class probabilities (3-class and 2-class softmax, respectively). For SQuAD v2 and NER, we use one-minus F1 between teacher-decoded answer strings and entity mention sets, respectively. Counterfactual data generation. Variants are generated offline using GPT-4.1-mini with prompts preserving task labels (sentiment, entailment, answer spans, or entity annotations) while varying superficial features such as register, syntax, and vocabulary. For NER, entity surface forms are preserved while context is paraphrased. Each sample has K∈[2,6]K∈[2,6] variants (mean ≈5≈ 5); samples with K<2K<2 or label-inconsistent variants are removed. Hardware and seeds. Experiments use 2 NVIDIA RTX A4500 GPUs (20 GB each). The Qwen SQuAD v2 teacher uses batch size 4 with gradient accumulation of 8. Each configuration is evaluated over five seeds 42,43,44,45,46\42,43,44,45,46\, reporting mean and standard deviation. B.5. Counterfactual Environment Types This subsection lists the paraphrase environments used to generate counterfactual variants. Each environment targets a distinct spurious surface feature while preserving labels. All environments are applied to training samples, and label-inconsistent variants are discarded. The number of variants varies across experiments. Table 5 shows several examples of environments data generation. MNLI (7 environments). We generate variants by perturbing lexical overlap (high/low_overlap), negation patterns (negation), irrelevant context (distractor), syntactic heuristics (syntactic_shortcut_break), length bias (length_decorrelation), and voice structure (voice_alternation) while preserving labels. SST-2 (8 environments). Variants perturb sentiment-related shortcuts through negation invariance, lexical overlap (lexical_paraphrase_low_overlap), domain and register shifts (domain_transposition, social_media_register), sentiment cue removal, intensifier/hedge changes, and comparative reframing while preserving polarity. SQuAD v2 (6 environments). Variants target QA shortcuts through distractor insertion, question paraphrasing, answer-position shifts, decoy answer spans, entity substitution, and voice alternation while preserving the gold answer. NER (7 environments). Variants include 3 programmatic transforms (lowercasing, typographical noise, and punctuation perturbation) and four LLM-based transforms targeting entity substitution, emerging entities, social-media register, and syntax variation. Entity annotations are preserved for all. Distance Metric ID OOD 1 OOD 2 MNLI Benchmark MNLI-m HANS SNLI Jensen–Shannon (DJSD_JS) 82.60 ± 0.28 53.35 ± 1.38 78.68 ± 0.34 Symmetric KL (DSymKLD_SymKL) 81.38 ± 0.40 53.01 ± 2.12 77.43 ± 0.74 L2L_2 Distance 82.38 ± 0.27 52.76 ± 1.03 78.58 ± 0.33 SST-2 Sentiment Benchmark SST-2 CR Yelp Jensen–Shannon (DJSD_JS) 92.28 ± 0.24 88.21 ± 0.77 91.70 ± 0.49 Symmetric KL (DSymKLD_SymKL) 92.01 ± 0.54 87.85 ± 1.63 91.73 ± 0.44 L2L_2 Distance 92.01 ± 0.37 88.30 ± 1.48 91.55 ± 0.55 Extractive QA Benchmark SQuAD-v2 NewsQA NQ Span F1 26.60 ± 0.70 25.21 ± 0.22 37.35 ± 0.96 Jaccard Distance 27.03 ± 1.13 25.36 ± 0.23 37.78 ± 0.29 Levenshtein Distance 27.10 ± 1.51 24.65 ± 0.23 38.19 ± 1.38 NER Benchmark CoNLL-03 WNUT17 OntoNotes Entity Set F1 91.62 ± 0.22 46.04 ± 1.11 78.56 ± 0.30 Entity Set Jaccard 91.62 ± 0.33 45.80 ± 0.57 78.46 ± 0.37 Entity Flip Rate 91.62 ± 0.33 45.80 ± 0.57 78.46 ± 0.37 Table 7: Sensitivity of IWD to distance metric formulations across tasks on DeBERTa-v3 (N=5kN=5k). Variant Strategy MNLI-M MNLI-M HANS SNLI & Metric (ID) (ID) (OOD1) (OOD2) JS Distance (Primary Metric) 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 All Variants 83.311 84.121 54.420 78.784 Original Variant 82.170 83.197 52.076 75.706 Random Variant 82.292 82.994 52.643 77.027 L2L_2 Distance All Variants 84.014 84.875 55.036 79.617 Original Variant 82.934 83.868 55.626 77.037 Random Variant 82.638 83.390 52.923 78.195 Table 8: Ablation study evaluating environment variant selection strategies during student distillation across different distance metrics using DeBERTa-v3 on ID and OOD benchmarks. Appendix C C. Ablation Experiments and Results Sensitivity to Weight Temperature (τ). Table 6 evaluates the weighting temperature τ∈0.1,0.3,1.0,3.0,10.0,30.0τ∈\0.1,0.3,1.0,3.0,10.0,30.0\, which dictates the sharpness of sample re-weighting relative to teacher divergence (DiD_i). Classification tasks (MNLI, SST-2) and NER are robust across τ∈[0.3,3.0]τ∈[0.3,3.0], with ID and mean OOD metrics varying by less than 0.70.7 and 1.51.5 points, respectively. In contrast, extractive QA (SQuAD v2) is highly sensitive due to high variance in teacher answer-span fluctuations, benefiting dramatically from aggressive downweighting (τ=0.1τ=0.1). Lowering τ from 30.030.0 to 0.10.1 yields substantial OOD gains of +7.6+7.6 F1 on NewsQA (29.2029.20) and +9.8+9.8 F1 on Natural Questions (45.3245.32). Thus, while our default setting of τ=1.0τ=1.0 is optimal for classification and sequence tagging, tasks with high divergence variance like QA benefit from smaller temperatures (τ∈0.1,0.3τ∈\0.1,0.3\). Sensitivity to Distance Metric Selection. The optimal teacher invariance distance metric depends directly on the structure of the output space. For classification tasks (MNLI, SST-2), bounded and symmetric Jensen–Shannon divergence (DJSD_JS) achieves peak performance by preventing extreme target probability tails from overly skewing sample weights compared to Symmetric KL or L2L_2 distance. For structured prediction, metrics directly targeting output boundary shifts yield superior signals: Levenshtein distance performs best for Extractive QA (38.1938.19 F1 on NQ), while Entity Set F1 achieves top accuracy for NER (46.0446.04 F1 on WNUT17). Aligning the metric to the output modality—probability divergences for classification and structural metrics for span extraction—maximizes IWD’s ability to isolate non-causal prediction variance. Environment Variant Utilization. Table 8 compares three student training presentation strategies using teacher divergence weights (DiD_i) computed over the full set of variants: (1) All Variants (distilling on every generated environment state), (2) Original Variant (distilling on clean inputs only), and (3) Random Variant (selecting a single random variant per sample). Distilling across All Variants consistently achieves top performance across all benchmarks. Restricting student inputs to the Original or Random variant causes substantial drops (e.g., losing up to 4.24.2 percentage points on HANS). Evaluating across multiple distance measures (JSJS, Symmetric KLKL, L2L_2) confirms this trend is metric-agnostic, with JS Distance providing the most effective weighting signal overall. Weight Normalisation. We compare default batch-mean normalisation (w~i=|ℬ|⋅wi/∑j∈ℬwj w_i=|B|· w_i/ _j w_j) against unnormalised weighting across MNLI and SST-2 (Table 9). Performance is effectively equivalent: on MNLI, ID accuracy shifts by <0.17<0.17 points and mean OOD performance differs by ≤0.07≤ 0.07 points across DeBERTa-v3 and Qwen-2.5. On SST-2, unnormalised weighting shows minor average OOD gains (+0.63+0.63 on DeBERTa-v3, +0.41+0.41 on Qwen-2.5) but fluctuates across individual datasets (e.g., +1.32+1.32 on FPB vs. −0.22-0.22 on SemEval-2017 for DeBERTa-v3), reflecting single-seed variance. We retain batch-mean normalisation because it: (i) decouples the effective learning rate from the absolute scale of DiD_i, stabilizing τ tuning across tasks; and (i) strictly preserves the per-sample relative importance weights required by our theoretical bound. 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 Δ -0.03 -0.18 +0.25 Qwen-2.5 Batch-Mean 81.87 53.19 78.39 None 81.70 53.20 78.23 Δ -0.17 +0.02 -0.16 Backbone Norm ID SemEval CR Yelp FPB Task: SST-2 DeBERTa-v3 Batch-Mean 91.40 89.27 85.64 91.38 90.19 None 91.74 89.05 86.70 91.74 91.51 Δ +0.34 -0.22 +1.06 +0.36 +1.32 Qwen-2.5 Batch-Mean 92.89 74.19 73.40 78.45 83.58 None 93.00 73.91 72.34 79.13 85.87 Δ +0.11 -0.28 -1.06 +0.68 +2.29 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 w_i % wi w_i % wi w_i Effective (Mean ± Std) ≥0.9≥ 0.9 <0.5<0.5 N MNLI 0.87±0.100.87± 0.10 40.8% 0.0% 4,931 (98.6%) SST-2 0.95±0.080.95± 0.08 78.0% 0.0% 4,966 (99.3%) SQuAD v2 0.86±0.200.86± 0.20 63.5% 4.4% 4,748 (95.0%) NER 0.48±0.060.48± 0.06 0.0% 58.7% 4,931 (98.6%) Table 10: Distribution of per-sample IWD weights across training data at τ=1.0τ=1.0 and batch-mean normalization (N=5,000N=5,000 per task). Effective sample size (∑iwi)2/∑iwi2( _iw_i)^2/ _iw_i^2 measures information retention relative to uniform weighting, where values near N indicate gentle re-weighting rather than aggressive filtering. Weight Distribution Across Tasks. Table 10 reports IWD weight statistics across N=5,000N=5,000 samples at τ=1.0τ=1.0. High effective sample sizes across all tasks (95.0%95.0\%–99.3%99.3\%) confirm that IWD performs smooth continuous re-weighting rather than aggressive sample filtering. Weight profiles naturally adapt to task complexity: classification tasks (SST-2, MNLI) retain high mean weights (0.950.95 and 0.870.87), whereas SQuAD v2 exhibits higher variance (±0.20± 0.20) and 4.4%4.4\% low-weight samples (wi<0.5w_i<0.5) due to teacher sensitivity on complex extractive spans. On NER, higher overall perturbation sensitivity yields a lower mean weight (0.48±0.060.48± 0.06), yet a 98.6%98.6\% effective sample size confirms that relative importance ordering is preserved while maintaining overall data capacity.