Patent Description:
Neural networks are machine learning models that employ one or more layers of nonlinear computation units to predict an output for a received input.

Recurrent neural networks generate, for a current time step, an output that is informed by other outputs generated for one or more previous time steps. Some recurrent neural networks employ gated activation units. Such recurrent neural networks may be referred to as gated recurrent neural networks.

Gated activation units maintain a type of memory by implementing functions that control how much information generated in a previous time step should be remembered and how much should be forgotten. Common gated activation units include Long Short Term Memory units (LSTM units), Gated Recurrent Units (GRUs), several variants thereof.

In general, a gated activation unit updates a current hidden state using at least the previous hidden state and a current input. Updating the hidden state generally involves one or more linear transformations and one or more nonlinear activations. Each linear transformation can use a weight matrix and a bias vector. Training a gated recurrent neural network thus involves learning the weight matrices and bias vectors for each gated activation unit.

<FIG> illustrates a prior art LSTM unit <NUM>. The LSTM unit <NUM> maintains both a hidden state, ht, as well as a cell state, ct. The LSTM unit <NUM> implements a forget gate, which determines how much of the previous hidden state to forget; and an input gate, which determines which values of the cell state to update, and an output gate, which determines which values to output.

The operations of the LSTM unit <NUM> can be defined by the following equations, in which the previous hidden state ht-<NUM> corresponds to ht-<NUM> <NUM>, a previous cell state ct-<NUM> corresponds to ct-<NUM> <NUM>, and the current input xt corresponds to xt <NUM>. In this context, "*" refers to pointwise multiplication, "+" refers to pointwise addition, and "σ" is a sigmoid activation function. The notation "Wx[ht-<NUM>, xt]" refers to matrix multiplication of a matrix Wx by a vector of ht-<NUM> concatenated with xt. Some literature splits each matrix Wx into two matrices, W and U, in which case W is multiplied with ht-<NUM> and U is multiplied by xt. The operations of LSTM unit <NUM> are defined by the following equations: <MAT> <MAT> <MAT> <MAT> <MAT> <MAT>.

A forget gate <NUM> corresponds to Equation (<NUM>) and generates ft <NUM>; an input gate <NUM> corresponds to Equation (<NUM>) and generates it <NUM>; and an output gate <NUM> corresponds to Equation (<NUM>) and generates ot <NUM>. A tanh layer <NUM> corresponds to Equation (<NUM>) and generates a vector of candidates c_bart <NUM> for adding to the current cell state. A current cell state ct <NUM> is given by Equation (<NUM>). A last tanh layer <NUM> uses the computed cell state ct <NUM> to generate the current hidden state ht <NUM> according to Equation (<NUM>).

Each of Equations (<NUM>), (<NUM>), (<NUM>), and (<NUM>) specifies performing a matrix operation between a respective weight matrix Wn for the corresponding layer and the current input vectors xt and ht-<NUM>. The result is then added to a respective bias vector bn to the result. The result of these calculations is then fed through nonlinear activation functions σ and tanh to generate a final output vector ht for time step t.

<FIG> illustrates a prior art gated recurrent unit (GRU) <NUM>. A main difference between the GRU <NUM> and the LSTM unit <NUM> is that the GRU <NUM> effectively merges the LSTM cell state and hidden state into just a hidden state. Therefore, the GRU <NUM> receives as input a previous hidden state and outputs only a current hidden state for a given input xt.

The operations of the GRU <NUM> are generally defined by the following equations, in which the previous hidden state ht-<NUM> corresponds to ht-<NUM> <NUM> and the current input xt corresponds to xt <NUM>. <MAT> <MAT> <MAT> <MAT>.

A reset gate <NUM> corresponds to Equation (<NUM>) and generates rt <NUM>. An update gate <NUM> corresponds to Equation (<NUM>) and generates zt <NUM>. A tanh layer <NUM> corresponds to Equation (<NUM>) and generates h_bart <NUM>. The final hidden state ht <NUM> is then computed according to Equation (<NUM>).

As can be seen, the GRU <NUM> is somewhat simpler than the LSTM unit <NUM> in that it implements fewer gates and activation functions. But like the LSTM <NUM>, the GRU <NUM> also uses a number of matrix operations using a weight matrix Wn and a current input xt, and a respective bias vector bn is then added to the result. The result of these calculations is also then fed through respective nonlinear activation functions σ or tanh.

