Normalization and effective learning rates in reinforcement learning
Authors: Clare Lyle
Zeyu Zheng
Khimya Khetarpal
James Martens
Hado van Hasselt
Razvan Pascanu
Will Dabney
Venue: NeurIPS 2024 poster (2024)
Keywords: continual learning
reinforcement learning
optimization
plasticity
Primary area: reinforcement_learning
TL;DR: We propose a method to stabilize training in nonstationary tasks, improving performance on several benchmarks and obtaining insight into the role of implicit effective learning rate schedules in RL.

Abstract: Normalization layers have recently experienced a renaissance in the deep reinforcement learning and continual learning literature, with several works highlighting diverse benefits such as improving loss landscape conditioning and combatting overestimation bias. However, normalization brings with it a subtle but important side effect: an equivalence between growth in the norm of the network parameters and decay in the effective learning rate. This becomes problematic in continual learning settings, where the resulting learning rate schedule may decay to near zero too quickly relative to the timescale of the learning problem. We propose to make the learning rate schedule explicit with a simple re-parameterization which we call  Normalize-and-Project (NaP), which couples the insertion of normalization layers with weight projection, ensuring that the effective learning rate remains constant throughout training. This technique reveals itself as a powerful analytical tool to better understand learning rate schedules in deep reinforcement learning, and as a means of improving robustness to nonstationarity in synthetic plasticity loss benchmarks along with both the single-task and sequential variants of the Arcade Learning Environment. We also show that our approach can be easily applied to popular architectures such as ResNets and transformers while recovering and in some cases even slightly improving the performance of the base model in common stationary benchmarks.

