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

Hugging Face Daily Papers Papers
knowledge-distillation llm efficiency training kl-loss offline-kd

Summary

This paper presents a practical study on making knowledge distillation training for LLMs more efficient, introducing offline top-K logits caching and a fused chunked KL loss that reduces memory spikes and enables longer contexts on a single GPU.

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.
Original Article
View Cached Full Text

Cached at: 08/10/26, 02:15 PM

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

Source: https://huggingface.co/papers/2608.03796

Abstract

Smalllanguagemodelsareoftentheonlyoptionfordeploymentundertightlatency,cost,andon-premisesconstraints,buttheyarerarelytrainedfromscratch:acompressedmodelisusuallyrecoveredthroughknowledgedistillation(KD).Thisrecoverysteplargelydecidesthefinalquality,yetitisexpensive.Wepresentapractitioner’sstudyofhowtomakedistillationtrainingefficient,organisedaroundtwosystemscontributions.First,weshowthatofflineKD(cachingtheteacher’stop-Klogitsonceandtrainingthestudentagainstthecache)matchesonlinedistillationatnear-identicaltraininglosswhileremovingtheteacherfrommemory,runningabout29\%fasterperiteration,andreachingupto41\%higherthroughputonasingleH200GPU.Second,weintroduceafused,chunkedKLlossthatnevermaterialisesthefullvocabulary-sizedlogittensor,makingpeakmemorylinearinthesequencelength.Thisremovesthememoryspikethatotherwisecapscontextlengthandletsustrainatfourtimesthecontext(32{,}768tokens)onasingleGPU.Aseparateoutput-head-onlytoybenchmarkisolatesthelosskernelandconfirmsitsmemoryanditeration-ratescalingfrom4Kto256Ktokens.Togetherthesemakelarge-scalehealingandhundredsofablationsaffordable.Wealsoreportsupportingablationsonlossdesignandsequencepacking.Wereleaseourchunked-lossimplementation:https://github.com/CompactifAI/Full-Chunked-KL-Loss.

View arXiv pageView PDFProject pageGitHub3Add to collection

Get this paper in your agent:

hf papers read 2608\.03796

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper2

#### MultiverseComputingCAI/Hypernova-60B-2605 Text Generation• 59B• Updated3 days ago • 1.45k • 15 #### MultiverseComputingCAI/LittleLamb Text Generation• 0.3B• Updated3 days ago • 1.6k • 9

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2608.03796 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2608.03796 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

KDFP: A first-principles approach to knowledge distillation in large language models

arXiv cs.LG

KDFP: A first-principles approach to knowledge distillation in large language models Source: https://arxiv.org/html/2610.10854 Ryan Swift & Konstantinos PsounisAffiliation:Thomas Lord Department of Computer ScienceAffiliation:University of Southern CaliforniaAffiliation:Los Angeles, CA 90007, USAEmail:[{ryanswif,kpsounis}@usc.edu](mailto:) ###### Abstract Knowledge distillation is an established technique for improving the capabilities of small, efficient student models by training them with the representations of larger, more capable teacher models. Much of the recent work in the distillation of large language models (LLMs) has focused on distilling abilities learned during post-training, such as instruction following, chain-of-thought reasoning, and tool usage. This has left a large research gap in general knowledge distillation for LLMs, which is essential for developing efficient and private systems suitable for deployment on edge devices. We take a first-principles approach, evaluating previous lessons from prior works and conducting new explorations to develop a distillation methodology suitable for modern LLMs. We present KDFP, a novel methodology for white-box general knowledge distillation in LLMs. We demonstrate that KDFP outperforms existing methods by 1.6% – 4.9% across 9 benchmarks while increasing training efficiency by up to 99.1% through ephemeral parameter reduction.