Some aspects of the architecture of gated activation units lend themselves to parallel processing techniques. In particular, none of the operations of Equations (<NUM>)-(<NUM>) for the LSTM unit depend on each other, and none of the operations of Equations (<NUM>)-(<NUM>) for the GRU depend on each other. Furthermore, because these operations involve matrix multiplications and point-wise addition, multiple processing devices, e.g., the streaming multiprocessors (SMs) of a graphics processing unit (GPU) can be used to compute partial results of the matrix operations, and then partial results can be combined.

<NPL> describes techniques for accelerating speech recognition on GPUs. The general approach involves arranging the calculations and implementing a number of optimizations, including converting vector-matrix multiplications to matrix-matrix multiplications and fusing elementwise operations.

This specification describes techniques for performing the operations of gated activation units on parallel processing hardware.

Thus in one aspect a method of implementing a neural network on a parallel processing device, comprises receiving a plurality of weight matrices of a gated activation unit of the neural network. The gated activation unit has two or more layers, each layer defining operations comprising: (i) a matrix operation between a weight matrix for the layer and concatenated input vectors and (ii) a nonlinear activation operation using a result of the matrix operation. The matrix operation typically comprises a matrix multiplication but in principle other matrix operations may also be implemented. The method interleaves rows of the plurality of weight matrices by assigning groups of corresponding rows to respective thread blocks. Each thread block comprises a computation unit for execution by an independent processing unit of a plurality of independent processing units of a parallel processing device. As described later, such an approach can provide advantages including removing the necessity of performing a thread block synchronization step, which might otherwise require writing a set of results to a storage area from where they can be accessed by a next stage of processing.

There is also described a method of processing data using a neural network implemented in this way. The processing comprises: receiving, by each thread block, input vectors; and generating, by each thread block, a respective portion of a current state vector. The generating includes performing a plurality of partial matrix operations using one or more groups of corresponding rows of the plurality of weight matrices assigned to the thread block and a concatenation of the input vectors. The generating further includes performing a plurality of nonlinear activation operations using respective results of the plurality of partial matrix operations. Optionally multiple portions of the current state vector may be generated at least partially in parallel.

In some implementations interleaving the rows may comprise assigning at least one row from every weight matrix of the plurality of weight matrices to each thread block, and/or assigning a same number of rows to all thread blocks. This can facilitate efficient use of the available processing hardware.

The gated activation unit may comprise, for example, a gated recurrent unit or a long short term memory unit of the neural network, or another recurrent neural network unit.

In some implementations the operations of a layer may include a bias addition operation between a result of a matrix operation for the layer and a bias vector for the layer; this can improve performance of the neural network. More particularly this may comprise distributing portions of the bias vector to each thread block of the plurality of thread blocks and adding a portion of the bias vector to a result of a matrix operation.

The method further comprises interleaving rows of the plurality of weight matrices by warps such that all warps receive some rows from every weight matrix and corresponding rows from every weight matrix are assigned to the same warp. Within a row assigned to a warp each thread is assigned a number of row values given by M / (number of threads per warp) where M and the number of threads per warp are both integers. Some implementations may additionally or alternatively comprise interleaving values of the plurality of weight matrices by threads such that all threads receive some values from every weight matrix and corresponding values from every weight matrix are assigned to the same thread.

In a related aspect a system is configured to implement a neural network on a parallel processing device as described above.

A parallel processing device can generate all output values of a gated activation unit in parallel or substantially in parallel and without requiring any block synchronizations. A model can be trained to exactly fit the available processing hardware, which can result in optimal load balancing and resource utilization. A gated recurrent neural network can use the parallel processing techniques described in this specification to compute outputs in real-time or faster than real-time even for networks that have very high throughput requirements, e.g., autoregressive neural networks in which values of a signal, more particularly output data representing a distribution of values of a signal, are generated based upon previously generated values of the signal. Such networks can include audio-, video-, image-, or text-generation neural networks.

This specification describes how a neural network system can improve throughout and reduce latency by interleaving the matrix operations of a gated activation unit over multiple independent processing units of a parallel processing device. The examples described below will commonly refer to the independent processing units being streaming multiprocessors (SMs) having multiple processing cores and the parallel processing device being a graphics processing unit (GPU). However, the same techniques can also be implemented on other hardware devices that implement true thread parallelization with multiple independent processing units. Such devices include single instruction, multiple data (SIMD) processors generally, tensor processing units (TPUs), or other application-specific integrated circuits. In addition, where the examples mention the use of a GPU, this does not necessarily imply that graphics data is being processed or produced.

