Memory-efficient backpropagation through time

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a recurrent neural network on training sequences using backpropagation through time. In one aspect, a method includes receiving a training sequence including a respective input at each of a number of time steps; obtaining data defining an amount of memory allocated to storing forward propagation information for use during backpropagation; determining, from the number of time steps in the training sequence and from the amount of memory allocated to storing the forward propagation information, a training policy for processing the training sequence, wherein the training policy defines when to store forward propagation information during forward propagation of the training sequence; and training the recurrent neural network on the training sequence in accordance with the training policy.

BACKGROUND

This specification relates to training recurrent neural networks.

Some neural networks are recurrent neural networks. A recurrent neural network is a neural network that receives an input sequence and generates an output sequence from the input sequence. In particular, a recurrent neural network can use some or all of the internal state of the network from a previous time step in computing an output at a current time step.

SUMMARY

In general, one innovative aspect of the subject matter described in this specification can be embodied in methods performed by one or more computers for training a recurrent neural network on a plurality of training sequences using backpropagation through time that include the actions receiving a training sequence including a respective input at each of a number of time steps; obtaining data defining an amount of memory allocated to storing forward propagation information for use during backpropagation; determining, from the number of time steps in the training sequence and from the amount of memory allocated to storing the forward propagation information, a training policy for processing the training sequence, wherein the training policy defines when to store forward propagation information during forward propagation of the training sequence; and training the recurrent neural network on the training sequence in accordance with the training policy.

Other embodiments of this aspect can include one or more of the following optional features. In some implementations, training the recurrent neural network on the training sequence includes forward-propagating the inputs in the training sequence from a first time step in the sequence to a last time step in the sequence; during the forward propagating, storing forward propagation information in accordance with the training policy; and backpropagating gradients from the last time step in the sequence to the first time step in the sequence, including determining, for each time step, whether additional forward propagation information is necessary to backpropagate the gradient for the time step and, if so, regenerating the additional forward propagation information using the stored forward propagation information.

In some implementations, the forward propagation information includes hidden states. In some implementations, the forward propagation information includes internal states. In some implementations, the forward propagation information includes hidden states and internal states. In some implementations, the training policy defines, for each time step, whether to store a hidden state, an internal state, or neither for the time step. In some implementations, determining the policy includes determining a policy that balances a trade-off between caching of forward propagation information and re-computation of forward propagation information during backpropagation.

In some implementations, determining the policy includes, for each time slot, determining a computational cost of storing an associated piece of forward propagation information based on the number of time steps and the amount of memory allocated to storing the forward propagation information, and adding to the training policy a set of pieces of forward propagation information that are determined to have the lowest computational cost; and—training the recurrent neural network on the training sequence in accordance with the training policy includes storing the set of pieces of forward propagation information that have been added to the training policy.

In some implementations, the methods include the actions of determining a computational cost of training the recurrent neural network on the training sequence in accordance with the policy. In some implementations, the methods include the actions of providing data identifying the computational cost for presentation to a user.

The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages. By employing the memory-efficient backpropagation through time techniques described in this specification, recurrent neural networks (RNNs) can be trained using training techniques that are less memory-intensive and efficiently utilize the memory allocated to train the RNNs. Accordingly, the time required to train an RNN can be reduced when the amount of memory available for training the RNN is less than the amount of memory needed to store all forward propagation information associated with the RNN. An RNN can be trained with processing circuits that can perform faster processing but have lower memory capacities, such as graphic processing units that have lower memory capacities than many central processing units. In addition, by optimizing the computational cost of training RNNs the methods described herein are more efficient with regard to computational resources and may be performed on less powerful computers.

DETAILED DESCRIPTION

FIG. 1is a block diagram of an example neural network system100. The neural network system100is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

The neural network system100includes a recurrent neural network (RNN)101, an RNN training engine102, an RNN training memory103, and a policy engine104. The RNN training engine102trains the RNN101by, in part, storing selected forward propagation information112in the RNN training memory103. The RNN training engine102determines which forward propagation information112to store in the RNN training memory103based on a training policy114determined by the policy engine104.

The RNN101is an example neural network having multiple parameters that includes at least one recurrent neural network layer. At each time step of multiple time steps, the RNN101is configured to process an input from a sequence of inputs and, for each time step except for the first time step, a hidden state from a preceding time step in accordance with current values of the parameters to generate an output for the time step and an updated hidden state for the time step for use at a subsequent time step.

Examples of tasks that can be performed using an RNN101include translation, conversion of text to speech, and generating sequential data such a sequence of words or letters. The RNN101can be a conventional RNN or a long-short term memory (LSTM) RNN. Example LSTM RNNs are described in Graves,Generating Sequences with Recurrent Neural Networks.

The RNN training engine102trains the RNNs by adjusting the values of the parameters of the RNN101to optimize an optimization function that depends on a measure of error between outputs generated by the RNN for a number of time steps and target outputs.

The RNN training engine102can train the RNN using a training algorithm known as backpropagation through time (BPTT). Example BPTT algorithms are described in Rumelhart et al.,Learning Internal Representations by Error Propagation, and Webos,Backpropagation Through Time: What It Does and How to Do It.

To train an RNN101using BPTT on a given training sequence, the RNN training engine102evaluates the gradient of the parameters of the RNN101with respect an objective loss function to optimize the objective loss function. The RNN training engine102can do this this in two steps. First, the RNN training engine102performs a forward propagation by processing inputs to the RNN101at each time step while saving only selected forward propagation information112into training memory103as indicated by the training policy104. After the forward propagation, the RNN training engine102calculates a measure of error between the outputs of the RNN101for the number of time steps and a target set of outputs and uses the measure of error to backpropagate gradients of the parameters of the RNN101with respect to the objective function throughout the RNN101.

The forward propagation information112allows the RNN training engine102to continue the forward propagation from a last, e.g., the most recent, saved time step of the RNN101whose corresponding forward propagation information are saved in the training memory103without reprocessing the time steps of the RNN101before the last saved time step.

To backpropagate gradients from a particular time step in a sequence of time steps of the RNN102, the RNN training engine102can read the last saved forward propagation information112in the sequence from the training memory103. The RNN training engine102can perform another forward propagation over the sequence starting at the time step of the last saved forward propagation information112and ending at the particular time step from which the RNN training engine102is backpropagating gradients in order to obtain forward propagation information of the particular time step. During the repeated forward propagation, the RNN training engine102may store intermediate states as defined by the policy104.

After the RNN training engine102obtains the forward propagation information112of the particular time step from which the RNN training engine102is backpropagating gradients, the RNN training engine102can compute the gradients of the particular time step and backpropagate gradients from the particular time step. Backpropagating gradients backward across a single time step is described in more detail in LeCun et al.,Efficient BackProp.

The RNN training engine102can repeat the process described above for backpropagating gradients from a particular time step of the RNN101until the gradients are backpropagated to the beginning of the sequence of time steps of the RNN101that includes the particular time step. For example, the RNN training engine102can repeat the process starting from a last time step in a sequence of time steps of the RNN101to a first time step in the sequence to backpropagate the gradients throughout the sequence.

In some implementations, the forward propagation information112for a time step of the RNN101only includes the hidden state output of the RNN101, while in other implementations the forward propagation information112for a time step can include all internal states of the RNN101at a given time step.

The internal states of the RNN101for a time step are the activations of the RNN for the time step. For example, in a traditional RNN, the internal states of the RNN101for a time step includes the hidden state for the time step, because a traditional RNN processes the hidden state to generate the updated hidden state for a time step. In other RNNs, however, the internal states of the RNN for a time step can include activations for the time step other than and in addition to the hidden state for the time step. For example, in an LSTM, the internal states of the LSTM may include values of gates of the LSTM101during the time step in addition to the hidden state of the time step, because the LSTM first processes the hidden state to update its gates and then uses the values of the gates to generate the updated hidden state for the time step. In some implementations, a hidden state of an RNN101for a time step is also part of the internal states of the RNN101for the time step.

