TopKD: Top-Scaled Knowledge Distillation

• TopKD is a logit-based knowledge distillation framework that prioritizes the teacher model’s top-K logits to enhance semantic supervision.
• The Top-K Scaling Module amplifies the most informative logits using rank-sensitive factors and bias terms to preserve ground-truth signals.
• The Top-K Decoupled Loss aligns student outputs through contrastive and cosine-based measures, boosting performance across diverse architectures.

Top-scaled Knowledge Distillation (TopKD) is a logit-based knowledge distillation framework that treats the teacher model’s Top-K logits as a critical supervisory source rather than aligning the full output distribution uniformly. It is formulated as a simple, efficient, and architecture-agnostic method with two training-time components—a Top-K Scaling Module (TSM) and a Top-K Decoupled Loss (TDL)—and is designed to integrate into existing distillation pipelines without extra modules or architectural changes (Wang et al., 6 Aug 2025). The central premise is that the teacher’s highest-ranked logits encode semantically relevant supervision that conventional KL-divergence-based distillation underutilizes, especially when all classes are treated equally despite large variation in informativeness across logits (Wang et al., 6 Aug 2025).

1. Conceptual basis and definition of Top-K knowledge

TopKD arises from a critique of two dominant tendencies in the distillation literature. First, many recent methods emphasize feature-level transfer, which is resource-intensive and highly dependent on specific architectures. Second, conventional logit-based methods often rely on KL divergence over the entire teacher distribution, thereby allocating supervision uniformly across classes and potentially suppressing the most discriminative signals carried by the teacher’s highest logits (Wang et al., 6 Aug 2025).

Within this framework, Top-K knowledge denotes the information contained in the teacher model’s $K$ largest output logits. These logits correspond to the classes the teacher deems most probable or similar to the input. The method attributes special significance to them because they often encode hierarchical or semantic relations between classes; even when the top-1 prediction is incorrect, the Top-K set may still contain semantically related classes. The example given in the method description is that the top predictions for a seal image may include “otter,” a visually similar class, even if the teacher predicts incorrectly (Wang et al., 6 Aug 2025).

This definition makes TopKD a targeted intervention on the structure of the teacher output space. Rather than discarding the non-Top-K region outright, the method amplifies the Top-K region and then supervises the student with a loss that decouples different logit segments. A plausible implication is that TopKD is less a sparse-output approximation scheme than a reweighting and alignment strategy over logits.

2. Top-K Scaling Module

The Top-K Scaling Module is the mechanism that explicitly amplifies the teacher’s most informative logits. Its stated purpose is twofold: to prioritize the Top-K teacher logits and to correct bias by further boosting the ground-truth logit when the teacher’s top prediction is wrong (Wang et al., 6 Aug 2025).

For a teacher logit $z_i$, the scaled logit $z_i'$ is defined as

$$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$

where $z_i$ is the teacher’s logit for class $i$, $\mathcal{I}_{\text{top}}$ is the index set of Top-K logits, $y_g$ is the ground-truth class index, $w_i$ is a rank-dependent scaling factor, and $\Delta$ is a bias term proportional to the mean difference between Top-K and non-Top-K logits (Wang et al., 6 Aug 2025).

The role of $z_i$0 is to impose a rank-sensitive amplification, with higher-ranked logits receiving stronger emphasis. The role of $z_i$1 is to preserve or restore supervision around the ground-truth class when the raw teacher ranking is imperfect. In effect, TSM modifies the teacher signal before any alignment loss is computed, so that subsequent optimization is driven by a teacher representation in which semantically salient and task-relevant logits are made more prominent.

This module also clarifies that TopKD is not equivalent to simply taking the Top-K teacher entries and ignoring the rest. The non-Top-K logits remain present in the scaled representation; they are relatively down-emphasized rather than removed.

3. Top-K Decoupled Loss

