Computationally efficient softmax loss gradient backpropagation

1. A computation unit, comprising:

a first circuit to traverse gradient loss elements g pn of a softmax function and normalized output elements p n of the softmax function and produce an accumulation C, wherein the accumulation C is produced by element-wise multiplying the gradient loss elements g pn with the corresponding normalized output elements p n and summing the results of the element-wise multiplication;

a second circuit operatively coupled to the first circuit to element-wise subtract the accumulation C from each of the gradient loss elements g pn and produce modulated gradient loss elements g pn′ ; and

a third circuit operatively coupled to the second circuit to traverse the modulated gradient loss elements g pn′ and produce gradient loss elements g xn for a function preceding the softmax function, wherein the gradient loss elements g xn are produced by element-wise multiplying the modulated gradient loss elements g pn′ with the corresponding normalized output elements p n ,

wherein the first, second, and third circuits comprise a set of one or more computation units,

wherein at least one of the computation units comprises a multi-lane, multi-stage computation pipeline,

wherein the gradient loss elements g pn are converted from a first data format to a second data format using precision upconvert,

wherein, at first and second stages of the multi-lane, multi-stage computation pipeline, the accumulation C is element-wise subtracted in the second data format from corresponding ones of the gradient loss elements g pn to produce corresponding ones of the modulated gradient loss elements g pn′ in the second data format,

wherein, at a third stage of the multi-lane, multi-stage computation pipeline, the normalized output elements p n are converted from the first data format to the second data format using precision upconvert,

wherein, at fourth and fifth stages of the multi-lane, multi-stage computation pipeline, corresponding ones of the normalized output elements p n are element-wise multiplied in the second data format with corresponding ones of the modulated gradient loss elements g pn′ to produce corresponding ones of the gradient loss elements g xn in the second data format, and

wherein the corresponding ones of the gradient loss elements g xn are converted from the second data format to the first data format using precision downconvert.

2. The computation unit of claim 1 , further configured to:

generate a Jacobian matrix by determining a partial derivative of the normalized output elements p n of the softmax function over unnormalized input elements x n to the softmax function;

decompose the Jacobian matrix into diagonal entries and off-diagonal entries;

further decompose the off-diagonal entries into a vector of the normalized output elements p n and a transpose p n T of the vector of the normalized output elements p n ;

element-wise multiply the gradient loss elements g pn with the corresponding normalized output elements p n of the transpose p n T ; and

sum the results of the element-wise multiplication to produce the accumulation C.

3. The computation unit of claim 1 , wherein the gradient loss elements g pn of the softmax function are determined for a current classification module using backpropagation based on a loss function.

4. The computation unit of claim 3 , wherein the current classification module is part of a neural network.

5. The computation unit of claim 4 , wherein the neural network is a deep neural network.

6. The computation unit of claim 1 , wherein the number of multiplication operations required to produce the gradient loss elements g xn of the function preceding the softmax function are reduced from n 2 to n, where n is the dimensionality of the gradient loss elements g pn of the softmax function, the gradient loss elements g xn of the function preceding the softmax function, the normalized output elements p n of the softmax function, and the unnormalized input elements x n to the softmax function.

7. The computation unit of claim 1 , wherein the first circuit receives the gradient loss elements g pn of the softmax function and the normalized output elements p n of the softmax function in a first format and stores the accumulation C in a second format, with the second format having more bits than the first format.

8. The computation unit of claim 7 , wherein the first format is Brain floating-point format 16 (bfloatl6) and the second format is floating point 32 (FP32).

9. The computation unit of claim 8 , wherein the second circuit converts the gradient loss elements g pn into the second format from the first format and produces the modulated gradient loss elements g pn′ in the second format.

10. The computation unit of claim 9 , wherein the third circuit converts the normalized output elements p n into the second format from the first format and produces the gradient loss elements g xn for the function preceding the softmax function in the second format.

11. The computation unit of claim 10 , wherein the third circuit converts the gradient loss elements g xn into the first format from the second format and makes them available in the first format for downstream processing.

12. A re-configurable processor, comprising:

a first circuit to traverse gradient loss elements g pn of a softmax function and normalized output elements p n of the softmax function and produce an accumulation C, wherein the accumulation C is produced by element-wise multiplying the gradient loss elements g pn with the corresponding normalized output elements p n and summing the results of the element-wise multiplication;

a second circuit operatively coupled to the first circuit to element-wise subtract the accumulation C from each of the gradient loss elements g pn and produce modulated gradient loss elements g pn′ ; and

a third circuit operatively coupled to the second circuit to traverse the modulated gradient loss elements g pn′ and produce gradient loss elements g xn for a function preceding the softmax function, wherein the gradient loss elements g xn are produced by element-wise multiplying the modulated gradient loss elements g pn′ with the corresponding normalized output elements p n ,

wherein the first, second, and third circuits comprise a set of one or more computation units,

