Understanding plasticity in neural networks
Plasticity, the ability of a neural network to quickly change its predictions
in response to new information, is essential for the adaptability and
robustness of deep reinforcement learning systems. Deep neural networks are
known to lose plasticity over the course of training even in relatively simple
learning problems, but the mechanisms driving this phenomenon are still poorly
understood. This paper conducts a systematic empirical analysis into plasticity
loss, with the goal of understanding the phenomenon mechanistically in order to
guide the future development of targeted solutions. We find that loss of
plasticity is deeply connected to changes in the curvature of the loss
landscape, but that it often occurs in the absence of saturated units. Based on
this insight, we identify a number of parameterization and optimization design
choices which enable networks to better preserve plasticity over the course of
training. We validate the utility of these findings on larger-scale RL
benchmarks in the Arcade Learning Environment.
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, Will Dabney
# Understanding plasticity in neural networks

###### Abstract

Plasticity, the ability of a neural network to quickly change its predictions in response to new information, is essential for the adaptability and robustness of deep reinforcement learning systems. Deep neural networks are known to lose plasticity over the course of training even in relatively simple learning problems, but the mechanisms driving this phenomenon are still poorly understood. This paper conducts a systematic empirical analysis into plasticity loss, with the goal of understanding the phenomenon mechanistically in order to guide the future development of targeted solutions. We find that loss of plasticity is deeply connected to changes in the curvature of the loss landscape, but that it often occurs in the absence of saturated units or divergent gradient norms. Based on this insight, we identify a number of parameterization and optimization design choices which enable networks to better preserve plasticity over the course of training. We validate the utility of these findings in larger-scale learning problems from the Arcade Learning Environment.

Machine Learning, Reinforcement Learning, Machine Learning, Machine Learning, Machine Learning, Machine Learning, Machine Learning, Machine Learning

## 1 Introduction

It is a widely observed phenomenon that after training on a non-stationary objective, neural networks exhibit a reduced ability to solve new tasks (Lyle et al., 2021; Nikishin et al., 2022; Dohare et al., 2021). This loss of plasticity occurs most robustly when the relationship between inputs and prediction targets changes over time, and the network must learn to 'overwrite' its prior predictions (Lyle et al., 2021). While such scenarios are relatively rare in supervised learning, they are baked into the way that deep reinforcement learning (RL) agents are trained. Understanding how plasticity is lost, and whether this loss can be mitigated, is crucial if we wish to develop deep RL agents which can rise to the challenge of complex and constantly-changing environments. Existing methods to promote trainability act on a wide variety of potential mechanisms by which plasticity might be lost, including resetting of layers (Nikishin et al., 2022) and activation units (Dohare et al., 2021), and regularization of the features (Kumar et al., 2020; Lyle et al., 2021). While all of these works observe performance improvements, it is unlikely that they are all obtaining these improvements by the same mechanism. As a result, it is difficult to know how to improve on these interventions to further preserve plasticity.

This paper seeks to identify the mechanisms by which plasticity loss occurs. We begin with an analysis of two interpretable case studies, illustrating the mechanisms by which both adaptive optimizers and naive gradient descent can drive the loss of plasticity. Prior works have conjectured, implicitly or explicitly, that a variety of network properties might cause plasticity loss: we present a falsification framework inspired by the study of causally robust predictors of generalization (Dziugaite et al., 2020), and leverage this framework to show that loss of plasticity cannot be uniquely attributed to any of these properties. While difficult to characterize explicitly, we provide evidence that the curvature of the loss landscape induced by new tasks on trained parameters is a crucial factor determining a network's plasticity, particularly in value-based reinforcement learning algorithms.

We conclude by completing a broad empirical analysis of methods which aim to improve the ability of a network to navigate the loss landscape throughout training, ranging from architectural choices to regularization schemes. We find that architectural choices which have been conjectured to smooth out the loss landscape, such as categorical output representations and normalization layers, provide the greatest improvements to plasticity, while methods which perturb the parameters or provide other forms of regularization tend to see less benefit. To test the generality of these findings, we apply the best-performing intervention, layer normalization, to a standard DQN architecture and obtain significant improvements in performance on the Arcade Learning Environment benchmark. We conclude that controlling the loss landscape sharpness and optimizer stability present highly
promising avenues to improve the robustness and usability of deep RL methods.

## 2 Background

It has long been observed that training a network first on one task and then a second will result in reduced performance on the first task (French, 1999). This phenomenon, known as catastrophic forgetting, has been widely studied by many works (2017). This paper concerns itself with a different phenomenon: in certain situations, training a neural network on a series of distinct tasks can result in worse performance on later tasks than what would be obtained by training a randomly initialized network of the same architecture.

### Preliminaries

