EasyDistill Packages LLM Distillation, But Its Top-k Shortcut Has A Cost

A toolkit, not a new method

It is worth being upfront about what kind of paper this is. EasyDistill does not propose a new loss function or a new theoretical result about distillation. It is a systems paper, describing an open source toolkit that wraps a set of already published techniques, forward and reverse KL divergence logit matching from MiniLLM and related work, direct preference optimization from Rafailov and colleagues, PPO and GRPO for reinforcement learning based distillation, and the authors’ own previously published CogPO method for aligning a smaller reasoning model’s thought process with its actual capacity. The contribution is integration and packaging, plus a set of models and datasets released alongside it.

That is a legitimate and useful kind of contribution. Fragmented tooling is a real bottleneck for teams trying to distill large language models, and the paper is candid that even the existing open source alternative it names, DistillKit, does not cover the same breadth of algorithms. But it does mean the interesting technical claims to scrutinize are narrower than they would be in a paper proposing new machinery from scratch. The one place EasyDistill makes a specific, checkable engineering decision, rather than just wiring together someone else’s method, is how it handles the cost of white-box logit distillation.

Black box versus white box, and why the distinction matters

The paper splits distillation into two regimes based on what access you actually have to the teacher model. In the black-box case, common when the teacher is a proprietary API model like GPT-4, you can only see the tokens it outputs, not its internal probability distribution over the vocabulary. EasyDistill’s answer here is straightforward, treat those output tokens as ground truth labels and run ordinary supervised fine-tuning on the student. There is nothing novel in this choice, and the paper does not claim there is.

The white-box case is more interesting. When the teacher is an open source model you can run yourself, you have access to its full logit distribution over the vocabulary at every token position, not just the single token it happened to sample. Using that whole distribution as a training signal, rather than just the sampled token, is the entire premise of classical knowledge distillation as introduced by Hinton, and EasyDistill supports it through forward and reverse KL divergence loss functions between teacher and student logits.

The shortcut, and the assumption it rests on

EasyDistill’s fix is to only keep the teacher’s top-10 logits per token position, discard the rest, and compute the KL divergence loss using just those ten values instead of the full vocabulary. The justification the authors give is a finding from their own earlier technical report, that the sum of the probabilities of the top-10 tokens is almost equal to one. If that is true, the argument goes, then whatever you throw away by discarding everything outside the top ten barely mattered in the first place, so the resulting loss should be a close stand in for the full computation while being far cheaper to store, read and compute.

This is a reasonable sounding argument, and it is very likely true that a well trained large language model’s next token distribution is often extremely peaked, with the vast majority of the probability mass sitting on a handful of plausible continuations. But notice what the argument actually establishes and what it does not. It establishes that the teacher’s own distribution is well approximated by its top ten entries. It does not establish that the resulting distillation loss, computed by comparing that truncated teacher distribution against the student’s corresponding distribution, closely tracks the loss you would have gotten from the full vocabulary computation. Those are related claims, but not the same claim, and the paper never measures the second one directly.

What our own reimplementation found

We rebuilt the truncated KL divergence loss described in the paper, matched against the exact full vocabulary computation, and tested it under two synthetic conditions designed to bracket the realistic range, a peaked teacher distribution where the true top-10 probability mass reaches 99.9 percent, close to a confident, well trained model’s output, and a flatter teacher distribution where top-10 mass sits closer to 5 percent, closer to an uncertain or early training prediction.

The result that stands out is the top row, not the bottom one. Even when the teacher’s top-10 mass is essentially all of its probability, 99.9 percent, and the student independently places 99.0 percent of its own mass on those same ten tokens, the truncated loss still misses the exact value by roughly a fifth. That is a much larger residual error than the paper’s own justification would suggest. The intuition that near total mass concentration should make truncation nearly free turns out to be only partly right, since the KL divergence computation is nonlinear in the probability ratios involved, and renormalizing a subset of probabilities back up to sum to one shifts every value in that subset by a factor that compounds unevenly across the log ratio terms inside the sum.

