Training neural networks with explicit learning rate schedules

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training neural network with explicit learning rate schedules. In one aspect, a method includes: obtaining a training dataset for training a neural network parametrized by a set of network parameters; and training the neural network on the training dataset over a number of training iterations, including, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterion is satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated values of the network parameters.

Description

TRAINING NEURAL NETWORKS WITH EXPLICIT LEARNING RATESCHEDULESCROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. Provisional Patent Application No. 63 / 650,896, titled “TRAINING NEURAL NETWORKS WITH EXPLICIT LEARNING RATE SCHEDULES”, filed on May 22, 2024, which is hereby incorporated by reference in its entirety.TECHNICAL FIELD

[0002] This disclosure relates generally to methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training neural networks with explicit learning rate schedules.BACKGROUND

[0003] This disclosure relates to neural networks. Neural networks are machine learning models that employ 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 neural 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 training system implemented as computer programs on one or more computers in one or more locations that can train a neural network with explicit learning rate schedules, e g., by layer-wise normalizing the neural network and periodically normalizing a set of network parameters of the neural network during training. These and other features of the training system described herein are summarized below.

[0005] According to a first aspect, a method performed by one or more computers is provided, the method including: obtaining a training dataset for training a neural network parametrized by a set of network parameters, wherein the neural network includes: one or more nonlinear layers thateach apply a nonlinear transformation over a respective normalized layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer; and training the neural network on the training dataset over a number of training iterations, including, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterion is satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

[0006] In some implementations of the method, at each training iteration, determining whether the criterion is satisfied at the training iteration includes: determining whether a threshold number of training iterations has elapsed since the network parameters were last normalized; or determining whether a norm over the updated values of the network parameters is greater than a threshold norm.

[0007] In some implementations of the method, at each training iteration: the subset of training data includes a network input, and training the neural network on the subset of training data to generate the updated values of the network parameters includes: processing the network input using the neural network, in accordance with the current values of the network parameters, to generate a network output for the network input; computing, with respect to the network parameters, a gradient of an objective function when evaluated on the network output; and augmenting, using the gradient of the objective function, the current values of the network parameters to generate the updated values of the network parameters.

[0008] In some implementations of the method, at each training iteration: the subset of training data further includes a target network output for the network input, and the objective function evaluates an error between: (i) the network output, and (ii) the target network output, for the network input.

[0009] In some implementations of the method, the training dataset was generated using the neural network.

[0010] In some implementations of the method, the training dataset is a data stream including a respective subset of training data at each of a number of timesteps, and each training iteration is a respective timestep in the data stream.

[0011] In some implementations of the method, the neural network includes one or more additional layers each parametrized by a respective subset of the network parameters, and at each training iteration, normalizing the updated values of each of the network parameters includes, for each additional layer: computing a norm over the updated values of the network parameters in the respective subset; and for each network parameter in the subset: normalizing the updated value of the network parameter with the norm to generate a normalized version of the updated value; and scaling the normalized version of the updated value with a respective reference norm for the additional layer.

[0012] In some implementations of the method, the method further includes, at a first of the training iterations: computing, as the respective reference norm for each additional layer, a norm over the current values of the network parameters in the respective subset.

[0013] In some implementations of the method, each normalization layer is parametrized by a respective subset of the network parameters, and the subset includes: (i) a respective scaling vector, and (ii) a corresponding offset vector.

[0014] In some implementations of the method, for each normalization layer, applying the normalization operation over the respective layer input to the normalization layer includes: computing a mean over a number of elements of the layer input; shifting each of the elements of the layer input by the mean to generate a shifted version of the layer input; computing a norm of the shifted version of the layer input; normalizing the shifted version of the layer input with the norm to generate a normalized version of the layer input; computing an elementwise product between: (i) the respective scaling vector, and (ii) the normalized version of the layer input; and offsetting the elementwise product with the respective offset parameter.

[0015] In some implementations of the method, at each training iteration, normalizing the updated values of the network parameters includes, for each normalization layer: computing a norm over the updated values of the scaling and offset vectors in the respective subset; normalizing the updated values of the scaling and offset vectors with the norm to generate normalized versions of the updated values; and scaling the normalized versions of the updated values with a respective reference norm for the normalization layer.

[0016] In some implementations of the method, each normalization layer is a layer normalization layer or a batch normalization layer.

[0017] In some implementations of the method, each nonlinear layer is a Rectified Linear Unit (ReLU), a Gaussian Error Linear Unit (GELU), a hyperbolic tangent activation, a sigmoid activation, or a softmax activation.

[0018] In some implementations of the method, the method further includes: obtaining the neural network including the one or more nonlinear layers; and for each of the one or more nonlinear layers: determining whether a normalization layer immediately precedes the nonlinear layer; and when a normalization layer does not immediately precede the nonlinear layer, inserting the respective normalization layer into the neural network.

[0019] In some implementations of the method, the neural network is configured to process a network input including one or more of: an image, an audio waveform, or a text sequence.

[0020] In some implementations of the method, the neural network is configured to process a network input including an observation characterizing a state of an environment being interacted with by an agent, and generate a network output defining an action to be performed by the agent in response to the observation.

[0021] In some implementations of the method, the method further includes, after training the neural network on the training dataset over the training iterations: parametrizing the neural network with the trained values of the network parameters; receiving a network input; and processing the network input using the neural network, in accordance with the trained values of the network parameters, to generate a network output for the network input.

[0022] According to a second aspect, a method performed by one or more computers is provided, the method including: receiving a network input; and processing the network input, using a neural network, to generate a network output for the network input, where the neural network has been trained via the method of the first aspect in any of its abovementioned implementations.

[0023] According to a third aspect, a method performed by one or more computers is provided, the method including: receiving a network input; and processing the network input, using a neural network parametrized by a set of network parameters, to generate a network output for the network input, wherein the neural network includes: one or more nonlinear layers that each apply a nonlinear transformation over a respective layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer, and wherein the neural network has been trained on a training dataset over a number of trainingiterations, including, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterion is satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

[0024] According to a fourth aspect, a system including one or more non-transitory computer storage media is provided, the one or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the method of any one of the first, second, or third aspects in any one of its abovementioned implementations.

[0025] According to a fifth aspect, a system including one or more computers and one or more storage devices communicatively coupled to the one or more computers is provided, the one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the method of any one of the first, second, or third aspects in any one of its abovementioned implementations.

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

[0027] Normalization layers have recently regained popularity in deep reinforcement learning and continual learning due to several benefits during training, including improved loss landscape conditioning and mitigation of overestimation bias. However, applying normalization to a neural network at the layer level can introduce a subtle but consequential side effect - an implicit coupling between the growth in the norm of the neural network’s parameters and a decay in the effective learning rate. In continual learning settings, this dynamic can be particularly harmful. As training progresses, the effective learning rate of the neural network may decrease too rapidly, far outpacing the natural timescale of learning involved for a given task. This premature learning rate decay often leads to undertraining, diminished model adaptability, and stagnation in performance, both during and after training. Moreover, it results in inefficient use of computational resources, as the neural network continues training with vanishingly small updates that contribute little to actual learning of the given task.

[0028] To address the problem of implicit and uncontrolled decay in effective learning rates, this specification introduces a training algorithm referred to herein as “Normalize-and-Project” (or “NaP”), which explicitly stabilizes the learning rate schedule of a neural network through a reparameterization strategy. NaP functions by coupling the insertion of normalization layers into the neural network with periodic projections of the neural network’s weights, ensuring that the effective learning rate of the neural network remains constant over the course of training, even as the norm of the neural network’s parameters evolves.

[0029] This approach yields several important technical benefits. First, by maintaining a consistent effective learning rate, NaP prevents the premature stagnation of a neural network during training, a phenomenon that often occurs in reinforcement learning, continual learning, and nonstationary environments, thereby preserving learning capacity of the neural network over long time horizons. Second, NaP enhances robustness of the neural network to distributional drift and nonstationarity, as demonstrated on synthetic plasticity loss benchmarks and both single-task and sequential variants of the Arcade Learning Environment (“ALE”). Third, NaP generalizes across a variety of neural network architectures, e.g., including residual networks (“ResNets”) and transformers, enabling broad applicability of NaP with little or no changes to the neural network architecture. Additionally, NaP can recover and improve the baseline performance of a neural network in standard stationary benchmarks, highlighting its effectiveness as both a diagnostic tool for learning dynamics and a practical training enhancement.

[0030] 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 THE DRAWINGS

[0031] FIG. 1 A is a schematic diagram depicting an example of a normalized nonlinear unit.

[0032] FIG. IB is a schematic diagram depicting an example of a neural network with layer-wise normalization.

[0033] FIGs. 2A-2B are schematic diagrams depicting another example of a neural network with layer-wise normalization.

[0034] FIGs. 3A-3C are a schematic diagram depicting examples of a training system configured to train a neural network to perform a machine learning task via a “Normalize-and-Project” (or “NaP”) training algorithm.

[0035] FIG. 4 is a flow diagram of an example process for training a neural network to perform a machine learning task via a NaP training algorithm.

[0036] FIGs. 5A-5C are experimental plots depicting continual random-labels CIFAR training of a convolutional neural network (“CNN”), a CNN with layer-wise normalization, and a CNN with layer-wise normalization and weight projection.

[0037] FIGs. 6A-6D are experimental plots depicting results of a coupled networks experiment.

[0038] FIGs. 7A-7D are experiment plots depicting random label memorization robustness for a CNN and a multi-layer perceptron (“MLP”) trained with and without NaP.

[0039] FIGs. 8A-8J are experimental plots depicting learning curves of continual Atari agents.

[0040] FIGs. 9A-9B are experimental plots depicting aggerate effects of normalization on singletask Atari.

[0041] Like reference numbers and designations in the various drawings indicate like elements.DETAILED DESCRIPTIONI. Introduction

[0042] Many promising applications of deep learning, such as reinforcement learning (“RL”) and continual learning (“CL”), involve training a neural network on a problem which is in some way nonstationary. For this type of training to be effective, the neural network should maintain adaptability to new information as it becomes available. That is, the neural network should remain “plastic”. Loss of plasticity can present a barrier to performance improvement in deep learning, often resulting from accumulation of saturated activations and increased sharpness of the loss landscape. Some mitigation strategies involve resetting saturated activations and regularizing the neural network’s parameters towards their initial values. Such mitigation strategies can be realized by proactively reducing statistical drift of the pre-activations. In the supervised learning setting, this problem can be resolved by incorporating normalization layers into the neural network architecture. Indeed, as demonstrated in this specification, layer-wise normalization can be effective at stabilizing optimization in both RL and CL.

[0043] However, while effective, layer-wise normalization on its own may be insufficient to avoid plasticity loss in some neural network architectures. Part of the reason lies in a subtle property of layer-wise normalization - a normalization layer can cause the subnetwork preceding it to become scale-invariant, meaning the layer’s effective learning rate (“ELR”) can depend on the norm of its parameters. Thus, when the norm of the network parameters grows, as it typically does in neural networks trained without regularization, the ELR shrinks.

[0044] For reference, the (nominal) learning rate is a hyperparameter in neural network training that determines the step size used by an optimizer (e.g., stochastic gradient descent, Adam, etc.) to update the neural network’s parameters in the gradient direction of an objective function. If the learning rate is too high, the updates can be unstable and cause divergence. If the learning rate is too low, the convergence is slow, with a risk of the optimizer trapping itself in suboptimal minima. A learning rate schedule refers to a predefined or adaptive plan for adjusting the learning rate over the course of training. Generally, the goal of the learning rate schedule is to improve convergence speed and final model performance. The effective learning rate is the actual magnitude of the network parameter updates during training, which may differ from the nominal learning rate due to various architectural or algorithmic features. In such cases, even if the nominal learning rate remains fixed, the actual update size can vary significantly, which in turn, can cause significant variation in the learning dynamics. This can be particularly important in settings like RL and CL, where long-term adaptability is desirable.

[0045] To address some, or all, of these abovementioned challenges, this specification introduces a training system implementing a training algorithm referred to as “Normalize-and-Project” (or “NaP”), which ensures that the learning rate of an optimizer reflects the ELR of the neural network. Broadly, application of NaP to training of a neural network can be decomposed into two operations: (i) the insertion of normalization layers prior to nonlinearities in the neural network, and (ii) the periodic projection of the neural network’s weights onto a fixed-norm radius throughout training. NaP avoids implicit learning rate decay, allowing the training system to train a neural network to state-of-the-art performance in a variety of nonstationary learning problems, e.g., even without incorporating regularization intended to prevent plasticity loss.

[0046] In general, the training system implements layer-wise normalization to improve a neural network’s plasticity during training. Particularly, when coupled with adaptive optimizers, e.g., Adam or RMSProp, whose step size is independent of the gradient norm, saturated activations canstill receive sufficient gradient signal to induce non-trivial changes in their pre-activation parameters, providing the opportunity for them to recover as a natural byproduct of optimization. Moreover, layer-wise normalization can combat the implicit learning rate decay induced by paramet er- wise norm growth in RL settings. For example, some components of an RL value function can involve learning a sufficiently small ELR. Thus, an optimization process which does not reach this ELR can underfit the value function in ways that inhibit performance improvement, which layer-wise normalization can mitigate.

[0047] An empirical evaluation of NaP was performed in several experiments. The results of which demonstrated that NaP can be applied to a variety of neural network architectures and training datasets for improving performance of the model on a machine learning task. Examples include 400M transformer models trained on the C4 dataset and vision models trained on CIFAR- 10 and ImageNet. Results of studies on single-task and sequential variants of the Arcade Learning Environment (“ALE”) are also provided herein, where NaP demonstrated robustness to task changes and outperformed a baseline Rainbow agent with initialized parameters after four hundred million (“400M”) training iterations and one hundred million (“100M”) optimizer steps.

[0048] These and other features relating to the neural networks, training system, and NaP training algorithm are described in more detail below.II. Background

[0049] Sec. II provides a background on trainability of neural networks, learning dynamics, and effective learning rates in the context of parameter-wise normalization. Note, to avoid confusion between different forms of normalization, “layer-wise normalization” refers to normalizing the pre-activation of a nonlinear layer (e.g., via the insertion of a normalization layer), while “parameter- wise normalization” refers to normalizing the parameters of a layer itself.

[0050] Maintaining trainability not only at initialization, but also over the course of optimization, is relevant in reinforcement and continual learning. Failure to do so is referred to as “loss of plasticity” or “plasticity loss”. This phenomenon presents a limiting factor to performance in several reinforcement learning tasks, along with continual learning and warm-starting neural network training. Plasticity loss can be further decomposed into two distinct components: (i) loss of trainability, and (ii) reduced generalization performance.

