Quantizing neural networks using approximate quantization function
Some embodiments of the invention provide a novel method for training a quantized machine-trained network. Some embodiments provide a method of scaling a feature map of a pre-trained floating-point neural network in order to match the range of output values provided by quantized activations in a quantized neural network. A quantization function is modified, in some embodiments, to be differentiable to fix the mismatch between the loss function computed in forward propagation and the loss gradient used in backward propagation. Variational information bottleneck, in some embodiments, is incorporated to train the network to be insensitive to multiplicative noise applied to each channel. In some embodiments, channels that finish training with large noise, for example, exceeding 100%, are pruned.
1. A method for transforming a neural network that uses floating point values into a neural network that uses quantized values, the method comprising:
receiving a definition of a floating-point neural network comprising a set of floating-point weight coefficients, wherein intermediate activation values of the floating-point neural network are floating-point values;
training the floating-point neural network to be a quantized neural network comprising a set of ternary weight coefficients, wherein intermediate activation values of the quantized neural network are quantized to a particular range, the training comprising:
propagating a set of inputs through the neural network to generate a set of outputs using an approximate quantization function that is a differentiable approximation to a step-wise quantization function, wherein the approximate quantization function is applied to intermediate activation values to approximate quantization of the intermediate activation values; and
performing a backpropagation operation that (i) backpropagates through the network a loss function value calculated based on the generated set of outputs to determine, for each of at least a subset of the weight coefficients, a rate of change in the calculated loss function value relative to a rate of change in the weight coefficient and (ii) adjusts the weight coefficient based on the determined.
2. The method of claim 1 , wherein the approximate quantization function is defined in terms of a temperature hyperparameter that determines a degree to which the approximate quantization function differs from the step-wise quantization function.
3. The method of claim 2 , wherein the temperature hyperparameter defines noise associated with the approximate quantization function.
4. The method of claim 3 , wherein:
the floating-point neural network comprises a plurality of layers, each layer comprising at least one computation node that uses activation values of computation nodes from previous layers as input values;
the noise associated with the approximate quantization function is used to generate a noise value associated with a particular computation node in the plurality of layers during training of the floating-point neural network to be a quantized neural network.
5. The method of claim 4 , wherein the particular computation node in the neural network is removed from the neural network based on the generated noise value associated with the particular computation node.
6. The method of claim 5 , wherein the particular computation node is removed when the generated noise value is above a threshold value.
7. The method of claim 6 , wherein the noise associated with the approximate quantization function is a multiplicative noise.
8. The method of claim 1 , wherein the quantized neural network is for execution by a neural network inference circuit, the method further comprising generating a set of program instructions for executing the quantized neural network on the neural network inference circuit.
9. The method of claim 8 , wherein the quantized intermediate activation values are one of 8-bit values and 4-bit values.
10. The method of claim 1 , wherein the floating-point intermediate activation values are stored using a variable position of a binary point used to represent the floating-point activation value.
11. The method of claim 10 , wherein the intermediate activation values that are quantized to a particular range use a fixed binary point position.
12. The method of claim 1 , wherein the step-wise quantization function receives values in a range of values and, for each value in the range of received values, outputs a particular value in a set of values that can be represented in a particular number of bits of information used by a neural network inference circuit used to implement the quantized neural network.
13. A non-transitory machine readable medium storing a program for execution by a set of processing units, the program for transforming a neural network that uses floating point values into a neural network that uses quantized values, the program comprising sets of instructions for:
receiving a definition of a floating-point neural network comprising a set of floating-point weight coefficients, wherein intermediate activation values of the floating-point neural network are floating-point values;
training the floating-point neural network to be a quantized neural network comprising a set of ternary weight coefficients, wherein intermediate activation values of the quantized neural network are quantized to a particular range, the training comprising:
propagating a set of inputs through the neural network to generate a set of outputs using an approximate quantization function that is a differentiable approximation to a step-wise quantization function, wherein the approximate quantization function is applied to intermediate activation values to approximate quantization of the intermediate activation values; and
performing a backpropagation operation that (i) backpropagates through the network a loss function value calculated based on the generated set of outputs to determine, for each of at least a subset of the weight coefficients, a rate of change in the calculated loss function value relative to a rate of change in the weight coefficient and (ii) adjusts the weight coefficient based on the determined rate.
14. The non-transitory machine readable medium of claim 13 , wherein the approximate quantization function is defined in terms of a temperature hyperparameter that determines a degree to which the approximate quantization function differs from the step-wise quantization function.
15. The non-transitory machine readable medium of claim 14 , wherein the temperature hyperparameter defines noise associated with the approximate quantization function.
16. The non-transitory machine readable medium of claim 15 , wherein:
the floating-point neural network comprises a plurality of layers, each layer comprising at least one computation node that uses activation values of computation nodes from previous layers as input values;
the noise associated with the approximate quantization function is used to generate a noise value associated with a particular computation node in the plurality of layers during training of the floating-point neural network to be a quantized neural network.
17. The non-transitory machine readable medium of claim 13 , wherein the quantized neural network is for execution by a neural network inference circuit, the program further comprising a set of instructions for generating a set of program instructions for executing the quantized neural network on the neural network inference circuit.
18. The non-transitory machine readable medium of claim 13 , wherein the step-wise quantization function receives values in a range of values and, for each value in the range of received values, outputs a particular value in a set of values that can be represented in a particular number of bits of information used by a neural network inference circuit used to implement the quantized neural network.
19. The non-transitory machine readable medium of claim 13 , wherein:
the floating-point intermediate activation values are stored using a variable position of a binary point used to represent the floating-point activation values; and
the intermediate activation values that are quantized to a particular range use a fixed binary point position.
20. The non-transitory machine readable medium of claim 13 , wherein the quantized intermediate activation values are one of 8-bit values and 4-bit values.