Normalization and effective learning rates in reinforcement learning

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.

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.

Questions

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 .

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.

Why not leaky ReLUs: We 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] ). Further, while preventing dead ReLUs is one advantage of layernorm, many other advantages have been noted as well [5], which we wished to benefit from.

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.

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:

• $\rho_l$ is the weight scale for a specific layer. We set this to be equal to its norm at initialization
• $\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.

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.

Learning rate for the bias and gain parameters We 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).

[1] Correcting Momentum in Temporal Difference Learning. Emmanuel Bengio, Joelle Pineau, Doina Precup. https://arxiv.org/abs/2106.03955

[2] Liu, Zhuang, et al. "Regularization Matters in Policy Optimization-An Empirical Study on Continuous Control." International Conference on Learning Representations.

[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).

[4] Lyle, Clare, et al. "Disentangling the causes of plasticity loss in neural networks." Third Conference on Lifelong Learning Agents (2024).

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

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.

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:

• h refers to a hidden unit’s pre-activation
• $\rho_l$ refers to the reference norm to which layer l is projected.
• Len is standard length function which returns the length of a tuple/vector as in python.

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.

... 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.

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.

“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.

We thank the reviewer for their engagement with the paper, and for their constructive comments. We address individual concerns below.

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.

Experiments. We have evaluated NaP on PPO agents trained with the Brax framework on Mujoco and have included these results in our general comment.

[1] Ilyas, Andrew, et al. "A Closer Look at Deep Policy Gradients." International Conference on Learning Representations.

We thank the reviewer for their detailed and helpful comments, and for their deep engagement with the paper. We address specific comments and questions individually as follows.

W1: “folk knowledge” … “interplay between layer normalization and NaP”:

We agree that further theoretical analysis into the dynamics of NaP is an exciting direction for further research. However, we emphasize that the theory explaining much of the “folk knowledge” referred to by the reviewer has been extensively developed in recent years, giving NaP a more rigorous grounding (see e.g. Martens et al. (2021)). Further, the study of optimization in scale-invariant networks and the resulting tying of the parameter norm and learning rate has been extensively developed by the works highlighted in Section 2.2, among others. These works show that the interaction between layer normalization and the parameter norm is precisely the effective learning rate. The interplay between the normalization and projection steps of NaP is therefore to keep the ELR exactly equal to the explicit learning rate schedule. Indeed, we motivate the weight projection step as a means of ensuring that the ELR does not unexpectedly shrink or grow over the course of training as a result of unintended changes in the parameter norm. We hope that future work may yield further insights into the implications of these works in the reinforcement learning context.

Regarding capacity, we show in Proposition 2 (Appendix A.7) shows that NaP does not change a particular notion of network expressivity studied extensively by Poole et al.

W2: "The empirical results [...] little further insight."

While we agree with the reviewer that the primary function of our supervised learning evaluations is to simply verify that NaP does not harm performance on standard benchmarks, we note that our experiments in the deep reinforcement learning settings do provide additional nuance to our understanding of NaP. In particular, we find that contrary to folk wisdom, decay in the effective learning rate can indeed be beneficial in reinforcement learning. Further, our detailed learning curves in Figure 6 illustrate how upticks in performance frequently occur when the network passes through a specific range of learning rates, suggesting that different learning rates are needed to learn different skills within a task. These insights could not have been deduced from the toy experiments, and demonstrate that NaP can be used to gain new insights into well-studied problem settings.

[Q1]

There are a variety of reasons for loss of plasticity in neural networks. Some of these involve saturated units, while others are more complex and involve collapse of the features or starvation dynamics. Proposition 1 notes a new mechanism by which layer normalization helps to protect against plasticity loss, but does not claim that this implies LN completely mitigates the plasticity loss phenomenon.

[Q2-4] We thank the reviewer for highlighting these sources of ambiguity. We will rewrite [Q2] to clarify that we mean that infinite-width dynamics (optimization/training are equivalent in this sentence) are well-characterized theoretically but don’t accurately capture empirically-observed dynamics in finite-width settings. The “input plasticity” referred to by Lee et al. in [Q3] subsumes the “warm-starting effect”, where networks generalize worse after being trained on nonstationary inputs. We will clarify this in our revisions. We can confirm that the gradient is w.r.t. the inputs of f in [Q4]. We will rewrite this statement to clarify that we are interested in computing the gradient of f w.r.t. Inputs $\theta$ evaluated at two particular values of $\theta$. I.e. $\nabla_\theta f(\theta)$ evaluated at $\theta_1$ vs $\nabla_\theta f(\theta)$ evaluated at ${\theta_2}$ where $\theta_1 = c \theta_2$ for some $c$.

[Q5] please see “general comment” (choice of learning rate schedules/schedule misspecification).

[Q6] We understand the reviewer’s confusion, as we did not have space in the paper to include per-task learning curves in Figure 4. Our comment refers to the fact that even in the first task, some strategies, in particular those which add noise to the training process, result in slower learning and `shallower’ learning curves. We will be sure to add these results to the supplementary material in our revisions and refer to them in this section of the paper.

**[Q7] ** On projected parameters, adamw has no effect because its parameter scaling step is immediately undone by the projection step (though we might see slightly more complex behaviour with L2 regularization). However, because transformers have a number of nonlinear layers (in particular positional encodings) which are not invariant under parameter projection, some form of weight decay is likely helpful to improve stability in these layers where we can’t project.

[Q8] please see "general comment" (continual atari schedule).

The shared reply addresses many of my concerns, and I will be increasing my score. However, I have a few followups:

• Re Q2-Q4: "The “input plasticity” referred to by Lee et al. in [Q3] subsumes the “warm-starting effect”, where networks generalize worse after being trained on nonstationary inputs."

  It is true that the warm-starting paper (Ash and Adams, 2019) investigates loss of generalization. It is also true that the problem studied in that paper has changes in the input distribution. But I am struggling to see how input plasticity would "subsume" warm-starting. I think this would require further explanation to make this connection clear. One alternative is to separately reference papers that individually investigate loss of trainability and loss of generalization.

• Re Q5/Q8 and general reply: I am pleased to see that the specifics of the learning rate schedule does not necessarily require knowledge of the parameter norm growth. I am also generally surprised by the Mujoco results you presented, especially the effectiveness of learning rate schedules. I understand LR schedules are somewhat common in supervised learning, I do not know how common they are in RL. I hope that the future revision of this paper reflect the surprising effectiveness of learning rate schedules in addition to the proposed method (or otherwise points to previous work that investigates LR schedules in RL).

• Re Q7: Thanks, I did not realize that the positional encoding layers would be trainable. The standard transformer architecture introduced by Vaswani et al. (2017) used non-trainable positional encodings. Could you comment on the specifics of the positional encoding used?