Sparse Teacher Support Cuts LLM Distillation Cost Without Losing Accuracy

Forward‑KL distillation of large language models traditionally requires dense teacher and student scores over the entire vocabulary at every supervised token, a process that is both compute‑ and memory‑heavy. A new study from the Grafy Research Lab demonstrates that discarding all but the teacher’s highest‑probability tokens—down to just 0.01% of the vocabulary—leaves generation quality unchanged while dramatically reducing training overhead.

What is "teacher support" and how is it reduced?

In standard next‑token distillation the loss at position $t$ compares the full teacher distribution $p_{t}$ with the student distribution $q_{t}$ over a vocabulary of size $V$:

The researchers introduce a retained‑support objective that keeps only the top‑$k$ tokens with the highest teacher probability at each position. The rank $k_{\rho}(t)$ is defined as the smallest $k$ such that the cumulative probability of the top‑$k$ tokens reaches a target mass $\rho$:

By sweeping $\rho$ from 100% down to 0.01%, they examine how aggressive truncation affects downstream performance.

How much of the vocabulary can be removed without hurting performance?

The main experiments distill a 4 B‑parameter Qwen3 teacher into a 1.7 B student. The vocabulary contains 151,936 tokens. Keeping only $k\!=\!15$ tokens—exactly 0.01% of the vocab—retains at least 90% of the teacher’s probability mass for 99.0% of supervised positions, and 95% for 97.9% of positions. Crucially, on the primary benchmark MMLU‑Pro there is no observable degradation compared with full‑vocabulary distillation. A secondary benchmark, GPQA‑Diamond, shows the same trend.

When the same method is applied to larger pairs—Qwen3 8 B → 4 B and Nemotron Nano 12 B → 9 B—the results are consistent for Qwen3 models, but Nemotron Nano requires a slightly larger support of 0.1% (≈150 tokens) to match full‑vocabulary performance; dropping below that prevents the model from emitting the required reasoning‑closure tag.

What training speed and memory gains are realized?

Caching the retained teacher scores for $k\!=\!15$ eliminates the need to run the teacher model during each training step. This cache reduces step time by 38%. Because the student only needs to compute logits for the $k$ retained tokens, peak memory usage drops by 24% compared with a dense teacher‑forward pass. Importantly, these savings are achieved without any measurable loss in MMLU‑Pro accuracy.

Does the approach generalize across model families and supervision styles?

The paper evaluates two supervision regimes: full reasoning‑trace supervision and answer‑only supervision. In both settings, the sparse‑support strategy holds up for Qwen3 models. Cross‑family tests with Nemotron Nano reveal a modest sensitivity to model architecture, suggesting that the optimal support fraction may vary with model size and tokenization scheme. Nonetheless, the overall conclusion—that a tiny fraction of the vocabulary can carry sufficient information for effective distillation—appears robust.

How can creators and researchers try this technique with Grafy?

Grafy’s Reader app lets you upload the arXiv PDF, ask the model to highlight the retained‑support methodology, and retrieve the exact $k$‑value calculations without leaving the workspace. For hands‑on experimentation, the Code studio provides a ready‑to‑run Jupyter‑style environment where you can implement the rank‑$k$ loss, cache teacher logits, and benchmark training speed on your own models. If you prefer an interactive exploration, Chat lets you query multiple LLM back‑ends side‑by‑side about the trade‑offs of different $\rho$ settings. Together, these tools let developers prototype sparse‑support distillation pipelines without building the infrastructure from scratch.

---

Source: https://arxiv.org/abs/2611.14850

Frequently asked questions

Does retaining only 0.01% of the vocabulary affect the student’s ability to generate correct answers?

In the Qwen3 4 B → 1.7 B experiments, the student’s performance on MMLU‑Pro was indistinguishable from a full‑vocabulary baseline, indicating that the essential information for correct answer generation is preserved even with such extreme truncation.

Why does Nemotron Nano need a larger support fraction than Qwen3 models?

The Nemotron Nano experiments showed that dropping below 0.1% of the vocabulary prevented the model from emitting the required reasoning‑closure token, suggesting that its token distribution is less concentrated than Qwen3’s, so a slightly larger support is needed to capture critical probabilities.

Can the cached teacher‑logit approach be combined with other distillation tricks like higher temperature or reverse KL?

The paper focuses on forward‑KL with temperature $\tau\!=\!1$, but the authors note that the caching mechanism is orthogonal to other distillation variations. In principle, you could apply a higher temperature or switch to reverse KL while still benefiting from the reduced compute and memory footprint.

---

Source: Grafy Research Lab — Effective Teacher Support for Reasoning Distillation