To obtain a particular item of the forward propagation information112, the RNN training engine101can store the internal state in the RNN training memory103during forward propagation and retrieve the stored item during backpropagation. However, consistent application of this approach can lead to excessive storage costs especially for RNNs that execute for a large number of time steps.

Alternatively, the RNN training engine102can obtain a particular item of the forward propagation information112by performing a forward propagation of all or part of time steps of the RNN101to reconstruct the particular item. For example, if the RNN training engine102has not stored any forward propagation information112in the RNN training memory103, the RNN training engine102can obtain a hidden state value for an ith time step of the RNN101by performing a forward propagation from the first time step of the RNN101to a (i−1)th time step of the RNN101. If the RNN training engine101has stored the hidden state for a time step of the RNN101, the RNN training engine102can process the hidden state to generate the internal states for the time step using only one forward propagation, i.e., the propagation through the particular time step.

Thus, the RNN training engine102can use stored forward propagation information112to re-generate forward propagation information112that are not stored. This approach can reduce the amount of memory needed to train an RNN. However, consistent application of this approach can make training the RNN101computationally costly due to the computational costs associated with repeated forward propagations.

The policy engine104receives as inputs an amount of memory113allocated to storing forward propagation information112for use during backpropagation and a number of time steps111during which the RNN is being trained101. The policy engine104uses those inputs to determine a training policy114. The training policy114defines when to store forward propagation information during forward propagation of the training sequence. In other words, the training policy114defines which items in the forward propagation information112to store in the RNN training memory103.

The RNN training engine102uses the training policy114to determine which items of items in the forward propagation information112to store in the RNN training memory103and which items in the forward propagation information112to re-generate during backpropagation.

Determining a training policy114is described in greater detail below with reference toFIG. 3.

FIG. 2is a flow diagram of an example process200for training a recurrent neural network. For convenience, the process200will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system100ofFIG. 1, appropriately programmed in accordance with this specification, can perform the process200.

The system receives a training sequence including a respective input to the RNN at each of a number of time steps (210) and obtains data defining an amount of memory allocated to storing forward propagation information for use during backpropagation (220). The system generally samples the training sequence from a set of training data for use in training the RNN. The amount of allocated memory is usually supplied by a computer application that manages memory resources for the system, such as an operating system or other memory management application. The computer application managing memory resources for the system allocates a certain amount of memory to the system for storing forward propagation. The amount of memory is an amount in a particular unit of memory, e.g., bytes, and can be memory available for use by the system in a single storage device or in multiple physical storage devices in one or more physical locations.

The system determines a training policy (230) from the number of time steps in the training sequence and from the amount of memory allocated to storing the forward propagation information. In some implementations, the system determines a training policy that balances a trade-off between caching of forward propagation information and re-computation of forward propagation information during backpropagation. Determining a training policy114is described in greater detail below with reference toFIG. 3.

The system then trains the recurrent neural network (240) on the training sequence in accordance with the training policy.

In some implementations, the system forward propagates the inputs in the training sequence from a first time step in the sequence to a last time step in the sequence. During the forward propagation, the system stores some forward propagation information in accordance with the training policy. After the forward propagation, the system backpropagates gradients from the last time step in the sequence to the first time step in the sequence. During backpropagation, for each time step, the system determines whether additional forward propagation information items are necessary to backpropagate the gradient for the time step, i.e., based on determining whether the system has stored all the forward propagation information necessary to backpropagate the gradient for the time step. If the system determines that additional forward propagation information items are necessary to backpropagate the gradient for the time step, the system re-generates the additional forward propagation information items using the stored forward propagation information.

The system can re-generate a particular item of the forward propagation information by performing a forward propagation of all or part of time steps of the RNN. For example, the system can generate a hidden state value for an ith time step of the RNN101by performing a forward propagation from the first time step of the RNN to a (i−1)th time step of the RNN. The system can also process a hidden state for an ith time step to generate the internal states for the (i+1)th time step using only one forward propagation, i.e., the propagation through the (i+1)th time step.

