Recovering Plasticity of Neural Networks via Soft Weight Rescaling
This addresses a fundamental optimization issue in deep learning that affects training stability and adaptability, though it appears incremental as it builds on known causes of plasticity loss.
The paper tackles the problem of plasticity loss in neural networks, where unbounded weight growth harms learning and generalization, by proposing Soft Weight Rescaling (SWR) to scale down weights during training, which improves performance in warm-start, continual, and single-task learning on image classification benchmarks.
Recent studies have shown that as training progresses, neural networks gradually lose their capacity to learn new information, a phenomenon known as plasticity loss. An unbounded weight growth is one of the main causes of plasticity loss. Furthermore, it harms generalization capability and disrupts optimization dynamics. Re-initializing the network can be a solution, but it results in the loss of learned information, leading to performance drops. In this paper, we propose Soft Weight Rescaling (SWR), a novel approach that prevents unbounded weight growth without losing information. SWR recovers the plasticity of the network by simply scaling down the weight at each step of the learning process. We theoretically prove that SWR bounds weight magnitude and balances weight magnitude between layers. Our experiment shows that SWR improves performance on warm-start learning, continual learning, and single-task learning setups on standard image classification benchmarks.