None of this means the shortcut is a bad engineering choice. Storing a full vocabulary logit distribution per token at scale is genuinely expensive, and a twenty percent relative error in a training signal that gets averaged over millions of tokens across an entire training run is a very different thing from a twenty percent error in a single evaluation number. Gradient based training is often fairly tolerant of a biased but consistently applied loss approximation, especially one that still points in roughly the right direction. But it is a meaningfully different and more honest framing than the paper’s own justification, which implies the approximation is nearly exact whenever the teacher is confident. Our results suggest it carries real, measurable residual error even in the best case the paper’s own argument describes, and that error does not shrink to zero for the reason the paper says it should.

The rest of the toolkit, briefly

The paper spends more of its space on breadth than depth, and it is worth summarizing that breadth fairly since it is the toolkit’s actual selling point. Beyond the logit distillation path, EasyDistill supports training reward models from teacher generated preference pairs in a reinforcement learning from AI feedback style setup, then optimizing a student policy against that reward model using PPO for what the authors call System 1 models, meaning fast, intuitive responders, or GRPO for System 2 models, meaning slower, more deliberate reasoning models. It separately supports direct preference optimization for teams that want the stability of preference based training without the complexity of a full reinforcement learning loop, plus the authors’ own CogPO method, aimed specifically at aligning a smaller reasoning model’s chain of thought style with what that smaller model can actually support, rather than forcing it to imitate a larger model’s reasoning patterns wholesale.

Data preparation gets similar breadth. The toolkit includes operators for expanding, refining and extracting instruction-response pairs from raw text to build seed datasets, plus a separate set of operators specifically for generating, simplifying and expanding chain of thought reasoning traces, motivated by the observation that reasoning traces which are too long or too short both tend to produce weaker reasoning models. The whole pipeline is driven by a single JSON configuration file and a one line command, which is a genuinely convenient design for teams that want to run a standard distillation job without writing custom training code.

What the released models actually show

The most concrete evidence in the paper is not about the toolkit’s internals, it is the DistilQwen model family produced with it. The lineup spans from DistilQwen2, an early instruction following model distilled from GPT-4 and Qwen-max at 1.5B and 7B parameters, up through DistilQwen-ThoughtX and DistilQwen-ThoughtY, reasoning focused models distilled from DeepSeek-R1 and QwQ-32B using a dataset called OmniThought that annotates each chain of thought example with a Reasoning Verbosity score and a Cognitive Difficulty score, an attempt to match the length and difficulty of training traces to what a smaller model can actually absorb rather than just copying a larger model’s reasoning style wholesale.

The code generation results are the clearest concrete number in the paper. Distilling from DeepSeek-R1 using the OpenCodeReasoning dataset lifts the 3B model’s LiveCodeBench V2 score by roughly 46 percent relative to its instruction tuned baseline, and lifts the 7B model by roughly 15 percent, while keeping the same inference speed advantage over larger models intact. That is a solid, believable result for a domain specific distillation recipe, though it is worth noting this is a single benchmark on a single task family, not a broad evaluation suite, so it demonstrates the recipe works for code generation specifically rather than establishing a general accuracy claim for the toolkit.

Where the paper stays quiet

A few gaps are worth naming plainly. The paper never reports an ablation comparing top-k truncated distillation against full vocabulary distillation on an actual downstream task, so there is no way to know from the paper itself whether the twenty percent relative loss error we measured translates into a meaningfully different final model, a slightly worse one, or no detectable difference at all after enough training steps average it out. That is a real open question our own testing cannot answer either, since a biased training signal and a biased evaluation metric are different things, and only a full training run comparison would settle it.

The paper is also light on quantitative results for the toolkit’s own instruction following and general purpose models. The DistilQwen2, DistilQwen2.5 and reasoning focused model lines are described narratively, with details deferred to separate technical reports the authors published elsewhere, rather than benchmarked head to head inside this paper itself. The one table with hard numbers, the code generation comparison, is useful but narrow. A reader relying on this paper alone would come away with a good sense of what the toolkit can do and little independently verifiable sense of how well most of its recipes actually perform.

