Training a neural network using layer-wise losses

Abstract

The present disclosure relates to methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a neural network using layer-wise losses.

Description

TECHNICAL FIELD

[0001] The present disclosure relates to training a neural network using layer-wise losses. BACKGROUND

[0002] This specification relates to training a neural network.

[0003] A neural network is a machine learning model that employs one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters. SUMMARY

[0004] This specification describes a system implemented as computer programs on one or more computers in one or more locations that train a neural network that processes a network input to generate a network output. In particular, the systems described in this specification train a neural network using layer-wise losses so that weight updates for the layers of the neural network can be computed in parallel for each layer in the neural network.

[0005] Particular embodiments of the subject matter described in this specification can be implemented so as to realize one or more of the following advantages.

[0006] This specification describes techniques for training a neural network using layer-wise updates, e.g., updates based on a matching loss for a transfer function of a layer of the neural network. Training using the techniques allows a system to take multiple gradient steps independently and in parallel for all local layer-wise problems. Training a neural network in this way results in a neural network that outperforms a neural network trained using conventional backpropagation techniques and competes with a neural network trained using second-order methods, and in some cases outperforms a neural network trained using second-order methods, while consuming fewer computational resources than these second-order methods, i.e., because second-order methods require careful tuning for the task at hand, e.g., through computationally expensive hyperparameter search. Because the local problems are independent of each other, the internal update can run in parallel, making it much faster than running multiple forward-backward steps. In comparison to second-order methods, the described techniques are significantly easier to implement and scale to larger networks, because second-order methods typically rely on computing inverses and do not scale well when the number of parameters is large.

[0007] Furthermore, training using the described techniques allows a system to efficiently parallelize training and to train layers independently in parallel. Because the devices assigned to each layer are primarily focused on computing local, internal updates, training can be easily distributed across multiple devices.

[0008] In other words, the described techniques exploit parallelism in order to improve the quality of network training relative to conventional backpropagation with minimal additional computational overhead.

[0009] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF DRAWINGS

[0010] Figure 1 An example training system is shown.

[0011] Figure 2 is a flow diagram of an example process for performing a training step during training of a neural network.

[0012] Figure 3 is a flow diagram of an example process for performing an update iteration based on pre-activation to minimize a squared local loss.

[0013] Figure 4 is a flow diagram of an example process for performing an update iteration based on post-activation to minimize a squared local loss.

[0014] Figure 5 is a flow diagram of an example process for performing an update iteration to minimize a local matching loss.

[0015] Figure 6 is a flow diagram of an example process for performing an update iteration to minimize a bi-Bregman divergence loss.

[0016] In the various drawings, like reference numbers and designations indicate like elements. DETAILED DESCRIPTION

[0017] Figure 1 An example training system 100 is shown. Training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, where the systems, components, and techniques described below can be implemented.

[0018] System 100 trains a neural network 110 configured to perform a particular machine learning task on training data 130. That is, neural network 110 is configured to process a network input 112 to generate a network output 114 for the network input 112 for the particular machine learning task.

[0019] Neural network 110 can be trained to perform any kind of machine learning task, that is, can be configured to receive any kind of numerical data input and generate any kind of score, classification, or regression output based on that input.

[0020] In some cases, the neural network 110 is a neural network configured to perform an image processing task (i.e., receive an input image) and process the input image (i.e., process intensity values of pixels of the input image) to generate a network output for the input image. For example, the task can be image classification, and the output generated by the neural network for a given image can be a score for each of a set of object classes, where each score represents an estimated likelihood that the image contains an object belonging to that class. As another example, the task can be image embedding generation, and the output generated by the neural network can be a numerical embedding of the input image. As yet another example, the task can be object detection, and the output generated by the neural network can identify locations in the input image that depict objects of a particular type. As yet another example, the task can be image segmentation, and the output generated by the neural network can assign each pixel of the input image to a class from a set of classes.

[0021] As another example, if the input to the neural network 110 is a feature of an impression context for a particular advertisement, the output generated by the neural network can be a score representing an estimated likelihood that the particular advertisement will be clicked.

[0022] As another example, if the input to the neural network 110 is a feature of an impression context for a particular advertisement, the output generated by the neural network can be a score representing an estimated likelihood that the particular advertisement will be clicked.

