Loss of plasticity refers to the progressive inability of a model to adapt to new tasks and poses a fundamental challenge for continual learning. While this phenomenon has been extensively studied in homogeneous neural architectures, such as multilayer perceptrons, its mechanisms in structurally heterogeneous, attention-based models such as Vision Transformers (ViTs) remain underexplored. In this work, we present a systematic investigation of loss of plasticity in ViTs, including a fine-grained diagnosis using local metrics that capture parameter diversity and utilization. Our analysis reveals that stacked attention modules exhibit increasing instability that exacerbates plasticity loss, while feed-forward network modules suffer even more pronounced degradation. Furthermore, we evaluate several approaches for mitigating plasticity loss. The results indicate that methods based on parameter re-initialization fail to recover plasticity in ViTs, whereas approaches that explicitly regulate the update process are more effective. Motivated by this insight, we propose ARROW, a geometry-aware optimizer that preserves plasticity by adaptively reshaping gradient directions using an online curvature estimate for the attention module. Extensive experiments show that ARROW effectively improves plasticity and maintains better performance on newly encountered tasks.
Figures & tables
Figure 1: (a) ViTs exhibit plasticity loss in task-incremental CIFAR-100, with AAT degrading over the task stream. (b) ARROW, a geometry-aware optimizer with an online low-rank curvature estimate, stabilizes optimization rather than exhibits oscillatory updates under distribution shift. Both ARROW and TRAC outperform the vanilla ViT, and ARROW’s advantage over TRAC increases in later tasks (inset).
Figure 2: Architecture of the ViT.
Figure 3: Performance comparison regarding the AAT of the MLP and different ViT architectures. All the models exhibit significant performance loss as the task sequence progresses, indicating a broad phenomenon of plasticity loss. The solid line and shaded region represent the mean and standard deviation over multiple runs, respectively.
Figure 4: (a) The effective rank of features produced by different components in the 2nd block (left) and the 8th block (right) of the ViT. (b) Δ -heatmap (left) of the effective rank differences between the 1st block and the 12th block in ViT. The right figure shows the weight magnitude variation from the 3rd and 11th blocks.
Figure 5: (a) The effective rank (left) and weight magnitude (right) of one of the attention heads. The green curve represents the V matrix, which is most unstable. The blue and red curves represent the Q and K , respectively. (b) The effective rank of CLS token across sampled tasks. The values on embedding layers, shallow and deep blocks are tracked sequentially. Detailed component labels of ViT are provided in the Appendix.
Figure 6: (a) Performance comparison between vanilla ViT and ViT with CBPs. For CBP1–CBP5, the maturity thresholds are 400, 500, 1200, 600, and 1200. The replacement rates are 8×10−6 , 1×10−5 , 3×10−6 , 5×10−6 , and 8×10−7 . CBP2 and CBP3 apply CBP to the first block, whereas CBP1, CBP4, and CBP5 apply it to all blocks. (b) Performance comparison between vanilla ViT and representative plasticity loss mitigation methods. (c) Weight magnitude of the 1st block (i) and 12th block (iii), and the corresponding effective rank of the 1st block (ii) and 12th block (iv).
α
W=10
W=20
W=30
1
–
59.51 ( β =1, SGD)
–
0.5
–
58.08 ( β =1, SGD)
59.51 ( β =1, SGD)
0.05
–
60.45 ( β =1, SGD)
–
0.01
–
60.51 ( β =1, RMS)
–
0.003
–
60.51 ( β =1, RMS)
–
0.001
58.50 ( β =1, SGD)
60.41 ( β =0.9, RMS)
58.75 ( β =1, SGD)
Table 1: Hyperparameter study of ARROW. Left: Performance across different α , β , window sizes W and warm-up modes. Right: Fine-grained tuning with fixed W=20 and RMS-like warm-up on CIFAR100 (10 tasks).
Method
Location
α
β
AAT (%)
ARROW
Last 1
1.0×10−3
0.9
64.21
ARROW_b0
ALL
1.0×10−3
0
59.85
ARROW_a1
ALL
1.0
0.9
59.31
ARROW_6
Last 6
1.0×10−3
0.9
63.94
ARROW_1
ALL
1.0×10−3
0.9
63.96
Table 2: Ablation study results of ARROW with controlling terms of α,β , and effected blocks.
CIFAR100
ImageNet-R
Method
10 Task
20 Task
25 Task
20 Task
40 Task
50 Task
Baseline
57.21
66.53
70.93
28.16
36.94
40.61
CBP
57.43
66.18
67.09
27.27
36.35
41.96
NaP
58.62
66.93
71.46
28.20
37.35
42.06
L2P
60.17
67.54
71.13
30.83
35.06
40.88
TRAC
63.58
73.32
72.19
27.98
37.17
41.05
Table 3: Comparison of plasticity performance (AAT %) on CIFAR100 and ImageNet-R across different length of task streams.
Appendix figures & tables28 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Stable rank and effective rank on first block (column left) and second block (column right) of ViT. The rows of sub-figures correspond with normalization 1, attention output, residual structure 1, normalization 2, FFN’s FC layer 1, FFN’s FC layer 2, and the residual structure 2 of vanilla ViT.
Figure 8: Stable rank and effective rank on 3rd block (column left) and 4th block (column right) of ViT. The rows of sub-figures correspond with normalization 1, attention output, residual structure 1, normalization 2, FFN’s FC layer 1, FFN’s FC layer 2, and the residual structure 2 of vanilla ViT.
Figure 9: Stable rank and effective rank on 11th block (column left) and 12th block (column right) of ViT. The rows of sub-figures correspond with normalization 1, attention output, residual structure 1, normalization 2, FFN’s FC layer 1, FFN’s FC layer 2, and the residual structure 2 of vanilla ViT.
Figure 10: Stable rank and effective rank on embedding layer (column left) and classification head (column right) of vanilla ViT.
Figure 11: Effective rank (left) and weight magnitude (right) of attention projections and FFN of ViT, block 1. The blue curves represent the metrics of attention projection. The red and green curves represent the first and second FC layers in FFN.
Figure 12: Effective rank (left) and weight magnitude (right) of attention projections and FFN of ViT, block 2. The blue curves represent the metrics of attention projection. The red and green curves represent the first and second FC layers in FFN.
Figure 13: Effective rank (left) and weight magnitude (right) of attention projections and FFN of ViT, block 3. The blue curves represent the metrics of attention projection. The red and green curves represent the first and second FC layers in FFN.
Figure 14: Effective rank (left) and weight magnitude (right) of attention projections and FFN of ViT, block 11. The blue curves represent the metrics of attention projection. The red and green curves represent the first and second FC layers in FFN.
Figure 15: Effective rank (left) and weight magnitude (right) of attention projections and FFN of ViT, block 12. The blue curves represent the metrics of attention projection. The red and green curves represent the first and second FC layers in FFN.
Figure 16: The effective rank of 12 attention heads in ViT, block 1. The green curve represents the V matrix, which is the most unstable one. The blue and red curves represent the Q and K , respectively.
Figure 17: The effective rank of 12 attention heads in ViT, block 2. The green curve represents the V matrix, which is the most unstable one. The blue and red curves represent the Q and K , respectively.
Figure 18: The effective rank of 12 attention heads in ViT, block 3. The green curve represents the V matrix, which is the most unstable one. The blue and red curves represent the Q and K , respectively.
Figure 19: The effective rank of 12 attention heads in ViT, block 11. The green curve represents the V matrix, which is the most unstable one. The blue and red curves represent the Q and K , respectively.
Figure 20: The effective rank of 12 attention heads in ViT, block 12. The green curve represents the V matrix, which is the most unstable one. The blue and red curves represent the Q and K , respectively.
Figure 21: The weight magnitude of 12 attention heads in ViT, block 1. The green curve represents the weight magnitude of V matrix. The blue and red curves represent the weight magnitude of Q and K , respectively.
Figure 22: The weight magnitude of 12 attention heads in ViT, block 2. The green curve represents the weight magnitude of V matrix. The blue and red curves represent the weight magnitude of Q and K , respectively.
Figure 23: The weight magnitude of 12 attention heads in ViT, block 3. The green curve represents the weight magnitude of V matrix. The blue and red curves represent the weight magnitude of Q and K , respectively.
Figure 24: The weight magnitude of 12 attention heads in ViT, block 11. The green curve represents the weight magnitude of V matrix. The blue and red curves represent the weight magnitude of Q and K , respectively.
Figure 25: The weight magnitude of 12 attention heads in ViT, block 12. The green curve represents the weight magnitude of V matrix. The blue and red curves represent the weight magnitude of Q and K , respectively.
Figure 26: The effective rank of each component of cls token in ViT during continual learning. The order of components starts from embedding layers, 12 ViT blocks, and classification head. The sampled tasks are task 1 (blue), task 51 (orange), task 101 (green), task 151 (red), and task 201 (purple).
Figure 27: The FAU (blue) and FDU (orange) in activation function of FFN in ViT’s 12 blocks. The significant decrease of FAU and the rise of FDU demonstrate the loss of adaptability in FFN module.
Figure 28: The accuracy curves of Upper bound (red), vanilla ViT (brown), and ViT with CReLU (orange) on first 25 tasks of stream. The AAT of CReLU is 0.617, which is less than vanilla ViT (0.638) and Upper bound (0.732), indicating the worse performance of CReLU. Moreover, the cost of CReLU is double than other approaches.
a1w20b1
a1e-1w20b01rms
a5e-1w20b1
a5e-2w20b1
a1e-2w20b1
a3e-3w20b1
a1e-3w20b1
a1e-4w20b1
task1
0.595
0.576
0.632
0.607
0.603
0.607
0.595
0.595
task2
0.658
0.659
0.652
0.648
0.669
0.648
0.658
0.658
task3
0.639
0.655
0.600
0.642
0.651
0.642
0.639
0.639
task4
0.657
0.667
0.620
0.661
0.665
0.661
0.657
0.657
task5
0.609
0.625
0.602
0.613
0.602
0.613
0.609
0.609
task6
0.583
0.606
0.562
0.627
0.629
0.627
0.583
0.583
Appendix
Table 4: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
a1e-2w20b1rms2trac
a1e-3w20b12trac
a3e-3w20b1rms_2trac
a5e-2w20b1rms_2trac
a1e-3w20b0.9rms2trac
task1
0.587
0.612
0.593
0.619
0.611
task2
0.644
0.647
0.643
0.638
0.665
task3
0.673
0.666
0.668
0.670
0.669
task4
0.686
0.653
0.665
0.658
0.672
task5
0.644
0.651
0.643
0.646
0.654
task6
0.644
0.639
0.644
0.649
0.649
Appendix
Table 5: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
c1e-2-3w20b1rmspertask
c1e-2-3w20b1rms2trcapertask
trac
baseline
task1
0.620
0.614
0.603
0.632
task2
0.669
0.658
0.653
0.652
task3
0.655
0.677
0.656
0.606
task4
0.647
0.680
0.683
0.617
task5
0.621
0.658
0.656
0.575
task6
0.614
0.646
0.648
0.542
Appendix
Table 6: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
a1w20b0.8
a5e-1w20b0.9
a1e-2w20b0.5
a3e-3w20b1rms
a1e-3w20b0.9
task1
0.595
0.632
0.607
0.576
0.595
task2
0.658
0.652
0.648
0.659
0.658
task3
0.639
0.600
0.642
0.655
0.639
task4
0.657
0.620
0.661
0.667
0.657
task5
0.609
0.607
0.613
0.625
0.609
task6
0.583
0.563
0.627
0.606
0.583
Appendix
Table 7: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
a5e-1w30b1
a1e-2w20b1rms
a1e-3w20b0.9rms
a1e-3w10b1
task1
0.595
0.576
0.576
0.634
task2
0.658
0.659
0.659
0.663
task3
0.639
0.655
0.655
0.635
task4
0.657
0.677
0.667
0.654
task5
0.609
0.625
0.625
0.587
task6
0.583
0.606
0.606
0.575
Appendix
Table 8: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
a1e-3w30b1
a1e-2w20b0.3rms
task1
0.632
0.576
task2
0.648
0.659
task3
0.633
0.655
task4
0.610
0.667
task5
0.595
0.625
task6
0.579
0.606
Appendix
Table 9: The accuracies of different combinations of hyperparameters for ARROW on task stream and corresponding AAT.
The loss of plasticity - the ability of a network to learn new information after having already learned older information - is a fundamental challenge in creating artificial neural networks capable of continual learning. Although this phenomenon has been known for decades, it has mostly been studied in older, relatively small architectures and rarely in natural-language domains. To determine whether loss of plasticity remains a problem in the modern transformer-based LLM paradigm, we study plasticity loss in GPT-style Transformer models trained on a multilingual continual learning problem. Consistent with prior work, we find evidence of plasticity loss across models ranging from 5M to 314M non-embedding parameters, as measured by deterioration on a held-out Vietnamese probing task. We further find that the onset of plasticity loss follows a predictable scaling law, growing sublinearly with model size. These results suggest that larger models may delay the measurable effects of plasticity loss, but that increasing parameter count alone is likely to be insufficient to completely prevent it. We also find evidence of plasticity loss under stationary multilingual training, challenging the view that the phenomenon is exclusive to continual learning with abrupt task changes. Overall, our results suggest that even large Transformer language models trained on natural-language will eventually lose the ability to efficiently adapt to new data after sufficiently long training, in both continual and stationary settings.
J. Fernando Hernandez-Garcia, Tomás Figliolia, Beren Millidge
Transfer learning with models pretrained on ImageNet has become a standard practice in computer vision. Transfer learning refers to fine-tuning pretrained weights of a neural network on a downstream task, typically unrelated to ImageNet. However, pretrained weights can become saturated and may yield insignificant gradients, failing to adapt to the downstream task. This hinders the ability of the model to train effectively, and is commonly referred to as loss of neural plasticity. Loss of plasticity may prevent the model from fully adapting to the target domain, especially when the downstream dataset is atypical in nature. While this issue has been widely explored in continual learning, it remains relatively understudied in the context of transfer learning. In this work, we propose the use of a targeted weight re-initialization strategy to restore neural plasticity prior to fine-tuning. Our experiments show that both convolutional neural networks (CNNs) and vision transformers (ViTs) benefit from this approach, yielding higher test accuracy with faster convergence on several image classification benchmarks. Our method introduces negligible computational overhead and is compatible with common transfer learning pipelines.
Xander Coetzer, Arné Schreuder, Anna Sergeevna Bosman
Department of Computer Science, University of Pretoria, Pretoria, South Africa
Continual training of deep neural networks under non-stationarity often leads to a progressive loss of plasticity, eventually limiting further learning. We relate plasticity to the empirical Neural Tangent Kernel, and identify dynamical isometry (the condition that layer-wise Jacobian singular values remain close to one) as a key mechanism for preserving plasticity in continual learning. We revisit a class of networks that are almost-everywhere isometric while remaining universal Lipschitz function approximators, demonstrating that near-dynamical isometry is compatible with expressive nonlinear representations. For general architectures, we propose an efficient isometry-promoting regularization scheme and identify a novel mechanism by which it can reactivate dormant ReLU units. Building on this, we introduce AdamO, an Adam-style adaptive optimizer that decouples isometry regularization from gradient updates, analogous to AdamW. We further reinterpret prior plasticity-preserving approaches through the lens of dynamical isometry, showing that they target only a partial measure of isometry. Across supervised and reinforcement-learning continual-learning benchmarks designed to induce plasticity loss, our methods consistently match or outperform existing approaches.