System and method for knowledge distillation between neural networks
Systems and methods for knowledge distillation provide supervised training of a student network with a teacher network, including inputting a batch to the teacher network, inputting the batch to the student network, generating a teacher activation map at a layer of the teacher network, generating a student activation map at a layer of the student network corresponding to the layer of the teacher network, generating a pairwise teacher similarity matrix based on the teacher activation map, generating a pairwise student similarity matrix based on the student activation map, and minimizing a knowledge distillation loss defined as a difference between the pairwise teacher similarity matrix and the pairwise student similarity matrix.
1. A knowledge distillation system for supervised training of a computer-implemented student network with a computer-implemented teacher network, the system comprising:
a processor;
a memory in communication with the processor, the memory storing instructions that, when executed by the processor cause the processor to:
input a batch to the teacher network;
input the batch to the student network;
generate a teacher activation map at a layer of the teacher network;
generate a student activation map at a layer of the student network corresponding to the layer of the teacher network;
generate a pairwise teacher similarity matrix based on the teacher activation map;
generate a pairwise student similarity matrix based on the student activation map;
minimize a knowledge distillation loss defined as a difference between the pairwise teacher similarity matrix and the pairwise student similarity matrix, and
train the student network based on the knowledge distillation loss.
2. The knowledge distillation system of claim 1 , wherein the pairwise teacher similarity matrix is generated based at least in part on dot products of the teacher activation map.
3. The knowledge distillation system of claim 1 , wherein the pairwise student similarity matrix is generated based at least in part on dot products of the student activation map.
4. The knowledge distillation system of claim 1 , wherein the knowledge distillation loss is defined as a summation of a mean element-wise squared difference between the pairwise teacher similarity matrix and the pairwise student similarity matrix.
5. The knowledge distillation system of claim 1 , wherein the knowledge distillation loss penalizes differences in the L2-normalized outer products of the teacher network activation map and the student network activation map.
6. The knowledge distillation system of claim 1 , wherein the knowledge distillation loss L SP is based on:
ℒ
SP
(
G
T
,
G
S
)
=
1
b
2
∑
(
l
,
l
′
)
∈
𝒥
G
T
(
l
)
-
G
S
(
l
′
)
F
2
where b is a size of the batch, G T is the pairwise teacher similarity matrix, G S is the pairwise student similarity matrix, collects the (l,l′) layer pairs of layer l of the teacher network and layer l′ of the student network, and ∥⋅∥ F is a Frobenius norm.
7. The knowledge distillation system of claim 1 , wherein the memory stores instructions that, when executed by the processor cause the processor to: minimize a total loss for training the student network, the total loss defined as a cross-entropy loss and a product of the knowledge distillation loss and a hyperparameter.
8. The knowledge distillation system of claim 1 , wherein the teacher network is trained with privileged data that is inaccessible to the student network.
9. The knowledge distillation system of claim 1 , wherein the layer of the teacher network and the layer of the student network correspond to a same depth.
10. The knowledge distillation system of claim 1 , wherein the layer of the teacher network and the layer of the teacher network are at an end of a same block.
11. A method for supervised training of a computer-implemented student network with a computer-implemented teacher network, the method comprising:
inputting a batch to the teacher network;
inputting the batch to the student network;
generating a teacher activation map at a layer of the teacher network;
generating a student activation map at a layer of the student network corresponding to the layer of the teacher network;
generating a pairwise teacher similarity matrix based on the teacher activation map;
generating a pairwise student similarity matrix based on the student activation map;
minimizing a knowledge distillation loss defined as a difference between the pairwise teacher similarity matrix and the pairwise student similarity matrix, and
training the student network based on the knowledge distillation loss.
12. The method of claim 11 , wherein the generating the pairwise teacher similarity matrix is based at least in part on dot products of the teacher activation map.
13. The method of claim 11 , wherein the generating the pairwise student similarity matrix is based at least in part on dot products of the student activation map.
14. The method of claim 11 , wherein the knowledge distillation loss is defined as a summation of a mean element-wise squared difference between the pairwise teacher similarity matrix and the pairwise student similarity matrix.
15. The method of claim 11 , wherein the knowledge distillation loss penalizes differences in the L2-normalized outer products of the teacher network activation map and the student network activation map.
16. The method of claim 11 , wherein the knowledge distillation loss L SP is based on:
ℒ
SP
(
G
T
,
G
S
)
=
1
b
2
∑
(
l
,
l
′
)
∈
𝒥
G
T
(
l
)
-
G
S
(
l
′
)
F
2
where b is a size of the batch, G T is the pairwise teacher similarity matrix, G S is the pairwise student similarity matrix, collects the (l,l′) layer pairs of layer 1 of the teacher network and layer l′ of the student network, and ∥⋅∥ F is a Frobenius norm.
17. The method of claim 11 , further comprising: minimizing a total loss for training the student network, the total loss defined as a cross-entropy loss and a product of the knowledge distillation loss and a hyperparameter.
18. The method of claim 11 , wherein the teacher network is trained with privileged data that is inaccessible to the student network.
19. The method of claim 11 , wherein the layer of the teacher network and the layer of the student network correspond to a same depth.
20. A non-transitory computer readable medium comprising a computer readable memory storing computer executable instructions thereon that when executed by a computer cause the computer to perform the method of claim 11 .