The Top-K Decoupled Loss complements TSM by specifying how the student is aligned to the scaled teacher logits. Its stated purpose is to provide a tailored, decoupled alignment of student and teacher logits, with emphasis on intra-instance structural similarity and especially on the Top-K dimensions (Wang et al., 6 Aug 2025).

TopKD combines two losses. The first is a contrastive loss for instance-level alignment:

$z_i$2

where $z_i$3 and $z_i$4 are student and teacher logits, $z_i$5 is a temperature, and $z_i$6 denotes batch indices (Wang et al., 6 Aug 2025). This term is described as enabling instance-level matching and improving discrimination for each example.

The second is a cosine-similarity-based decoupled loss. Cosine similarity is written as

$z_i$7

The teacher and student logits are partitioned into three segments—Positive Top-K, Negative Top-K, and Non-Top-K—and weighted differently:

$z_i$8

where $z_i$9 is the scaled teacher logit for sample $z_i'$0, and $z_i'$1 are tunable weights for positive and negative Top-K components; the reported experimental defaults are $z_i'$2 and $z_i'$3 (Wang et al., 6 Aug 2025).

The full TopKD objective is

$z_i'$4

This formulation differs from standard KL-based logit matching in that it does not enforce a single uniform alignment criterion over the entire output simplex. Instead, it decouples segments of the logit space and assigns stronger supervision to the most informative regions.

4. Integration properties and training behavior

TopKD is presented as a plug-and-play framework. Both TSM and TDL are lightweight, operate at the output layer, require no changes to the student backbone, and can be added on top of both logit-based and feature-based distillation pipelines (Wang et al., 6 Aug 2025). The method is active only during training and introduces negligible computational overhead, with no increase in inference cost (Wang et al., 6 Aug 2025).

The intended synergy between the two components is explicit. TSM strengthens the semantic signal of Top-K and ground-truth logits, while TDL uses this modified signal to compute focused alignment losses. The contrastive term provides instance-level supervision, and the cosine-based decoupling is meant to preserve structural relations across distinct logit regions rather than only value-wise agreement (Wang et al., 6 Aug 2025).

This design also underlies the claim of architecture-agnostic applicability. Because the method acts only on logits, it does not depend on matching hidden representations, intermediate feature shapes, or auxiliary projectors. That property is central to its reported deployment across CNNs, lightweight mobile backbones, and Vision Transformers (Wang et al., 6 Aug 2025).

A further operating characteristic is the sensitivity of performance to the choice of $z_i'$5. Sensitivity analysis reports optimal performance for $z_i'$6–$z_i'$7, while too large a $z_i'$8 dilutes knowledge quality (Wang et al., 6 Aug 2025). This indicates that TopKD is not predicated on exhaustive retention of the teacher distribution; its effect depends on isolating a relatively small high-value region of the logit vector.

5. Empirical profile

TopKD is evaluated on CIFAR-100, ImageNet, STL-10, and Tiny-ImageNet, with architectures including Wide ResNets, ResNets of various depths, VGG, MobileNet, ShuffleNet, and multiple Vision Transformers such as DeiT-Ti, T2T-ViT7, PiT-Ti, and PVT-Ti (Wang et al., 6 Aug 2025). The reported results state that TopKD consistently surpasses state-of-the-art distillation methods, including strong feature-based methods such as CRD, ReviewKD, SimKD, and FCFD, as well as advanced logit-based methods such as KD, DKD, DKD+LS, WTTM, and DOT (Wang et al., 6 Aug 2025).

Setting	Reported result	Context
CIFAR-100, WRN-40-2 / ResNet8x4	76.88%	Compared to DKD+LS at 75.93%
ImageNet, ResNet-50 $z_i'$9 MobileNetV1	73.51% Top-1 accuracy	Higher than most existing methods
STL-10, Tiny-ImageNet transfer	Higher downstream performance	Outperformed CRD, DOT, FCFD, etc.