**Temporal difference learning.** Plasticity loss naturally arises under non-stationarity; we will focus our analysis on temporal difference (TD) learning with neural networks, an setting known to induce significant non-stationarity. We assume the standard reinforcement learning problem of an agent interacting with an environment \(\mathcal{M}\), with observation space \(\mathcal{S}\), action space \(\mathcal{A}\), reward \(R\) and discount factor \(\gamma\), with the objective of maximizing cumulative reward. Networks trained via temporal difference learning receive as input sampled _transitions_ from an agent's interaction with the environment, of the form \(\tau_{t}=(s_{t-1},a_{t},r_{t},s_{t})\), where \(s_{t-1},s_{t}\in\mathcal{S}\), \(a_{t}\in\mathcal{A}\), and \(r_{t}=R(s_{t})\) some reward. We let \(\theta^{\prime}\) denote the _target parameters_; in practice, \(\theta^{\prime}\) is usually an outdated copy of \(\theta\) from a previous timestep, but other choices include setting it to be equal to the current parameters, or using a moving average of past values. The network \(f:\Theta\times\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}\) is trained to minimize the temporal difference error

\[\ell(\theta,\tau_{t})=\left[f(\theta,s_{t-1},a_{t})-\Box(r_{t}+\gamma f(\theta ^{\prime},s_{t},a^{\prime}))\right]^{2} \tag{1}\]

where \(\Box\) denotes a stop-gradient, \(\gamma<1\) is the discount factor, and \(a^{\prime}\) is chosen based on the variant of TD learning used. Crucially, the regression target \(r_{t}+\gamma f(\theta^{\prime},s_{t},a^{\prime})\) depends on the parameters \(\theta^{\prime}\) and changes as learning progresses. This nonstationarity occurs even if the policy and input distribution are fixed, meaning that we can study the role of nonstationarity independent of the agent's exploration strategy. We will use the shorthand \(\ell(\theta)\) for the expectation of this loss over some input distribution.

**Loss landscape analysis.** We will be particularly interested in the study of the structure of the loss landscape traversed by an optimization algorithm. We will leverage two principal quantities in this analysis: the Hessian of the network with respect to some loss function, and the gradient covariance. The Hessian of a network \(f\) at parameters \(\theta\) with respect to some loss \(\ell(\theta)\) is a matrix defined as

\[H_{\ell}(\theta)=\nabla_{\theta}^{2}\ell(\theta)\in\mathbb{R}^{d\times d} \tag{2}\]

where \(d=|\theta|\) is the number of parameters. Of particular relevance to optimization is the eigenspectrum of the Hessian \(\Lambda(H_{\ell}(\theta))=(\lambda_{1}\geq\cdots\geq\lambda_{d})\). The maximal eigenvalue, \(\lambda_{1}\), can be interpreted as measuring the sharpness of the loss landscape (Dinh et al., 2017), and the condition number \(\lambda_{1}/\lambda_{d}\) has significant implications for convergence of gradient descent optimization in deep neural networks (Gilmer et al., 2022).

We will also take interest in the covariance structure of the gradients of different data points in the input distribution, a property relevant to both optimization and generalization (Fort et al., 2019; Lyle et al., 2022). We will estimate this covariance structure by sampling \(k\) training points \(\mathbf{x}_{1},\ldots,\mathbf{x}_{k}\), and computing the matrix \(C_{k}\in\mathbb{R}^{k\times k}\) defined entrywise as

\[C_{k}[i,j]=\frac{\langle\nabla_{\theta}\ell(\theta,\mathbf{x}_{i}),\nabla_{ \theta}\ell(\theta,\mathbf{x}_{j})\rangle}{\|\nabla_{\theta}\ell(\theta, \mathbf{x}_{i})\|\|\nabla_{\theta}\ell(\theta,\mathbf{x}_{j})\|}. \tag{3}\]

If the off-diagonal entries of \(C_{k}\) contain many negative values, this indicates interference between inputs, wherein the network cannot reduce its loss on one region without increasing its loss on another. If the matrix \(C_{k}\) exhibits low rank (which, given a suitable ordering \(\sigma\) of the data points \(\mathbf{x}_{\sigma(1)},\ldots,bx_{\sigma(k)}\) will yield a block structure) then the gradients are degenerate and largely colinear, which can indicate either generalization when their dot product is positive, or interference when their dot product is negative.

### Defining plasticity

The study of plasticity has concerned neuroscience for several decades (Mermillod et al., 2013; Abbott and Nelson, 2000), but has only recently emerged as a topic of interest in deep learning (Berariu et al., 2021; Ash and Adams, 2020; **7**). Classical notions of complexity from the computational learning theory literature (Vapnik, 1968; Bartlett and Mendelson, 2002) evaluate whether a hypothesis class contains functions that capture arbitrary patterns, but are agnostic to the ability of a particular search algorithm, such as gradient descent, to find them. A billion-parameter neural network architecture might have the _capacity_ to represent a rich class of functions, but if all of its activation units are saturated then it cannot be trained by gradient descent to realize this capacity.

