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

Weaknesses:

Regarding the computational and memory costs of our approach, please refer to our general response.

Avoid backward pass: Theory suggests that incorporating the DDN layer's backward pass should provide additional gradient information for updating the features’ parameters in the backbone neural network. While we have observed that this additional gradient information improves stability, it has not been critical for the algorithm's convergence. Therefore, omitting the backward pass can be a practical choice in situations where computational resources are a concern.

Specifically, in our experiments with ImageNet, the current implementation of the DDN backward pass exceeds GPU memory limits. As a result, we opted not to use it for these experiments. However, we acknowledge the importance of the DDN backward pass for gradients to be technically correct, and plan to address these computational and memory challenges in future work to ensure that it can be effectively utilised even for large-scale datasets such as ImageNet.

Questions:

1. Minority collapse refers to the training dynamics observed in standard approaches when dealing with significant class imbalance. In such cases, minority classes tend to collapse into a single vertex of a simplex ETF, resulting in a degenerate solution. However, under certain conditions (e.g., when the number of features $d$ exceeds the number of classes $C$), a simplex ETF can still represent an optimal classifier configuration, even in imbalanced settings (see Theorem 1 of Yang et al. [67]). The primary challenge is that a learned classifier may struggle to achieve this configuration easily. It’s important to note that the assumption for neural collapse to occur is zero misclassification error during the TPT, which is clearly not met in the case of minority collapse. Substantial research on imbalance training and long-tail classification tasks employs fixed ETF approaches to address problems associated with conventional training methods. We believe that our approach, which induces the nearest ETF direction rather than an arbitrary one, will yield even better and more stable results compared to fixed ETF methods. We have included preliminary results (found in the attached pdf) on CIFAR100LT with an imbalanced ratio of 0.01 (following Yang et al. [65]) using a VGG network to support our claims. We have not considered graph neural networks in this work. Our method applies to the structure of the classifier solution as defined by the loss function. Exploration of alternative architectures and loss functions is left for future work.

2. Running the Riemannian optimisation problem every $k$ steps is indeed a viable extension of our method and can potentially reduce computational demands. We conducted experiments where $k$ was set to the number of steps per epoch. Our findings indicate that optimising once per epoch is insufficient for guiding the network to the nearest ETF and achieving rapid convergence, as the features can change significantly within an epoch. Thus, we opted to perform the optimisation at every step. Hence, we present a trade-off between achieving faster and more stable convergence versus higher computational cost. Introducing $k>1$ as an additional hyperparameter requires careful tuning for each specific case. Another advantage of our approach is that, as training progresses, the target ETF stabilises because the features approach convergence. At this point, the Riemannian optimisation problem yields similar solutions. To eliminate the additional computational cost, we can fix the simplex ETF determined by the Riemannian optimisation and continue to converge towards this fixed ETF. However, determining the exact epoch at which this stability is achieved remains heuristic and lacks theoretical guarantees, thus requiring case-by-case consideration.

3. Our approach is particularly well-suited for optimisers that dynamically adjust step sizes based on the current model state, such as AGD. In contrast, using methods like SGD or Adam with our approach requires careful tuning of the learning rate. This may involve developing a new type of scheduling that considers the objective value of the Riemannian optimisation problem to ensure that the learning rate is appropriately decayed as the solution approaches the ETF. Current learning rate schedulers are typically derived heuristically and optimised for conventional training scenarios, which may not account for the specific needs of our approach. Therefore, adapting these schedulers for our approach requires additional analysis, which is beyond the scope of this paper and will be addressed in future work.

Nit: $U_{init}$ is defined in the paragraph for Hyperparameter Selection and Riemannian Initialisation Schemes in lines 230-235 in our paper.

Nit: Thank you. This is a typo. The Frobenius norm should be taken wrt the feature means and not the features, so it should be $||\bar{H}||_F$.

Weakness 1 with question 1 & 2:

Eq. 8 is optimised as a bi-level optimisation problem. At each gradient update step, we first solve the inner optimisation problem to obtain the nearest ETF solution from the Riemannian problem. This gives the classifier weights directly. Subsequently, we perform the outer optimisation by optimising the rest of the network using stochastic gradient descent.

Proposition 1 restates a known result for gradient computation in differentiable optimisation problems, utilising the implicit function theorem for implicit differentiation (Gould et al. [21], Figure 2 and Section 4.1 provide the underlying intuition).

