A method for a cross-platform distillation framework includes obtaining a plurality of training samples. The method includes generating, using a student neural network model executing on a first processing unit, a first output based on a first training sample. The method also includes generating, using a teacher neural network model executing on a second processing unit, a second output based on the first training sample. The method includes determining, based on the first output and the second output, a first loss. The method further includes adjusting, based on the first loss, one or more parameters of the student neural network model. The method includes repeating the above steps for each training sample of the plurality of training samples.

TECHNICAL FIELD

This disclosure relates to a cross-platform distillation framework.

BACKGROUND

In machine learning, training a model is a time and resource intensive process, especially as models typically perform better when they are trained on large data sets. However, there are drawbacks with using large models (i.e., models trained on large data sets), such as inflexibility and difficulty in deploying large models on smaller devices. Smaller models, which are easier to deploy, can be used for specific contexts (i.e., trained on context-specific data). However, such smaller models may not perform as well as large models, as the training process is less robust. Knowledge distillation is the process of transferring the knowledge from a generally larger trained model (i.e., a “teacher model”) to a generally smaller model (i.e., a “student model”) such that the smaller model can perform without significant loss of performance compared to the large model, while still maintaining the benefits of a smaller model.

SUMMARY

One aspect of the disclosure provides for a computer-implemented method for a cross-platform distillation framework. The computer-implemented method when executed by data processing hardware causes the data processing hardware to perform operations including obtaining a plurality of training samples. The operations include generating, using a student neural network model executing on a first processing unit, a first output based on a first training sample of the plurality of training samples. The operations further include generating, using a teacher neural network model executing on a second processing unit, a second output based on the first training sample of the plurality of training samples, the second processing unit remote from the first processing unit. The operations include determining, based on the first output and the second output, a first loss. The operations also include adjusting, based on the first loss, one or more parameters of the student neural network model. The operations include after adjusting the one or more parameters of the student neural network model, generating, using the student neural network model, a third output based on a second training sample of the plurality of training samples. The operations include generating, using the teacher neural network model, a fourth output based on the second training sample of the plurality of training samples. The operations further include determining, based on the third output and the fourth output, a second loss. The operations include readjusting, based on the second loss, the one or more parameters of the student neural network model.

Implementations of the disclosure may include one or more of the following optional features. In some implementations, the first processing unit and the second processing unit each include a respective tensor processing unit. The operations may further include transmitting a remote procedure call (RPC) to the teacher neural network model to generate each output. Further, the first output, the second output, the third output, and the fourth output may each include a respective logit.

In some implementations, the operations further include determining, based on the loss, a gradient and adjusting one or more parameters of the student neural network model by applying the gradient to the student neural network model. The teacher neural network model may include a trained model. In some implementations, the first processing unit belongs to a first entity and the second processing unit belongs to a second entity different from the first entity.

In some implementations, the first training sample includes an unlabeled training sample. In these implementations, the operations may further include generating, based on the second output from the teacher neural network model, a label for the first training sample. In these implementations, the operations may further include generating, using a second student neural network model executing on a third processing unit, a fifth output based on the labeled first training sample, determining, based on the label and the fifth output, a third loss, and adjusting, based on the third loss, the second student neural network model.

Another aspect of the disclosure provides a system for a cross-platform distillation framework. The system includes data processing hardware and memory hardware in communication with the data processing hardware. The memory hardware stores instructions that when executed on the data processing hardware cause the data processing hardware to perform operations. The operations include generating, using a student neural network model executing on a first processing unit, a first output based on a first training sample of the plurality of training samples. The operations further include generating, using a teacher neural network model executing on a second processing unit, a second output based on the first training sample of the plurality of training samples, the second processing unit remote from the first processing unit. The operations include determining, based on the first output and the second output, a first loss. The operations also include adjusting, based on the first loss, one or more parameters of the student neural network model. The operations include after adjusting the one or more parameters of the student neural network model, generating, using the student neural network model, a third output based on a second training sample of the plurality of training samples. The operations include generating, using the teacher neural network model, a fourth output based on the second training sample of the plurality of training samples. The operations further include determining, based on the third output and the fourth output, a second loss. The operations include readjusting, based on the second loss, the one or more parameters of the student neural network model.

