Preserving Plasticity in Continual Learning with Adaptive Linearity Injection

TL;DR

AdaLin addresses plasticity loss in continual learning by adaptively injecting linearity into neuron activations via a per-neuron learnable parameter and a gradient-based gating mechanism. The method fuses a base activation φ with a tunable linear term controlled by α_i and gated by g(x), preserving gradient flow especially in saturated regions, and it applies across architectures from MLPs to ResNets with negligible overhead. Empirically, AdaLin improves average online and test accuracy on plasticity benchmarks (e.g., Permuted MNIST, Random Label MNIST, Shuffled CIFAR-10, Class-Split CIFAR-100) and in class-incremental CIFAR-100 and off-policy RL tasks, outperforming activation-based and regularization baselines; neuron-level adaptation is shown to be essential. The work highlights a practical route to lifelong learning by revisiting activation design and suggests directions for scaling to larger models and integrating with complementary forgetting mitigation strategies.

Abstract

Loss of plasticity in deep neural networks is the gradual reduction in a model's capacity to incrementally learn and has been identified as a key obstacle to learning in non-stationary problem settings. Recent work has shown that deep linear networks tend to be resilient towards loss of plasticity. Motivated by this observation, we propose Adaptive Linearization (AdaLin), a general approach that dynamically adapts each neuron's activation function to mitigate plasticity loss. Unlike prior methods that rely on regularization or periodic resets, AdaLin equips every neuron with a learnable parameter and a gating mechanism that injects linearity into the activation function based on its gradient flow. This adaptive modulation ensures sufficient gradient signal and sustains continual learning without introducing additional hyperparameters or requiring explicit task boundaries. When used with conventional activation functions like ReLU, Tanh, and GeLU, we demonstrate that AdaLin can significantly improve performance on standard benchmarks, including Random Label and Permuted MNIST, Random Label and Shuffled CIFAR-10, and Class-Split CIFAR-100. Furthermore, its efficacy is shown in more complex scenarios, such as class-incremental learning on CIFAR-100 with a ResNet-18 backbone, and in mitigating plasticity loss in off-policy reinforcement learning agents. We perform a systematic set of ablations that show that neuron-level adaptation is crucial for good performance and analyze a number of metrics in the network that might be correlated to loss of plasticity.

Figures (17)

• Figure 1: (a) Baseline comparison of activation functions and their vulnerability to plasticity loss in the random-label MNIST task. (b) Improved performance when combining AdaLin with these activations.
• Figure 2: Illustration of AdaLin’s adaptive activation mechanism. (Left) A neural network structure where a parameter $\alpha$ is learned per each neuron that modulates its activation between linear and non-linear components. Here, ${h}^i_{\text{out}}$ and $z_i$ represent the post-activation and pre-activation values of neuron $i$, respectively. $\phi(\cdot)$ is a base activation function, and $g(\cdot)$ is a gating function computed from the gradient of $\phi(\cdot)$ (i.e., it automatically adjusts based on how close $\phi(\cdot)$ is to saturation). (Right) Plots demonstrating the behavior of the learned activation function, where AdaLin smoothly interpolates between a standard nonlinearity (e.g., $\tanh(x)$, red) and a linear function (green). This injection of linearity is crucial in saturated regions (where the derivative of $\phi(\cdot)$ is small), helping sustain gradient flow and preserve plasticity. In the top figure, $z_i$ falls in the non-saturated region of $\tanh(x)$ (red), so no linearity is added. In the bottom figure, $z_{i+1}$ lies in a saturated region (green), prompting the gating function $g(\cdot)$ to enable linear injection based on $\alpha_{i+1}$.
• Figure 3: Derivative of AdaLin with $\tanh$ as the base activation function under varying $\alpha$. For $\alpha=0$, this gives the original $\tanh$ function and we see that when x lies in the saturated region of $\tanh$ (shaded in gray), the derivative is zero, limiting the gradient flow through the corresponding neuron. As $\alpha$ takes a non-zero value, a linear term is injected, ensuring a nonzero derivative for all possible input values.
• Figure 4: Average Task Accuracy across five plasticity benchmarks (Permuted MNIST, Random Label MNIST, Shuffle Label CIFAR-10, Class-Split CIFAR-100, and Random Label CIFAR-10) for various methods, all using ReLU as the base activation except the Deep Linear model. Overall, AdaLin consistently maintains plasticity and often outperforms other approaches, particularly in tasks where transferring knowledge can improve learning (e.g., Permuted MNIST, Shuffled Label CIFAR-10, and Class-Split CIFAR-100). Methods like Scratch and Shrink & Perturb struggle due to frequent resetting, while baseline ReLU models degrade severely in random-label tasks. Although Deep Linear does not exhibit plasticity loss, its limited expressivity leads to lower overall performance.
• Figure 5: Test accuracy for tasks where generalization is particularly relevant (Permuted MNIST, Shuffle CIFAR-10, and Class-Split CIFAR-100). AdaLin achieves the highest test accuracy in Class-Split CIFAR-100 and performs strongly on Shuffle CIFAR-10 and Permuted MNIST. Deep Fourier struggles with generalization over time in Shuffle CIFAR-10, while Deep Linear, despite avoiding plasticity loss in training, fails to generalize effectively in this setting.