[0023] As another example, if the input to the neural network 110 is a feature of an impression context for a particular advertisement, the output generated by the neural network can be a score representing an estimated likelihood that the particular advertisement will be clicked.

[0024] As another example, if the input to the neural network 110 is a sequence of text in one language, the output generated by the neural network can be a score for each of a set of pieces of text in another language, where each score represents an estimated likelihood that the piece of text in the other language is a correct translation of the input text into the other language.

[0025] As another example, the task can be an audio processing task. For example, if the input to the neural network 110 is a sequence representing a spoken utterance, the output generated by the neural network can be a score for each of a set of text snippets, each score representing an estimated likelihood that the text snippet is a correct transcription of the utterance. As another example, the task can be a keyword spotting task, in which, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase (a "hotword") was spoken in the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can identify the natural language in which the utterance was spoken.

[0026] As another example, the task can be a natural language processing or understanding task operating on a sequence of text in a certain natural language, e.g., an entailment task, a paraphrasing task, a text similarity task, a sentiment task, a sentence completion task, a syntax task, etc.

[0027] As another example, the task can be a text-to-speech task, in which the input is text in a natural language or features of text in a natural language, and the network output is a spectrogram or other data defining audio of the text spoken in the natural language.

[0028] As another example, the task can be a health prediction task, in which the input is electronic health record data for a patient, and the output is a prediction related to the future health of the patient, e.g., a predicted treatment that should be prescribed for the patient, a likelihood that the patient will experience an adverse health event, or a predicted diagnosis for the patient.

[0029] As another example, the task can be an agent control task, in which the input is an observation characterizing a state of an environment, and the output defines an action to be performed by an agent in response to the observation. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility, or a control system controlling a different kind of agent.

[0030] The training data 130 includes a set of training inputs, and, for each training input, a label. The label for a given training input specifies the network output that should be generated by performing the machine learning task on the given training input, i.e., is the target output that should be generated by the neural network 110 after training.

[0031] The neural network 110 can have any appropriate architecture that allows the neural network 110 to perform the particular machine learning task, i.e., to map network inputs of the type and dimension required by the task to network outputs of the type and dimension required by the task.

[0032] As one example, when the input is an image, the neural network 110 can be a convolutional neural network (e.g., a neural network with a ResNet architecture, an Inception architecture, an EfficientNet architecture, etc.) or a Transformer neural network (e.g., a Vision Transformer).

[0033] As another example, when the input is text, features of a medical record, audio data, or other sequential data, the neural network 110 can be a recurrent neural network (e.g., a neural network based on long short-term memory (LSTM) or gated recurrent units (GRU)) or a Transformer neural network.

[0034] As another example, the neural network can be a feedforward neural network that includes multiple fully connected layers, e.g., an MLP.

[0035] In general, however, the neural network 110 includes multiple layers 116A-N, each layer having respective weights.

[0036] In particular, each of the multiple layers 116A-N is configured to receive a layer input and apply the respective weights of the layer to the layer input to generate a pre-activation of the layer. How the layer 116A-N applies the weights to the layer input depends on the type of neural network layer. For example, a convolutional layer computes a convolution between the weights and the layer input. As another example, a fully connected layer computes a product between the weights of the layer and the layer input.

[0037] Each of the multiple layers 116A-N is then configured to apply a transfer function of the layer to the pre-activation to generate a post-activation (i.e., a layer output of the layer), and then provide the post-activation to one or more other layers of the neural network that are configured to receive input from the layer in accordance with the neural network architecture. The transfer function of any given layer is an element-wise non-linear function, and different layers can have different transfer functions. Examples of transfer functions include ReLU, leaky ReLU, Tanh, and ArcTan. That is, another example of a transfer function is the identity function for linear layers that do not have an activation function.

[0038] The neural network 110 can have additional layers and components that do not have weights, e.g., normalization layers, pooling layers, residual connections, etc.

[0039] Accordingly, to train the neural network 110, the training system 100 repeatedly updates the weights of the multiple layers 116-N using the training data 130 at different training steps to minimize a task loss function. The task loss function can be any appropriate differentiable loss function that is suitable for the particular task. Examples of task loss functions include cross-entropy loss, squared error loss, negative log-likelihood loss, etc.