Normalization and effective learning rates in reinforcement learning
Clare Lyle
Zeyu Zheng
Khimya Khetarpal*
James Martens
Hado van Hasselt!
Razvan Pascanu
Will Dabney
Abstract
Layer normalization has demonstrated remarkable effectiveness preventing plasticity loss continual and reinforcement learning (RL), though the precise reasons for this effectiveness remain mysterious. In this work, identify new mechanisms by which layer normalization can help - and hinder training in neural networks, and leverage these insights to improve the robustness of gradient- based optimization algorithms to nonstationarity. Our analysis reveals surprising ability of layer normalization revive dormant ReLU units, along with under- appreciated vulnerability to unconstrained decay of the effective learning rate (ELR), which can drive loss of plasticity in long-running nonstationary tasks. Motivated by these findings, we propose Normalizz-ann-PPooeet (NaP), simple protocol designed provide the numerous benefits of normalization while ensuring that the effective learning rate remains constant throughout training. To do so, NaP couples the insertion of normalization layers with weight projection. This technique mitigates loss of plasticity in two challenging continual learning problems; a sequential supervised learning task, and continual variant of the Arcade Learning Environment. Further, by using NaP to explicitly control the effective learning rate in deep RL agents, we find that in fact the implicit ELR decay induced by parameter norm growth in these agents is critical to their ability to achieve competitive performance, suggesting the common practice of using constant learning rates deep RL may be far from optimal.
Introduction
Many of the most promising application areas of deep learning, in particular deep reinforcement learning (RL., require training on a problem which is in some way nonstationary. In order to perform well on a nonstationary problem, the neural network must maintain its ability adapt new information, i.c. it must remain plastic. Several recent works have shown that loss of plasticity can present major barrier to performance improvement in RL and in continual learning (Dohare et 2021 Lyle et al.) 2021 Nikishin et al. 2022]. These works have proposed variety of explanations for plasticity loss such as the accumulation of saturated ReLU units and increased sharpness of the loss landscape [Lyle et al. 2 along with mitigation strategies, such as resetting dead units (Sokar parameters towards their initial values [Kumar et al. 2023] Many of these explanations and their corresponding mitigation strategies center around reducing drift in the distribution pre-activations (Lyle et al.) 2024), a problem which has historically been resolved in the supervised learning setting, and more recently in continual learning and RI. [Hussing et al. 2024 Ball et al. 2023). by incorporating normalization layers into the network architecture.
ct al. 2023 and regularizing
GGoogle DeepMind, Correspondence clarelyye@google..om
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

While effective, normalization on its own is insufficient to avoid loss of plasticity [Lee et al. 2023]. Part of the reason for this lies in a subtle property of normalization: a normalization layer causes the subnetwork preceding it to become scale-invariany which means that the layer's effective learning rate (ELR) now depends on the norm of its parameters [Van Laarhoven 2017]. partieular, when the norm of the parameters grows, as it typically does in neural networks trained without regularization, the effective learning rate shrinks. Weight decay can address this problem, but runs the risk of either failing to fully mitigate the norm growth or over-regularizing the model to the point of slowing down learning and thus requires careful tuning, adding additional complexity to the training protocol.
The aim of this work is to provide principled strategy for the use of layer normalization in non- stationary learning problems. To do so, we first identify two critical properties of layer normalization: its ability to facilitate the revival of dormant neurons, and its vulnerability to vanishing gradients as the parameter norm grows. In partieular, we show that layer normalization buffers against dormant neurons due to its effect of mixing gradient signal between units. This analysis motivates Normalize- and-Project (NaP), a method which inserts normalization layers prior each nonlinearity in a network architecture and maintains constant per-layer parameter norm. A network trained with the NaP protocol stands to benefit from the increased robustness to unit saturation provided by LayerNorm's s gradient-mixing property, while also avoiding its vulnerability to vanishing gradients under growing parameter norms. Indeed, we observe empirically that NaP virtually eliminates loss of plasticity in multiple challenging non-stationary learning problems, including sequential training games from the Arcade Learning Environment (ALE) [Bellemare et al.] 2013 Abbas et al., 2023] We further confirm that NaP can be seamlessly integrated into standard computer vision and sequence modeling baselines from standard benchmarks without hindering performance.
We further leverage Normalize-and-Proie Pro-et an analytical tool to explain the dramatic performance degradations induced by weight decay in deep RL agents. By making the effective learning rate explicit, NaP reveals that the implicit learning rate decay induced by parameter norm growth in Rainbow [Hessel et al., 2018) agents is in fact critical to their performance in the ALE benchmark: certain components of the value function require a sufficiently small ELR in order to be learned, and an optimization process which does not reach this value will therefore underfit the value function in ways that can inhibit performance improvement. Indeed, while demonstrate that the implicit schedule induced by parameter norm growth outperforms a constant effective learning rate, an explicit schedule with more aggressive learning rate decay outperforms the implicit one on the ALE benchmark. These findings are at odds the common folk wisdom that the continual nature of RL makes it unsuitable for learning rate decay, and demonstrate the untapped potential of step size schedules to accelerate deep reinforcement learning.
2 Background and related work
We begin by providing background on trainability and its loss in nonstationary learning problems. We additionally give an overview of neural network training dynamics and effective learning rates.
2.1 Training dynamics and plasticity in neural networks
Early work on neural network initialization centered around the idea of controlling the norm of the activation vectors [LeCun et al.] 2002 Glorot and Bengio, 2010 He et al.) [015)] using informal arguments. More recently, this perspective has been formalized and expanded [Poole et al. 2016 Daniely et al. 2016 Martens et al., 2021] to include the inner- inner-products between pairs of activation vectors (for different inputs to the network). 'he function that describes the evolution of these inner-products determines the network's ' gradients at initialization up to rotation, and this in turn determines trainability (which was shown formally in the Neural Tangent Kernel regime by Xiao ct al. 2020) and Martens et al. 2021)) I). A variety of initialization methods have been developed to ensure the network avoids "shattering" [Balduzzi al. 2017 or collapsing gradients (Poole et al.
2016
Martens et al. 2021 Zhang et al. 2001b]]
Once training begins, learning dynamics can be ceaa-charecterizeed in the infinite-width limit by the neural tangent kernel and related quantities [Jacot et al. 2018 Yang 2019). although in practice optimization dynamics diverge significantly from this limit Forrt et al.) 2020]. A number of beneficial
phenomena emerge in the width, finite-step-size regime, such as the self-stabilizatioo of
gradient descent [Lewkowyez et al. 2020 Cohen et al. 2021
Agarwala et al. 2022 and implicit

---

regularization Barrett and Dherin 2020 Smith et al. 2020] as result of the non- non-linear training
dynamics. However, particularly y in non-stationary learning problems, neural networks can also be
vulnerable to loss of plasticity (Sodhani et al. 2020 Dohare et al., 2021 Nikishin et al. 2022 Abbas
et al.] 2023 Lyle et al. [024] as they drift away from the initial parameters. This phenomenon has been shown to present a limiting factor to performance in number of RI. tasks (Igl et al.) 2021 Lyle
ct al. 2021 Nikishin et al. 2023]. along with continual learning and warm-starting neural network
training Berariu et al. 202 Ash and Adams 2020].
variety of architectural choices can accelerate the training of extremely deep networks, including residual connections [He et al., 2016] and normalization layers [loffe and Szegedy (2015] Ba et al. 2016]: these methods have also been demonstrated help networks maintain plasticity in reinforce- ment [Bal et al.] [202]] Lyle et al.] 2023] and continual [Kumar et al. 2023) learning. Some additional works have aimed to replicate the benefits of normalization layers via normalization of the network parameters (Salimans and Kingma| [0166] Arpit et al. (0166) though layer normalization (LayerNorm) remains standard practice in most domains. A variety of analyses highlight LayerNorm's effect on gradient moments as critical factor in its efficacy (Xu et al. 2019 Xiong et al. 2020).
2.2 Effective learning rates
, normalization of the type performed by BatchNorm and LayerNorm introduces scale ieaae-ivvaiance into the layers to which it applied, where by a scale-invariant function f we mean ff fcfl, x) f(0, x) for any positive scalar с > 0. This leads to the gradient scaling inversely with the parameter norm, i.c. V f(l) c f The intuition behind this property simple: changing the direction of large vector requires greater perturbation than changing the direction of a small vector. This motivates the concept of an 'effective learning rate', which provides a scale-invarian notion of optimizer step size. In the following definition, we take the approach of Kodryan et al. [2022] and assume an implicit 'reference norm' of size for the parameters.
As noted by several prior works Van Laarhoven 2017 Li and Arora 2020 Li et al. 2020b
Definition 1 (Effective learning rate). Consider scale-invariant function f. parameters O and update function 0z+1 + 0, + rpg(O:) for some update effective learning rate in as follows:
contattffucttoo function g. Letting p T[. we then define the
npp
(fg()))
(o.(
IPP if. = vof(0))
(1)
where, letting 0 = 0 াতা we then have f(ō + (0) = f(O + +g))))
Thus by reducing the parameter norm, weight decay can have the dual effect of increasing the effective learning rate (Van Laarhoven [0117 Hoffer al.
2018 a property which has been extensively analyzed [Arora et al. 2018 Li and Arora 2020 Li et al. 2020a]. This perspective justifies the decoupling of weight decay and gradient accumulation in AdamW [Loshchilov and Hutter 2019
along with the application of learning rate schedules to the weight decay parameter [Xie et al. 2024
It also motivates scaling the norm of the updates be proportional to the parameter norm variety of optimizers [Liu et al. 2021 You et al. 2017 2020). The perspective of updates as rotations of the parameters has been applied by recent analyses the equilibrium dynamics of optimization scale-
invariant networks trained with weight decay (Wan et al. [021 Kosson et al. 2024). More recently, the work Lobacheva et al. [202T] and Kodryan et al. [2022] has studied the training properties of networks trained with parameters constrained to the unit sphere, training regime we expand upon in this work. similar approach, referred to as weight standardization, has been demonstrated to reduce the need for normalization layers in diffusion models [Karras et al., 2024
3 Analysis of normalization layers and plasticity
Although widely used and studied, the precise reasons behind the effectiveness of layer normalization remain mysterious. In this section, we provide some new insights into how normalization can help neural networks to maintain plasticity by facilitating the recovery of saturated nonlinearities, and highlight the importance of controlling the parameter norm in networks which incorporate normalization layers. We leverage these insights to propose Normalize-andd-Proeer Project, a simple training protocol to maintain important statistics of the layers and gradients throughout training.
3

---

Dormant units under different optimizers
First layer (ggd)
Socond layer (sgd)
First layer (adam)
Second layor (adam)
First layer (sgd momentum)
Second layer (sgd momentum)
1:00
1.00
1.00
0.75
0.75
0.75
0.75
0.75
0.75
risnorm
layororm
0.50
0.50
0.50
0.50
0.50
0.50
Fraction
0.25
0.25
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
20
40
40
20
2.0
40
40
Staps (x100)
Steps (x105)
Steps (X2 06)
Staps (x100)
Staps (x100)
Staps (x100)
Figure 1: Accumulation of dead units in iterated random label memorization task. The network is trained memorize random labels of MNIST dataset which are re-randomized every 1000 optimizer steps. Networks with normalization layers are able to recover from spikes in the number of dead units.
3.1 Normalization and ReLU revival
It is widely accepted that achieving approximately mean-zero, unit-variance pre-activations (assuming suitable choices of activation functions) is useful to ensure a network is trainable at initialization [e.g Martens et al. 2021]. J. and many neural network initialization schemes aim to maintain this property as the depth of the network grows. Indeed, is easy to show that in extreme cases large deviations of these statistics from their initial values can lead to variety of network pathologies including saturated units and low numerical rank of the empirical neural tangent kernel [Xiao et al. 2020]. Layer normalization not only guarantees that activations unit-norm, mean-zero initialization, but also that they stay that way over the course of training even if the data distribution changes, assuming no scale or offset parameters.
Beyond re-normalizing the pre-activation statistics, layer normalization also introduces a dependency between units in given layer via the mean subtraction and division by standard deviation transfor- mations, which translates to correlations the gradients of their corresponding weights. This mixing step allows gradients to propagate through pre-activation even if the unit is saturated, provided layer normalization is applied prior to the nonlinearity, property we highlight in Proposition[] We will use the notation SRMS (h) [[ill to specify the RMSNorm transform, and let f'(x) Oz / (x) refer to the scalar derivative of any f at scalar T.
Proposition 1. Consider two indices and of feature embedding ((fRMs(h)) such that o' (fRMs(h)) * 0, and h4 h2 # O. Then we have
(RRus(())) = -6' (D'JRMs (1)))) |/lp hohhy + 40. dh;
In contrast, for post-activation normalization the gradient is zero whenever p'(h,) 0, i.e.
d d'(h:) = 0 (aRRMs((())) = -d [0b)) (hi))))))) 0.
(2)
dh,
In other words, normalization effectively gives dead ReLU units a second chance at life - rather than immediately decaying zero, the gradients propagated the incoming weights of saturated ReLU will take on non-zero values, which depend on the gradients of the mean and variance that particular layer. These gradients will be much smaller than those that would typically backpropagate to the unit, but if an optimizer such as Adam or RMSProp is used to correct for the gradient norm, then the unit may still be able take nontrivial steps, which have a chance at propelling it back into the activated regime. This property is also naturally inherited by layer normalization, which can be viewed the composition of RMSNorm with a centering transform. An illustration of normalization allowing the network to revive dead units is given in Figure [] We include the full derivation of Proposition D in Appendix A.2 and we demonstrate the effect this can have on theoretical model of neural network training in Appendix C.4 However, while layer normalization can help the network to recover from saturated nonlinearities, introduces a new source of potential saturation which must be carefully considered, which is something we will do in the next section.

---

Parameter nom,, Jacobian norm, and performance on continual leaming task
Jacobian norm
Perameter norm
20
Norm
250
LayerNorm
Layerlorrm+
Weight Project
200
400
200
400
200
400
Num label resets
Figure 2: Continual random-labels CIFAR training: simple feedforward network architecture (No Norm) exhibits rapid growth in its parameter norm and the norm of its gradients, whereas the otherwise-identicl network with layer normalization sees parameter norm growth coupled with reduction in the norm of its gradients and reduced performance later tasks. Constraining the parameter norm of this network maintains the performance a random initialization.
3.2 Parameter norm and effective learning rate decay
While the output of iave-invariian function is insensitive to scalar multiplication of the parameters, its gradient magnitude scales inversely with the parameter norm. Consequently, growth in the norm decline in the network's sensitivity to changes in these parameters.
of the parameters corresponds In sense this is preferable, as the glacially slow but stable regime of vanishing gradients is easier to recover from than the unstable exploding gradient regime. However, if the parameter norm grows indefinitely then the corresponding reduction in the effective learning rate will eventually cause noticeable slowdowns in learning indeed, show in Figure 2 that it is quite easy to induce this type of situation. To do so, we take a small base convolutional neural network architecture (detailed in Appendix B..4 and train it random labels of the CIFAR-10 10 dataset, akin the classic setting of Zhang et al. 2021a). We then re-randomize these labels and continue training, repeating this process 500 times. When we apply this process to network without normalization layers, the Jacobian norm grows to unstable values as the parameter norm increases; in contrast, an equivalent architecture with normalization layers
sharp decline in the Jacobian norm the parameter norm increases. In both cases, the end result is similar: increased parameter norm accompanies reduced performance on new tasks.
While this particular problem is artificial, it is a real and widely observed underlying phenomenon that the magnitudes of neural network parameters tend to increase over the course of training [Nikishin ett al. 2022 Abbas et al. 2023). In a supervised learning problem, where one is using a fixed training budget, the ELR decay induced by growing parameter norms might be desirable and help protect against larg learning rates [Arora et al. 20188] Salimans and Kingmal (2066). Allowed to continue extremes, however, ELR decay becomes problematic [Lyle al.j [2024] Fortunately, this problem admits an obvious solution: re-scaling the parameters to have nontrivial ELR. Since LayerNorm
induces scale-invariance, this will not change the function computed by the network, but will change
its training dynamics. We demonstrate the utility of this strategy in Figure 2
3.3 Normalize-and-Proge
We conclude from the above investigation that normalizing network's 's pre-activations and fixing the parameter norm presents a simple but effective defense against loss plasticity. In this section, we propose principled approach to combine these two steps which we call Normalize-and- ano-Project. Our goal for NaP to provide a flexible recipe which can be applied to essentially any architecture, and which improves the stability of training, motivated by but not limited to non-stationary problems. Our approach can be decomposed into two steps: the insertion of normalization layers prior to nonlinearities in the network architecture, and the periodic projection of the network's 's weights onto a fixed-norm radius throughout training, along with a corresponding update to the per-layer learning rates into the optimization process. Algorithm provides overview of NaP.
Layer normalization. in order to benefit from the robustness to saturated nonlinearities outlined in Proposition]] D we propose adding layer normalization prior to every nonlinearity in the network. While might seem extreme, this proposal is not too far removed from standard practice. For example, Vaswani et al. [017]] apply normalization after every two fully-connected layers, and recent results suggesting that adding normalization to the key and query matrices attention heads (Henry
et al. 2020 can improve performance and the robustness of optimization WWorsmman et al. 2023
5

---

Algorithm NaP: Normalize-anddProjje
Input: network N, input for nonlinearity Φi in network do if Φ1 not already normalized then
Weight
if Wi is weight parameter then
Wi
p.W. IW
Φt t +1 LayerNorm
else
end if end for
T1. pu + W/, d t- len(
compute 0' update(0)
σ1 t oiV o0 +
for parameter Wi in network do compute Wi (eight(rrject(WW) end for
14 - 14 V ||o11
end
Step 1: re- normalize each activation function's inputs in each forward pass
Step 2: periodically re-normalize the weights
Projection stop
W
Pre-activations after
Pra-activations before normalization
normalization
Optimization step
Weight projection. As discussed in Section 3.2 once we have incorporated layer normalization we must take care to ensure that these normalization layers will not saturate, i.e. that the network's effective learning rate will not decay to zero over the course of training due to growth of the parameters. We propose disentangling the parameter norm from the effective learning rate by enforcing constant norm on the weight parameters of the network, allowing scaling of the layer outputs to depend only on the learnable scale and offset parameters. This approach is similar to that proposed by Kodryan et al. [2022]. but importantly takes care to treat the scale and offset parameters separately from weights. For simplicity, remove bias terms these are made redundant by the learnable offset parameters. In order to maintain constant parameter norm, we rescale the parameters of each layer to match their initial norm periodically throughout training the precise frequency not important as long as the parameter norm does not meaningfully grow to point of slowing optimization between projections. For example, we find that Rainbow agents an interval of 1000 steps and step produce nearly identical empirical results. In principle we could constrain the parameter norm to any arbitrary value, but the choice of fixing the initial norm makes learning rate transfer easier when adapting an existing architecture. We do not project the final linear output layer, as it is not
Scale and offset parameters. We find it is absolutely critical to normalize the weight parameters, as these represent the bulk of trainable parameters in the network. The learnable scale and offset parameters, assuming they are included in the network must be dealt with differently. Whereas rescaling the parameters of a linear map that immediately precedes a LayerNorm transform does not change the output of the function, the scale and offset terms will be first passed through nonlinearity before entering the next LayerNorm and so will not necessarily share this property. For these parameters, we can proceed in one of three ways. In the case of homogeneous activations such as ReLUs, we can normalize the concatenated scale-offset vector as described in in Appendix [D] This can require some effort to implement due to the dependency between the scale and offset parameters and may not be worthwhile for small training runs indeed, most of our empirical results did not require this step, though we include it in our analysis of smaller networks in Figure
A more general solution which requires less implementation overhead and which applies equally to Ron-ReLU activations is regularize the scale and offset parameters their initial values, strategy which we employ in our continual learning evaluations in Section Assuming suitable initial values, this approach encourages the mean and variance of the pre-activations toward values where the nonlinearity does not saturate. Finally, they can be allowed to drift unconstrained from their initialization, a choice we find umproblematic in standard benchmarks for supervised learning. For more detailed discussion on this choice, we refer to Appendix D
'We observed in many of our experiments that removing trainable scale and offset parameters often has little effect on network performance. In Rainbow agents, for example, removing the trainable offset parameter even improves performance in several environments.

---

Coupling learning rates and parameter norms
Layer-wise
Layer-wise rescaling
Global rescaling
rescaling
All learning
Global
0.5
With projection
0.0
1000
2000
3000
1000
2000
3000
1000
2000
3000
5000
10000
Without projection
step
step
step
step
Figure 3: We run a 'coupled networks' experiment as described in the text. networks exhibit similar learning curves, as seen by the rightmost subplot, however there is small but visible gap between the learning curves obtained by NaP and an unconstrained network with fixed learning rates. Using a global learning rate schedule almost entirely closes this gap, but does not induce precise equivalence in the dynamics obtained by layer-wise rescaling (leftmost).
4 Understanding effective learning rate dynamics
NaP constrains the network's effective learning rate follow an explicit rather than implicit schedule. In this section, we explore how this property affects network training dynamics, demonstrating how implicit learning rate schedules due to parameter norm growth can be made explicit and be leveraged to improve the performance of NaP in deep RL. domains.
4.1 Learning rate and parameter norm equivalence
We begin our study of effective learning rates by illustrating how the implicit learning rate schedule induced by the evolution of the parameter norm can be translated to an explicit schedule in NaP. We study a small CNN described in Appendix B. B.4 with layer normalization prior to each nonlinearity trained on CIFAR-10 with the usual label set. We train two 'twin' scale-invariant networks with the Adam optimizer in tandem: both networks see the exact same data stream and start from the same initialization, but the per-layer weights of one projected after every gradient step to have constant norm, while the other is allowed to vary the norms of the weights. We then consider three experimental settings: in the first, we re-scale the per-layer learning rates of the projected network so that the explicit learning rate is equal to the effective learning rate of its twin. In the second, we re-scale the global learning rate based the ratio of parameter norms between the projected and unprojected network, but do not tune per-layer. In the third, we do no learning rate re-scaling. We see in F the shapes of the learning curves for all networks except for the constant-ELLR variant are quite similar, with the global learning rate scaling strategy producing smaller gap than the no-rescaling strategy. By construction, the dynamics of the per-layer rescaling network and its twin are identical. Because global learning rate schedules are standard practice and induce dynamics that are quite close those obtained by parameter norm growth in Figure 3 we take this approach in the remainder of the paper, leaving layer- specific learning rates and schedules for future work.
4.2 Implicit learning rate schedules in deep RL
When taken to extremes, learning rate decay will eventually prevent the network from making nontrivial learning progress. However, learning rate decay plays an integral role in the training of many modern architectures, and is required to achieve convergence for stochastic training objectives (unless the interpolation applies or Polyak averaging is employed). In this section we will show that, perhaps unsurprisingly, naive application of NaP with a constant effective learning rate can sometimes harm performance in settings where the implicit learning rate schedule induced by parameter norm growth was in fact critical to the optimization process. More surprising is that the domain where this phenomenon is most apparent is where common wisdom would suggest learning rate decay would be undesirable: value-based deep reinforcement learning.
RL involves a high degree of nonstationarity. As result algorithms such as DQN and Rainbow often use constant learning rate schedule. Given that layer normalization has been widely observed to improve performance in these algorithms, and that parameter norm tends to increase significantly in deep RL tasks, one might at first believe that the performance improvement offered by layer normalization is happening in spite of the resulting implicit learning rate decay. A wider view of the literature, however, reveals that several well-known algorithms such as AlphaZero [Schrittwieser et al. 2020), along with many implementations of popular methods such as Proximal Policy Opti-
7

---

Implicit LR schedules Atari
Normalized retum
allen
Jam eddesoond
space Invaders
phoenix
fishing_derby
200
25
50
100 frame (CIM)
200
100
200
100
200
100
200
100 frame (XIM)
200
frame (XIM)
frame (XIM)
frame (x1M)
Figure 4: Without an explicit learning rate schedule, a Rainbow trained with NaP may fail to make any performance improvement: while the implicit schedule induced by the parameter norm is clearly important to performance, in several games this significantly outperformed by simple linear schedule terminating halfway through training. Intriguingly, we see a characteristic sharp improvement near the the decay schedule several (though not all, c.g. fishing derby) games.
mization [Schulman et al.) 2017], incorporate some form of learning rate decay, suggesting that a constant learning rate not always desirable. Indeed, Figure shows that constraining the parameter norm to induce a fixed ELR in the Rainbow agent frequently results worse performance compared to unconstrained parameters. This is particularly striking given that many of the benefits supposedly provided by layer normalization, such as better conditioning of the loss landscape [Lyle et al. [023] and mitigation of overestimation bias (Ball et al. 2023). should be independent of the effective learning rate. Instead, these properties appear to either be irrelevant for optimization or dependent on reductions in the effective learning rate to materialize.
We can close this gap by introducing learning rate schedule (linear decay from the default 6.25 10-5 to 10-6 roughly proportional to the average parameter norm growth across games). We further observe in Appendix C.T that when we vary the endpoint the learning rate schedule, we often obtain corresponding x-axis shift in the learning curves, suggesting that reaching a particular learning rate was necessary to master some aspect of the game. We conclude that, while beneficial, the implicit schedule induced by the parameter norm is not necessarily optimal for deep RL agents, and it is possible that a more principled adaptive approach could provide still further improvements.
5 Empirical Evaluations
We now validate the utility of NaP empirically. Our goal in this section is to validate two key properties: first, that NaP does not hurt performance stationary tasks; second, that NaP can mitigate plasticity loss under a variety of both synthetic and natural nonstationarities
5.1 Robustness to nonstationarity
We begin with a continual classification problem which has been widely used in works studying loss of plasticity: memorization of iteratively re-randomized labels of image dataset. In this task, each input from the dataset assigned a uniform random label; the network is then trained on this set of labels for a fixed duration, after which a set of random labels are generated and optimization begins again. Full details can be found in Appendix [..4 There is no shared structure between tasks in this problem setting, so performance is solely determined by trainability and not transfer between tasks. We evaluate our approach on a variety of sources of nonstationarity. using two architectures: small CNN, and fully-connected MLP (see Appendix [344 for details). We consider a number of methods designed to maintain plasticity including Regenerative regularization (Kumar al. [202)), Shrink and Perturb [Ash and Adams 2020), ReDo (Sokar et al. 2002), Ieaky ReLU units (inspired by the success of concatenated ReLU activations [Abbas et al. 2023)). L2 regularization, and random Gaussian perturbations to the optimizer update, heuristic form of Langevin Dynamics. We track the average online accuracy over the course of training for 20M steps, equivalent to 200 data relabelings, using a constant learning rate. We find varying degrees of efficacy in these approaches, with regenerative regularization and ReDO tending to perform the best. When we apply NaP on top of the same suite methods in Figure 5 we observe near-monotont improvements (with the exception of ReDO, where the reset mechanism does not take normalization into account) in performance and a significant reduction in the gaps between methods, with the performance curves of the different methods nearly indistinguishable in the MLP. Further, observe constant increasing slopes in

---

Continual random label memorization robustness
Average online sccuracy
cnn Baseline
mlp Ba seline
cnn With NaP
mlp With NaP
2.0
1.0
1.0
1.0
0.5
0.0
0.0
0.0
0.0
5000
10000
5000 Steps (x2000)
10000
0
5000
10000
5000
20000
Steps (x2000)
Steps (x2000)
Steps (x2000)
Ro.D0
Regenerative rog
Nokey updates
leaky activation
reg
Figure 5: Robustness to nonstationarity: see that without NaP, there is wide spread in the effectiveness of various muasticity-prssevvi methods across architectures. Once incorporate NaP, however, the gaps between these methods shrink significantly and almost uniformly improves over the unconstrained baseline.
CIFAR-I0
ImageNet-Il
C4
Pile
WikiText
Lambada
STQA
PIQA
NaP
94.64
77.26
45.7
47.9
45,4
56.6
44.2
68.8
Baseline
94.65
77.08
44.8
47.4
44.2
54.1
43.5
67.3
Norm only
94.47
77.45
44.9
47.6
44.3
53.6
43.8
Table 1: Left: Top-1 prediction accuracy on the test sets of CIFAR-10 and ImageNet-Ik. Right: per-token
accuracy of 400M transformer model pretrained the C4 dataset, evaluated on variety of language benchmarks. See Appendix for more results with variation measures.
the online accuracy, suggesting that the difference between methods has more to do with their effect on within-task performance than plasticity loss once the parameter and layer norms have been constrained.
5.2 Stationary supervised benchmarks
Having observed remarkable improvements in synthetic tasks, we now confirm that NaP does not interfere with learning more widely-studied, natural datasets.
Large-scale image classífication. We begin by studying the effect of NaP on two well-established benchmarks: V-Gl6e network [Simonyan and Zisserman 2014) on CIFAR-10, and a ResNet- 50 [He et al.) 2016) on the ImageNet- Ik dataset. We provide full details in Appendix BB4 In Table [[ obtain comparable performance in both cases using the same learning rate schedule as the baseline.
Natural language: we evaluate the effect of NaP on a 400M-parameter transformer architecture (details in Appendix B.3) trained on the C4 dataset [Raffel et al., 2020). Table | shows that our approach does not interfere with performance on this task, where we match final per-token accuracy of the baseline. When evaluating the pre- rre-trined network on variety of other datasets, we find that NaP slightly outperforms baselines in terms of performance on a variety benchmarks, including WikiText- 103, Lambada [Paperno et al. 2016), PIQA (Bisk et al. 2020 SocialIQA (Sap et al. 2019
and Pile [Gao et 2020
5.3 Deep reinforcement learning
Finally, we evaluate our approach on setting where maintaining plasticity is critical to performance: RL on the Arcade Learning Environment. We conduct full sweep over 57 Atari 2600 games comparing the effects of normalization, weight projection, and learning rate schedules on a Rainbow agent [Hessel al. 2018). In the RHS of ( we plot the spread of scores, along with estimates of the Mean and IQM of four agents: standard Rainbow, Rainbow + LayerNorm, Rainbow + NaP without explicit LR schedule, and Rainbow NaP with the LR schedule described in Section 4.2 We find that NaP with a linear schedule outperforms the other methods.
We also consider the sequential setting of Abbas et al. [2023 In this case, we consider an idealized setup where we reset the optimizer state and schedule every time the environment changes, using a cosine schedule with warmup described in Appendix B.2 To evaluate NaP on this regime, we train on each of 10 games for 20M frames, going through this eycle twice. We not reset parameters of the continual agents between games, but do reset the optimizer. We plot learning curves for the

---

Sequential Atari performance
Aggregate scores an sing la task Atari
allen
attantis
baxing
breakout
centipede
IdM
Berom
Rainbox-
Scheduled ind Unsched
10
10
10
10 space invaders
10 tennia
double dunk
freeway
pong
Median
Renum
Naibbow
Schedaled Na7-
Unscteduied Nar
10
10 Fainbow
10
10 Liyoftar
WP (reset) Rainbow (reset)
Normalized Score
(reset)
Figure 6: Left: We visualize the learning curves of continual atari agents on sequential ALE training (i.e. 200M frames). Each game is played for 20M frames, and agents pass sequentially from one to another, repeating all ten games twice for total of 400M training frames. Solid lines indicate performance the second visit to each game, and dotted lines indicate performance of randomly initialized network the game. Even in its second visit to each game, NaP performs comparably the randomly initialized networks, whereas the standard rainbow agent exhibits poor performance on all games in the sequential training regime. Right: aggregate effects of normalization on single-task atari, computed via the approach of Afgrraa et al. 2021 Bars indicate 95% confidence intervals over 4 seeds and 57 environments.
second round of games in the LHS of Figure 6 finding that NaP significantly outperforms baseline Rainbow agent with and without layer normalization. Indeed, even after 200M steps the networks trained with NaP make similar learning progress to a random initialization.
Discussion
This paper has demonstrated that maintaining plasticity in the face of nonstationary training objectives can be achieved through careful normalization the network's 's features and parameters. While there many factors contributing to the efficacy of normalization in maintaining plasticity, identified two non-obvious factors: the effect of normalization on the network's ' ability to revive saturated units, and role of the parameter norm in determining the effective learning rate of networks with normalization layers. With these insights in hand, we proposed Normalize-and-Projec a simple protocol which consists of adding layer normalization prior to nonlinearities in the network and periodically re-scaling the per-layer weights back to their initial norms.
Beyond improving performance in non-stationary supervised learning problems, NaP also presents a powerful tool for understanding the role of the effective learning rate on training dynamics. By applying NaP to reinforcement learning problems, to the ability of value- based deep RL agents to improve their performance on Atari tasks, explaining why approaches such weight decay often struggle provide the same performance benefits in reinforcement learning as they do in supervised problems. This finding opens the door to a number of further questions: why is ELR decay so critical to performance RL. What features of the environments require sufficiently low ELR to learn, and why? In what settings is the implicit learning rate schedule yielded by the parameter norm sub-optimal for learning progress, and can better schedules be determined automatically, rather than being proscribed prior to the start of training?
revealed the crucial importance of ELR decay
The development of even more effective normalization strategies is further promising avenue for future work. While we did not observe pathological behaviour in the unnormalized and unregularized scale and offset parameters of networks trained on single tasks, regularization of these parameters was critical to maintain performance in the sequential Atari domain and suggests that these parameters can interfere with learning if left unconstrained. Further, while layer normalization is not observed to impede network expressivity in visual domains, where the scale of the input does not usually earry fask-relevant information, in proprioceptive domains normalization of features may increase the difficulty of learning, as it removes valuable spatial information from the input. Applying NaP to these domains will require careful design of feature embedding layers or an alternate normalization strategy which does not remove information about the Euclidean distance between inputs.
10

---

Acknowledgements
We would like thank Vincent Roulet for his helpful and detailed feedback on earlier versions of this manuscript.
References
Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marlos C Machado. Loss of plasticity in continual deep reinforcement learning. ICML, 2023,
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 293200, 2021.
Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. arXiv preprint arXiv:22/0.04880 2022.
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In International Conference on Learning Representations, 2018.
Devansh Arpit, Yingbo Zhou, Bhargava Kota, and Venu Govindaraju. Normalization propagation: A parametric technique for removing internal covariate shift in deep networks. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1168-1176. New York, New York, USA, 20-22 Jun 2016. PMLR.
Jordan Ash and Ryan P Adams. On warm-starting neural network training. Advances in Neural Information Processing Systems, 33:3884-333 -394, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv: 160700640 2016.
David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: resnets are the answer, then what is the question? In International Conference Machine Learning, pages 342-350. PMLR, 2017.
Philip J. Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 1577-1594.
PMLR, 23-29 Jul 2023.
David GT Barrett and Benoit Dherin. 17Xiv:2009.11162 2020.
Implicit gradient regularization. arXiv preprint
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Tudor Berariu, Wojeiech Czarnecki, Soham De, Jorg Bornschein, Samuel Smith, Razvan Pascanu, and Claudia Clopath. A study on the plasticity of neural networks. arXiv preprint 20t:::/06600000 2021.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint 103.00000300065, 2021.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. In Advances Neural Information Processing Systems, volume 29, 2016.

---

Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. Advances in Neural Information Processing Systems,
2020.
Shibhansh Dohare, A Rupam Mahmood, and Richard Sutton. Continual backprop: Stochastic gradient descent with persistent randomness. arXiv preprint arXiv:2108. 06325, 2021.
Shibhansh Dohare, Juan Hemandez-Garria Parash Rahman, Richard Sutton, and A Rupam Mahmood. Loss of plasticity in deep continual learning. JMLR, 2023.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-5861, 2020.
Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax-a differentiable physics engine for large scale rigid body simulation. arXiv preprint 1xXiv:2/06.1328/ 2021.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101 0000077 2020.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and stafistics, pages 256.256. JMLR Workshop and Conference Proceedings, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-10344 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, and Yuxuan Chen. Query-key normalization for transformers. arXiv preprint arXiv:2010.044 04245, 2020.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry. Norm matters: efficient and accurate normalization schemes in deep networks. Advances in Neural Information Processing Systems, 31, 2018.
Marcel Hussing, Claas Voeleker, Igor Gilitschenski, Amir-massoud Farahmand, and Eric Faton. Dissecting deep rl with high update ratios: Combatting value overestimation and divergence. arXiv preprint 2xXv:22403.05996 2024.
Maximilian Igl. Gregory Farquhar, Jelena Luketina, Wendelin Bochmer, and Shimon Whiteson. Transient non-stationarity and generalisation in deep reinforcement learning. In International Conference on Learning Representations, 2021. URL htttppse/gppenevve...eteffoeeeee Qun8fv4q$by
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456 pmlr, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Tero Karras, Miika Aittlla, Jaakko I Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEF/CVE Conference Computer Vision and Pattern Recognition, pages 24174-24184 2024.
12

---

Maxim Kodryan, Ekaterina Lobacheva, Maksim Nakhodnov, and Dmitry P Vetrov. Training scale- invariant neural networks on the sphere can happen in three regimes. Advances in Neural Informa- lion Processing Systems, 14148-144470 2022.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. In Arxiv, 2024.
Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy, Maintaining plasticity via regenerative regularization. arXiv preprint arXiv:2308. 17958, 2023.
Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-50. Springer, 2002.
Hojoon Lee, Hanseul Cho, Hyunseung Kim, Daehoon Gwak, Joonkee Kim, Jaegul Choo, Se- Young Yun, and Chulhee Yun. Plastic: Improving input and label plasticity for sample efficient reinforcement learning. In seirth- Conference on Neural Information Processing Systems, 2023.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl- Dock-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218. 2020.
Xiang Li, Shuo Chen, and Jian Yang. Understanding the disharmony between weight normalization family and weight decay. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 4715-4722, 2020a.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In 8th International Conference Learning Representations, ICLR 2020, 2020.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33:14544 14555, 2020b.
Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue. Learning by turning: Neural architecture aware optimisation. In International Conference on Machine Learning, pages 6748- 6758. PMLR, 2021.
Ekaterina Lobacheva, Maxim Kodryan, Nadezhda Chirkova, Andrey Malinin, and Dmitry P Vetrov. On the periodic behavior of neural network training with batch normalization and weight decay. Advances in Neural Information Processing Systems, 34:21545-21556 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference on Leurning Representations, 2019. URL httt:tttt////eeeeeeee ntwp:/rmm/t/tfiitttt
BYy7
Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. In International Conference Learning Representations, 2021.
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. In International Conference on Machine Learning, 2023.
Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado Hasselt, Razvan Pascanu, James Martens, and Will Dabney. Disentangling the causes of plasticity loss in neural networks. arXiv preprint arXiv:2402. 18762, 2024.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenbolz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint 10.002210.00766. 2021.
Exgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning. In International Conference on Machine Learning, pages 16828-1668 16847. PMLR, 2022.
13

---

Evgenii Nikishin, Junhyuk Oh, Georg Ostrovski, Clare Lyle, Razvan Pascanu, Will Dabney, and André Barreto. Deep reinforeement learning with plasticity injection. International Conference on Learning Representations, 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv:/606.06600 2016.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24652-24663, 2020.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In D. Lee, M. Sugiyama, U. Luxburg. Guyon, and R. Gareett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. httphhttss/////peeeeeeeee
neressnnuriiie www.ieuurrrcccmm ceeprrpps.cpppppe/2001666/ pttpppppe///////eleeeee// 1445:00113360022eccc5543321222 Paper pdf
John Quan and Georg Ostrovski. DQN Zoo: Reference implementations of DQN-based agents, 2020.
URL http://github. http:tuuww..oom//eemiin/dd//dgaa
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 558555555,, 2020.
Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jaseha Sohl- Sohl-Dickstein. On the expressive power of deep neural networks. In international conference on machine learning, pages 2847-2854. PMLR, 2017.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv:1904.09722 2019.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839))604-669, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777 2017.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014.
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, pages 9058-999 9968- PMLR, 2020.
Shagun Sodhani, Sarath Chandar, and Yoshua Bengio. Toward training recurrent neural networks for lifelong learning. Neural computation, 32(1):1-35, 2020.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phe- nomenon in deep reinforcement learning. ICML, 2023.
Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1006.05330, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
14

---

Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learn- ing dynamics of normalized neural network using sgd and weight decay. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 6380-6391. Curran Associates, Inc., 2021. URL htttps::/pprceeeddnngg hww.:/eeeeeedinnnn..ee......e.e cc/ wwwppppppp....e/eeec. fwwes/llaalaeeee wwtpes//pppee//2///////.... patpr///22/ffilee 2268655500110054465533344-----
pdf
Mitchell Wortsman, Peter Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. /4322, 2023.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and gener- alization in deep neural networks. In International Conference Machine Learning, pages 10462-10472. PMLR, 2020.
Zeke Xie, Zhiqiang Xu, Jingzhao Zhang, Issei Sato, and Masashi Sugiyama. On the overlooked pitfalls of weight decay and how to mitigate them: A gradient-norm perspective. Advances in Neural Information Processing Systems, 36, 2024.
Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 10524-10533. PMLR, 13-18 Jul 2020. URL https: https://procediings.ntt mlr prress/v19/xxioggg
html
Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems, 32, 2019.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint
2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. UtR.htttss:/ppperevaae
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115. 2021a.
Guodong Zhang, Aleksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference on Learning Representations, 2021b.
15

---

A Derivations
A.1 Notation
Analysis of a network's 's training dynamics depends on characterizing the evolution of (pre-)ot )activations and gradients in the forward and backward passes respectively. We lay out basic notation for fully- connected layers first, and then note additional details which must be considered for convolutional, skip connection, and attention layers. We will write the parameters of a layer as 0,, and use f to denote neural network.
Fully-connected layers: We write the forward pass through network Rdo + Rdd composition of layer-wise computations f IRdi-1 + Rd of the form:
a' = f'a'tt = + h' w'ddd w 8
(3)
where & denotes a nonlinearity and fLN is a (possibly absent) normalization operator. We let a² denote the network inputs. We will refer to forward pass through subset of network with the notation f = f'a o fla, and use interchangeably f fl We will refer to the full set of network parameters by vee( (0' (2))
Convolutional layers: since convolution can be viewed as a parameterization of matrix with particular symmetry, express these layers identically to the fully-connected layers, with change in semantics such that W' is the matrix representation of the convolutional parameters 0,, where we write the embedding of 0, into matrix as Weonv (0r). For simplicity, we ignore the choice of padding.
JUN JuN(N)) + n²), = W and W = Weonv (0:)
44)
Skip-connect layers: in some network architectures, nonlinearity is placed the outputs of two subnetworks to produce a function of the form
di Σ LEL
or f' f'= -
(fIN
Σ au ) LEL
(5)
for some index set L. The choice of whether to apply normalization to the sum of the subnetwork outputs each output individually depends on the desired signal propagation properties of the network De and Smith 2020
A.2 Derivation of Proposition 1
The result described in proposition follows straightforwardly from the chain rule. For pre-activation RMSNorm have
d dh, *(fums(h)) =d' (f(rms(h)), dh,
d frms (h)
(6)
d
dh; fRMs(()))
(7)
dhi
1/2
Σ
h
hj
Σ
(8)
2
3/2 dhi
Σ
h2
hi
3/2 h4
(9)
Σ 22
dhg *(fuss(h)), J'(nns (1))s (1)))
d
(10)
||
16

---

A.3 Gradients and signal propagation
One perspective which can provide some additional insight into our approach to consider the following decomposition of the gradient being backpropagated through layer.
V Vw, ((0) = d., [(0) Ov, (LNN
(11)
Oa D hI (hx)))
(12)
With this decomposition, obtain the following interpretation of some common pathologies.
Saturated nonlincarities: a saturated nonlinearity implies that Da. and the gradient is exactly (in the case of ReLU) or very close (e.g. tanh) zero. As result, U will be zero for the affected coordinates and the corresponding parameters will remain frozen. NaP addresses this problem by using Layer or RMSNorm prior any nonlinearity the network.
2 is vulnerable to Saturated normalization layers: the normalization transformation x + Tz]
saturation ||>|||| grows, in the sense that for fixed-norm update we will have
x + v |
lim |=
0
(13)
and so the term Vn JIN (hg)) will vanish. In this case, while the optimizer will be able to update the parameters, these updates will have diminishing effect on the network's 's output.
Vanishing and exploding gradients: divergence and disappearance of the activations ag-1 and backpropagated gradients da, L(0) are well-known pathologies which can make networks untrainable. However, even networks which start training from a well-tuned initialization may still encounter exploding gradients due parameter norm growth over time [Dohare al. 202T Wortsman et al.)
2023], or vamishing gradients and activations, such as in the case of saturated (i.c. dead/dormant)
ReL.U units [Sokar et al. 2023
A.4 Details of NaP
Guiding principles: in general, the goal of NaP is to avoid dramatic distribution shifts in the pre- activation and parameter norms, and to ensure that the network can perform updates to parameters even if nonlinearity is saturated. With these in mind, there are two key properties that network designer should aim to maintain:
1. All parametric functions entering nonlinearity should have normalization layer that at ensures the gradients of all units' parameters are correlated. If there are no parameters between nonlinearities (as sometimes the case in e.g. resnets) normalization is not essential.
2. Based on our investigations in Appendix C.4 L2 normalization of the pre-activations is crucial to obtain the positive benefits of layer normalization, while centering does not have noticeable effects on the network's 's robustness to unit saturation. As a result, applying at least RMSNorm is crucial prior nonlinearities, but the choice of whether not to incorporate centering is to the designer's 's discretion.
Batch normalization layers: by default, we put layer normalization prior to batch normalization an architecture already incorporates batch normalization prior to nonlinearity. This preserves the property batch norm that individual units have mean zero across the batch, which may not be the case if layer normalization is applied after. We also always omit offset parameters if layernorm is succeeded by batchnorm, as these offset parameters will be zeroed out by batchnorm.
Skip-connect layers: provided that layer normalization is applied to the outputs of a linear transfor- mation prior nonlinearity, NaP is agnostic whether normalization is applied prior to or after a residual connection's 's outputs are added to the output of a layer. In particular, if we have layer of the form Φ(a) + a2) and di and 02 are the outputs of some subnetwork of the form Φi (.I.N (h. 1)) )) and Φ2( (h2)) where Φ1 and 22 are (possibly trivial) activation functions, then the relevant parameters will already benefit from Proposition T and is not necessary to add an additional normalization layer prior to the activation o.
17

---

Attention layers:
unlike linear layers, attention layers do not typically include bias terms. To analogize this common practice in NaP, we omit the offset and mean subtraction components of the LayerNorm transform, obtaining
MHA(WQ.,, W -
(RMM (Wx)))
(14)
where crucially fRMS is not applied along the token axis as this can lead to leakage of information during training on next-token- token-prediction objectives. similar problem also prevents us from using normalization directly on the QK product matrix, along with the observation that the intuition of normalizing vectors so that their product is equal to the cosine similarity is lost once the dot products have already occurred [Henry et al., 2020]. Empirically, we find that the scale parameters dg and σ don't seem to be strictly necessary for expressivity, and that networks can even form selective attention masks for in-context learning without using these parameters to further saturate softmax.
A.5 Dynamics of NaP
Weight projection non-interference. NaP incorporates projection onto the ball of constant norm after each update step. A natural question is whether this projection step might simply be the inverse of the update step, leaving the parameters of the network constant. Fortunately, we note that the normalization layers have the effect of projecting gradients onto a subspace which is orthogonal to the current parameter values, i.e.
(V</RMIS(()) (x))x 0.
(15)
We also note that except for extreme situations such as Neural Collapse (Papyan et al. 2020 real- world gradient updates are almost never colinear with the parameters, meaning that even without normalization layers this problem would be unlikely.
Another concern that arises from the constraints we place on the weights and features is the possibility that these constraints will limit the network's 's expressivity. Normalization does remove the ability to distinguish colinear inputs of differing norms, meaning that the inputs x and ox will map to the same output for all a; however, since many data preprocessing pipelines already normalize inputs, we argue this is not a significant limitation. Indeed, under a more widely- notion of expressivity, the number of activation patterns (Raghu al. 2017], NaP does limit expressivity all. While straightforward, we provide a formal statement and proof of this claim in Appendix A.7
Layer normalization and parameter growth: In fact, if we incorporate normalization layers into the network we might expect an even more aggressive decay schedule. Recall that in a scale have (V)f(0).0) 0) O. Thus we know that the gradient at cach time step
invariant network, will be orthogonal to the current parameters. In an idealized setting where we use the update Voē(0,) rule 0(+1 + 0, + a (V this would result in the parameter norm growing at rate @(f), corresponding to effectively linear learning rate decay.
A.6 Scale-invariance and layer-wise gradient norms
One benefit of NaP is that, because we normalize layer outputs, we limit the extent to which divergence in the norm of one layer's parameters can propagate to the gradients of other layers. For e.g. linear homogeneous activations such as ReLUUs the gradient of some objective function for some input with respect to the parameters of particular layer contains depend multilinearly on the norm of each matrix. In particular, in the simplified setting of deep linear network where f(0,x) x)
of matrix products whose norm will
II W' we recall Saxe al. (2013)
V Tw:
W
x
W
(16)
A
In particular, with 0' = W1 ..... cWk ..... WL, for k + l we would have
= f = cV wt (0xx)
(17)
18

---

The situation changes little if we add ReLU nonlincarities to the network. In this case, we use the notation Da, (x) to denote the diagonal matrix indicating whether a' [i]]
T x D (x))*)
Vwif(0;x3)
II tb. (x) K<1
(18)
[wr'((';;)) cx
(19)
we incorporate a normalization layer at the end of the network (for simplicity we consider RMSNorm here, but similar argument applies to standard LayerNorm), the scale-invariance of the resulting output means that the norms of each layer's s gradients are independent of the norms of the other layers' parameters, i.e.
W JAMS [(; x) WRJRSS [(0';x) whenever 0' = (Wi,...,, CW WI., k +
(20)
This property is appealing as means that growth or decay of the norm single layer will not interfere with the dynamics of the others. However, it does mean that a layer's effective learning rate will still be sensitive to scaling, which motivates our use of renormalization. It also does not help to avoid saturated units, motivating
of layer normalization prior to nonlinearities.
A.7 Expressivity of NaP
Finally, we discuss the effect normalization and weight projection on a notion expressivity known the number of activation patterns (Raghu et al. 2017] exhibited by neural network. This quantity relates to the complexity the function class a network can compute, giving the following result the corollary that NaP doesn't interfere with this notion of expressivity.
Proposition 2. Let f be a fully-connected network with ReLU nonlinearities. Let be the function computed by f after applying NaP. Then the activation pattemn particular architecture f and parameter be Ao, we have
A, = - A)()
(21)
Further, the decision boundary maxiedous f((x) is preserved under NaP.
Proof. We apply inductive argument on each layer. In particular, when & is ReLU nonlinearity we have
A()(h)) A(o(Fnus(hi))
@ffus (h)) <(h)
||b||
f(x) m// ([xx)(((()]]] f(x)
which trivially results in identical activation patterns in the normalized and unnormalized networks. is worth noting that one distinguishing factor from a standard ReLU network that the resulting scaling factor will be different for each x. Thus while the activation patterns will be the same, the two different inputs x y might have different scaling factors, which will be nonlinear function of the input. NaP networks, even with ReLU activations, thus do not have the property of being piecewise linear.
A.8 Rescaling scale/offset parameters (linear homogeneous networks)
We observe that for any > 0, letting d(x) = max(x, 0) we have:
fIN(o va)) = JLN (Lo((ooW /1)
(22)
SIN(o(eW (ox(e(Waxx+ (ax n4)) by homogeneity ReLU FLN(Ó(W(coOs) qp)))
(23)
(24)
19

---

Then we obtain analogous effective learning rates, letting
Doy Ve,(ea, cp; Jeo = V Vol(co,
(25)
then have equivalent updates for |0 + 12l1 and Tle =²
Jix((co + NoJoo )x ep McDon))
fIN((Co e² [oo)) + ep + è gcu)
(26)
fIN((Co + e² 1 -- + tep + 2 -u))
(27)
= -IN(((( + 9o)x + p 9p) = fLN((O + 90)xx + 9u)
(28)
Finally, note that the above also holds if we apply linear transformation W to the output of the scale-offset transform, since
fW - cp)) = fLX (eW (ox + +)) fLN(W 1))
(29)
and so the effective learning rate scales precisely as we had previously for linear layers, but now with respect to the joint norm of the scale and offset parameters o.
B Experiment details
B.1 Toy experiment details
We conduct variety of illustrative experiments toy problem settings and small networks.
in Figures 1 and 2, we use DQN-style network which consists of two sets of two convolutional layers with 32 and 64 channels respectively. We then apply max pooling and flatten the output, feeding through a 512-umit hidden linear layer before applying a final linear transform to obtain the output logits. The network uses ReLU nonlincarities. When NaP is applied, we add layer normalization prior to each nonlinearity.
Network:
B.2 RL details
Single-task atari: We base RL experiments off of the publicly available implementation of the Rainbow agent [Hessel et al. 2018) DQN Zoo [Quan and Ostrovski, 2020]. We follow the default hyperparameters detailed in this codebase. In our implementation, we add normalization layers prior to each nonlinearity except for the final softmax. We train for 200M frames on the Atari 57 suite [Bellemare al. 2013). We also allow for a learning rate schedule, which explicitly detail in cases where non-constant learning rates are used.
Sequential ALE: we use the same rainbow implementation as for the single-task results, using cosine decay learning rate for all variants. We restart the cosine decay schedule at every task change for all agents. Our cosine decay schedule uses an init value of 10-5, a peak value of the default LR for Rainbow (0.000625), 1000 warmup steps after the optimizer is reset, and end-value equal to 10-6 as in the single-task settings. We choose cosine decay due to its popularity in supervised learning, and to highlight the versatility of NaP to different LR schedules. We follow the game sequence used by Abbas et al. 2023). training for 20M frames per game.
B.3 Language details
Sequence memorization: we set dataset size of 1024 and a sequence length of 512. We use a vocabulary size of 256, equivalent to ASCII tokenization. We use the adam optimizer, and train all networks for a minimum of 10 000 steps. We reset the dataset every 1000 optimizer steps, generating set of 1024 random strings of length 512. We use baseline a transformer architecture Vaswann et al. 2017 Raffel et al. 2020 consisting of 4 attention blocks, with heads and dmodel equal 256. We
batch size 128.
In-context learning: our in-context learning experiments use the same overall setup as the sequence memorization experiments, with identical architectures and baseline optimization algorithm. In this
20

---

case we train on a dataset consisting of 4096 randomly generated strings, in which the final 100 tokens are a contiguous subsequence of the first 412, selected uniformly at random from indices in [1. 31 312].
our natural language experiments on 400M parameter transformer architecture based on the same backbone as the previous two tasks, this time consisting of 12 blocks with 12 heads and model dimension 1536. We use the standard practice of learning rate warmup followed by cosine decay, setting a peak learning rate of 2 10 which is reached after a linear warmup of 1000 steps. We use a batch size of 128 and train for 30,000 steps. We use a weight decay parameter of 0.1 with the adamw optimizer.
Natural language:
B.4 Vision details
CIFAR-10 Memorization: We consider three classes of network architecture: a fully-connected multilayer perceptron (MLP), for which we default to a width of 512 and depth of 4 in our evaluations; a convolutional network with i convolutional layers followed by two fully connected layers, for which we default to depth four, 32 channels, and fully-connected hidden layer width of 256. In all networks, we apply layer normalization before batch normalization if both are used at the same time. By default, we typically do not use batch normalization. We use ReLU activations
Our continual supervised learning domain is constructed from the CIFAR-10 dataset, from which we construct a family of continual classification problems. Each continual classification problem characterized by transformation function on the labels. For label transformations, we permute classes (for example, all images with the label 5 will be re-assigned the label 2), and random label assignment, where each input is uniformly at random assigned a new label independent of its class in the underlying classification dataset. Figures in the main paper concern random label assignments, as this is a more challenging task which produces more pronounced effects.
In our figures in the main paper, we run a total of 20M steps and a total of 200 random target resets. All networks are trained using the Adam optimizer. We conducted a sweep over learning rates for the different architectures, settling on 10 as 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.
VGGNet and ResNet-50 baselines: use the standard data augmentation policies for the CIFAR- 10 and ImageNet- 1k experiments. In our ImageNet-Ik k experiments, we use the ResNetv2 architecture variant, with label smoothing parameter of 0.1, weight decay 10-4, and an optimizer we use SGD with a cosine annealing learning rate schedule, and Nesterov momentum with decay rate 0.9, We use a batch size of 256. Our CIFAR-10 experiments use a VGG-Net architecture. We use the sgd optimizer with batch size of 32, and set learning rate schedule which starts at 0.025 and decays by a factor of 0.1 iteratively through training. We use Nesterov momentum with decay rate 0.9. We train for total of 400 000 steps.
C Additional experimental results
C.1 Arcade learning environment
Our choice of learning rate schedule in the paper is motivated by an attempt to roughly approximate the shape of the implicit schedule obtained by parameter growth on average across games in the suite, with a slightly smaller terminal point than would typically be achieved by the parameter norm alone. We consider learning rate schedules which linearly interpolate between the initial lcarning rate of 0, 0,000625 and a learning rate of 10 which is roughly equivalent to increasing the parameter norm by factor of 60. We explore the importance of the duration of this decay in Figure 7, where we conduct sweep over subset of the full Atari 57 suite (selected by sorting alphabetically and selecting every third environment). We observe that faster decay schedules result in better performance initially, but often plateau at a lower value. Decaying linearly over the entire course of training, in contrast, exhibits slow initial progress but often picks up significantly towards the end of training. We conclude that in many games, reducing the learning rate is necessary for performance improvements in this agent, and the linear decay over the entire training period doesn't 't give the agent sufficien time to take advantage of the finer-grained updates to its predictions that a lower learning rate affords.
21

---

Rainbow + NaP with LR schedules asterix
alien
bank_heist
berzerk
Normaizzd retur
10
50
1
0
100
100
200
100
breakout
crazy_ climber
double dunk
freeway
Norallzeed return
20
15
10
1.25
10-
10
1.00
5
0.75
0.50
100
200
100
200
200
100
200
gravitar
jamesbond
kung_fu_master master
name_the _hmm_gameee______
1.0
Normaized
100
2
50
1
0.0
100
200
100
200
100
200
100
200
pong
riverraid
seaquest
space space_invaders
20
2
Nomalized
1.0
2
10
0-
0
0
100
200
100
200
0
100
200
0
100
200
tennis
mp____e n_dwwn
wizard of wor
LR schedule
Narmaizzd return
15
0.00015
2.0
10
100.0
3.5
10
0.00010
25.0
200.0
5-
0.00005
50.0
0.5
0
0,00000
100
200
100
200
100
200
100
200
Million frames
Million frames
Million frames
Million frames
Figure We see that faster LR decays typically accompany fast initial progress followed by plateaus. Termi- nating the linear schedule halfway through training strikes the best balance of the four settings considered for overall progress.
C.2 Non-stationary MNIST
We include additional network statistics from the experiments shown in Figure 5jin Figure g
Linearized units: given a large batch, what fraction of the ReLU units in the penultimate layer are either 0 for all units or nonzero for all units. This gives a slightly more nuanced take on the amount of computation performed in the penultimate layer than the feature rank.
we compute the numerical rank of the penultimate layer features by sampling a batch matrix of d-dimensional feature vectors. 01,.... Fa. We then compute the numerical rank as Σ 100//01 > 0.01).
Feature rank:
of data and computing the singular values of the
Parameter and gradient norm: these both computed in the standard way by flattening out the set of parameters per-parameter gradients and computing the 2-norm of this vector.
C.3 Non-stationary sequence modeling
We find additionally that NaP is capable of improving the robustness of sequence models to nonstation- aritics, while also not interfering with the formation of in-context learning circuits. We demonstrate
22

---

Continual random label memorization robustness
baseline
mlp baseline
with NaP
mip with NaP
Lincorzed units
200
400
100
200
200
400
400
200
Stops (X2000)
Stops (X2000)
Steps (X2000)
Steps (x2000)
500
Featurn rank
200
255.5
400
100
255.0
300
0
Steps (x2000)
Steps (x2000)
Steps (x2000)
Steps (x2000)
1e7
1.5
Parameter norm
400000
1200
4000
1.0
3900
200000
0.5
1100
3800
3700
0.0
1000
200
400
200
400
200
400
0
200
400
(x2000)
Stops (x2000)
Steps (x2000)
(K2000)
2000
125
150
Gradient
1000
100
5000
100
500
50-
200
200
400
Steps (x2000)
Steps (x2000)
Steps (x2000)
Steps (x2000)
ReDO
Regenerative
Nolsy updates
leaky ReLU activation
Shrink & Perturb
L2 reg
Figure 8: We plot a variety of additional network statistics the continual CIFAR-10 experiment shown in Figure
the latter point in Figure where we train a small transformer model on a dataset of the form Sp Q sy- where Sp is prefix string of length 100 and s, is suffix string of length 50 which is contiguous subset of the prefix string. We use a fixed dataset, such that in theory the network could memorize all strings, though we sufficiently large dataset that this is not possible to achieve within the training budget. We train four protocols using next-token prediction: with and without weight projection, and with and without normalization (the networks are small enough that normalization is not critical for training stability). We evaluate accuracy on the final 48 tokens of s, as training progresses, and observe that all networks very quickly learn to identify the starting point of the suffix string and copy the relevant subset of the prefix. We observe that normalization accelerates learning of both the retrieval component of the accuracy and the memorization component of the accuracy, and that the weight projection step in fact also accelerates this process.
In Figure 10, we see that normalization and weight projection can also help to improve robustness a nonstationary random string memorization task, a sequence-modelling analogue of random label memorization in CIFAR-10. We observe that in this case, allowing learnable scale to evolve unregularized can somewhat slow down learning compared omitting this parameter, and that weight projection recovers similar dynamics to learning with large weight decay factor.
23

---

Memorization retrieval
small transformers
With LN Scale
No LN Scalee
With LN Scale
No LN Scale
1.50
1.00
0.23
0.25
Necttval
Necioval
Menvebaao
0.35
Memortiaton
0.20
0.50
0.50
0.30
0.03
9.05
0.00
0.50
2500
25.00
2503
2500
Stop
Stop
Figure 9: Demonstration that applying normalization removing learnable scale parameter does not prevent the network from learning to copy previously-observe subsequences.
Continual random string memorization
No weight projection
Weight prajection
wright projection
We Weight projection
Par 201n
WD 0.0
Finaaacuursy
100000
WD 0.1
Falbe
50000
10000
0.0
40300
S0300
80000
40000
80000
Figure 10: Random string memorization: unregularized networks exhibit plasticity loss when trained to memorize sequence of random strings, while weight projection and weight decay improve robustness this nonstationarity.
C.4 ReLU revival experiments
Many previous works have noted that adaptive optimizers are particularly damaging to network plasticity [Dohare al., 2021 Lyle al. 2023 Dohare et al., (023). The primary mechanism underlying this is due to the sudden distribution shift in gradient moments due to changes in the learning objective when the gradient norm increases, adaptive optimizers are slow to catch up and can take enormous update steps when this occurs. Layer normalization mitigates this effect due to two facts: first, the gradients of negative preactivations are nonzero, and second, all nonzero gradients are treated essentially the same by adaptive optimizers (up to e). As result, networks with layer norm can still perform significant updates to parameters feeding into 'dead' units, meaning that these networks have good chance of turning on again later.
We illustrate this with simple experimental setting, where model optimizer updates with isotropic -iissiaandistibbuteed gradient signals and perform a (truncated at zero) random walk. Formally we look at the time evolution of the system:
V + max(v v, 0) Zt,
(0,
(30)
model the evolution of features under gradient descent trajectory. To simulate the steps taken by an adaptive optimizer like RMSProp or Adam, where updates to each parameter have fixed norm, we modify this process slightly as follows:
Vi v sign (max 0) ① Oz), 2, NO.
(31)
To simulate layer normalization, compute the dot product between 2, and the Jacobian V [|v|| to simulate gradient descent:
| (
Ve = Vi-1 + be Vy max
2, ~~ ~(0,1)
(32)
and analogously compute the sign of this update model RMSProp-style optimizers:
v V. = V + sign Z V max | :)):
2: ~(0,7)
(33)
In Figure we simulate each these processes 1000 steps, and track the number of negative (i.e. 'dead') indices. We observe that layer normalization doesn't 't avoid dead units in GD, but does reduce
24

---

Dormant units under different optimizers
First loyer (ggd)
Socond layer (sod)
First layer (adam)
Second layor (adam)
First layer (sgd momentum
Second layer (sgd momentum)
1.00
1.00
0.75
0.75
0.75
0.75
0.75
0.75
norm Emnenerm rmrnoom
0.50
0.50
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.25
0.25
0.00
0.00
0.00
20
40
40
20
2.0
40
Staps (x100)
Steps (x195)
Steps (X1 06)
Stups (x100)
Staps (x200)
Staps (x100)
Figure 11: Simple MLP model with dead unit recovery after sudden changes in the classification task. We
Gaussian gradients ReLU model with adaptive optimizer
Percertage inactive units
RMSNorm ("weight projection")
With RMSNorm ("weight projection"))
No RMSNorm
With RMSNorm
1.0
1.0
1.0
0.5
0.5
0.5
0.5
normalization
Center activations
LayerNorm
0.0
0.0
0.0
0.0
Training steps
Training steps
Training steps
Training steps
Gaussian gradients ReLU model with gradient descent
Percenage inactive units
No RM5Norm
With RMSNorm ("weight projection")
No RMSNorm
With RMSNorm
("weight projection')
1.0
1.0
1.0
1.0
0.5
0.5
0.5
0.5
Unconstrained
RMSNorm
LayerNorm
0.0
0.0
0.0
0.0
0
1000
2000
1000
2000
1000
2000
1000
2000
Training
Training
Training steps
Training steps
Figure 12: Accumulation of gradients in random walk model: backpropagated 'gradients' are isotropic random Gaussian vectors and updates computed by taking the product of these vectors with the layer jacobian. We see that the centering transform actually does relatively little to reduce risk dead units, and can in fact produce a tikeer-take-all' effect wherein one large
the rate at which they accumulate. We also observe that layer normalization does avoid monotonic increases in the number of dead units in model of RMSProp. Intuitively, this makes sense: rather than freezing once they reach negative value, parameters continue updating once they become negative with equally large steps as they did before, making escape from the dead zone more probable. One visualization of few trajectories each setting is shown in Figure
C.5 Stationary supervised benchmarks
We provide the results with standard deviations in Table 2 and Table
CIFAR-10
ImageNet-1k
NaP
94.64 (0.12)
77.26 (0.04)
Baseline
94.65 (0.08)
77.08 (0.11)
Norm only
94.47 (0.18)
77.45 (0.08)
Table 2: Top-1 prediction accuracy on the test sets CIFAR-10 10 and ImageNet- 1k. Numbers in parentheses are
standard deviations.
25

---

C4
Pile
WikiText
Lambada
SIQA
PIQA
NaP
45.7 (0.0)
47.9 (0.1)
(0.1)
(0.4)
44.2 (0.2)
68.8 (0.7)
Baseline
44.8
(0.1)
(0.0)
54.1 (0.2)
43.5 (0.6)
67.3 (0.2)
Norm only
(0.9)
47.6 (0.1)
44.3 (0.0)
53.6
43.8 (0.6)
67.1 (0.4)
Table 3: Per-token accuracy of a 400M transformer model pretrained on the C4 dataset, evaluated on a variety of language benchmarks. Numbers in parentheses are standard deviations.
D A note on scale and offset parameters
One loose end from our presentation of NaP is what to do with the learnable scale and offset terms, which are by default projected and may drift from their initial values. Most supervised tasks too short for this drift to present problems, and adding layer-specific regularization or normalization adds additional engineering overhead to an experiment. However, in deep RL or in the synthetic continual tasks we present in Figure [5] this is real concern. In the case of homogeneous activations such as ReLU, the scale and offset parameters can be viewed identically to the weight and bias terms and normalized accordingly, noting that now all that matters is the relative ratio of the scale and the offset. To account for this, we propose treat the joint set a, μ as a single parameter be normalized. This resolves the issues involved with normalizing parameter to an initial value of zero, and can be shown not to change the network output (see Appendix A.8 for a derivation of this fact). With non-homogeneous nonlinearities, however, this property will not hold, and we suggest in the general case to use mild weight decay towards the initial values of and 0 for the scale and offset terms respectively. These two approaches can be summarized in the following two update rules:
(0.1) Vool2 + |M
Uuorm( μu)
and Uffecay (o, (1) = (ao + (1 - a,))
(34)
Most of our evaluations on single tasks do not use any regularization or projection of the scale and offset parameters, but we do include a regularization-
approach in our evaluations on the sequential ALE in Section In general, if it did not appear that the scale/offset drift was causing problems, we did not introduce additional complexity by adding regularization or normalization. Indeed, many cases simpler solution was to omit these parameters entirely; for example we observed that removing offsets was beneficial in several, though not all, games in the Arcade Learning Environment.
E
Additional RL Experiments
E.1
Actor-critic algorithms
We additionally investigate the effect of NaP on actor-critic methods, using the Brax implementations of SAC and PPO to evaluate the effects of the different design choices which contributed to the best-performing Rainbow agent in continuous control domains.
Experiment details: we base our experiments on the Brax [Freeman et al., 2021], using the im- plementation available at httpp:///githubbec http://ittuub.oogooggglee/e// com/google/bbray We use networks of depth 4 with width 1024. When layer normalization is added to the architecture, we add it after each hidden layer's outputs in the actor and after all but the final two layer's 's outputs for the critic (we observe reduced performance when adding normalization to the penultimate layer as well, which we conjecture is due to the importance of euclidean distance in the final layer outputs in the continuous control tasks). All normalization layers have learnable scale and bias terms, which are not regularized or normalized in experiments. We used either constant learning rate schedule, or linear schedule which decayed over 90% of the training budget to a final value 1e-6. 6. For each agent we considered two learning rates (Ie-3 and le-4) and plotted whichever achieved the highest final performance.
Effect of normalization and learning rate decay on performance: we observe an environment- dependent ranking of the relative benefits from different design choices in Figure T3 In ant (leftmost figure), the linear learning rate schedule accounted for most of the improvement experienced by the full NaP approach. In combination with learning rate decay, layer normalization also tended to induce
26

---

SIP OIRI3STG CORERNEES EITINE OLENRRS
Fpe happo?
FPR Turruee
301 huanaa Blundo
Conpunt
Comsantt
c000
Doctert Poaect Contani LA
230
Figure 13: Ablations on learning rate schedule, weight projection, and layer normalization in PPO agents trained on mujoco environments.
Effect hoP ablattons doad units SAC
Hidden
Hiddon
6.60
3.5
Irsecl Corstart
Leeer
8.02
0.5
0.2
Layatiom Consunt
Layattom Uner 23
Layettam Prccet Canpantik
0.00
Figure 14: A closer look at the effect of NaP on dead units a SAC agent trained on the Ant environment.
some improvement whether not weight projection was used, but the effect size of this was much smaller than was observed in Atari. Weight projection had relatively little effect on performance, though the combination of layer norm, projection, and linear learning rate decay consistently among the top-perforning methods across environments.
Effect of normalization and learning rate decay dead units: in Figure 14 track the accumulation of dead units across layers in the network virtually effect on dead units, with linear learning rate decay slowing the accumulation but layer normalization providing the greatest benefit, a finding consistent with observations from prior works.
SAC agent. We observe that weight projection has
E.2 Detailed learning curves on Atari
For completeness, we include learning curves for all Rainbow variants in Figure 15
27

---

Roinhow agent on atari assault
alien
amidar
asterix
asteroids
Nemaitzdd
150
0.1
50
atlan tis
bank_heist heist
battle__ore _one
beam rider
berzerk
Nomaa
0
bowling
boxing
breakout
centipede
chogper_command command
Momma
0.50
0.75
0.50
0.00
0-
crazy climber
defender
demon, attack
double
enduro
7.5
50
5.0
derbiy
freeway
frostbite
gopher
gravitar
Normalin
1.0
1.0
0.5
0.0
hero
hockey
jamesbond
kangaroo
Nomalted
5.0
25
lung NJ moster
montezuma, revenge
pacman
name: this game
phoenix
Norm feed
0.5
1.0
0.5
0.0
pitfall
pong
private eye
gbert
Noma
0.04
3,0025
0.5
road runner
robotank
seaquest
skcing
Nemated
10
0.1
space_inee invaders
stat_gunner
surround
Lime_pplot
Nomoatzd
1.0
15
20
20
0.5
1.0
tutantham
up. nn
venture
wor
Nomatted
400
10
200
0-
yars_revange
Normaooe
Raintew lsyemorm
Reinbow
Rainow leyamorm seight provecios
200
200 Trames x30'6
Trames $10"5
Figure 15: Rainbow agents on atari.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: we don't have an explicit limitations section, but we do highlight several examples where our method under-performs baselines, particularly without appropriate LR schedules.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
29

---

Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
30

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We include as many details of the experiments as we can, but have not obtained permission open-source source the code itself yet. While the descriptions should be sufficient to reproduce the smaller-scale results and the RI. results, which apply small modifications otherwise open-source or easy-to-implement models, it is possible that replicating the transformer results precisely would be challenging due to the many minor implementation details that can slightly change results in these models.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
cublic/guides www.e/euummmiiinoooooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https cnips.cc/publin hwtpabllccgggidsdeeeeeeeee gwwdeg/ddd/g/e https///oooesuuaassisooooooiooi
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
where relevant we include details and standard deviations (e.g. the inter- quartile mean and mean of deep RL agents), however many experiments are more illustrative (c.g. the twin experiments in Figure 2), and are not amenable to error bars.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
31

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g,, Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: We will add these for the camera-ready but it was difficult to locate precise details of all experiments prior the deadline.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http::/hhharpp...cppp////cccccca hwtp///wb.cc/thhccccceeeell..eee
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: include a broader impact statement.
32

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: train small models standard benchmarks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, batasets.papersstt paperswi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
34


Reviews:
[{"review_id": "ARNIvhLxfP", "reviewer": "Reviewer_6ZG9", "summary": "When training a neural network with layer normalization, an increase in the norm\nof the parameters can lead to a lower effective learning rate. This paper makes\nthe observation that, when layer normalization is used, periodic projection is\nenough to overcome this vanishing step-size. The idea is verified empirically,\nand involves essentially no tuning due to the norm being set to initialization\nand the periodicity of the projection does not significantly impact the\nperformance.\n\n\n\nDecision: This paper provides a clear articulation of a relatively straightforward\nobservation. However, the implications of the observation seem wide ranging:\nfrom non-stationary problems to reinforcement learning and even supervised\nlearning. Although the method is relatively simple and lacks \"novelty\", I think\nthis paper merits acceptance. It could be further improved by providing a more\nnuanced investigation of the interplay between layer normalization and\nnormalize-and-project.\n\nAfter rebuttal: I have updated my score to reflect the additional results in the shared reply and the discussion with authors (6 -> 7)", "questions": "- Section 1:  \"We first show that when coupled with adaptive optimizers [..] saturated units still receive sufficient gradient signal to make non-trivial changes to their parameters\"\n  Is this claim at odds with the growing empirical evidence that loss of plasticity often occurs (without normalization layers)?\n\n- Section 2.1: \"learning dynamics can be well-characterized in the\n  infinite-width limit [...], although in practice optimization dynamics diverge\n  significantly from the infinite-width limit\"\n\n  What is the distinction between learning dynamics in the infinite-width limit and the optimization dynamics in the infinite width limit? I am not sure what the takeaway of this discussion should be.\n\n- Section 2.1: \"Plasticity loss can be further decomposed into two distinct components [Lee et al., 2023]\"\n  While I agree that loss of plasticity can be thought of as loss of trainability and loss of generalization, the provided reference discusses input and label plasticity. I do not see the distinction between trainability and generalization discussed in the paper.\n\n- Section 2.2: \"$\\nabla f(c\\theta) = \\frac{1}{c} \\nabla f(\\theta)$\"\n  There is some ambiguity here in what the gradient operator is with respect to. For example, $\\nabla_\\theta f(c\\theta)  = c \\nabla_{c\\theta} f(c\\theta)$. If the parameter norm is large, then the parameter is not defined as the effectively normalized parameter but the \"raw\" parameter $c \\theta$.\n  The intuition provided is not quite as simple as the authors make it seem: it assumes that the perturbations are independent of the scale.\n\n- Section 4.2: The schedule proposed seems completely ad-hoc and the Appendix does not describe this in sufficient detail. Could the authors comment on how \"schedule misspecification\" may impact performance? For example, does this Atari-based schedule harm performance on other RL tasks? If so, then the use of the schedule seems to use relatively priviledged knowledge for a continual learning algorithms (it requires running an algorithm on the entire continual learning problem to track the parameter growth)\n\n- Section 5.1: \"Further, we observe constant or increasing slopes in the online accuracy, suggesting that the difference between methods has more to do with their effect on within-task performance \"\n\n  Can you clarify how the slope of the online accuracy has to do within-task performance rather than loss of plasticity? It could be that the effect of NAP is redundant in conjunction with some of these methods. Regenerative regularization, for example, likely keeps the norm close to that of NAP.\n\n- Section 5.2: The results on language tasks is interesting, because some degree of weight decay is usually used to train transformers (AdamW is a common choice, which your experiments also use). What effect does NAP have in conjuction with weight deacy? This is obviously related to my previous question, but NAP+AdamW seems redundant.\n\n- Section 5.3: I find the use of knowledge of the game switch to be against the\n  spirit of sequential ALE. The results are interesting, and impressive with\n  this caveat. But the results would be much more impressive if evaluated\n  without this type of human intervention.\n\n### Minor:\n- Section 3.3: There is a reference to an Algorithm -1 which does not point to anything.", "limitations": "See above", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem of reduced effective learning rate is clearly articulated. The solution is simple and straightforward, translating to improved performance across a wide range of problems: non-stationary supervised learning, reinforcement learning and even stationary learning tasks.\n- The toy experiments presented before the main experiments are convincing. I particularly like the result with the coupled networks,", "weaknesses": "- Some of the arguments are ad-hoc, and rely on \"folk knowledge\" specific to deep learning. Layer normalization is indeed commonly used, and it is a good starting point, but understanding the interplay of layer-normalization and the normalize-and-project (in particular, to what degree it affects capacity) would provide more insight.\n- The empirical results, while comprehensive, provide little further insight in addition to the toy examples presented earlier.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D62Rv2sZ33", "reviewer": "Reviewer_KJMZ", "summary": "Normalization layers improve various aspects of deep RL and continual learning, such as loss landscape conditioning and reducing overestimation bias. However, normalization can inadvertently decrease the effective learning rate as network parameters grow. This effect is problematic in continual learning where the learning rate can decay too quickly. This work proposes a re-parameterization method called Normalize-and-Project (NaP) to maintain a consistent effective learning rate throughout training, improving performance in both stationary and nonstationary environments.\n\n - Normalization layers stabilize optimization by conditioning the loss landscape and mitigating overestimation bias.\nThey create a scale-invariance in the network parameters, leading to a decline in the effective learning rate as the parameter norm increases.\nNormalize-and-Project (NaP) - NaP is a protocol combining normalization layers with weight projection to keep the effective learning rate constant. It involves inserting normalization layers before nonlinearities and periodically rescaling the network’s weights to maintain a fixed norm.\n-The paper explores how normalization layers affect a network’s plasticity and introduce the concept of effective learning rates.\nIt shows that normalization can lead to implicit learning rate decay, which can be beneficial or harmful depending on the learning context.\n -The paper evaluates NaP on various architectures and datasets, including RL tasks and benchmarks like CIFAR-10 and ImageNet.\nNaP demonstrates improved robustness to nonstationarity and maintains performance in stationary settings.\n-Loss of plasticity is a major barrier in RL and continual learning.Normalization layers help maintain plasticity by reducing parameter norm growth and stabilizing gradients.\n-NaP can be easily integrated into existing architectures like ResNets and transformers.It provides a framework for better understanding and managing learning rate schedules in nonstationary problems.\n\nThe paper validates NaP through experiments on synthetic tasks and large-scale benchmarks, showing consistent improvements in performance and stability.It highlights the importance of controlling parameter norms and maintaining effective learning rates to mitigate plasticity loss.", "questions": "No questions.", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Originality\nInnovative Method: The introduction of Normalize-and-Project (NaP) is a novel approach that addresses the challenge of effective learning rate decay in deep reinforcement learning and continual learning.\nCreative Combinations: Combining normalization with weight projection to maintain a consistent learning rate is both original and practical.\nNew Application: Applying these concepts to nonstationary reinforcement learning settings is a novel contribution.\n\nQuality\nTheoretical Rigor: The paper provides some theoretical insights into the relationship between parameter norms, learning rates, and network plasticity.\nEmpirical Validation: Robust experiments across various architectures and datasets convincingly demonstrate NaP's effectiveness.\nReproducibility: Clear methodology and detailed descriptions ensure that the work can be reproduced and built upon.\n\nClarity\nClear Writing: The paper is well-written and logically structured, making it easy to follow the arguments and results.\nGood Coverage: Background and related work sections provide context and help readers of all levels grasp the contributions.\n\nSignificance\nBroad Impact: NaP has the potential to significantly impact deep reinforcement learning and continual learning by addressing a fundamental challenge.\nPractical Applicability: The method can be easily integrated into existing architectures, enhancing its real-world relevance.\nAdvancing Understanding: Theoretical insights advance the field's understanding of learning rates and network plasticity, suggesting new research directions.\nOverall Assessment\nThe paper combines originality with practical impact, offering a novel solution to a key challenge in reinforcement learning and continual learning. Its good theoretical and empirical work, clear exposition, and significant contributions make it a valuable addition to the field.", "weaknesses": "Lack of Theoretical RL Analysis: The paper primarily focuses on the impact of normalization on network parameters and effective learning rates without delving into the theoretical implications from a reinforcement learning (RL) perspective, particularly regarding policy learning.\nPotentially include a section discussing how the NaP method affects the learning of policies in RL. Explain how the steps in network parameter space translate to changes in policy space and the potential impact on policy optimization. This addition would provide a more comprehensive theoretical foundation and align the method's analysis with RL-specific objectives.\n\nExperiments\nAdd a continuous control tasks or more complex environments, as in environment with more actions,  beyond the Atari suite, to add evidence to the generality and robustness of NaP across different RL scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their engagement with the paper, and for their constructive comments. We address individual concerns below.\n\n**Lack of theoretical RL analysis:** We agree with the reviewer that translating insights on the optimization and trainability of neural networks to policy optimization is an interesting area for future work. Given the relatively nascent state of the field’s understanding of the relationship between optimization dynamics and policy learning, particularly in actor-critic settings (see e.g.[1]), we are not currently aware of an existing framework in which we could insert our method to study its impact on policy learning and would be interested in hearing if the reviewer had a particular flavour of result in mind. One result from our paper which could provide an interesting foundation for future investigation in this direction is Figure 6 in Appendix C, which shows that inflection points in the agent’s learning curve across different schedules tend to occur when the schedules reach a particular learning rate, suggesting that some environments require decaying the learning rate to a particular value for the network to be expressive enough for an optimization step to translate to an improved policy. Additionally, we are interested in exploring whether the optimal ELR for the actor and critic may differ in policy gradient algorithms. \n\n**Experiments.**  We have evaluated NaP on PPO agents trained with the Brax framework on Mujoco and have included these results in our general comment.\n\n[1]  Ilyas, Andrew, et al. \"A Closer Look at Deep Policy Gradients.\" International Conference on Learning Representations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XDzaqOApBW", "reviewer": "Reviewer_vdAM", "summary": "This paper explores the use of normalization layers in deep reinforcement learning and continual learning, as well as their impact on the effective learning rate. Although normalization layers offer a variety of benefits in stabilizing optimization and improving the loss landscape, they also introduce a significant side effect: the growth of the network parameter norm is equivalent to the decay of the effective learning rate. In continual learning environments, the implicit decay of the learning rate due to the increase in parameter norm may drop too quickly relative to the timescale of the learning problem, which is detrimental.\n\nTherefore, this paper proposes a new method called Normalize-and-Project (NaP). The NaP method consists of two main steps: inserting normalization layers in the network architecture before non-linear activation functions; and regularly projecting the network's weights onto a fixed norm radius during the training process, along with corresponding updates to the learning rates for each layer in the optimization process. This approach ensures that the effective learning rate remains constant throughout the training.\n\nFurthermore, this paper validates the effectiveness of the NaP method through a series of experiments, demonstrating its potential to enhance performance and robustness across different learning environments and settings.", "questions": "1. Please clarify the standard mathematical definition of layer normalization and its relationship with the notation used in the various mathematical theories presented in this paper. Specifically, it would be helpful to understand how the notation for layer normalization relates to the symbols used in Definition 1， Proposition 1 and Algorithm 1, and to have a clear explanation of the transition from \\(\\theta_t\\) to \\(h\\) in the context of the paper's mathematical framework.\n\n2. Elucidate the strong coupling relationship between NaP and reinforcement learning. Particularly, please supplement the performance of NaP on typical continuous tasks such as Gym Mujoco, because on these types of tasks, a decaying learning rate schedule design is often not required. This could potentially better demonstrate the advantages of NaP, as it may allow for a more consistent learning rate that could be beneficial in environments where the learning problem does not necessitate a reduction in learning rate over time.", "limitations": "The authors have not discussed the potential negative impacts of this paper, but I believe this work indeed has no significant negative societal impacts.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Firstly, this paper addresses a significant issue within the machine learning community, namely, the plasticity of neural network learning. It offers novel and intriguing insights into the widely-used layer normalization technique, revealing its advantages in controlling the growth of parameter norms and its disadvantages in continual learning settings. Thus, the paper provides innovative insights.\n\nSecondly, the paper presents a concise and effective method that builds upon the commonly used layer normalization by incorporating a parameter norm projection step, making it more suitable for continual learning scenarios. Consequently, the method proposed in this paper has universality and practicality.\n\nLastly, through a series of experiments, the paper demonstrates that the NaP method does not affect learning in static tasks and has distinct advantages in continual supervised learning and reinforcement learning tasks.\n\nAdditionally, the paper is supported by theoretical analysis and ablation study evidence, which solidify and complete the arguments presented in the paper.", "weaknesses": "1. This paper lacks a unified mathematical expression framework; the mathematical notation for many theoretical conclusions is disjointed and abrupt, increasing the cognitive cost for readers. For instance, throughout the paper, the specific definition of layer normalization is not explicitly provided, which leaves me somewhat confused about whether normalization is applied to the parameters of each layer itself or to the inputs of the layer. Moreover, in Definition 1, parameters are denoted by \\(\\theta_t\\), but in Proposition 1, a new symbol \\(h\\) is introduced without clarifying its relationship to \\(\\theta_t\\); furthermore, in Algorithm 1, the left column only uses WeightProject(\\(W_l\\)), while the right column defines WeightProject(\\(W_l\\), \\(\\rho_l\\)) and employs an undefined function \\(len\\). Therefore, I suggest the authors carefully review the mathematical language they have used.\n\n2. Although the title of this paper indicates a focus on reinforcement learning, more than half of the experiments in the paper are based on supervised learning or artificial experiments, and the discussion and theory in this paper do not seem to be specific to reinforcement learning. Therefore, the authors need to discuss the strong coupling relationship of this scheme with reinforcement learning. In addition, there are some logically inconsistent aspects in this paper. On one hand, the paper argues that the decay of the learning rate is inappropriate in the supervised learning setting, which supports the design concept of NaP; on the other hand, in the reinforcement learning setting, learning rate decay is needed, but NaP cannot provide an independent solution, that is, it still requires some learning rate scheduling strategies. Therefore, I am more skeptical about the universality of the NaP method for the RL field.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful engagement with the manuscript and for their helpful comments. We will ensure to take these into account in our revisions.\n\n**Mathematical notation.** We thank the reviewer for highlighting this. We will be sure to review the mathematical notation and improve it in our updates based on these suggestions. We had left out a formal definition of layer normalization because of its widespread use, but can include LN explicitly in the same section where we define RMS-norm, and will also include the following clarifications:\n-  h refers to a hidden unit’s pre-activation\n- $\\rho_l$ refers to the reference norm to which layer l is projected.\n- Len is standard length function which returns the length of a tuple/vector as in python.\n\n**Strong coupling relationship of this scheme with reinforcement learning.**  While we agree with the reviewer that NaP can be applied outside of reinforcement learning problems, we do not view the generality of our approach as a weakness but rather as a strength: NaP is not only a highly effective tool for understanding and improving reinforcement learning algorithms, but also more broadly applicable to other nonstationary learning problems. We emphasize RL in the title because it is both the regime where NaP is most beneficial, and also the regime where NaP offers the most insight as a diagnostic tool. In particular, NaP yields significant performance improvements in both single- and multi-task variants of Atari, and further demonstrates the surprising and previously unknown importance of ELR decay in this regime when used as a diagnostic tool. Figure 6 in the Appendix offers particularly intriguing hints at the possibility that learning certain components of a value function require dropping the learning rate below some critical threshold. We will be sure to emphasize this point in our revisions to make the relationship between NaP and reinforcement learning clearer.\n\n**... the paper argues that the decay of the learning rate is inappropriate in the supervised learning setting, which supports the design concept of NaP.** We emphasize that we **do not** claim that learning rate decay is always bad for network training dynamics – rather, we argue that *unintended* ELR decay due to parameter norm growth can slow down learning in some cases, particularly in long training runs such as those used in continual learning benchmarks. Some learning rate decay can be appropriate for supervised learning, and is indeed a standard part of neural network training. Our emphasis is that for long and/or nonstationary training problems, parameter norm growth can induce *excessive* decay that slows down learning if it is not controlled for in some way. We will endeavour to make this point clearer in our revisions, and thank the reviewer for highlighting this ambiguity in our original text.\n\n**on the other hand, in the reinforcement learning setting, learning rate decay is needed, but NaP cannot provide an independent solution, that is, it still requires some learning rate scheduling strategies.** As our previous point has hopefully clarified, the main message of this paper is not that ELR decay is bad, but that depending on growth in the parameter norm to decay the ELR is sub-optimal compared to tuning the schedule explicitly. The strength of NaP is that it allows the learning rate schedule to be set based on prior knowledge of the problem, rather than depending on incidental growth in the parameter norm. In the case of single-task RL, for example, the natural parameter norm growth is too mild, and we find that a more aggressive schedule is beneficial. By contrast, in the sequential atari benchmark, the natural decay is too extreme and a cyclic schedule that periodically re-warms the learning rate is preferrable instead, an intuitive solution given the cyclic nature of the problem.\n\n**“Skeptical … universality of the NaP method for the RL field.”** We find that the basic principle of “start at the default learning rate used by the baseline, then decay by 2 orders of magnitude” is a relatively robust recipe for not just Atari but also the continuous control tasks “Ant” and “Humanoid” with PPO agents (see our “General Comment” for numerical details from these experiments). We believe that even better LR schedules which depend on properties like the signal-to-noise ratio in the gradients and the local curvature of the loss landscape can likely provide further benefits, and are an exciting direction for future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i8peLUnKhe", "reviewer": "Reviewer_AEQp", "summary": "This work attempts to improve the optimization for deep reinforcement learning by inserting additional layer norms into the architecture and performing weight projection steps that constrain the magnitude of the matrix weights. The paper discusses how the implicit effective learning rate schedule resulting from the growth of weight norms in standard optimization (without weight decay) affects reinforcement learning. The authors perform experiments across several domains, including non-stationary learning and standard deep reinforcement learning, generally obtaining improvements with their method.", "questions": "### Related work\nControlling some notion of an effective learning rate is not new to the broader field of deep learning. There are various optimizers that attempt to do this including LARS [1], LAMB [2], Nero [3], RVs [4]. Over time weight decay has also been shown to modulate the effective learning rate, bringing it towards a specific value, especially for SGD [5] and AdamW [4] (which should result in a very similar effect to projecting the weights, see discussion in [4]). It would be interesting to see a comparison or at least a discussion of these related works and techniques. Nero, RVs and the forced weight normalization of [6] use weight projections that should work very similarly to the proposed method (although on a finer granularity).\n\nThese optimizers and weight decay in general would probably also need a learning rate schedule similar to NaP. I find it weird that this is supposedly not standard practice in DeepRL. Even for a stationary distribution the learning rate needs to be decreased with a schedule to obtain good results when weight decay is used on top of stochastic optimization. This can be seen e.g. in the original AdamW work [7], where the optimal weight decay for standard CIFAR ResNet training is zero when a fixed learning rate is used but not with a cosine schedule.\n\n* [1]: You, Yang, Igor Gitman, and Boris Ginsburg. \"Large batch training of convolutional networks.\" arXiv preprint arXiv:1708.03888 (2017).\n* [2]: You, Yang, et al. \"Large batch optimization for deep learning: Training bert in 76 minutes.\" arXiv preprint arXiv:1904.00962 (2019).\n* [3]: Liu, Yang, Jeremy Bernstein, Markus Meister, and Yisong Yue. \"Learning by turning: Neural architecture aware optimisation.\" In International Conference on Machine Learning, pp. 6748-6758. PMLR, 2021.\n* [4]: Kosson, Atli, Bettina Messmer, and Martin Jaggi. \"Rotational equilibrium: How weight decay balances learning across neural networks.\" arXiv preprint arXiv:2305.17212 (2023).\n* [5]: Wan, Ruosi, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. \"Spherical motion dynamics: Learning dynamics of neural network with normalization, weight decay, and sgd.\" arXiv preprint arXiv:2006.08419 (2020).\n* [6]: Karras, Tero, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. \"Analyzing and improving the training dynamics of diffusion models.\" In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174-24184. 2024.\n* [7]: Loshchilov, Ilya, and Frank Hutter. \"Decoupled weight decay regularization.\" arXiv preprint arXiv:1711.05101 (2017).\n\n### Experimental Procedure\nThe paper proposes two tricks, additional normalization and the weight projection. However the effects of each one individually are not evaluated very well, only their combination. For example in Figure 1, the effects of only weight projections without the additional normalization layers is not shown.\n\nThere is little comparison with simpler baselines. For example, if preventing dead ReLUs are the main advantage of layernorm, why not consider leaky ReLUs as a baseline?\n\nThe hyperparameter tuning is not very clear. In general the weight projection change is quite significant and should have the learning rate etc tuned separately from the baseline for a fair comparison.\n\n### Clarity\nAlgorithm 1: This is quite unclear in the main body. Why are rho, mu, sigma not defined? Should these be arguments to the function along with W? Why is rho an argument but no value is provided when it is called? How do theta and theta prime relate to W? Is the semicolon subscript used to denote all layers or a subset of the matrix?\n\nThe whole description of the bias and gain handling is a bit unclear. I find it somewhat surprising that the “drift” causes problems in some of your settings. I wonder if this is due to the learning rate being too high for these parameters. One important property of weight decay is that it effectively defines a second learning rate for the bias and gain parameters by scaling the effective learning rate of the other weights (as shown / discussed in [4]). I wonder if this aspect is missing from NaP leading to some of these issues with gains and biases. Scaling the norm of the sphere you project onto might give you a similar effect.\n\n**Minor suggestions and notes:**\n\nL196: Algorithm reference is undefined\n\nL224: This forward reference to Figure 2 is unclear since the relevant information for it has not been included and the figure caption only states “as described in text”.\n\nL273: I disagree that the benefits of better conditioning should be independent of the effective learning rate in general. It is quite clear that if you set the effective learning rate sufficiently high or low you won’t learn anything, regardless of conditioning.\n\nL278: This corresponds nicely with standard training, see e.g. example for ResNet AdamW mentioned above.\n\nEQ14: This seems wrong. Is it supposed to be an inner product and a scalar zero?\n\nL637: I don’t think this is correct, linear growth would result in something more like 1/x decay in the effective learning rate schedule, not something like 1-x which is typically meant by linear decay.\n\nFigure 10: The caption is unclear or cut off.\n\nEQ33: You don’t use the d term here like you do in the algorithm", "limitations": "Yes, no concerns here", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The paper is well written overall with relatively clear figures and sufficient details (hyperparameters etc) to allow reproduction\n+ The topic and proposed method are of relevance to the community\n+ The proposed NaP method demonstrates good performance across varied experiments\n+ The method is seemingly novel to the reinforcement field (but this not my background)\n+ Variance quantification for at least some experiments", "weaknesses": "-  The NaP method seems to be closely related to various deep learning optimization methods that are not discussed in the paper (see details in question section). There should ideally be a comparison or at the very least a discussion of similar methods.\n- The experimental procedure could be slightly more rigorous / complete in some places (see details in question section)\n- Some parts of the paper are a bit unclear (details below)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their extremely helpful comments, in particular introducing us to a number of work on optimization dynamics in scale-invariant networks of which we were not aware. We address individual concerns below.\n\n**W1**: We thank the reviewer for the recommended citations, and will be sure to include them in our related work. Implementing these methods as baselines was unfortunately outside of the scope of the rebuttal period, but we have read the papers with great interest and plan on including them in our revisions. As the reviewer mentions they do not come from a reinforcement learning background, we want to emphasize that there is an extensive “graveyard” of tried-and-true supervised learning methods which have been unable to provide analogous benefits in reinforcement learning, including such basic methods as weight decay [3] and batch normalization [2].  We emphasize that translating an insight or technique from supervised learning to RL is often a highly nontrivial task (see e.g. [1] on the challenges of using momentum-based optimizers in RL), and hope that the reviewer takes this into account when assessing the significance of the contribution. \n\n\n### Questions\n\n**Related work.** We thank the reviewer for the references, and will be sure to discuss them in the context of related work to the paper in our revisions. While it is an interesting direction for further work to explore these methods in RL, doing so was outside of scope for the rebuttal period. We reiterate that it is common for techniques standard in deep learning to fail to yield benefits in deep RL, and while it is likely possible that methods which are similar to our approach .\n\n**Ablating WP and LN**: We agree with the reviewer that the paper would benefit from including baselines where WP and LN are studied in isolation. An important point to emphasize is that without normalization layers, weight projection no longer carries th esemantics of constraining the effective learning rate. In our preliminary evaluations, we found weight projection alone without normalization to be a weak baseline and did not include it in our final results, however we will be sure to include this data in our revisions.\n\n**Why not leaky ReLUs**:\nWe evaluate against a number of dead-unit-mitigating baselines in the experiments in Figure 4, including leaky ReLUs. While leaky ReLUs in particular provide some benefits, they do not typically completely mitigate plasticity loss (see e.g. [4] ).\nFurther, while preventing dead ReLUs is one advantage of layernorm,  many other advantages have been noted as well [5], which we wished to benefit from. \n\n**Hyperparameter tuning**: For standard baselines (CIFAR/Imagenet, Rainbow) we use the default LR from the baseline method we compare against and set the LR schedule to end ~2 OOMs below this value (rounding to the nearest power of 10). We found that optimal LRs and schedules for NaP often involve slightly smaller initial learning rates and more aggressive decay than their counterparts (especially if no or weak weight decay or L2 regularization is used) but that it is typically more robust to the particular starting/ending values of the learning rate when using a schedule compared to a fixed value.\n\n**Algorithm 1 notation**: We thank the reviewer for highlighting this, and will clarify the notation for Algorithm 1 in our revisions to the paper. Specifically:\n- $\\rho_l$ is the weight scale for a specific layer. We set this to be equal to its norm at initialization\n- $\\mu_l$ and $\\sigma_l$ are the (possibly vacuous) layernorm offset and scale terms respectively, which are learned parameters in the network. To make this clearer and to address another reviewer’s concern about the absence of an explicit definition of LayerNorm, we will include these terms in our LayerNorm expression which will be added to Section 3.\n\n**description of the bias and gain handling**: In most settings involving fewer than O(10^8) optimizer steps (i.e. all of the supervised learning benchmarks and single-task RL), we did not notice a significant effect from the choice of strategy for dealing with the bias and gain, and found even in the sequential ALE that the obvious solution of mild weight decay worked out of the box. We therefore did not devote much space in the paper to discuss how best to deal with these parameters as this choice does not appear to be practically significant. We will be sure to emphasize this more in our revisions. \n\n**Learning rate for the bias and gain parameters**\nWe thank the reviewer for the pointer to the reference. We agree that in principle a similar phenomenon should be at play here, with NaP behaving similarly to the RV-AdamW method in [4 (Reviewer’s reference)]. While we unfortunately did not log the relative values of the scale and offset parameters in our continual atari experiments, we do observe that the total norm of these parameters does noticeably increase over the course of training if they are allowed to evolve unimpeded (the total parameter norm grows from ~135 to ~330, and all of this growth is attributable solely to the scale/offset terms).  \n\n[1] Correcting Momentum in Temporal Difference Learning. Emmanuel Bengio, Joelle Pineau, Doina Precup. https://arxiv.org/abs/2106.03955\n\n[2] Liu, Zhuang, et al. \"Regularization Matters in Policy Optimization-An Empirical Study on Continuous Control.\" International Conference on Learning Representations.\n\n[3] Salimans, Tim, and Durk P. Kingma. \"Weight normalization: A simple reparameterization to accelerate training of deep neural networks.\" Advances in neural information processing systems 29 (2016).\n\n[4]  Lyle, Clare, et al. \"Disentangling the causes of plasticity loss in neural networks.\" Third Conference on Lifelong Learning Agents (2024).\n\n[5] Nauman, Michal, et al. \"Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning.\" Forty-first International Conference on Machine Learning.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their careful engagement with the work and thoughtful reviews. Reviewers generally agreed on the “wide ranging” implications of our findings [6ZG9], along with the “good theoretical and empirical work”[KJMZ], and “novel and intriguing insights”[vdAM], agreeing that while developing on a strong body of existing work in deep learning theory, our findings were “novel to the reinforcement [learning] field” and our method “demonstrates good performance across varied experiments”[AEQp].

We have addressed specific comments in individual rebuttals, and respond to shared concerns in this general note.
**Notation:** we clarify the following notation for Algorithm 1, which we will incorporate in our revisions:

**Choice of learning rate schedules:** 
While any particular learning rate schedule may seem arbitrary, we emphasize that we chose each task's schedule with relatively little tuning -- in the case of supervised benchmarks we swept over number of training steps and initial learning rate and used the benchmark's default schedule. Our choices for RL were inspired by parameter norm growth curves, but the principles underlying them can be applied more broadly (and indeed, were applied to our Mujoco experiments that follow with relative ease *without* any prior knowledge of the baseline parameter norm growth).

**Sensitivity to schedule mis-specification:**
Just as using a too-large or too-small learning rate can result in poor performance in deep RL, we find analogously that using a schedule that decays too quickly, or which starts at too high of a value and decays too slowly, can hurt performance. In general, we find that schedule misspecification is slightly more forgiving than standard learning rate misspecification where the learning rate is held fixed throughout training. By letting training pass through a wider range of learning rates, we increase the probability that the network will spend at least some time in the optimal range for the given problem. 

**Continual atari schedule:** The purpose of this experiment was to show that *there exists* a learning rate schedule under which we do not see loss of plasticity on the sequential ALE. We think further work exploring how to adaptively set the learning rate in response to environment nonstationarity is a particularly exciting direction to eliminate the need for this human intervention. In this particular setting, we expect that using a sufficiently good changepoint detection algorithm as a trigger for schedule resets would be sufficient to recover the performance of the handcrafted schedule.

**Mujoco experiments:** Multiple reviewers have requested evaluations on continuous control domains such as mujoco to complement our results in Atari. While we did not have sufficient time during the rebuttal period to conduct and hyperparameter-tune long-running experiments, we found that using the highly parallel brax library enabled sufficiently speedy iteration time that we could get reasonable baselines running fairly quickly, and provided a reasonable set of default hyperparameters. We have included results from training PPO a handful of popular continuous control benchmarks in the following table, in order to highlight the versatility of our method, based on the hyperparameters outlined in [this iPython notebook](https://github.com/google/brax/blob/main/notebooks/training.ipynb), using a 4-layer, 1024 width MLP as the network architecture for both actor and critic networks. 

The results here suggest that while different design choices contribute more or less significantly to performance depending on the environment (for example, the ant environment benefits most from a learning rate schedule, whereas humanoid-stand benefits most from layer normalization), the general recipe of normalization + LR schedule consistently outperforms the standard baseline, and is more robust to varying the learning rate. We do not observe significant parameter norm growth in any of these environments due to the comparatively short training times relative to Atari (even in the longest-running task humanoid-stand, parameters increased by less than an order of magnitude). We also note that due to the large scale of the value function in these environments, using layer norm and weight projection can make it difficult for the critic network to accurately predict the value. We therefore found it beneficial in these domains to omit the normalization in the final layer of the critic network to allow it to scale its outputs to the target magnitude.

| PPO                                            | ant      | hopper   | humanoid | humanoid-standup |
|------------------------------------------------|----------|----------|----------|------------------|
| Baseline                                       | 4286     | 2527     | 6182     | 33093            |
| + LR Schedule                                  | **7154** | **3400** | 8366     | 33810            |
| + LR Schedule + LayerNorm                      | **7336** | 2574     | **8473** | **45793**        |
| + LR Schedule + Layer Norm + Weight Projection | **7186** | 2232     | **8677** | **50563**        |
| + LR Schedule + Weight Projection              | **6843** | 1276     | 7315     | 33982            |
| + LayerNorm                                    | 2764     | 2652     | 7080     | 39224            |
| + LayerNorm + Weight Projection                | 3379     | 2458     | 6911     | 31367            |
| + Weight Projection                            | 4629     | 922      | 4470     | 34256            |
|                                                |          |          |          |                  |

Meta-review: Adoptation of layer normalization implies that the growth of parameter norm will incur the decay of the effective learning rate, which will have detrimental effects on picking up new knowledge in the later stages of the training process. The authors solves this problem by periodically renormalizing the weights. The reviewers appreciated the simpleness and novelty of the proposed algorithm, as well as the robustness of the experiments. Although multiple concerns, including ones on the efficacy of NaP outside of Atari domains and others on the ambiguities in the experimental procedures, were raised, the authors have adequately addressed them during the discussion period. It is recommended to accept the paper.
Decision: Accept (poster)