As illustrated in Figure 1 of our paper, two streams of information contribute to the final loss function. Proposition 1 enables us to backpropagate through both streams of information, theoretically enhancing the stability of our solution. The key insight is that incorporating the second stream of gradient information from the Riemannian optimisation allows us to account for changes in the features $H$ and adjust the network’s parameters accordingly. In practice, we have observed that the gradient magnitude of the second stream (bottom in Figure 1) is relatively small compared to the first stream (top in Figure 1). Coupled with the additional computational cost of computing the DDN gradients, this contribution is somewhat restricted in this case.

Weakness 2: Please refer to our general response for computational cost discussions.

Weakness 3:

Ablation studies: Our goal is to find the nearest simplex ETF of the feature means while considering all the features from the network. In the UFM case, where full-batch gradient updates are performed, all features are available, so additional techniques such as exponential moving averages are not necessary. However, in real-world scenarios with stochastic updates, optimising towards the nearest ETF for the features in a given batch can lead to highly variable results, as we do not have the full picture, and the feature means are constantly changing. By incorporating the described techniques, we ensure that we are moving towards a stable ETF target, accounting only for changes in feature weights after each gradient step.

Question 3:

Fixed ETF results: During VGG training with the fixed ETF case, we observed significant variability in the solutions, with CIFAR-10 showing the largest variability. This variation highlights the impact that fixing to a predefined ETF direction can have on the quality of the solution space. While ResNets, with their residual connections, can somewhat mitigate this issue, VGG-type networks appear to struggle more with fixed ETF directions. The intuition is that depending on the initialisation seed, we may start closer to or further from the chosen ETF solution, leading to results where fixed ETF performs either comparably or inferiorly. In contrast to fixed ETFs, the standard approach can sometimes be more robust and perform better. However, this comes with increased memory and computational costs due to the need to learn the classifier. This is because, with the standard method, the classifier has the flexibility to adapt and better match the features, eventually reaching a neural collapse solution in practice. Our approach, on the other hand, is robust to initialisation seeds by moving towards the closest ETF solution from any starting point and achieves such convergence more quickly.

Weaknesses:

Misleading results: The reviewer correctly observed that, particularly on smaller datasets, the performance converges to be approximately equivalent by the end of training. Any observed deviations are likely due to random effects. This is to be expected, as all properly trained ETF solutions (i.e., with non-random labels) should theoretically yield similar results. Our work leverages the symmetries inherent in the optimal solution geometry to reach an ETF solution in fewer iterations. Thus, the primary advantage of our approach lies in its faster convergence while maintaining comparable generalisation ability. Notably, in some cases, our method demonstrates superior generalisation performance because, within the given training time, other approaches may have yet to reach the optimal max-margin solution.

Non-competitive results: In our experiments, we utilised the AGD optimiser. While AGD achieves SoTA results for VGG architectures, it is not yet fully tuned for ResNet architectures. However, it has been demonstrated that state-of-the-art results can be obtained for ResNet with prolonged training (see [6]). Our focus, however, is not on achieving state-of-the-art results but rather on illustrating the convergence trajectory. Due to computational constraints and the required scale of experiments, we terminated training at 200 epochs.

Figure 4 no error bars: Fixed and replaced with 5 seed runs (found in the attached pdf).

Cost: Please see tables and discussion in the general response.

Questions:

Optimal ETF solution: As previously discussed, all ETF solutions are equivalent up to rotations and permutations. Our Riemannian optimisation approach ensures that with each gradient update step, we move towards the nearest ETF solution, effectively breaking these symmetries and achieving faster convergence. Consequently, while our solution is optimal wrt the Riemannian problem, it is also optimal for the learning problem, although no better than any other ETF solution.

Compare two ETFs (ours vs fixed): In theory, we expect the testing accuracy of the two solutions to be equivalent. However, in practice, due to the finite amount of training and the complexity of tasks such as those with many classes (e.g., ImageNet), we observe that faster convergence to the ETF solution can lead to better generalisation performance. However, we are careful not to make any theoretical claims in this regard.

Proximal term: Solving the original problem in Eq. 6 yields a non-unique solution due to the rank deficiency of matrix $M$, resulting in a family of possible solutions. To achieve a unique solution, we introduce a proximal term to make the problem strongly convex, as shown in Eq. 7. This approach stabilises the solution. Additionally, the DDN backward pass for Eq. 6 cannot be computed deterministically because of the singularity of the gradients. By incorporating the proximal term, we obtain non-singular gradients, which allows us to compute Prop. 1.

