Decoupled Knowledge Distillation

In plain words

Teaching a small model by copying a big model's answers works, and the answers that matter are the big model's scores for every possible class. The usual recipe mixes two kinds of information into one term: how sure the teacher is about the correct class, and how it ranks the wrong classes. This paper pulls them apart. Its main finding is that the wrong-class part carries most of the useful hints, and that the usual mix quietly shrinks that part whenever the teacher is confident. The fix is a new loss with a separate knob for each part.

The new loss costs about the same to compute as the classic one and adds no network pieces, unlike methods that copy the teacher's internal features. The authors report that it matches or beats those feature methods on image classification. On object detection it does not do so alone, and it works best combined with a feature method. It is a practical option for anyone who wants to distill a model cheaply and has the teacher's output scores.

TL;DR

The classic KD loss can be rewritten as a target-class term plus a non-target-class term multiplied by (1 - p_t), where p_t is the teacher's probability for the correct class §3.1 eq. 6. DKD removes that multiplier and gives each term its own weight, alpha and beta §3.3 eq. 7. It adds no parameters and costs about the same as KD §4.2. On CIFAR-100 and ImageNet it beats KD and is comparable to or better than several feature-based methods §4.1; on MS-COCO, DKD alone trails feature methods §4.1. All numbers are author-reported.

Problem

Logit distillation (matching the student's output scores to the teacher's) is cheap, but the authors note that it was generally weaker than feature distillation, which needs extra modules and memory during training §1. They argue the classic KD loss is a coupled formulation with two defects §1 §3.1:

• Suppression: the non-target term is weighted by (1 - p_t), so a teacher that is confident on the correct class shrinks the term that carries the most useful information §3.2 §3.3.
• No balance: the target-class and non-target-class terms share one weight, so their relative importance cannot be tuned §3.3.

The derivation omits the softmax temperature T, which the method then uses §3.1 footnote.

Method

Notation. Class probabilities come from a softmax over logits z eq. 1. Two views are used: a two-way "binary" distribution b = [p_t, p_\t] (target class vs. all others) §3.1, and the distribution over the C-1 non-target classes, renormalised without the target eq. 2.

Reformulation. Splitting the classic KL loss gives eq. 5:

\mathrm{KD} = \mathrm{KL}(\mathbf{b}^{\mathcal{T}} \| \mathbf{b}^{\mathcal{S}}) + (1 - p_t^{\mathcal{T}})\,\mathrm{KL}(\hat{\mathbf{p}}^{\mathcal{T}} \| \hat{\mathbf{p}}^{\mathcal{S}})

Writing the first term as TCKD and the second as NCKD gives eq. 6:

\mathrm{KD} = \mathrm{TCKD} + (1 - p_t^{\mathcal{T}})\,\mathrm{NCKD}

The full derivation is in app. A.1.

Findings that motivate the method §3.2:

1. Each term on its own: TCKD alone can hurt the student, while NCKD alone is comparable to or better than full KD tab. 1.
2. TCKD helps more on harder training data: with strong augmentation tab. 2, with label noise tab. 3, and on ImageNet tab. 4. The paper reads TCKD as carrying knowledge about how hard each sample is §3.2.
3. Using NCKD only on the more confidently predicted half of the training set gives a higher student accuracy than using it on the other half tab. 5. The paper reads this as the suppression described above §3.2.

DKD §3.3 eq. 7. The coupled factor (1 - p_t) is replaced by a constant beta, and TCKD and NCKD get separate weights:

\mathrm{DKD} = \alpha\,\mathrm{TCKD} + \beta\,\mathrm{NCKD}

The per-sample KL terms are computed in the same way as in KD. Algorithm 1 is a PyTorch-style sketch of the same loss alg. 1. The pseudocode below follows it and the official code (see the Paper vs code section).

# DKD loss for one batch  [§3.3] [eq. 7] [alg. 1]
# Input:  student_logits, teacher_logits: (B, C); target: (B,) class ids
# Output: scalar loss
# kl(p, q) = KL(p || q): p is the teacher, q is the student