FIG. 3is a flow diagram of an example process300for determining which forward propagation information to store. For convenience, the process300will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system100ofFIG. 1, appropriately programmed in accordance with this specification, can perform the process300.

The system obtains a count of time steps during which a recurrent neural network is trained (302) and an amount of memory allocated to storing forward propagation information for use during backpropagation (304). The count of time steps during which an RNN is trained is generally equal to the number of inputs in a batch of input sequences for an RNN.

The system determines forward propagation information that are candidates for storage (306). The candidate forward propagation information can include hidden states for time steps of the RNN and/or internal states for time steps of the RNN.

The system identifies one or more strategies for storing forward propagation information from the training memory (308). A strategy includes a sequence of candidates for storage that the system will store if the system is following the strategy.

The system determines a computational cost for each strategy (310). The computational cost of a strategy is an estimated computational cost, e.g., in terms of the number of forward propagations needed to generate all forward propagation information necessary to perform backpropagation during training of the recurrent neural network, of training the RNN if the system follows the strategy.

In some implementations, the computational cost of a strategy that includes saving a hidden state of a time step after choosing not to save hidden states for y time steps is given by the following equation for Q1(t,m,y), where t is a number of time steps over which backpropagation is performed and m is a number of available memory units:
Q1(t,m,y)=y+C(y,m)+C(t−y,m−1)

The computational cost of a strategy that includes saving internal states of a time step after choosing not to save internal states for y time steps is given by the following equation for Q2(t,m,y), where t is a number of time steps over which backpropagation is performed, m is a number of available memory units, and a is a ratio of size of internal states for time steps to the size of hidden states for time steps:
Q2(t,m,y)=y+C(y−1,m)+C(t−y,m−α)

In the above equations, C(t,m) is the optimal cost, i.e., the lowest possible computational cost, of backpropagating over a sequence of t time steps given the amount of allocated memory equal to m. The system can compute the C(t,m) values in the equations using dynamic programming with the following boundary conditions: C(t,1)=½t(t+1) and C(0,m)=0

The system selects the strategy having the lowest computational cost (312) and determines a position of next storage and a type of forward propagation information from the selected strategy (314).

The system can select the forward propagation information having the type i (i.e., i=1 for hidden states and i=1 for internal states) and being at a position y in a sequence of time steps in a manner that produces a lowest computational cost strategy and determine to store the selected forward propagation information. In other words, the system selects i and y in accordance with argmini,yQi(t,m,y).

The system can repeat process300to generate a training strategy for a sequence of time steps of an RNN. For example, the system can use the process300to determine a position and a type of selected forward propagation information that the system will first save. After saving the selected forward propagation information, the system can divide sequence of time steps into two subsequences, a first subsequence including time steps before the time step of the selected forward propagation information as well as the time step of the selected forward propagation information and a second subsequence including time steps after the time step of the selected forward propagation information. The system can perform process300on each subsequence given the count of time steps of each subsequence using a divide-and-conquer strategy to select a position and type of next forward propagation information to save.

Determining a training policy with the technique described above that uses a recursive divide and conquer approach with dynamic programming can be performed using an algorithm laid out in the pseudocode below:

Executing a computed training policy can be performed using an algorithm laid out in the pseudocode below:

FIG. 4is an operational example400that illustrates saving forward propagation information. The operational example400illustrates saving a hidden state at a position y and dividing a sequence of time steps including the time step at the position y into two subsequences: a first subsequence that includes the portion of the sequence before and including the position y and a second subsequence after position y.

A training engine configured to train an RNN can perform the division technique illustrated in the operational example400on subsequences in a recursive manner and using a divide-and-conquer approach. By doing this, the training engine can determine the computational cost of a strategy that involves saving forward propagation information at a position y and following an optimal, i.e., lowest computational cost, training policy after storing the forward propagation information at the position y. The system can use the determined computational cost to perform a process for determining which forward propagation information to store, e.g., the process300ofFIG. 3.