The Gradient That Top-k Distillation Never Computes
cs.LG, cs.CL
Submitted: 2026-10-11
Updated: 2026-10-11
Code: https://github.com/KurbanIntelligenceLab/top-k-distillation
License: http://creativecommons.org/licenses/by/4.0/
The gist: Many deployed distillation pipelines show the student a truncated view of the teacher: keep the top-k probabilities, renormalize, and minimize a divergence on what remains.
Terminology
Abstract
Many deployed distillation pipelines show the student a truncated view of the teacher: keep the top-k probabilities, renormalize, and minimize a divergence on what remains. Renormalization is assumed to account for the missing mass. We prove that it removes the visibility of tail mass, not the freedom to move it. The renormalized objective depends only on log-odds among the retained tokens, so at each position its gradient along the directions that control tail mass is exactly zero; on a real teacher that component measures 2 - 3 times10-15 of the gradient. The full-vocabulary gradient equals the truncated gradient plus a residual weighted by the tail masses, and training acts only on the first term. Without the cross-entropy term, a top-k student's tail mass drifts away from the teacher's while a full-vocabulary control moves slightly toward it; with it, the gap shrinks about fourfold but persists. On seven open-weight teachers, the retained set determines the teacher's distribution only from k=d+1 (d the hidden dimension): recovery is exact from there on all seven, and commonly cached sizes sit below it. In our off-policy runs the practical effect is real but smaller than the choice of teacher-student pair. Renormalization is a normalization, not a safeguard. The code is available at https://github.com/KurbanIntelligenceLab/top-k-distillation.
Sources
- Distillation Scaling Laws
- Stealing Part of a Production Language Model
- DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence
- Logits of API-Protected LLMs Leak Proprietary Information
- Representation Degeneration Problem in Training Natural Language Generation Models
- Why do small language models underperform? Studying Language Model Saturation via the Softmax Bottleneck
- Distilling the Knowledge in a Neural Network
- Liger Kernel: Efficient Triton Kernels for LLM Training
- The Geometry of Last-Layer Model Stealing
- DistiLLM: Towards Streamlined Distillation for Large Language Models
- DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs
- On the Generalization of Knowledge Distillation: An Information-Theoretic View
- Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
- Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation
- FitNets: Hints for Thin Deep Nets
- Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
- Geometric Limits of Knowledge Distillation: A Minimum-Width Theorem via Superposition Theory
- A Survey of On-Policy Distillation for Large Language Models
- The Illusion of Certainty: Decoupling Capability and Calibration in On-Policy Distillation
- KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks