Efficient softmax computation
AssigneeNVIDIA Corporation
InventorsJacob Robert Stevens, Rangharajan Venkatesan, Steve Haihang Dai, Brucek Khailany
Life of this patent
Estimating the term…
Abstract
Solutions improving efficiency of Softmax computation applied for efficient deep learning inference in transformers and other neural networks. The solutions utilize a reduced-precision implementation of various operations in Softmax, replacing ex with 2x to reduce instruction overhead associated with computing ex, and replacing floating point max computation with integer max computation. Further described is a scalable implementation that decomposes Softmax into UnNormalized Softmax and Normalization operations.
Claims25
A system comprising:
- one or more processors;
- and
- logic that when applied to the one or more processors computes an unnormalized Softmax vector from an input vector by:
- raising elements of the input vector to powers of two;
- computing an integer vector maximum of the input vector;
- and
- logic that when applied to the one or more processors transforms the unnormalized Softmax vector into a normalized Softmax vector.
An artificial neural network comprising:
- one or more feed-forward layers;
- and
- one or more Softmax layers coupled to the one or more feed-forward layers;
- at least one of the Softmax layers configured to compute an unnormalized Softmax vector from an input vector by:
- raising elements of the input vector to powers of two;
- and
- computing an integer vector maximum of the input vector.
A transformer artificial neural network comprising:
- a self-attention layer;
- and
- an encoder-decoder attention layer;
- each of the self-attention layer and the encoder-decoder attention layer comprising a Softmax layer configured to generate an unnormalized Softmax vector from an input vector by:
- raising elements of the input vector to powers of two;
- and
- computing an integer vector maximum of the input vector.
A non-transitory computer-readable storage medium, the computer-readable storage medium including instructions that when executed by a computer, cause the computer to perform neural network inference by:
- formulating a vector of 2x valued elements in a Softmax computation, where each x is an element of an input vector to the Softmax computation;
- and
- computing an integer maximum value of x in the input vector.
A method executed in an artificial neural network layer, the method comprising:
- executing first machine instructions to formulate a vector of 2x valued elements in a Softmax computation, where each x is an element of an input vector to the Softmax computation;
- and
- executing second instructions to compute an integer maximum value of x in the input vector.
Dependencies are inferred from the claim text.