Paper: Towards Perpetually Trainable Neural Networks (ICLR-2024)

Reviewer comment (Novelty):
Originality aspects: While the paper researches extensively into reasons of loss and plasticity, the remedies are based on well-known methodologies (L2 regularisation, layer normalisation) and therefore the message is mostly limited to identifying the remedies as opposed to novel methodological contributions towards solving this problem (see Q7).

Author response:
**Concern 1.** While the paper researches extensively into reasons of loss and plasticity, the remedies are based on well-known methodologies (L2 regularisation, layer normalisation) and therefore the message is mostly limited to identifying the remedies as opposed to novel methodological contributions towards solving this problem (see Q7)

**AR:** We agree that the components of the most effective strategy we identify are already widely known. However, we disagree that this reduces the significance of the contribution: first, as noted in other responses, we emphasize that the primary contribution of this work is the observation that plasticity loss can be decomposed into largely-independent mechanisms, and that mitigation strategies on these mechanisms can be combined. The observation that mechanisms of plasticity loss can be studied in isolation and then combined together significantly reduces the combinatorial complexity of the search for even better plasticity-preserving training algorithms. Second, we note that this unoriginal combination outperforms a number of more “novel” methods which we evaluated and which have been previously published. Showing that a simple existing method can be adapted to outperform newer more complex approaches is in our opinion an important regularizer in the complexity of machine learning algorithms; indeed, we tried a variety of more "novel" approaches (e.g. the normalization strategy in Appendix D, Figure 21) which did not significantly outperform the simple combination of LN + L2 and omitted these from the main paper to simplify the exposition of the work. Finally, the framework we use to select these interventions (the mechanisms presented in Figure 1) can be leveraged to find even more effective solutions: for example, while we use L2 regularization due to its simplicity, it is possible that a more careful variant of weight normalization or regularization towards the initial parameter values may interfere less with convergence speed in single tasks.

Reviewer comment (Experiments):
In the introduction, the paper discusses at length the problem of nonstationarity and continual learning at the same time. Although these problems are connected, it is important to distinguish between these two. With the problem of nonstationarity alone, one can see the catastrophic forgetting problem entangled with the plasticity loss. For example: “Our empirical evaluations aim to split the philosophical difference between these perspectives. In the following sections, we will evaluate the ability of a learning algorithm to maintain plasticity in a continual supervised learning task, where we fix some image classification dataset (in our case, CIFAR10) and at fixed intervals apply a random transformation to the inputs and labels.” The difference can be illustrated as follows: imagine that a model is trained on 180 degrees rotated MNIST symbols, but later needs to work on non-rotated ones, continuing from the rotated model checkpoint. If the performance drops due to the problems in Figure 1, one might accept catastrophic forgetting but not the loss of performance. To address this, two actions are suggested: (1) disentangling the distribution nonstationarity problem in the introduction (2) calculating the baseline of last task and all tasks performance: in Figure 6, how would the proposed changes affect the performance if the whole WiLDS dataset were learned in a stationary distribution fashion? How would the changes affect the performance on just the last task?

Author response:
**Q1.** Distinguish nonstationarity vs continual learning. With the problem of nonstationarity alone, one can see the catastrophic forgetting problem entangled with the plasticity loss. 2 suggestions: (1) disentangling distribution nonstationarity problem in the introduction (2) calculating the baseline of the last task and all tasks performance in Figure 6, include baseline of training on whole distribution + baseline of training random init each time dataset reshuffled.

**AR:** We use the terms “nonstationary” and “continual” learning interchangeably in this work, and the reviewer raises a valid concern that continual learning problems often evaluate the learner on not just forward transfer, but also on backward transfer to previously seen data (or in its absence, catastrophic forgetting). We agree that this is confusing, and will update the manuscript to ensure that when we evaluate methods on forward-looking performance, we use the term “nonstationary learning”. We have also implemented a baseline for the WiLDS benchmark which resets the network parameters each time the dataset changes, which is included in our revisions (Appendix D, Figure 18).