$U_{prox}$: An obvious initialisation scheme for both $U_{prox}$ and $U_{init}$​ is to set them as either canonical or randomly orthogonal directions. However, both approaches are sensitive to the initial parameter settings of the network. We found that the most effective strategy is to solve the Riemannian problem in Eq. 6 at the first gradient update step (corresponding to the first epoch in GD and the first mini-batch in SGD) to obtain a solution from the family of solutions. This solution, denoted $U*$, is then used as the initialisation for both $U_{init}$ and $U_{prox}$ in Eq. 7. We then solve Eq. 7 and perform the first backward pass. The intuition behind this initialisation scheme is to start with a solution that is both feasible and optimal for the original problem formulation. Empirically, this approach has yielded the best and most stable results.

$\delta$ proximal param: $\delta$ is a tradeoff between the distance of the feature means from an ETF and the distance between the target direction ($U_{prox}$) and the currently optimised direction. We have found that a broad range of $\delta$ values provides stable and good performance. Therefore, $\delta$ can be set within this range, provided it is not too small (e.g., $\delta\ll 10^{-7}$, which essentially solves Eq. 6) and not too large (e.g., $\delta\gg 10$, where the proximal term dominates and results in a fixed ETF case).

$\tau$ param in CE: Since we normalise features, the cross-entropy loss will have a non-zero lower bound (Thm. 1 of Yaras et al. [67]). $\tau$ adjusts this lower bound, bringing it closer to zero. According to Yaras et al., including this parameter results in more defined NC metrics while keeping accuracy unchanged. They empirically found that setting $\tau=5$ yields optimal results for real-world datasets and architectures, and we adopt this approach in our work.

Clarification on Prop.1: The results in the proposition are currently presented wrt the unnormalised feature means matrix $\bar{H}$. However, this choice does not affect the result or its derivation. For consistency, we will update the proposition to use the normalised values, i.e., $\tilde{H}$. Gould et al. [21] derived their result for vector functions with vector variables, whereas we have generalised this result to matrix functions and matrix variables as a natural extension.

Error bars in Figures: All experiments were conducted five times. The results for UFM (Fig 2, 3) are highly stable, with minimal deviation, and the ranges are easily visible due to the plots' linewidth. $P_{CM}$, defined in Eq. 16, represents the cosine margin of each specific feature.

STL VGG range: We applied our method to various datasets and architectures without case-by-case tuning. We suspect that the initial variation observed in the STL case is related to the relatively large batch size used for that dataset. However, our method effectively reduces the variance within just a few epochs compared to the fixed ETF.

Overhead Cost: Please refer to the tables in our general response and the discussion around computational concerns.

Standard Procedure: Our new experiments show that the standard method, which excludes feature and weight normalisation, achieves similar performance (found in the attached pdf). However, we observe that the learned features are less well-aligned with a max-margin classifier compared to when normalisation schemes are included.

Image Baselines are not SOTA: Since our results overlap with those previously reported in the AGD paper, it is demonstrated that state-of-the-art results can be achieved by extending training over more epochs using the AGD optimiser (Figure 4 in Berstein et al. (2023)). However, the primary focus of our work is not to establish a new state-of-the-art result but to study the neural collapse phenomenon. We originally included results using ResNet50 on ImageNet (Figure 4 in our paper), and in the attached pdf, we also included ImageNet results after five runs with error bars.

Fixed ETF procedure: The reviewer’s observation that the canonical projection results in a sparse weight matrix would have been true if we were working with canonical orthogonal frames. For canonical simplex equiangular tight frames, the sparsity is observed in the semi-orthogonal matrix $U$ (i.e., $U$ is the identity matrix), not in the weight matrix itself, which remains dense (as shown in Equation 1). To address this concern, we conducted additional experiments with the fixed method on CIFAR-10 using ResNet-18, where we replaced the canonical rotation with a random orthogonal rotation (drawn from the Haar measure). Our results indicate that the performance remains superior and that the fixed and canonical methods yield nearly identical outcomes (found in the attached pdf).

Remark: We have updated our table results to reflect the chosen ResNet and VGG architectures.

Limitations:

We will include the results in the supplementary material that test whether computing the DDN gradient affects the stability of our solution on small-scale problems.