[0040] In particular, at each training step, the system 100 performs a forward pass through the neural network or a backward pass through the neural network to determine layer inputs and target pre-activations or post-activations for each layer. The system 100 then performs multiple local update iterations for each layer to update the weights of the layer using the layer inputs and the target pre-activations or post-activations.

[0041] Reference will be made to the drawings throughout the description, like reference numerals being used for like elements throughout. Figures 2-4 The execution of a training step will be described in more detail.

[0042] In some implementations, the system 100 distributes the training of the neural network 100 across multiple devices.

[0043] In particular, the system 100 can distribute the training of the neural network 100 across multiple devices 118A-118N. Each device can be, for example, a CPU, GPU, TPU, or other ASIC, FPGA, or other computer hardware configured to perform the operations required to compute the layer outputs of at least one of the computational layers 116A-N and to compute the gradient of the loss function.

[0044] The system 100 can distribute the training of the neural network 100 in any of a variety of configurations. For example, as shown in FIG. 1, the system 100 can assign each of the layers 116A-116N to a different one of the devices 118A-118N. As another example, the system 100 can assign different partitions of the layers (which can include multiple layers) to each of the devices 118A-118N. Figure 1

[0045] By distributing the training across devices, the system 100 can ensure that sufficient computational resources are available to perform the local update steps for each of the layers 116A-116N in parallel at each training step. By performing the local update steps in parallel, the system 100 achieves the benefits of multiple update steps while minimizing the additional computational overhead required to perform multiple steps, i.e., instead of a single update step as performed by conventional first-order optimizers.

[0046] Following training, the trained student neural network 110 is deployed by the training system 100 or a different inference system 170 on one or more computing devices to perform inference, i.e., to generate new network outputs 114 for a machine learning task for new network inputs 112.

[0047] Figure 2 is a flowchart of an example process 200 for performing a training iteration during the training of a neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of Figure 1 , appropriately programmed, can perform the process 200. ​

[0048] The system can repeatedly perform iterations of the process 200 to repeatedly update the network parameters until a termination criterion is satisfied, e.g., until a threshold number of iterations of the process 200 have been performed, until a threshold amount of wall-clock time has elapsed, or until the values of the network parameters have converged.

[0049] The system obtains a batch comprising one or more training inputs and a respective label for each training input (step 202). The system will typically obtain different training inputs at different iterations, e.g., by sampling a fixed number of inputs from a larger set of training data at each iteration. The label for each training input identifies a target output for the training input that should be generated by performing a particular machine learning task on the training input.

[0050] The system performs a forward pass through the neural network in order to generate a respective training network output for each training input in the batch (step 204). That is, the system processes each training network input through each layer in the neural network to generate a training output for that network input. As part of performing the forward pass, the system determines, for each training input in the batch and for each layer of the neural network, a respective layer input to the layer that is generated during processing of the training input.

[0051] The system performs a backward pass through the neural network for each training input using the training output for that training input and the label for the training input to determine, for each layer of the neural network and for each training input, an estimated target for the neural network layer (step 206).

[0052] In some implementations, the estimated target is an estimated target pre-activation. For example, given an estimated gradient descent (GD) target pre-activation a m may be satisfied:

[0053]

[0054] where, is the current pre-activation of the layer, is the layer input to the layer, W m is the weight of the layer, and γ is a constant greater than zero representing an activation learning rate, is a task loss evaluated at the training output for the training input and the label for the training input, and represents a gradient with respect to .

[0055] As another example, an estimated double-mirrored descent (Double MD) target pre-activation a m may be satisfied:

[0056]

[0057] wherein, is the current pre-activation of the layer, is the layer input of the layer, W m is the weight of the layer, and γ is a constant greater than zero representing an activation learning rate, is the task loss evaluated at the training output for the training input and the label of the training input.

[0058] In some other implementations, the estimation target is an estimation target post-activation.

[0059] As one example, given an estimation target GD post-activation y m can satisfy:

[0060]

[0061] wherein, is the current post-activation of the layer, and f m is the transfer function of the layer m.

[0062] As another example, given an estimation target mirror-descent (MD) post-activation y m can satisfy:

[0063]

[0064] wherein, and f m is the transfer function of the layer m.

[0065] In any of the above implementations, the system can compute the corresponding target by backpropagating the gradient of the task loss through the neural network using regular techniques and reusing the pre-activations or post-activations from the forward step or recomputing them during the backward step.