This aspect may include one or more of the following optional features. In some implementations, the first processing unit and the second processing unit each include a respective tensor processing unit. The operations may further include transmitting a remote procedure call (RPC) to the teacher neural network model to generate each output. Further, the first output, the second output, the third output, and the fourth output may each include a respective logit.

In some implementations, the operations further include determining, based on the loss, a gradient and adjusting one or more parameters of the student neural network model by applying the gradient to the student neural network model. The teacher neural network model may include a trained model. In some implementations, the first processing unit belongs to a first entity and the second processing unit belongs to a second entity different from the first entity.

In some implementations, the first training sample includes an unlabeled training sample. In these implementations, the operations may further include generating, based on the second output from the teacher neural network model, a label for the first training sample. In these implementations, the operations may further include generating, using a second student neural network model executing on a third processing unit, a fifth output based on the labeled first training sample, determining, based on the label and the fifth output, a third loss, and adjusting, based on the third loss, the second student neural network model.

DETAILED DESCRIPTION

Distillation is the process of transferring knowledge from a teacher model (e.g., a large trained model) to a student model (e.g., a smaller untrained model) during training of the student model. Distillation provides for the student model to be trained more quickly (and on less training data) than the teacher model while maintaining performance that is substantially the same as the teacher model. One distillation technique, known as online distillation, includes feeding a training sample to both the teacher model and the student model, obtaining a respective output based on the training sample from each of the teacher model and the student model, determining a loss from a combination of the teacher output and the student output, and adjusting the student model based on the loss. Current techniques for online distillation require the teacher model and the student model to be hosted on the same platform/framework (e.g., the same processor).

Another distillation technique, known as offline distillation, includes labeling a set of training data using the teacher model, and then training the student model using the teacher-labeled training data. In offline distillation, the teacher model and the student model can be stored in different frameworks. However, offline distillation requires the storage, and transfer, of the teacher-labeled set of training data, which can be prohibitively large, making offline distillation not scalable.

Implementations herein are directed to a cross-platform distillation framework. In particular, the distillation techniques of the current disclosure allow for a teacher model stored on or using a first framework to transfer knowledge to a student model stored on or using a second framework during training of the student model. A training module may transmit training samples to each of the teacher model and the student model. In some implementations, the training module transmits the training samples to the teacher model through a remote procedure call (RPC). The student model may then generate a first output based on the training sample and the teacher model may generate a second output based on the training sample. The training module may then generate a loss based on the first output and the second output. In some implementations, the training module adjusts one or more parameters of the student model based on the loss. The training module may repeat these steps for each training sample in the plurality of training samples.

Referring toFIG.1, in some implementations, an example cross-platform distillation system100includes a cloud environment140in communication with one or more user devices110via a network112. The cloud environment140may be a single computer, multiple computers, or a distributed system having scalable/elastic resources142including computing resources144(e.g., data processing hardware) and/or storage resources146(e.g., memory hardware). The cloud environment140may be configured to store a large amount of data for use in big data analytics. A data store150(i.e., a remote storage device) may be overlain on the storage resources146to allow scalable use of the storage resources146by one or more of the clients (e.g., the user device10) or the computing resources144. The data store150is configured to store a plurality of training samples152,152A-N associated with a set of training data152.

The cloud environment140is configured to obtain the training samples152of the plurality of training samples152for a training module201from, for example, a user device110via the network112. For example, if the training module201is configured to train one or more neural network models for speech recognition, the training samples152may be a recording of an utterance (i.e., snippet of speech) of the user using the client device110. In some implementations, the training samples152may be labeled. In these implementations, the training samples152may include an utterance with a ground truth transcription of the utterance. The user device110may correspond to any computing device, such as a desktop workstation, a laptop workstation, or a mobile device (i.e., a smart phone). The user device110includes computing resources (e.g., data processing hardware) and/or storage resources (e.g., memory hardware).

