Extreme language model compression with optimal sub-words and shared projections

Patent metadata

Field	Value
Publication number	US-12260340-B2
Application number	US-202318471866-A
Country	US
Kind code	B2
Filing date	Sep 21, 2023
Priority date	Jan 22, 2020
Publication date	Mar 25, 2025
Grant date	Mar 25, 2025

Abstract

Provided is a knowledge distillation technique for training a student language model that, relative to a larger teacher language model, has a significantly smaller vocabulary, lower embedding dimensions, and/or hidden state dimensions. Specifically, aspects of the present disclosure are directed to a dual-training mechanism that trains the teacher and student language models simultaneously to obtain optimal word embeddings for the student vocabulary. In some implementations, this approach can be combined with learning shared projection matrices that transfer layer-wise knowledge from the teacher language model to the student language model. Example experimental results have also demonstrated higher compression efficiency and accuracy when compared with other state-of-the-art compression techniques, including the ability to compress the BERT BASE model by more than 60×, with only a minor drop in downstream task metrics, resulting in a language model with a footprint of under 7 MB.

First claim

What is claimed is:

1. A computing system for training a machine-learned model, the computing system comprising:
one or more processors; and
one or more non-transitory computer-readable media that collectively store:
a first language model comprising one or more transformer layers, wherein the first language model includes a plurality of first language model parameters, wherein each first language model parameter of the plurality of first language model parameters is associated with at least one transformer layer of the one or more transformer layers of the first language model;
a second language model comprising one or more transformer layers, wherein the second language model includes a plurality of second language model parameters, wherein each second language model parameter of the plurality of second language model parameters is associated with at least one transformer layer of the one or more transformer layers of the second language model, wherein the one or more transformer layers of the second language model are of a different dimension than the one or more transformer layers of the first language model, and
instructions that, when executed by the one or more processors, cause the computing system to perform operations, the operations comprising:
projecting the first language model parameters into a shared space with the second language model parameters; and
training the second language model using a loss function based on a comparison of the projected first language model parameters and the second language model parameters.

2. The computing system of claim 1 , wherein the loss function is a mean square error loss function.

3. The computing system of claim 1 , the instructions further comprising:
evaluating a second loss function based on the first language model parameters and the second language model parameters; and
training the second language model based on the second loss function.

4. The computing system of claim 1 , wherein the first language model is a teacher language model and the second language model is a student language model.

5. The computing system of claim 4 , wherein a dimension of the one or more transformer layers of the second language model is smaller than a dimension of the one or more transformer layers of the first language model.

6. The computing system of claim 5 , wherein projecting the first language model parameters into the shared space with the second language model parameters includes down projecting the first language model parameters into the dimension of the one or more transformer layers of the second language model.

7. The computing system of claim 1 , wherein the first language model is a student language model and the second language model is a teacher language model.

8. The computing system of claim 7 , wherein a dimension of the one or more transformer layers of the second language model is larger than a dimension of the one or more transformer layers of the first language model.

9. The computing system of claim 8 , wherein projecting the first language model parameters into the shared space with the second language model parameters includes up projecting the first language model parameters into the dimension of the one or more transformer layers of the second language model.

10. A computer-implemented method for training a machine-learned model, the method comprising:
projecting a plurality of first language model parameters into a shared space with a plurality of second language model parameters, wherein:
a first language model comprises one or more transformer layers, wherein the first language model includes a plurality of first language model parameters, wherein each first language model parameter of the plurality of first language model parameters is associated with at least one transformer layer of the one or more transformer layers of the first language model;
a second language model comprises one or more transformer layers, wherein the second language model includes the plurality of second language model parameters, wherein each second language model parameter of the plurality of second language model parameters is associated with at least one transformer layer of the one or more transformer layers of the second language model, wherein the one or more transformer layers of the second language model are of a different dimension than the one or more transformer layers of the first language model; and
training the second language model using a loss function based on a comparison of the projected first language model parameters and the second language model parameters.

11. The method of claim 10 , wherein the loss function is a mean square error loss function.

12. The method of claim 10 , the method comprising:
evaluating a second loss function based on the first language model parameters and the second language model parameters; and
training the second language model based on the second loss function.

13. The method of claim 10 , wherein the first language model is a teacher language model and the second language model is a student language model.

14. The method of claim 13 , wherein a dimension of the one or more transformer layers of the second language model is smaller than a dimension of the one or more transformer layers of the first language model.

15. The method of claim 14 , wherein projecting the first language model parameters into the shared space with the second language model parameters includes down projecting the first language model parameters into the dimension of the one or more transformer layers of the second language model.

16. The method of claim 10 , wherein the first language model is a student language model and the second language model is a teacher language model.

17. The method of claim 16 , wherein a dimension of the one or more transformer layers of the second language model is larger than a dimension of the one or more transformer layers of the first language model.

18. The method of claim 17 , wherein projecting the first language model parameters into the shared space with the second language model parameters includes up projecting the first language model parameters into the dimension of the one or more transformer layers of the second language model.

19. A computing system for executing a student machine-learned model having parameters aligned with those of a teacher machine-learned model, the computing system comprising:
one or more processors; and
one or more non-transitory computer-readable media that collectively store:
a student language model comprising one or more transformer layers, wherein the student language model includes a plurality of student language model parameters;
wherein the student language model was trained using a loss function based on a comparison of the plurality of student language model parameters and a plurality of teacher language model parameters in a shared space, wherein the plurality of student language model parameters were projected into the shared space, the plurality of teacher language model parameters were projected into the shared space, or the plurality of student language model parameters and the plurality of teacher language model parameters were projected into the shared space; and
instructions that, when executed by the one or more processors, cause the computing system to perform operations, the operations comprising:
generating an output sequence by processing an input sequence using the student machine-learned model.

20. The computing system of claim 19 , w

Assignees

• Google Llc

Inventors

• Song Yang
• Gupta Raghav
• Zhou Dengyong
• Zhao Sanqiang

Classifications

• G06N3/096

  Transfer learning · CPC title

• G06N3/0495

  Quantised networks; Sparse networks; Compressed networks · CPC title

• G06N3/09

  Supervised learning · CPC title

• G06N3/045

  Combinations of networks · CPC title

• G06F40/284 Primary

  Lexical analysis, e.g. tokenisation or collocates · CPC title