The method is reported to perform strongly in both homogeneous and heterogeneous teacher-student settings on CIFAR-100. On ImageNet, it matches or exceeds feature-based and logit-based baselines across setups including ResNet-34/ResNet-18 and ResNet-50/MobileNetV1 (Wang et al., 6 Aug 2025). Students trained with TopKD also show stronger representation transferability on downstream classification tasks over STL-10 and Tiny-ImageNet.

The ablation evidence is similarly structured. Both TSM and TDL contribute additive gains, and their combination performs best. t-SNE plots and logit correlation matrix visualizations are reported to show improved feature discriminability and stronger knowledge-transfer fidelity. The method is also described as mitigating the “bigger-teacher is worse” issue in large-capacity teacher scenarios (Wang et al., 6 Aug 2025). For Vision Transformers, TopKD improves student performance across various ViT-like architectures, which is used as evidence for its versatility.

6. Relation to adjacent distillation paradigms

TopKD belongs to a broader line of work that treats distillation as a problem of selective rather than uniform knowledge transfer, but its selectivity is specifically logit-centric. Knowledge Condensation Distillation (KCD), for example, argues that standard KD transfers redundant knowledge because knowledge has different value to the student at different stages; it dynamically estimates sample value, uses an Expectation-Maximization framework to condense a compact knowledge set, and explicitly characterizes TopKD as focusing on transferring only the most informative entries of the teacher’s output distribution, such as top-k logits or classes (Li et al., 2022). This suggests that TopKD occupies the static, output-entry-focused end of a broader family of student-adaptive selection strategies.

A separate line of work complicates any universal interpretation of “Top-K knowledge.” In neural machine translation, empirical analysis finds that the benefit of KD largely comes from the teacher’s top-1 predictions, that retaining top-k predictions for $$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$0 does not improve the benefit, and that a method emphasizing top-1 information through hierarchical ranking loss and iterative KD yields improved BLEU on WMT benchmarks (Zhang et al., 2023). By contrast, TopKD reports that in its vision setting the best $$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$1 lies roughly between 3 and 10, with larger values diluting knowledge quality (Wang et al., 6 Aug 2025). The plausible implication is that the effective granularity of “top” information is task-dependent rather than universal.

Another distinction concerns the difference between emphasizing Top-K logits and truncating teacher distributions to Top-K entries. In LLM pre-training, naive sparse knowledge distillation based on caching Top-K probabilities is shown to provide biased estimates of the teacher distribution, leading to suboptimal performance and calibration; the corresponding gradient becomes $$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$2 with $$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$3, producing over-confident students that cannot match tail probabilities (Anshumann et al., 21 Mar 2025). TopKD does not describe such Top-K truncation. Instead, it rescales teacher logits and applies decoupled losses while retaining non-Top-K regions in the training objective (Wang et al., 6 Aug 2025). This distinction is central to avoiding a common misconception that all Top-K-based distillation methods are sparse approximations of the teacher distribution.

DeepKD provides yet another selective variant, using a dynamic top-k mask that filters low-confidence non-target logits and gradually increases $$z_i' = \begin{cases} z_i \cdot w_i + \Delta, & \text{if } i \in \mathcal{I}_{\text{top} \cup \{y_g\}} \ z_i, & \text{otherwise} \end{cases}$$4 according to a curriculum-learning schedule (Huang et al., 21 May 2025). Compared with such dynamic masking, TopKD uses a fixed Top-K emphasis within a decoupled logit-alignment framework. The shared theme across these methods is that distillation quality depends on identifying which parts of the teacher signal are informative, but they operationalize that principle differently: TopKD through scaled logits and decoupled alignment, KCD through sample-value condensation, TIE-KD through top-1-focused ranking supervision, and DeepKD through curriculum-based masking.

Within contemporary KD research, TopKD is therefore best understood as a logit-centric framework that reasserts the value of teacher output distributions—not in their entirety, and not as a uniformly weighted target, but as a structured space in which the highest-ranked logits carry privileged supervisory content (Wang et al., 6 Aug 2025).