Systems and methods for distilled BERT-based training model for text classification
Embodiments described herein provides a training mechanism that transfers the knowledge from a trained BERT model into a much smaller model to approximate the behavior of BERT. Specifically, the BERT model may be treated as a teacher model, and a much smaller student model may be trained using the same inputs to the teacher model and the output from the teacher model. In this way, the student model can be trained within a much shorter time than the BERT teacher model, but with comparable performance with BERT.
1. A method for distilling knowledge from a first neural network to train a second neural network, the method comprising:
receiving a plurality of training samples corresponding to a first set of pre-defined classes from a given dataset;
retrieving the first neural network that is pre-trained to classify input samples into the first set of pre-defined classes;
obtaining a first plurality of classifications by feeding the plurality of training samples to the first neural network;
transforming, using an out-of-distribution (OOD) sample generation module, one or more of the plurality of the training samples into one or more out-of-distribution (OOD) training samples, wherein the transforming further includes:
computing, using a term frequency-inverse document frequency model, inter-class word importance probabilities of one or more words of a training sample in the plurality of training samples, wherein the inter-class word importance probabilities indicate that the one or more words distinguish between the first plurality of classifications;
computing, using a discriminator model, in-distribution word importance probabilities of the one or more words of a training sample in the plurality of training samples, wherein the in-distribution word importance probabilities indicate contributions of the one or more words to a classification in the first plurality of classifications;
identifying a set of the one or more words within the training sample based on the inter-class word importance probabilities and the in-distribution word importance probabilities; and
replacing the set of the one or more words within the training sample with one or more random words;
generating a second set of classes by adding an out-of-distribution class to the first set of pre-defined classes; and
training the second neural network defined with the second set of classes based on the plurality of training samples, the one or more out-of-distribution training samples and the first plurality of classifications from the first neural network.
2. The method of claim 1 , wherein the first neural network includes any combination of a bidirectional encoder representation from transformers (BERT) model and embeddings from language models (ELMO).
3. The method of claim 1 , wherein the second neural network has a smaller size than the first neural network, and the second neural network is implementable on a central processing unit.
4. The method of claim 1 , further comprising:
training, using a customer dataset, the first neural network to classify input samples into the first set of pre-defined classes, wherein the customer dataset includes the plurality of training samples.
5. The method of claim 1 , wherein the training the second neural network defined with the second set of classes comprises:
generating a second plurality of classification outputs by feeding the plurality of training samples to the second neural network;
computing a knowledge distillation loss between the first plurality of classifications and the second plurality of classification outputs; and
using backpropagation on the second neural network by the knowledge distillation loss to update parameters for the second neural network.
6. The method of claim 5 , further comprising:
generating one or more additional classification outputs by feeding the one or more out-of-distribution training samples to the second neural network;
computing a loss metric between the one or more additional classification outputs and a classification distribution corresponding to the added out-of-distribution class; and
incorporating the loss metric into the knowledge distillation loss.
7. The method of claim 1 , wherein the training the second neural network defined with the second set of classes further comprises:
preprocessing the plurality of training samples or the one or more out-of-distribution training samples by adding a Gaussian noise component before feeding the plurality of training samples or the one or more out-of-distribution training samples to the second neural network.
8. The method of claim 1 , wherein the training the second neural network defined with the second set of classes further comprises:
generating a number of reference class vectors corresponding to the first set of pre-defined classes; and
determining whether an input sample belongs to the added out-of-distribution class based on whether a vector representation of the input sample is orthogonal to the number of reference class vectors.
9. The method of claim 1 , wherein the training the second neural network defined with the second set of classes further comprises:
training the second neural network using the plurality of training samples having a first feature dimension;
in response to receiving an input sample having the first feature dimension, using a Gaussian distribution based sparsification vector to reduce the first feature dimension to a second feature dimension; and
generating, via the second neural network, an output based on the input sample having the second feature dimension.
10. A system for distilling knowledge from a first neural network to train a second neural network, the system comprising:
a communication interface that receives a plurality of training samples;
a memory containing machine readable medium storing machine executable code; and
one or more processors coupled to the memory and configurable to execute the machine executable code to cause the one or more processors to:
receive a plurality of training samples corresponding to a first set of pre-defined classes from a given dataset;
retrieve the first neural network that is pre-trained to classify input samples into the first set of pre-defined classes;
obtain a first plurality of classifications by feeding the plurality of training samples to the first neural network;
transform, using an out-of-distribution (OOD) sample generation module, one or more of the plurality of the training samples into one or more out-of-distribution training samples, wherein the transformation further includes:
computing, using a term frequency-inverse document frequency model, inter-class word importance probabilities of one or more words of a training sample in the plurality of training samples, wherein the inter-class word importance probabilities indicate that one or more words distinguish between the first plurality of classifications;
computing, using a discriminator model, in-distribution word importance probabilities of the one or more words of a training sample in the plurality of training samples, wherein the in-distribution word importance probabilities indicate contributions of the one or more words to a classification in the first plurality of classifications;
identifying a set of the one or more words within the training sample based on the inter-class word importance probabilities and the in-distribution word importance probabilities; and
replacing the set of the one or more words within the training sample with one or more random words;
generate a second set of classes by adding an out-of-distribution class to the first set of pre-defined classes; and
train the second neural network defined with the second set of classes based on the plurality of training samples, the one or more out-of-distribution training samples and the first plurality of classifications from the first neural network.
11. The system of claim 10 , wherein the first neural network includes any combination of a bidirectional encoder representation from transformers (BERT) model and embeddings from language models (ELMO).
12. The system of claim 10 , wherein the second neural network has a smaller size than the first neural network, and the second neural network is implementable on a central processing unit.
13. The system of claim 10 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to:
train, using a customer dataset, the first neural network to classify input samples into the first set of pre-defined classes, wherein the customer dataset includes the plurality of training samples.
14. The system of claim 10 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to train the second neural network defined with the second set of classes by:
generating a second plurality of classification outputs by feeding the plurality of training samples to the second neural network;
computing a knowledge distillation loss between the first plurality of classifications and the second plurality of classification outputs; and
using backpropagation on the second neural network by the knowledge distillation loss to update parameters for the second neural network.
15. The system of claim 14 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to:
generate one or more additional classification outputs by feeding the one or more out-of-distribution training samples to the second neural network;
compute a loss metric between the one or more additional classification outputs and a classification distribution corresponding to the added out-of-distribution class; and
incorporate the loss metric into the knowledge distillation loss.
16. The system of claim 10 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to train the second neural network defined with the second set of classes by:
preprocessing the plurality of training samples or the one or more out-of-distribution training samples by adding a Gaussian noise component before feeding the plurality of training samples or the one or more out-of-distribution training samples to the second neural network.
17. The system of claim 10 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to train the second neural network defined with the second set of classes by:
generating a number of reference class vectors corresponding to the first set of pre-defined classes; and
determining whether an input sample belongs to the added out-of-distribution class based on whether a vector representation of the input sample is orthogonal to the number of reference class vectors.
18. The system of claim 10 , wherein the one or more processors are configurable to execute the machine executable code to cause the one or more processors to train the second neural network defined with the second set of classes by:
training the second neural network using the plurality of training samples having a first feature dimension;
in response to receiving an input sample having the first feature dimension, using a Gaussian distribution based sparsification vector to reduce the first feature dimension to a second feature dimension; and
generating, via the second neural network, an output based on the input sample having the second feature dimension.
19. A non-transitory computer readable medium storing instructions thereon, that when executed by a computing device cause the computing device to perform operations comprising:
receiving a plurality of training samples corresponding to a first set of pre-defined classes from a dataset;
retrieving a first neural network that is pre-trained to classify input samples into the first set of pre-defined classes;
obtaining a first plurality of classifications by feeding the plurality of training samples to the first neural network;
transforming, using an out-of-distribution (OOD) sample generation module, one or more of the plurality of the training samples into one or more out-of-distribution (OOD) training samples, wherein the transforming further includes:
computing, using a term frequency-inverse document frequency model, inter-class word importance probabilities of one or more words of a training sample in the plurality of training samples, wherein the inter-class word importance probabilities indicate that the one or more words distinguish between the first plurality of classifications;
computing, using a discriminator model, in-distribution word importance probabilities of the one or more words of a training sample in the plurality of training samples, wherein the in-distribution word importance probabilities indicate contributions of the one or more words to a classification in the first plurality of classifications;
identifying a set of the one or more words within the training sample based on the inter-class word importance probabilities and the in-distribution word importance probabilities; and
replacing the set of the one or more words within the training sample with one or more random words;
generating a second set of classes by adding an out-of-distribution class to the first set of pre-defined classes; and
training a second neural network defined with the second set of classes based on the plurality of training samples, the one or more out-of-distribution training samples and the first plurality of classifications from the first neural network.
20. The non-transitory computer readable medium of claim 19 , wherein the first neural network includes any combination of a bidirectional encoder representation from transformers (BERT) model and embeddings from language models (ELMO).