[0051] Parameter-wise normalization introduces scale-invariance into the neural network layers to which it is applied, whereby a scale-invariant function satisfies f c9') = (0) for any positive scalar c > 0. This leads to the gradient scaling inversely with the parameter- wise norm,— c-1Vgf (0). Intuitively, this property implies that changing the direction of a large vector involves a greater perturbation than changing the direction of a small vector. This motivates the concept of an “effective learning rate” (or “ELR”), which provides a scale-invariant notion of optimizer step size. Note, in the following definition of the ELR, an implicit “reference norm” of size one is assumed.

[0052] Consider a scale-invariant function ( ) parametrized by a set of parameters (0), and an update function 0t+1«- 0f+where p is the (nominal) learning rate, and t indexes a training iteration. Let pl= || 0f||-1be the inverse of the parameter- wise norm where || ... || denotes the Euclidean norm. The effective learning rate (?t) of the scale-invariant function can be defined as follows:

[0053] Then, letting 0 = 0 / ||01| be the normalized set of parameters, the above definition of the ELR implies that:

[0054] which suggests that regularization of the parameter- wise norm can have the dual effect of increasing the effective learning rate. For example, scale-invariant neural networks can be trained with network parameters constrained to the unit sphere, a training regime elaborated upon herein.III. Examples of Layer-wise Normalized Neural Networks

[0055] Sec. Ill provides examples of how a neural network, particularly the nonlinear layers of the neural network, can be layer-wise normalized to maintain plasticity by facilitating the recovery of saturated activations. This section also highlights the effects of controlling the parameter-wise norm of a lay er- wise normalized neural network. These insights are leveraged in Sec. IV to introduce the “Normalize-and-Project” training algorithm, which is a training algorithm that maintains the statistics of a layer-wise normalized neural network and gradients throughout training by periodically normalizing the network parameters of the neural network.

[0056] FIG. 1 A is a schematic diagram depicting an example of a normalized nonlinear unit 50. The normalized nonlinear unit 50 is an example of a nonlinear unit with layer-wise normalization, e.g., a layer-wise normalized nonlinear unit. The normalized nonlinear unit 50 is an example of a module of a system, e.g., a neural network or a system including the neural network, implemented as computer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0057] The normalized nonlinear unit 50 is configured to: receive a block input (x) 22; and process the block input 22 to generate a block output (y) 24 for the block input 22. Note, the block output 24 may also be referred to as the “activation” of the normalized nonlinear unit 50. The operations of the normalized nonlinear unit 50 can be represented concisely as y =where feis a function representing the parametric model of the normalized nonlinear unit 50, parametrized by a set of block parameters (0). The normalized nonlinear unit 50 is configured to process the block input 22, in accordance with a respective current value of each block parameter in its set of block parameters, to generate the block output 24. In other words, the block output 24 depends on the block input 22 and the current values of the block parameters.

[0058] In general, the block input 22 and the block output 24 can each be a respective tensor, such as a zeroth-order tensor (i.e., a scalar), a first-order tensor (i.e., a vector), a second-order tensor (i.e., a matrix), a third-order tensor, or a higher-order tensor of more than three orders. For example, the block input 22 and the block output 24 can each be an embedding vector of a token, a matrix including a sequence of token embedding vectors representing a text piece, a third-order tensor representing an image, a fourth-order tensor representing a video or a batch of images, a fifth-order tensor representing a batch of videos, and so on.

[0059] The normalized nonlinear unit 50 is a sequence of layers including a linear layer 52, a normalization layer 54, and a nonlinear layer 56. Each of the linear layer 52, normalization layer 54, and nonlinear layer 56 is configured to receive: a respective layer input; and process the respective layer input to generate a respective layer output for the respective layer input.

[0060] The linear layer 52 is a first layer in the normalized nonlinear unit 50. The linear layer 52 is configured to: receive the block input 22 as its layer input; and apply a linear transformation over the block input 22 to generate, as its layer output, a pre-activation ( / i) 32 of the normalized nonlinear unit 50. The application of the linear transformation over the block input 22 can be expressed concisely as:h = Wx, (la)

[0061] where W is a weight tensor of the linear layer 52.

[0062] In general, the linear layer 52 can be any type of linear neural network layer, and the weight tensor can be any tensor representation of the linear layer 52, e.g., imposed by one or more symmetries and / or constraints on the weight tensor. For example, the linear layer 52 can be a linear fully-connected layer, a linear convolutional layer, a linear projection layer, a linear embedding layer, an identity mapping layer, or other linear neural network layer. Correspondingly, the weight tensor can be a dense weight tensor, a convolutional weight tensor, a projection weight tensor, an embedding weight tensor, an identity tensor, or other weight tensor having one or more symmetries and / or constraints.

[0063] Note, the normalization layer 54 typically makes a bias term in Eq. (la) redundant, but one can be implemented if desired. For example, more generally, the linear layer 52 can be configured to: receive the block input 22 as its layer input; and apply an affine transformation over the block input 22 to generate the pre-activation 32 as its layer output.

[0064] The normalization layer 54 follows (e.g., immediately follows) the linear layer 52 in the normalized nonlinear unit 50. The normalization layer 54 is configured to: receive the preactivation 32 as its layer input; and apply a normalization operation over the pre-activation 32 to generate, as its layer output, an updated pre-activation (h) 34 of the normalized nonlinear unit 50. The application of the normalization operation over the pre-activation 32 can be expressed concisely as:

[0065] where J\f is a normalization operator of the normalization layer 54, y is a scaling tensor of the normalization layer 54, and ? is a shift tensor of the normalization layer 54. For example, the normalization layer 54 can be a layer normalization (“LayerNorm”) layer, a batch normalization (“BatchNorm”) layer, a weight normalization (“WeightNorm”) layer, a root mean square normalization (“RMSNorm”) layer, a group normalization (“GroupNorm”) layer, or other normalization layer.

[0066] Thus, to apply the normalization operation over the pre-activation 32, the normalization layer 54 can be configured to: apply the normalization operator to the pre-activation 32 to generate a normalized pre-activation ( ) of the normalized nonlinear unit 50; and apply an affine transformation over the normalized pre-activation to generate the updated pre-activation 34 of thenormalized nonlinear unit 50. The application of the normalization operator to the pre-activation 32 can be expressed as h — JV'(h), and the application of the affine transformation over the normalized pre-activation can be expressed as h = yh -I- (3. In general, the normalization operator can implement any normalization scheme, such as a statistical normalization, a Euclidean normalization, or an Lp normalization. For example, the normalization operator can perform a whitening transformation on the pre-activation 32, such that the normalized pre-activation has zero mean and unit variance.

[0067] In some implementations, the pre-activation 34 is a sequence of pre-activated token embeddings, and the normalization layer 54 does not apply the normalization operator along a token axis of the sequence. In other words, the normalization operator is applied element-wise to each pre-activated token embedding in the pre-activation 34. For example, if the normalization layer 54 is a LayerNorm layer, the normalization operator can be applied to each pre-activated token embedding q E h in the pre-activation 34 to generate a respective normalized pre-activated token embedding q 6 h in the normalized pre-activation as:

[0068] where q = (qq, q2, ... , qH) is the pre-activated token embedding, q = (qq, q2, ... , qH) is the normalized pre-activated token embedding, H is the number of features in the pre-activated token embedding, / z is the mean over each feature (qq) in the pre-activated token embedding, cr is the standard deviation of each of the features in the pre-activated token embedding, and e > 0 is a constant having a positive value to avoid division by zero, e.g., a small positive value of at most about IO’6, IO’7, IO’8, IO’9, IO’10, 10’11, IO’12, or less.

[0069] The nonlinear layer 56 follows (e.g., immediately follows) the normalization layer 52 in the normalized nonlinear unit 50. The nonlinear layer 56 is configured to: receive the updated pre- activation 34 as its layer input; and apply a nonlinear transformation over the updated pre- activation 34 to generate, as its layer output, the block output 24 of the normalized nonlinear unit 50. The application of the nonlinear transformation over the updated pre-activation 34 can be expressed concisely as:

[0070] where < > is a nonlinear function of the nonlinear layer 56 that implements the nonlinear transformation. For example, the nonlinear function of the nonlinear layer 56 can be an activation function, such as a rectified linear unit, a hyperbolic tangent activation, a sigmoid activation, a leaky rectified linear unit, a Gaussian error linear unit, a swish activation, a mish activation, a softmax activation, a sparsemax activation, an entmax activation, and so on. As another example, the nonlinear function of the nonlinear layer 56 can be an attention mechanism, such as a dotproduct attention mechanism, an additive attention mechanism, a cosine similarity attention mechanism, a soft attention mechanism, a hard attention mechanism, a global attention mechanism, a local attention mechanism, and so on.

[0071] In general, the nonlinear layer 56 has no trainable parameters. Thus, the set of block parameters 6 = {W, Y, / 3} parametrizing the normalized nonlinear unit 50 includes the weight tensor (l / F) of the linear layer 52, the scaling tensor (y) of the normalization layer 54, and the shift tensor ( / ?) of the normalization layer 54.

[0072] In some implementations, the normalized nonlinear unit 50 further includes a BatchNorm layer between the normalization 54 and nonlinear 56 layers, and the normalization layer 54 is a normalization layer of a different type than the BatchNorm layer. For example, if the BatchNorm layer is present, the normalization layer 54 can be a LayerNorm layer, a WeightNorm layer, a RMSNorm layer, or a GroupNorm layer. The BatchNorm layer proceeds (e.g., immediately proceeds) the normalization layer 54 and precedes (e.g., immediately precedes) the nonlinear layer 56 in the normalized nonlinear unit 50. In some of these implementations, the shift tensor ( / ?) of the normalization layer 54 is omitted if the BatchNorm layer is present.

