Paper: Curvature Explains Loss of Plasticity (ICLR-2024)

Reviewer comment (Theory):
I find the current figures not sufficiently convincing. First of all, the srank is a normalized quantity, and therefore it cannot tell us the curvature of the loss. This is because the curvature is closely related to the absolute scale of the eigenvalues -- for a fixed sharpness of the local landscape, one can have an arbitrary srank and vice versa. Therefore, I am not convinced that the srank can be interpreted as "curvature"

Author response:
>Definition of curvature, scale of eigenvalues and rank as a normalized quantity.

Please see the shared reply, [Curvature Definitions], for more details

There are several definitions of curvature, but the Hessian describes the local curvature of the function. The srank of the Hessian gives the number of curvature directions and we are arguing this property of the curvature explains loss of plasticity. 

The fact that the rank is a normalized quantity (normalized meaning it ignores the scale of the singular values) means that it provides information about the relative distribution of the singular values. Our findings suggest that this normalized quantity is useful for explaining plasticity.

Reviewer comment (Evaluation):
The figures are not clearly convincing. In particular, what the authors are trying to argue is that the batch error increase is correlated with srank decrease -- but there is no single figure that precisely shows this. The authors need to include a figure whose x-axis is the srank and the y-axis is the batch error, and perform a rigorous correlation analysis to convince me.

Author response:
>include a figure whose x-axis is the srank and the y axis is the batch error

Please see shared reply [Correlation Plots] and updated manuscript (Appendix A.1, D.3 and D.4) with correlation plots.

Reviewer comment (Evaluation):
I think it would be valuable to discuss the recent results in https://arxiv.org/abs/2309.16932 (mainly section 4.4). For example, the authors say that "This suggests that loss of plasticity only occurs with the non-linear activations," which does not seem true in light of the result in the related work. Also, throughout, the authors suggest that L2 regularization can possibly mitigate the loss of plasticity, whereas the results in the above work show that L2 reg. exacerbates it. The authors should explain the contradictions and why they occur.

Author response:
>Discussing “Symmetry Leads to Structured Constraint of Learning”

First, the paper that you link to was uploaded to arxiv after the ICLR submission deadline. We do appreciate the reference, as the paper looks quite interesting and relevant. Discussion of the linked paper, however, should not be relevant to the reviewing process as per the reviewer FAQ. That being said, we are not sure what conclusions to draw regarding linear networks in that paper. Furthermore, L2 regularization has been shown to mitigate plasticity loss in several other works, which we reference in Section 2.1, but it requires careful hyperparameter tuning.

Reviewer comment (Theory):
Why is the stable rank based on the singular values but not the eigenvalues? The singular values are only positive, but, in my opinion, the negative eigenvalues in the Hessian could play an important role here.

Author response:
>Why is the stable rank based on the singular values but not the eigenvalues?

The stable rank is always defined with respect to the singular values because the magnitude of the singular values is the magnitude of curvature in the corresponding. The stable rank measures the number of singular values that account for a large fraction of the sum of singular values.

Reviewer comment (Experiments):
The paper doesn't evaluate across different architectures, and only considers extremely small MLPs. Given that prior work has noted that an intervention's efficacy often heavily depends on the architecture to which it is applied, it is important to evaluate an intervention on a range of architectures to get a full view of its robustness.

Author response:
>The paper doesn't evaluate across different architectures, and only considers extremely small MLPs. [...] The paper would benefit from evaluations in other architectures … Does it generalize to convolutional networks / ResNets / sequence models?

Please see the shared reply [Neural Network Architecture] and updated manuscript (Appendix D.2, D.4)
We speculate that a resnet would mitigate some plasticity loss because its residual connections provide a deep linear subnetwork, but we leave this for future work.

Reviewer comment (Theory):
Saying that the curvature of the neural network is reduced because of a reduction in srank might go against existing notions of curvature in the literature on deep neural network training, which often characterize the curvature by the maximal Hessian eigenvalue. In particular, a reduction in srank of the features is often caused by the largest eigenvalue growing faster than the remainder. It is possible that a similar phenomenon is happening here, wherein the largest eigenvalue of the Hessian is growing faster than the other eigenvalues. A visualization of the eigenvalue distribution at the start vs at the end of training would be useful. I would be inclined to call a reduction in the srank an increase in the 'degeneracy' of the Hessian, rather than a reduction in curvature.

Author response:
>reduction in srank might go against existing notions of curvature in the literature on deep neural network training, which often characterize the curvature by the maximal Hessian eigenvalue.

See shared reply, under [Curvature Definitions]