On such parallel processing devices, control over thread parallelization can be provided by program abstractions that define how threads are assigned to be executed by the multiple independent processing units. For clarity of presentation, this specification uses the terminology of common GPU program abstractions, but equivalent program abstractions that control how threads are scheduled on independent processing units can be used for other systems that are not GPUs.

A thread block, or for brevity, a block, is a group of threads that are executed by a single SM. Threads in a block can coordinate by making use of shared memory of the SM. Communication between threads in a block is therefore typically orders of magnitude faster than communicating with threads in other blocks.

A warp is a group of threads within a block and in some cases represents the smallest assignable unit of computation for a GPU. Threads in a warp typically execute instructions in lockstep. Thus, threads within a warp can, for example, fetch data from memory together. If the data required for each thread is stored in the same memory segment, all threads within a warp can read from memory with only one memory transaction. Common warp sizes are <NUM>, <NUM>, or <NUM> threads, to name just a few examples.

During execution, each block is assigned to be executed by one respective SM. The threads within the block execute on the SM until the block is complete, at which point another block can be executed by the SM. For SMs that have multiple processing cores, sometimes referred to as signal processors or execution lanes, each processing core can execute one thread. GPUs commonly have between <NUM> and <NUM> SMs, and each SM can have between <NUM> and <NUM> processing cores. Therefore, the GPU can therefore often execute hundreds of threads in parallel.

In order for blocks to make their computed data available to other blocks, the blocks must coordinate by performing a synchronization, or for brevity, a sync. Syncs are expensive performance-wise because executing the sync requires stopping execution of one or more blocks to wait. A synchronization requires all blocks to write their data to a place where another block can access it. The synchronization location can be in RAM or in the local memory of a single block.

In this specification, where a gated activation unit has two or more layers having two or more respective weight matrices, interleaving the matrix operations means that the majority of blocks of a parallel processing device, if not all blocks, will operate on some data from each of the two or more weight matrices.

<FIG> is a diagram that illustrates the segments of data required to compute the output of a single prior art GRU. As described above, a GRU has three weight matrices: an update weight matrix Wz <NUM>, a reset weight matrix Wr <NUM>, and a state weight matrix Wh <NUM>. Each of the weight matrices has a corresponding bias vector: an update bias vector bz <NUM>, a reset bias vector br <NUM>, and a state bias vector bh <NUM>. A GRU also receives input vectors that include a previous hidden state vector ht-<NUM> <NUM> and the current input vector xt <NUM>. Finally, a GRU has three bias vectors: an update bias vector bz <NUM>, a reset bias vector br, and a state bias vector bh <NUM>. Each matrix has N rows and M columns. Meanwhile, each of the input vectors <NUM>, <NUM> and the bias vectors, <NUM>, <NUM>, and <NUM>, have M/<NUM> data elements.

The data elements shown in <FIG> can be used to compute the output vectors of a GRU according to Equations (<NUM>)-(<NUM>) above.

<FIG> illustrates how the matrix operations of a GRU can be interleaved. The matrix operations being interleaved means that each block will receive some rows of each matrix.

As shown in <FIG>, one or more of the first rows 401z of the update weight matrix Wz <NUM> are assigned to block <NUM>. The next one or more rows 402z of the update weight matrix Wz <NUM> are assigned to a different block, block <NUM>.

Rows being assigned to a particular block means both that (<NUM>) the data from the rows will be copied into memory of the SM assigned to execute the block, and (<NUM>) that one or more threads belonging to the block will perform the corresponding matrix operations for the rows when the one or more threads are executed by the SM.

The actual matrix values, and therefore their corresponding matrix operations, can be assigned by thread or by warp. For example, each of the threads in a block can be assigned one or more values from one of the rows of the update weight matrix Wz <NUM>.

In some implementations, the system stores each matrix value with half precision, which allows each register to store two matrix values. Thus, each thread in the block can allocated a single register to store two adjacent row values.