wherein at least one of the computation units comprises a multi-lane, multi-stage computation pipeline,

wherein the gradient loss elements g pn are converted from a first data format to a second data format using precision upconvert,

wherein, at first and second stages of the multi-lane, multi-stage computation pipeline, the accumulation C is element-wise subtracted in the second data format from corresponding ones of the gradient loss elements g pn to produce corresponding ones of the modulated gradient loss elements g pn′ in the second data format,

wherein, at a third stage of the multi-lane, multi-stage computation pipeline, the normalized output elements p n are converted from the first data format to the second data format using precision upconvert,

wherein, at fourth and fifth stages of the multi-lane, multi-stage computation pipeline, corresponding ones of the normalized output elements p n are element-wise multiplied in the second data format with corresponding ones of the modulated gradient loss elements g pn′ to produce corresponding ones of the gradient loss elements g xn in the second data format, and

wherein the corresponding ones of the gradient loss elements g xn are converted from the second data format to the first data format using precision downconvert.

13. The re-configurable processor of claim 12 , further configured to:

generate a Jacobian matrix by determining a partial derivative of the normalized output elements p n of the softmax function over unnormalized input elements x n to the softmax function;

decompose the Jacobian matrix into diagonal entries and off-diagonal entries;

further decompose the off-diagonal entries into a vector of the normalized output elements p n and a transpose p n T of the vector of the normalized output elements p n ;

element-wise multiply the gradient loss elements g pn with the corresponding normalized output elements p n of the transpose p n T ; and

sum the results of the element-wise multiplication to produce the accumulation C.

14. The re-configurable processor of claim 12 , wherein the gradient loss elements g pn of the softmax function are determined for a current classification module using backpropagation based on a loss function.

15. The re-configurable processor of claim 14 , wherein the current classification module is part of a neural network.

16. The re-configurable processor of claim 15 , wherein the neural network is a deep neural network.

17. The re-configurable processor of claim 12 , wherein the number of multiplication operations required to produce the gradient loss elements g xn of the function preceding the softmax function are reduced from n 2 to n, where n is the dimensionality of the gradient loss elements g pn of the softmax function, the gradient loss elements g xn of the function preceding the softmax function, the normalized output elements p n of the softmax function, and the unnormalized input elements x n to the softmax function.

18. A computer-implemented method, comprising:

traversing, by a first circuit, gradient loss elements g pn of a softmax function and normalized output elements p n of the softmax function and producing an accumulation C, wherein the accumulation C is produced by element-wise multiplying the gradient loss elements g pn with the corresponding normalized output elements p n and summing the results of the element-wise multiplication;

element-wise subtracting, by a second circuit, the accumulation C from each of the gradient loss elements g pn and producing modulated gradient loss elements g pn′ ; and

traversing, by a third circuit, the modulated gradient loss elements g pn′ and producing gradient loss elements g xn for a function preceding the softmax function, wherein the gradient loss elements g xn are produced by element-wise multiplying the modulated gradient loss elements g pn′ with the corresponding normalized output elements p n ,

wherein the first, second, and third circuits comprise a set of one or more computation units,

wherein at least one of the computation units comprises a multi-lane, multi-stage computation pipeline,

wherein the gradient loss elements g pn are converted from a first data format to a second data format using precision upconvert,

wherein, at first and second stages of the multi-lane, multi-stage computation pipeline, the accumulation C is element-wise subtracted in the second data format from corresponding ones of the gradient loss elements g pn to produce corresponding ones of the modulated gradient loss elements g pn′ in the second data format,

wherein, at a third stage of the multi-lane, multi-stage computation pipeline, the normalized output elements p n are converted from the first data format to the second data format using precision upconvert,

wherein, at fourth and fifth stages of the multi-lane, multi-stage computation pipeline, corresponding ones of the normalized output elements p n are element-wise multiplied in the second data format with corresponding ones of the modulated gradient loss elements g pn′ to produce corresponding ones of the gradient loss elements g xn in the second data format, and

wherein the corresponding ones of the gradient loss elements g xn are converted from the second data format to the first data format using precision downconvert.

19. The computer-implemented method of claim 18 , further comprising:

generating a Jacobian matrix by determining a partial derivative of the normalized output elements p n of the softmax function over unnormalized input elements x n to the softmax function;

decomposing the Jacobian matrix into diagonal entries and off-diagonal entries;

further decomposing the off-diagonal entries into a vector of the normalized output elements p n and a transpose p n T of the vector of the normalized output elements p n ;

element-wise multiplying the gradient loss elements g pn with the corresponding normalized output elements p n of the transpose p n T ; and

summing the results of the element-wise multiplication to produce the accumulation C.

20. The computer-implemented method of claim 18 , wherein the gradient loss elements g pn of the softmax function are determined for a current classification module using backpropagation based on a loss function.