Source-linked AI summary
Understanding plasticity in neural networks
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, Will Dabney
TL;DR
Plasticity loss limits neural networks’ ability to adapt to new information, especially in the non-stationary objectives inherent to deep reinforcement learning. This paper systematically analyzes its mechanisms, finding a central connection to loss-landscape curvature and identifying design choices that better preserve plasticity, with validation on larger-scale RL benchmarks.
Problem
Plasticity loss is well documented in neural networks, but its mechanisms remain poorly understood despite its importance for adapting to changing deep reinforcement-learning environments.
Method
The paper combines interpretable case studies, a falsification framework, broad intervention experiments, and Arcade Learning Environment validation to study plasticity mechanisms and preservation.
Results
Loss of plasticity is closely connected to loss-landscape curvature, while architectural choices such as categorical outputs and normalization provide the greatest improvements and layer normalization improves standard DQN performance on the Arcade Learning Environment.
Takeaways & Limitations
Controlling loss-landscape sharpness and optimizer stability is a promising route toward improving the robustness and usability of deep reinforcement-learning methods.
Takeaways & Limitations
Two-hot encoding can reduce plasticity loss at the cost of learned-policy stability and requires substantially different optimizer hyperparameters, limiting its suitability as a plug-in intervention.
Abstract
from arXiv · showhide
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.
1. Introduction
Plasticity loss is especially relevant to deep reinforcement learning because non-stationary objectives require networks to overwrite prior predictions. The paper investigates its mechanisms and finds that loss-landscape curvature is central, motivating targeted architectural and optimization choices.
- Plasticity loss is most robust when changing input-target relationships require networks to overwrite prior predictions.Such non-stationarity is uncommon in supervised learning but inherent to deep reinforcement learning.
- Existing interventions improve performance, but their differing mechanisms make further improvement difficult to target.Examples include resetting layers or units and regularizing features.
- The paper uses interpretable case studies and a falsification framework to investigate candidate mechanisms of plasticity loss.The framework shows that loss of plasticity cannot be uniquely attributed to the proposed network properties.
- Loss-landscape curvature is identified as a crucial factor determining plasticity, particularly in value-based reinforcement learning.
- Architectures conjectured to smooth the loss landscape provide the greatest plasticity improvements, while parameter perturbations and other regularization tend to help less.Layer normalization substantially improves performance when applied to a standard DQN on the Arcade Learning Environment benchmark.
2. Background
The paper distinguishes plasticity loss from catastrophic forgetting and defines plasticity as the ability to adapt to new learning signals. It studies this ability through temporal-difference learning and loss-landscape quantities including Hessian spectra and gradient covariance.
- Plasticity loss concerns worse performance on later tasks than a same-architecture randomly initialized network, unlike catastrophic forgetting’s reduced performance on earlier tasks.
- Temporal difference learning: Temporal-difference learning is non-stationary because its regression target depends on changing target parameters, even with fixed policy and input distribution.
- Loss landscape analysis: Loss-landscape analysis examines the Hessian eigenspectrum and gradient covariance, whose sharpness, conditioning, and interference relate to optimization.Negative off-diagonal gradient-covariance entries indicate interference between inputs.
- Defining plasticity: Plasticity broadly means a neural network’s ability to learn new things, while representational capacity alone does not guarantee trainability by gradient descent.Saturated activation units can prevent optimization from realizing the functions an architecture can represent.
- Defining plasticity: The paper operationalizes plasticity as the expected improvement obtained by optimizing sampled objectives from a checkpoint, relative to a baseline loss.
- Defining plasticity: Plasticity loss along a training trajectory is the change in plasticity from the initial checkpoint, independent of the probe task’s baseline difficulty.
3. Methodology and Motivating Questions
The methodology combines probe-task measures of adaptability with controlled reinforcement-learning environments and experiments targeting mechanisms, causes, and interventions. It evaluates loss-landscape hypotheses and tests methods intended to preserve plasticity.
- Motivating questions: The experiments separate causal pathways of plasticity loss and evaluate hypotheses about its root causes.
- Measuring plasticity: Plasticity interventions are evaluated using a generic distribution of target functions representing possible future learning signals.
- Measuring plasticity: A challenging sinusoidal prediction objective perturbs network predictions in effectively uniform random directions over an RL replay buffer.
- RL testbed: The block-MDP testbed uses ten states and ten actions with CIFAR-10 or MNIST observations, spanning true-label, random-label, and sparse-reward variants.
- RL testbed: The environments vary reward smoothness and density while isolating non-stationarity from policy and target-network updates independently of state-visitation changes.
- Mechanisms and causes: The analysis tests whether optimization dynamics sharpen the loss landscape and evaluates weight norm, feature rank, dead units, and learning curves as candidate explanations.
- Mitigation: Interventions focus on architectural and optimization methods known to increase loss-landscape smoothness, using the same plasticity evaluation across the classification-MDP testbed.
4. Two Simple Studies on Plasticity
The studies show that non-stationarity can destabilize adaptive optimization and that gradient-based training can make the loss landscape less amenable to learning arbitrary new objectives.
- 4.1. Optimizer instability and non-stationarity: Abrupt task changes destabilize adaptive optimizers when second-order moment estimates lag behind sudden gradient changes.Increasing epsilon and using a more aggressive decay rate for the second-moment estimate avoids the observed catastrophic instability.
- 4.2. Loss landscape evolution under non-stationarity: Even without saturated units, optimization can reduce plasticity by moving parameters into regions that are harder to optimize toward arbitrary new objectives.This motivates analyzing optimization dynamics and loss-landscape structure beyond unit saturation.
- 4.2. Loss landscape evolution under non-stationarity: Gradient descent produces faster-growing Hessian spectral outliers than equal-norm random parameter perturbations.Both processes increase Hessian spectral norm, but the outliers grow significantly faster after gradient descent.
- 4.2. Loss landscape evolution under non-stationarity: Gradient descent also induces negative gradient interference between inputs, which is absent from the random-walk comparison.The comparison uses coupled trajectories with identical starting parameters and equal-norm updates differing only in direction.
5. Explaining Plasticity Loss
The analysis finds that common network statistics do not provide robust causal explanations of plasticity loss, while training increasingly impairs navigation of the probe-task loss landscape.
- 5. Explaining Plasticity Loss: Candidate explanations based on feature rank, saturated units, parameter norm, and weight-matrix rank do not identify robust causal relationships.The paper instead finds evidence that plasticity loss is associated with changes in the network’s loss landscape.
- 5.2. Falsification of prior hypotheses: Across learning problems, each tested statistic can correlate either positively or negatively with plasticity, weakening its causal interpretation.The experiments span tasks, observation spaces, optimizers, and random seeds, and include 128 DQN agents.
- 5.2. Falsification of prior hypotheses: Weight norm correlates positively with plasticity loss for CIFAR-10 observations but slightly negatively for MNIST observations.Analogous correlation reversals occur for the other variables considered.
- 5.3. Loss landscape evolution during training: Plasticity loss reflects slower probe-task training rather than convergence to higher loss plateaus.Later checkpoints produce shallower learning-curve slopes and increasing variance, indicating greater difficulty navigating the loss landscape.
- 5.3. Loss landscape evolution during training: Increasing difficulty navigating the loss landscape drives plasticity loss in the studied problem.The evidence includes increasing loss-landscape sharpness and, in mini-batch training, increasing interference between minibatches.
6. Solutions
Scaling reduces plasticity loss but does not eliminate it, even at sizes exceeding a single GPU. Across tested interventions, parameterizations that smooth the loss landscape provide the strongest preservation, while two-hot encoding introduces stability and tuning trade-offs.
- 6.1. The role of scaling on plasticity: Scaling consistently reduces plasticity loss across target update frequencies, but cannot completely eliminate it even when the architecture exceeds a single GPU’s memory.This result holds across several architecture-dataset combinations on simple classification-inspired problems.
- 6.1. The role of scaling on plasticity: Sufficiently large, overparameterized networks may avoid plasticity loss on sufficiently simple tasks, whereas limited expressivity or poor task-inductive-bias alignment still reduces adaptation.The latter pattern is observed for small networks and MLPs trained on CIFAR-10.
- 6.2. Interventions in toy problems: The intervention study compares MLP, CNN, ResNet-18, and small transformer architectures using resets, normalization, perturbation, and regularization methods.The protocol repeats 100 iterations of 1000 training steps.
- 6.2. Interventions in toy problems: Smoothing the loss landscape through network parameterization is the most effective tested approach for preserving plasticity, sometimes outperforming final-layer resets.The comparison is summarized after 100 training iterations across architecture-intervention combinations.
- 6.2. Interventions in toy problems: Two-hot encoding reduces plasticity loss but can destabilize learned policies and requires optimizer hyperparameters substantially different from regression encoding.These trade-offs make it unsuitable as a general plug-in intervention in the considered protocols.
- 6.3. Application to larger benchmarks: Layer normalization robustly improves performance across 57 Arcade Learning Environment games without additional hyperparameter tuning.Improved environments tend to show less degenerate gradient covariance and better-conditioned optimization structure under layer normalization.
7. Related Work
Related work distinguishes plasticity loss from catastrophic forgetting and situates it within trainability, continual learning, and non-stationary reinforcement-learning research. Existing trainability analyses largely study stationary supervised objectives, whereas this paper focuses on changing prediction targets in RL.
- Trainability: Trainability research studies initialization and architecture choices that prevent vanishing or exploding gradients, including ResNets’ identity-biased mappings.These works address gradient pathologies rather than the paper’s non-stationary-target setting.
- Trainability: Related supervised-learning analyses examine edge-of-stability, catapult dynamics, and linear mode connectivity under stationary objectives.The paper contrasts those settings with plasticity loss driven by non-stationary prediction targets.
- Continual learning: Continual learning covers input-distribution and covariate shifts, while this paper focuses on non-stationary targets as a driver of performance plateaus in reinforcement learning.The work builds on studies linking plasticity loss to deep RL and on research about resetting and distillation.
8. Conclusions
The paper connects plasticity preservation to stabilizing the loss landscape and contrasts this with how unsuitable objectives can impair adaptation in non-stationary prediction. It argues that preserving plasticity may become more important as reinforcement-learning tasks grow more complex.
- 8. Conclusions: Unsuitable objectives can hurt adaptation to new learning signals in non-stationary prediction problems, unlike suitable pretraining objectives that can improve later-task adaptation and generalization.The paper presents this as a contrast between two learning regimes.
- 8. Conclusions: Plasticity loss may not limit performance on many relatively small benchmark environments, but its importance is conjectured to increase with reinforcement-learning task complexity.The paper notes that the divide between large-model pretraining and non-stationary RL may shrink as tasks become more complex.
- 8. Conclusions: Stabilizing the loss landscape is identified as a crucial direction for promoting plasticity and may also benefit optimization and generalization.The paper calls for further work disentangling memorization and generalization in plasticity.
A. Experiment details
The experiments use classification-inspired and reinforcement-learning settings with multiple neural architectures, optimization analyses, and intervention comparisons. Supplementary figures extend the intervention analysis to categorical output encodings on CIFAR-10.
- Experimental settings: The MNIST memorization experiment repeatedly randomizes labels on a fixed subset of 5000 images while logging accuracy on 4096 states.It uses a two-hidden-layer MLP with width 1024 and an Adam optimizer.
- Experimental settings: The Brownian-motion experiment trains on a Q-learning loss using stochastic gradient descent with learning rate 0.001 and batch size 512.The environment uses an MNIST observation space.
- Optimization analyses: Hessian eigenvalue density is approximated using a Gaussian-smoothed Lanczos estimate from a single large transition batch.The implementation follows Ghorbani et al. (2019).
- Architectures: The architecture set includes MLP, CNN, ResNet-18, and a scratch-trained vision transformer with specified widths, convolutional structures, patch size, and dropout.The MLP width sweep uses factors 1, 2, 4, 8, 12, and 16 over a base width of 16.
- Atari protocol: Atari training runs for 200 million frames with optimizer updates every four environment steps, replay buffer size 100,000, epsilon-greedy exploration, and layer normalization after hidden layers.The benchmark protocol uses epsilon = 0.1.
- Output encodings: CIFAR-10 intervention analyses compare regression and categorical output encodings and report significant categorical-encoding benefits for transformer, CNN, and MLP architectures.The analysis repeats the Figure 6 intervention comparison.
B.1. Detailed intervention analysis
Intervention analysis finds that resetting the last layer, two-hot outputs, and layer normalization improve plasticity, while several perturbation and regularization methods do not. Layer normalization also stabilizes loss-landscape sharpness, though weight-perturbing methods can interfere with primary-task learning.
- Resetting the last layer, two-hot output representations, and layer normalization improve plasticity, whereas shrink and perturb, weight decay, and optimizer-state resetting do not.
- Categorical output representations can improve trainability from initialization and combine with layer normalization, resetting, and shrink and perturb.
- Weight-perturbing interventions often interfere with primary-task learning, particularly under the harder reward structure.
- Layer normalization prevents the sharpness increase observed in the default convolutional network trained with regression loss.
B.2.1. TRAINING ACCURACY
In the classification MDP experiments, TD losses can look similar even when easy- and hard-task performance differs substantially. Layer normalization and two-hot encodings support more stable learning, but stable predictions do not always yield optimal policies.
- TRAINING ACCURACY: Easy- and hard-task performance differs dramatically across many agents, while their TD losses behave more similarly.The passage attributes this partly to fewer rewards producing smaller TD-loss changes and easier learning problems.
- TRAINING ACCURACY: Layer normalization and two-hot encodings enable significantly more stable learning according to temporal-difference losses.TD losses also indicate whether an agent’s predictions are diverging.
- TRAINING ACCURACY: Stable TD losses do not always correspond to optimal policies.
- TRAINING ACCURACY: Probe-task learning curves from different training checkpoints illustrate the subtlety of measuring plasticity loss.
B.3. Qualitative findings in DDQN
Qualitative DDQN analyses associate learning progress with less degenerate, more diagonal-weighted gradient covariance structures, while Hessian spectra show no obvious performance correlation. The role of gradient degeneracy remains unresolved.
- Qualitative findings in DDQN: Gradient covariance structures vary substantially across environments, networks, and random seeds.
- Qualitative findings in DDQN: Learning-progress seeds in Freeway show heavier diagonal gradient covariance, whereas non-learning seeds preserve initially degenerate gradient structure.
- Qualitative findings in DDQN: Networks with layer normalization generally exhibit a slight bias toward less degenerate gradients.
- Qualitative findings in DDQN: The analysis does not establish whether gradient degeneracy is a symptom or a cause of performance plateaus.
- Qualitative findings in DDQN: The Hessian spectrum shows no obvious correlation with agent performance in the reported single-seed visualizations.