[0066] For each layer, the system then performs a number of update iterations to determine a final updated weight for the layer using (i) the layer input generated for the training input of the layer and (ii) the estimation target for the training input of the layer for each training input and each layer (step 208).

[0067] For a given layer, at each update iteration, the system computes the gradient of the weight of the layer with respect to the local per-layer loss and uses the gradient to update the current weight of the layer. Then, the system uses the updated weight as the final updated weight for the given layer after performing the last training iteration, i.e., the weight to be used for performing the next iteration of the process 200.

[0068] In particular, once the forward and backward passes are performed, the system can perform multiple update iterations independently and in parallel for each layer, as the layer inputs and the estimation targets remain fixed and reused at each update iteration, ensuring that no information from any other layer is needed to perform the multiple update iterations.

[0069] For example, respective devices can be assigned to perform updates for each layer, and each device can perform update iterations for the layer assigned to that device in parallel with each other device.

[0070] In some implementations, each device includes a copy of each of the neural network layers, and is assigned to perform updates for a respective set of one or more of the layers. In these implementations, each device can independently perform the forward and backward passes, and then, after performing step 206, (i) provide the final updated weights for access by the hardware devices performing operations of the other neural network layers, and (ii) obtain the final updated weights of the other neural network layers of the plurality of neural network layers for use in performing forward and backward passes through the neural network, i.e., at the next iteration of process 200.

[0071] In some other implementations, each device includes only a copy of the layers assigned to that device. In these implementations, to perform the forward pass, each device receives the layer inputs for the layers assigned to that device, processes the layer inputs using the corresponding layers according to the layers’ weights, and then provides the layer outputs to the device to which the next layer in the network architecture is assigned.

[0072] By performing multiple update iterations, i.e., instead of a single update iteration, the system can improve the quality of the training process relative to first-order training techniques. By ensuring that the update iterations are local to each layer and that the update iterations are performed in parallel for all layers, the system ensures that the additional training quality is achieved with minimal additional computational overhead relative to first-order training techniques.

[0073] Figure 3 is a flowchart of an example process 300 for performing an update iteration to minimize a squared local loss based on pre-activations of a given layer. For convenience, process 300 will be described as performed by a system of one or more computers located in one or more locations. For example, a training system, such as training system 100 of Figure 1 may perform process 300.

[0074] The system can perform a fixed number T of update iterations for a given layer at each iteration of the training process, i.e., at each iteration of process 200.

[0075] Before any iteration of the process 300, the system obtains, for each training input, the layer input for the training input and the estimated GD target preactivation for the training input, i.e., as a result of the forward and backward passes described above with reference to Figure 2 the process 200.

[0076] The system identifies the current weights for the layer (step 302). For the first update iteration, the current weights are the weights at the end of the previous iteration of the process 200. For each subsequent iteration, the current weights are the weights at the end of the previous update iteration, i.e., the updated weights after the previous iteration of the process 300.

[0077] The system uses the layer input for the training input in the batch and the estimated target GD preactivation for the training input in the batch to compute, from the current weights for the particular neural network layer, the gradient of the weights of the given neural network layer with respect to the squared local loss (step 304).

[0078] In particular, the squared local loss includes two terms: (i) a squared loss between the preactivation generated from the updated weights and the GD target preactivation, and (ii) a regularization term that penalizes the layer for the difference between the current weights and the updated weights. For example, the squared local loss for layer m can satisfy:

[0079]

[0080] where, is the updated weights for the layer, is the layer input for the layer, a m is the GD target preactivation for the layer input, W m is the current weights for the layer, and η is a constant greater than zero that controls the trade-off between the minimized loss and the regularization term.

[0081] To compute the gradient of the loss at a given update iteration, the system computes a new preactivation by applying the current weights to the layer input, and computes the difference between the new preactivation and the estimated GD target preactivation. The system then computes the gradient based on the difference. In particular, the gradient is equal to:

[0082]

[0083] Thus, the system keeps the layer input for the training input and the estimated target preactivation for the training input fixed across all update iterations, ensuring that performing the update iterations does not require any additional backward pass through the neural network and forward pass through the neural network, and, thus, the update iterations can be performed independently and in parallel for each layer.

[0084] The system updates the current weights of the particular neural network layer using the gradient (step 306). For example, the system can subtract the gradient from the current weights to generate updated weights.

