🤖 machine learning

Preserving Plasticity in Continual Learning via Dynamical Isometry

This paper establishes that dynamical isometry is a key mechanism for preserving plasticity in continual learning, proposing a novel regularization scheme and the AdamO optimizer to maintain near-isometric conditions, thereby preventing the progressive loss of learning capacity in deep neural networks.

Original authors: Andries Rosseau, Robert Müller, Ann Nowé

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andries Rosseau, Robert Müller, Ann Nowé

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Burnout" of AI

Imagine a student who is incredibly smart and learns a new subject every day. On Monday, they learn to play chess. On Tuesday, they learn to speak French. On Wednesday, they learn to code.

In a perfect world, this student would get better and better at everything. But in reality, deep learning AI models often suffer from "plasticity loss." This is like the student getting so tired or "burned out" that they can no longer learn new things. They might remember how to play chess, but when you ask them to learn French, they struggle to make progress. Their brain has become rigid; it can't bend to fit new shapes of information.

The authors of this paper ask: Why does this happen, and how do we stop it?

The Solution: Keeping the "Brain" Flexible (Dynamical Isometry)

The authors discovered that the secret to keeping an AI flexible is a mathematical concept called Dynamical Isometry.

The Analogy of the Water Hose:
Imagine the AI is a complex system of pipes (layers) carrying water (information) from a source to a destination.

The Problem: As the AI learns, the pipes sometimes get too narrow (squeezing the water too much) or too wide (letting the water spray out and lose pressure). If the water pressure changes wildly as it travels through the pipes, the AI gets confused and stops learning new tasks.
The Solution (Isometry): "Isometry" simply means "same measure." The authors want the pipes to be perfectly sized so that the water pressure stays exactly the same from the start of the pipe to the end. No squeezing, no spraying. Just a steady, consistent flow.

When the AI keeps this "steady flow" (Dynamical Isometry), it stays flexible. It can handle a new task (like learning French) just as easily as the first one (chess) because its internal "pipes" haven't gotten clogged or distorted.

How They Fixed It

The paper proposes three main ways to keep these pipes in perfect shape:

1. The "Perfectly Balanced" Architecture
They looked at a special type of AI design where the math is built to be perfectly balanced from the start. They used a specific type of "switch" (called GroupSort) that rearranges information without changing its volume.

Analogy: Imagine a dance troupe where the dancers constantly swap places but never leave the stage or change the total number of people. The energy of the group stays exactly the same, no matter how they move. This ensures the AI never loses its "vitality."

2. The "Self-Correcting" Rule (Regularization)
For standard AI designs (which aren't perfectly balanced by default), the authors added a special rule to the training process.

Analogy: Imagine a teacher walking around a classroom. Every time a student leans too far forward or backward (representing the AI's weights getting too big or too small), the teacher gently nudges them back to the center. This rule constantly checks the AI's "posture" and forces it to stay balanced, preventing it from getting rigid.

3. Waking Up the "Sleeping" Neurons
In standard AI, some parts of the brain often go to sleep (called "dead ReLUs") and stop working.

Analogy: The authors found that their "self-correcting rule" acts like a gentle shake. Because the rule forces the active parts of the brain to stay balanced, it accidentally pushes the "sleeping" parts to wake up and join the party again. This keeps the whole brain active and ready to learn.

The New Tool: AdamO

To make all this work smoothly, they created a new tool called AdamO.

Analogy: Think of standard AI training as a car where the engine (learning) and the brakes (regularization) are connected to the same pedal. If you press hard, you mess up both.
AdamO separates the pedals. It lets the engine run fast to learn new things, while the brakes gently and independently keep the car on the straight and narrow path. This allows the AI to learn quickly without losing its balance.

Did It Work?

The authors tested this on two types of challenges:

Supervised Learning: Teaching the AI to recognize images or solve puzzles where the rules keep changing.
Reinforcement Learning: Teaching an AI to play video games where the game levels change constantly.

The Result: In every test, the AI using their "steady flow" method (Dynamical Isometry) kept learning new things much longer than standard AI. It didn't get "burned out." It matched or beat all the other top methods, proving that keeping the internal "pipes" balanced is the key to a never-ending learner.

Summary

The paper argues that AI loses its ability to learn new things because its internal structure gets distorted over time. By forcing the AI to maintain a perfect balance (Dynamical Isometry) using new rules and a new training tool (AdamO), we can keep the AI flexible, awake, and ready to learn forever.

Technical Summary: Preserving Plasticity in Continual Learning via Dynamical Isometry

1. Problem Statement

Continual learning in deep neural networks under non-stationary conditions often suffers from a progressive loss of plasticity. This phenomenon manifests as a network's diminishing ability to adapt to new tasks, leading to degraded performance on subsequent data distributions. While specific failure modes are well-documented—such as "dying ReLUs," primacy bias in reinforcement learning (RL), and the growth of weight norms causing sharp loss landscapes—these are often studied in isolation. The authors argue that the common denominator across these failure modes is a deterioration in the efficiency of gradient descent to update the network in all relevant output-space directions. This is mathematically expressed as an increase in the anisotropy (condition number) of the empirical Neural Tangent Kernel (NTK).

2. Theoretical Framework

The paper establishes a functional definition of plasticity and links it to the geometry of the parameter-function map:

Plasticity Definition: Plasticity is defined not merely as the ability to minimize loss, but as the efficiency of eliminating the "excess achievable loss" (the gap between current performance and the optimal loss for the function class) within a fixed resource budget (e.g., optimization steps).
NTK and Anisotropy: Under first-order gradient descent, the local progress is governed by the Rayleigh coefficient of the empirical NTK. Task-agnostic plasticity requires the NTK to be isotropic (proportional to the identity matrix) so that gradient descent can make comparable progress in all output-space directions.
Dynamical Isometry as a Surrogate: Exact NTK isotropy is unattainable for finite networks. However, the authors propose dynamical isometry as a tractable, task-agnostic surrogate. A network is dynamically isometric if the singular values of the layer-wise input-output Jacobians remain concentrated near one. This ensures that perturbations propagate through the network without systematic expansion or contraction, keeping the NTK spectrum well-conditioned in expectation over an isotropic input prior.

3. Methodology

The authors propose a multi-faceted approach to preserve dynamical isometry and, consequently, plasticity:

Architectural Existence Proof: The paper revisits GroupSort networks (specifically MaxMin activations) combined with orthonormal linear layers. These networks are proven to be almost-everywhere (a.e.) isometric. Despite being isometric, they remain universal Lipschitz function approximators, demonstrating that strict isometry constraints are compatible with expressive nonlinear representations.
Regularization for General Architectures: For standard networks using element-wise activations (e.g., ReLU) where exact isometry is impossible, the authors introduce an efficient isometry-promoting regularizer. This penalty drives the singular values of weight matrices toward one by minimizing the Frobenius norm deviation of the Gram matrix ( $W^\top W - I$ or $WW^\top - I$ ).
ReLU Revival Mechanism: A novel contribution is the identification of a mechanism by which this regularizer reactivates "dead" ReLU units. Because the regularizer couples the rows of weight matrices to maintain orthogonality, the movement of active neurons under task loss induces a rotation of inactive weight vectors. This can push previously dormant units back into regions of positive pre-activation, effectively reviving them without explicit task-driven gradients.
AdamO Optimizer: To prevent the mixing of task gradients and regularizer gradients within adaptive moment estimates (which can lead to unstable rescaling), the authors introduce AdamO. Analogous to AdamW's decoupled weight decay, AdamO decouples the isometry regularization step from the adaptive gradient update. It updates the Adam moments using only the task gradient, then applies a separate step for the isometry penalty.

4. Experimental Results

The proposed methods were evaluated across supervised and reinforcement learning continual learning benchmarks designed to induce plasticity loss:

Supervised Learning: On benchmarks including Random-Label Memorization (testing raw optimization capacity), Permuted MNIST/CIFAR, and Label-Shuffled CIFAR, the authors' methods (specifically AdamO with ReLU or MaxMin) consistently matched or outperformed existing baselines like Normalize-and-Project (NaP), Spectral Norm regularization, and ReDo. The methods demonstrated superior ability to maintain training accuracy and generalization over long sequences of tasks.
Reinforcement Learning: In continual MinAtar and Octax environments with random channel permutations and task switching, the AdamO-based approaches maintained consistent performance over 20 cycles (1.2 billion steps total). In contrast, standard PPO baselines exhibited significant performance degradation over time.
Diagnostics: The experiments confirmed that the proposed methods successfully maintained lower condition numbers, higher effective NTK rank, and sustained ReLU activity compared to baselines, validating the theoretical link between isometry and plasticity.

5. Key Contributions

Functional Definition of Plasticity: A formalization of plasticity in terms of the NTK geometry and the efficiency of eliminating excess achievable loss.
Dynamical Isometry as a Target: Identifying layer-wise dynamical isometry as a principled, task-agnostic mechanism to preserve plasticity.
Compatibility with Expressivity: Proving that almost-everywhere isometric architectures (GroupSort) are universal approximators, resolving concerns that isometry constraints limit model capacity.
Practical Training Tools:
- An efficient isometry regularizer for general architectures.
- A mechanism explaining and leveraging the revival of dead ReLU units via orthogonality constraints.
- AdamO, a decoupled optimizer specifically designed to handle isometry regularization without corrupting adaptive gradient statistics.

6. Significance and Claims

The paper claims that preserving dynamical isometry provides a unified geometric explanation for various plasticity loss phenomena and offers a robust solution for continual learning. By keeping the network near the isometric manifold, the system retains the capacity to make nontrivial progress from current parameters regardless of the current objective.

The authors position their work as a bridge between initialization theory (where isometry is known to aid signal propagation) and continual learning (where isometry preservation prevents plasticity loss). They modestly note that while their methods outperform existing approaches, they do not claim to solve all aspects of continual learning (such as catastrophic forgetting, which they distinguish from plasticity loss). Future directions suggested include applying these principles to deep reinforcement learning agents and large language models, where maintaining isometry could stabilize training in deeper architectures and improve continual adaptation.