Plasticity, broadly construed, refers to a network's ability to learn new things. Learning can be quantified as either the reduction in the model's training loss after optimization, or as the performance of the model on held-out test points. Studies of plasticity in both supervised and reinforcement learning have observed reduced generalization performance
as a result of overfitting to limited data early in training (Ash and Adams, 2020; Berariu et al., 2021; Igl et al., 2021). However, reinforcement learning tasks often lack a natural notion of a test set as the data gathered by the agent is not generally independent and identically distributed, and many works have identified impaired ability to even reduce the learning objective on the training distribution (Dohare et al., 2021; Lyle et al., 2021; Nikishin et al., 2022). This work will leverage the formulation of Lyle et al. (2021), who define plasticity as the ability of a network to update its predictions in response to a wide array of possible learning signals on the input distribution it has been trained on. This formulation is applicable to learning problems which do not admit a straightforward train-test split, as is the case in many deep RL environments.

Concretely, we consider an optimization algorithm \(\mathcal{O}:(\theta,\ell)\mapsto\theta^{*}\) which takes initial parameters \(\theta\in\Theta\) and some objective function \(\ell:\Theta\rightarrow\mathbb{R}\), and outputs a new set of parameters \(\theta^{*}\). The parameters \(\theta^{*}\) need not be an optimum: \(\mathcal{O}\) could, for example, run gradient descent for five steps. In order to measure the flexibility with which a network can update its predictions under this optimization algorithm, we consider a distribution over a set of loss functions \(\mathcal{L}\) each defined by some learning objective. For example, consider a distribution over regression losses

\[\ell_{f,\mathbf{X}}(\theta)=\mathbb{E}_{x\sim\mathbf{X}}[(f(\theta,\mathbf{x} )-g_{\omega}(\mathbf{x}))^{2}] \tag{4}\]

where \(g_{\omega}\) is induced by a random initialization \(\omega\) of a neural network. In order to match the intuition that more adaptable networks should have greater plasticity, we set a baseline value \(b\) to be the loss obtained by some baseline function (e.g. if \(\ell\) is a regression loss on some set of targets, we set \(b\) to be the variance of the targets), and then define plasticity to be the difference between the baseline and the expectation of the final loss obtained by this optimization process after starting from an initial parameter value \(\theta_{t}\) and optimizing a sampled loss function \(\ell\) subtracted from the baseline \(b\).

\[\mathcal{P}(\theta_{t})=b-\mathbb{E}_{\ell\sim\mathcal{L}}[\ell(\theta_{t}^{* })]\text{ where }\theta_{t}^{*}=\mathcal{O}(\theta_{t},\ell) \tag{5}\]

We then define the loss of plasticity over the course of a trajectory \((\theta_{t})_{t=0}^{N}\) as the difference \(\mathcal{P}(\theta_{t})-\mathcal{P}(\theta_{0})\). We note that this definition of plasticity loss is independent of the value of the baseline \(b\), i.e. the difficulty of the probe task for the network, allowing us to measure the relative change in performance of checkpoints taken from a training trajectory.

## 3 Methodology and motivating questions

The following sections will present a series of experiments which tease apart different causal pathways by which plasticity loss occurs and evaluate the predictive power of a range of hypotheses concerning the root causes thereof. We now outline the experimental methodology and research questions underpinning this investigation.

### Measuring plasticity

In order to determine whether a candidate intervention preserves plasticity, we must first set out a consistent standard by which we will measure plasticity. Given a suitably generic class of target functions inducing the losses \(\ell\), equation 4 characterizes the adaptability of the network to arbitrary new learning signals from the unknown set of possible future tasks. We therefore construct a set of regression targets to sample over which correspond to this uniform prior over possible future update directions. A different distribution over future target functions might give different numerical results; however, we believe that a uniform distribution captures a more universal notion of plasticity.

In our empirical evaluations, we will set \(\mathbf{X}\) to be the set of transitions gathered by an RL agent and stored in some replay buffer, and \(f\) to be the neural network architecture. Given some offset \(a\in\mathbb{R}\), we will apply the transformation \(g(x)=a+\sin(10^{5}f(\mathbf{x};\omega_{0}))\), with \(\omega_{0}\) sampled from the same distribution as \(\theta_{0}\), to construct a challenging prediction objective which measures the ability of the network to perturb its predictions in random directions sampled effectively uniformly over the input space. Because the mean prediction output by a deep RL network tends to evolve away from zero over time as the policy improves and the reward propagates through the value function, we will set \(a\) to be equal to the network's mean prediction in order not to bias the objective in favour of random initializations, which have mean much closer to zero. The optimizer \(\mathcal{O}\) will be identical to that used by the network on its primary learning objective, and we found running this optimizer for a budget of two thousand steps enabled reasonably efficient iteration time while also providing enough opportunity for most random initializations to solve the task.

