Robustness aware norm decay for quantization aware training and generalization
A method includes obtaining a plurality of training samples, determining a minimum integer fixed-bit width representing a maximum quantization of an automatic speech recognition (ASR) model, and training the ASR model on the plurality of training samples using a quantity of random noise. The ASR model includes a plurality of weights that each include a respective float value. The quantity of random noise is based on the minimum integer fixed-bit value. After training the ASR model, the method also includes selecting a target integer fixed-bit width greater than or equal to the minimum integer fixed-bit width, and for each respective weight of the plurality of weights, quantizing the respective weight from the respective float value to a respective integer associated with a value of the selected target integer fixed-bit width. The operations also include providing the quantized trained ASR model to a user device.
1 . A computer-implemented method when executed on data processing hardware causes the data processing hardware to perform operations comprising:
obtaining a plurality of training samples, each respective training sample of the plurality of training samples comprising:
a respective speech utterance; and
a respective textual utterance representing a transcription of the respective speech utterance;
determining a minimum integer fixed-bit width representing a maximum quantization of an automatic speech recognition (ASR) model, the ASR model comprising a plurality of weights, each respective weight of the plurality of weights comprising a respective float value;
training the ASR model on the plurality of training samples using a quantity of random noise, the quantity of random noise based on the minimum integer fixed-bit width;
after training the ASR model, selecting a target integer fixed-bit width greater than or equal to the minimum integer fixed-bit width;
for each respective weight of the plurality of weights, quantizing the respective weight from the respective float value to a respective integer associated with a value of the selected target integer fixed-bit width; and
providing the quantized trained ASR model to a user device.
2 . The method of claim 1 , wherein the maximum quantization level comprises 4-bit quantization.
3 . The method of claim 1 , wherein the maximum quantization level comprises 2-bit quantization.
4 . The method of claim 1 , wherein the random noise is drawn from a uniform distribution of noise.
5 . The method of claim 1 , wherein training the ASR model using the quantity of random noise comprises, for each respective channel of each respective tensor of the ASR model:
determining a respective maximum value for the respective channel of the respective tensor; and
adding, to the respective channel of the respective tensor, a uniform distribution of noise based on the respective maximum value and the minimum integer fixed-bit width.
6 . The method of claim 5 , wherein the uniform distribution of noise represents the entire range of noise the ASR model receives due to quantization up to the minimum integer fixed-bit width.
7 . The method of claim 5 , wherein adding the uniform distribution of noise comprises scaling the uniform distribution of noise based on the respective maximum value.
8 . The method of claim 7 , wherein scaling the uniform distribution of noise is further based on a sensitivity of the respective channel to scaling.
9 . The method of claim 1 , wherein training the ASR model using the quantity of random noise comprises adding, during forward propagation of the ASR model, the quantity of random noise.
10 . The method of claim 1 , wherein:
the ASR model further comprises a plurality of activations, each activation of the plurality of activations comprising a respective float value; and
for each respective activation of the plurality of activations, the operations further comprise quantizing the respective activation from the respective float value to the respective integer with the value of the selected target fixed-bit width.
11 . A system comprising:
data processing hardware; and
memory hardware in communication with the data processing hardware, the memory hardware storing instructions that when executed on the data processing hardware cause the data processing hardware to perform operations comprising:
obtaining a plurality of training samples, each respective training sample of the plurality of training samples comprising:
a respective speech utterance; and
a respective textual utterance representing a transcription of the respective speech utterance;
determining a minimum integer fixed-bit width representing a maximum quantization of an automatic speech recognition (ASR) model, the ASR model comprising a plurality of weights, each respective weight of the plurality of weights comprising a respective float value;
training the ASR model on the plurality of training samples using a quantity of random noise, the quantity of random noise based on the minimum integer fixed-bit width;
after training the ASR model, selecting a target integer fixed-bit width greater than or equal to the minimum integer fixed-bit width;
for each respective weight of the plurality of weights, quantizing the respective weight from the respective float value to a respective integer associated with a value of the selected target integer fixed-bit width; and
providing the quantized trained ASR model to a user device.
12 . The system of claim 11 , wherein the maximum quantization level comprises 4-bit quantization.
13 . The system of claim 11 , wherein the maximum quantization level comprises 2-bit quantization.
14 . The system of claim 11 , wherein the random noise is drawn from a uniform distribution of noise.
15 . The system of claim 11 , wherein training the ASR model using the quantity of random noise comprises, for each respective channel of each respective tensor of the ASR model:
determining a respective maximum value for the respective channel of the respective tensor; and
adding, to the respective channel of the respective tensor, a uniform distribution of noise based on the respective maximum value and the minimum integer fixed-bit width.
16 . The system of claim 15 , wherein the uniform distribution of noise represents the entire range of noise the ASR model receives due to quantization up to the minimum integer fixed-bit width.
17 . The system of claim 15 , wherein adding the uniform distribution of noise comprises scaling the uniform distribution of noise based on the respective maximum value.
18 . The system of claim 17 , wherein scaling the uniform distribution of noise is further based on a sensitivity of the respective channel to scaling.
19 . The system of claim 11 , wherein training the ASR model using the quantity of random noise comprises adding, during forward propagation of the ASR model, the quantity of random noise.
20 . The system of claim 11 , wherein:
the ASR model further comprises a plurality of activations, each activation of the plurality of activations comprising a respective float value; and
for each respective activation of the plurality of activations, the operations further comprise quantizing the respective activation from the respective float value to the respective integer associated with the value of the selected target fixed-bit width.
21 . A computer-implemented method when executed on data processing hardware causes the data processing hardware to perform operations comprising:
obtaining a plurality of training samples;
determining a minimum integer fixed-bit width representing a maximum quantization of a model, the model comprising a plurality of weights, each respective weight of the plurality of weights comprising a respective float value;
training the model on the plurality of training samples using a quantity of random noise, the quantity of random noise based on the minimum integer fixed-bit width;
after training the model, selecting a target integer fixed-bit width greater than or equal to the minimum integer fixed-bit width;
for each respective weight of the plurality of weights, quantizing the respective weight from the respective float value to a respective integer associated with a value of the selected target integer fixed-bit width; and
providing the quantized trained model to a user device.
22 . The method of claim 21 , wherein the model comprises an automated speech recognition model.
23 . The method of claim 21 , wherein the model comprises a large language model.
24 . The method of claim 21 , wherein the model comprises an image processing model.