Knowledge Distillation: Teaching a Small Network to Mimic a Large One

Knowledge distillation is a common method for improving the training accuracy of a neural network and speeding up convergence. It is frequently used to fine-tune models after compression, and to train small networks in NAS.

The idea is straightforward: treat the outputs of some high-accuracy teacher model as labels for a student model, letting the student mimic the teacher’s behavior.

The knowledge distillation process — Knowledge distillation: use the teacher network’s soft outputs to teach the student network.

Classic Distillation: Soft Labels and Temperature

The most common implementation adds a distillation loss between the teacher’s and student’s outputs to the student network’s own classification loss, forming a new loss. Let $\mathcal{H}$ be the cross-entropy, $T$ the softmax temperature, $z_s$ the student output, $z_t$ the teacher output, $y$ the label, and $\alpha,\beta$ the weights of the two loss terms:

\mathcal{L}(x;W) = \alpha\,\mathcal{H}\big(y,\ \sigma(z_s; T{=}1)\big) + \beta\,\mathcal{H}\big(\sigma(z_t; T{=}\tau),\ \sigma(z_s; T{=}\tau)\big)

The first term is the student’s ordinary classification loss against the ground-truth labels; the second makes the student’s soft output approximate the teacher’s soft output. The temperature-scaled softmax is:

\sigma(z^i; T)=\frac{\exp(z^i/T)}{\sum_j \exp(z^j/T)}

When $T=1$ , this is the ordinary softmax; when $T>1$ , the probability distribution over classes becomes smoother and easier for the student to mimic—these are the soft labels and soft predictions in the figure above.

Distilling at Intermediate Layers: FitNets and FSP

Beyond transferring knowledge at the final output layer, FitNets proposes hints training, transferring knowledge between an intermediate layer of the teacher and the student networks to give the student “hints.”

Transferring knowledge at an intermediate layer — FitNets: transferring knowledge between intermediate layers of the teacher and student.

The hint loss is as follows, where $r$ is an extra layer used to match dimensions, $W_r$ are its parameters, and $u_h$ and $v_g$ are the intermediate-layer outputs of the teacher and student respectively:

\mathcal{L}_{HT}(W_{Guided}, W_r)=\tfrac{1}{2}\big\lVert u_h(x; W_{Hint})-r\big(v_g(x; W_{Guided}); W_r\big)\big\rVert^2

After hints training, the whole student network is then trained properly through output-layer distillation.

W_s^{*} = \underset{W_s}{\arg\min}\ \mathcal{L}_{KD}(W_s)

After hints training, perform another round of distillation at the output layer.

Another work (A Gift from Knowledge Distillation) proposes transferring knowledge using the FSP matrix—each element of the FSP matrix is obtained from the inner product between two layers’ feature maps.

Computing the FSP matrix — FSP matrix: formed from inner products between feature maps.

Similar to FitNets, it also transfers knowledge at intermediate layers, but across multiple layers (rather than just a single intermediate layer). This method suits scenarios where the teacher and student have similar structures but different depths, ensuring the FSP matrices have matching dimensions so that an L2 loss can be computed.

Transferring knowledge across multiple layers — Transferring FSP matrices across multiple layers.

The Teacher Assistant: Keeping the Gap Manageable

Given that the gap between teacher and student may be too large for knowledge to transfer well, TAKD inserts a teacher assistant between the two: knowledge first flows from the teacher to the assistant, then from the assistant to the student. This extra intermediate step makes the transfer smoother.

Using a teacher assistant to make distillation smoother — TAKD: adding a teacher assistant to make the teacher→student knowledge transfer smoother.

Learning from Each Other: DML

All of the above have the student learn from the teacher in one direction. DML (Deep Mutual Learning) instead has a group of networks learn from each other: two differing network architectures complement one another during training, each achieving better results. During training, each network’s loss includes not only the cross-entropy between its output and the labels, but also the KL divergence with the other network’s output.

Algorithm: Deep Mutual Learning

Input:      Training set X, label set Y, learning rates γ_{1,t} and γ_{2,t}.
Initialize: Models Θ_1 and Θ_2 to different initial conditions.
Repeat:
    t = t + 1
    Randomly sample data x from X.
    1: Update the predictions p_1 and p_2 of x for the current mini-batch.
    2: Compute the stochastic gradient and update Θ_1:
           Θ_1 ← Θ_1 + γ_{1,t} · ∂L_{Θ_1} / ∂Θ_1
    3: Update the predictions p_1 and p_2 of x for the current mini-batch.
    4: Compute the stochastic gradient and update Θ_2:
           Θ_2 ← Θ_2 + γ_{2,t} · ∂L_{Θ_2} / ∂Θ_2
Until: convergence

DML: each of the two networks serves as the other’s “soft labels,” with a loss combining cross-entropy and the mutual KL divergence.

The losses $L_{\Theta_1}$ and $L_{\Theta_2}$ of the two networks are defined as:

L_{\Theta_1}=L_{C_1}+D_{KL}(p_2\,\Vert\,p_1),\qquad L_{\Theta_2}=L_{C_2}+D_{KL}(p_1\,\Vert\,p_2)

where $L_{C_1}$ is the cross-entropy loss and $D_{KL}$ is the KL divergence (Kullback-Leibler divergence):

D_{KL}(p_2\,\Vert\,p_1)=\sum_{i=1}^{N}\sum_{m=1}^{M}p_2^m(x_i)\log\frac{p_2^m(x_i)}{p_1^m(x_i)}

References

Hinton, Geoffrey, Vinyals, Oriol, Dean, Jeff. Distilling the Knowledge in a Neural Network. arXiv:1503.02531, 2015.
Gou, Jianping, et al. Knowledge Distillation: A Survey. arXiv:2006.05525, 2020.
Romero, Adriana, et al. FitNets: Hints for Thin Deep Nets. arXiv:1412.6550, 2014.
Yim, Junho, et al. A Gift from Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning. CVPR, 2017.
Mirzadeh, Seyed Iman, et al. Improved Knowledge Distillation via Teacher Assistant. AAAI, 2020.
Zhang, Ying, et al. Deep Mutual Learning. CVPR, 2018.

Technology

2021 · 01 · 30