Reviewer comment (Experiments):
The rank of the Hessian with respect to the last two layers of the network is likely to be highly correlated with the rank of the penultimate feature layer. Figures 1 and 2, for example, show that the rankings given by the srank of the Hessian and the srank of the features largely agree with each other. Similar issues in the correlation between the feature srank and the error over time also emerge in the Hessian srank, for example the leaky-ReLU network in Figure 2, middle column. The evaluation of the srank of the features in Figure 1 appears to have either been performed with a very narrow network or a very small batch size, making it difficult to evaluate the significance of these measurements. I think the paper would benefit from a closer examination of the difference between the Hessian rank and the feature rank in these networks.

Author response:
>I think the paper would benefit from a closer examination of the difference between the Hessian rank and the feature rank in these networks.

The Hessian is a second-order property of the objective function, describing the local curvature and its rank counts the number of directions of curvature. In contrast, the feature representation is a zeroth-order property of the neural network. For example, if the labels are shuffled between tasks then the feature rank at the end of one task and at the beginning of the subsequent task would be the same. Thus, the feature rank cannot distinguish between some changes in the optimization landscape. There is no a-priori reason to believe that the feature rank is related to curvature or the Hessian rank, and indeed previous papers have also found that feature rank is not a consistent explanation for plasticity loss [2].

Reviewer comment (Evaluation):
One of the biggest issues that emerges from the paper is the observation that the Hessian srank often exhibits similar inconsistencies to other quantities, like the srank of the features, which were used as justification to dismiss these quantities as good explanations of plasticity. For example, weight decay and regenerative regularization both exhibit similar Hessian srank, but weight decay exhibits greater plasticity loss in Figure 4.

Author response:
>For example, weight decay and regenerative regularization both exhibit similar Hessian srank, but weight decay exhibits greater plasticity loss in Figure 4.

Although the Hessian srank decreases in the regularized experiment, the rank is still substantially larger than the unregularized setting, even compared to the linear baseline. We do not claim that weight decay exhibits plasticity loss in Figure 4, as the confidence interval is large we conclude that weight decay is unstable. See the new correlation plots for further evidence.

Reviewer comment (Theory):
The choice to study the rank of the Hessian, as opposed to the rank of the empirical NTK or the features, is not given sufficient motivation. Each of these quantities reveals something meaningful about the loss landscape, and it is not immediately obvious that the Hessian should be preferred over the other two.

Author response:
>The choice to study the rank of the Hessian, as opposed to the rank of the empirical NTK or the features, is not given sufficient motivation. 

The motivation is twofold. First, previously proposed explanations of plasticity have focused on zeroth (feature rank) and first (gradient/update norm) order properties. Second, there is a growing literature finding that curvature (which the Hessian describes locally) strongly influences learning dynamics. It is thus a natural hypothesis that the Hessian may help explain loss of plasticity.

Reviewer comment (Theory):
While the proposed regularizer is reasonable, it feels very much orthogonal to the curvature of the network. The causal connection between preserving the distribution of the parameters and maintaining the rank of the hessian remains vague, and the benefits of the regularizer on the rank appear incidental.

Author response:
>While the proposed regularizer is reasonable, it feels very much orthogonal to the curvature of the network.

Please see shared reply, [Connecting Curvature and Regularization]

Reviewer comment (Experiments):
Because of the limited set of evaluations, it is difficult to say whether the proposed regularizer will be useful in other contexts, where it may be necessary for the distribution of weights to shift in ways that increase the Wasserstein distance from initialization, or where the regularizer may slow down learning too much to be practical. It would be useful to evaluate on other datasets and to also show the learning curves; for example, the worse online learning results for the wasserstein regularization coupled with the better end-of-training results suggest that not only does the loss start out much higher, but that it stays higher than the other methods for a long time -- in other words, it is resulting in much slower learning.

Author response:
> for example, the worse online learning results for the wasserstein regularization coupled with the better end-of-training results suggest that not only does the loss start out much higher, but that it stays higher than the other methods for a long time -- in other words, it is resulting in much slower learning.

We found that our proposed Wasserstein regularizer does not result in slower learning than the regen regularizer (See Appendix D.5 for inter-task learning curves). In general, the Wasserstein regularizer is smaller than the L2 and regen regularizer. Thus, there is a larger set of parameters within a set distance with respect to the Wasserstein distance.

