Distillation Losses

Knowledge distillation loss functions

Overview

This module provides loss functions for knowledge distillation. These losses enable training a smaller “student” network to mimic a larger “teacher” network.

Loss Categories: - Output-based: SoftTarget, Logits, Mutual - compare final predictions - Feature-based: Attention, FitNet, Similarity, ActivationBoundaries - compare intermediate representations

Output-Based Losses

These losses compare the final output predictions between student and teacher networks.

Every output-based loss takes pred and teacher_pred, student and teacher logits (B, C).

source

SoftTarget

def SoftTarget(
    pred:torch.Tensor, teacher_pred:torch.Tensor, T:float=5, # Temperature for softening
    **kwargs
)->torch.Tensor:

Knowledge distillation with softened distributions (Hinton et al.)

pred, teacher_pred = torch.randn(4, 10), torch.randn(4, 10)
SoftTarget(pred, teacher_pred)

source

Logits

def Logits(
    pred:torch.Tensor, teacher_pred:torch.Tensor, **kwargs
)->torch.Tensor:

Direct logit matching between student and teacher

source

Mutual

def Mutual(
    pred:torch.Tensor, teacher_pred:torch.Tensor, **kwargs
)->torch.Tensor:

KL divergence between student and teacher

source

DecoupledKD

def DecoupledKD(
    pred:torch.Tensor, teacher_pred:torch.Tensor, T:float=4, # Temperature for softening
    alpha:float=1.0, # Weight for target-class KD (TCKD)
    beta:float=8.0, # Weight for non-target-class KD (NCKD)
    target:torch.Tensor | None=None, # Ground-truth labels (B,)
    normalize:bool=False, # NKD mode: softmax only over non-target classes
    **kwargs
)->torch.Tensor:

Decoupled Knowledge Distillation (Zhao et al. CVPR 2022). With normalize=True: Normalized KD (Yang et al. ICCV 2023).

Feature-Based Losses

These losses compare intermediate feature representations, enabling the student to learn internal representations similar to the teacher.

Every feature-based loss takes fm_s and fm_t, dicts {layer_name: feature map} for student and teacher, zipped in order.

source

Attention

def Attention(
    fm_s:dict[str, torch.Tensor], fm_t:dict[str, torch.Tensor], p:int=2, # Power for attention computation
    **kwargs
)->torch.Tensor:

Attention transfer loss (Zagoruyko & Komodakis)

fm_s = {'block': torch.randn(4, 8, 16, 16)}
fm_t = {'block': torch.randn(4, 16, 16, 16)}  # attention pools the channels away, so the widths may differ
Attention(fm_s, fm_t)

source

ActivationBoundaries

def ActivationBoundaries(
    fm_s:dict[str, torch.Tensor], fm_t:dict[str, torch.Tensor], m:float=2, # Boundary margin
    **kwargs
)->torch.Tensor:

Boundary-based knowledge distillation (Heo et al.)

source

FitNet

def FitNet(
    fm_s:dict[str, torch.Tensor], fm_t:dict[str, torch.Tensor], **kwargs
)->torch.Tensor:

FitNets: direct feature map matching (Romero et al.)

source

Similarity

def Similarity(
    fm_s:dict[str, torch.Tensor], fm_t:dict[str, torch.Tensor],
    pred:torch.Tensor, # Student predictions (unused, for API consistency)
    p:int=2, # Normalization power
    **kwargs
)->torch.Tensor:

Similarity-preserving knowledge distillation (Tung & Mori)

See Also

Loss Selection Guide

Loss Type Best For Complexity
SoftTarget Output General distillation, logit matching Low
DecoupledKD Output Fine-grained logit distillation, dark knowledge emphasis Medium
DecoupledKD normalize=True Output Improved dark knowledge transfer (NKD, ICCV 2023) Medium
Logits Output Direct logit regression Low
Mutual Output KL divergence matching Low
Attention Feature When attention patterns matter Low
FitNet Feature Intermediate feature matching Medium
Similarity Feature Relational knowledge transfer Medium
ActivationBoundaries Feature Boundary-aware matching Medium

Tests live in nbs/tests/test_losses.ipynb.