[0073] In some implementations, the normalized nonlinear unit 50 represents a normalized version of a neural network layer, e.g., a layer-wise normalized neural network layer, where the neural network layer implements the linear transformation of the linear layer 52, followed by the nonlinear transformation of the nonlinear layer 56. That is, the neural network layer implements a neuron of the form y = (f>(Wx). Here, the normalized nonlinear unit 50 can be generated by inserting the normalization layer 54 between the linear 52 and nonlinear 56 layers.

[0074] As one example, the normalized nonlinear unit 50 can represent a normalized version of a fully-connected layer, e.g., a fully-connected layer with layer-wise normalization. In this case, the weight tensor of the linear layer 52 can be a dense weight tensor (e.g., having no non-zero entries),and the nonlinear function of the nonlinear layer 56 can be a rectified linear unit or a Gaussian error linear unit.

[0075] As another example, the normalized nonlinear unit 50 can represent a normalized version of a convolutional layer, e.g., a layer-wise normalized convolutional layer. In this case, the weight tensor of the linear layer 52 can be a convolutional weight tensor (e g., having translational equivariance symmetry), and the nonlinear function of the nonlinear layer 56 can be a rectified linear unit, a leaky error linear unit, a swish activation, or a mish activation.

[0076] As yet another example, the normalized nonlinear unit 50 can represent a normalized version of an attention layer, e.g., a layer-wise normalized attention layer. In this case, the weight tensor of the linear layer 52 can be a projection weight tensor (e.g., including a query, key, and / or value projection tensor), and the nonlinear function of the nonlinear layer 56 is an attention mechanism.

[0077] In this case, the weight tensor of the linear layer 52 can be represented in block tensor form as:W = [IVQWKlVv], (2a)

[0078] where IVQ is a query weight tensor of the attention layer, WKis a key weight tensor of the attention layer, and Wy is a value weight tensor of the attention layer. The pre-activation 32 can be represented in block tensor form as: h = [Q K 7] = [WZQWKX Wyx], (2b)

[0079] which includes the query (Q), the key (A), and the value (V) of the attention layer upon projection of the block input 22. With zero shift (3 = 0, the updated pre-activation 34 can be represented in block tensor form as:

[0080] which includes an updated query (Q) and an updated key (A). Note, the normalization layer 54 is not applied to the value (V) of the attention layer, e.g., as this may lead to leakage of information during training on next-token-prediction objectives. Here, y = [yQyK] is the scaling tensor of the normalization layer 52 in block tensor form, where yQis a query scaling matrix for the query, and yKis a key scaling matrix for the key. In some implementations, one or both of the query and key scaling matrices can be identity matrices. The block output 24 can then be represented as:(h) = Attention((?, K, V) = softmax

[0081] where softmax is the softmax activation, and dkis the dimension of the query and key.

[0082] As yet another example, the normalized nonlinear unit 50 can represent a normalized version of an activation layer, e.g., a layer-wise normalized activation layer. That is, the activation layer implements a neuron of the form y = 0(x). For example, the normalized nonlinear unit 50 may follow a skip connection or may be implemented as a gate in a gated layer block, e.g., a gated multi-layer perceptron. In this case, the linear layer 52 can be an identity mapping layer, e.g., such that the weight tensor of the linear layer 52 is an identity tensor, and the nonlinear function of the nonlinear layer 56 can be an activation function of the activation layer.

[0083] Note, achieving approximately mean-zero, unit-variance of the updated pre-activation 34 can be useful to ensure a neural network is trainable at initialization, and some neural network initialization schemes aim to maintain trainability as the depth of the neural network grows. For example, large deviations of the pre-activation statistics from their initial values can lead to network pathologies such as saturated activations and low numerical rank of the Neural Tangent Kernel (“NTK”). Layer-wise normalization can not only guarantee that pre-activations are meanzero, unit-variance at initialization, but also maintain these statistics over the course of training, even if the data distribution changes. This property improves robustness to network pathologies that cause loss of plasticity, as well as conditioning of a neural network’s gradients in reinforcement learning.

[0084] Beyond re-normalizing the pre-activation statistics, the normalization layer 54 can also introduce a dependency between the linear 52 and nonlinear 56 layers. This mixing step allows gradients to propagate through the pre-activation 32 even if the nonlinear layer 56 is saturated. In other words, the normalization layer 54 can give dead activation functions a second chance at life - rather than immediately decaying to zero, the gradients flowing through a dead activation function will instead take small, non-zero values.

[0085] FIG. IB is a schematic diagram depicting an example of a neural network 100 with layerwise normalization, e.g., a lay er- wise normalized neural network. The neural network 100 is 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 are implemented.

[0086] The neural network 100 includes a normalized network 150. The normalized network 150 includes a sequence of normalized nonlinear units 50.1 to 50. L each configured in accordance with FIG. 1A. Here, I = 1,2, ... , L indexes a normalized nonlinear unit 50.1, and L denotes the total number of normalized nonlinear units 50 in the normalized network 150, i.e., the depth of the normalized network 150. For example, the normalized network 150 can include 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 50, 100, 150, 200, 250, 500, 1000, or more normalized nonlinear units 50. The normalized network 150 can include any of multiple types of normalized nonlinear units 50. For example, the normalized nonlinear units 50.1 to 50.L can include normalized versions of one or more fully-connected layers, one or more convolutional layers, one or more attention layers, and / or one or more activation layers.

[0087] Each normalized nonlinear unit 50.1 is configured to: receive a respective block input (x;) 22.1; and process the block input 22.1 to generate a respective block output (yt) 24.1 for the block input 22.1. The operations of the normalized nonlinear unit 50.1 can be represented concisely as yt= where is the respective function representing the parametric model of the normalized nonlinear unit 50.1, parametrized by its respective set of block parameters (0 . For each normalized nonlinear unit 50.1 preceding the last normalized nonlinear unit 50.L, the block output 24.1 generated by the normalized nonlinear unit 50.1 is the block input 22.1 to the next normalized nonlinear unit 50. (1+1) in the normalized network 150, that is, xt+1= yt.

[0088] In this example, the neural network 100 includes one normalized network 150. However, more generally, the neural network 100 can include any number of normalized networks 150 corresponding to different subnetworks and / or layer blocks of the neural network 100, e.g., interconnected to one another by one or more skip connections, as shown in FIGs. 2A-2B.

[0089] The neural network 100 is configured to: receive a network input (X) 102; and process the network input 102 to generate a network output (F) 104 for the network input 102. The operations of the neural network 100 can be represented concisely as ais a function representing the parametric model of the neural network 100, parametrized by a set of network parameters (0). The network parameters (0) include the respective set of block parameters of each normalized nonlinear unit 50.1 of the neural network 100. That is, the respective set of block parameters (0;) of each normalized nonlinear unit 50.1 is a respective subsetc 0 of the set of network parameters of the neural network 100. The neural network 100 is configured to process the network input 102, in accordance with a respective current value of each network parameter inits set of network parameters, to generate the network output 104. In other words, the network output 104 depends on the network input 102 and the current value of each network parameter.

[0090] The neural network 100 is configured to perform a machine learning task on the network input 102 to generate the network output 104, where the neural network 100 has been trained to perform the machine learning task via a “Normalize-and-Project” (or “NaP”) training algorithm. Implementations of the neural network 100 can be configured through the NaP training algorithm to perform any kind of machine learning task. That is, a training system 200 implementing the NaP training algorithm can train the neural network 100 to receive any kind of network input 102 and process the network input 102 to generate any kind of network output 104, e.g., an action, a policy, a score, a classification, or a regression output, based on the network input 102. Examples of the training system 200 and NaP training algorithm are described in more detail in Section IV. Examples of machine learning tasks that the training system 200 can train the neural network 100 to perform are described in more detail in Sec. V.

[0091] The neural network 100 can be implemented in any appropriate location, e.g., on a user device (e.g., a mobile device), or on one or more computers in a data center, etc. In some implementations, users can interact with the neural network 100, e.g., by providing a query by way of an interface, e.g., a graphical user interface, or an application programming interface (“API”). For example, a user can provide a query to the neural network 100, where the query includes: (i) a request to process a network input 102, and (ii) the network input 102. Upon receiving the query, the neural network 100 can process the network input 102, responsive to the request, and provide the network output 104 resulting from the machine learning task performed on the network input 102 to the user, e.g., for implementation on a user device of the user, or for storage in a data storage device. In some cases, the neural network 100 can transmit the network output 104 to a user device of the user, e g., by way of a data communication network (e.g., the Internet).

[0092] Note, while the output of a scale-invariant function is insensitive to scalar multiplication of its parameters, the gradient magnitude scales inversely with the paramet er- wise norm. For the neural network 100 with layer-wise normalization, growth in the parameter-wise norm ||0|| corresponds to a decline in the neural network 100’s sensitivity to changes in its network parameters. In some cases, this may be preferable, as the slow but stable regime of vanishing gradients can be easier to recover from than the unstable exploding gradient regime. However, if the paramet er- wise norm grows indefinitely, then the corresponding reduction in the effectivelearning rate can eventually cause noticeable slowdowns in learning. Instead, if the network parameters are re-normalized periodically so they recover the parameter- wise norm at initialization, the neural network 100 can maintain performance of a random initialization even after hundreds of training iterations.

[0093] FIG. 2A is a schematic diagram depicting another example of the neural network 100 with layer-wise normalization. Here, the neural network 100 includes an input subnetwork 110, a residual network 120, an output subnetwork 130, a skip connection 140, and an additive gate 142. For example, the neural network 100 can be a Transformer, or a Transformer-based model such as a Vision Transformer (“ViT”). As another example, the neural network 100 can be a ResNet, or a Transformer-ResNet hybrid model.

[0094] The input subnetwork 110 is a normalized network 150 configured in accordance with FIG. IB. The input subnetwork 110 is configured to: receive the network input (A) 102; and process the network input 102 to generate an input embedding (z) representing the network input 102. In some implementations, the input subnetwork 110 is an embedding network. For example, the input subnetwork 110 can be a linear embedding network, a lookup table-based embedding network, a subword or byte-level embedding network, a sequence-to-sequence embedding network, a Transformer-based embedding network, a convolutional-based embedding network, a recurrentbased embedding network, a long short-term memory (“LSTM”)-based embedding network, a graph-based embedding network, or a hybridization thereof.

[0095] The residual network 120 follows the input subnetwork 110. The residual network 120 may be the core or backbone of the neural network 100. The residual network 120 is configured to: receive the input embedding (z) representing the network input 102; and process the input embedding to generate a residual embedding (z) for the network input 102, e.g., an embedding representing the network output 104.

[0096] The additive gate 142 follows the residual network 120. The additive gate 142 is connected to the output of the residual network 120 and the input of the residual network 120 via the skip connection 140. The additive gate 142 is configured to combine the input (z) and residual (z) embeddings to generate an output embedding (z) for the network input 102, e.g., an updated embedding representing the network output 104. For example, the output embedding z = z + z can be a summation or linear combination of the input and residual embeddings.

[0097] The output subnetwork 130 follows the additive gate 142. The output subnetwork 130 is a normalized network 150 configured in accordance with FIG. IB. The output subnetwork 130 is connected to the output of the additive gate 142. The output subnetwork 130 is configured to: receive the output embedding (z); and process the output embedding to generate the network output (K) 104 for the network input 102. In some implementations, the output subnetwork 130 is a projection network. For example, the output subnetwork 130 can be a softmax-based projection network, a fully-connected projection network, an attention-based projection network, a contrastive or distance-based projection network, a decoder-based projection network, a graphbased projection network, or a hybridization thereof.

[0098] FIG. 2B is a schematic diagram depicting an example of the residual network 120 including a sequence of residual layer blocks 125.1 to 125. R. Each residual layer block 125. r is a normalized network 150 configured in accordance with FIG. IB. The sequence of residual layer blocks 125.1 to 125. R are arranged into a residual configuration, where r = 1,2, ..., R indexes a residual layer block 125. r, and R denotes the total number of residual layer blocks, i.e., the depth of the residual network 120. The residual network 120 includes a respective skip connection 140. r and a respective additive gate 142. r for each residual layer block 125. r. The additive gate 142. r follows the residual layer block 125. r. The additive gate 142. r is connected to the output of the residual layer block 125. r and the input of the residual layer block 125. r via the skip connection 140. r. Generally, a residual configuration enables the neural network 100 to implement deep models, e.g., including 10, 100, 1,000, 10,000, or more normalized nonlinear units 50, that are easier to train and approach better accuracy with increasing numbers of the normalized nonlinear units 50. For example, the residual network 120 can include at least about 2, 5, 10, 20, 50, 100, 200, 500, 1,000, or more residual layer blocks.

[0099] Each residual layer block 125. r is configured to: receive a respective block input (^r); and process the respective block input to generate a respective block output ( r). The additive gate 142. r is configured to combine the block input (xr) and block output (4 >r) of the residual layer block HO.r to generate the block input (xr+i) for the next residual lay block 125. (r+1) in the residual network 120. For example, the block input xr+1for the next residual lay block 125. (r+1) can be a summation of the block input and block output of the residual layer block 125. r. The block input = z for the first residual layer block 125.1 in the residual network 120 is the input embedding to the residual network 120. The residual embedding xR+ yR= z generated bythe residual network 120 is the combination of the block input (xR) and block output (-^> ) of the last residual layer block 125.R in the residual network 120.IV. Examples of Training Systems and Normalize-and-Project Training Algorithms

[0100] Sec. Ill described how layer-wise normalizing the neural network 100’s pre-activations and fixing the parameter- wise norm can mitigate loss of plasticity of the neural network 100. Sec. IV outlines a training algorithm involving a combination of layer-wise and parameter-wise normalization, referred to herein as “Normalize-and-Project” (or “NaP”). The NaP training algorithm can be applied to a variety of neural network architectures to improve the stability of training, including those involving nonstationary problems, e.g., reinforcement and continual learning. Broadly, the NaP training algorithm includes two operations: (i) the insertion of normalization layers 54 prior to nonlinear layers 46 in a neural network 100, and (ii) the periodic projection of the neural network 100’s weight tensors onto a fixed-norm radius throughout training. In some implementations, each periodic projection of the weight tensors is accompanied by a respective projection of the scaling and shift tensors of the normalization layers 54 themselves, e.g., corresponding to an update to the per-layer effective learning rates during the training process.

[0101] FIG. 3 is schematic diagram depicting an example of the training system 200. The training system 200 is configured to train a neural network 100 to perform a machine learning task via the NaP training algorithm. The training system 200 is 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 are implemented.

[0102] The training system 200 receives the neural network (J7) 100 and obtains a training dataset (D) 210 for training the neural network 100 to perform the machine learning task.

[0103] The neural network 100 includes one or more nonlinear layers 56.1 to 56.L and a respective linear layer 52.1 preceding each nonlinear layer 56.1. For example, if the linear layer 52.1 immediately precedes the nonlinear layer 56.1, the linear layer 52.1 and nonlinear layer 56.1 can implement a neuron of the form{(lV{x), where VFj is the weight tensor of the linear layer 52.1 and (pi is the nonlinear function of the nonlinear layer 56.1. Here, I = 1,2, ..., L indexes a nonlinear layer 56.1, and L denotes the total number of nonlinear layers 56 in the neural network 100. For example, the neural network 100 can include 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 50, 100, 150, 200, 250, 500, 1,000, or more nonlinear layers 56.

[0104] The training dataset 210, given as D = includes a number of samples (dq) 214.1to 214. Q. Here, q indexes each sample 214. q in the training dataset 210, and Q is the total number of samples 214 in the training dataset 210, e.g., corresponding to the data size of the training dataset 210. For example, the training dataset 210 can include at least about 100, 1,000, 10,000, 100,000, 1,000,000, or more samples 214.

[0105] In general, for each nonlinear layer 56.1 in the neural network 100, the training system 200 identifies a respective normalization linear unit 54.1 that the nonlinear layer 56.1 belongs to or generates the normalization linear unit 54.1 by layer-wise normalizing the nonlinear layer 56.1. A few edge cases are also described below.

[0106] The training system 200 determines whether a normalization layer 54 immediately precedes the nonlinear layer 56.1.

[0107] In some implementations, if a normalization layer 54 does not immediately precede the nonlinear layer 56.1, the training system 200 inserts a respective normalization layer 54.1 immediately before the nonlinear layer 56.1 to generate the respective normalized nonlinear unit 50.1 in the neural network 100 including the respective linear layer 52.1, the respective normalization layer 54.1, and the nonlinear layer 56.1.

[0108] In some implementations, if a normalization layer 54 does not immediately precede the nonlinear layer 56.1, the training system 200 determines whether a different nonlinear layer 56 immediately precedes the nonlinear layer 56.1, e.g., the respective linear layer 52.1 is absent or equivalent to an identity transformation. In this implementation, no trainable parameters are between the nonlinear layer 56.1 and the different nonlinear layer 56, and the training system 200 may determine to not insert a normalization layer 54 immediately before the nonlinear layer 56.1. This situation can sometimes occur, for example, in ResNet configurations of the neural network 100. Here, the training system 200 inserts a normalization layer 54 each time a linear layer 52 immediately precedes a nonlinear layer 56 in the neural network 100.

[0109] In some implementations, if a normalization layer 54 does immediately precede the nonlinear layer 56.1, the training system 200 determines whether the normalization layer 54 is a BatchNorm layer.

[0110] If the normalization layer 54 is a BatchNorm layer, the training system 200 inserts a respective normalization layer 54.1 immediately before the BatchNorm layer to generate the respective normalized nonlinear unit 50.1 in the neural network 100, where the normalization layer54.1 is a normalization layer of a different type than the BatchNorm layer. For example, the normalization layer 54.1 can be a LayerNorm layer, a WeightNorm layer, a RMSNorm layer, or a GroupNorm layer.

[0111] The normalized nonlinear unit 50.1 is parametrized by a respective set of block parameters &i — [WbYi, Pi} including the respective weight tensor (VF;) of the linear layer 52.1, the respective scaling tensor (y of the normalization layer 54.1, and the respective shift tensor ( / ?j) of the normalization layer 54.1.

[0112] The training system 200 returns a set of network parameters (0) of the neural network 100, where the set of network parameters, given as 0 = {0 f=1, includes the respective set of block parameters of each normalized nonlinear unit 50.1 in the neural network 100.

[0113] The training system 200 then trains the neural network 100 on the training dataset 210 over a number of training iterations (t), e.g., corresponding to a sequence of training iterations. Here, t = 1,2, ... , T indexes a training iteration, and T denotes the total number of training iterations. For example, the training system 200 can perform at least about 100, 1,000, 10,000, 100,000, 1,000,000, or more training iterations.

[0114] In some implementations, the training system 200 initializes a respective initial value (01) of each network parameter in the set of network parameters. For example, the training system 200 can perform a random initialization by drawing the initial values from a probability distribution, e.g., a uniform or normal distribution. The training system 200 can also perform a Xavier initialization, Glorot initialization, He initialization, LeCun initialization, orthogonal initialization, sparse initialization, or any appropriate initialization.

[0115] At each training iteration t = 1,2, ... , T, the training system 200 is configured to perform the following operations:

[0116] The training system 200 receives a respective current value (0f) of each network parameter in the set of network parameters of the neural network 100. The current set of network parameters, given as 0f= {0 }f=1, includes a respective current set of block parameters (0 ) for each normalized nonlinear unit 50.1 in the neural network 100. The training system 200 then parametrizes the neural network 100 with the respective current value of each network parameter T7= '(0t), e.g., in preparation for evaluating the neural network 100 under the current set of network parameters.

[0117] The training system 200 samples, from the training dataset 210, a batch of training data 212. t at the training iteration. For example, the training system 200 can randomly sample the batch 212. t from the training dataset 210 or according to an algorithm. The batch of training data 212. t is a subset of training data including a subset of samples214.t.1 to 214. t.N from the training dataset 210. Here, n = 1,2, ..., N indexes a sample 214. t.n in the batch of training data 212. t, and N denotes the total number of samples 214.t in the batch of training data 212. t, e.g., corresponding to the batch size of the batch of training data 212. t. For example, the batch of training data 212. t can include at least about 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 500, 1000, or more samples 214.t.

[0118] The training system 200 then trains the neural network 100 on the batch of training data 212.t to generate a respective updated value (0t+ 1) of each network parameter in the set of network parameters at the training iteration. The updated set of network parameters, given as 0t+1= {0;t+1}f=i, includes a respective updated set of block parameters (0+1) for each normalized nonlinear unit 50.1 in the neural network 100. The training system 200 can use any of a variety of machine learning techniques to generate the updated values of the network parameters at the training iteration. For example, the training system 200 can use an unsupervised learning technique, a semi-supervised learning technique, a supervised learning technique, a continual learning technique, or a reinforcement learning technique. Examples of each are described below with reference to FIGs. 3A-3C.

[0119] The training system 200 then determines whether a criterion is satisfied at the training iteration. In general, the criterion specifies whether the updates values of the network parameters should be normalized at the training iteration. When the criterion is not satisfied at the training iteration, the training system 200 does not normalize the updates values of the network parameters. When the criterion is satisfied at the training iteration, the training system 200 normalizes the respective updated values of the respective set of block parameters of each normalized nonlinear unit 50.1 in the neural network 100, which can include some, or all, of the updated values of the set of block parameters, as described in more detail below.

[0120] To determine whether a criterion is satisfied at the training iteration, the training system 200 can determine whether a threshold number (Tth) of training iterations has elapsed since the set of network parameters were last normalized. The threshold number of training iterations can be about 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 500, 1000, ormore training iterations, e.g., depending on a threshold frequency of parameter-wise normalization during the NaP algorithm. For example, the threshold number of training iterations can be a default value of the training system 200. As another example, the training system 200 can receive the threshold number of training iterations as input, e.g., with the neural network 100 and / or training dataset 210. As yet another example, the training system 200 can derive the threshold number of training iterations based on other hyperparameters implemented during training, e.g., such that the threshold number of training iterations is inversely proportional to the learning rate Tthoc j?-1.

[0121] Alternatively, or in addition, the training system 200 can determine whether a norm (||0t+i||) over the respective updated value of each network parameter in the set of network parameters is greater than a threshold norm (Cth) for the set of network parameters. The threshold norm for the set of network parameters can have a value of about 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 500, 1000, or more, e.g., depending on the total number and / or maximum values of the network parameters. For example, the threshold norm can be a default value of the training system 200. As another example, the training system 200 can receive the threshold norm as input, e.g., with the neural network 100 and / or training dataset 210. As yet another example, the training system 200 can derive the threshold norm based on the training dynamics, e.g., such that threshold norm is the value at which the effective learning rate is near-zero, e.g., less than a threshold effective learning rate.

[0122] The training system 200 can individually normalize the respective updated set of block parameters of each normalized nonlinear unit 50.1 in the neural network 100, e.g., corresponding to a per-layer parameter- wise normalization.

[0123] The updated set of block parameters of the normalized nonlinear unit 50.1 is given as:

[0124] which includes a respective updated weight tensor (l / F(f+1) of the linear layer 52.1, a respective updated scaling tensorof the normalization layer 54.1, and a respective updated shift tensor ( ?+1) of the normalization layer 54.1.

[0125] The training system 200 normalizes the updated weight tensor of the linear layer 52.1 to a respective reference norm (C;) for the normalized nonlinear unit 50.1 as:

[0126] where CLcorresponds to the radius of a sphere the updated weight tensor is projected onto. For example, the reference norm for the normalized nonlinear unit 50.1 can be the norm of the initial weight tensor Ct=|| at the initial training iteration. Alternatively, the reference norm for the normalized nonlinear unit 50.1 can have a fixed value of about 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 500, 1000, or more, e.g., depending on the total number and maximum desired values of the weight tensor.

[0127] In some implementations, the training system 200 normalizes the updated scaling tensor and the updated shift tensor of the normalization layer 54.1 to a respective reference size (£>;) for the normalized nonlinear unit 50.1 as:

[0128] For example, the reference size for the normalized nonlinear unit 50.1 can be the size of the scaling tensor D, = size(y;), the size of the shift tensor D, = size( ?;), or the size of a tensor product between the two Dt= size(yj= size(yj)size( / ?z), e.g., a product between their respective sizes. Alternatively, the reference size for the normalized nonlinear unit 50.1 can be a fixed value of about 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 500, 1000, or more.

[0129] In some implementations, the training system 200 regularizes the updated scaling tensor and the updated shift tensor of the normalization layer 54.1 by applying a respective weight decay (a for the normalized nonlinear unit 50.1 to the updated scaling and shift tensors as:

[0130] where the weight decay, given as 0 < aL< 1, is a positive scalar between zero, and H is an identity tensor. For example, when initializing the initial values of the network parameters, the training system 200 can initialize the scaling tensor as the identity tensor y* = H, and the shift tensor as a zero tensor ?* = 0. The weight decay will then drive the scaling and shift tensors to these initial values.

[0131] In some implementations, the training system 200 determines whether to normalize or regularize the updated scaling and shift tensors of the normalization layer 54.1 based on the nonlinear transformation applied by the nonlinear function of the nonlinear layer 56.1. For example, the training system 200 can normalize the updated scaling and shift tensors according toEq. (3c) if the nonlinear function is an activation function, and regularize the updated scaling and shift tensors according to Eq. (3d) if the nonlinear function is not an activation function.

[0132] In some implementations, the training system 200 both normalizes and regularizes the updated scaling and shift tensors of the normalization layer 54.1. For example, the training system 200 can first normalize the updated scaling and shift tensors according to Eq. (3c), and subsequently regularize the updated scaling and shift tensors according to Eq. (3d). Alternatively, the training system 200 can first regularize the updated scaling and shift tensors according to Eq. (3d), and subsequently normalize the updated scaling and shift tensors according to Eq. (3c).

[0133] The training system 200 repeats this process at the next training iteration (t + 1), and the next training iteration (t + 2), and so on, until the final training iteration is reached t = T, or a termination condition is satisfied at a training iteration of t < T, e.g., an objective function 22O.t characterizing a performance of the neural network 100 on the batch of training data 212.t satisfies a threshold at the training iteration.

[0134] FIG. 3A is a schematic diagram depicting an example of the training system 200A for training the neural network 100 via the NaP algorithm in an unsupervised, semi-supervised, or supervised learning setting. Example experiments performed on the neural network 100 in the supervised learning setting are described in Sec. VI.

[0135] In this case, the training dataset 210 is a stationary training dataset 210A that can be unlabeled, partially labeled, or labeled. For example, the training system 200A can be provided with the training dataset 210A, or retrieve the training dataset 210A from an online repository.

[0136] Here, at each training iteration, each sample 214. t.n in the batch of training data 212. t at the training iteration includes a respective network input (X^) 102. t.n for the neural network 100. In some implementations, e.g., in the semi-supervised or supervised learning settings, the sample 214. t.n further includes a respective target output (K„*) 104*. t.n for the network input 102. t.n. For example, the target output 104*. t.n can be a ground truth output that the neural network 100 is expected to generate upon processing the network input 102. t.n.

[0137] For each sample 214. t.n in the batch of training data 212.t at the training iteration, the training system 200A is configured to perform the following operations:

[0138] The training system 200A processes the network input 102. t.n of the sample 214. t.n, using the neural network 100 parametrized by the respective current value of each network parameter at the training iteration, to generate a respective network output (Vj ) 104. t.n for the sample 214.t.n.The operations of the neural network 100 can be represented concisely as Y = 1F(X ; 0f), where J?(0t) is the function representing the parametric model of the neural network 100, parametrized by the current set of network parameters at the training iteration.

[0139] The training system 200A computes a respective objective function term (Lj 222.t.n for the sample 214. t.n that depends on the network output 104. t.n. In some implementations, the objective function term 222. t.b further depends on the target output 104*.t.n of the sample 214.t.n. For example, the objective function term 222.t.n can measure a respective error between the target 104*.t.n and network 104*.t.n outputs of the sample 214.t.n. For example, the objective function term 222. t.n can include a Euclidean distance loss, a Manhattan distance loss, a Chamfer distance, a Harsdorf distance, or other distance, similarity, or error metric.

[0140] The training system 200A generates an objective function (FT) 22O.t including the respective objective function term 222.t.n for each sample 214. t.n in the batch of training data 212. t. For example, the objective function 220. t can include a summation of the objective function terms 222.t such that each sample 214.t.n is weighted equally. Alternatively, the objective function 220. t can include a weighted combination of the objective function terms 222. t, e.g., a weighted linear combination, such that different samples 214. t are emphasized and deemphasized.

[0141] The training system 200A computes, with respect to the set of network parameters of the neural network 100, a gradient 224.t of the objective function 22O.t at the training iteration. The gradient 224. t of the objective function 22O.t can be expressed as:

[0142] For example, the training system 200A can compute the gradient 224. t of the objective function 220. t with respect to the set of network parameters using backpropagation, e.g., implemented via automatic differentiation.

[0143] The training system 200A then augments, using the gradient 224. t of the objective function 22O.t, the respective current value of each network parameter at the training iteration to generate the respective updated value of the network parameter at the training iteration. For example, to update the current values of the set of network parameters using the gradient 224. t of the objective function 22O.t, the training system 200A can use an optimizer implementing a (e.g., stochastic) gradient descent method with a particular learning rate, learning rate schedule, and / or weight decay. Examples of optimizers include adaptive and non-adaptive optimizers such as stochastic gradient descent (“SGD”) with Momentum, Nesterov Accelerated Gradient (“NAG”), Implicitupdates, Adam, RMSProp, AdaGrad, among others. In some implementations, the training system 200A uses an optimizer that corrects for the norm of the gradient 224. t, e.g., by normalizing the gradient 224. t by its respective norm at each training iteration.

[0144] As a particular example of an optimizer, the training system 200 can generate the updated values of the network parameters as:

[0145] where 7] is a learning rate. For example, the learning rate can be a fixed learning rate having a value of about 10’1, 10'2, 10'3, 10'4, 10'5, or less, or in a range between any of these two values. In other examples, the learning rate is reduced from a learning rate schedule specifying a respective current value ( ) of the learning rate at each training iteration, e.g., according to a time-based, step-based, or exponential-based learning rate schedule.

[0146] FIG. 3B is a schematic diagram depicting an example of the training system 200B for training the neural network 100 via the NaP algorithm in a continual learning setting, e.g., an online learning setting or an online continual learning setting. Example experiments performed on the neural network 100 in the continual learning setting are described in Sec. VII.

[0147] In this case, the training dataset 210 is a data stream 210B, e.g., a data stream that the training system 200 receives and processes in real-time. Here, each training iteration is a respective timestep in the data stream 210B, where the data stream 210B includes the respective batch of training data 212. t at the respective timestep. In some implementations, the data stream 210B is a nonstationary data stream, e.g., a data stream having statistical properties that changes dynamically at each timestep, e.g., due to a different distribution of the samples 214.t in the batch of training data 212.t at the timestep.

[0148] The training system 200B trains the neural network 100 on the batch of training data 212.t at the timestep in the same manner as the training system 200A trains the neural network 100 on the batch of training data 212.t at the corresponding training iteration, only in real-time. Hence, the description above for the training system 200A holds equally well for the training system 200N. For example, the training system 200B can implement the continual learning in the unsupervised, semi-supervised, or supervised learning settings.

[0149] FIG. 3C is a schematic diagram depicting an example of the training system 200C for training the neural network 100 via the NaP algorithm in a reinforcement learning setting. Examples of reinforcement learning tasks the training system 200C can train the neural network 100 to perform the NaP algorithm are described in Sec. V. Example experiments performed on the neural network 100 in the reinforcement learning setting are described in Sec. VI.

[0150] In this case, the training system 200C trains the neural network 100 to process a network input 102 including an observation 232 characterizing a state (s) 231 of an environment 230 being interacted with by an agent 235, and generate a network output 104 defining an action (a) 234 to be performed by the agent 235 in response to the observation 232.

[0151] In some implementations, the neural network 100 is configured to: receive a network input 102 including the observation 232; and process the network input 102 to generate a network output 104 including a policy (n) for selecting actions to be performed by the agent 235 in response to the observation 232. The training system 200C then selects the action 234 to be performed by the agent 235 in accordance with the policy. For example, the policy can be a score distribution over possible actions to be performed by the agent 235 in response to the observation 232, e.g., a Gaussian distribution, a softmax distribution, or a beta distribution. The training system 200C can then sample the action 234 to be performed by the agent 235 from the score distribution.

[0152] In some implementations, the neural network 100 is configured to: receive a network input 102 including the observation 232; and process the network input 102 to generate a network output 104 including the action 234 to be performed by the agent 235 in response to the observation 232. That is, the policy is a deterministic policy, and the neural network 100 directly outputs the action234 upon processing the observation 232.

[0153] Here, the training dataset 210 is a replay buffer 210C that was generated using the neural network 100 to control the agent 235 in the environment 230. That is, the respective batch of training data 212.t at each training iteration is a respective interaction trajectory between the agent235 and the environment 230 while the agent 235 was being controlled by the neural network 100. In some implementations, the replay buffer 210C is a nonstationary replay buffer, e.g., a replay buffer having statistical properties that changes dynamically at each training iteration, e g., due to being updated by the neural network 100 at the training iteration.

[0154] To obtain a batch of training data 212. t for a training iteration, the training system 200C executes a simulation 240. t of the agent 235 interacting with the environment 230 while controlledby the neural network 100. Each sample 214. t.n in the batch of training data 212. t, given as djj =(s^, dn>rn>sn+i )> isarespective interaction between the agent 235 and the environment 230 at a respective simulation timestep (n) in the simulation 24O.t. The sample 214.t.n includes: (i) a respective state (s ) 231.t.n of the environment 230 at the simulation timestep; (ii) a respective action (a ) 234. t.n performed by the agent 235 in response to an observation 232.t.n characterizing the state 231. t.n of the environment 230 at the simulation timestep; (iii) a respective reward (r7f) 236. t.n received by the agent 235 in response to performing the action 234. t.n in the state 231. t.n of the environment 230 at the simulation timestep; and (iv) a respective state231 ,t.(n+l) of the environment 230 at the next simulation timestep in response to performing the action 231. t.n in the state 231.t.n of the environment 230 at the simulation timestep.

[0155] At each simulation timestep in the simulation 240. t, the training system 200C is configured to perform the following operations:

[0156] The training system 200C receives, from the simulation 24O.t, the observation 232.t.n at the simulation timestep. The training system 200C forwards, to the neural network 100, a network input 102. t.n including the observation 232. t.n at the simulation timestep. The training system 200C processes the network input 102.t.n, using the neural network 100, to generate a network output 104.t.n defining the action 234.t.n to be performed by the agent 235 at the simulation timestep. The training system 200C forwards, to the simulation 240. t, the action 234.t.n to be performed by the agent 235 at the simulation timestep. The simulation 24O.t then simulates the agent 235 performing the action 234.t.n in the state 231.t.n of the environment 230 at the simulation timestep to determine the state 23 l.t.(n+l) of the environment 230 at the next simulation timestep and the reward 236.t.n for the agent 235 at the simulation timestep. The training system 200C then returns the sample 214. t.n including the state 231.t.(n+l) at the next simulation timestep, and the action 234.t.n, state 231. t.n, and reward 236. t.n at the current simulation timestep.

[0157] To train the neural network 100 on a batch of training data 212. t at a training iteration, the training system 200C samples the batch of training data 212.t from the replay buffer 210C and generates the objective function 22O.t evaluated on the batch of training data 212.t. For example, the objective function 22O.t can include an expected return, a policy gradient objective, a value function loss, an actor-critic objective, or a Kullback-Leibler divergence. The training system 200C then generates the updated values of the network parameters of the neural network 100 in the same manner as the training system 200B. Particularly, the training system 200C computes thegradient 224.t of the objective function 22O.t with respect to the set of network parameters of the neural network 100, e.g., using backpropagation. The training system 200C then augments, using the gradient 224. t of the objective function 220. t, the respective current value of each network parameter at the training iteration to generate the respective updated value of the network parameter at the training iteration, e.g., according to Eq. (4b) and / or using a gradient descent method such as Implicit updates, Momentum, Adam, RMSProp, or AdaGrad.

[0158] Note, the procedure described above is just one example of how the training system 200C can train the neural network 100 via the NaP algorithm in the reinforcement learning setting. However, there are many different implementations of the training system 200C depending on how the batch of training data 212.t for the training iteration is generated by the neural network 100, e.g., where the environment 230 is a real-world environment or a simulation environment, where the batch 212.t is generated continually, on-policy, off-policy, and so on. Examples of some of these implementations are described below with further examples described in Sec. V.

[0159] In some implementations, the training system 200C periodically (e.g., continually) generates the respective batch of training data 212. t at each training iteration using the neural network 100, and subsequently trains the neural network 100 on the batch of training data 212.t at the training iteration. In this case, the training system 200C can implement the NaP algorithm in an “on-policy” reinforcement learning setting, such that the neural network 100 learns (e.g., only learns) from the batch of training data 212.t that it generated under the current values of the set of network parameters at the current training iteration. In other words, the training system 200C does not reuse previous batches 212. (t-1), 212. (t-2), etc., that were generated by the neural network 100 under previous values of the network networks at the previous training iterations. For example, to update the replay buffer 210C at each training iteration, the training system 200C can overwrite the batch of training data 212. (t-1) at the previous training iteration with the batch of training data 212.t generated by the neural network 100 at the current training iteration.

[0160] These implementations of the training system 200C can be useful for nonstationary (e.g., highly nonstationary) settings. For example, the training system 200C can implement the NaP algorithm in a continual reinforcement learning setting, where the agent 235 interacts with the environment 230 over a single interaction trajectory, and the neural network 100 is continually updated during the single interaction trajectory. In this case, the respective batch of training data 212. t at each training iteration is a respective fragment of the single interaction trajectory betweenthe agent 235 and the environment 230, while the agent 235 was being controlled by the neural network 100 during the respective fragment of the single interaction trajectory. For example, each training iteration can correspond to a respective subset of simulation timesteps elapsed in a single simulation 240 executed by the training system 200C to generate the single interaction trajectory. As another example, the environment 230 can be a real-world environment of the agent 235, and each training iteration can correspond to a respective subset of real-world timesteps elapsed in the real-world environment.

[0161] In some implementations, the training system 200C generates the respective batch of training data 212.t for each training iteration using the neural network 100, while the neural network 100 is parametrized by the initial values of the network parameters. In other words, the replay buffer 210C is fixed over the training iterations. In this case, the training system 200C can implement the NaP algorithm in an “off-policy” reinforcement learning setting, such that the neural network 100 learns from batches 212 that were generated by the neural network 100 under different values of the network networks. Here, there may be no ordering and / or correlations between the batches 212 in the replay buffer 210C, e.g., such that the training system 200C can sample each batch of training data 212. t from the replay buffer 210C at random.

[0162] FIG. 4 is a flow diagram of an example process 300 for training a neural network 100 to perform a machine learning task via a NaP training algorithm. For convenience, the process 300 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 200 of FIG. 3, appropriately programmed in accordance with this specification, can perform the process 300.

[0163] The training system 200 receives the neural network 100 and obtains a training dataset 210 for training the neural network 100 to perform the machine learning task (302).

[0164] The training system 200 inserts, into the neural network 100, a respective normalization layer 54.1 immediately before each unnormalized nonlinear layer of the neural network 100 to generate a respective normalized nonlinear unit 50.1 (304). Particularly, for each nonlinear layer 56.1 in the neural network 100, the training system 200 determines whether a normalization layer 54 immediately precedes the nonlinear layer 56.1. If not, the neural network 100 inserts the respective normalization layer 54.1 immediately before the nonlinear layer 56.1 to generate the respective normalized nonlinear unit 50.1 in the neural network 100.

[0165] In some implementations of the process 300, the training system 200 initializes a respective initial value of each network parameter in a set of network parameters parametrizing the neural network 100, where the set of network parameters includes a respective set of block parameters parametrizing each normalized nonlinear unit 50.1 in the neural network 100.

[0166] At each of multiple training iterations t = 1,2, ... , T, the training system 200 performs operations (306. t) to (318.t):

[0167] The training system 200 receives a respective current value of each network parameter in the set of network parameters at the training iteration (306.t)

[0168] The training system 200 parametrizes the neural network 100 with the respective current value of each network parameter in the set of network parameters (308.t).

[0169] The training system 200 samples, from the training dataset 210, a batch of training data 212.t for the training iteration (3 lO.t).

[0170] The training system 200 trains the neural network 100 on the batch of training data 212.t to generate a respective updated value of each network parameter in the set of network parameters at the training iteration (312. t). For example, as described above with reference to FIGs. 3A-3C, the training system 200 can train the neural network 100 on the batch of training data 212. t using an unsupervised learning technique, a semi -supervised learning technique, a supervised learning technique, a continual learning technique, or a reinforcement learning technique.

[0171] The training system 200 determines whether a criterion for normalizing the set of network parameters of the neural network 100 is satisfied at the training iteration (314.t). For example, as described above with reference to FIGs. 3A-3C, the training system 200 can determine whether a threshold number of training iterations has elapsed since the set of network parameters were last normalized and / or whether a norm over the respective updated value of each network parameter in the set of network parameters is greater than a threshold norm for the set of network parameters.

[0172] When the criterion for normalizing the set of network parameters of the neural network 100 is not satisfied at the training iteration, the training system 200 breaks to operation (306.t+l) and the operations repeat at the next training iteration (t + 1).

[0173] When the criterion for normalizing the set of network parameters of the neural network 100 is satisfied at the training iteration, the training system 200 breaks to operation (316.t).

[0174] The training system 200 normalizes the updated values of the set of network parameters of the neural network 100 (316.t). For example, as described above with reference to FIGs. 3A-3C,the training system 200 can perform a per-layer parameter- wise normalization of the neural network 100 to normalize the respective updated values of the respective set of block parameters of each normalized nonlinear unit 50.1 in the neural network 100. This can include some, or all, of the updated values of the set of block parameters. For example, the training system 200 can normalize a respective updated weight tensor of the normalized nonlinear unit 50.1. As another example, the training system 200 can normalize the updated weight tensor and a respective updated scaling and shift tensor of the normalized nonlinear unit 50.1. In some implementations, the training system 200 regularizes the updated scaling and shift tensors alternatively, or in addition, to normalizing the updated scaling and shift tensors.

[0175] The training system 200 then breaks to operation (306. t+1) and the training system 200 repeats the operations at the next training iteration (t + 1).V. Examples of Machine Learning Tasks

[0176] The training system 200 can train the neural network 100 via the NaP algorithm to perform any of multiple types of image processing tasks. For the image processing task, the training system 200 trains the neural network 100 receive: a network input 102 including an input image; and process respective intensity values of pixels of pixels of the input image to generate network output 104 for the input image. The input image may be a real-world image. The input image may be an image of a real-world environment captured by a camera.

[0177] For example, the image processing task can be an object classification task, and the network output 104 generated by the neural network 100 can include a respective score for each object category in a set of object categories, where the score represents an estimated likelihood that the input image depicts an object belonging to the category. As another example, the image processing task can be an image embedding generation task, and the network output 104 generated by the neural network 100 can be a numeric embedding of the input image. As yet another example, the image processing task can be an object detection task, and the network output 104 generated by the neural network 100 can identify locations in the input image at which particular types of objects are depicted. As yet another example, the image processing task can be an image segmentation task and the network output 104 generated by the neural network 100 can assign each pixel of the input image to a category from a set of categories. Other examples of image processing tasks include depth prediction, optical flow prediction, and so on.

[0178] The training system 200 can train the neural network 100 via the NaP algorithm to perform any of multiple types of audio processing tasks.

[0179] For example, the audio processing task can be an audio transcription task where, if the network input 102 to the neural network 100 includes a sequence representing a spoken utterance, the network output 104 generated by the neural network 100 can include a respective score for each piece of text in a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. As another example, the audio processing task can be a keyword spotting task where, if the network input 102 to the neural network 100 includes a sequence representing a spoken utterance, the network output 104 generated by the neural network 100 can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance. As another example, the audio processing task can be a language recognition task where, if the network input 102 to the neural network 100 includes a sequence representing a spoken utterance, the network output 104 generated by the neural network 100 can identify the natural language in which the utterance was spoken. As another example, the audio processing task can be an audio classification task, e.g., identifying an audio source that generated the sequence of audio, identifying the speaker in the audio, and so on.

[0180] The training system 200 can train the neural network 100 via the NaP algorithm to perform any of multiple types of score or classification tasks. For the score or classification task, the training system 200 trains the neural network 100 to process a network input 102, and generate a network output 104 including a score or classification output based on the network input 102.

[0181] For example, if the network input 102 to the neural network 100 includes image or features that have been extracted from images, the network output 104 generated by the neural network 100 for a given image can include a respective score for each object category in a set of object categories, with the score representing an estimated likelihood that the given image depicts an object belonging to the object category.

[0182] As another example, if the network input 102 to the neural network 100 includes Internet resources (e.g., web pages), documents, portions of documents, or features extracted from Internet resources, documents, or portions of documents, the network output 104 generated by the neural network 100 for a given Internet resource, document, or portion of a document can include a respective score for each topic in a set of topics, with the score representing an estimated likelihood that the given Internet resource, document, or document portion is related to the topic.

[0183] As another example, if the network input 102 to the neural network 100 includes features of an impression context for a particular advertisement, the network output 104 generated by the neural network 100 can include a score representing an estimated likelihood that the particular advertisement will be clicked on.

[0184] As another example, if the network input 102 to the neural network 100 includes features of a personalized recommendation for a user, e.g., features characterizing the context for the recommendation, e.g., features characterizing previous actions taken by the user, the network output 104 generated by the neural network 100 can include a respective score for each content item in a set of content items, with the score representing an estimated likelihood that the user will respond favorably to being recommended the content item.

[0185] As another example, if the network input 102 to the neural network 100 includes input text in one language, the network output 104 generated by the neural network 100 can include a respective score for each piece of text in a set of pieces of text in another language, with the score representing an estimated likelihood that the piece of text in the other language is a proper translation of the input text into the other language.

[0186] As another example, if the network input 102 to the neural network 100 includes a spoken utterance, a sequence of spoken utterances, or features derived from one of the two, the network output 104 generated by the neural network 100 can include a respective score for each piece of text in a set of pieces of text, with the score representing an estimated likelihood that the piece of text is the correct transcript for the utterance or sequence of utterances.

[0187] The training system 200 can train the neural network 100 via the NaP algorithm to perform any of multiple types of reinforcement learning tasks. For the reinforcement learning task, the training system 200 trains the neural network 100 to process a network input 102 including an observation 232 characterizing a state 231 of the environment 230 being interacted with by the agent 235, and generate a network output 104 defining an action 234 to be performed by the agent 235 in response to the observation 232.

[0188] In some implementations, the environment 230 is a real -world environment, the agent 235 is a mechanical agent interacting with the real-world environment, e.g., a robot, an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment 230, and the actions 234 performed by the agent 235 are actions taken by the mechanical agent in the real-world environment to perform the task. For example, the mechanical agent may be interactingwith the environment 230 to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment 230, move an object of interest to a specified location in the environment 230, or navigate to a specified destination in the environment.

[0189] In these implementations, the observations 232 can include one or more of images, object position data, or sensor data to capture observations as the agent 235 interacts with the environment 230, e.g., sensor data from an actuator, an image sensor, a distance sensor, or a position sensor located on the agent. For example, in the case of a robot, the observations 232 may include data characterizing the current state of the robot, e.g., data including one or more of joint positionjoint velocityjoint force, joint torque, or joint acceleration (e.g., gravity-compensated torque feedback), and a global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle, the observations 232 can similarly include one or more of the position, linear velocity, angular velocity, force, torque, or acceleration of one or more parts of the agent 235, and the global or relative pose of the one or more parts of the agent 235. The observations 232 may be defined in 1, 2, or 3 dimensions, and may be absolute and / or relative observations. The observations 232 may also include, for example, sensed electronic signals such as motor current or a temperature signal; and / or image or video data for example from a camera or a LIDAR sensor, e.g., data from sensors of the agent 235 or data from sensors that are located separately from the agent 235 in the environment 230.

[0190] In these implementations, the actions 234 may be control signals to control the robot or other mechanical agent, e.g., torques for the joints of the robot or higher-level control commands, or the autonomous or semi-autonomous land, air, sea vehicle, e.g., torques to the control surface or other control elements, e.g., steering control elements of the vehicle, or higher-level control commands. The control signals can include, for example, the position, velocity, force, torque, and / or acceleration data for one or more joints of a robot or parts of another mechanical agent. Alternatively, or in addition, the control signals may include electronic control data such as motor control data, or more generally data for controlling one or more electronic devices within the environment 230, where the control of which influences the observation 232 of the state 231 of the environment 230. For example, in the case of an autonomous or semi -autonomous land, air, or sea vehicle, the control signals may define actions to control navigation e.g., steering, and movement e.g., braking and / or acceleration of the vehicle.

[0191] In some implementations, the environment 230 is implemented in a simulation 240 of the above-described real -world environment, and the agent 235 is implemented as one or more computers interacting with the environment 230 in the simulation 240. For example, the simulation 240 may be a simulation of a robot or vehicle and the training system 200 trains the neural network 100 on the simulation 240 and then, once trained, the neural network 100 can be used in the real- world.

[0192] In some implementations, the environment 230 is a real-world manufacturing environment for manufacturing a product, such as a chemical, biological, or mechanical product, or a food product. As used herein a “manufacturing” a product also includes refining a starting material to create a product, or treating a starting material, e.g., to remove pollutants, to generate a cleaned or recycled product. The manufacturing plant may include multiple manufacturing units such as vessels for chemical or biological substances, or machines, e.g., robots, for processing solid or other materials. The manufacturing units are configured such that an intermediate version or component of the product is moveable between the manufacturing units during manufacture of the product, e.g., via pipes or mechanical conveyance. As used herein manufacture of a product also includes manufacture of a food product by a kitchen robot.

[0193] The agent 235 may include an electronic agent configured to control a manufacturing unit, or a machine such as a robot, that operates to manufacture the product. That is, the agent 235 may include a control system configured to control the manufacture of the chemical, biological, or mechanical product. For example, the control system may be configured to control one or more of the manufacturing units or machines, or to control movement of an intermediate version or component of the product between the manufacturing units or machines.

[0194] As one example, a task performed by the agent 235 may include a task to manufacture the product or an intermediate version or component thereof. As another example, a task performed by the agent 235 may include a task to control, e.g., minimize, use of a resource such as a task to control electrical power consumption, or water consumption, or the consumption of any material or consumable used in the manufacturing process.

[0195] The actions 234 performed by the agent 235 may include control actions to control the use of a machine or a manufacturing unit for processing a solid or liquid material to manufacture the product, or an intermediate or component thereof, or to control movement of an intermediate version or component of the product within the manufacturing environment, e.g., between themanufacturing units or machines. In general, the actions 234 performed by the agent 235 may be any actions that influence the observation 232 of the state 231 of the environment 230, e.g., actions configured to adjust any of the sensed parameters described below. These may include actions to adjust the physical or chemical conditions of a manufacturing unit, or actions to control the movement of mechanical parts of a machine or joints of a robot. The actions 234 performed by the agent 235 may include actions imposing operating conditions on a manufacturing unit or machine, or actions that result in changes to settings to adjust, control, or switch on or off the operation of a manufacturing unit or machine.

[0196] The rewards 236 may relate to a metric of performance of the task. For example, in the case of a task that is to manufacture a product, the metric of performance may include a metric of a quantity of the product that is manufactured, a quality of the product, a speed of production of the product, or to a physical cost of performing the manufacturing task, e.g., a metric of a quantity of energy, materials, or other resources, used to perform the task. In the case of a task that is to control use of a resource the metric of performance may include any metric of usage of the resource.

[0197] In general, the observations 232 of the states 231 of the environment 230 may include any electronic signals representing the functioning of electronic and / or mechanical items of equipment. For example, an observation 232 of a state 231 of the environment 230 may be derived from observations made by sensors sensing a state of the manufacturing environment, e.g., sensors sensing a state or configuration of the manufacturing units or machines, or sensors sensing movement of material between the manufacturing units or machines. As some examples such sensors may be configured to sense mechanical movement or force, pressure, temperature; electrical conditions such as current, voltage, frequency, impedance; quantity, level, flow / movement rate or flow / movement path of one or more materials; physical or chemical conditions e.g. a physical state, shape or configuration or a chemical state such as pH; configurations of the units or machines such as the mechanical configuration of a unit or machine, or valve configurations; image or video sensors to capture image or video observations of the manufacturing units or of the machines or movement; or any other appropriate type of sensor. In the case of a machine such as a robot, the observations from the sensors may include observations of position, linear velocity, angular velocity, force, torque, and / or acceleration, or pose of one or more parts of the machine, e g., data characterizing the current state of the machine or robot or ofan item held or processed by the machine or robot. The observations may also include, for example, sensed electronic signals such as motor current or a temperature signal, or image or video data for example from a camera or a LIDAR sensor. Sensors such as these may be part of or located separately from the agent in the environment.

[0198] In some implementations, the environment 230 is the real-world environment of a service facility including multiple items of electronic equipment, such as a server farm or data center, e.g., a telecommunications data center, or a computer data center for storing or processing data, or any service facility. The service facility may also include ancillary control equipment that controls an operating environment of the items of equipment, for example environmental control equipment such as temperature control, e.g., cooling equipment, or air flow control or air conditioning equipment such as a heater, a cooler, a humidifier, or other hardware that modifies a property of air in the real-world environment. The task may include a task to control, e.g., minimize, use of a resource, such as a task to control electrical power consumption, or water consumption. The agent 235 may include an electronic agent configured to control operation of the items of equipment, or to control operation of the ancillary, e.g., environmental, control equipment.

[0199] In some implementations, the environment 230 is the real-world environment of a power generation facility e.g., a renewable power generation facility such as a solar farm or wind farm. The task may include a control task to control power generated by the facility, e.g., to control the delivery of electrical power to a power distribution grid, e.g., to meet demand or to reduce the risk of a mismatch between elements of the grid, or to maximize power generated by the facility. The agent 235 may include an electronic agent configured to control the generation of electrical power by the facility or the coupling of generated electrical power into the grid. The actions 234 performed by the agent 235 may include actions to control an electrical or mechanical configuration of an electrical power generator such as the electrical or mechanical configuration of one or more renewable power generating elements e.g., to control a configuration of a wind turbine or of a solar panel or panels or mirror, or the electrical or mechanical configuration of a rotating electrical power generation machine. Mechanical control actions may, for example, include actions that control the conversion of an energy input to an electrical energy output, e.g., an efficiency of the conversion or a degree of coupling of the energy input to the electrical energy output. Electrical control actions may, for example, include actions that control one or more of a voltage, current, frequency, or phase of electrical power generated.

[0200] In the case of a task to control the delivery of electrical power to the power distribution grid, the metric of performance can relate to a measure of power transferred, or to a measure of an electrical mismatch between the power generation facility and the grid such as a voltage, current, frequency or phase mismatch, or to a measure of electrical power or energy loss in the power generation facility.

[0201] In the case of a task to maximize the delivery of electrical power to the power distribution grid, metric of performance may relate to a measure of electrical power or energy transferred to the grid, or to a measure of electrical power or energy loss in the power generation facility.

[0202] In general, the observations 232 of the states 231 of the environment 230 may include any electronic signals representing the electrical or mechanical functioning of power generation equipment in the power generation facility. For example, an observation 232 of a state 231 of the environment 230 may be derived from observations made by any sensors sensing a physical or electrical state of equipment in the power generation facility that is generating electrical power, or the physical environment of such equipment, or a condition of ancillary equipment supporting power generation equipment. Such observations may thus include observations of wind levels or solar irradiance, or of local time, date, or season. Such sensors may include sensors configured to sense electrical conditions of the equipment such as current, voltage, power, or energy; temperature or cooling of the physical environment; fluid flow; or a physical configuration of the equipment; and observations of an electrical condition of the grid e.g., from local or remote sensors. An observation 232 of a state 231 of the environment 230 may also include one or more predictions regarding future conditions of operation of the power generation equipment such as predictions of future wind levels or solar irradiance or predictions of a future electrical condition of the grid.

[0203] As another example, the environment 230 may be a chemical synthesis or protein folding environment, such that each state 231 of the environment 230 is a respective state of a protein chain or of one or more intermediates or precursor chemicals, and the agent 235 is a computer system for determining how to fold the protein chain or synthesize the chemical. In this example, the actions 234 performed by the agent 235 are possible folding actions for folding the protein chain or actions for assembling precursor chemicals / intermediates and the result to be achieved may include, e.g., folding the protein so that the protein is stable and so that it achieves a particular biological function or providing a valid synthetic route for the chemical. As another example, theagent 235 may be a mechanical agent that indirectly performs or controls the protein folding actions, or chemical synthesis steps, e.g., by controlling synthesis steps selected by the system automatically without human interaction. The observations 232 of the states 231 of the environment 230 may include direct or indirect observations of a state of the protein or chemical / intermediates / precursors and / or may be derived from simulation. Thus, the training system 200 may be used to train the neural network 100 to automatically synthesize a protein with a particular function such as having a binding site shape, e.g., a ligand that binds with sufficient affinity for a biological effect that it can be used as a drug. For example, e g., it may be an agonist or antagonist of a receptor or enzyme; or it may be an antibody configured to bind to an antibody target such as a virus coat protein, or a protein expressed on a cancer cell, e.g., to act as an agonist for a particular receptor or to prevent binding of another ligand and hence prevent activation of a relevant biological pathway.

[0204] In a similar way, the environment 230 may be a drug design environment, such that each state 231 of the environment 230 is a respective state of a potential pharmaceutically active compound, i.e., a drug, and the agent is a computer system for determining elements of the pharmaceutically active compound and / or a synthetic pathway for the pharmaceutically active compound. The drug / synthesis may be designed based on a reward derived from a target for the pharmaceutically active compound, for example in simulation. The agent 235 may be, or may include, a mechanical agent that performs or controls synthesis of the pharmaceutically active compound; and hence a process as described herein may include making such a pharmaceutically active compound.

[0205] For example, the environment 230 may be an in-silico drug design environment, e.g., a molecular docking environment, and the agent 235 may be a computer system for determining elements or a chemical structure of the drug. The drug may be a small molecule or biologic drug. An observation 232 of a state 231 of the environment 230 may be an observation of a simulated combination of the drug and a target of the drug. An action 234 performed by the agent 235 may be an action to modify the relative position, pose or conformation of the drug and drug target (or this may be performed automatically) and / or an action to modify a chemical composition of the drug and / or to select a candidate drug from a library of candidates. The rewards 236 may be defined based on one or more of: a measure of an interaction between the drug and the drug target, e.g., of a fit or binding between the drug and the drug target; an estimated potency of the drug; an estimatedselectivity of the drug; an estimated toxicity of the drug; an estimated pharmacokinetic characteristic of the drug; an estimated bioavailability of the drug; an estimated ease of synthesis of the drug; and one or more fundamental chemical properties of the drug. A measure of interaction between the drug and drug target may depend on, e.g., a protein-ligand bonding, van der Waal interactions, electrostatic interactions, and / or a contact surface region or energy; it may include, e.g., a docking score. Following identification of elements or a chemical structure of a drug in simulation, a method involving the training system 200 training the neural network 100 to control the agent 235 may further include making the drug. The drug may be made partly or completely by an automatic chemical synthesis system.

[0206] In some applications, the agent 235 may be a software agent i.e., a computer program, configured to perform a task. For example, the environment 230 may be a circuit, or an integrated circuit design, or routing environment, and the agent 235 may be configured to perform a design or routing task for routing interconnection lines of a circuit or of an integrated circuit e.g., an ASIC. The rewards 236 may then be dependent on one or more routing metrics such as interconnect length, resistance, capacitance, impedance, loss, speed, or propagation delay; and / or physical line parameters such as width, thickness or geometry, and design rules. The rewards 236 may also or instead include one or more rewards relating to a global property of the routed circuitry e.g., component density, operating speed, power consumption, material usage, a cooling requirement, level of electromagnetic emissions, and so forth. The observations 232 may be, e.g., observations of component positions and interconnections; the actions 234 may include component placing actions, e.g., to define a component position or orientation and / or interconnect routing actions, e.g., interconnect selection and / or placement actions. The task may be, e.g., to optimize circuit operation to reduce electrical losses, local or external interference, or heat generation, or to increase operating speed, or to minimize or optimize usage of available circuit area. A method involving the training system 200 training the neural network 100 to control the agent 235 may further include making the circuit or integrated circuit to the design, or with interconnection lines routed as determined by the method.

[0207] In some applications, the agent 235 is a software agent, and the environment 230 is a real- world computing environment. In one example, the agent 235 manages distribution of tasks across computing resources e.g., on a mobile device and / or in a data center. In these applications, the observations 232 of the states 231 of the environment 230 may include observations of computingresources such as compute and / or memory capacity, or Internet-accessible resources; and the actions 234 performed by the agent 235 may include assigning tasks to particular computing resources. The rewards 236 may be configured to maximize or minimize one or more of: utilization of computing resources, electrical power, bandwidth, and computation speed.

[0208] In another example, the software agent manages the processing, e.g., via one or more real- world servers, of a queue of continuously arriving jobs. The observations 232 may include observations of the times of departures of successive jobs, or the time intervals between the departures of successive jobs, or the time a server takes to process each job, e.g., the start and end of a range of times, or the arrival times, or time intervals between the arrivals, of successive jobs, or data characterizing the type of job(s). The actions 234 may include actions that allocate particular jobs to particular computing resources; the reward(s) may be configured to minimize an overall queueing or processing time or the queueing or processing time for one or more individual jobs, or in general to optimize any metric based on the observations.

[0209] As another example, the environment 230 may include a real -world computer system or network, the observations 232 may include any observations characterizing operation of the computer system or network, the actions 234 performed by the software agent may include actions to control the operation e.g. to limit or correct abnormal or undesired operation e.g. because of the presence of a virus or other security breach, and the rewards 236 may include any metric(s) that characterizing desired operation of the real-world computer system or network.

[0210] In some applications, the environment 230 is a real -world computing environment, and the software agent manages distribution of tasks / jobs across computing resources e.g., on a mobile device and / or in a data center. In these implementations, the observations 232 may include observations that relate to the operation of the computing resources in processing the tasks / jobs, the actions may include assigning tasks / jobs to particular computing resources, and the rewards 236 may relate to one or more metrics of processing the tasks / jobs using the computing resources, e.g., metrics of usage of computational resources, bandwidth, or electrical power, or metrics of processing time, or numerical accuracy, or one or more metrics that relate to a desired load balancing between the computing resources.

[0211] In some applications, the environment 230 is a data packet communications network environment, and the agent 235 is part of a router to route packets of data over the communications network. The actions 234 performed by the agent 235 may include data packet routing actions, andthe observations 232 of the states 231 of the environment 230 may include, e.g., observations of a routing table which includes routing metrics such as a metric of routing path length, bandwidth, load, hop count, path cost, delay, maximum transmission unit (“MTU”), and reliability. The reward 236 may be defined in relation to one or more of the routing metrics, e.g., configured to maximize one or more of the routing metrics.

[0212] In some other applications, the environment 230 is an Internet or mobile communications environment, and the agent 235 is a software agent which manages a personalized recommendation for a user. The observations 232 of the states 231 of the environment 230 may include previous actions taken by the user, e.g., features characterizing these; the actions may include actions recommending items such as content items to a user. The rewards 236 may be configured to maximize one or more of an estimated likelihood that the user will respond favorably to being recommended the (content) item, a suitability unsuitability of one or more recommended items, a cost of the recommended item(s), and a number of recommendations received by the user, optionally within a time span.

[0213] As a further example, the actions 234 performed by the agent 235 may include presenting advertisements, the observations 232 of the states 231 of the environment 230 may include advertisement impressions or a click-through count or rate, and the rewards 236 may characterize previous selections of items or content taken by one or more users.

[0214] In some cases, the observations 232 may include textual or spoken instructions provided to the agent 235 by a third-party (e.g., an operator of the agent). For example, the agent 235 may be an autonomous vehicle, and a user of the autonomous vehicle may provide textual or spoken instructions to the agent 235 (e.g., to navigate to a particular location).

[0215] As another example, the environment 230 may be an electrical, mechanical, or electromechanical design environment, e.g., an environment in which the design of an electrical, mechanical, or electro-mechanical entity is simulated. The environment 230 may be implemented in a simulation 240 of a real-world environment in which the entity is intended to work. The task may be to design the entity. The observations 232 may include observations that characterize the entity, i.e., observations of a mechanical shape or of an electrical, mechanical, or electromechanical configuration of the entity, or observations of parameters or properties of the entity. The actions 234 may include actions that modify the entity, e.g., that modify one or more of the observations. The rewards 236 may include one or more metrics of performance of the design ofthe entity. For example, the rewards 236 may relate to one or more physical characteristics of the entity such as weight or strength or to one or more electrical characteristics of the entity, such as a measure of efficiency at performing a particular function for which the entity is designed. The design process may include outputting the design for manufacture, e.g., in the form of computer executable instructions for manufacturing the entity. The process may include making the entity according to the design. Thus, the design of an entity may be optimized, e.g., by reinforcement learning, and then the optimized design output for manufacturing the entity, e.g., as computer executable instructions; an entity with the optimized design may then be manufactured.

[0216] As previously described, the environment 230 may be a simulated environment implemented in a simulation 240. Generally, in the case of a simulated environment, the observations 232 may include simulated versions of one or more of the previously described observations or types of observations, and the actions 234 may include simulated versions of one or more of the previously described actions or types of actions. For example, the simulated environment may be a motion simulation environment, e.g., a driving simulation or a flight simulation, and the agent may be a simulated vehicle navigating through the motion simulation. In these implementations, the actions 234 may be control inputs to control the simulated user or simulated vehicle. Generally, the agent 235 may be implemented as one or more computers interacting with the simulated environment.

[0217] The simulated environment may be implemented in a simulation 240 of a particular real- world environment 230 and agent 235. For example, the neural network 100 may be used to select actions 234 in the simulated environment during training or evaluation of the neural network 100 and, after training, or evaluation, or both, are complete, may be deployed for controlling a real- world agent in the particular real-world environment that was the subject of the simulation 240. This can avoid unnecessary wear and tear on and damage to the real-world environment or real- world agent and can allow the control neural network to be trained and evaluated on situations that occur rarely or are difficult or unsafe to re-create in the real-world environment. For example, the training system 200 can partly train the neural network 100 using a simulation 240 of a mechanical agent in a particular real -world environment, and afterwards deploy the neural network 100 to control the real mechanical agent in the particular real-world environment. Thus, in such cases, the observations 232 of the simulated environment relate to the real-world environment, and theactions 234 selected in the simulated environment relate to the actions to be performed by the mechanical agent in the real-world environment.

[0218] In some implementations, the observations 232 of the environment 230 are observations of a real-world environment in which a human is performing a task, e.g., an image observation from an image sensor and / or a language observation from a speech recognition system; and the actions 234 performed by the agent 235 are language actions that control (instruct) the human, e.g., using natural language or images, to perform actions in the real-world environment to perform the task. A language action may be an action that outputs a natural language sentence, e.g., by defining a sequence of language tokens, e.g., words or wordpieces, to be emitted at sequential timesteps.

[0219] Thus, the agent 235 may include a user interface device such as a digital device (a “digital assistant”), e.g., a smart speaker or smart display or other device, e.g., with a natural language input and / or output, that controls (instructs) a human user to perform a task. In general, such a digital device can be a mobile device with a natural language interface to receive natural language requests from a human user and to provide natural language responses. It may also include a visionbased input e.g., a camera and / or display screen. The digital device may include a language model or language generation neural network system either stored locally, or accessed remotely, or both. The user interface device may include, e.g., a mobile device, a keyboard (and optionally display), or a speech-based input mechanism, e.g., to input audio data characterizing a speech waveform of speech representing the input from the user in the natural or computer language, and to convert the audio data into tokens representing the speech in the natural or computer language, e.g., representing a transcription of the spoken input. The user interface device can also include a text or speech-based output, e.g., a display and / or a text-to- speech subsystem.

[0220] Thus, in some implementations, the actions 234 performed by the agent 235 contribute to performing the task. A monitoring system, e.g., a video camera system, may be provided for monitoring the action (if any) which the user actually performs at each timestep in case, e.g., due to human error, it is different from the action 234 which the training system 200 instructed the user to perform. The monitoring system can be used to determine whether the task has been completed. Training data may be collected to record the actions which the user actually performed based on the instruction. The reward 236 for an action 234 may be generated, for example, by comparing the action the user took with a corpus of data showing a human expert performing the task, e.g., using techniques known from imitation learning, or in some other way, e.g., using a trained rewardmodel. The training system 200 configured in this way can learn how to guide a human to perform a task, e.g., avoiding difficult to perform actions.

[0221] Optionally, in any of the above implementations, the observation 232 at any given timestep may include data from a previous timestep that may be beneficial in characterizing the state 231 of the environment 230 at the given timestep, e.g., the action 234 performed by the agent 235 at the previous timestep, the reward 236 received by the agent 235 at the previous timestep, or both.VI. Experiments

[0222] This utility of NaP was validated empirically in several experiments, the results of which are presented in the Sec. VI. In general, the goal of the experiments was to validate two properties: (i) NaP maintains or improves performance of a neural network on stationary tasks; and (ii) NaP mitigates plasticity loss and improves performance of the neural network on a variety of both synthetic and natural nonstationaiy tasks.A. Effective Learning Rate

[0223] NaP constrains a neural network’s effective learning rate (“ELR”) to follow an explicit schedule, rather than an implicit schedule. In Sec. VI. A, the inventors explored how this property affects network training dynamics, demonstrating how implicit learning rate schedules due to parameter-wise norm growth can be made explicit and leveraged to improve the performance of NaP in deep reinforcement learning domains.

[0224] The inventors began by considering a convolutional neural network (“CNN”) and training it on random labels of the CIFAR-10 dataset. Details of the neural network architecture of the CNN are provided in Sec.VII.D. The inventors then re-randomized these labels and continued training, repeating this process five hundred times. FIGs. 5A, 5B, and 5C are experimental plots of an example of the CNN, the example of the CNN with layer-wise normalization, and the example of the CNN with layer-wise normalization and weight projection, respectively. When applied to a network without normalization layers, the Jacobian norm grows to unstable values as the paramet er- wise norm increases; in contrast, an equivalent architecture with normalization layers sees a sharp decline in the Jacobian norm as the parameter- wise norm increases. In both cases, the result was similar - increased parameter- wise norm accompanies reduced performance.However, as shown in FIGs. 5A-5C, re-normalizing the network parameters post task maintained the performance of a random initialization even after hundreds of training iterations.

[0225] The inventors continued their study of effective learning rates by illustrating how the implicit learning rate schedule induced by the evolution of the paramet er- wise norm can be translated to an explicit schedule in NaP. The inventors used the CNN with layer-wise normalization trained on CIFAR-10. The inventors trained two ‘twin’ scale-invariant networks with the Adam optimizer in tandem. Both networks processed the same data stream and started from the same initialization, but the per-layer weights of one were projected after every training iteration to have constant norm, while the other was allowed to vary the norms of the weights.

[0226] FIGs. 6A-6D are experimental plots of three coupled networks experimental settings that were considered. FIG. 6A shows the loss of a layer-wise rescaling experiment, where the inventors re-scaled the per-layer learning rates of the projected network so that the explicit learning rate was equal to the effective learning rate of its twin. FIG. 6B shows the loss in a global rescaling experiment, where the inventors re-scaled the global learning rate based on the ratio of parameterwise norms between the projected and unprojected network, but did not tune on per-layer bases. FIG. 6C shows the loss of a no-rescaling experiment, where the inventors did not implement learning rate re-scaling. FIG. 6D shows the learning curves of the experiments together. As shown in FIGs. 6A-6D, the shapes of the learning curves for all networks except for the constant-ELR variant were quite similar, with the global learning rate scaling strategy producing a smaller gap than the no-rescaling strategy. By construction, the dynamics of the per-layer rescaling network and its twin were identical. Since global learning rate schedules induce dynamics that are quite close to those obtained by parameter-wise norm growth, this approach was utilized for the remainder of the experiments.B. Nonstationarity

[0227] The inventors began their experiments on nonstationarity with a continual classification problem described in Sec. VII. D. The inventors evaluated their approach on several sources of nonstationarity using two example architectures: a CNN, and a fully-connected multi-layer perceptron (“MLP”). They first evaluated several methods designed to maintain plasticity including Regenerative regularization, Shrink and Perturb, ReDo, leaky ReLU units, L2 regularization, and random Gaussian perturbations to the optimizer update (a heuristic form ofLangevin Dynamics). They tracked the average online accuracy over the course of training for twenty million steps, equivalent to two hundred data relabelings, using a constant learning rate. The inventors found varying degrees of efficacy in these approaches in the base network architectures, with regenerative regularization and ReDO tending to perform the best. When they applied the same suite of methods to networks with NaP, shown in FIGs. 7A-7D, near-monotonic improvements in performance were observed. FIGs. 7A, 7B, 7C, and 7D are experiment plots of label memorization robustness for the CNN baseline, the MLP baseline, the CNN trained with NaP, and the MLP trained with NaP, respectively. As shown in FIGs. 7A-7D, a significant reduction in the gaps between methods was observed using NaP, with the performance curves of the different methods nearly indistinguishable for the MLP. Further, the inventors observed constant or increasing slopes in the online accuracy, suggesting that the difference between methods had more to do with their effect on within-task performance than on plasticity loss once the parameter and layer norms had been constrained.C. Stationary Benchmarks

[0228] Having observed improvements in synthetic tasks, the inventors went on to confirm that NaP did not interfere with learning on more widely-studied, natural datasets.Table 1: Top-1 prediction accuracy on the Table 2: Per-token accuracy of a 400M test sets of CIFAR-10 and ImageNet-lk. transfonner model pretrained on the C4 dataset.

[0229] Large-scale image classification. The inventors began by studying the effect of NaP on two well-established benchmarks: a VGG16-like network on CIFAR-10, and a ResNet-50 on the ImageNet-lk dataset. Full details are provided in Sec. VII. D. Table 1 reports the results of the experiments on each network. Comparable performance was achieved in both cases using the same learning rate schedule as the baseline.

[0230] Natural language. The inventors then turned to language tasks, training a four hundred million (“400M”) parameter transformer model on the C4 dataset. Full details are provided in Sec. VII. C. Table 2 reports the results of the experiment, showing that NaP did not interfere with performance on this task, where final performance was matched in terms of training accuracy. When evaluating the pre-trained network on a variety of other datasets, the inventors found that NaP slightly outperformed baselines in terms of performance on a variety of benchmarks, including WikiText-103, Lambada, PIQA, SociallQA, and Pile.D. Deep Reinforcement Learning

[0231] Lastly, the inventors evaluated NaP on a setting where maintaining plasticity can impact performance: reinforcement learning (“RL”) on the Arcade Learning Environment (“ALE”). The inventors conducted a full sweep over fifty-seven Atari 2600 games, comparing the effects of layer-wise normalization, weight projection, and learning rate schedules on a Rainbow agent.

[0232] The inventors also analyzed sequential ALE training. In this case, the inventors considered a setup where optimizer state and schedule were reset every time the environment changed, using a cosine schedule with warmup. The details are provided in Sec. VII.B. To evaluate NaP in this regime, the agents were trained on ten games each for 20M frames, going through this cycle twice. Parameters of the continual agents were not reset between games; however, the optimizer was rest. The learning curves for the second round of games are shown in FIGs. 9A-9J, demonstrating that NaP significantly outperformed a baseline Rainbow agent with and without layer-wise normalization. Indeed, even after 200M steps the networks trained with NaP made similar learning progress to a random initialization.

[0233] FIGs. 8A, 8B, 8C, 8D, 8E, 8F, 8G, 8H, 81, and 8J are experimental plots of the learning curves of continual Atari agents on twelve Atari games, including alien, atlantis, boxing, breakout, centipede, double_dunk, freeway, pong, space_invaders, and tennis, respectively. The learning curves are visualized on the sequential ALE training, i.e., 200M frames. Each game was played for 20M frames, and agents passed sequentially from one to another, repeating all ten games twice for a total of 400M frames. Solid lines indicate performance on the second visit to each game, and dotted lines indicate performance of a randomly initialized network on the game. Even in its second visit to each game, NaP performed comparably to the randomly initialized networks, whereas the standard Rainbow agent exhibited poor performance on all games in the sequential training regime.

[0234] FIGs. 9A and 9B are experimental plots depicting aggerate effects of normalization on single-task Atari in terms of the Interquartile Mean (“IQM”) and medium, respectively. Bars indicate 95 % confidence intervals over four seeds and fifty-seven environments. FIGs. 9A-9B show the spread of scores, along with estimates of the Mean and IQM of four agents: standard Rainbow, Rainbow + LayerNorm, Rainbow + NaP without a learning rate schedule, and Rainbow + NaP with a learning rate schedule. The inventors found that NaP with a linear learning rate schedule outperformed the other methods.VII. Experiment DetailsA. Toy Experiment Details

[0235] The inventors conducted a variety of illustrative experiments on toy problem settings and small networks. In FIGs. 5A-5D, they used a DQN-style CNN network which included two sets of two convolutional layers with 32 and 64 channels, respectively. The inventors then applied max pooling and flattened the output, feeding through a 512-unit hidden linear layer before applying a final linear transform to obtain the output logits. The network used ReLU nonlinearities. When NaP was applied, layer-wise normalization was added prior to each nonlinearity.B. Reinforcement Learning Details

[0236] Single-task Atari: the reinforcement learning experiments were based off the publicly available implementation of the Rainbow agent in DQN Zoo. The inventors followed the default hyperparameters detailed in this codebase. In the NaP implementation, the inventors added normalization layers prior to each nonlinearity except for the final softmax. The agents were trained for 200M frames on the Atari 57 suite. The reinforcement learning experiments allowed a learning rate schedule, which was explicitly detailed in cases where non-constant learning rates were used.

[0237] Sequential ALE: the same Rainbow agent for the single-task experiments was used for the sequence ALE experiments, where a cosine decay learning rate was employed for all variants. The cosine decay schedule was restarted at every task change for all agents. The cosine decay schedule used an initial value of 10x, a peak value of the default learning rate for Rainbow (0.000625), 1,000 warmup steps after the optimizer is reset, and end-value equal to 106as in the single-task settings. Cosine decay was chosen due to its popularity in supervised learning, and tohighlight the versatility of NaP on different learning rate schedules. The inventors followed the game sequence as described by Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marios C. Machado, “Loss of plasticity in continual deep reinforcement learning,” In Conference on Lifelong Learning Agents, pp. 620-636, PMLR (2023). The agents were trained for 20M frames per game in the game sequence.C. Language Details

[0238] Sequence memorization: the inventors set a dataset size of 1024 and a sequence length of 512. They used a vocabulary size of 256, equivalent to ASCII tokenization. The Adam optimizer was implemented, and all networks were trained for a minimum of 10,000 iterations. The dataset was reset every 1,000 iterations, generating a new set of 1024 random strings of length 512. The inventors used as a baseline transformer model including 4 attention blocks, with 8 heads and a model dimension equal to 256. A batch size of 128 was employed.

[0239] In-context learning: the in-context learning experiments used the same overall setup as the sequence memorization experiments, with identical architectures and baseline optimization algorithm. In this case, the inventors trained on a dataset including 4096 randomly generated strings, in which the final one hundred tokens were a contiguous subsequence of the first four hundred and twelve, selected uniformly at random from indices in [1,312],

[0240] Natural language: the natural language experiments were run on a 400M parameter transformer model based on the same backbone as the previous two tasks, this time including 12 attention blocks with 12 heads and a model dimension 1536. Learning rate warmup followed by cosine decay was employed, with a peak learning rate of 2 ■ I O4, which was reached after a linear warmup of 1,000 steps. A batch size of 128 was used and training was performed for 30,000 steps. A weight decay parameter of 0.1 was used with the Adam optimizer.D. Vision Details

[0241] CIFAR-10 Memorization: the inventors considered two classes of neural network architectures: (i) a fully-connected MLP with a width of 512 and depth of 4; (ii) a CNN with k convolutional layers followed by two fully-connected layers, with a depth of 4, 32 channels, and fully-connected hidden layer width of 256.

[0242] The continual supervised learning domain was constructed from the CTFAR-10 dataset, generating a family of continual classification problems. Each continual classification problem was characterized by a transformation function on the labels. For label transformations, the classes were permuted (e.g., all images with the label 5 were re-assigned the label 2), and random label assignment, where each input was uniformly at random assigned a new label independent of its class in the underlying classification dataset.

[0243] The inventors ran a total of 20M steps and a total of 200 random target resets. All networks were trained using the Adam optimizer. A sweep over learning rates for the different architectures was performed, settling on 104as this provided a reasonable balance between convergence speed and stability in all architectures, to ensure that all networks could at least solve the single-task version of each label and target transformation.

[0244] VGG-Net and ResNet-50 baselines: the inventors used the standard data augmentation policies for the CIFAR-10 and ImageNet-lk experiments. In the ImageNet-lk experiments described herein, the inventors used the ResNetv2 architecture variant, with a label smoothing parameter of 0.1, and a weight decay of 104. For an optimizer a SGD with a cosine annealing learning rate schedule was employed, and Nesterov momentum with a decay rate of 0.9. A batch size of 256 was used. The CIFAR-10 experiments used a VGG-Net architecture. Here, the inventors used the SGD optimizer with a batch size of 32, and set a learning rate schedule which started at 0.025 and decayed by a factor of 0.1 iteratively through training. Nesterov momentum with decay rate 0.9 was used and training was performed for a total of 400,000 steps.VIII. Conclusion

[0245] This specification introduced the NaP training algorithm which can substantially mitigate the loss of plasticity of a neural network by constraining the norm of the neural network’ s network parameters and placing normalization layers before each nonlinearity in the neural network. Implementations of the NaP training algorithm can maintain constant pre-activation norms of the neural network, protect against saturated nonlinearities in the neural network, and make the effective learning rate of the neural network explicit in the training process. NaP was leveraged in experiments to demonstrate its efficacy controlling learning rate decay schedules and improving performance of neural networks in deep learning settings. Particularly, NaP was shown to improveboth the performance and robustness of neural network to nonstationarity in deep reinforcement and continual learning settings.SCOPE OF THE DISCLOSURE

[0246] This 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.

[0247] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine- readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively, or in addition, the program instructions can be encoded on an artificially-generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0248] The term “data processing apparatus” refers to data processing hardware and encompasses 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 also 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.

[0249] 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. 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 interconnected by a data communication network.

[0250] 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, 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 or computers.

[0251] 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.

[0252] 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 randomaccess 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.

[0253] 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.

[0254] 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 that is used by 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 from the user in return.

[0255] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and computeintensive parts of machine learning training or production, i.e., inference, workloads.

[0256] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework.

[0257] 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.

[0258] 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., an HTML page, 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 the server from the device.OTHER EMBODIMENTS

[0259] In addition to the embodiments of the attached claims and the embodiments described above, the following numbered embodiments are also innovative.

[0260] 1. A method performed by one or more computers, the method comprising: obtaining a training dataset for training a neural network parametrized by a set of network parameters, wherein the neural network comprises: one or more nonlinear layers that each apply a nonlinear transformation over a respective normalized layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer; and training the neural network on the training dataset over a plurality of training iterations, comprising, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterionis satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

[0261] 2. The method of embodiment 1, wherein at each training iteration, determining whether the criterion is satisfied at the training iteration comprises: determining whether a threshold number of training iterations has elapsed since the network parameters were last normalized; or determining whether a norm over the updated values of the network parameters is greater than a threshold norm.

[0262] 3. The method of any one of embodiments 1-2, wherein at each training iteration: the subset of training data comprises a network input, and training the neural network on the subset of training data to generate the updated values of the network parameters comprises: processing the network input using the neural network, in accordance with the current values of the network parameters, to generate a network output for the network input; computing, with respect to the network parameters, a gradient of an objective function when evaluated on the network output; and augmenting, using the gradient of the objective function, the current values of the network parameters to generate the updated values of the network parameters.

[0263] 4. The method of embodiment 3, wherein at each training iteration: the subset of training data further comprises a target network output for the network input, and the objective function evaluates an error between: (i) the network output, and (ii) the target network output, for the network input.

[0264] 5. The method of any one of embodiments 1-4, wherein the training dataset was generated using the neural network.

[0265] 6. The method of any one of embodiments 1-4, wherein the training dataset is a data stream comprising a respective subset of training data at each of a plurality of timesteps, and each training iteration is a respective timestep in the data stream.

[0266] 7. The method of any one of embodiments 1-6, wherein: the neural network comprises one or more additional layers each parametrized by a respective subset of the network parameters, and at each training iteration, normalizing the updated values of each of the network parameters comprises, for each additional layer: computing a norm over the updated values of the network parameters in the respective subset; and for each network parameter in the subset: normalizing the updated value of the network parameter with the norm to generate a normalized version of theupdated value; and scaling the normalized version of the updated value with a respective reference norm for the additional layer.

[0267] 8. The method of embodiment 7, further comprising, at a first of the training iterations: computing, as the respective reference norm for each additional layer, a norm over the current values of the network parameters in the respective subset.

[0268] 9. The method of any one of embodiments 7-8, wherein each normalization layer is parametrized by a respective subset of the network parameters, and the subset comprises: (i) a respective scaling vector, and (ii) a corresponding offset vector.

[0269] 10. The method of embodiment 9, wherein for each normalization layer, applying the normalization operation over the respective layer input to the normalization layer comprises: computing a mean over a plurality of elements of the layer input; shifting each of the plurality of elements of the layer input by the mean to generate a shifted version of the layer input; computing a norm of the shifted version of the layer input; normalizing the shifted version of the layer input with the norm to generate a normalized version of the layer input; computing an elementwise product between: (i) the respective scaling vector, and (ii) the normalized version of the layer input; and offsetting the elementwise product with the respective offset parameter.

[0270] 11. The method of any one of embodiments 9-10, wherein at each training iteration, normalizing the updated values of the network parameters comprises, for each normalization layer: computing a norm over the updated values of the scaling and offset vectors in the respective subset; normalizing the updated values of the scaling and offset vectors with the norm to generate normalized versions of the updated values; and scaling the normalized versions of the updated values with a respective reference norm for the normalization layer.

[0271] 12. The method of any one of embodiments 1-11, wherein each normalization layer is a layer normalization layer or a batch normalization layer.

[0272] 13. The method of any one of embodiments 1-12, wherein each nonlinear layer is a Rectified Linear Unit (ReLU), a Gaussian Error Linear Unit (GELU), a hyperbolic tangent activation, a sigmoid activation, or a softmax activation.

[0273] 14. The method of any one of embodiments 1-3, further comprising: obtaining the neural network comprising the one or more nonlinear layers; and for each of the one or more nonlinear layers: determining whether a normalization layer immediately precedes the nonlinearlayer; and when a normalization layer does not immediately precede the nonlinear layer, inserting the respective normalization layer into the neural network.

[0274] 15. The method of any one of embodiments 1-14, wherein the neural network is configured to process a network input comprising one or more of: an image, an audio waveform, or a text sequence.

[0275] 16. The method of any one of embodiments 1-14, wherein the neural network is configured to process a network input comprising an observation characterizing a state of an environment being interacted with by an agent, and generate a network output defining an action to be performed by the agent in response to the observation.

[0276] 17. The method of any one of embodiments 1-16, further comprising, after training the neural network on the training dataset over the plurality of training iterations: parametrizing the neural network with the trained values of the network parameters; receiving a network input; and processing the network input using the neural network, in accordance with the trained values of the network parameters, to generate a network output for the network input.

[0277] 18. A method performed by one or more computers, the method comprising: receiving a network input; and processing the network input, using a neural network, to generate a network output for the network input, wherein the neural network has been trained via the respective method of any one of embodiments 1-17.

[0278] 19. A method performed by one or more computers, the method comprising: receiving a network input; and processing the network input, using a neural network parametrized by a set of network parameters, to generate a network output for the network input, wherein the neural network comprises: one or more nonlinear layers that each apply a nonlinear transformation over a respective layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer, and wherein the neural network has been trained on a training dataset over a plurality of training iterations, comprising, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterionis satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

[0279] 20. One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the respective method of any one of embodiments 1-19.

[0280] 21. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of embodiments 1-19.

[0281] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination 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 may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0282] 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 may 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.

[0283] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can beperformed 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 shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

[0284] What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers, the method comprising: obtaining a training dataset for training a neural network parametrized by a set of network parameters, wherein the neural network comprises: one or more nonlinear layers that each apply a nonlinear transformation over a respective layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer; and training the neural network on the training dataset over a plurality of training iterations, comprising, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterion is satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

2. The method of claim 1 , wherein at each training iteration, determining whether the criterion is satisfied at the training iteration comprises: determining whether a threshold number of training iterations has elapsed since the network parameters were last normalized; or determining whether a norm over the updated values of the network parameters is greater than a threshold norm.

3. The method of any preceding claim, wherein at each training iteration: the subset of training data comprises a network input, and training the neural network on the subset of training data to generate the updated values of the network parameters comprises:processing the network input using the neural network, in accordance with the current values of the network parameters, to generate a network output for the network input; computing, with respect to the network parameters, a gradient of an objective function when evaluated on the network output; and augmenting, using the gradient of the objective function, the current values of the network parameters to generate the updated values of the network parameters.

4. The method of claim 3, wherein at each training iteration: the subset of training data further comprises a target network output for the network input, and the objective function evaluates an error between: (i) the network output, and (ii) the target network output, for the network input.

5. The method of any preceding claim, wherein the training dataset was generated using the neural network.

6. The method of any of claims 1-4, wherein the training dataset is a data stream comprising a respective subset of training data at each of a plurality of timesteps, and each training iteration is a respective timestep in the data stream.

7. The method of claim 1, wherein: the neural network comprises one or more additional layers each parametrized by a respective subset of the network parameters, and at each training iteration, normalizing the updated values of each of the network parameters comprises, for each additional layer: computing a norm over the updated values of the network parameters in the respective subset; and for each network parameter in the subset: normalizing the updated value of the network parameter with the norm to generate a normalized version of the updated value; andscaling the normalized version of the updated value with a respective reference norm for the additional layer.

8. The method of claim 7, further comprising, at a first of the training iterations: computing, as the respective reference norm for each additional layer, a norm over the current values of the network parameters in the respective subset.

9. The method of any of claims 7-8, wherein each normalization layer is parameterized by a respective subset of the network parameters, and the subset comprises: (i) a respective scaling vector, and (ii) a corresponding offset vector.

10. The method of claim 9, wherein for each normalization layer, applying the normalization operation over the respective layer input to the normalization layer comprises: computing a mean over a plurality of elements of the layer input; shifting each of the plurality of elements of the layer input by the mean to generate a shifted version of the layer input; computing a norm of the shifted version of the layer input; normalizing the shifted version of the layer input with the norm to generate a normalized version of the layer input; computing an elementwise product between: (i) the respective scaling vector, and (ii) the normalized version of the layer input; and offsetting the elementwise product with the respective offset parameter.

11. The method of any of claims 9-10, wherein at each training iteration, normalizing the updated values of the network parameters comprises, for each normalization layer: computing a norm over the updated values of the scaling and offset vectors in the respective subset; normalizing the updated values of the scaling and offset vectors with the norm to generate normalized versions of the updated values; and scaling the normalized versions of the updated values with a respective reference norm for the normalization layer.

12. The method of any preceding claim, wherein each normalization layer is a layer normalization layer or a batch normalization layer.

13. The method of any preceding claim, wherein each nonlinear layer is a Rectified Linear Unit (ReLU), a Gaussian Error Linear Unit (GELU), a hyperbolic tangent activation, a sigmoid activation, or a softmax activation.

14. The method of any preceding claim, further comprising: obtaining the neural network comprising the one or more nonlinear layers; and for each of the one or more nonlinear layers: determining whether a normalization layer immediately precedes the nonlinear layer; and when a normalization layer does not immediately precede the nonlinear layer, inserting the respective normalization layer into the neural network.

15. The method of any preceding claim, wherein the neural network is configured to process a network input comprising one or more of: an image, an audio waveform, or a text sequence.

16. The method of any of claims 1-14, wherein the neural network is configured to process a network input comprising an observation characterizing a state of an environment being interacted with by an agent and generate a network output defining an action to be performed by the agent in response to the observation.

17. The method of any preceding claim, further comprising, after training the network parameters on the training dataset over the plurality of training iterations: parametrizing the neural network with the trained values of the network parameters; receiving a network input; and processing the network input using the neural network, in accordance with the trained values of the network parameters, to generate a network output for the network input.

18. A method performed by one or more computers, the method comprising:receiving a network input; and processing the network input, using a neural network, to generate a network output for the network input, wherein the neural network has been trained via the respective method of any preceding claim.

19. A method performed by one or more computers, the method comprising: receiving a network input; and processing the network input, using a neural network parametrized by a set of network parameters, to generate a network output for the network input, wherein the neural network comprises: one or more nonlinear layers that each apply a nonlinear transformation over a respective layer input to the nonlinear layer; and for each nonlinear layer, a respective normalization layer that: (i) immediately precedes the nonlinear layer, and (ii) applies a normalization operation over a respective layer input to the normalization layer, and wherein the neural network has been trained on a training dataset over a plurality of training iterations, comprising, at each training iteration: obtaining current values of the network parameters; parametrizing the neural network with the current values of the network parameters; sampling, from the training dataset, a subset of training data; training the neural network on the subset of training data to generate updated values of the network parameters; determining whether a criterion is satisfied at the training iteration; and when the criterion is satisfied at the training iteration, normalizing the updated value of each of the network parameters.

20. One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the respective method of any of claims 1-19.

21. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any of claims 1-19.

Images