The cloud environment140executes the training module201for performing cross-platform distillation using a teacher neural network model210(also referred to herein as just the teacher210) executing on a first platform205and a student neural network model260(also referred to herein as just the student260) executing on a second platform255. The student260includes one or more parameters261(e.g., weights or the like). Here, the teacher210may be a trained model that the training module201is configured to implement in a cross-platform distillation training process to transfer knowledge from the teacher210to one or more students260during training of the student(s)260. As used herein, the term “platform” can refer to any suitable computing environment for executing a neural network model and/or a machine learning model, such as a processing unit of a computer and/or a framework (such as tensorflow, jax, pytorch, MXNet, etc.). In some implementations the first platform205is a first processor (e.g., a first tensor processing unit (TPU)) and the second platform255is a second processor (e.g., a second TPU) different from the first processor. The first processor and second processor may belong a group of computers in a computing environment (e.g., cloud environment140). In additional implementations, the first platform205is a first framework and the second platform255is a second framework. In these implementations, the teacher210is a model trained and inferred on the first framework and the student is trained in the second framework. Alternatively, the first processor and second processor may belong to distinct entities. For example, the first processor may belong to a server environment (i.e., cloud environment140) capable of storing and hosting a large trained neural network model (i.e., teacher210) while the second processor belongs to a device (i.e., client device110) with less processing power and less memory that is configured to deploy and train a smaller neural network model (i.e., student260). In this example, the second processor can train the smaller neural network model through distillation without having access to the large trained neural network model. Or put another way, the owner of the teacher210may allow training of one or more students260without having to allow access to the teacher210(e.g., for security or proprietary purposes).

The training module201is configured to obtain training samples152(e.g., from the user device110, the data store150, remote servers, or other modules of the cloud environment140). The training module201, in some implementations, transmits one or more training samples152of the plurality of training samples152to each of the teacher210and the student260. The training module201may transmit the training sample152to the teacher210via a remote procedure call (RPC)202. The training module may then obtain respective outputs215,265from the teacher210and student260. In some implementations the training module201may implement a loss function280to determine a loss285based on the outputs215,265(i.e., a representative of a difference between the outputs215,265). The distillation process of the training module201is discussed in greater detail below (FIG.2andFIG.4).

The system ofFIG.1is presented for illustrative purposes only and is not intended to be limiting. For example, although only a single example of each component is illustrated, the system100may include any number of components110,112,140,150,201,205,210,255,260, and280. Further, although some components are described as being located in a cloud computing environment140, in some implementations, some or all of the components may be hosted locally on the client device110. Further, in various implementations, some or all of the components, are hosted locally on user device110, remotely (such as in the cloud computing environment140), or any combination thereof.

FIG.2illustrates a schematic view200of an example cross-platform distillation training process. Here, the training module201is configured to perform the cross-platform distillation. The training process may begin when the training module201obtains a first training sample152from a data store150. The training sample152may belong to a set and/or a plurality of training samples152stored at the data store150. The training sample152may be based on the type of neural network models that are being trained by the training module201in the cross-platform distillation training process. For example, when the neural network models are adapted for speech recognition, the training samples152include speech samples, such as utterances recorded at a client device. In another example, when the neural network models are adapted for natural language processing, the training samples152are transcripts of text. In yet other examples, the training samples152may include images.

Once the training module201obtains the first training sample152, the training module may then transmit the training sample152to each of a student neural network model260and a teacher neural network model210. In some implementations, the training module201transmits the training sample152as a remote procedure call (RPC)202, causing the teacher210, via teacher platform205, to generate a teacher output215based on the training sample152. In some implementations, the teacher platform205of the teacher210is different than the student platform255of the student260. For example, the teacher platform205of the teacher210is a processor unit, such as a tensor processing unit running on a first computer, and the student platform255of the student260is a different processing unit, such as a different tensor processing unit running on a second computer. In some implementations, the teacher210and student260execute on respective platforms205,255that are in a shared computer network, such as two computers in a cloud computing environment (e.g., cloud environment140). In other implementations, the teacher210and student260execute on respective platforms205,255that are in a remote computer networks. For example, the teacher210executes in a teacher platform205belonging to a first entity (e.g., a server of an organization or business), while the student260executes in a student platform255belonging to a second entity (e.g., a client device of a customer of the organization) adapted to communicate with the first entity (e.g., via one or more networks). The platforms205,255can be any suitable platform and/or processing unit for executing a neural network model such as a tensor processing unit, a graphics processing unit, a computer processing unit, etc.