def dkd_loss(student_logits, teacher_logits, target, alpha, beta, T):
    # Temperature-softened class probabilities  [alg. 1]
    p_student = softmax(student_logits / T)              # (B, C)
    p_teacher = softmax(teacher_logits / T)              # (B, C)
    # Two-way view: target class vs. all others  [§3.1] [eq. 5]
    pt_student = pick(p_student, target)                 # (B,)
    pt_teacher = pick(p_teacher, target)                 # (B,)
    binary_student = stack(pt_student, 1 - pt_student)   # (B, 2)
    binary_teacher = stack(pt_teacher, 1 - pt_teacher)   # (B, 2)
    # TCKD: per-sample KL on the two-way view  [eq. 5] [alg. 1]
    tckd = kl(binary_teacher, binary_student)            # (B,)
    # Non-target classes only, renormalised  [§3.1] [eq. 2]
    student_others = drop_target(student_logits, target) / T   # (B, C-1)
    teacher_others = drop_target(teacher_logits, target) / T   # (B, C-1)
    # NCKD: per-sample KL on non-target classes; no (1 - p_t) factor  [§3.3]
    nckd = kl(softmax(teacher_others), softmax(student_others))  # (B,)
    # Weight the two parts separately, average over the batch, then
    # scale by T^2  [alg. 1]
    loss = alpha * tckd + beta * nckd
    return mean(loss) * T ** 2

Training objective. The paper trains with cross-entropy on the labels plus the DKD term, and ramps the DKD term up linearly at the start of training app. A.2:

# Training objective for CIFAR-100 runs  [app. A.2]
# Input: image, target, epoch (the trainer counts epochs from 1)

def training_loss(student, teacher, image, target, epoch):
    # Teacher is frozen; no gradients flow into it
    # [code: mdistiller/distillers/DKD.py#L66-L68]
    student_logits = student(image)
    teacher_logits = teacher(image)
    # Cross-entropy on the labels, weight 1.0  [app. A.2]
    loss_ce = cross_entropy(student_logits, target)
    # Linear ramp of the DKD term over the first 20 epochs  [app. A.2]
    # [code: mdistiller/distillers/DKD.py#L72]
    ramp = min(epoch / 20, 1.0)
    loss_dkd = ramp * dkd_loss(
        student_logits, teacher_logits, target, alpha, beta, T
    )
    return loss_ce + loss_dkd

Key values: alpha = 1.0 and T = 4 on CIFAR-100; beta depends on the teacher (see Hyperparameters). The derivation omits T §3.1 footnote, and the experiments use T = 4 on CIFAR-100 and T = 1 on ImageNet and detection app. A.2.

Hyperparameters

Values are from the paper's supplement unless marked "code". Where the paper is silent, the table says "not stated".

Results

All results are author-reported. Accuracy is top-1 (%) on the validation split unless stated. Differences are as printed in the paper, or computed from the table values where marked.

CIFAR-100, same teacher and student architecture (average of 5 runs) tab. 6

The text reports 1 to 2 points over KD for same-series pairs §4.1.

CIFAR-100, different architectures (average of 5 runs) tab. 7

The text reports 2 to 3 points over KD for different-series pairs §4.1.

Ablation of alpha and beta (ResNet32x4 teacher, ResNet8x4 student, CIFAR-100) §4.1

This table has no caption in the paper, so it is cited by section.

The coupled form (beta replaced by 1 - p_t, i.e. KD) gives 73.63 §4.1.

The paper says TCKD's contribution is stable for alpha around 1.0 §4.1.

Reformulation ablation (CIFAR-100, weights all 1.0, see Pitfalls) tab. 1

Deltas are over the student baseline. "TCKD + NCKD" is the classic KD form.

Harder training data (NCKD is used in all rows; the TCKD column is ticked where it is added) §3.2:

• AutoAugment, teacher ResNet32x4 at 81.29% tab. 2: ResNet8x4 73.82 to 75.33 (+1.51) with TCKD; ShuffleNet-V1 77.13 to 77.98 (+0.85) with TCKD.
• Symmetric label noise on the training set, teacher ResNet32x4 and student ResNet8x4 tab. 3: noise 0.1, 70.99 to 70.96 (-0.03); noise 0.2, 67.55 to 68.03 (+0.48); noise 0.3, 64.62 to 65.26 (+0.64).
• ImageNet, teacher ResNet-34 and student ResNet-18 tab. 4: 70.71 to 71.03 (+0.32).
• NCKD only, split by teacher confidence p_t tab. 5: all samples 74.26; the "0-50%" subset 74.23; the "50-100%" subset 73.96. The text reads "0-50%" as the top half by p_t §3.2.

CIFAR-100, teacher larger than the student (average of 5 runs) tab. 11 tab. 12

Students are WRN-16-2 in both tables tab. 11 tab. 12. The table caption for the second says "different series", but it includes WRN-16-4, which is a WRN.

ImageNet, ResNet-34 to ResNet-18 (average of 3 runs) tab. 8

ImageNet, ResNet-50 to MobileNet-V1 (average of 3 runs) tab. 9

Against KD: +1.04 top-1 (ResNet-34 pair) and +3.47 (ResNet-50 pair), computed from the table values.

MS-COCO, Faster R-CNN with FPN (AP on val2017, average of 3 runs) tab. 10

DKD alone beats KD on every pair. The paper says the best detection results come from combining DKD with a feature method §4.1.

Other reported results

• Transferability, linear probe on features from WRN-16-2 tab. 13: STL-10 accuracy, baseline 69.7, KD 70.9, FitNet 70.3, CRD 71.6, ReviewKD 72.4, DKD 72.9. Tiny-ImageNet, baseline 33.7, KD 33.9, FitNet 33.5, CRD 35.6, ReviewKD 36.6, DKD 37.1.
• Training cost: the paper reports DKD as having about the same complexity as KD, with no extra parameters, and a better speed/accuracy trade-off than the feature methods §4.2 fig. 2. The exact timings are only in the figure, and this card does not extract them.
• Visualisations: t-SNE features are more separable for DKD than for KD fig. 3; the difference between student and teacher logit correlation matrices is smaller for DKD fig. 4 §4.2.

Paper vs code

Official repository: megvii-research/mdistiller (MEGVII Research), linked in the abstract and in the paper. Cloned at commit a08d46f10d6102bd6e3f258ca5ac880b020ea259, MIT licence.

Discrepancies found:

• ImageNet alpha. The paper sets alpha = 0.5 for all ImageNet experiments app. A.2. The ImageNet DKD configs do not set alpha, so they take the library default of 1.0 dkd.yaml dkd.yaml cfg.py. Unless a run overrides it, the code does not match the paper here.
• Warmup scope. The paper's CIFAR-100 section says a 20-epoch linear warmup is used "for all experiments" app. A.2. The code applies a 20-epoch ramp to the DKD term on CIFAR-100 cfg.py DKD.py. On ImageNet the config uses a 1-epoch ramp dkd.yaml. The detection code adds no ramp rcnn.py.

Checked and matching:

• CIFAR-100 values match the paper: batch 64, 240 epochs, LR 0.05 or 0.01 by student, decay at 150/180/210, momentum 0.9, weight decay 5e-4, T = 4, alpha = 1.0 res32x4_res8x4.yaml res32x4_shuv1.yaml.
• Per-teacher beta on CIFAR-100 matches app. A.3: 2.0 for ResNet56 and ResNet110 res56_res20.yaml; 6.0 for WRN-40-2 and VGG13 vgg13_mv2.yaml; 8.0 for ResNet32x4 and ResNet50 (the library default, cfg.py).
• ImageNet beta (0.5 and 2.0), T = 1, batch 512, LR 0.2, 100 epochs and the 30/60/90 decay match app. A.2 dkd.yaml dkd.yaml.
• Detection alpha = 1.0, beta = 0.25, T = 1 match app. A.2 [code: detection/model/config.py#L18-L20]. DKD is applied to the R-CNN head on the same sampled proposals, with teacher features from the teacher's backbone, as described in app. A.6 rcnn.py.
• The loss matches Algorithm 1: alpha TCKD + beta NCKD, times T^2 alg. 1 DKD.py.

Implementation details the paper does not show:

• The code computes NCKD by subtracting 1000 from the target logit (after dividing by T) before the softmax DKD.py. This is a numerical stand-in for removing the target class. It is exact only when the target logit is far below the others on the T-scaled scale, which holds in practice but is not guaranteed.
• The code's KL uses size_average=False and divides by the batch size DKD.py. This is the batch mean of the per-sample KL, which the paper's pseudocode does not spell out.

Pitfalls

• Warmup is a loss ramp, not an LR warmup, in code. The epoch counter starts at 1 trainer.py, so epoch 1 gives a DKD weight of 1/20 on CIFAR-100. Whether an implementer should ramp the learning rate instead is not stated.
• Beta is per teacher and has no general rule. The paper's suggested guide is the gap z_t - z_max between the target logit and the largest non-target logit, averaged over training samples app. A.3. The fit uses a few teachers and is called a conjecture app. A.3 §5. The paper says the relation between performance and beta is "not fully investigated" §5. Beta ranges from 0.5 (ImageNet same-series) to 8.0 across the configs above.
• TCKD alone is unreliable at T = 4. The paper finds TCKD alone can hurt, for example -3.87 on ResNet8x4 tab. 1. It attributes this to a large gradient on non-target logits at high T. With TCKD alone, T = 1 gives 73.24 on WRN-40-2 to WRN-16-2 (vs 73.26 baseline) app. A.5 tab. A.3. The TCKD-only rows are ablations, not a recommended setting.
• Loss weights in the analysis are non-standard. The §3 experiments set all loss weights to 1.0 (CE, KD, TCKD, NCKD). The CRD default that the paper cites uses 0.1 CE and 0.9 KD app. A.4. Results in tab. 1 are not directly comparable to CRD-default KD numbers.
• Baseline KD numbers come from elsewhere. Compared-method results are "reported in their original papers or reproduced by previous works" app. A.2. Only the ImageNet "KD (authors' run)" row is an authors' run tab. 8 tab. 9. Expect a fresh KD run to differ from the quoted KD rows.
• Mask-based NCKD. The -1000 offset in the code is a numerical trick DKD.py. The TCKD term uses log of a softmax, not log_softmax DKD.py. A very small target probability can underflow to zero there, and the code has no epsilon guard.
• Detection needs features. The paper finds DKD alone is clearly below ReviewKD on detection, because logits carry no localisation signal §4.1 §5. It suggests combining DKD with feature distillation §4.1.
• Efficiency claims rely on a figure. The cost comparison fig. 2 is a plot, and the text gives no numbers. Its caption says training time is per batch fig. 2.
• Several numbers are in the supplement. The transfer probe, the detection setup and the beta table are in the appendix, not the main text. Read app. A.2 to app. A.7 before reproducing.

Lineage

• Builds on: classic knowledge distillation (Hinton et al.), cited as [kd] in the paper §1 §3.1. DKD keeps the same teacher-student setup and changes only the logit loss. The classic formula is not compiled in this wiki.
• Compared against: logit methods KD (and a re-run KD on ImageNet) tab. 8 tab. 9; feature methods FitNet, RKD, CRD, OFD, ReviewKD, AT and FGFI, as listed in the result tables tab. 6 tab. 7 tab. 8 tab. 9 tab. 10 tab. 13. Related work also covers teacher-assistant and early-stopped teachers §2 §4.2.
• Superseded by: none recorded in this wiki.