### Environments

The experimental framework we consider, and which will be touched on in each of the following sections is as follows. We construct a simple MDP analogue of image classification, i.e. the underlying transition dynamics are defined over a set of ten states and ten actions, and the reward and transition dynamics depend on whether or not the action taken by the agent is equal to the state's latent label. We construct three variants of a block MDP whose state space can be given by the discrete set \(\{0,\dots,9\}\) and whose observation space is given by either the CIFAR-10 or MNIST dataset.

**True-label:** each state \(s\) of the MDP produces an observation from that class in the underlying classification dataset. Given action \(a\), the reward is the indicator function \(\delta_{a=s}\). The MDP then randomly transitions to a new state.

**Random-label:** follows the same dynamics as the previous environment, but each image is assigned a random label
in \(\{0\dots 9\}\), and the observation from an MDP state \(i\) is sampled from images with (randomized) label \(i\).

**Sparse-reward:** exhibits the same observation mapping as _true-label_. The reward is equal to \(\delta_{a=s=9}\). The MDP transitions to a random state if \(a\neq s\) and otherwise to \(s+1\).

We design these environments to satisfy two principal desiderata: first, that they present visually interesting prediction challenges with varying degrees of reward smoothness and density, and second that they allow us to isolate non-stationarity due to policy and target network updates independent of a change in the state visitation distribution. In the true-label and random-label variants, the transition dynamics do not depend on the agent's action, whereas in the sparse environment the policy influences the state visitation distribution. The different reward functions allow us to compare tasks which are aligned with the network's inductive bias (in the true-label task) and those which are not (the random-label task).

### Outline of experiments

The experiments presented in Sections 4 and 5 aim to answer a fundamental question: **what happens when neural networks lose plasticity?** Section 4 constructs two experimental settings which illuminate phenomena driving two very different forms of plasticity loss. The first constructs a non-stationary learning problem that induces extreme forms of instability in adaptive optimizers, leading to plasticity loss. The second identifies a bias in the dynamics of gradient descent which leads to a progressive sharpening of the loss landscape of not just the current task, but also new tasks, by contrasting a gradient descent trajectory with a random walk in parameter space.

Section 5 asks **what properties _cause_ plasticity loss?** Disentangling cause from correlation is a notoriously difficult problem throughout science and economics. We put a number of quantities which have been conjectured to drive plasticity loss to the test, evaluating the consistency of quantities such as weight norm, feature rank, and the number of dead units in the network across the range of tasks outlined in Section 3.2. Crucial to this investigation is the notion of a causally robust predictor: a quantity which causally influences plasticity should exhibit consistency across a variety of experimental settings, such as different learning environments or network architectures. We follow up the largely negative results of these experiments with a qualitative analysis of learning curves on the probe tasks described in Section 3.1 that emphasizes the critical role of the loss landscape in plasticity.

Section 6.2 addresses the question: **how can we mitigate plasticity loss?** It evaluates the effectiveness of a range of interventions on the network architecture and on the optimization protocol, focusing on methods known to increase the smoothness of the loss landscape, applying the same evaluation protocol as described in this section in order to measure plasticity across our classification MDP testbed.

## 4 Two simple studies on plasticity

We begin with some interpretable examples of learning problems where plasticity loss occurs. These examples illustrate how the design of optimizers can interact with nonstationarity to produce instabilities that drive plasticity loss in one of the examples above, and explore how the dynamics of gradient-based optimizers might affect more subtle properties of the loss landscape.

### Optimizer instability and non-stationarity

The robustness of existing optimizers across a wide range of datasets and network architectures has played a key role in the widespread adoption of deep learning methods. For example, the Adam optimizer (Kingma & Ba, 2015) with a learning rate of \(10^{-3}\) will often yield reasonable initial results on a range of network architectures from which the practitioner can iterate. However, when the assumptions on stationarity underlying the design of this optimizer no longer hold, the optimization process can experience catastrophic divergence, killing off most of the network's ReLU units. We can see an example of this in a simple non-stationary task in Figure 1. A two-hidden-layer fully-connected neural network is trained to memorize random labels of MNIST images (full details provided in Appendix A.1). After a fixed training budget, the labels are re-randomized, and the network continues training from its current parameters. This process quickly leads a default Adam optimizer to diverge, saturating most of its ReLU units and resulting in trivial performance on the task that a freshly initialized network could solve perfectly.

Figure 1: Abrupt task changes can drive instability in optimizers which depend on second-order moment estimates for adaptive learning rate scaling. Setting these estimators to be more robust to small gradient norms and to update moment estimates more quickly mitigates this issue.
The mechanism of this phenomenon emerges when we consider the update rule for Adam, which tracks a second-order estimate \(\hat{v}_{t}\) along with a first-order moment estimate \(\hat{m}_{t}\) of the gradient via an exponential moving average

\[u_{t}=\alpha\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}+\bar{\epsilon}}+\epsilon}. \tag{6}\]

Gradients tend to have norm proportional to the training loss. When the loss changes suddenly, as is the case when the perfectly-memorized MNIST labels are re-randomized (or when the target network is updated in an RL agent), \(\hat{m}_{t}\) and \(\hat{v}_{t}\) will no longer be accurate estimates of their moment distributions. Under the default hyperparameters for deep supervised learning, \(\hat{m}_{t}\) is updated more aggressively than \(\hat{v}_{t}\), and so the updates immediately after a task change will scale as a large number divided by a much smaller number, contributing to the instability we observe in Figure 1. In this instance, the solution is simple: we simply increase \(\epsilon\) and set a more aggressive decay rate for the second-moment estimate, and the network avoids catastrophic instability. Intriguingly, a large value of \(\epsilon\) is frequently used in deep RL algorithms such as DQN (Mnih et al., 2015) relative to the default provided by optimization libraries, suggesting that the community has implicitly converged towards optimizer hyperparameters which promote stability under nonstationarity.

### Loss landscape evolution under non-stationarity

Even when optimization is sufficiently stable to avoid saturated units, prior work still observes reductions in network plasticity (Lyle et al., 2021). The causes of this phenomenon are more difficult to tease apart; neural network initializations have been tuned over several decades to maximize trainability, and many properties of the network change during optimization which could be driving the loss of plasticity. A natural question we can ask in RL is whether the optimization dynamics followed by a network bias the parameters to become less trainable, or whether the loss of plasticity is a natural consequence of any perturbation away from a carefully chosen initialization distribution.

We frame this question as a controlled experiment, in which we compare the evolution of two coupled updating procedures: one follows gradient-based optimization on a non-stationary objective (full details in Appendix A.1); the second follows a random walk, where we add a Gaussian perturbation to the parameters with norm equal to the size of the gradient-based optimizer update. Both trajectories start from the same set of randomly initialized parameters and apply updates of equal norm; the only difference is the direction each step takes. We evaluate how the structure of the local loss landscape with respect to a probe task evolves in both networks by comparing the Hessian eigenvalue distribution, and by comparing the covariance structure \(C_{k}\) of gradients on sampled inputs, with \(k=512\) equal to the batch size used for training. We compute the Hessian matrix for a regression loss towards a perturbation \(\epsilon\sim\mathcal{N}(0,1)\) of the network's current output, i.e. \(\ell(\theta)=[f_{\theta}(\mathbf{X})-\Box f_{\theta}(\mathbf{X})+\epsilon]^{2}\) where \(\Box\) indicates a stop-gradient, to obtain a proxy for how easily the network can update its predictions in arbitrary directions; we do not evaluate the Hessian or gradient structure of the primary learning objective as these will trivially differ between the trajectories.

We observe that the spectral norm of the Hessian of both processes increases over time; however, the outliers of the spectrum grow significantly faster in the network trained with gradient descent. Additionally, the network trained with gradient descent begins to exhibit negative interference between gradients, a phenomenon not observed in the Brownian motion. In other words, the inductive bias induced by gradient descent can push the parameters towards regions of the parameter space where the local loss landscape is less friendly to optimization towards arbitrary new objectives than what would be obtained by blindly perturbing randomly initialized parameters.

## 5 Explaining plasticity loss

While in some instances it is straightforward to deduce the cause of plasticity loss, most learning problems induce com

Figure 2: Evolution of the gradient and Hessian under gradient-based optimization compared to random perturbation of the parameters. Top: the density of the spectrum of the Hessian over different values of \(\lambda\) exhibits a larger outlier peak after gradient descent. Bottom: gradient descent induces more gradient interference between inputs and greater curvature of the loss landscape.
plex learning dynamics that make it difficult to determine root causes. This section will show that a number of plausible explanations of plasticity loss, including the rank of the network's features, the number of saturated units, the norm of its parameters, and the rank of the weight matrices, do not identify robust causal relationships. We provide some evidence supporting the hypothesis that plasticity loss arises due to changes in the network's loss landscape, and conclude with a discussion of the potential trade-offs that must be faced between preserving a trainable gradient structure and accurately predicting a value function.

### Experimental setting