Reviewer comment (Presentation):
Figure 2, top left: should the axes have values? Does it also make sense to report fraction, not the absolute number?

Author response:
**Q3.** Figure 2, top left: should the axes have values? Does it also make sense to report fraction, not the absolute number?
**AR:** Having reviewed the figure, the only missing value is the x-axis in the middle subplot, which corresponds to training steps. While switching between fraction and absolute number of dead units does not change the overall trend, we agree that since the total number of units isn’t obvious from the figure, it is clearer to set the y-axis ticks to be a fraction. We will make this change in our future revisions.

Reviewer comment (Experiments):
“Mild L2 regularization (we use a value of 1e-5) is sufficient to resolve this issue in all three architectures.” Why was this choice of hyperparameter made?

Author response:
**Q4.**  “Mild L2 regularization (we use a value of 1e-5) is sufficient to resolve this issue in all three architectures.” Why this choice of the hyperparameter?
**AR:**  The value of 1e-5 was chosen because it was the first value we tried and seemed to work reasonably well. In response to this and other reviews, we ran a more detailed sweep over L2 penalties for the MLP and found that in fact a slightly smaller value of 1e-6 results in higher per-task performance while still avoiding plasticity loss in later tasks, this can be seen in Figure 16 in Appendix D.

Reviewer comment (Novelty):
The statement that “while saturated units have received much attention as a factor in plasticity loss, the accumulation of effectively-linear units has not previously been studied in the context of continual learning” may not fully address such accumulation in continual learning. As mentioned in Q1, the task involves learning under nonstationarity but does not tackle whether and how previously accumulated knowledge should be used or discarded. Could the added aspects of continual learning be contrasted with previous works, such as Montufar (2013), in this analysis?

Author response:
**Q6.** “While saturated units have received much attention as a factor in plasticity loss, the accumulation of effectively-linear units has not previously been studied in the context of continual learning.” Issue with saying “continual” rather than “nonstationary” learning. Would it be possible to contrast the added aspects of continual learning to previous works, such as  uMontufar (2013), in this analysis?
**AR:** Montufar (2013) describes the piecewise linear structure of a neural network assuming ReLU activations, and uses this structure to understand the expressivity of the model. In particular, the focus is on counting the number of linear regions, and providing a specific construction for deep architectures that maximizes the potential number of regions that you get. In this work, **we are interested in learnability rather than expressivity**. We will update the paper to clarify this. In our revisions to Appendix D, we provide a simple demonstration that a completely linearized network can encounter training difficulties (Figures 19 & 20). However, in addition to this example, it is possible to reason more generally about the effect of linearized units on training dynamics with a simple thought experiment, which we include in a child comment.

Reviewer comment (Experiments):
Do different types of layer normalisation (e.g., batch normalisation, layer normalisation) provide different impacts? Did the authors compare these different remedies?

Author response:
**Q7:** Do different types of layer normalisation (batch normalisation, layer normalisation provide different impact?) Did the authors compare between those different remedies?

**AR:** We provide one data point on the interaction between layer normalization and batch normalization in Figure 4 of the main paper. However, we also conducted preliminary experiments contrasting the effects of different normalization choices by isolating the effects of mean subtraction and division by standard deviation for both layer and batch normalization, and also studying the interaction between performing different normalization operations on different axes. We did not originally include these in the paper due to space constraints, but we have added a section discussing them and visualizing these results in Figure 21 of Appendix D.

Reviewer comment (Novelty):
Section 2.2 (Loss of Plasticity) needs to be much more specific about the contributions in the prior work, and to more clearly delineate how the contributions of this work differs. In its current form, this section only summarizes how prior work characterize the concept of plasticity in a very abstract manner, without mentioning any specific differences between this work and prior work. Many of the ideas presented in this paper are very similar to ones already existing in prior work (especially Lyle 2021, Kumar 2023a, Abbas 2023), and it would be helpful for understanding the novel contributions of this paper beyond prior work if it were described in more specific terms.

