Guiding Neural Collapse: Optimising Towards the Nearest Simplex Equiangular Tight Frame

Summary

The paper uses Riemannian optimization to guide the final layer weights (the linear classifier) toward the nearest simplex ETF orientation. In particular, consider the two common approaches of training a deep classifier network: 1. The standard training strategy where the final layer weights are updated by backpropagation. 2. The final layer weights are fixed as a simplex ETF (which has been well-studied in previous works). The proposed approach leverages the duality between penultimate layer features and the final layer weights (to form a simplex ETF orientation) and gradually guides the latter to an optimal simplex ETF per training step.

Strengths

1. The proposed approach frames the gradual transition of weights to a simplex ETF as a Riemannian optimization problem, which can be differentiated. Thus, allowing for an end-end training pipeline. The combination of these techniques is novel to the neural collapse setting. 2. The experimental results are presented for the simple UFMs as well as practical networks and datasets to showcase the convergence benefits.

Weaknesses

The authors do not provide numerical data for the extra memory and step-time that is required by the extra deep declarative layer. A brief discussion is presented in Section 5 but I believe further details would strengthen the paper. For instance: - By what percentage does the step time and memory increase when adding this layer? - When should one avoid the backward pass through this layer and consider only the forward pass? - What is the dependence on the memory and step time growth with the feature dimension and the number of classes? Maybe a simple UFM-based analysis should suffice. see more questions below.

Questions

1. How effective is the proposed approach in settings with imbalanced classes [1] ? More generally, for settings where the simplex ETF might not be an ideal configuration (for instance: graph neural networks, see [2] ). A brief discussion on these topics can further strengthen the paper. 2. Instead of running the optimization layer to select the final layer weights at every step, what if we do it after every $k$ step? Can we potentially reduce the majority of the computational overheads while improving the convergence? 3. What is the convergence behavior when employing SGD/Adam instead of the AGD approach? nit: Where is $U_{init}$ defined? nit: line 117, below eq (5), is the formulation of $\widetilde{H}$ correct? shouldn't the denominator be $||\overline{H}||_F$ ? References [1] Fang, C., He, H., Long, Q., & Su, W. J. (2021). Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences [2] Kothapalli, Vignesh, Tom Tirer, and Joan Bruna. "A neural collapse perspective on feature evolution in graph neural networks." Advances in Neural Information Processing Systems 36 (2024).

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

N/A

Summary

This paper proposed a novel algorithm for neural network training. The algorithm is motivated by the recent discovery on the neural collapse phenomenon, which demonstrates that the last layer of neural network classifier will converge to a specific structure named simplex ETF. The authors propose to guide the network parameters to the ETF structure via explicitly penalizing on the distance to the ETF, and further address the non-uniqueness of the solution via adding a proximal term. Experimental results on various neural network architecture and real world datasets are presented, and the proposed algorithm can universally improve the training and testing accuracy over the standard training.

Strengths

The proposed algorithm is novel and well motivated, and it shows universal and significant improvement over multiple choices of network architecture and datasets. The contribution of this work is solid, it helps the community to understand the benefit of the neural collapse phenomenon, and can potentially improve the standard paradigm of neural network training.

Weaknesses

1. The presentation should be improved, see questions for detail. In general the authors should give more detailed information about how the algorithm is implemented. 2. Although the accuracy on train and test dataset exhibits significant improvement within the fixed number of training epochs, the proposed algorithm are much more complicated to compute. Therefore it makes more sense to compare the running time and computational cost with Standard and Fixed ETF. 3. Proper ablation study is missing. The authors add many additional techniques, such as exponential moving average, stratified batch sampling, deep declarative layer to improve the training. It is not clear how much the improvement indeed comes from the nearest ETF optimization.

Questions

1. How is equation 8 being optimized? Is it using Lagrangian multiplier method? A pseudo code of the proposed algorithm will be very helpful. 2. I found the Proposition 1 hard to follow. The authors should explain what is the implication of this proposition and how does it helps to the stability of training. The current statement is confusing to distinguish the main result, and notations such as $D_y$ and $\Lambda$ are not properly introduced. This proposition should be improved. 3. In Table 1 and 2, fixed ETF on CIFAR10 with VGG seems to have much worse performance than others. Do you have insights about what is going on here?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have discussed the limitation properly in the paper. A more detailed discussion with empirical results on the computational cost will be helpful.

Summary

