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

Source

Source name: arXiv
Source kind: paper
Authors / organization: Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús
Venue: arXiv: cs.CL
Published: 2026-08-04
Collected: 2026-08-05
URL: http://arxiv.org/abs/2608.03796v1
PDF: https://arxiv.org/pdf/2608.03796v1
DOI: UNKNOWN

Plain-Language Brief

The researchers present a study to make knowledge distillation training efficient for small language models deployed under tight constraints. They propose two methods: offline knowledge distillation (caching teacher's top-$K$ logits) that matches online distillation with faster iteration speeds and higher throughput, and a fused, chunked KL loss that reduces memory usage by never materializing the full vocabulary-sized logit tensor. The study confirms the scaling of these methods from 4K to 256K tokens.

What the Document Says

Small language models are often deployed under tight constraints and are usually recovered through knowledge distillation (KD).
Offline KD (caching teacher's top-$K$ logits) matches online distillation at near-identical training loss.
Offline KD removes the teacher from memory, runs faster per iteration, and reaches higher throughput on a single GPU.
A fused, chunked KL loss is introduced that makes peak memory linear in sequence length.
The fused, chunked KL loss lets them 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.

Key Concepts Mentioned

knowledge distillation
small language models
offline KD
fused, chunked KL loss

Claims Made by the Document

Offline KD matches online distillation at near-identical training loss.
Offline KD runs faster per iteration and achieves higher throughput.
The fused, chunked KL loss makes peak memory linear in sequence length.
The fused, chunked KL loss lets them train at four times the context on a single GPU.

Evidence Used

Matches training loss of online distillation
Runs faster per iteration and achieves higher throughput
Makes peak memory linear in sequence length
Lets them train at four times the context on a single GPU

Methods / Approach

UNKNOWN

Limitations Stated by the Authors

UNKNOWN

Limitations Noted by the Agent

Agent interpretation — not from the source.

UNKNOWN

Open Questions

UNKNOWN

Notes for Future Q&A

May include interpretation. Labeled where it does.

UNKNOWN

Raw Abstract / Excerpt

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 \emph{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.