We train a set of DQN agents on each environment-observation space combination in the classification MDP set described in Section 3.2, and evaluate the ability of each network to fit a randomly generated set of target functions as described in Section 2.2 after a fixed number of training steps. In the experiments shown here, we run the DQN agents with a target network update period of 1,000 steps; as mentioned previously, this is the principal source of non-stationarity in the true-label and random-label tasks. Every 5000 steps, we pause training, and from a copy of the current parameters \(\theta_{t}\) we train the network on a set of new regression problems to probe its plasticity. We log the loss at the end of 2,000 steps of optimization, sampling 10 different random functions, then resume training of the RL task from the saved parameters \(\theta_{t}\). We consider two network architectures: a fully-connected network (MLP) and a convolutional network architecture (CNN). Full details of the environments are included in Appendix A.2.

### Falsification of prior hypotheses

Prior work has proposed a number of plausible explanations of why neural networks may exhibit reduced ability to fit new targets over time. Increased weight norm (Nikishin et al., 2022), low rank of the features or weights (Kumar et al., 2020; Gulcehre et al., 2022), and inactive features (Lyle et al., 2021; Dohare et al., 2021) have all been discussed as plausible mechanisms by which plasticity loss may occur. However, the explanatory power of these hypotheses has not been rigorously tested. While a correlation between a particular variable and plasticity loss can be useful for diagnosis, only a causal relationship indicates that intervening on that variable will necessarily increase plasticity.

This section will seek to answer whether the above candidate explanations capture causal pathways. Our analysis is based on a simple premise: that for a quantity to exhibit explanatory power over plasticity loss, it should exhibit a consistent correlation across different experimental interventions (Buhlmann, 2020). If, for example, parameter norm is positively correlated with plasticity in one observation space and negatively correlated in another, then it can be ruled out as a causal factor in plasticity loss. To construct this experiment, we train 128 DQN agents under a range of tasks, observation spaces, optimizers, and seeds. Over the course of training, we log several statistics of the parameters and activations, along with the plasticity of the parameters at each logging iteration.

In Figure 3, we show scatterplots illustrating the relationship between plasticity and each statistic, where each point in the scatterplot corresponds to a single training run. We see that for each of four quantities, there exists a learning problem where the quantity positively correlates with plasticity, and one in which it exhibits a negative correlation. In many learning problems the correlation between plasticity loss and the quantity of interest is nonexistent. In all cases we note that the correlation with plasticity is already quite weak; even so, the ability to reverse the sign of this correlation is a further mark against the utility of these simple statistics as causal explanations of plasticity. For example, we see a positive correlation between weight norm and plasticity loss in environments which use CIFAR-10 observations, but a slight negative correlation in environments which sample observations from MNIST. A similar reversal happens with respect to feature rank across environments.

### Loss landscape evolution during training

If the simple statistics we have considered thus far lack explanatory power, how should we characterize plasticity loss? One open question is whether the reduced ability to fit arbitrary new targets arises because the optimization process gets caught in local optima, or whether it arises due to overall slow or inconsistent optimization progress. To answer this question, we turn our attention towards the learning curves obtained by networks when we ask them to fit new target functions. We study these learning curves primarily because they convey precisely the ease or difficulty of navigating the loss landscape. In particular, the learning curve tells us whether optimization is getting trapped in bad minima (in which case the learning curve would hit an early plateau at a large loss value), or whether the network has greater difficulty reducing the loss enough to find a minimum in the first place (corresponding to a flatter slope).

We show in Figure 4 the learning curves obtained by an optimization trajectory from parameters \(\theta_{t}\) on the probe task from different timesteps \(t\) of training on the RL task. We see that parameters from early training checkpoints quickly attain low losses, but that the slopes of these learning curves are monotonically increasing with the parameter age \(t\). Of particular note is the increasing variance of the curves: in the full-batch case, this non-monotonicity is associated with increasing loss landscape sharpness (Cohen et al., 2021). In the mini-batch optimization setting, we observed both
increasing interference between minibatches as well as non-monotonicity in the loss even on the minibatch on which the gradient was computed. In short, we see that it is increasing difficulty of navigating the loss landscape, rather than poor local minima, that appears to drive plasticity loss.

## 6 Solutions

Thus far, we have demonstrated that neural networks can be demonstrated to lose plasticity even in a task as simple as classifying MNIST digits, assuming that a degree of non-stationarity is introduced into the optimization dynamics. We now turn our attention to means of reducing or reversing this loss of plasticity. Section 6.1 will evaluate whether scaling alone can eliminate plasticity loss. Section 6.2 will evaluate the effects of a variety of interventions on plasticity. We test the applicability of these findings to larger scale tasks in Section 6.3.

### The role of scaling on plasticity