As described above, on some parallel processing devices, a warp is the smallest assignable unit of work. Therefore, individual rows can be assigned to respective warps belonging to the block. The rows that are assigned to a single warp can be assigned in any appropriate way, e.g., continuously or grid-strided, so that the matrices are interleaved at the warp level in addition to being interleaved at the block level. In other words, every warp is assigned some values from all three weight matrices. In addition, the matrices can also be interleaved at the thread level, meaning that every thread can be assigned some values from all three weight matrices.

Within a row assigned to a warp, each thread is assigned a number of row values given by M / # threads per warp. For example, if M is <NUM>, each thread in a warp will be assigned <NUM> values per row assigned to the warp. If the values are stored in half precision, this can require using only <NUM> registers.

In order to exactly fit the model to the available processing hardware, the system can train a model in which M is a multiple of the number of SMs times the number of warps per block. For example, if there are <NUM> SMs and <NUM> warps per block, the system can train a model having a size that is a multiple of <NUM>, e.g., <NUM>, <NUM>, or <NUM>, to name just a few examples. In this example, the size of a model refers to the number of columns in each of the weight matrices.

Interleaving the rows among the blocks can require that corresponding rows have a same block or warp distribution. Thus, for every row n of a first weight matrix assigned to a block m, the block m is also assigned the same row n from the other two weight matrices. Similarly, if the rows are interleaved among the blocks and the warps, for every row n of a first weight matrix assigned to a warp m, the warp m is also assigned the same row n from the other two weight matrices.

As shown in <FIG>, one or more of the first rows 401r of the reset weight matrix Wz <NUM> are also assigned to block <NUM>, and the next one or more rows 402r of the reset weight matrix Wz <NUM> are assigned to a different block, block <NUM>. Similarly, one or more of the first rows <NUM> of the state weight matrix Wh <NUM> are assigned to block <NUM>, and the next one or more rows <NUM> of the state weight matrix Wh <NUM> are assigned to a different block, block <NUM>.

<FIG> illustrates how interleaving the matrices among the blocks can allow an output value for all GRU operations to be computed without a block synchronization. In this example, for simplicity it is assumed that only a single row from each matrix has been assigned to each block. Thus, the following rows are assigned to block <NUM><NUM>: row <NUM>401z of the update weight matrix Wz <NUM>, row <NUM>401r of the reset weight matrix Wr <NUM>, and row <NUM><NUM> of the state weight matrix Wh <NUM>. And the following rows are assigned to block <NUM><NUM>: row <NUM>402z of the update weight matrix Wz <NUM>, row <NUM>402r of the reset weight matrix Wr <NUM>, and row <NUM><NUM> of the state weight matrix Wh <NUM>.

The input vectors xt <NUM> and ht-<NUM> <NUM> can be broadcast in full to all blocks. Therefore, both blocks <NUM> and <NUM> hare copies of xt <NUM> and ht-<NUM> <NUM>.

Because the bias vectors bz, br, and bh are relatively small, they can either be broadcast in full to all blocks or also interleaved in the same way that the rows were interleaved, either by block, warp, or thread. In this example, the bias vectors are interleaved, and thus, block <NUM><NUM> has a copy of just the first element 312a of the update bias vector bz <NUM>, the first element 322a of the reset bias vector br <NUM>, and the first element 332a of the state bias vector <NUM>. And therefore, block <NUM> has a different set of elements from the bias vectors, having just the second element 312b of the update bias vector bz <NUM>, the second element 322b of the reset bias vector br <NUM>, and the second element 332b of the state bias vector <NUM>.

Once block <NUM><NUM> has this information, the threads of block <NUM> can implement all of the operations of the GRU defined by Equations (<NUM>)-(<NUM>) above without performing a block synchronization. In other words, no threads of block <NUM> need to interact with any threads from any other blocks. In this simple example, the interleaving allows block <NUM><NUM> to generate the first element <NUM> of the output vector ht. Similarly, block <NUM><NUM> can use the data assigned to block <NUM><NUM> to generate the second element <NUM> of the output vector ht. In a more realistic scenario, each block would generate multiple values of the output vector, and each block would generate roughly the same number of elements of the output vector ht.

And thus, in some implementations all values of the output vector ht can be computed without performing any block synchronizations. The values of the output vector ht can also be computed largely in parallel with optimal or nearly optimal load balancing, particularly if the system trains a model to exactly fit the available processing hardware. In this context, computing the output vector in parallel means that substantially all SMs of the processing hardware execute blocks concurrently to generate values of the output vector without performing block synchronizations between them.