One of the key aspects of neural collapse (NC) is that the penultimate class feature means form a simplex Equiangular Tight Frame (ETF). The main idea of this paper is to leverage this insight and improve training by further encouraging this property during training. The authors suggest doing this by solving a Riemannian optimization at a given iteration. The way it works is that the classifier weights are set to the nearest simplex ETF obtained by solving this inner Riemannian optimization problem. The classifier weights are dynamically updated during training using this Riemannian opitmization problem at each iteration (rather than trained using gradient descent) using a "deep declarative node" this allows gradients to propagate through the Riemannian optimization. They show that this approach indeed speeds up training convergence and improves training stability. Their experiments include both synthetic and real-world classification tasks and architectures. Overall the authors present a nice idea and it is a well-written paper. However, there are a few issues related to the experiments that I outline below. From my viewpoint, the value of this paper and their method (to me) is less the improved test accuracy and more the improved stability and speed of convergence. It's important to note that this speed up also comes at an additional cost (i.e. in performing the Riemannian optimization). Therefore, the improvements to stability or speed of convergences should be weighted against this caveat. I think it would help to highlight this tradeoff more upfront and make that more clear/transparent.

Strengths

This is a thoughtful and well-written paper. The authors suggest a nice idea to leverage this insight of NC in deep learning and their approach has clear benefits. It is a nice idea and very well executed. There are clear improvements to the current methods; e.g., their improvement upon [74] by solving the inner Riemannian optimization instead of requiring the model backbone to do the work of matching to a chosen fixed ETF. The theory and the idea is very compelling. The implementation is good and well explained. Beyond the theory and the novelty of the idea, the main strength of the paper is the value added wrt convergence speed in terms of the number of epochs required for the network to converge. Good work.

Weaknesses

The main points of concern for me are in regards to the experiments and how the results are reported in the paper. Table 2 looks good but is a bit misleading particularly when comparing the ranges of the test top-1 accuracy. The results are still interesting but it's not such a strong/clear winner; that is, when looking at the ranges, it's not so obvious. The authors point this out and clarify that the advantages are speed to convergence and decreased variability which I agree are definite plusses. The test top-1 accuracies reported in Table 2 aren't competitive with what can be obtained on these benchmark datasets, particularly for the Resnet models. For example, looking at 200 epochs or training, STL on ResNet50 should be able to achieve 85-90% test accuracy, even for Resnet18 the test top-1 accuracy for STL should be upwards of 75%. Similarly, for CIFAR100 on Resnet50, the test accuracies aren't competitive. It'd be interesting to see if these claims about variability still hold when giving the baselines adequate chance to be competitive. For Figure 4, also no error bars. Understanding compute restraints, it would be nice to see similar multiple seed runs for ImageNet experiments. Finally, one thing that is not reported here is an estimate of compute cost. Their method requires additional compute for each iteration. Perhaps when compared on this axis their implicit ETF and the Standard training method would be more fairly compared. The authors do mention this in the limitations section.

Questions