Before considering sophisticated methods to address plasticity loss, we must first answer the question of whether this is simply a disease of small networks. In the context of the impressive successes of large models and the resultant'scaling laws' phenomenon (Kaplan et al., 2020), it is entirely plausible that plasticity loss, like many other challenges, vanishes in the limit of infinite computation. We find that while plasticity loss is easiest to induce in extreme forms in small networks, scaling a CNN to the limit of a single GPU's memory is insufficient to eliminate plasticity loss even in the simple classification tasks described in the previous section. We visualize the relationship between network width and plasticity loss in Figure 5.

These observations suggest that plasticity loss is unlikely to be the limiting factor for sufficiently large networks on sufficiently simple tasks. However, for tasks which do not align with the inductive bias of the network (as in the MLPs trained on CIFAR-10), or for which the network is not sufficiently expressive (as is the case for the small networks of any architecture), we see a reduction in the ability to fit new targets over time. Because we typically cannot guarantee a priori that a learning problem will fall in the first category, we therefore turn our attention to other design choices which

Figure 4: Plasticity loss corresponds to slower training progress, rather than higher plateaus, in the networks studied in this paper. We plot learning curves on a new target fitting task starting from network checkpoints at different points in training. This figure illustrates a CNN trained on the true-label MDP described in Section 6 with a CIFAR-10 observation space.

Figure 5: We observe a consistent decline in plasticity loss across different target update frequencies as a result of scaling in several architecture-dataset combinations; however, even when scaling the architecture to the point where it no longer fits on a single GPU, we are still unable to completely eliminate plasticity loss on these simple classification-inspired problems.

Figure 3: Results of our experimental falsification design: for any variable we consider, it is possible to construct a set of learning problems in which the variable exhibits either a positive or a negative correlation with plasticity. For example, weight norm and weight rank exhibit differing correlation signs depending on the observation space, while feature rank and sparsity depend on the reward structure of the environment.
might further insure networks against plasticity loss.

### Interventions in toy problems

In this section we evaluate the effect of a variety of interventions on plasticity loss. We evaluate interventions on the same task used in Section 5.1, training for 100 iterations of 1000 steps. We consider four architectures: a multi-layer perceptron (MLP), a convolutional neural network (CNN) without skip connections, a ResNet-18 (He et al., 2016), and a small transformer based on the Vision Transformer (ViT) architecture (Dosovitskiy et al., 2020).

We consider the following interventions: **resetting** the last layer of the network at each target network update, a simplified variant of the scheme proposed by Nikishin et al. (2022); **resetting the network optimizer state** at each target network update; adding **layer normalization**(Ba et al., 2016) after each convolutional and fully-connected layer of the CNN and the MLP; performing **Shrink and Perturb**(Ash and Adams, 2020): multiplying the network weights by a small scalar and adding a perturbation equal to the weights of a randomly initialized network; leveraging a **two-hot** encoding, which presents a distributional formulation of scalar regression wherein the network outputs a categorical probability distribution over fixed support and minimizes a cross-entropy loss with respect to an encoding of a regression target which distributes mass across two adjacent bins of the support; **spectral normalization** of the initial linear layer of the CNN and the MLP (Gogianu et al., 2021); and **weight decay,** setting the \(\ell_{2}\) penalty coefficient to \(10^{-5}\).

These methods were chosen to be representative samples of a number of approaches to mitigating plasticity loss: resetting the optimizer state and last layer temporarily remove a source of poor conditioning from the optimization process; layer normalization and residual connections tend to make networks more robust to optimizer choices; weight decay and spectral normalization both regularize the parameters of the network in different ways; shrink and perturb applies a perturbation to the current parameters without significantly changing the decision boundary (though we note that for regression tasks this will still influence the scale of the network outputs, and so may not be suitable ).

We visualize our key takeaways in Figure 6, which compares plasticity loss after 100 iterations of training on each of the architecture-intervention combinations. Overall, explicitly constructing a network parameterization which smooths out the loss landscape is the most effective means of preserving plasticity of all approaches we have considered, and has a greater effect on plasticity than resetting the final layer of the network. We visualize some learning curves of networks with and without layer normalization in Figure 17 in the supplementary material.

We note that while the two-hot encoding does demonstrate significant reductions in plasticity loss, it does so at the cost of stability of the learned policy in several instances we considered. Additionally, this intervention required significantly different optimizer hyperparameters from the regression parameterization, suggesting that while it can be a powerful tool to stabilize optimization, it might not be suitable as a plug-in solution to mitigate plasticity loss in an existing protocol.

### Application to larger benchmarks

