Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models
cs.CL
Submitted: 2026-09-29
Updated: 2026-09-29
Comments: Accepted at NeurIPS 2026
Code: https://github.com/openlm-research/open_llama
License: http://creativecommons.org/licenses/by/4.0/
The gist: Knowledge Distillation (KD) trains a smaller-capacity student model to imitate a larger-capacity teacher model by matching output distributions, implicitly assuming the teacher to be a reliable
Terminology
Abstract
Knowledge Distillation (KD) trains a smaller-capacity student model to imitate a larger-capacity teacher model by matching output distributions, implicitly assuming the teacher to be a reliable oracle. In large language models (LLMs), this assumption often fails: teacher predictions can exhibit high entropy and hallucinations, causing standard KD to degrade well-calibrated student priors. We propose CaRE-KD, a confidence-gated distillation framework that replaces static objectives with uncertainty-adaptive optimization. CaRE-KD has two components: a token-level loss (CaRE-Divergence) that adaptively switches between Forward and Reverse KL divergence based on teacher--student confidence, and a batch-level epistemic rejection mechanism (Revival) that suppresses updates when the teacher is more uncertain than the student. We provide a gradient-level analysis showing how this dual-granularity design induces a conditional calibration mechanism that prior static divergences cannot reproduce. Empirically, across eight teacher--student pairs and eleven benchmarks spanning instruction following, chat alignment, code generation, and mathematical reasoning, CaRE-KD delivers consistent gains over strong baselines (Skewed-KL, α -- β divergence). Highlights include up to +3.2 average ROUGE-L on instruction-following tasks, +2.1 pass@1 on MBPP, +1.7 accuracy on GSM8k, and +1.8 accuracy on CollegeMath over the strongest baseline, with consistent gains in LLM-as-a-judge factuality (up to +2.5 per task over Skewed-RKL). Revival further acts as a principled, loss-agnostic plug-in that systematically strengthens existing distillation objectives by filtering epistemically unreliable teacher supervision.
Sources
- Sequence-Level Knowledge Distillation
- Why Language Models Hallucinate
- Large Language Models are overconfident and amplify human bias
- DistiLLM: Towards Streamlined Distillation for Large Language Models
- Language Models (Mostly) Know What They Know
- Bayesian Active Learning for Classification and Preference Learning
- Distilling the Knowledge in a Neural Network
- OPT: Open Pre-trained Transformer Language Models
- Gemma 2: Improving Open Language Models at a Practical Size
- Qwen2 Technical Report
- DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence
- Evaluating Large Language Models Trained on Code
- Program Synthesis with Large Language Models
- Qwen2.5 Technical Report
- Training Verifiers to Solve Math Word Problems
- MathScale: Scaling Instruction Tuning for Mathematical Reasoning
- Patient Knowledge Distillation for BERT Model Compression
- Revisiting Knowledge Distillation for Autoregressive Language Models
- Knowledge Distillation Based on Transformed Teacher Matching
- MiniLLM: On-Policy Distillation of Large Language Models
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering