Knowledge Distillation from A Stronger Teacher

In plain words

Knowledge distillation trains a small "student" model to copy a large "teacher" model. The usual way is to make the student's output probabilities match the teacher's exactly. That goes wrong when the teacher is much stronger: the student cannot follow the teacher's exact numbers, and forcing it to try can hurt training.

The paper's idea is to care less about exact values and more about the pattern. If the teacher thinks class A is more likely than class B, and the student agrees on that ordering with a similar spread of confidence, the student has learned what matters. The method measures this with a correlation score, and it looks at the pattern in two directions: across the classes for each example, and across the examples for each class. Training then uses the ordinary label loss plus these two correlation terms.

The payoff, as the authors report it, is better students from strong teachers, with training cost close to standard distillation, on image classification and on two dense prediction tasks, object detection and semantic segmentation.

TL;DR

• Replaces the exact KL match of standard knowledge distillation with Pearson correlation between student and teacher softmax outputs §3.1.
• Two terms: inter-class (per instance, across classes) and intra-class (per class, across instances) §3.1 §3.2.
• Total loss: classification cross-entropy plus the two weighted correlation losses eq. 10.
• Author-reported gains are largest when the teacher is stronger or trained with stronger strategies §4.2 tab. 3 tab. 4.
• Needs only the logits of teacher and student; the appendix gives a short PyTorch implementation app. A.1.

Problem

Vanilla knowledge distillation minimises the KL divergence between the teacher's and student's softmax outputs §2 eq. 1. The paper argues that this exact match is the wrong target when the teacher is strong. Its observations:

• With a stronger teacher, the student's accuracy under vanilla KD can fall below training without KD. The paper points to its Figure 1 for this §1 fig. 1.
• The gap between teacher and student predictions grows when the teacher is trained with a stronger strategy, such as label smoothing, mixup or auto-augmentation. The paper shows this as the KL discrepancy of ResNet-18 and ResNet-50 trained with two strategies §2.1 fig. 2.
• A ResNet-18 student's outputs change little between the two strategies, while the ResNet-50 teacher's do. The authors read this as a capacity limit on the student §2.1.

The paper's claim is that the exact-match objective is "way too overambitious" for these cases, and that preserving the relations between predictions is enough §2.1 §3.

Method

Overview. The student is trained with the usual classification loss plus two relation losses that compare its softmax outputs with the teacher's (which are fixed; the paper calls the teacher pre-trained §2). No feature maps are matched, and no extra layers are added §4.2 app. A.2.

1. Turn each model's logits into per-instance class probabilities with a softmax §2 eq. 2. The paper's notation includes a temperature τ in this step; see Pitfalls for what the code does.
2. Inter-class relation: for each instance, measure the Pearson correlation between the student's and the teacher's C class probabilities, convert it to a distance, and average over the batch of B instances §3.1 eq. 6 eq. 7 eq. 8.
3. Intra-class relation: for each class, measure the Pearson correlation between the student's and the teacher's B instance probabilities for that class, and average over the C classes §3.2 eq. 9.
4. Total loss: α times the classification loss, plus β times the inter-class loss, plus γ times the intra-class loss eq. 10. The weights are given in the Hyperparameters table.

Why Pearson: the paper requires the match to be invariant to positive scale and shift of each prediction vector §3.1 eq. 4 eq. 5. Pearson correlation has this property, and cosine similarity has only the scale part. The paper's ablation favours Pearson over cosine tab. 10.

# DIST relation loss  [§3.1] [§3.2] [eq. 8] [eq. 9] [app. A.1]
# Input:  student_logits, teacher_logits: (B, C)
# Output: scalar relation loss

def dist_loss(student_logits, teacher_logits, beta, gamma):
    # Per-instance class probabilities  [§2] [eq. 2]
    student_probs = softmax(student_logits, dim=1)  # (B, C)
    teacher_probs = softmax(teacher_logits, dim=1)  # (B, C)

    # Inter-class: correlate the C class scores of each instance  [eq. 8]
    inter_loss = 1 - mean(pearson(student_probs, teacher_probs))

    # Intra-class: correlate the B instance scores of each class  [eq. 9]
    intra_loss = 1 - mean(
        pearson(student_probs.transpose(0, 1),
                teacher_probs.transpose(0, 1)))

    # Weighted sum of the two relation terms  [§3.2] [eq. 10]
    return beta * inter_loss + gamma * intra_loss


# Pearson correlation of each matching row pair  [§3.1] [eq. 6] [eq. 7]
def pearson(a, b, eps=1e-8):
    # Centre each row  [code: segmentation/losses/dist_kd.py#L8-L9]
    a_centred = a - mean(a, dim=1, keepdim=True)  # (N, K)
    b_centred = b - mean(b, dim=1, keepdim=True)  # (N, K)

    # Cosine similarity of the centred rows, with eps against zero norms
    # [code: segmentation/losses/dist_kd.py#L4-L5]
    dot = sum(a_centred * b_centred, dim=1)
    norms = norm(a_centred, dim=1) * norm(b_centred, dim=1)
    return dot / (norms + eps)

The total training loss is written in the paper as α·L_cls + β·L_inter + γ·L_intra eq. 10. In the paper, L_cls is the cross-entropy between the student's predictions and the ground-truth labels §2.

Hyperparameters

Loss weights (paper, §4.1):

KD temperature: the paper says the default is τ=4 and that ImageNet uses τ=1 instead of τ=4 §4.1. Table 6's caption gives τ=1 for KD on detection tab. 6. See Pitfalls.

DIST temperature: not stated in the paper. The appendix code applies the softmax to raw logits with no temperature app. A.1 fig. 4.

Training strategies, image classification tab. 1:

Source: tab. 1. The paper's abbreviations are as in its caption. Strategy A1 uses the same pre-trained models as CRD §4.1.

Teachers and students §4.1 §4.2 tab. 2 tab. 4:

• ImageNet baselines: teachers are torchvision ResNet-34 and ResNet-50 tab. 2.
• Stronger-teacher sweep: ResNet-18 and ResNet-34 students with ResNet-34, 50, 101 and 152 teachers, all under B1 tab. 3.
• Stronger strategies: ResNet-50 teacher at 80.1% trained by the cited timm work, Swin-L teacher at 86.3% pre-trained on ImageNet-22K tab. 4.
• Detection: Cascade Mask R-CNN X101 teaches Faster R-CNN R50; RetinaNet X101 teaches RetinaNet R50 §4.3.
• Segmentation: DeepLabV3-R101 (78.07 mIoU) teaches DeepLabV3-R18 and PSPNet-R18 on Cityscapes §4.4 tab. 7.

Detection and segmentation training schedules: the paper says it uses "standard" training strategies following prior work §4.3 and gives no numbers. Not stated in paper. The segmentation code defaults are in Paper vs code.

Results

Author-reported. All numbers are single runs except the CIFAR-100 rows of tab. 5, which carry standard deviations (the number of runs is not stated).

ImageNet, baseline strategy B1 tab. 2 (Top-1 / Top-5 %):

ImageNet, stronger teachers tab. 3 (ResNet-18 student, baseline 69.76 %, top-1):

ResNet-34 student (baseline 73.31 %, top-1):

ImageNet, stronger training strategies tab. 4 (top-1 %; the Swin-T student uses B3, the others B2):

CIFAR-100 tab. 5 (top-1 %, mean ± std). Columns are teacher → student pairs in the order shown in the paper's header:

Setup: strategy A1, teacher and student from the CRD-style pre-trained models §4.1 tab. 1. The paper's checklist says standard deviations are reported for CIFAR-100 only [checklist].

COCO object detection tab. 6 (box AP, COCO validation set):

The DIST + mimic row adds a mean-squared error between the FPN features of teacher and student §4.3. The paper says this term is a feature loss, so the detection results do not isolate DIST.

Cityscapes semantic segmentation tab. 7 (val mIoU %, ImageNet pre-trained; teacher DeepLabV3-R101 at 78.07):

Ablations (ImageNet, ResNet-18 student, ResNet-34 teacher; §4.5 app. A.3):

Source: tab. 8. Training with the KD loss alone, without the classification loss tab. 9:

The student's standalone accuracy is 69.76 tab. 9. Cosine similarity versus Pearson, with and without label smoothing tab. 10:

Training speed, ResNet-18 student with ResNet-34 teacher, 8 V100 GPUs, batches per second tab. 11: KD 14.28, RKD 11.11, SRRL 12.98, CRD 8.33, DIST 14.19.

Paper vs code

The official repository is hunto/DIST_KD (Apache-2.0), linked from the abstract. It contains the segmentation code. The classification code is in a separate repository that the README links as a submodule (image_classification_sota); it was not cloned here. The paper's detection code is in MasKD/mmrazor; it was not cloned either. The clone used for this card is at commit 2418f18.