The teacher210and the student260may each generate a respective output215,265based on the first training sample152. The outputs215,265may be in any suitable form of an output of a neural network model and/or machine learning model such as a logit, a probability density function, a sigmoid function, etc. The training module201may then obtain the respective outputs215,265from the teacher210and the student260. The training module may then implement a loss function280to determine a loss285based on the outputs215,265. In some implementations, the teacher210and the student260may run in parallel (e.g., synchronously) to produce outputs215,265at the same time or nearly the same time. Accordingly, the loss function280obtains the outputs215,265at the same time or nearly the same time to generate the loss285. During such synchronous cross-distillation, student260is generally adjusted based on the loss280for a set of outputs215,265prior to the next set of outputs215,265being generated (i.e., the teacher210generates outputs215synchronously with the outputs265of the student260). In other implementations, the teacher210and the student260execute asynchronously (i.e., the teacher210may generate outputs215independent of the student260generating outputs265). In these implementations, the loss function280may not generate a loss285until the respective pair of outputs215,265generated based on the first training sample152is received. The loss function280may execute on the student platform255. Alternatively, the loss function280executes in a separate environment remote from either platform205,255.

The training module201may adjust the one or more parameters261(e.g., weights) of the student260based on the loss285. In some implementations, the training module201generates a gradient based on the loss285. The training module201may adjust the one or more parameters261of the student260by applying the gradient to the student260. Once the training module201adjusts the one or more parameters261of the student based on the loss285, the training module201may restart the cross-platform distillation training process with a second training sample152. The training module201can continue the cross-platform distillation training process for each training sample152of the set of training samples152. In some implementations, the training module201does not transmit the next training sample152until the process has been completed for the current training sample152. In other words, the training module201executes the cross-platform distillation training process sequentially or synchronously for each training sample152of the set of training samples152such that training module201completes the training process for each training sample152(i.e., generates outputs215,265, calculates a loss285, and adjusts one or more parameters of the student260based on the loss285) before continuing with the next training sample152.

FIG.3is a schematic view300of another example cross-platform distillation training process. Here, the training module201may be configured to obtain a labeled training sample152,152L from a teacher neural network model210executing on a teacher platform205based on an unlabeled training sample152,152U. In some implementations, the training module201transmits the unlabeled training sample152U as a remote procedure call (RPC), causing the teacher210, via teacher platform205, to generate a teacher output215based on the unlabeled training sample152U. The teacher output215may be any suitable output from a neural network model and/or a machine learning model, such as a logit. The training module201may implement a label module301to generate the labeled training sample152L based on the teacher output215of the teacher210and the unlabeled training sample152U. In some implementations, the labeled training sample152L merely includes the unlabeled training sample152U with the teacher output215as the label. In other implementations, the label module301modifies or adjusts the teacher output215to an appropriate label for machine learning/neural network training using a labeled training sample152L.

The training module201may then transmit the labeled training sample152L to a student neural network model260executing on a student platform255. The student260may generate a student output265based on the labeled training sample152L. The student output265may be in any suitable form as an output from a neural network model and/or machine learning model, such as a logit. The training module201may implement a loss function280to determine a loss285based on the student output265and the labeled training sample152L. For example, the loss function280may compare the label of the labeled training sample152L to the student output265to determine the loss285. The training module201may then adjust one or more parameters of the student260based on the loss285.