Author response:
**Concern 2:** Section 2.2 needs to be much more specific about the contributions in the prior work, and to more clearly delineate how the contributions of this work differs. 
**AR:** We thank the reviewer for highlighting this pitfall, and will incorporate into our revisions. As an overview, while our work shares many similarities with that of Lyle et al. (we consider overlapping mitigation strategies and use similar benchmarks), we believe that this paper makes a significant contribution to the literature beyond that provided by prior work. In particular: 

a) We show that **mitigation strategies to independent mechanisms of plasticity loss can be combined** to significantly outperform single interventions. Prior work did not consider that interventions may operate on independent mechanisms and benefit from combination. This insight not only explains the falsification results of Lyle et al., who show that single mechanisms fail to explain plasticity loss in all settings, but also provides a more efficient framework for searching for new plasticity preserving optimization methods. The approach we follow, of isolating individual mechanisms, identifying effective interventions on these mechanisms, and then combining the best of each class, significantly reduces the combinatorial complexity of searching over training improvements.

b) Our combined strategies exhibit significant performance improvements over the single interventions studied by Lyle et al. and also shows significant robustness to particularly difficult sequential supervised learning tasks in which which individual mitigation strategies exhibit declining performance.

c) We provide a finer-grained view into the mechanisms of plasticity loss than that provided by previous works, identifying the importance of controlling for preactivation shift and unbounded norm growth. The benefits of layer normalization have been validated in many prior works, but these works **did not provide an explanation of its benefit**. This mechanistic view also sheds light onto the falsification experiments of Lyle et al., giving intuition behind the inefficacy of weight norm in networks with saturated units.

d) we evaluate a much **more comprehensive spectrum of mitigation strategies** on a significantly broader set of tasks. We perform a thorough evaluation of varying normalization layer combinations and configurations, consider recently proposed interventions such as ReDO (Sokar et al., 2023) and Regenerative Regularization (Kumar et al., 2023) and further evaluate a range of strategies to control or mitigate pathological loss landscape curvature.

e) We show that in fact the “loss landscape pathologies” identified by Lyle et al. (2023) were largely attributable to the norm of the network outputs (see point 3 below) rather than more subtle properties of the loss landscape. Further, **we show that curvature-aware optimization is not sufficient to mitigate plasticity loss, contradicting the implications of Lyle et al.**

Reviewer comment (Experiments):
In Section 3.2, the paper claims that another reason for the loss of plasticity is the linearization of units. However, there are no experiments in this section. Therefore, there is no evidence that this phenomenon happens in practice. Furthermore, the paper itself noted "recovery from this state is possible", so even if it does happen in practice, the network could potentially quickly recover from this state, making it even more important to empirically understand to what extent this "plight" is a problem in practice.

Author response:
**Concern 3** (Sec 3.2): No experiments on linearization of units. Furthermore…  the network could potentially quickly recover from this state

**AR:** Experiments studying the linearization of units are included in the right hand side of Figure 4, which illustrates how networks trained with non-saturating activations accumulate zombie units in tandem with reduced performance on future tasks, and this accumulation is more pronounced in networks that lose plasticity. While the network could in theory recover from unit linearization, we see in many cases that it does not. Particularly striking is the case of residual networks without layer normalization, where we see that nearly all units are computing linear functions, and for which training curves often never pick up. This can be observed in Figures 19 and 20 in Appendix D. We also note that Appendix C already studies the relationship between plasticity and the variance of the slope of nonlinearities on the preactivation distribution, a quantity which is closely connected to the number of linearized units. We think that further investigation into the effect of these linearized units on optimization dynamics is an exciting direction for future research, potentially drawing on the analysis of Poole et al. and Martens et al. to characterize the feature geometry and gradient dynamics in this network state, and hope that our findings serve as inspiration for further analysis.

