System and method for knowledge-preserving neural network pruning

1. A computer-implemented method, comprising:

obtaining a pre-trained machine learning model trained based on a plurality of general-purpose training data;

training a task-specific machine learning model by tuning the pre-trained machine learning model based on a plurality of task-specific training data corresponding to a task, wherein the trained task-specific machine learning model is a teacher network;

constructing a student network based on the teacher network; and

executing two or more rounds of fine-tuning of the student network, each round of fine-tuning comprising performing:

network pruning on the student network, and

knowledge distillation from the teacher network to the pruned student network based on a plurality of loss functions corresponding to one or more layers of the teacher network, wherein the one or more layers comprise an embedding layer of the teacher network.

2. The method of claim 1 , wherein the performing knowledge distillation comprises:

feeding the plurality of task-specific training data into the student network and the teacher network; and

adjusting parameters of the student network to minimize a difference between an output of an embedding layer of the student network and an output of the embedding layer of the teacher network in response to the plurality of task-specific training data.

3. The method of claim 1 , wherein the pre-trained machine learning model is a pre-trained natural language neural network.

4. The method of claim 1 , wherein the performing knowledge distillation comprises:

constructing the plurality of loss functions corresponding to the one or more layers of the teacher network using mean-square errors.

5. The method of claim 1 , wherein the one or more layers of the teacher network of the teacher network comprise:

a self-attention layer;

a feedforward layer; and

an output layer.

6. The method of claim 1 , wherein the plurality of loss functions comprise at least one of:

a loss function based on a mean-square error of a difference between outputs of the embedding layer of the teacher network and a corresponding embedding layer of the student network;

a loss function based on a mean-square error of a difference between attention matrices of a layer of the teacher network and a corresponding layer of the student network;

a loss function based on a mean-square error of a difference between hidden representations generated by the student network and the teacher network; or

a soft cross-entropy loss between the teacher network and logits of the student network.

7. The method of claim 1 , wherein the performing network pruning on the student network comprises:

reducing a number of activated neurons of the student network by setting corresponding parameters as zeros.

8. The method of claim 1 , wherein the teacher network and the student network share a same neural network architecture.

9. The method of claim 1 , wherein the pre-trained machine learning model is a pre-trained residual network for computer vision tasks.

10. The method of claim 9 , wherein the one or more layers of the teacher network comprise:

a weight layer; and

an output layer.

11. The method of claim 1 , wherein the constructing a student network based on the teacher network comprises:

initializing the student network by copying the task-specific machine learning model.

12. A system comprising one or more processors and one or more non-transitory computer-readable memories coupled to the one or more processors and configured with instructions executable by the one or more processors to cause the system to perform operations comprising:

obtaining a pre-trained machine learning model trained based on a plurality of general-purpose training data;

training a task-specific machine learning model by tuning the pre-trained machine learning model based on a plurality of task-specific training data corresponding to a task, wherein the trained task-specific machine learning model is a teacher network;

constructing a student network based on the teacher network; and

executing two or more rounds of fine-tuning of the student network, each round of fine-tuning comprising performing:

network pruning on the student network, and

knowledge distillation from the teacher network to the pruned student network based on a plurality of loss functions corresponding to one or more layers of the teacher network, wherein the one or more layers comprise an embedding layer of the teacher network.

13. The system of claim 12 , wherein the performing knowledge distillation comprises:

feeding the plurality of task-specific training data into the student network and the teacher network; and

adjusting parameters of the student network to minimize a difference between an output of an embedding layer of the student network and an output of the embedding layer of the teacher network in response to the plurality of task-specific training data.

14. The system of claim 12 , wherein the pre-trained machine learning model is a pre-trained natural language neural network, and the performing knowledge distillation comprises:

constructing the plurality of loss functions corresponding to the one or more layers of the teacher network using mean-square errors.

15. The system of claim 12 , wherein the performing network pruning on the student network comprises:

reducing a number of activated neurons in the student network by setting corresponding parameters as zeros.

16. The system of claim 12 , wherein the pre-trained machine learning model is a pre-trained residual network for computer vision tasks, and the one or more layers of the teacher network comprise:

a weight layer; and

an output layer.

17. A non-transitory computer-readable storage medium configured with instructions executable by one or more processors to cause the one or more processors to perform operations comprising:

obtaining a pre-trained machine learning model trained based on a plurality of general-purpose training data;

training a task-specific machine learning model by tuning the pre-trained machine learning model based on a plurality of task-specific training data corresponding to a task, wherein the trained task-specific machine learning model is a teacher network;

constructing a student network based on the teacher network; and

executing two or more rounds of fine-tuning of the student network, each round of fine-tuning comprising performing:

network pruning on the student network, and

knowledge distillation from the teacher network to the pruned student network based on a plurality of loss functions corresponding to one or more layers of the teacher network, wherein the one or more layers comprise an embedding layer of the teacher network.

18. The storage medium of claim 17 , wherein the performing knowledge distillation comprises:

feeding the plurality of task-specific training data into the student network and the teacher network; and

adjusting parameters of the student network to minimize a difference between an output of an embedding layer of the student network and an output of the embedding layer of the teacher network in response to the plurality of task-specific training data.

19. The storage medium of claim 17 , wherein the pre-trained machine learning model is a pre-trained natural language neural network, and the performing knowledge distillation comprises:

constructing the plurality of loss functions corresponding to the one or more layers of the teacher network using mean-square errors.

20. The storage medium of claim 17 , wherein the constructing a student network based on the teacher network comprises:

initializing the student network by copying the task-specific machine learning model.