In some implementations, the training module201executes sequentially for each unlabeled training sample152U,152from a plurality of unlabeled training samples152U. In other words, the training module201may transmit a first unlabeled training sample152U to the teacher210to obtain a first labeled training sample152L (based on a first teacher output215based on the first unlabeled training sample152U) from the label module301. The training module201may then transmit the first labeled training sample to the student260such that the student260generates a first student output265. The training module201may then determine a first loss285, via loss function280, based on the student output265and the first labeled training sample152L. The training module201may then adjust one or more parameters of the student260based on the loss285. In some implementations, the training module201generates a gradient based on the loss285. The training module201may then adjust the one or more parameters of the student260by applying the gradient to the student260.

After the training module201adjusts one or more parameters of the student260based on the first loss285, the training module201may then restart the training process with a second unlabeled training sample152U from the plurality of training samples152U. In this manner, the labeled training samples152L are not stored in memory for an extended period of time. Instead, the labeled training samples152L are used to train the student260once they are generated.

In some implementations, the training module201may be adapted to store a number of labeled training samples152L until the student260is ready to process the labeled training samples152L. In these implementations, the training module201may store the labeled training samples152L at a memory that is convenient for the training process. For example, the training module201can store the labeled training samples152L in a memory of the student platform255. Alternatively, the training module201can store the labeled training samples152L in a memory of a cloud computing network that has a large amount of free capacity.

As described above with relation toFIG.2, the teacher platform205of the teacher210may be different from the student platform255of the student260. For example, the teacher platform205of the teacher210is a processor unit, such as a tensor processing unit running on a first computer, and the student platform255of the student260is a different processing unit, such as a different tensor processing unit running on a second computer. In some implementations, the teacher210and student260execute on respective platforms205,255that are in a shared computer network, such as two separate computers in a cloud computing environment (e.g., cloud environment140). In other implementations, the teacher210and student260execute on respective platforms205,255that are in a remote computer networks. For example, the teacher210executes in a teacher platform205belonging to a first entity (e.g., a server), while the student260executes in a student platform255belonging to a second entity (e.g., a client device) adapted to communicate with the first entity. The platforms205,255can be any suitable platform and/or processing unit for executing a neural network model such as a tensor processing unit, a graphics processing unit, a computer processing unit, etc.

FIG.4is a schematic view of a sequence diagram400for a cross-platform distillation framework. The sequence begins at step402when a training module201obtains a training sample (e.g., training sample152) from a data store150. At step404and406, the training module201transmits the training sample152to a student neural network model (“student”)260and a teacher neural network model (“teacher”)210, respectively. The teacher210and student260may reside in separate platforms and/or processing units. In some implementations, the teacher210is a trained neural network model. The training module201may be adapted to transfer knowledge from the teacher210to the student260through the process of cross-platform distillation training. The teacher210and student260may be any suitable neural network model and/or machine learning model. The student260may be a smaller model that may be easier to deploy on a device such as a smart phone or personal computer with limited computational resources while the teacher210may be a large trained model that executes on a server. In some implementations, the teacher210and student260are the same type of neural network/model. For example, the teacher210may be trained neural network adapted for speech recognition and the student260is an untrained neural network configured to be deployed on a client device and adapted for local speech recognition on the client device. Further, the teacher210and student260may be any neural networks and/or models that are suitable for learning through distillation, such as computer vision, natural language processing, speech recognition, etc. Steps404and406may occur simultaneously or nearly simultaneously.

At step408, the student260may transmit a student output (e.g., student output265) to a loss function280. At step410, the teacher210may transmit a teacher output (e.g., teacher output215) to the loss function280. The student output265and the teacher output215may be of any suitable type of output from a neural network model and/or machine learning model, such as a logit. In some implementations, the teacher210and the student260operate in parallel (i.e., synchronously) such that the loss function280obtains the teacher output215and the student output265at the same time or near the same time. In alternative implementations, the student260and the teacher210may operate asynchronously. In these alternative implementations, the loss function may be adapted to store either output215,265until the corresponding output265,215is received before determining the respective loss (step412) based on the outputs215,265associated with the training sample152. In some implementations, the loss function280is a component of the training module201. Further, the training module201and the loss function280may execute in the same platform as the student260. Alternatively, the training module201and the loss function280may execute in a different platform than the student260and/or the teacher210. At step412the loss function280may determine a loss (e.g., loss285) based on the outputs215,265.

