Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models
Ayan Sengupta, Vaibhav Seth, Tanmoy Chakraborty
NeurIPS 2026
cs.CL
2026-09-29
CaRE-KD gates KL per token by teacher confidence and skips batches with unreliable teachers, beating static baselines across 11 benchmarks (MBPP +2.1, GSM8K +1.7).
Knowledge distillation trains a small student to match a large teacher's output distributions, which quietly assumes the teacher is a reliable oracle. With LLMs that assumption breaks: teachers produce high-entropy predictions and hallucinate. Forward KL is mean-seeking, so the student ends up spreading probability mass over the teacher's unreliable long tail, overwriting priors that were already well calibrated.
The paper measures this first (Figure 1). Distilling OpenLLaMA-3B from a 7B teacher with Skewed-RKL, supervision yields a median +4% gain when the teacher is accurate but only +0.5% when the teacher is uncertain about its own generation, and in those regimes the student's pre-distillation uncertainty is often lower than the teacher's. Standard KD forces agreement anyway; the authors call this failure mode the Fidelity Trap. Existing fixes (Reverse KL, Skewed KL, alpha-beta divergence) apply one fixed geometric preference for the entire run, so they cannot distinguish a token where the teacher is solid from one where it is fabricating.
CaRE-KD combines two granularities:
The gradient analysis isolates a selection term unique to the soft gate: when the teacher is comparatively uncertain, it explicitly pushes student confidence up. Proposition 3.2 proves that no global skew λ or fixed β can reproduce this token-varying gradient.
Eight teacher-student pairs and eleven benchmarks spanning instruction following, chat alignment, code and math, against FKL, RKL, SRKL, alpha-beta, plus MiniLLM, GKD, Distillm and ABKD.
| Comparison | Metric | Result |
| vs Forward KL (OPT-125M) | avg ROUGE-L / factuality | +3.2 / +3.0 |
| vs Skewed-RKL (MBPP) | pass@1 | 62.4 vs 60.3 (+2.1) |
| vs Skewed-RKL (GSM8K) | accuracy | 73.9 vs 72.2 (+1.7) |
| vs strongest baseline (CollegeMath) | accuracy | 46.1 (+1.8) |
| vs Skewed-RKL (factuality) | LLM judge | up to +2.5 per task |
The paper's own reading: the advantage is robustness, not universal first place. CaRE-Divergence has the best average on four of five instruction-following pairs, with an overall mean of 20.1 versus 19.8 for Skewed-RKL. Worst-case behavior is more telling: Forward KL trails the best by 4.1 on OpenLLaMA and Reverse KL by 3.0 on GPT2-large, while CaRE-KD's worst gap on any pair is 0.2. The losses are stated plainly: 17.0 vs 19.7 on Dolly for Gemma-2-2B, and a statistical tie with Distillm on OpenLLaMA (26.23 vs 26.41, smaller than per-seed variance). The stated scope: gains are clearest with unreliable teachers under heavy compression (GPT2-XL→base at 15x, OPT→125M at 22x); with a strong, well-calibrated teacher at 2.3x compression a static mode-seeking objective is already near optimal.
Isolating Revival (Table 2): it lifts static baselines on average but with high variance (RKL drops 1.78 on Self-Inst for GPT2-large), while paired with CaRE-Divergence it is reliably positive (one-sample t = 3.98, p < 0.001). Ablations favor 50-70% skip rates, an increasing selectivity schedule, and N = 3 MC samples.
On uncertainty profiles (Figure 4): the teacher's mean BALD is 0.52 on Self-Inst while the distilled student reaches 0.077; standard SRKL moves the student's uncertainty distribution by only 0.4% in KS distance, CaRE-Divergence alone by 2.4%, and adding Revival by 7%. The token-level loss is the main driver, with Revival amplifying it.
Distillation is the standard route to capable small models. This paper turns a vague worry that teachers are not always right into a measured failure mode, and the fix needs no new data and no auxiliary model: a changed loss plus one update rule. Revival is loss-agnostic and drops into existing pipelines. For teams compressing domain models or building on-device models this is a small, independently verifiable increment. Keep the framing honest: per-task gains mostly sit in the 1-3 point range, and the selling point is never losing badly rather than always winning.
The authors concede two points. Compute: BALD needs N stochastic forward passes per batch, about +42% wall-clock at N = 3 for instruction following and up to +170% on code; they explicitly retract calling the overhead modest and restrict the method to offline, one-time distillation with moderate-size teachers. Second, treating single-pass entropy as a BALD proxy relies on the Sharpness Hypothesis (cross-entropy-trained LLMs produce sharp individual samples), which may weaken in open-ended multimodal generation.
One mechanism claim remains unsettled: whether Revival actually selects hallucination-prone regions. Teacher factuality is 25.2 when Revival fires versus 28.6 when it does not (Wilcoxon p = 0.06; a complementary split gives p = 0.26), neither significant, so the authors fall back to the weaker claim that end-to-end student factuality improves. That is honest, but the mechanistic story is not closed. Also note where the gains live: small students under heavy compression. With strong, well-calibrated teachers the results are mostly ties, so extrapolating gains to distilling frontier models has no support in this paper.