Recovering Plasticity of Neural Networks via Soft Weight Rescaling

TL;DR

This paper introduces Soft Weight Rescaling (SWR), a method to prevent unbounded weight growth in neural networks, thereby recovering plasticity and improving learning performance without losing learned information.

Contribution

The paper proposes SWR, a novel technique that bounds weight magnitudes and maintains network plasticity, with theoretical proofs and empirical validation across various learning scenarios.

Findings

SWR effectively bounds weight magnitudes during training.

SWR improves performance in continual and warm-start learning.

SWR maintains learned information while enhancing plasticity.

Abstract

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.