[0085] Figure 4 is a flowchart of an example process 400 for performing update iterations to minimize a squared local loss based on a post-activation of a given layer. For convenience, the process 400 will be described as performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of FIG. 1) appropriately programmed can perform the process 400. Figure 1

[0086] The system can perform a fixed number T of update iterations for a given layer at each iteration of the training process, i.e., at each iteration of the process 200.

[0087] Before performing any iteration of the process 400, the system obtains, for each training input, a layer input for the training input and an estimated GD target post-activation for the training input, i.e., as a result of performing the forward pass and the backward pass described above with reference to Figure 2

[0088] The system identifies current weights for the layer (step 402). For the first update iteration, the current weights are the weights at the end of the previous iteration of the process 200. For each subsequent iteration, the current weights are the weights at the end of the previous update iteration, i.e., the updated weights after the previous iteration of the process 400.

[0089] The system computes, from the current weights of the particular neural network layer, a gradient of the weights of the given neural network layer with respect to the squared local loss using the layer inputs for the training inputs in the batch and the estimated GD target post-activations for the training inputs in the batch (step 404).

[0090] In particular, the squared local loss includes two terms: (i) a squared loss between the post-activations generated from the updated weights and the GD target post-activations, and (ii) a regularization term that penalizes the layer for the difference between the current weights and the updated weights. For example, the squared local loss for the layer m can satisfy:

[0091]

[0092] where y m is the GD target post-activation of the layer input, w m is the current weight of the layer, and η is a constant greater than zero that controls the trade-off between the minimized loss and the regularization term.

[0093] ​​To compute the gradient of this loss at a given update iteration, the system computes new pre-activations by applying the current weights to the layer inputs, and computes new post-activations by applying the transfer function to the new pre-activations, and then computes the difference between the new post-activations and the estimated GD target post-activations. The system then computes the gradient based on this difference. In particular, the gradient is equal to:

[0094]

[0095] where, is the transpose of the Jacobian of the transfer function f m .

[0096] Thus, the system keeps the layer inputs for the training inputs and the estimated target post-activations for the training inputs fixed across all update iterations, ensuring that performing the update iterations does not require any additional backward pass through the neural network and forward pass through the neural network, and thus, the update iterations can be performed independently and in parallel for each layer.

[0097] The system uses the gradient to update the current weights of the particular neural network layer (step 406). For example, the system can subtract the gradient from the current weights to generate updated weights.

[0098] Figure 5 is a flowchart of an example process 500 for performing update iterations to minimize the local matching loss for a given layer. For convenience, the process 500 will be described as performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of Figure 1 ) appropriately programmed can perform the process 500.

[0099] The system can perform the update iterations for a given layer a fixed number of times T at each iteration of the training process, i.e., at each iteration of the process 200.

[0100] Prior to performing any iteration of the process 500, the system obtains the layer inputs for the training inputs and the estimated MD target post-activations for the training inputs for each training input, i.e., as a result of performing the forward pass and backward pass described above with reference to Figure 2 .

[0101] The system identifies the current weights of the layer (step 502). For the first update iteration, the current weights are the weights at the end of the previous iteration of the process 200. For each subsequent iteration, the current weights are the weights at the end of the previous update iteration, i.e., the updated weights after the previous iteration of the process 500.

[0102] The system computes, for a given neural network layer, a gradient of the weights of the layer with respect to the local matching loss of the transfer function of the layer, based on the current weights of the given neural network layer, using the layer input of the training input in the batch and the estimated MD post-activation of the training input in the batch (step 504).

[0103] The matching loss of a transfer function f is a measure of the difference between the target output of the transfer function and the actual output of the transfer function. In particular, the matching loss L of a transfer function f is defined as: f

[0104]

[0105] where a is the target pre-activation.

[0106] The matching loss of various public transfer functions is shown in Table 1 below.

[0107]

[0108] Table 1

[0109] In particular, the local matching loss includes two terms: (i) a matching loss between the post-activation generated from the updated weights and the target MD post-activation, and (ii) a regularization term that penalizes the layer for the difference between the current weights and the updated weights. For example, the local matching loss of layer m can satisfy:

[0110]

[0111] where, is the updated weights of the layer, is the layer input of the layer, y m is the MD target post-activation of the layer input, w m is the current weights of the layer, is the matching loss of the transfer function f m of the layer, and η is a constant greater than zero that controls the trade-off between the minimization loss and the regularization term.