We now evaluate whether the benefits of layer normalization on plasticity in toy classification tasks translate to larger-scale benchmarks. We use the standard implementation of double DQN (Van Hasselt et al., 2016) provided by Quan and Ostrovski (2020), and evaluate three seeds on each of the 57 games in the Arcade Learning Environment benchmark (Bellemare et al., 2013). We use the RMSProp optimizer, \(\epsilon\)-greedy exploration, and frame stacking (Mnih et al., 2015). Full implementation details can be found in Appendix A.3. The only difference between the baseline implementation and our modification is the incorporation of layer normalization after each hidden layer in the network.

We see in Figure 7 that the introduction of layer normalization robustly improves performance across the benchmark. We emphasize that we did not perform any optimizer or other hyper parameter tuning. While this improvement cannot be definitively attributed to a reduction in plasticity loss from the evidence provided, it points towards the regularization of the optimization landscape as a fruitful direction towards more robust RL agents. We further observe that many of the environments where layer normalization offers a significant boost to performance are those where the gradient covariance structure of the default architecture is degenerate or where the Hessian is ill-conditioned, and the LN networks which obtain performance improvements tend to have correspondingly better behaved gradient covariance. We provide a hint into this phenomenon in Figure 7, and defer the complete evaluation over all 57 games to Ap

Figure 6: Effect of architectural and optimization interventions on plasticity loss. Colour indicates change in loss on challenge targets between initial and final epoch of training on RL task. Darker shading indicates less plasticity loss.
pendix B.3.

## 7 Related Work

**Trainability:** the problem of finding suitable initializations for neural networks to enable training has a long history (Glorot and Bengio, 2010; He et al., 2015; Sutskever et al., 2013). Without careful initialization and architecture design, it is common to run into the issue that gradients will either explode or vanish as the depth of the network grows (Yang and Schoenholz, 2017). ResNets (He et al., 2016) in particular are known to resolve many of these pathologies by biasing each layer's mapping towards the identity function, leading to better-behaved gradients (Balduzzi et al., 2017). Mean-field analysis (Yang and Schoenholz, 2017; Schoenholz et al., 2017; Yang et al., 2019), information propagation (Poole et al., 2016), and deep kernel shaping (Zhang et al., 2021b; Martens et al., 2021) have all been applied to study trainability in neural networks. A wealth of prior work additionally studies the role of loss landscape smoothness in generalization and performance (Li et al., 2018; Santurkar et al., 2018; Ghorbani et al., 2019; Park and Kim, 2022). Other works highlight the chaotic behaviour of early training periods (Jastrzebski et al., 2020), in particular the 'edge of stability' phenomenon (Cohen et al., 2021) and the 'catapault mechanism' (Lewkowycz et al., 2020), and relate closely to the observations grounding 'linear mode connectivity' (Frankle et al., 2020) to explain generalization and trainability in deep neural networks; however, these approaches all focus on supervised learning with a stationary objective.

**Resetting + continual learning:**(Zhang et al., 2021a), (Berariu et al., 2021)(Hadsell et al., 2020; Rolnick et al., 2019). (Ostapenko et al., 2019) class-incremental learning, differs from our setting because the input distribution changes, not the functional relationship. Tangarasa et al. (2020) propose a modified Hebbian learning rule. Studies of plasticity in task-shift continual learning usually focus on ability to learn under new input distributions (Rolnick et al., 2019), rather than new targets. Most related to our study is the identification of the _loss_ of plasticity as a potentially limiting factor in deep reinforcement learning (Lyle et al., 2021; Dohare et al., 2021). This study can be motivated by the rich literature studying the effect of resetting and distillation on performance (Fedus et al., 2020; Nikishin et al., 2022; Igl et al., 2021; Schmitt et al., 2018).

## 8 Conclusions

The findings of this paper highlight a divide between the study of curriculum learning and foundation models, which identify suitable early training objectives to accelerate learning and improve generalization on later tasks, and the phenomena we have identified concerning the loss of plasticity in non-stationary prediction problems. However, as reinforcement learning algorithms scale up to more complex tasks, the divide between these regimes shrinks. While it is possible that in many settings, plasticity loss is not a limiting factor in network performance and so need not be a concern for many of the relatively small environments used to benchmark algorithms today, we conjecture that as the complexity of the tasks to which we apply RL grows, so will the importance of preserving plasticity.

The findings of this paper point towards stabilizing the loss landscape as a crucial step towards promoting plasticity. This approach is likely to have many ancillary benefits, presenting an exciting direction for future investigation. A smoother loss landscape is both easier to optimize and tends to exhibit better generalization, and it is an exciting direction for future work to better disentangle the complementary roles of memorization and generalization in plasticity.

Figure 7: Layer normalization improves performance and changes the gradient covariance structure in DDQN agents. Top: Human-normalized improvement score (Wang et al., 2016) of adding layer normalization over the default double DQN agent. Bottom: Gradient covariance matrices for Freeway (left) and Kangroo (right). In environments where layer normalization significantly improves performance, it also induces weaker gradient correlation.