Reviewer comment (Evaluation):
In Section 3.4, the paper hypothesizes that another reason for the loss of plasticity is that "regression to targets which have a large mean relative to their variance is innately difficult for neural networks". However, the experiments do not necessarily support this hypothesis. Instead, the experiments show that using distributional losses seems to mitigate loss of plasticity, which is a finding that has already been presented in prior work. While the paper's hypothesis could potentially explain why distributional losses work well, there are a number of other potential explanations, and the paper provides no additional evidence for why the proposed hypothesis is the correct one.

Author response:
**Concern 4 (Sec 3.4):** "regression to targets which have a large mean relative to their variance is innately difficult for neural networks" hypothesis not supported by experiments. Show that using distributional losses seems to mitigate loss of plasticity, already known. No justification for why this is the right explanation, and not just that distributional losses are better.

**AR:** The sweep used in Section 3.4 over the two values of gamma provides evidence supporting our claim that it is the higher target mean, rather than some generic benefit of distributional losses, which drives performance gains. Note in particular that the distributional losses with gamma=0 do *not* improve plasticity over regression losses, suggesting that improvements are not coming from uniform superiority of distributional losses. To illustrate the importance of the target offset value at a finer-grained level of detail, we provide dose-response curves for a MLP trained to regress on targets with varying means and then fine-tuned on a new task in Appendix D, Figure 12. In particular, we see a monotone relationship between the magnitude of the pretraining regression problem offset and the error attained on later tasks.

Reviewer comment (Experiments):
The paper is not specific enough about the exact environments/datasets being evaluated. In particular, the paper says it "trained on the arcade learning environments", as well as "a dataset from the WiLDS benchmark". However, the arcade learning environment includes over 50 games, and the WiLDS benchmark includes 10 different datasets. For the RL experiment, it is unclear whether the experiment is only trained on one environment (if so, which one?), or if the results are averaged over training runs on many different environments (if so, which ones?). Furthermore, the paper should be more specific about which particular dataset it used from the WiLDS benchmark for its experiment.

Author response:
**Concern 5 (Sec. 5)**: For RL: what exact games are being evaluated? Which dataset from WiLDS benchmark? “Assuming the paper only evaluated on 1 RL environment, I believe this, in addition to only 1 distribution shift dataset, is not extensive enough to validate the efficacy of the proposed approach on realistic learning problems”

**AR:** While the arcade learning environment is only a single benchmark, it contains a rich diversity of games which require complementary algorithmic strengths in order to master; we emphasize that **we evaluate on all 57 games**, and so while we agree that more evaluations always provide more information, it was not clear to us what additional information we would expect to gain from an additional environment. If the reviewer has particular domains outside of Atari that they would like to see evaluated, we would be happy to discuss further. We use the iwildcam dataset from the WiLDS benchmark, as was described in Appendix B.

Reviewer comment (Experiments):
Assuming the paper only evaluated on one RL environment, this, in addition to only one distribution shift dataset, is not extensive enough to validate the efficacy of the proposed approach on realistic learning problems.

Author response:
**Concern 5 (Sec. 5)**: For RL: what exact games are being evaluated? Which dataset from WiLDS benchmark? “Assuming the paper only evaluated on 1 RL environment, I believe this, in addition to only 1 distribution shift dataset, is not extensive enough to validate the efficacy of the proposed approach on realistic learning problems”

**AR:** While the arcade learning environment is only a single benchmark, it contains a rich diversity of games which require complementary algorithmic strengths in order to master; we emphasize that **we evaluate on all 57 games**, and so while we agree that more evaluations always provide more information, it was not clear to us what additional information we would expect to gain from an additional environment. If the reviewer has particular domains outside of Atari that they would like to see evaluated, we would be happy to discuss further. We use the iwildcam dataset from the WiLDS benchmark, as was described in Appendix B.

Reviewer comment (Experiments):
The proposed approach is using layer norm with L2 regularization. It is unclear why the RL experiment only compares C51 with layer norm and C51 without layer norm. What about L2 regularization? Even if the L2 regularization doesn't help, or even hurts performance, it would be informative to the reader.