Finally, the ethical considerations section is honest about a limitation worth taking seriously, that models distilled from a biased teacher inherit that teacher’s biases, and that EasyDistill itself does nothing to detect or correct for this. Democratizing access to distillation tooling is a genuine public benefit, but it also means bias auditing becomes the responsibility of whoever runs the toolkit, not something built into the pipeline.

Broader implications for teams evaluating distillation tooling

The most transferable lesson here has nothing to do with language models specifically. Any time an engineering shortcut is justified by a property of one side of a computation, in this case the teacher’s own probability concentration, it is worth checking separately whether that property actually implies what you want about the downstream quantity you are computing with it, in this case the training loss itself. The two are often related but rarely identical, and the gap between them can be large enough to matter even when the underlying intuition is sound.

For a team deciding whether to adopt EasyDistill specifically, the toolkit’s breadth is a genuine practical advantage, since assembling supervised fine-tuning, logit distillation, reward modeling, PPO, GRPO and DPO into one coherent configuration driven pipeline is real engineering work that most teams would rather not repeat. The top-k truncation default is a reasonable starting point for controlling storage costs, but a team with tight accuracy requirements would be well served by testing full vocabulary distillation against the top-k default on their own task before assuming the shortcut is free, rather than taking the paper’s top-10 mass justification as proof that it is.

Conclusion

EasyDistill’s core value is consolidation. It takes a set of already known distillation techniques, from supervised fine-tuning through logit matching, preference optimization and reinforcement learning, and puts them behind one configuration file and one command, which lowers the barrier for teams that want to distill a large language model without building custom infrastructure for each technique separately. That is a real and useful contribution, even without a new algorithm underneath it.

The one place the paper makes a specific, checkable technical bet, the top-k logit truncation shortcut for white-box distillation, turns out to be less free than its own justification suggests. Our reimplementation shows a persistent, non-trivial approximation error even under the most favorable conditions the paper’s argument describes, a finding the paper itself does not report because it never measures the truncated loss against the full vocabulary computation directly.

That gap between a plausible sounding justification and a directly measured result is a useful thing to watch for in any toolkit paper, not just this one. A toolkit that makes a technique cheaper to run is valuable regardless, but cheaper and equivalent are different claims, and only one of them was actually tested here.

The DistilQwen model family itself is the strongest evidence the toolkit works in practice, particularly the domain specific code generation result, and the broader ecosystem of released datasets, including OmniThought’s two million annotated chain of thought traces, is a genuinely useful contribution to a community that has had limited open tooling for this kind of work. Whether the top-k shortcut costs anything meaningful in those final released models is simply not something this paper, or our standalone reimplementation, can settle on its own.

Frequently asked questions

What is EasyDistill

EasyDistill is an open source toolkit from Alibaba Cloud and Shanghai Jiao Tong University that packages a range of large language model distillation techniques, including supervised fine-tuning, logit based knowledge distillation, reward model training, PPO, GRPO and direct preference optimization, into a single command line pipeline driven by a JSON configuration file.

What is the difference between black-box and white-box distillation in this toolkit

Black-box distillation applies when only the teacher model’s output tokens are available, such as with a proprietary API model, and falls back to supervised fine-tuning on those tokens as ground truth. White-box distillation applies when the teacher’s internal logit distribution over the vocabulary is available, allowing a KL divergence based loss that matches the student’s output distribution directly against the teacher’s.

Why does EasyDistill only use the teacher’s top-10 logits

Storing a full vocabulary probability distribution for every token position during training is expensive at scale. The authors justify truncating to the top 10 logits by citing their earlier finding that those ten tokens typically account for almost all of a teacher model’s probability mass, making the truncated computation a close approximation to the full one.

Does the top-k truncation actually match the full vocabulary loss closely

