Method and system for training a neural network model using knowledge distillation
An agnostic combinatorial knowledge distillation (CKD) method for transferring trained knowledge of neural model from a complex model (teacher) to a less complex model (student) is described. In addition to training the student to generate a final output that approximates both the teacher's final output and a ground truth of a training input, the method further maximizes knowledge transfer by training hidden layers of the student to generate outputs that approximate a representation of a subset of teacher hidden layers are mapped to each of the student hidden layers for a given training input.
1 . A method of knowledge distillation from a teacher neural model having a plurality of teacher hidden layers to a student neural model having a plurality of student hidden layers, the method comprising:
training the teacher neural model, wherein the teacher neural model is configured to receive a training input and generate a teacher output for the training input; and
training the student neural model on a plurality of training inputs, wherein the student neural model is configured to receive inputs and generate a corresponding student output, comprising:
processing each training input using the teacher neural model to generate the teacher output for the training input, each of the plurality of teacher hidden layers generating a teacher hidden layer output to obtain a plurality of teacher hidden layer outputs;
mapping a subset of the plurality of teacher hidden layers to each of the plurality of student hidden layers;
calculating, for each of the plurality of student hidden layers, a representation of the teacher hidden layer outputs of the subset of the plurality of teacher hidden layers mapped to each of the plurality of student hidden layers; and
training the student to, for each of the training inputs, generate a student output that approximates the teacher output for the training input, wherein each of the plurality of student hidden layers, for each of the training inputs, is trained to generate a student hidden layer output that approximates the representation of the subset of the plurality of teacher hidden layers mapped to the each of the plurality of student hidden layers,
wherein the mapping comprises, for each of the student hidden layers, assigning an attention weight (ϵ ij ) to the teacher hidden layer output of each of the subset of the plurality of teacher hidden layers,
wherein each attention weight (ϵ ij ) is computed by:
ϵ
i
j
=
e
φ
i
j
∑
k
=
1
k
=
|
H
T
|
e
φ
ik
where
φ
i
j
=
Φ
(
h
i
s
,
h
j
T
)
is an energy score between an output generated by an i th hidden layer of the student
(
h
i
s
)
and an output generated by a j th hidden layer of the teacher
(
h
j
T
)
,
the energy score being indicative of a similarity between the two generated outputs,
φ
i
k
=
Φ
(
h
i
s
,
h
k
T
)
is an energy score between the output generated by the i th hidden layer of the student
(
h
i
s
)
and an output generated by the k th hidden layer of the teacher
(
h
k
T
)
,
the energy score being indicative of a similarity between the two generated outputs,
Φ(h S , h T ) is an energy function of an output generated by a hidden layer of the student and an output generated by a hidden layer of the teacher,
h
k
T
where k∈{1, . . . , |H T |} is a k th layer of the teacher that belongs to a set of all hidden layers of the teacher (H T ), and
|H T | is a size of the set H T representing a total number of hidden layers of the teacher.
2 . The method of claim 1 , further comprising training the student to, for each of the training inputs, generate the student output to approximate a ground truth of the training input.
3 . The method of claim 1 , wherein training the student to generate the student output that approximates the teacher output for the training input further comprises:
computing a knowledge distillation (KD) loss between the student output and the teacher output;
computing a standard loss between the student output and the ground truth;
computing a combinatorial KD (CKD) loss between each of the plurality of student hidden layers and the subset of the teacher hidden layers mapped to each of the plurality of student layer;
calculating a total loss as a weighted average of the KD loss, the standard loss and the CKD loss; and
adjusting parameters of the student to minimize the total loss.
4 . The method of claim 3 , wherein the CKD loss is computed by:
L
CKD
=
∑
h
i
s
∈
H
s
MSE
(
h
i
s
,
f
i
T
)
where
L CKD is the CKD loss,
MSE( ) is a mean-square error function,
h
i
s
is the output generated by the i th hidden layer of the student,
f
i
T
is an output generated by the subset of hidden teacher layers which are associated with or assigned to the i th hidden layer of the student by a mapping function M which is computed by
f
i
T
=
F
(
H
T
(
i
)
)
,
where
H
T
(
i
)
=
{
j
∈
M
(
i
)
}
,
F( ) is a fusion function that fuses/aggregates a subset of hidden teacher layers which are assigned to a particular hidden layer of the student model provided by a first
(
h
1
T
)
and third
(
h
3
T
)
hidden layers of the teacher,
an output of F( ) being mapped to a second hidden layer of the student
(
h
2
s
)
,
H S is a set of all hidden layers of the student
H T is the set of all hidden layers of the teacher,
H T (i) is a subset of hidden layers of the teacher selected mapped to the i th hidden layer of the student, and
M( ) is a mapper function that takes an index referencing to the hidden layer of the student and returns a set of indices for the teacher.
5 . The method of claim 4 , wherein the fusion function F( ) includes a concatenation operation followed by a linear projection layer.
6 . The method of claim 5 , wherein the fusion function F( ) is defined by:
F
(
h
1
T
,
h
3
T
)
=
m
u
l
(
W
,
[
h
1
T
;
h
3
T
]
)
+
b
where
“;” is a concatenation operator,
mul( ) is the matrix multiplication operation, and
W and b are learnable parameters.
7 . The method of claim 4 , wherein the mapping function M( ) defines a combination policy for combining the hidden layers of the teacher.
8 . The method of claim 7 , wherein the combination policy is any one of overlap combination, regular combination, skip combination, and cross combination.
9 . The method of claim 1 , wherein the mapping further comprises, defining a combination policy for mapping the teacher hidden layers to each of the plurality of student hidden layers.
10 . The method of claim 9 , wherein the combination policy is any one of overlap combination, regular combination, skip combination, and cross combination.
11 . The method of claim 3 , wherein the CKD loss is computed by:
L
C
K
D
*
(
H
S
,
H
T
)
=
∑
h
i
s
∈
H
s
MSE
(
h
i
s
,
f
i
*
T
)
where
L CKD* is the CKD loss,
MSE( ) is a mean-square error function,
h
i
s
is the output generated by the i th hidden layer of the student,
f
i
*
T
is a combined attention-based representation of the set of all hidden layers (H T ) of the teacher for an i th hidden layer of the teacher, and
H S is a set of all hidden layers of the student.
12 . The method of claim 11 , wherein the student output and the teacher output have identical dimensions
(
❘
"\[LeftBracketingBar]"
h
i
s
❘
"\[RightBracketingBar]"
=
❘
"\[LeftBracketingBar]"
h
i
T
❘
"\[RightBracketingBar]"
)
,
and
f
i
*
T
is computed by:
f
i
*
T
=
∑
h
j
T
∈
H
*
T
ϵ
i
j
h
j
T
where
ϵ ij is the attention weight, and the attention weight indicates how much the j th hidden layer of the teacher
(
h
j
T
)
contributes to the knowledge distillation process of the i th hidden layer of the student
(
h
j
s
)
,
h
j
T
is a j th hidden layer of the teacher, and
H* T is the set of all hidden layers from the teacher that is assigned to the i th hidden layer of the student.
13 . The method of claim 11 , wherein the student output and the teacher output have different dimensions
(
❘
"\[LeftBracketingBar]"
h
i
s
❘
"\[RightBracketingBar]"
≠
❘
"\[LeftBracketingBar]"
h
j
T
❘
"\[RightBracketingBar]"
)
,
and
f
i
*
T
is computed by:
f
i
*
T
=
∑
h
j
T
∈
H
T
ϵ
i
j
(
W
i
h
j
T
)
where
ϵ ij is the attention weight, and the attention weight indicates how much the j th hidden layer of the teacher
(
h
j
T
)
contributes to the knowledge distillation process of the i th hidden layer of the student
(
h
i
s
)
,
h
j
T
is a j th hidden layer of the teacher,
H T is the set of all hidden layers of the teacher that is assigned to the i th hidden layer of the student, and
W
i
∈
R
|
h
t
s
|
×
|
h
j
T
|
and is a weight value for the i th hidden layer of the teacher.
14 . The method of claim 13 , wherein a sum of all attention weights ϵ ij is 1.
15 . The method of claim 1 , wherein the energy function
Φ
(
h
i
s
,
h
j
T
)
is computed as the dot product of the output of the i th hidden layer of the student
(
h
i
s
)
and the output generated by the j th hidden layer of the teacher
(
h
j
T
)
by:
Φ
(
h
i
s
,
h
j
T
)
≡
<
h
i
s
,
h
j
T
>
.
16 . The method of claim 15 , wherein the energy function
Φ
(
h
i
s
,
h
j
T
)
is computed as the dot product of the output of the i th hidden layer of the student
(
h
i
s
)
and a weighted value of the output generated by the j th hidden layer of the teacher
(
W
i
h
j
T
)
by:
Φ
(
h
i
s
,
h
j
T
)
≡
<
h
i
s
,
W
i
h
j
T
>
.
17 . A method of knowledge distillation from a plurality of teacher neural models each having a plurality of teacher hidden layers to a student neural model having a plurality of student hidden layers, the method comprising:
inferring the plurality of teacher neural models, wherein each of the plurality of the teacher neural models is configured to receive an input and generate a teacher output; and
training the student neural model on a plurality of training inputs, wherein the student neural model is configured to receive inputs and generate a student output, comprising:
processing each training input using the plurality of teacher neural models to generate a plurality of teacher outputs for the training input, each of the plurality of teacher hidden layers of each of the plurality of teacher neural models generating a teacher hidden layer output;
mapping a subset of the plurality of teacher hidden layers of the plurality of teacher neural models to each of the plurality of student hidden layers;
calculating a representation of the teacher hidden layer outputs of the subset of the plurality of teacher hidden layers mapped to each of the plurality of student hidden layers; and
training the student neural model to generate a student output that approximates the plurality of teacher outputs for the training input, wherein each of the plurality of student hidden layers, for each of the training inputs, is trained to generate a student hidden layer output that approximates the representation of the subset of the plurality of teacher hidden layers mapped to the each of the plurality of student hidden layers,
wherein the mapping comprises, for each of the student hidden layers, assigning an attention weight (ϵ pq ) to the teacher hidden layer outputs of each of the subset of the plurality of teacher hidden layers,
wherein each attention weight (ϵ pq ) is computed by:
∈
pq
=
e
φ
pq
∑
i
=
1
K
e
φ
iq
where
φ
pq
=
Φ
(
T
p
(
x
i
q
)
,
S
(
x
i
q
)
)
;
x
i
q
∈
D
q
is an energy score between an inferred output of a p th teacher T p , represented by
T
p
(
x
i
q
)
,
and an inferred output of a student, represented by
S
(
x
i
q
)
,
the energy score being indicative of a similarity between the respective inferred outputs,
φ iq is an energy score between an inferred output of an i th teacher and the inferred output of the student, where is K is a total number of teachers,
Φ
(
T
p
(
x
i
q
)
,
S
(
x
i
q
)
)
is an energy function of the inferred output of T p and the inferred output of S,
training data samples
{
(
x
i
q
)
}
i
=
1
N
q
of training dataset D q are sent to T p and S to obtain the respective inferred outputs, where D q is the training dataset for the K number of teachers, and
N q is a number of training data samples in the training dataset D q for the q th teacher.
18 . The method of claim 17 , wherein training the student neural model to generate a student output that approximates the plurality of teacher outputs for the training input further comprises computing a weighted knowledge distillation (KD) loss for each teacher neural model, where the weighted KD loss is computed by:
L
K
D
*
q
=
∑
p
=
1
p
=
K
∑
x
i
q
∈
D
q
ϵ
p
q
L
K
D
(
T
p
(
x
i
q
)
,
S
(
x
i
q
)
)
where
L
K
D
*
q
is the weighted KD loss,
T
p
(
x
i
q
)
is the inferred output of the p th teacher T p ,
S
(
x
i
q
)
is the inferred output of the student S,
L KD is a loss function, and
∈ pq is the weight value of the teacher hidden layers of each p th teacher T p of the plurality of teacher neural models.
19 . The method of claim 17 , wherein
T
p
(
x
i
q
)
and
S
(
x
i
q
)
are vectors and the
Φ
(
T
p
(
x
i
q
)
,
S
(
x
i
q
)
)
function is a dot product of the
T
p
(
x
i
q
)
vector and the
S
(
x
i
q
)
vector.
20 . The method of claim 17 , wherein
T
p
(
x
i
q
)
and
S
(
x
i
q
)
are vectors and the
Φ
(
T
p
(
x
i
q
)
,
S
(
x
i
q
)
)
function is a neural network suitable for measuring a similarity of the
T
p
(
x
i
q
)
vector and the
S
(
x
i
q
)
vector.