Author response:
**Concern 6:** Why no L2 regularization in RL experiments?

**AR:**  Our evaluations on the image classification setting revealed that growth of the parameter norm only begins to interfere with learning when it increases by several orders of magnitude (e.g. from O(10^3) to O(10^7+). While the parameter norm of DQN-style agents on the Atari domain is known to grow, it does not increase by five orders of magnitude (c.f. Appendix C for an illustration of parameter growth in the game Seaquest). L2 regularization is known to slow down training in RL and make it difficult for agents to take off, so we did not anticipate seeing any benefits. For completeness, we provide preliminary results for the C51 agent with L2 regularization on a subset of environments from the Atari benchmark in our updates to the paper (we alphabetized the environments and selected every 5th game), which confirm that L2 regularization does not improve performance in the C51 agent.

Reviewer comment (Theory):
The use of classification loss in a regression problem increases the minimum loss that the network can achieve because there will be a residual loss for each data point. Even though plasticity might be maintained with a classification loss, it comes at the expense of the best possible performance.

Author response:
**AR:** We are not entirely sure what residual the reviewer is referring to: the lower bound on the cross-entropy loss of the target entropy, or approximation error induced by the binning representation. In both cases, we do not believe that this presents a barrier to the network’s ability to match the desired scalar targets. In the case of the lower bound on the cross entropy, we note that this does not influence optimization dynamics. In the case of expressivity of the binning trick, the distributional output representation does not reduce the expressivity of the network, as any scalar value which a regression model can output, provided it is within the bounded range on which the support of the two-hot distribution is defined, can be represented by a corresponding two-hot distribution (see child comment for details).

Reviewer comment (Evaluation):
The paper states at many points that previously proposed methods "slow learning in the single-task setting." This suggests that the proposed method has no issues in the single-task setting. However, this is not true — the proposed method can lead to sub-optimal performance in the single-task case (top-left panel of Figure 3, Figure 8).

Author response:
**AR:**  We agree with the reviewer on the importance of distinguishing between 'not reducing performance' and 'consistently exhibiting good performance', and that L2 regularization can slow down learning in some domains and architectures. We plan to include a more detailed discussion of tradeoffs between performance on a single task and stability in later revisions. The random label memorization task presents a particularly adversarial realization of this tension, as the primary task amounts to memorizing random noise, and L2 regularization tends to have a smoothing effect and makes it more difficult to memorize noisy functions. However, we also note that ultimately this slowdown does not seem to prevent the network from eventually converging on the task -- it just increases the number of optimizer steps necessary to achieve this. To back up this claim, we have re-run a subset of the experiments originally shown in Figure 8 where we provide a slightly longer training interval between label resets, and this longer interval allows both the regularized and unregularized networks to attain >98% accuracy on the task, which can be seen in Figure 16 of Appendix D.  We have also included a sweep over L2 penalty values and activation functions for the CNN and MLP architectures. These results show that the value of 1e-5 which we defaulted to in the paper is suboptimal for certain architecture-dataset combinations, and in some architectures a different value (1e-6 in this case) can recover the stability benefits while also exhibiting minimal performance degradation relative to the network’s unregularized counterpart. 

As a result, we do think that L2 regularization at a suitable dose is a valid tool for maintaining plasticity. Our perspective can be summarized as follows: **while L2 regularization can slow down training if applied too liberally, in can be titrated to minimally interfere with performance on the current task while still providing benefits to plasticity by containing otherwise-unbounded parameter growth**. We will ensure that this perspective is effectively conveyed in our future revisions.

Reviewer comment (Evaluation):
The paper does not provide any details on the number of runs for each experiment, nor does it indicate what the shaded region in each plot represents. It is unclear whether any of the results in the paper have statistical significance and whether they will withstand the test of time.

Author response:
**AR:** Most tasks run with 3-5 seeds per hyperparameter setting; we will specify these details in our revisions. Shaded regions refer to the standard deviation over the training runs exhibiting the listed hyperparameter values.