[0112] To compute the gradient of such loss at a given update iteration, the system computes a new pre-activation by applying the current weights to the layer input, a new post-activation by applying the transfer function to the new pre-activation, and the difference between the new post-activation and the estimated MD target post-activation. Then, the system computes the gradient based on the difference. In particular, the gradient is equal to:

[0113]

[0114] ​Accordingly, the system maintains the layer inputs for the training inputs and the estimated target post-activations for the training inputs fixed throughout all update iterations, ensuring that performing an update iteration does not require any additional backward pass through the neural network and forward pass through the neural network, and, thus, the update iterations can be performed independently and in parallel for each layer. Additionally, although different transfer functions can have different matching losses, computing the gradients only requires the values of the layer inputs and the difference between the post and the MD target post-activations, allowing the process 500 to be used for layers with a variety of different transfer functions.

[0115] The system updates the current weights of the particular neural network layer using the gradients (step 506). For example, the system can subtract the gradients from the current weights to generate updated weights.

[0116] Figure 6 is a flowchart of an example process 600 for performing update iterations to minimize a Bregman divergence-based loss for a given layer. For convenience, the process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of FIG. 1) appropriately programmed can perform the process 600. Figure 1

[0117] The system can perform the update iterations for a given layer a fixed number of times T at each iteration of the training process, i.e., at each iteration of the process 200.

[0118] Before performing any iteration of the process 600, the system obtains, for each training input, the layer inputs for the training input and the estimated double MD target pre-activations for the training input, i.e., as a result of performing the forward and backward passes described above with reference to the process 200. Figure 2

[0119] The system identifies the current weights of the layer (step 602). For the first update iteration, the current weights are the weights at the end of the previous iteration of the process 200. For each subsequent iteration, the current weights are the weights at the end of the previous update iteration, i.e., the updated weights after the previous iteration of the process 600.

[0120] The system computes, using the layer inputs for the training inputs in the batch and the estimated double MD target pre-activations for the training inputs in the batch, the gradients of the weights of the given neural network layer with respect to the local matching loss of the transfer function of the layer from the current weights of the layer (step 604).

[0121] In particular, the loss includes two terms: (i) a loss between the dual of the Bregman divergence between the post-activation generated from the updated weights and the post-activation generated from the double MD target pre-activations, and (ii) a regularization term penalizing the difference between the current weights and the updated weights. For example, the loss for the layer m can satisfy:​​

[0122]

[0123] in It is the dual of the Bregman divergence, and a m It is a dual MD target preactivation used for layer input.

[0124] To compute the gradient of the loss in a given update iteration, the system computes a new pre-activation by applying the current weights to the layer input and calculates the difference between the new post-activation and the estimated dual MD target pre-activation. The system then computes the gradient based on this difference. Specifically, the gradient equals:

[0125]

[0126] in It is the transfer function f m The transpose of the Jacobian determinant, and a m It is a double MD target preactivation of the layer input.

[0127] Therefore, the system keeps the layer inputs used for training inputs and the estimated target preactivations used for training inputs fixed in all update iterations, ensuring that performing update iterations does not require any additional backpasses through the neural network and forwardpasses through the neural network, and thus, update iterations can be performed independently and in parallel for each layer.

[0128] The system updates the current weights of a specific neural network layer using gradients (step 606). For example, the system can subtract the gradient from the current weights to generate updated weights.

[0129] Figures 3 to 6 The description describes the computation of the gradient of a single training input. When the batch includes multiple training inputs, the system can combine these gradients at each update iteration, for example, by averaging or summing these gradients, and then use the combined gradients to update the weights at the update iteration, i.e., using the combined gradients in steps 306, 406, 506, or 606 to update the current weights at the update iteration.

[0130] Furthermore, the above description describes the product between the layer input and the weight matrix (i.e., ...) Preactivation can be generated by computation. However, more generally, preactivation can be generated by computation of any linear transformation that depends on the current weights of the layer and the layer input. As another example, besides matrix-vector multiplication, the linear transformation can be a convolution between the kernel of weights and the layer input, i.e., for convolutional layers.

[0131] The specification uses the term“configured” in connection with systems and computer program components. For a system of one or more computers to be“configured to” perform particular operations or actions means that the system has installed on it software, firmware, hardware or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

