# Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

arXiv: https://arxiv.org/abs/2608.03796 · HTML: https://mlsanity.com/paper/2608.03796/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 15 of its 18 references that are in the map.
- First posted on arXiv 4 Aug 2026.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 15 of 15)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Compact Language Models via Pruning and Knowledge Distillation](https://mlsanity.com/paper/2407.14679/) (2024)
- [How to Train Long-Context Language Models (Effectively)](https://mlsanity.com/paper/2410.02660/) (2024)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [HELMET: How to Evaluate Long-Context Language Models Effectively and Thoroughly](https://mlsanity.com/paper/2410.02694/) (2024)
- [Cut Your Losses in Large-Vocabulary Language Models](https://mlsanity.com/paper/2411.09009/) (2024)
- [A Comparative Analysis of Task-Agnostic Distillation Methods for Compressing Transformer Language Models](https://mlsanity.com/paper/2310.08797/) (2023)

## Built on by (most cited first, 2 of 2)

- [The Gradient That Top-k Distillation Never Computes](https://mlsanity.com/paper/2610.15589/) (2026)
- [DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss](https://mlsanity.com/paper/2610.17644/) (2026)

## Abstract (arXiv)

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-K logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29% faster per iteration, and reaching up to 41% higher throughput on a single H200 GPU. Second, we introduce a fused, chunked KL loss that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32,768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.
