Multi-task learning using knowledge distillation
Methods, systems, and apparatus, including computer programs encoded on computer storage media for performing multi-task learning. In one method a system obtains a respective set of training data for each of multiple machine learning tasks. For each of the machine learning tasks, the system configures a respective teacher machine learning model to perform the machine learning task by training the teacher machine learning model on the training data. The system trains a single student machine learning model to perform the multiple machine learning tasks using (i) the configured teacher machine learning models, and (ii) the obtained training data.
1. A computer implemented method comprising:
obtaining a respective set of training data for each of a plurality of machine learning tasks;
for each of the machine learning tasks, configuring a respective teacher machine learning model to perform the machine learning task by training the teacher machine learning model on the training data for the task; and
training a single student machine learning model having a plurality of student machine learning model parameters to perform all of the plurality of machine learning tasks using (i) the configured teacher machine learning models, and (ii) the obtained training data, wherein training the single student machine learning model comprises:
for each of the plurality of machine learning tasks:
selecting one or more subsets from the set of training data for the machine learning task;
processing the selected subsets using the respective teacher machine learning model to generate respective teacher machine learning model outputs; and
training the single student machine learning model to perform the machine learning task using (i) the selected one or more subsets, and (ii) respective generated teacher machine learning model outputs, comprising, for each subset:
augmenting the subset with an identifier for the machine learning task;
processing the augmented subset using the student machine learning model to generate a student machine learning model output; and
adjusting values of the student machine learning model parameters to match the generated student machine learning model output to the respective generated teacher machine learning model output for the subset.
2. The method of claim 1 , wherein the teacher machine learning model outputs comprise soft target outputs.
3. The method of claim 1 , wherein the training data for each of the plurality of machine learning tasks comprises (i) an input text segment in an input language, and (ii) an output text segment in a target language that is different from the input language.
4. The method of claim 3 , wherein the plurality of machine learning tasks comprise translating an input text segment in an input language into a target language.
5. The method of claim 4 , wherein augmenting the subset with an identifier for the machine learning task comprises prepending each input text segment with a token identifying at least the target language.
6. The method of claim 3 , wherein selecting one or more subsets from the set of training data for the machine learning task comprises selecting one or more sub-word units from the input text segment.
7. The method of claim 6 , wherein each generated respective teacher machine learning model output comprises a probability distribution indicating a respective translation of the corresponding sub-word unit.
8. The method of claim 3 , wherein the training data comprises an equal distribution of text segments in different languages.
9. The method of claim 1 , wherein augmenting the subset with an identifier for the machine learning task comprises prepending the subset with a token identifier for the machine learning task.
10. The method of claim 1 , wherein the student machine learning model is smaller in size than the teacher machine learning models.
11. The method of claim 1 , wherein the student machine learning model is larger in size or the same size as the teacher machine learning models.
12. The method of claim 1 , wherein the size of each of the teacher machine learning models is independent of the student machine learning model.
13. A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:
obtaining a respective set of training data for each of a plurality of machine learning tasks;
for each of the machine learning tasks, configuring a respective teacher machine learning model to perform the machine learning task by training the teacher machine learning model on the training data for the task; and
training a single student machine learning model having a plurality of student machine learning model parameters to perform all of the plurality of machine learning tasks using (i) the configured teacher machine learning models, and (ii) the obtained training data, wherein training the single student machine learning model comprises:
for each of the plurality of machine learning tasks:
selecting one or more subsets from the set of training data for the machine learning task;
processing the selected subsets using the respective teacher machine learning model to generate respective teacher machine learning model outputs; and
training the single student machine learning model to perform the machine learning task using (i) the selected one or more subsets, and (ii) respective generated teacher machine learning model outputs, comprising, for each subset:
augmenting the subset with an identifier for the machine learning task;
processing the augmented subset using the student machine learning model to generate a student machine learning model output; and
adjusting values of the student machine learning model parameters to match the generated student machine learning model output to the respective generated teacher machine learning model output for the subset.
14. The system of claim 13 , wherein the teacher machine learning model outputs comprise soft target outputs.
15. The system of claim 13 , wherein the training data for each of the plurality of machine learning tasks comprises (i) an input text segment in an input language, and (ii) an output text segment in a target language that is different from the input language.
16. One or more non-transitory computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform operations comprising:
obtaining a respective set of training data for each of a plurality of machine learning tasks;
for each of the machine learning tasks, configuring a respective teacher machine learning model to perform the machine learning task by training the teacher machine learning model on the training data for the task; and
training a single student machine learning model having a plurality of student machine learning model parameters to perform all of the plurality of machine learning tasks using (i) the configured teacher machine learning models, and (ii) the obtained training data, wherein training the single student machine learning model comprises:
for each of the plurality of machine learning tasks:
selecting one or more subsets from the set of training data for the machine learning task;
processing the selected subsets using the respective teacher machine learning model to generate respective teacher machine learning model outputs; and
training the single student machine learning model to perform the machine learning task using (i) the selected one or more subsets, and (ii) respective generated teacher machine learning model outputs, comprising, for each subset:
augmenting the subset with an identifier for the machine learning task;
processing the augmented subset using the student machine learning model to generate a student machine learning model output; and
adjusting values of the student machine learning model parameters to match the generated student machine learning model output to the respective generated teacher machine learning model output for the subset.
17. The non-transitory computer-readable media of claim 16 , wherein the teacher machine learning model outputs comprise soft target outputs.
18. The non-transitory computer-readable media of claim 16 , wherein the training data for each of the plurality of machine learning tasks comprises (i) an input text segment in an input language, and (ii) an output text segment in a target language that is different from the input language.
19. The non-transitory computer-readable media of claim 18 , wherein the plurality of machine learning tasks comprise translating an input text segment in an input language into a target language.
20. The non-transitory computer-readable media of claim 19 , wherein augmenting the subset with an identifier for the machine learning task comprises prepending each input text segment with a token identifying at least the target language.