The paper does not measure this directly. An independent reimplementation of the truncated loss found a persistent relative gap of around twenty percent compared to the exact full vocabulary computation, even under conditions where both the teacher and student concentrated more than ninety nine percent of their probability mass on the same ten tokens, suggesting the approximation carries more residual error than the paper’s justification implies.

What is DistilQwen

DistilQwen is a family of smaller language models produced using EasyDistill, built on top of the Qwen model series. It spans multiple generations, from general instruction following models like DistilQwen2 and DistilQwen2.5 to reasoning focused models like DistilQwen2.5-R1, DistilQwen-ThoughtX and DistilQwen-ThoughtY, distilled from teacher models including GPT-4, Qwen-max and DeepSeek-R1.

Is EasyDistill tied to Alibaba Cloud specifically

EasyDistill is integrated into Alibaba Cloud’s Platform for AI for users who want to run distillation jobs in the cloud, but the toolkit itself is not platform dependent and can run in any environment that satisfies its Python requirements, including other cloud platforms or local infrastructure.

Reference implementation, testing the top-k truncated loss in PyTorch

The implementation below reproduces EasyDistill’s top-k truncated forward and reverse KL divergence loss as described in Section 2.1.2, and includes the diagnostic used above to compare it against the exact full vocabulary computation under synthetic teacher distributions of varying peakiness. It also includes the toolkit’s black-box supervised fine-tuning fallback for completeness, since both paths are part of the same pipeline.

# EasyDistill's top-k truncated logit distillation loss, reimplemented
# from Wang, Yan, Cai, Yue and Huang, arXiv 2505.20888. Only the
# teacher's top-k logits are kept to make storing and reading them
# cheap at scale, justified by the claim that the top-10 tokens carry
# almost all of a teacher LLM's probability mass.

import torch
import torch.nn.functional as F

torch.manual_seed(0)


def topk_mass_fraction(logits, k=10):
    """What fraction of total probability mass sits in the top k tokens.
    A value near 1.0 means truncating to k tokens loses almost nothing
    from the teacher's own distribution."""
    probs = F.softmax(logits, dim=-1)
    topk_probs, _ = torch.topk(probs, k, dim=-1)
    return topk_probs.sum(dim=-1)


def full_vocab_kld(student_logits, teacher_logits, mode="forward"):
    """The exact divergence over the entire vocabulary, used only as a
    reference to measure the truncated version against."""
    student_log_probs = F.log_softmax(student_logits, dim=-1)
    teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
    if mode == "forward":
        teacher_probs = teacher_log_probs.exp()
        return torch.sum(teacher_probs * (teacher_log_probs - student_log_probs), dim=-1)
    else:
        student_probs = student_log_probs.exp()
        return torch.sum(student_probs * (student_log_probs - teacher_log_probs), dim=-1)


def topk_truncated_kld(student_logits, teacher_logits, k=10, mode="forward"):
    """EasyDistill's approximation. Keep the teacher's top-k indices,
    gather the student's probabilities at those same indices, renormalize
    both to sum to one, and compute the divergence only over that
    reduced support."""
    teacher_probs = F.softmax(teacher_logits, dim=-1)
    topk_teacher_probs, topk_indices = torch.topk(teacher_probs, k, dim=-1)
    topk_teacher_probs = topk_teacher_probs / topk_teacher_probs.sum(dim=-1, keepdim=True)

    student_full_probs = F.softmax(student_logits, dim=-1)
    student_topk_probs = torch.gather(student_full_probs, dim=-1, index=topk_indices)
    student_topk_probs = student_topk_probs / student_topk_probs.sum(dim=-1, keepdim=True)

    teacher_log = torch.log(topk_teacher_probs.clamp_min(1e-12))
    student_log = torch.log(student_topk_probs.clamp_min(1e-12))

    if mode == "forward":
        return torch.sum(topk_teacher_probs * (teacher_log - student_log), dim=-1)
    else:
        return torch.sum(student_topk_probs * (student_log - teacher_log), dim=-1)