At step414, the training module201may obtain the loss285. At step416, the training module201may adjust one or more parameters261of the student260. In some implementations, the training module201generates a gradient based on the loss285. The training module201may then adjust the one or more parameters of the student260by applying the gradient to the student260. The cross-platform distillation training process may iterate through each training sample152of a set of training samples (e.g., plurality of training samples152) to complete training of the student260.

FIG.5is a flowchart of an exemplary arrangement of operations for a method500of a cross-platform distillation framework. The method500may be performed, for example, by various elements of the cross-platform distillation system100ofFIG.1. At operation502, the method500includes obtaining a plurality of training samples152. At operation504, the method500includes generating, using a student neural network model260executing on a first processing unit255, a first output265based on a first training sample152of the plurality of training samples152. At operation506, the method500includes generating, using a teacher neural network model210executing on a second processing unit205, a second output215based on the first training sample152of the plurality of training samples152, the second processing unit205remote from the first processing unit255. At operation508, the method500includes determining, based on the first output265and the second output215, a first loss285. At operation510, the method500includes adjusting, based on the first loss285, one or more parameters161of the student neural network model260. At operation512, the method500includes, after adjusting the one or more parameters161of the student neural network model260, generating, using the student neural network model260, a third output265based on a second training sample152of the plurality of training samples152. At operation514, the method500includes generating, using the teacher neural network model210, a fourth output215based on the second training sample152of the plurality of training samples152. At operation516, the method500includes determining, based on the third output265and the fourth output215, a second loss285. At operation518, the method500includes readjusting, based on the second loss285, the one or more parameters of the student neural network model260.

The computing device600includes a processor610, memory620, a storage device630, a high-speed interface/controller640connecting to the memory620and high-speed expansion ports650, and a low speed interface/controller660connecting to a low speed bus670and a storage device630. Each of the components610,620,630,640,650, and660, are interconnected using various busses, and may be mounted on a common motherboard or in other manners as appropriate. The processor610can process instructions for execution within the computing device600, including instructions stored in the memory620or on the storage device630to display graphical information for a graphical user interface (GUI) on an external input/output device, such as display680coupled to high speed interface640. In other implementations, multiple processors and/or multiple buses may be used, as appropriate, along with multiple memories and types of memory. Also, multiple computing devices600may be connected, with each device providing portions of the necessary operations (e.g., as a server bank, a group of blade servers, or a multi-processor system).

The storage device630is capable of providing mass storage for the computing device600. In some implementations, the storage device630is a computer-readable medium. In various different implementations, the storage device630may be a floppy disk device, a hard disk device, an optical disk device, or a tape device, a flash memory or other similar solid state memory device, or an array of devices, including devices in a storage area network or other configurations. In additional implementations, a computer program product is tangibly embodied in an information carrier. The computer program product contains instructions that, when executed, perform one or more methods, such as those described above. The information carrier is a computer-or machine-readable medium, such as the memory620, the storage device630, or memory on processor610.

The high speed controller640manages bandwidth-intensive operations for the computing device600, while the low speed controller660manages lower bandwidth-intensive operations. Such allocation of duties is exemplary only. In some implementations, the high-speed controller640is coupled to the memory620, the display680(e.g., through a graphics processor or accelerator), and to the high-speed expansion ports650, which may accept various expansion cards (not shown). In some implementations, the low-speed controller660is coupled to the storage device630and a low-speed expansion port690. The low-speed expansion port690, which may include various communication ports (e.g., USB, Bluetooth, Ethernet, wireless Ethernet), may be coupled to one or more input/output devices, such as a keyboard, a pointing device, a scanner, or a networking device such as a switch or router, e.g., through a network adapter.

The computing device600may be implemented in a number of different forms, as shown in the figure. For example, it may be implemented as a standard server600aor multiple times in a group of such servers600a, as a laptop computer600b, or as part of a rack server system600c.