[0132] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied

[0133] The term“data processing apparatus” refers to data processing hardware and includes all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can be or further include special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0134] A computer program, which can 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. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), in a single file dedicated to the program in question, or in multiple coordinated files (e.g., files that store one or more modules, sub programs, or portions of code). A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and are interconnected by a data communication network.

[0135] In this specification, the term“database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, an index database can include multiple collections of data, each of which can be organized and accessed differently.

[0136] Similarly, in this specification, the term“engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components that are installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer(s).

[0137] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0138] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name just a few.

[0139] Computer readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0140] 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 a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device of the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. 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 that is running a messaging application, and receiving responsive messages as a

[0141] Data processing apparatuses for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing the common and compute-intensive portions of machine learning training or production (i.e., inference) workloads.

[0142] A machine learning framework (e.g., a TensorFlow framework) can be used to implement and deploy machine learning models.

[0143] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0144] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., data files, such as HTML pages, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at a server from the device.

[0145] While this specification contains many specific implementation details, these should not be construed as limitations on the scope or range of any invention, but rather as descriptions of features that can be specific to certain embodiments. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination or in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features can be described above as acting in certain combinations and even initially claimed as such, one or more features from a claimed combination can in some cases be excised from the combination and the claimed combination can be directed to a subcombination or variation of a subcombination.

[0146] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing can be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0147] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the acts recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order or sequential order illustrated, to achieve desirable results. In certain circumstances, multitasking and parallel processing can be advantageous.

Claims

1. A method for training a neural network having a plurality of neural network layers, each of the plurality of neural network layers having a respective set of weights, the method comprising repeatedly performing, for each particular neural network layer of the plurality of neural network layers, operations comprising: obtaining a batch comprising one or more training inputs and a respective label for each training input; for each training input in the batch: performing a forward pass through the neural network on that training input to determine at least a layer input to the particular neural network layer and a training output for that training input, and performing a backward pass through the neural network using the training output for that training input and the label for that training input to determine an estimated target post-activation for the particular neural network layer; and performing a plurality of update iterations to determine final updated weights for the particular neural network layer, wherein performing each update iteration comprises: using the layer inputs for the training inputs in the batch and the estimated target post-activations for the training inputs in the batch, computing, from current weights for the particular neural network layer, gradients of a respective local matching loss with respect to the weights of the particular neural network layer, the respective local matching loss depending on a matching loss of a transfer function for the particular neural network layer, and updating the current weights for the particular neural network layer using the gradients, wherein the operations for each neural network layer of the plurality of neural network layers are assigned to and performed on a respective hardware device, and the operations are performed in parallel for each neural network layer of the plurality of neural network layers. the operations further comprise:

2. The method of claim 1, wherein, for each neural network layer, providing, by the respective hardware device for that neural network layer, the final updated weights for access by hardware devices performing operations for other neural network layers, and obtaining, by the respective hardware device for that neural network layer, the final updated weights for the other neural network layers of the plurality of neural network layers for use in performing forward and backward passes through the neural network. the batch comprises the same training inputs for all of the plurality of neural network layers.

3. The method of claim 1, wherein, the layer inputs and the estimated target post-activations for the particular neural network layer are fixed for each of the plurality of update iterations.

4. The method of claim 1, wherein, determining the estimated target post-activation for the particular neural network layer comprises backpropagating gradients of a final loss between the training output for the training input and the label for the training input.

5. The method of claim 1, wherein, the estimated target post-activation for the particular neural network layer is a mirror-descent target post-activation.

6. The method of claim 1, wherein, computing the gradients of the respective local matching loss with respect to the weights of the particular neural network layer comprises, for each training input in the batch:

7. The method of any one of claims 1-6, wherein, applying the current weights to the layer input for that training input to generate a predicted pre-activation for that training input; applying the transfer function to the predicted pre-activation to generate a predicted post-activation; and determining a difference between the predicted post-activation for that training input and the estimated target post-activation. computing the gradients further comprises, for each training input in the batch:

8. The method of claim 7, wherein, ​ computing a product of a layer input of the training input and the difference determined for the layer input.

9. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform the operations of the method of any one of claims 1-8.

10. A computer-readable storage medium encoded with instructions that when executed by one or more computers cause the one or more computers to perform the operations of the method of any one of claims 1-8.

Images