• Loss function. The paper's appendix listing app. A.1 fig. 4 and segmentation/losses/dist_kd.py implement the same inter-class and intra-class functions dist_kd.py. The code adds a 4-D path that reshapes (B, C, H, W) to per-pixel rows dist_kd.py. The paper does not describe that path.
• Default weights. The segmentation class defaults to β=1, γ=1 dist_kd.py. The paper gives these for detection and segmentation §4.1.
• Temperature. The segmentation code applies the softmax to raw logits dist_kd.py. The --kd-temperature argument is parsed but never used train_kd.py. The external classification code, fetched through WebFetch from hunto/image_classification_sota (lib/models/losses/dist_kd.py) and not verified against a clone, has a temperature τ (default 1.0) and scales both terms by τ² [external: image_classification_sota lib/models/losses/dist_kd.py]. The paper does not say which form its ImageNet and CIFAR results use.
• Auxiliary head (segmentation). For ResNet student backbones the code sets aux = True train_kd.py. The student loss then adds 0.4 times the DIST loss on the auxiliary outputs train_kd.py. The paper does not mention an auxiliary head.
• Segmentation schedule. The scripts pass only teacher, student, backbone, data and paths pspnet_res18.sh. Their values come from the defaults in train_kd.py: 40 000 iterations train_kd.py, learning rate 0.02 with polynomial decay (power 0.9, floor 1e-4) train_kd.py train_kd.py, momentum 0.9 train_kd.py, weight decay 5e-4 train_kd.py, total batch 16 split across GPUs train_kd.py, crop 512×1024 train_kd.py. This is inferred from code; the paper does not state these values.
• CIFAR-100 WRN-40-2 → WRN-40-1. Table 5 gives DIST 74.73±0.24 tab. 5. The repository README gives 74.43±0.24 for the same pair README.md. The other six CIFAR pairs match.

Pitfalls

• The KD temperature is contradictory. §4.1 says the default is τ=4, then says ImageNet uses τ=1 instead §4.1. The repository README labels the ImageNet KD baseline as "KD (T=4)" README.md. The detection table labels KD as τ=1 tab. 6. Check which temperature a reproduction uses for each baseline before comparing numbers.
• The DIST temperature is not stated. Equations 6 to 10 do not include τ. The appendix code and the segmentation code use raw logits. The external classification code includes τ and τ² (see Paper vs code). An implementation that follows the paper's notation literally is not the same as either code.
• The introduction overstates the Swin-T gain. It says 82.3% on Swin-T is "improving KD by 1%" §1. Table 4 gives KD 81.5 and DIST 82.3, a gain of 0.8 points tab. 4. The "1.0%" in §4.2 is measured against the student without distillation (81.3) §4.2.
• Teacher accuracy for ResNet-50. Table 2 gives 76.16 tab. 2. Table 3 and §2.1 give 76.13 tab. 3 §2.1.
• The detection claim applies to two-stage only. The paper says DIST plus mimic outperforms prior detection KD methods §4.3. In the one-stage table, DIST plus mimic reaches 40.1 AP against CWD's 40.8 tab. 6. In the two-stage table it reaches 41.8 against CWD's 41.7 tab. 6.
• Loss weights are not ablated. DIST uses β+γ=4 on ImageNet and CIFAR, while the KD baseline has β=1 §4.1. The ablations vary which relation terms are on, not their weights §4.5 app. A.3. A gain may partly come from the larger total weight on the distillation term. The paper does not separate the two effects.
• Intra-class term depends on batch composition. The intra-class relation correlates the instance scores of each class across the batch §3.2 eq. 9. The segmentation code applies it across all pixels in the batch dist_kd.py. The paper does not analyse small batches or how the batch is built.
• Degenerate rows. The eps guard in the Pearson computation is the only protection against zero-variance rows dist_kd.py. If a row has equal logits, its centred vector is zero, the correlation becomes 0 and that row contributes a loss of 1. The paper does not discuss this case. This is inferred from the code.
• Variance is reported for CIFAR-100 only. ImageNet, detection and segmentation numbers are single values tab. 2 tab. 4 tab. 6 tab. 7. The checklist says error bars are reported for CIFAR-100 only.
• Limitation. The paper says DIST is less effective on binary classification, because the inter-class relation then has only two terms app. A.6.
• Cost. The paper reports training speed close to vanilla KD on one setup tab. 11. It does not report memory use.

Lineage

• Builds on: standard knowledge distillation, which matches softmax outputs with a KL divergence §2 eq. 1 eq. 3.
• Related relation-based methods named in the appendix: CCKD and RKD (instance relations) and a locality-preserving loss app. A.2. Not carded.
• Compared against in the tables: FitNet, VID, RKD, PKT, CRD, OFD, SRRL, Review, FKD, CWD, IFVD, SKD and CIRKD tab. 2 tab. 5 tab. 6 tab. 7. Not carded.
• Superseded by: not known from this wiki.