Reviewer comment (Evaluation):
I am skeptical of the results which seem to indicate that a linear network is able to attain near-0 error on MNIST. I have trained linear classifiers on this dataset in the past and don't recall seeing anything exceed 80-90% accuracy. I am also extremely surprised at the results which seem to indicate that a 3-layer, < 128-width **linear** MLP is able to attain zero training error on label-shuffled CIFAR-10. For example, even the larger networks studied by Zhang et al. in the classic paper "Understanding deep learning requires rethinking generalization" required dozens to hundreds of epochs to converge on random CIFAR-10 labels. I have a difficult time believing that a task sufficiently difficult for an ImageNet-calibre architecture to require Could the authors confirm that they are indeed able to obtain zero error on image classification datasets with a linear model?

Author response:
>a linear network is able to attain near-0 error on MNIST(?)

Our experiments use a subset of the MNIST dataset (similar to [1,2]), and a linear network can achieve a low error on this subset. The non-linear networks achieve a lower error than the linear network on the first task (see Appendix D.5), but due to loss of plasticity this error increases to be much higher than the error of the linear network. We emphasize that it is striking that loss of plasticity still occurs on problems for which linear features can be predictive.

Reviewer comment (Evaluation):
How strong are the correlations between feature, gradient, and Hessian srank, and can the authors provide evidence that the Hessian srank is providing a meaningful notion of curvature or trainability not captured by the other two quantities?

Author response:
>correlations between feature, gradient, and Hessian srank, and can the authors provide evidence that the Hessian srank is providing a meaningful notion of curvature or trainability not captured by the other two quantities?

Please see shared reply and updated manuscript (Appendix A.1, D.3 and D.4). These results validate the fact that the number of curvature directions is important in explaining plasticity.

Reviewer comment (Experiments):
The paper would benefit from evaluations in other architectures, and on tasks that cannot be solved by a linear model, to validate the utility of the proposed regularizer. In particular, how does the regularization method perform on tasks where linear networks fail? Does it generalize to convolutional networks / ResNets / sequence models?

Author response:
>The paper doesn't evaluate across different architectures, and only considers extremely small MLPs. [...] The paper would benefit from evaluations in other architectures … Does it generalize to convolutional networks / ResNets / sequence models?

Please see the shared reply [Neural Network Architecture] and updated manuscript (Appendix D.2, D.4)
We speculate that a resnet would mitigate some plasticity loss because its residual connections provide a deep linear subnetwork, but we leave this for future work.

Reviewer comment (Experiments):
The authors' decision to measure the partial blockwise of the last two layers raises questions. It's unclear why early layers are seemingly disregarded. Don't they play a significant role in the loss of plasticity?

Author response:
>It's unclear why early layers are seemingly disregarded.

The goal of this paper is to compute, exactly, the Hessian rank since the main claim of our paper is that the curvature, as measured by the Hessian rank, explains loss of plasticity. This computation has a O(d^3) cost where d is the number of parameters. It turns out that we do not need the exact hessian, because layer-wise approximations provide a good approximation of the Hessian for stationary tasks [1].

W1 -> R2: Thank you for your reply. You express some concern whether the last two layers can accurately capture curvature for deeper networks. We have added results using a 13-layer neural network in Appendix F. Our results show that, with this 13 layer neural network, the sRank of the Hessian with respect to the last two layers exhibits the same trend with plasticity loss, thereby validating its use as an approximation. For further discussion on why the last 2 layers are effective, please see [Previous Layers].

You also raise a very interesting question: 1) can the Hessian sRank help identify and localize loss of plasticity to particular layers or regions of the network?

[Previous Layers] We first note that the previous layers still contribute to the input for the last 2 layers, which can influence the Hessian srank. This is why, for example, the Hessian sRank decreases even when resetting the last 2 layers (see Figure 20). Although the Hessian is taken with respect to the last two resetted layers, the outputs of the previous layers are used as input for the last 2 layers and drive the reduction in the srank.

This provides further evidence that the Hessian-based metrics we analyze provide a good surrogate for the curvature of the whole neural network. It is not clear, however, how this can be used to localize and identify particular weights with low plasticity. While your suggestion is very interesting as a potential pruning strategy, we leave this for future work.

Reviewer comment (Experiments):
For a comprehensive evaluation, it would be valuable if the authors could provide a layer-wise or partial block-wise Hessian rank for the experiments. Especially, it would be important to check it out for the CIFAR-10 experiments, since it utilizes a deeper network than Mnist experiments.

Author response:
> provide a layer-wise or partial block-wise Hessian rank

While we have argued that the layer-wise Hessian is a good approximation for stationary tasks, it is not a good approximation for non-stationary tasks like in our continual learning experiments.
This is because the layer-wise Hessian does not change when the target function changes if we are using a piecewise-linear activation function like ReLU.
Taking the last 2 layers is the simplest way to avoid this issue with the layerwise approximation, and our results demonstrate that it is sufficient to explain plasticity loss.