How do you know that the ETF that you steer towards via this Riemannian optimization process is better than the one that you would have arrived at naturally? You say "this process [provides] the algorithm with a starting point that is closer to an optimal solution rather than requiring it to learn a simplex ETF or converge towards an arbitrary one". How do you know that this is optimal? Optimal in what context? If I understand correctly, it's just the solution of the Riemannian optimization which means it forces the class means into an ETF. It's optimal wrt to the optimizaiton problem but not necessarily for the learning task? Is that correct? Do you do any, or is it possible to perform a comparison of these two resulting ETFs? How does the test accuracy of your 'encouraged' ETF compare to the one you would have obtained naturally? In Section 3.3. The Proximal Problem. I just don't see immediately why adding the proximal term guarantees uniqueness of the solution and how it stabilizes the Riemannian optimization problem. Can you add more detail or proof or reference to proof? On first reading, it was unclear to me exactly how is U_prox defined? And what is used in practice. Is it determined from the previous iteration? If I understand correctly, you tried two approaches: setting U_init = U_prox = canonical ETF. Or to set both equal to random orthogonal matrices from classical compact groups. It sounds like, in the end, you run training without the proximal matrix for one epoch. Then use the resulting U* to set U_init = U_prox = U* from that one epoch. Is that correct? How was this "warmup" approach validated? Did you experiment with various epochs? How stable were the outcomes of that analysis? You later mention (line 225) that the correct choice of these values is "crucial" so it seems important to understand. In the Section Hyperparameter Selection and Riemannian Initialization Schemes: You mention that algorithm convergence is robust to values of \delta and that the \delta reg term is a trade-off between the optimal solution's proximity to the feature means and its proximity to the given simplex ETF direction. Did you explore how and when to introduce this constraint? Or any exploration of how the solution varies with \delta? In section 3.4 General learning Setting: The role of the temperature \tau is bit unclear to me. And the reference to [67, Theorem 1] isn't very helpful. Perhaps a little more clarity as to the role \tau plays here? You state later in the Experiments section that you use \tau=5 according to Yaras et al [67]. This hyperparam choice is not very clear to me. (typo? clarification?) Proposition 1. There is notation discrepancy between what is stated in the Proposition and what is derived in the Appendix B. Namely, the Proposition is stated wrt \bar{H} but the derivation is carried out for \tilde{H}. I understand that \tilde{H} is the normalized (wrt Frobenius) matrix \bar{H} so perhaps it all works out with the normalization constant but the discrepancy there and comparing back with dimensionality of matrices in the original statement of Proposition 4.5 in Gould et al. [21] (from which this result follows) had me a bit confused. Are there error bars in Figure 2? I see them for plot (f) but not for the others? (clarification) What is depicted in Figure 2(c)? What is \mathcal{P}_{CM}? I think I somehow missed that. Are there error bars in Figure 3? Were multiple trials run for these experiments? Tables 1 and 2: The ranges for train and test top-1 accuracy values for STL on VGG seem very large. In regards to Figure 4, I'd recommend either performing more training runs for Imagenet on Resnet50. The results look very compelling but without error bars don't say much. Similarly, comparing the results in Figure 4 with those for the other real-world datasets (e.g. Cifar10, Cifar100, STL) those contained in the Appendix which do have error bars are arguably less convincing of the primarily claims of speed to convergence.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A The authors address any limitations.

Summary

This paper presents a novel approach to utilizing ETF geometry. Instead of fixing the weights or making them learnable, their approach dynamically adjusts the weights by solving a Riemannian optimization problem while allowing end-to-end training. They show that their approach outperforms both the fixed and learnable approaches in terms of convergence speed and generalization.

Strengths

Originality: The idea of dynamically adjusting weights is not new, but in the context of neural collapse (NC), it is a natural extension. Fully learnable weights do not provide the ETF structure, and fixed weights are too restrictive. The proposed approach is a good compromise between the two and combines the best of both worlds. Quality: The paper is well-written, and the proposed approach is carefully supported by theorems and experiments. Clarity: The paper is well-written and easy to follow. Significance: Their approach is general and could be applied to a range of problems. The authors applied it to synthetic UFMs and some standard image benchmarks (CIFAR-10, CIFAR-100, STL-10, ImageNet). The authors plan to release code upon acceptance.

Weaknesses

Overhead Cost: The proposed method computes the exponential moving average of the feature means, performs a Riemannian optimization, and computes the gradient of DDN. These components introduce overhead in terms of epoch time. The authors claimed in the paper that the gradient of DDN is not computed, and the Riemannian optimization overhead is negligible. This unsupported claim should be backed up by an additional experiment that reports these extra computation times. Standard Procedure: "To ensure fair method comparison," the authors include classifier weight normalization and feature normalization for the standard procedure. This is usually not the case when using CE loss (see Fig 2). The authors should justify this choice by providing the results without these normalizations for the standard procedure. Image Baselines Results are not SOTA: The reported results are not state-of-the-art. For example, ResNet-18 trained on CIFAR-10 only reaches 80.47%. It seems that these baselines are not well-tuned, and the gain of the proposed approach is not clear and could potentially fade away with a better-tuned baseline. Can the authors comment on this? Additionally, the authors should include the results using ResNet-50 on ImageNet, which should provide a stronger reference point. Fixed ETF Procedure: The authors only used the canonical simplex ETF for the fixed procedure. The weight matrix results in many zeros and could lead to poor performance when used as the fixed classifier because some neurons will be inactive. The authors should include the results using the fixed ETF with a non-canonical (i.e., projection on a random basis). Remarks: The authors should directly clarify in Tables 1 and 2 the ResNet architecture used (18 or 50).

Questions

See weaknesses section.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

For large-scale problems, the computational cost of the proposed approach could be a limitation due to the high memory cost of computing the backward pass of the Riemannian optimization. Therefore, the authors did not compute the gradient calculation when reporting their results for the image benchmarks. The authors claimed that they empirically observed no significant difference in performance for small-scale problems where the DDN gradient can be computed. Including these results in the supplementary material would also be beneficial. Moreover, we agree that future work should verify if this is true for large-scale problems.