<FIG> is a flow chart of an example process for using parallel processing hardware to compute a current state vector of a gated activation unit. Some of the actions described in <FIG> are performed by independent processing units of a parallel processing device, while some of the actions are preparatory actions that a deployment system having at least one central processing unit performs to load the appropriate data into memory of the independent processing units. For convenience, the process will be described as being performed by a system having all of these components, appropriately programmed in accordance with this specification.

The system receives a plurality of weight matrices of a gated activation unit of a neural network (<NUM>). As described above, the gated activation unit can be any appropriate neural network computation unit that receives a hidden state from a previous time step and a current input vector and generates a hidden state for a current time step. The gated activation unit can be, for example, a GRU or an LSTM unit. In some implementations, the system also receives, for each weight matrix, a respective corresponding bias vector.

The system interleaves rows of the plurality of weight matrices (<NUM>). As described above, interleaving the rows means assigning rows to thread blocks such that (<NUM>) every thread block gets some rows from every weight matrix, and (<NUM>) corresponding rows are assigned to the same thread block. The weight matrices are also interleaved by warp. Being interleaved by warp means that every warp gets some rows from every weight matrix and that corresponding rows are assigned to the same warp. They can be interleaved by thread, which means that every thread gets some values from every weight matrix and that corresponding values are assigned to the same thread.

The system receives input vectors for the gated activation unit (<NUM>). As described above, the input vectors include at least a current input vector and a hidden state for a previous time step. Some gated activation units have even more input vectors. For example, an LSTM also includes a cell state vector for a previous time step.

The system generates, by each thread block, a respective portion of a current state vector (<NUM>). After interleaving the rows of the weight matrices by blocks, warps, threads, or some combination of these, each computation unit has enough data to compute a corresponding portion of the hidden state for the current time step. In general, this requires performing, by each block, a number of linear operations on the weight matrix and bias vector values assigned to the block and then applying a nonlinear activation function.

A number of implementations have been described above that specifically refer to rows and columns of a matrix. However, the same techniques can be applied equally if rows are interpreted as columns and vice versa.

In this specification, computing values in parallel or substantially in parallel means independent processing units perform operations toward generating the values over time windows that at least partially overlap. Computing values in parallel or substantially in parallel does not require all values to be computed at precisely the same time or in lockstep. In addition, some values generated by a single independent processing unit may actually be generated serially by the processing unit, but over a time window that overlaps with the computation of other values by other independent processing units.

A computer program which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code) can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment.

To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and pointing device, e. g, a mouse, trackball, or a presence sensitive display or other surface by which the user can provide input to the computer. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone, running a messaging application, and receiving responsive messages from the user in return.

Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Intemet.

Claim 1:
A method implementing a neural network on a parallel processing device, the method comprising:
receiving a plurality of weight matrices (<NUM>, <NUM>, <NUM>) of a gated activation unit of the neural network, the gated activation unit having two or more layers, each layer defining operations comprising: (i) a matrix operation between a weight matrix for the layer and concatenated current input and previous hidden state vectors and (ii) a nonlinear activation operation using a result of the matrix operation;
interleaving rows of the plurality of weight matrices by assigning groups of corresponding rows (401z, 401r, <NUM>; 402z, 402r, <NUM>) to respective thread blocks (<NUM>, <NUM>), each thread block being a computation unit for execution by an independent processing unit of a plurality of independent processing units of a parallel processing device;
the method further comprising processing data using the neural network by:
receiving, by each thread block, input vectors (<NUM>, <NUM>) for a gated activation unit ;
generating, by each thread block, a respective portion of a current state vector (<NUM>, <NUM>) of the gated activation unit, including:
performing a plurality of partial matrix operations using one or more groups of corresponding rows of the plurality of weight matrices assigned to the thread block and a concatenation of the input vectors, the concatenation of the input vectors including at least a current input vector (<NUM>) and a hidden state vector for a previous time step (<NUM>), and
performing a plurality of nonlinear activation operations using respective results of the plurality of partial matrix operations; and
further comprising interleaving rows of the plurality of weight matrices by warps such that all warps receive some rows from every weight matrix and corresponding rows from every weight matrix are assigned to the same warp, wherein within a row assigned to a warp each thread is assigned a number of row values given by M / (number of threads per warp) where M is a number of columns of each weight matrix and where M and the number of threads per warp are both integers.