Reviewer comment (Experiments):
The choice of the Wasserstein distance measure as a regularizer adds complexity, especially when it comes to tuning for smaller sample sizes ($d$). Why not employ the simpler L2 regularization from the initial parameter? Does the norm of the neural network should be large enough to make accurate predictions for a large number of tasks?

Author response:
>The choice of the Wasserstein distance measure as a regularizer adds complexity, especially when it comes to tuning for smaller sample sizes  Why not employ the simpler L2 regularization from the initial parameter? 

The Wasserstein distance is not any more complicated than the L2 or regen (which is L2 regularization from the initial parameter). Like both L2 and regen, it uses a single hyperparameter and does not depend on the sample size. In Appendix C.5, we show that the Wasserstein regularizer is actually more robust to the hyperparameter than both L2 and regen.

Reviewer comment (Experiments):
The omission of experiments using permuted CIFAR-10 seems to be an oversight, given its relevance in evaluating plasticity.

Author response:
>The omission of experiments using permuted CIFAR-10 seems to be an oversight

Thank you for bringing this up as it was not made explicit in the manuscript.
Observation permuted CIFAR 10 is not a common benchmark for studying plasticity, because the CIFAR experiment uses a convolutional neural network. Observation permutations would destroy spatial regularity, and the translation invariance of convolutions would be an improper inductive bias.

Reviewer comment (Experiments):
A more rigorous comparison with related works is needed, specifically a comparison with widely used techniques like dropout and data augmentation.

Author response:
>Comparison with dropout, CReLU, layer norm, last layer reset

Thank you for your suggestions, please see the shared reply [Additional Experiments] and updated manuscript (Appendix D.2, D.3 and D.4)

Reviewer comment (Experiments):
A more rigorous comparison with related works is needed, specifically a comparison with respect to plasticity literature including CReLU activation [1], Layer Normalization [2], Last layer reset [3], or even combinations of these approaches [4].

Author response:
>Comparison with dropout, CReLU, layer norm, last layer reset

Thank you for your suggestions, please see the shared reply [Additional Experiments] and updated manuscript (Appendix D.2, D.3 and D.4)

Reviewer comment (Experiments):
[Significance of results] The presented arguments that existing explanations for the loss of plasticity are inconsistent seems to be insufficient. Although the proposed explanation related to the loss of curvature is intriguing, the authors should provide additional experiments (diverse model architectures and benchmarks) to argue that the proposed one is a consistent explanation for the plasticity loss. Moreover, it would be great if the authors could measure the consistency quantitatively.

Author response:
> the authors should provide additional experiments … Moreover, it would be great if the authors could measure the consistency quantitatively.

Please see the shared reply [Additional Experiments] and updated manuscript (Appendix A.1, D.3 and D.4).

Reviewer comment (Evaluation):
[Evaluation metric in continual learning setup] The paper seems to overlook an analysis on forgetting which is one of the most important criteria in the continual learning setup. It is recommended that the authors expand their discussions to encompass the concept of forgetting, in addition to delving deeper into plasticity.

Author response:
>Forgetting

Our work focuses on the problem of plasticity, and is explicitly isolated from the problem of forgetting. Our experiments, for example, measure error with respect to the task being trained and not on any previous tasks. This is by design; we want to cultivate a better understanding of continual learning by first explaining the phenomenon of plasticity.

Reviewer comment (Evaluation):
As mentioned in the weakness section, the significance of inconsistency in existing explanations is unclear. In Section 3 results, the rank of representation (top right of Figure 1) decreased, and the batch error increased when using tanh. As tanh is experiencing small decrease on the rank of representation, the batch error also showed relatively small increase. This seems to be a consistent explanation for the loss of plasticity. More explanations on this will strengthen the authors’ argument.

Author response:
> As tanh is experiencing small decrease on the rank of representation, the batch error also showed relatively small increase.

The feature rank is an inconsistent explanation because, for tanh, the increase in the batch error occurs and continues when the feature rank is constant. We have added correlation plots to show this more clearly, please see the shared reply [Correlation Plots] and the updated manuscript (Appendix A.1)

Reviewer comment (Evaluation):
Why is Wasserstein regularization better than regenerative regularization? As mentioned in Section 5, the order statistics is defined as the difference between two regularizations. It would be beneficial for the readers if the authors can provide explanations regarding the substantial difference in the results presented in Figure 5.

Author response:
>why is Wasserstein regularization better than regenerative regularization?

Please see shared reply, [Connecting Curvature and Regularization]