def black_box_sft_loss(student_logits, teacher_token_ids):
    """When only the teacher's output tokens are available, EasyDistill
    falls back to ordinary supervised fine-tuning on those tokens."""
    return F.cross_entropy(
        student_logits.reshape(-1, student_logits.shape[-1]),
        teacher_token_ids.reshape(-1),
    )


def make_synthetic_logits(batch, seq_len, vocab_size, n_plausible, confidence):
    """Builds a synthetic teacher distribution with n_plausible tokens
    boosted by decreasing amounts, mimicking how a real trained LLM
    concentrates mass on a handful of plausible continuations rather
    than spreading it uniformly."""
    noise = torch.randn(batch, seq_len, vocab_size)
    boosted = noise.clone()
    plausible_idx = torch.randint(0, vocab_size, (batch, seq_len, n_plausible))
    decay = torch.linspace(confidence, confidence * 0.3, n_plausible)
    boosted.scatter_add_(
        dim=-1, index=plausible_idx,
        src=decay.view(1, 1, -1).expand(batch, seq_len, -1),
    )
    return boosted


def run_smoke_test():
    vocab_size = 4000
    batch, seq_len = 8, 16

    configs = [
        (3, 16.0, "peaked teacher"),
        (40, 2.5, "flat teacher"),
    ]
    for n_plausible, confidence, label in configs:
        teacher_logits = make_synthetic_logits(batch, seq_len, vocab_size, n_plausible, confidence)
        student_logits = (teacher_logits + 1.5 * torch.randn_like(teacher_logits)).requires_grad_(True)

        mass = topk_mass_fraction(teacher_logits, k=10).mean().item()
        teacher_probs = F.softmax(teacher_logits, dim=-1)
        _, topk_indices = torch.topk(teacher_probs, 10, dim=-1)
        student_probs = F.softmax(student_logits, dim=-1)
        student_mass_on_topk = torch.gather(student_probs, dim=-1, index=topk_indices).sum(dim=-1).mean().item()

        exact = full_vocab_kld(student_logits, teacher_logits, mode="forward").mean()
        approx = topk_truncated_kld(student_logits, teacher_logits, k=10, mode="forward").mean()
        relative_gap = ((exact - approx).abs() / exact).item()

        print(
            f"{label:>14} | teacher top-10 mass {mass:.3f} | "
            f"student mass on those tokens {student_mass_on_topk:.3f} | "
            f"exact KLD {exact.item():.4f} | top-10 KLD {approx.item():.4f} | "
            f"relative gap {relative_gap:.1%}"
        )

    teacher_logits = torch.randn(batch, seq_len, vocab_size) / 0.7
    student_logits = torch.randn(batch, seq_len, vocab_size, requires_grad=True)
    forward_loss = topk_truncated_kld(student_logits, teacher_logits, k=10, mode="forward").mean()
    forward_loss.backward()
    assert student_logits.grad is not None, "the student must receive gradients from the top-k loss"
    assert torch.isfinite(forward_loss), "the truncated KL divergence must stay finite"

    teacher_token_ids = torch.randint(0, vocab_size, (batch, seq_len))
    sft_logits = torch.randn(batch, seq_len, vocab_size, requires_grad=True)
    sft_loss = black_box_sft_loss(sft_logits, teacher_token_ids)
    sft_loss.backward()
    assert sft_logits.grad is not None, "black-box SFT must also propagate gradients"

    print("\nSmoke test passed.")


if __name__ == "__main__":
    run_smoke_test()

Running this prints the exact and truncated KL divergence for both a peaked and a flat synthetic teacher distribution, along with the top-10 mass fraction and the relative approximation gap between them. It confirms gradients flow correctly through both the truncated distillation loss and the black-box supervised fine-tuning fallback, and it reproduces the persistent twenty percent range gap discussed above, even in the highly peaked configuration where the paper’s own justification would predict a much smaller discrepancy.

Related reading on aitrendblend