Guiding Neural Collapse: Optimising Towards the Nearest Simplex Equiangular Tight Frame
Authors: Evan Markou
Thalaiyasingam Ajanthan
Stephen Gould
Venue: NeurIPS 2024 poster (2024)
Keywords: neural collapse
equiangular tight frames
Riemannian optimisation
deep learning
Primary area: optimization_for_deep_networks

Abstract: Neural Collapse (NC) is a recently observed phenomenon in neural networks that characterises the solution space of the final classifier layer when trained until zero training loss. Specifically, NC suggests that the final classifier layer converges to a Simplex Equiangular Tight Frame (ETF), which maximally separates the weights corresponding to each class. By duality, the penultimate layer feature means also converge to the same simplex ETF. Since this simple symmetric structure is optimal, our idea is to utilise this property to improve convergence speed. Specifically, we introduce the notion of \textit{nearest simplex ETF geometry} for the penultimate layer features at any given training iteration, by formulating it as a Riemannian optimisation. Then, at each iteration, the classifier weights are implicitly set to the nearest simplex ETF by solving this inner-optimisation, which is encapsulated within a declarative node to allow backpropagation. Our experiments on synthetic and real-world architectures on classification tasks demonstrate that our approach accelerates convergence and enhances training stability.

Guiding Neural Collapse: Optimising Towards the Nearest Simplex Equiangular Tight Frame
Evan Markou
Thalaiyasingam Ajanthan
Australian National University
Australian National University & Amazon
evan markou@anu. edu.
chalaiyasingam. a ajanthan@anu. edu. au
Stephen Gould Australian National University stephen. gouldBanu. edu. au stephen.goulddanu.
Abstract
Neural Collapse (NC) is a recently observed phenomenon neural networks that characterises the solution space of the final classifier layer when trained until zero training loss. Specifically, NC suggests that the final classifier layer converges to a Simplex Equiangular Tight Frame (ETF), which maximally separates the weights corresponding to each class. By duality, the penultimate layer feature means also converge to the same simplex ETF. Since this simple symmetric structure is optimal, our idea is to utilise this property to improve convergence speed. Specifically, introduce the notion of nearest simplex ETF geometry for the penultimate layer features at any given training iteration, by formulating it as a Riemannian optimisation. Then, each iteration. the classifier weights are implicitly set the nearest simplex ETF by solving this inner-optimisation, which is encapsulated within declarative node to allow backpropagation. Our experiments synthetic and real-world architectures for classification tasks demonstrate that our approach accelerates convergence and enhances training stability
1 Introduction
While modern deep neural networks (DNNs) have demonstrated remarkable success in solving diverse machine learning problems 22
34 38
the fundamental mechanisms underlying their
training process remain elusive. recent years, considerable research efforts have focused on delineating the optimisation trajectory and characterising the solution space resulting from the optimisation process training neural networks [72 17 41. One such finding is that gradient descent algorithms, when combined with certain loss functions, introduce an implicit bias that often favours max margin solutions, influencing the learned representations and decision boundaries. 44
57
33
27
21
54
70 31 49].
In this vein, Neural Collapse (NC) recently observed phenomenon neural networks that characterises the solution space of the final classifier layer in both balanced [50 76
74
28
48 63
43 32
and imbalanced dataset settings [19 55 [.. Specifically, NC suggests that the final classifier
layer converges to Simplex Equiangular Tight Frame (ETF), which maximally separates the weights corresponding to each class, and by duality, the penultimate layer feature means converge to the classifier weights, i.e., to the simplex ETF (formal definitions are provided in Appendix A). This simple, symmetric structure shown to be the only set of optimal solutions for a variety of loss functions when the features are also assumed to be free parameters, i.e., Unconstrained Feature
Code available attps//ggthhcoocccmmmmmmmddodode
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

28 75 Nevertheless, even in realistic large-scale deep networks, this phenomenon observed when trained convergence, even after attaining zero training error.
Models (UFMs) B2 19
74
76
Since we can characterise the optimal solution space for the classifier layer, a natural extension is to leverage the simplex ETF structure of the classifier weights to improve training. To this end, researchers have tried fixing the classifier weights to a canonical simplex ETF, effectively reducing the number of trainable parameters [76]. However, in practice, this approach does not improve the convergence speed as the backbone network still needs to do the heavy lifting of matching feature means to the chosen fixed simplex ETF.
In this work, we introduce mechanism for finding the nearest simplex ETF to the features at any given training iteration. Specifically,, the nearest simplex ETF is determined by solving Riemannian optimisation problem. Therefore, our classifier weights are dynamically updated based on the penultimate layer feature means each iteration, Le., implicitly defined rather than trained using gradient descent. Additionally, by constructing this inner-optimisation problem as a deep declarative node [23), we allow gradients to propagate through the Riemannian optimisation facilitating end-to- end learning. Our whole framework significantly speeds up convergence to a NC solution compared to the fixed simplex ETF and conventional learnable classifier approaches. We demonstrate the effectiveness of our approach on synthetic UFMs and standard image classification experiments.
Our main contributions are follows:
1. We introduce the notion of the nearest simplex ETF geometry given the penultimate layer features. Instead of selecting a predetermined simplex ETF (canonical or random), implicitly fix the classifier as the solution to a Riemannian optimisation problem.
2. To establish end-to-end learning, we encapsulate the Riemannian optimisation problem of determining the nearest simplex ETF geometry within a declarative node. This allows for efficient backpropagation throughout the network.
3. We demonstrate that our method achieves an optimal neural collapse solution more rapidly compared to fixed simplex ETF methods conventional training approaches, where learned linear classifier is employed. Additionally, our method ensures training stability by markedly reducing variance in network performance.
2 Related Work
Neural Collapse and Simplex ETFs. Zhu et al. [76 proposed fixing classifier weights to simplex ETF, reducing parameters while maintaining performance. Simplex ETFs effectively tackle imbalanced learning, as demonstrated by Yang et al. (67], where they fix the target classifier to an arbitrary simplex ETF, relying on the network's over-parameterisation to adapt. Similarly, Yang et al. [68] addressed class incremental learning by fixing the target classifier to a simplex ETF They advocate adjusting prototype means towards the simplex ETF using a convex combination, smoothly guiding backbone features into the targeted simplex ETF. However, these methods did not yield any benefits regarding convergence speed. The work most relevant to ours is that of Peifeng et al. [5]]. who argued about the significance of feature directions, particularly in long-tailed tailed learning scenarios. They compared their method against a fixed simplex ETF target, formulating their problem enable the network to learn feature direction through a rotation matrix. Additionally, they efficiently addressed their optimisation using trivialisation techniques [39 40). However, they did demonstrate any improvements in convergence speed over the fixed simplex ETF, achieving only a minimal increase in test accuracy. Fixing a classifier is not a recent concept, as it has been proposed prior to the emergence of neural collapse [52] 58 30) Most notably, Pernici et al. [52] demonstrated improved convergence speed by fixing the classifier to a simplex structure only on ImageNet while maintaining comparable performance on smaller-scale datasets. In contrast, our method shows superior convergence speed compared to both a fixed simplex ETF and a learned classifier across both small and large-scale datasets.
Optimisation on Smooth Manifolds. Our optimisation problem involves orthogonality constraints, characterised by the Stiefel manifold [1]. Due to the nonlinearity of these constraints, efficiently solving such problems requires leveraging Riemannian geometry [8] A multitude of works are dedicated to solving such problems by either transforming existing classical optimisation techniques
into Riemannian equivalent algorithms 73 20 64 [55
by carefully designing penalty functions

---

to address equivalent unconstrained problems (66 65 In our approach, we opt for retraction-basee Riemannian optimisation algorithm [T to optimally handle orthogonality constraints.
Implicit Differentiable Optimisation. In neural network setting, end-to-end architectures commonplace. To backpropagate solutions to optimisation problems, we rely on machinery from implicit differentiation. Pioneering works [ demonstrated efficient gradient backpropagation when dealing with solutions convex optimisation problems. This concept was independently introduced generalised version by Gould al. [23 24 to encompass any twice-differentiable optimisation problem. key advantage of Deep Declarative Networks (DDNs) lies in their ability efficiently solve problems any scale by leveraging the problem's 's underlying structure [25]. Our setting involves utilising an qquaaity-oomstrained declarative node to efficiently backpropagate through the network.
Optimising Towards the Nearest Simplex ETF
In this section, we introduce our method determine the nearest simplex ETF geometry and detail how we can dynamically steer the training algorithm to converge towards this particular solution.
3.1 Problem Setup
Let us first introduce key notation that will be useful when formulating our optimisation problem.
Simplex ETF Mathematically, general simplex ETF is collection of points in RCC specified by the columns of a matrix
M M- C
(I)
Here, a E R2 denotes an arbitrary scale factor, 1e is the C-dimensional vector of ones, and U Rdx (with d 2 C) represents semi-orthogonal matrix (U / Ic). Note that there many simplex ETFs in RC as the rotation U varies, and M is rank-deficient. Additionally, the
standard simplex ETF with unit Frobenius norm is defined as: M vC-I (Ic - dlcl)).
Mean of Features. Consider a classification dataset D {x ys)
1....,Ny .N} where the data
a; E X and labels V/ E = (1...., C). Suppose, is the number of samples correspond to label then De ne N. Let consider scenario where have collection of features defined as,
H hhe.i 1 c C, 1. i ne] RdxN
Here, each feature may originate from a nonlinear compound mapping of input data through neural network, denoted as, hygd (a;) for the data sample (1) Now, for the final layer, our decision variables (weights and biases) are represented as W 2 [wi...., welT E RCxd and b E C and the logits for the i-th sample is computed as,
whe (x.) Wh b where o0
In UFMs, the features assumed to be free variables, which serves as rough approximation for neural networks and helps derive theoretical guarantees. Additionally, we define the global mean and per-class mean of the features } as:
ha
hag
h
Σ hejs (1SesC), C) iml
(4)
N
The
and the globally centred feature mean matrix as,
H
h ha. he ha] E RdxC
(5)
Finally, we scale the feature mean matrix to have unit Frobenius norm, i.e., A H|| which will be used in formulation below.
3

---

3.2 Nearest Simplex ETF through Riemannian Optimisation
Once we obtain the feature means, our objective is to calculate the nearest simplex ETF based on these means and subsequently adjust the classifier weights W to align with this particular simplex ETF. The rationale means any given iteration. This approach aims expedite convergence during the training process by providing the algorithm with a starting point that closer to an optimal solution rather than requiring it to learn a simplex ETF direction converge towards an arbitrary one.
identify and establish simplex ETF that closely corresponds the feature
To find the nearest simplex ETF geometry, we solve the following Riemannian optimisation problem
minimize UESC
H - UM
UM-UM]
(6)
where Stg {x RdxC x Ic}. Here, M the standard simplex ETF with unit
Frobenius norm, and the of the orthogonality constraints Sig forms compact Stiefel manifold 2 embedded Euclidean space.
3.3 Proximal Problem
The solution to the Riemannian optimisation problem, denoted as U is not unique since a component of U lies in the null space of M. As simplex ETFs reside in (C' - >dimensional space, the matrix M is rank-one deficient. Consequently, we are faced with a family of solutions, leading to challenges in training stability, as we may oscillate between multiple simplex ETF directions. We address this issue by introducing proximal term the problem's s objective function. This guarantees the uniqueness of the solution and stabilises the training process, ensuring that our problem converges to solution closer the previous one.
So, the original problem in Equation Gls transformed into:
minimize UESI
H HM
-
U prxx |2
Here, Upmox represents the proximal target simplex ETF direction, and 8 > 0 serves as the proximal coefficient, handling the trade-off between achieving the optimal solution's proximity the feature means and its proximity to given simplex ETF direction. In fact, one can perceive our problem generalisation predetermined fixed simplex ETF solution. This is
formulation in Equation evident when considering that if we significantly increase 6, the optimal direction U would converge towards the fixed proximal direction proxx
3.4 General Learning Setting
Our problem formulation, following the general deep neural network architecture in Equation can be seen as a bilevel optimisation problem as follows:
N log
exp (ve (a U²)),)
minimize L(D:O,U)) &
N
Σ exp (vh (x, U+);)
(8)
8 U U prox
subject to U E arg aig
H UM
2
UESI
where (æ, U*) = TMU classifier weights are as W MU" and the bias is set centring. Furthermore, M is the standard simplex ETF, M is its normalised version, and H is the normalised centred feature matrix. The temperature parameter 7 > 0 controls the lower bound of the cross-entropy loss when dealing with normalised features, defined in 69 Theorem 1].
ho) with h; = the (xx:). Here, vi denotes the logits, where the MU*hG account for feature
3.5 Handling Stochastic Updates
In practice, we use stochastie gradient descent updates, and, as such, adjustments to our computations necessary. With each gradient update now based mini-batch, implement two key

---

classifier
L
H
W -UM
normalised
H
U*
DDN (Eq.)
7
means
U U prox
Figure 1: Schematic of our proposed architecture for optimising towards the nearest simplex ETF. The classifier weights W = U M are an implicit function of the CNN features H. Note that the parameters of the CNN are updated via two gradient paths from the loss function C, direct path (top) and an indirect path through U (bottom).
changes. First, rather than directly optimising the problem of finding the nearest simplex ETF geometry concerning the feature means of the mini-batch, we introduce an exponential moving average operation during the computation of the feature means. This operation accumulates statistics and enhances training stability throughout iterations. Formally, at time step I, we have the following equation, where or E IR represents the smoothing factor:
H. = a + (1-a) H.
Second, we employ stratified batch sampling to guarantee that all class labels are represented in the mini-batch. This ensures that we avoid degenerate solutions when finding the nearest simplex ETF geometry, as optimisation problem requires input feature means for all C' classes. In cases where the number of classes exceeds the chosen batch size, we compute the per-class feature mean for the class labels present in the given batch. For the remaining class labels, we set their feature mean the global mean of the batch. We repeat this process for each training iteration until we have sampled examples belonging to the missing class labels. At that point, update the feature mean of those missing class labels with the new feature statistics. We reserve this method only for cases where the batch size is smaller than the number of labels since can introduce instability during early iterations.
3.6 Deep Declarative Layer
We can backpropagate through the Riemannian optimisation problem update the feature means using a declarative node [23] Then, the features are updated from both the loss and the feature means through uuto-differentiation. The motivation for developing the DDN layer lies in recognising that, despite the presence of a proximal term, abrupt and sudden changes to the classifier may occur as the features are updated. These changes can pose challenges for backpropagation, potentially disrupting the stability and convergence of the training process. Incorporating an additional stream of gradients through the feature means to account for such changes, as depicted Figure [ assists in stabilising the feature updates during backpropagation.
To efficiently backpropagate through the optimisation problem, we employ techniques described in Gould et al. [23) utilising the implicit function theorem to compute the gradients. In our case, we have scalar objective function f R + IR and matrix constraint function J IR - CCxC Since we have matrix variables, we use vectorisation techniques [46 to avoid numerically dealing with tensor gradients. More specifically, we have the following:
Proposition (Following directly from Proposition 4.5 in Gould et al. (3)). Consider the opti- misation problem Equation 7 Assume that the solution exists and that the objective function f and the constraint function J are twice differentiable in the neighbourhood of the solution. If the rank(A A) CiC+1) and G is non-singular then:
Dy(U)= A
(AG AAAG
AT)- -
where,
A = rvech DD. H U)) E R ciC+1 xdC
B = Dve HU JH, U)) E dCxd
G rvoc(Dip JH, U)) - A rvech(Di J(H, e dCCxd
5

---

Here, the double dot product symbol (:) denotes a tensor contraction on appropriate indices between the Lagrange nultiplier matrix and a fourth-order tensor Hessian. Also, rvec(.) and and rvech(.) refer to the row-major vectorisation and half-vectorisation operations, respectively. To find the Lagrange we solve the following equation where we have vectorised the matrix
multiplier matrix A E RC
as X LER C(C41)
A Du (H.
Alternatively, for more efficient computation of the identity G, we can utilise the embedded gradient field method as defined in Birtea et al. 19 TO]. Therefore, we obtain:
= Dvv f(H, U)) La ((U) E RddCxdC
where E(U) - 1 (Dv (H, + U Du ((H U1). and & here denotes Kronecker product.
detailed derivation of each identity in the proposition can be found in Appendix B
4
Experiments
In our experiments, we perform feature normalisation onto a hypersphere, a common practice in training neural networks, which improves representation and enhances model performance [71 (6)
66 35]. We find that combining classifier weight normalisation with feature normalisation accelerates convergence [69]. Given that simplex ETFs are inherently normalised, we include classifier weight normalisation in our standard training procedure to ensure fair method comparisons.
53 42 61
Experimental Setup. In this study, we conduct experiments on three model variants. First, the standard method involves training model with learnable classifier weights, following conventional practice. Second, in the fixed ETF method, we set the classifier to predefined simplex ETF. In all experiments, we choose the simplex ETF with canonical direction. In Appendix C we also include additional experiments for fixed simplex ETFs with random directions generated from a Haar measure 477. Last, our implicit ETF method, where we set the classifier weights on- on-the-fly as the simplex ETF closest to the current feature means.
We repeat experiments on each method five times with distinct random seeds and report the median values alongside their respective ranges. For reproducibility and to streamline hyperparameter tuning. employed Automatic Gradient Descent (AGD) ( Following the authors' recommendation, set the gain/momentum parameter to 10 to expedite convergence, aligning it with other widely used optimisers like Adam (36) and SGD. Our experiments real datasets run for 200 epochs with batch size 256; for the UFM analysis, we run 2000 iterations.
Our method underwent rigorous evaluation across various UFM sizes and real model architectures trained on actual datasets, including CIFARIO (37], CIFARI00 [37]. STL10 [14), and ImageNet- 1000 [15,, implemented on ResNet (29] and VGG [56] architectures. More specifically,
trained
CIFAR10 on 8 and VGG13, CIFAR 100 and STL10 on ResNet50 and VGG13, and ImageNet- 1000 ResNet50. The input images were preprocessed pixel-wise by subtracting the mean and dividing by the standard deviation. Additionally, standard data augmentation techniques were applied, including random horizontal flips, rotations, and crops. All experiments were conducted using Nvidia RTX3090 and A100 GPUs
Hyperparameter Selection and Riemannian Initialisation Schemes. We solve the Riemannian optimisation problem defined in Equation 7 using a Riemannian Trust-Region method [ from pyManopt [60]. We maintain a proximal coefficient set 10-3 consistently across all experiments. It is worth mentioning that algorithm convergence is robust to the precise value of 5. In our problem, determining values for Uni and Uprox crucial. We explored several methods to initialise these parameters. One approach involved setting both towards the canonical simplex ETF direction. This means initialising them as a partial orthogonal matrix where the first C rows and columns form an identity matrix while the remaining d - C' rows are filled with zeros. Another approach is to initialise both of them as random orthogonal matrices from elassical compact groups, selected according to a Haar measure [7]. In the end, the approach that yielded the most stable results at initialisation was

---

12
1.30
ImpictET
10
LE lawer bound
TPT Phose: Standard
08
Phaae Ficed
0.
siandard
1.00
04
Implot
Thcorotical ETF Margin
standard
0.2
TPT Phase Standare
fixed ETF
0.0
TPT Phase - sed
TPT Phase implet
0.90
Theonetical ETT Hargin
implicd ETF
1200
1000
2000
400
600
800
3000
(a) CE loss
(b) Avg Cos. Margin
(c) PcM at EoT
0.8
Rlad
standard
standard
ETF
0.025
shander
Implict
TPT Prose Standard
IPI Phase Read
1.0
12
IPI Puase- Feed
TPT Phose Standard
Implct
impiott
Standard
TPT Prase Implic
0.8
TPT Phase Implck
0.035
0.4
0.5
0.010
0.2
6.0
0.0
0.000
1600
2000
1600
(d) NC1
(e) NC3
W H E--Euunoomm
Figure 2: UFM-10 results. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
to employ either of the aforementioned initialisation methods to solve the original problem without the proximal term in Equation 6 We then used the obtained U to initialise both Unii and U prox for the problem in Equation 7 This process was carried out only for the first gradient update the first epoch. In subsequent iterations, we update these parameters to the U obtained from the previous time step. Importantly, the proximal term held fixed during each Riemannian optimisation.
Regarding the calculation of the exponential moving average of the feature means, we have found that employing& decay policy on the smoothing factor yields optimal results. Specifically, we set 2/ 1), where T represents the number of iterations. Additionally, we include a thresholding value of 10-1 such that if a falls below this threshold, we a to be equal to the threshold. This precaution ensures that a does not diminish throughout the iterations, thereby guaranteeing that the newly calculated feature means contribute sufficient statistics to the exponential moving average.
Finally, in our experiments, we set the temperature parameter r to five. This choice aligns with the findings discussed by Yaras al. [69], highlighting the influence of the temperature parameter value on the extent of neural collapse statistics with normalised features.
Unconstrained Feature Models (UFMs). Our experiments on UFMs, which provide a controlled setting for evaluating the effectiveness of our method, done using the following configurations:
UFM-10: 10: a 10-class UFM containing 1000 features with dimension of 512.
UFM- 100: 100-class UFM containing 5000 features with a dimension of 1024.
UFM-200: a 200-class UFM containing 5000 features, with a dimension 1024.
UFM-1000: a 1000-class UFM containing 10000 features, with dimension of 1024.
Results. We present the results for the synthetic UFM-10 case in Figure 2 The CE loss plot demonstrates that fixing the classifier weights to a simplex ETF achieves the theoretical lower bound of Yaras al. 66 Thm. 1], indicating the attainment of a globally optimal solution. We also visualise the average cosine margin per epoch and the cosine margin distributions of each example at the end of training, defined in Zhou al. 774. The neural collapse metrics, NC1 and NC3, which measure the features' within-class variability. and the self-duality alignment between the feature means and the classifier weights 776,, are also plotted. Last, we depict the absolute difference of the classifier and feature means norms to illustrate their convergence towards equinorms. as described in Papyan et al. [50]. comprehensive description of the metrics can be found in Appendix [A] Collectively, the plots indicate the superior performance of our method in achieving neural collapse (NC) solution
7

---

Sonc
Standerz
Phose Slandad
TPT Pheae Stander
Priss Standard
Theae hood
(99) Phoo hed
19T Thoat hood
Procs Paot
(a) UFM-10
(b) UFM-100
UFM-200
(d) UFM-1000
Figure 3: The evolution of convergence measured in top-1 accuracy of the UFM as we increase the number of classes, plotted for the first 800 epochs. We omit the rest of the epochs as all methods have converged and have identical results.
Table 1:
Train top-1 accuracy results presented median with indices indicating the range of values from five random seeds. Best results are bolded.
Train accuracy at epoch 50
Train accuracy epoch 200
Dataset
Network
Standard
Fixed ETF Implicit ETF
Standard
Fixed ETF
Implicit ETF
ResNet18
87.42
89.7
86.89
88.6
88.71 89.6
96.69 98.6
97.18
97.9
98.09 98.6
CIFAR10
86.1
84.7
88.5
96.5
95.6
97.9
VGG13
97.0 93.59
85.8 76.66
95.69
96.2
99.15
99.8
79.36
99
99.56
99.7
90.7
53.9
95.2
97.9
59.6
99.4
CIFAR100
ResNet50
58.47
59.6
53.9
63.93
65.2
61.1
72.15
70.1
74.1
98.6 95.87
91.34
90.4
96.96
97.3
96.2
VGG13
82.00
84.0
81.14
$1.9
88.39
89.4
99.6 99.34
94.55
95.3
98.92
99.0
80.5
76.0
86.9
99.2
92.5
98.8
ResNet50
90.7 83.86
86.76 $6.8
93.54
95.3
99.42
99.9
98.38
99
99.72
99.9
STLI0
81.5
91.3
99.0
98.
98.0
VGG13
82.66
73.7
90.7
83.60
85.1
65.6
90.14
93.5
69.2
100.0 100
99.92
90.9
99.8
100.0
100
100
ImageNet
ResNet50
58.35
58.0
59.1
58.0
70.44 68.3
74.09 73.8
74.5
77.20 77.3
83.09 83.
83.6
88.01
88.5
87.5
faster than other approaches. In Figure demonstrate under the UFM setting that increase the number of classes, our method maintains constant performance and converges at the same rate, while the fixed ETF and the standard approach require time to reach the interpolation threshold.
Numerical results for the top-1 train and test accuracy are reported in Tables D and respectively. The results provided for snapshots taken epoch 50 and epoch 200. It is evident that our method achieves faster convergence speed compared to the competitive methods while ultimately converging to the same performance level. Additionally, it is noteworthy that method exhibits the smallest degree of variability across different runs, as indicated by the range values provided. Finally, in Figure we present qualitative results that confirm faster and reach peak performance earlier than the standard and fixed ETF methods on ImageNet. It's important to note that the standard method with AGD is reported to converge the same testing accuracy (65.5%) at epoch 350, as shown Bernstein et al. (.. Figure 4]. At epoch 200, the authors exhibit a testing accuracy of approximately 51%. Since have increased the gain parameter on AGD compared to the results reported in the original paper, we report final 60.67% testing accuracy for the standard method, whereas our method reaches peak convergence approximately epoch 80. We note that the ImageNet results reported in Tables and as well as Figure are generated solely by solving the Riemannian optimisation problem without considering its gradient stream on the feature updates, due to computational constraints. We discuss the computational requirements of our method in Section We also present qualitative results for all the other datasets and architectures in Appendix
solution' 's ability to converge much
Discussion: Limitations and Future Directions
Our method involves two gradient streams updating the features, as depicted in Figure Interestingly. empirical observations small-scale datasets (see Figure [5] indicate that even without the back- propagation through the DDN layer, the performance remains comparable, rendering the gradient calculation of the DDN layer optional. In Figure 15c we observe a strong impact of the DDN layer gradient on the atomic feature level, with more features reaching the theoretical simplex ETF margin by the end of training. To reach consensus the exact effectt of the DDN gradient on the learning
8

---

Table 2: Test top-1 accuracy results presented as median with indices indicating the range of values from five random seeds. Best results are bolded.
Test accuracy at epoch 50
Test accuracy at epoch 200
Dataset
Network
Standard
Fixed ETF
Implicit ETF
Standard
Fixed ETF
Implicit ETF
CIFAR1O
ResNet18
80.47
79.4
82.6
80.63 $1.8
79.4
81.76 82.0
81.4
84.8 83.97
83.2
84.53
84.9
83.7
84.78 85.
84.3
VGG13
86.70
89.4
70.99
80.7
88.30 85.7
91.5 90.34
72.48
90.5
90.98 91.5
83.7
50.7
87.4
89.1
56.9
90.6
CIFAR100
ResNet50
46.3 45.91
45.6 45.37
49.3 48.38
61.9 51.29
48.03
47.7
50.52
51.2
VGGI3
60.82
59.3
60.10
57.1
62.39
61.8
63.6
68.1 67.54
66,4
62.78
63.4
61.6
67.14
66.8
67.6
ResNet50
55.41
58.3
58.11
60.0
59.2 57.56
65.2 63.60
64.33
66.5
62.75
63.0
STL10
64.8
57.1
56.4
61.8
63.9
60.7
VGG13
66.09
72-a
66.15 67-7
68.78
71.3
81.61
81.9
79.53
80.3
79.94
56.2
81.3
ImageNet
ResNet50
53.3 52.64 52.3
58.85
59.1
63.40
63.8
60.7 60.02 59.8
61.47
62.0
65.36
65.7
58.3
63.2
61.4
65.0
55
02
-33
00
Sone
-1.
ingi FT1
325
030
0.75
125
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) Pom at EoT
foad ETC
0.35
mander impiceii
food FIF implin
0:20
2.0
(g) Test Top-1 Acc.
120
350
350
122
250
(d) NC1
(e) NC3
(f) W H -Equinorm
Figure 4: ImageNet results on ResNet-50. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
process, further experiments on large-scale datasets are needed. However, on large-scale datasets with large d and C, such as ImageNet, computing the backward pass of the Riemannian optimisation challenging due the memory inefficiency of the current implementation of DDN gradients. This limitation is an area we aim to address in future work. Note that in all other experiments, we use the full gradient computations, including both direct and indirect components, through the DDN layer. We summarise the GPU memory requirements for each method across various datasets in Table 3
Our discussion so far has focused convergence speed in terms of the number of epochs required for the network to converge. However, it also important to consider the time required per epoch. In our case, training progresses, the time taken by the Riemannian optimization quickly becomes almost negligible compared to the network's total forward pass time, while it approaches the standard and fixed ETF training forward times, shown in Figure 5a However, DDN gradient computation increases considerably when the feature dimension d and the number of classes C' increase and starts dominate the runtime for large datasets such ImageNet. Nevertheless, for ImageNet, we do not compute the DDN gradients and still outperform other methods, We plan to explore ways to expedite the DDN forward and backward pass in future work.
Conclusion
In this paper, we introduced novel method for determining the nearest simplex ETF to the penul- timate features of neural network and utilising as our target classifier at each iteration. This contrasts with previous approaches, which either fix to specific simplex ETF or allow the network

---

Table 3: GPU memory (in Gigabytes) during training.
Standard
Fixed ETF
Implicit ETF
Model
w/o DDN Bwd
w/ DDN Bwd
UFM-10
1.5
1.5
1.6
UFM-100
1.7
1.7
1.7
10.7
UFM-200
1.7
1.7
1.8
70.3
UFM-1000
1.9
1.9
2.9
N/A
ResNet18 CIFARIO 10
2.2
2.2
2.2
2.3
ResNet50 CIFARIO
2.6
2.6
2.7
2.8
ResNet50 CIFAR
2.6
2.6
2.7
18.9
ResNet50 ImageNet
27.5
27.2
27.8
N/A
stersdard
102
500
implicit ETT
Impilit
400
median impliii
307
100
301
100
125
150
(a) Forward pass times milliseconds.
(b) Forward and backward times (log) millisecs.
Figure 5: CIFARIO0 computational cost results ResNet-50. In (a), plot the forward pass time for each method. For the implicit ETF method, which has dynamic computation times, we also include the mean and median time values. In (b), plot the computational cost for each forward and backward pass across methods. For the implicit ETF forward pass, we have taken its median time. The notation is as follows: S/F Standard Forward Pass, S/B Standard Backward Pass, F/F Fixed ETF Forward Pass, F/B = Fixed ETF Backward Pass, I/F Implicit ETF Forward Pass, and I/B Implicit ETF Backward Pass.
to learn it through gradient descent. Our method involves solving Riemannian optimisation problem facilitated by a deep declarative node, enabling backpropagation through this process.
We demonstrated that our approach enhances convergence speed across various datasets and architec- tures while also reducing variability stemming from different random initialisations. By defining the optimal structure of the classifier and efficiently leveraging its rotation invariance property to find the one closest to the backbone features, we anticipate that our method will facilitate the creation of new architectures and the utilisation of new datasets without necessitating specific learning tuning of the classifier's s structure.
References
[1] P.-. Absil, C. G. Baker, and K. A. Gallivan. Trust-region methods on Riemannian man- ifolds. Foundations of Computational Mathematics, 7(3):303-330, 2007. doi: 10.1007/ s (0208-005-0119-9)
[2] P.-. Absil, Mahony, and R. Sepulchre. Optimization Algorithms on Matrix Manifolds. Princeton University Press, 2008. ISBN 9780691 132983. URL http:::ttt.:wsseeeeee org/
stable/j.ctt7mmel
[3] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances neural information processing systems, 32, 2019.
10

---

[4] Brandon Amos and J Zico Kolter. Opmnet: Differentiable optimization as a layer in neural networks. International Conference Machine Learning, pp- 136-145, PMLR, 2017.
[5] Tina Behnia, Ganesh Ramachandra Kini, Vala Vakilian, and Christos Thrampoulidis. On the implicit geometry of cross-entropy parameterizations for label-imbalancee data. In International Conference Artificial Intelligence and Statistics, pp. (0815-10833) PMLR, 2023.
[6] Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic Gradient Descent: Deep Learning without Hyperparameters. arXiv:2304.05187, 2023,
[7] Petre Birtea and Dan Comänescu. Geometrical dissipation for dynamical systems. Communica- tions in Mathematical Physics, 316:375-394 2012.
[8] Petre Birtea and Dan Comânescu. Hessian operators on constraint manifolds. Journal of Nonlinear Science, 25:1285-1308 2015.
[9] Petre Birtea, Ioan Caşu, and Dan Comânescu. First order optimality conditions and steepest descent algorithm on orthogonal stiefel manifolds. Optim. Lett., 13(8):1773-1791, November 2019.
[10] Petre Birtea, Ioan Caşu, and Dan Comânescu. Second order optimality on orthogonal stiefel
manifolds. Bulletin des Sciences Mathématiques, 161:102868, 2020. ISSN 0007-4497. doi: http:///...rr//11010111/o..la..a 102868. URLhttps:: /tttpps:/www... science/ /ttaac/e/atttaeeee///////0472220
/ww. http://wwwwscccccccc.c www.ssiiieeeerrct....mm com/
[11] Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
[12] Kwan Ho Ryan Chan, Yaodong Yu, Chong You. Haozhi Qi, John Wright, and Yi Ma. Deep networks from the principle of rate reduction. arXiv preprint arXiv:2010.14765 2020.
[13] Lingling He Changqing Xu and Zerong Lin. Commutation matrices and commutation tensors. Linear and Multilinear Algebra, 1721-17771 2020. doi: URL httt:httt///ooii1 org/ 10. 1000/00000000000 00000008 2018. 1556242
[14] Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsuper- vised feature learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudik (eds.). Proceed- ings of the Fourteenth International Conference on Artificial Intelligence and Statistics, vol- ume 15 of Proceedings of Machine Learning Research, pp. 215-223, Fort Lauderdale, FL, USA, 11-13 Apr 2011. PMLR. URL htRp:ttp:: hatttss:/pproceedingg
mlr press
wwtpp/wssesseeeeeeee.e.p..eeeees vvt///11///////// www.tsaaaallaaaaeee html
[15] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognition, 2009. CVPR 2009. IEEE Conference on, pp. 248-255, 255. IEEE, 2009. URL
ieee org/
abbtract/document/fee
[16] Jiankang Deng, Jia Guo, Niannan Xue, and Stefanos Zafeiriou. Areface: Additive angular margin loss for deep face recognition. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognition, pp. 4690-4699, 2019.
[17] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675-1685. PMIR, 2019.
[18] Alan Edelman, Tomás A Arias, and Steven T Smith. The geometry of algorithms with or- thogonality constraints. SIAM journal on Matrix Analysis and Applications, 20(2):303-333 1998.
[19] Cong Fang, Hangfeng He, Qi Long, and Weijie Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 11843))e2003091118, 2021. doi: 10.10073/p 2103091 URL
https:// pnas. wwwpppan...p.ppp.ppppppp.ppp..pp httaaa/rr///ooab///.ooot 1073 2a130911000911

---

[20] Bin Gao, Xin Liu, Xiaojun Chen, and Ya-xiang Yuan. A new first-order algorithmic framework for optimization problems with orthogonality constraints. SIAM Journal on Optimization, 28 (1):302-332, 2018
[21] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julies Implicit regularization of discrete gradient dynamics in linear neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Gareett (eds.). Advances in Neural Information Processing Sys- tems, volume 32. Curran Associates, Inc., 2019. URL URL / neurips cc/ pdf
[22] Ian J. Goodfellow, Yoshua Yoshua and Aaron Courville. Deep Learning. MIT Press, Cambridge, MA, USA, 2016. http://ww>: http://wwwww.p.plaaaooooooooooo. org
[23] S. Gould, R. Hartley, and D. Campbell. Deep declarative networks. IEEE Trunsactions Puttern Analysis Machine Intelligence, 44(08):3988 4004,, aug 2022. ISSN 1939-3539. doi: 10.1109/TPAMI.2211.1 200444444
[24] Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint 16X.v:600.00544/ 2016.
[25] Stephen Gould, Dylan Campbell, Itzik Ben-Shabal Chamin Hewa Koneputugodage. and Zhiwei Xu. Exploiting problem structure in deep declarative networks: Two case studies. arXiv preprint vrXiv:2202.12404 2022.
[26] Florian Graf, Christoph Hofer, Mare Nicthammer, and Roland Kwitt. Dissecting supervised contrastive learning. In International Conference on Machine Learning, pp. 3821-3830. PMLR, 2021.
[27] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent linear convolutional networks. Advances in neural information processing systems, 1, 2018.
[28] X.Y. Han, Vardan Papyan, and David L. Donoho. Neural collapse under MSE loss: Proximity to and dynamics on the central path. In International Conference on Learning Representations, 2022. http:ttt////peeeee wewpeneeeeiie..... net/ www.ooodduwwwooiorr 1ib.wwHHbbvW__HH.
[29] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In Proceedings of 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 11. pp. 770-778. IEEE, June 2016. doi: 10.1109/CVPR.2016.90 URL http://ieeexploree leee. org/ document/77804455
[30] Elad Hoffer, Itay Hubara, and Daniel Soudry. Fix your classifier: the marginal value of training the last weight layer. arXiv preprint arXiv: 1400000440, 2018.
[31] Meena Jagadeesan, Ilya Razenshteyn, and Suriya Gunasekar. Inductive bias of multi-channel linear convolutional networks with bounded weight norm. In Conference on Learning Theory, pp. 2276-2325. PMLR, 2022.
[32] Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie J Su. An unconstrained layer-peeled perspective on neural collapse. arXiv preprint 10.02:211..02796, 2021.
[33] Ziwei Ji, Miroslav Dudik, Robert E. Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 2109-2136 PMLR, 09-12 Jul 2020. URL https: // proceedings mlr wwwpppesssss///eppeeppp j120a. html
[34] John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas
12

---

Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure predic- tion with alphafold. Nature, 596:583 589, 2021. URI. wwtp::ttps/// /tapps://api... semanticscoolar
org/Corpus :
[35] Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in neural information processing systems, 33: 1661-18677, 2020.
[36] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diega, CA, USA, 2015.
[37] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[38] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521((553):436 2015.
[39] Mario Lezcano Casado. Trivializations for gradient- based optimization on manifolds. Advances in Neural Information Processing Systems, 32, 2019.
[40] Mario Lezcano-Casade and David Martunez-Rubio. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In International Conference on Machine Learning, pp. 3794-3803, PMLR, 2019.
[41] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
[42] Weiyang Liu, Yandong Wen, Zhiding Yu, Ming Li, Bhiksha Raj, and Le Song. Sphereface: Deep hypersphere embedding for face recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 212-220, 2017.
[43] Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Computational Harmonic Analysis, 59:224-241 2022.
[44] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020. URL https
//openreview, bet//formmidd=SJlllBgl
[45] Jan R. Magnus and H. Neudecker. The elimination matrix: Some lemmas and applications.
SIAM Journal Algebraic Discrete Methods, 1(4):422-44)) 1980. doi: 10.1 URL WRRLhttps://// h
[46] Jan R. Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and econometrics / Jan Rudolph Magnus and Heinz Neudecker: Wiley Series in Probability and Statistics. Wiley, Hoboken, NJ, third edition. edition, 2019. ISBN 1-1 1-119-54121-2.
[47] Francesco Mezzadri. How to generate random matrices from the classical compact groups arXiv preprint math-ph/0609050, 2006.
[48] Dustin G. Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. http:::///xxxvv.orp//p
CORR, abs/201 1. 1169, 2020.
11619
[49] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv: 2014.
[50] Vardan Papyan, X. Y. Han, and David L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Science, 117 (40):24652-22) 246633, October 2020. doi: 10.
[51] Gao Peifeng, Qianqian Xu, Peisong Wen, Zhiyong Yang, Huiyang Shao, and Qingming Huang. Feature directions matter: Long-tailed learning via rotated balanced representation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 27542-27563. PMLR, 23-29 Jul 2023. mlr. wwa.rppress/2200 wttps/sssv2222eeeeeee html html
URL htttps://prcceeeee
13

---

[52] Federico Pernici. Matteo Brumi, Claudio Baecchi, and Alberto Del Bimbo. Regular polytope networks. IEEE Transactions on Neural Networks and Learning Systems, (3(9):4373-4381 2021.
[53] Rajeev Ranjan, Carlos D Castillo, and Rama Chellappa. L2-constrained softmax loss for discriminative face verification. arXiv preprint nate::703009555 2017.
[54] Ohad Shamir. Gradient methods never overfit on separable data. Journal of Machine Learning Research, 22(85):1-24 2021.
[55] Jonathan W. Siegel. Accelerated optimization with orthogonality constraints. Journal of Computational Mathematics, 39(2):207-226, 2020. ISSN 1991-7139. doi: hups://doi.org/10 4208/jem. 1tt.:11111m2201-0244400 URL http::/globaa http:///ss.aaatttot/ooaaaaeaaaaa
j jcm/ 11/22.337.....
[56] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for larle image recognition. CoRR, abs/1409.1556, 2014. http://p/p/p...o//////aa/a/
1556
[57] Daniel Soudry, Elad Hoffer, and Nathan Srebro. The implicit bias of gradient descent on separable data. In International Conference on Learning Representations, 2018. URL https
[58] Korawat Tanwisuth, Xinjie Fan, Huangjie Zheng, Shujian Zhang, Hao Zhang, Bo Chen, and Mingyuan Zhou. A prototype-oriented framework for unsupervised domain adaptation. Ad- vances in Neural Information Processing Systems, 34:17194-17208, 2021.
[59] Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. Advances in Neural Information Processing Systems, 11527225-272335 2022.
[60] J. Townsend, N. Koep, and S. Weichwald. PyManopt: a Python toolbox for optimization on manifolds using automatic differentiation. Journal of Machine Learning Research, 17(137):
1-5, 2016. Utthttps://wwa pymanopt.org
[61] Hao Wang, Yitong Wang, Zheng Zhou, Xing Ji, Dihong Gong, Jingchao Zhou, Zhifeng Li, and Wei Liu. Cosface: Large margin cosine loss for deep face recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5265-5274, 2018.
[62] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International conference on machine learning, pp. 9929-9939 PMLR, 2020.
[63] E Weinan and Stephan Wojtowytsch. On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers. Mathematical and Scientific Machine Learning, pp. 270-290. PMLR, 2022.
[64] Zaiwen Wen and Wotao Yin. A feasible method for optimization with orthogonality constraints. Mathematical Programming, 142(1):397-434, 2013.
[65] Nachuan Xiao and Xin Liu. Solving optimization problems over the stiefel manifold by smooth exact penalty function. arXiv preprint arxiv:2/10.088966, 2021.
[66] Nachuan Xiao, Xin Liu, and Kim-Chuan Toh. Dissolving constraints for riemannian optimiza- tion. Mathematics of Operations Research, 49(1):366-397, 2024.
[67] Yibo Yang, Shixiang Chen, Xiangtai Li. Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need learnable classifier at the end of deep neural network? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (cds.). Advances in Neural Information Processing Systems, 2022. URL htt.tttt:t//pereeiiiiiii net/
forum?1d=A6Enx Xc
[68] Yibo Yang, Haobo Yuan, Xiangtai Li, Jianlong Wa, Lefei Zhang, Zhouchen Lin, Philip H.S. Torr, Bernard Ghanem, and Dacheng Tao. Neural collapse terminus: unified solution for class incremental learning and its variants. arXiv pre pre-print, 2023.
14

---

[69] Can Yaras, Peng Wang, Zhihui Zhu, Laura Balzano, and Qing Qu. Neural collapse with normalized features: A geometric analysis over the riemannian manifold. In S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 11547-11560. Curran Associates, Inc., 2022. Utt.https://proceeedngg newreeiinnnnnnnnneneneeeeee. cc/ wwwppppp..... fol.ffllesppaaeee tttee//ppperr22//fDDeee 155cccd1.88//ccc77cccc92244aa233 hatp://eaaaeaeeeeppeeee-ooreeeee Conference. pdf
[70] Chong You, Zhihui Zhu, Qing Qu, and Yi Ma. Robust recovery via implicit bias of discrepant learning rates for double parameterizaaion. Advances in Neural Information Processing Systems, 13:17733-17744, 2020.
[71] Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning diverse and discriminative representations via the principle of maximal coding rate reduction. Advances in Neural Information Processing Systems, 33:9422-9434, 2020.
[72] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. URL. htttpss://ppennreeiexx http:ppawrweeeeeeeee///orrraaaaa
[73] Hongyi Zhang and Suvrit Sra. First-order methods for geodesically convex optimization. In Conference on learning theory, pp. 1617-1638. PMLR, 2016.
[74] Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning, pp. 27 27179-27202. PMLR, 2022.
[75] Jinxin Zhou, Chong You, Xiao Li, Kangning Liu, Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: neural collapse perspective. Advances in Neural Information Processing Systems, 35:31 1697-317770 2022.
[76] Zhihui Zhu, Tianyu DING, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. geometric analysis of neural collapse with unconstrained features. A. Beygelz- imer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL
15

---

Appendices
A Background
Following the definitions the global mean and class mean of the penulimaae--ayer features { in Equation 4
here we introduce the within-class and between-class covariances,
1 N
Ew
A
hhe ho) (he hee)
1=1
(10)
ΣB
- Σ
1 C
(he ha) (he hc)
We proceed by expanding the four key properties of the last-layer activations and classifiers, as empirically observed by Papyan et al. [50] at the terminal phase of training (TPT), where we have achieved zero classification enror and continue towards zero loss.
NC1 Variability Collapse: Throughout training, feature activation variability diminishes they converge towards their respective class means.
NC1 & 1 Ir (ww>ip)) C
(11)
where denotes the Moore-Penross Penrose inverse.
NC2 Convergence to Simplex ETF: The class-mean activation vectors, centred around their global mean, converge to uniform norms while simultaneously maintaining equal-sized and maximally separable angles between then 2
WW
1 VC-1
(
NC2
(12)
[W I
NC3 Convergence to Self-duality: The feature class-means and linear classifiers eventually align in a dual vector space up to some scaling.
WH
1
Ic (1c
NC3
(13)
|WH|
VC-T
NC4 Simplification to Nearest Class-Center (NCC): The network classifier tends to select the class whose mean is closest (in Fuclidean distance) to a given deepnet activation.
NC4
arg max (we, h) + b arg min || h l2
(14)
Since the NC2 and NC3 metrics involve normalised matrices, it is not immediately evident whether the linear classifier and the class-mas -las-mmean activations are equinorm. Consequently, we introduce the following definition:
W H H-Equinorm = [aquinorm Hequinorm!
(15)
where Mequinonm
stde he No (2) and W Vequimomm ave, mella))
std. (Il w. (|2)
avg. h. (allla)
Finally, measure the extent of the variability collapse of each feature separately, define the cosine margin metric as follows:
(w, wG, he, ha)
CMa Ozyig cos a fejigh
where cos Oc.is
(16)
$$c
|wy we 1l2l he, hellz
Note that the maximal angle for simplex ETF vector collection {ve]c defined as such:
(Ve, ve)
forc dl (- for + d
1,
(17)
Tvellel/ver11 ll2
Therefore, we can calculate the theoretical simplex ETF cosine margin as C-1-
"atteemttcaalyy, have defined NC2 the collapse of the linear classifiers simplex ETE.
16

---

B DDN Gradients
The original problem we are trying solve, as defined in Section 3.2 is the following:
minimize Uest
H
UM
(18)
P
or equivalently expanded as such:
u(U) arg min UER4xC
: H VC-1
A UM
U UMM
C-1
(19)
subject to UC-
Here, denote the double-dot operator the Frobenius inner product, i.e, A B TrAAB)
E A,
To compute the first and second-order gradients the objective function and constraints, respectively, we need to consider matrices as variables. Utilising matrix differentials and vectorised derivatives, as defined in [46), simplifies the computation of second-ore objective derivatives and all constraint derivatives. For the objective function, we have:
df(H.U)) = VC-1 dUM
du LU:UM+U:dUN
C-1
2
HM dU + C-1
UM: dU
VC
(20)
2 UM :dU
HM
-C-1
C-1
DU)(H,U) -0-1
2
HM +
2
UM E Bdxc
C-1
2
dDu(H, C-1
dUM
2 C-1 rvec(dUM)
-c- 11 M )d)eeuU (21)
2
drvec Du UH
2 Ia o M) E Rddxxdo
(voc(Dit..(HHHU) C -
where we have defined here the row-major vectorisation method, i.e., rveo(A) = vec(AT), and we have used the property:
rvec (ABc( = (A 8 C. (r)eec(BB).
(22)
A similar proof follows for the second-order partial gradient of the objective. We omit the proof here and just declare the result:
2 B A rvec (uu = VC-1
1a [ e ddCx
(23)
Regarding the gradients of the constraint function, we have:
17

---

dJHHU -c) (1U) dU
drveeJ(HHU)
U)) rvec(U dU)
(Ic U drvec + +UT @ Ic) drvee
(Ic U divec + (U" @ Ic) divecU
(24)
Kco(U" @ Ic) + (U" @ Ic) dr)drve U
(Koc + (U=)(U @ lc)
= (Koc + (U @ Ic) drvecU
(ve U)) = (Kec + (U & Ic) E RCxxC
where we have defined the commutation matrix following identities for any given A E m and B E R2xq:
as K Kmns as a matrix which satisfies the two
K vec (A) vec( A )
(25)
K (A & B K
B & A
The gradient in Equation 24 contains redundant constraints because of the symmetrical nature of the orthogonality constraints. To retain only the non-redundant constraints, we must undertake half-vectorisation procedure. Given that we already possess the fully vectorised gradients (which are
simpler to compute in this scenario), require elimination matrix, Le E IR e(C+1) x such that
Le rvec(A ()) rvveh(A).
(26)
We can define an elimination matrix explicitly as follows 45
Ln = Σ
uec Eij) = Σ (ui 8 ej BeT),
E(us
(27)
22
$29
where
[1 in position (j - 1) + 100 -
Ui 0 elsewhere
(28)
Hence,
A (rech(DuJJHHi Lo(Kce Ic))(U & I e R
ciC+
xdC
(29)
We have the following useful Corollary for the second-order gradient of the constraint function.
Corollary (Magnus & Neudecker [46). Let o be a twice differentiable real-valued function of an q matrix X. Then, the following two relationships hold between the second differential and the Hessian matrix of o at X:
& d(x = Tr(A(X) ([XXX + Hb(x)= 1 (A -
With the above corollary established, have,
((, = -UUU)-UU--
dJHHU -dUV T
(30)
&J du dU = 2(UU)
(HH = 2e, (dU) = TT(e,e, (dU),
18

---

and hence, we get:
rvec (H,))( = & (ezee)) La (e,e,) E RaCxdc
(31)
Then, we repeat the process with the elimination matrix to eliminate the redundant constraints. However, while this is one way to compute the derivative, a more efficient approach exists, namely the embedded gradient vector field method, which we outline in the following subsection. For complete overview the method, we recommend readers to follow through works of Birtea et al.
[0), Birtea & Comânescu [8] 71
Finally, the gradients for the proximal problem Equation 7 are straightforward compute, with the only change in gradients being the added proximal terms in:
2 HM + 8(U -Upox) Bdx
Duf(H.U)
UM
K K-1
-K-1
(32)
2 (re K Ia (M + &JJa RdCxdc
B.1 Implicit Formulation of Lagrange Multipliers Differentiable Manifolds
An issue arises in Proposition D with the expression for G,
G rrvce(Diu f(H U)) - A rvech Div (HHU)) U)) RdCxxd
(33)
where both the calculation of the Lagrange multiplier matrix A (solved via linear system) and the construction of the fourth-order tensor representing the second-order derivatives of the constraint function are complex and challenging. However, by recognising the manifold structure of the problem, we can reformulate Equation 33 in simpler and more computationally efficient way. The embedded gradient vector field method offers such solution [8]
For general Riemannian manifold (M,g), g), we can define the Gram matrix for the smooth functions fi...., fo, fi..... hr (M,g) + R as follows,
(hh, Vfi
(Vhr, Vfi)
Gram (f (hammn)
Rexr
(34)
(h1 V
(V,, V
In our problem, we are working with compact Stiefel manifold, which is an embedded sub- manifold RdxC and we identify the isomorphism (via vec) between RaxC' and RRC Stiefel manifold Std. (U E RdxC UU- Ic] can be characterised by set of constraint functions, Js: RdCC + R, as follows:
UU=
1 jo(u)
2 |ua
1S8SC,
(35)
(p) = (up: ug):
ISpegsC.
We consider smooth cost function f Sig. R and define f: BdxC R smooth extension (or prolongation) of f. The embedded gradient vector field (after applying the isomorphism) is defined on the open set [R C gd formed with the regular leaves of the constrained function
j
.dC
R
((C+))
and is tangent to the foliation of this function. The vector field takes the
form 9
0f(u)= = vf(u)
Σ oo(u)Vis(u)- - Σ
(pp(z)))pp(())
(36)
1555C
ISP
where Je Op are the Lagrange multiplier functions defined such [7]
19

---

det (Gram (h (f
det (Gram, (u))
0z(u)=
det Gram Gram 1ra Ut....Jc.j12...... 1.)) (1)
det (Gram(u))
(37)
det (Gram (h
1.0)
(1)
det (Grampo (u)
Ang (u)
(1.. JC......... det
(ram (f1- Jc- (u)
det (Gram(u))
The Gram matrix the denominator of the Lagrange multiplier functions can be defined as a block matrix as follows,
A CT
Gram(u)=
(38)
C
B
where
(Vi, Vńn
(Vio, Víi)
=
(4i. Vie)
Vic. Vic)
(Vi12. Vi12)
(Vic-ic,
B
(V/12. Vjc-1.c
(Vic-1.o, Vic-1c)
(VJ1, Viie)
(Vic, Viv2
c=
(Vi Vjc-1.) (Vio, Vic-1o)
For Gram,(u) and Grampo (u), appropriate column of the Gram(u) (u) based on the index. For instance, to define Gram,(u), we replace the column with (Vf(w) Vj. (
similarly define block decompositions by replacing the
It has been proved [9] that if U E SIg is critical point of the function j. ie., 0f(u) = 0, then (.(w), pg (w) become the classical Lagrange multipliers.
Proposition 2 (Lagrange multiplier functions for Stiefel Manifolds). The Lagrange multiplier functions are described as function of the constraint functions in Equation 35 in the following way:
0f Tu,
(u) = (V)()), Vis(u))
(u), U2
(39)
1
0f
05 Örip
pp (V)( (u)) 2
(u), Uhp
(u), un
Oeg
Proof. We take the definition of the Lagrange multiplier functions in Equation 37 denominator, we observe that the Gram matrix is of the form:
Starting with the
Ic
Gram(u)
0
21
(40)
since for each block matrix (A, B, C') of the Gramian, we get for each of their elements:
VAgri (V).(u), (u)( dar
VBBr, (V), Vine(u)) 28 + 28
20

---

VCe (Vi. (u)) = 25 = 0
where 66 is defined as the Kronecker delta symbol, i.e., 6ig [i j]. Thus, the determinant of this Gram matrix is:
det det(Graame (Gram(u) u)) = det
Ic
0
21 2Ic
(et((c)) det (2I
0:0-1)
(41)
Now, let us turn our focus to the numerator of the Lagrange multiplier function os (u).
(VS(u), (i)
(Vf(u), Via(u))
det(Gram,(uu) = det
(Vf(u), Vio(w))
(Vf(u), V)12(u))
0
2
(42)
(VS(u), Vic-1.c(u))
Ic-1
(VS(u). Vi.( det
21 2
2 2(G-") (Vf(u), Vjs(u)).
Similarly, for (u), we get:
det (Gram = (V)(u), (u)) det Ic
21 2
(43)
= C(C-11
f(u). (u))
Finally, get the result by combining the determinants' results for each Lagrange multiplier function and computing the gradients of the constraint functions.
Similarly to the gradient vector field, we can show that the Hessian for constraint manifold (e.g., Stiefel manifold, (us)f(u) StS St. + R) can be given such [0
Hess Hessf(u) = Hessf(u)- Σ
o,(u) (e))j((j(()))
15:50
(44)
Σ
(psss))))sjp((1s))
ISPEESC
[T Stu XTuSt
We can transform the results into matrix form the following way. First, we denote,
((( A vec (VVf(u)) Rdxx; OF(U)= & vec ((f(u)) E _xx
(45)
We then introduce symmetric matrix E(U) & [u)] E RCXC, where also include the Lagrange multiplier function's symmetrical component. For the Stiefel manifold, the matrix form of the embedded gradient vector field (after some simple computation) is given by,
0f(U) = -
(46)
21

---

where -(U)= = (Vf(o) -VV)(
Regarding the Hessian in Equation 44 write it in matrix form, first need compute the Hessian matrices the constraint functions as follows
0a
0g
0₂
(U))
0g
La
0g
Ia 2 (e, @ e,,):
0a
0d
0a
0,
(47)
0a
0.
Ia
0,
[Hessjpq(U)
la @ (e, Beg + +eg 1e000 0e,)
Ia
0a
0.
0g
0g 0a
0,8
Finally, the matrix form of the Hessian of the cost function f. ::t - IR is given by,
Hes)f(U) (Hess Ia 8 E(U) (( ststx
(48)
where from there, we can re-define the expression G in Equation 33 as follows,
G = rvec(Diu (H, Ia 0 E(U) E RCdC
(49)
C Additional Experimental Results
055
Food
-10
-02
66
100
200
10000
25300
30300
6300 5000
Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PaM EoT
Soed
Standard
mokcit 277
2.0
200
(g) Test Top-1 Acc.
100
61
(d) NC1
(8) NC3
(1) W H Equinorm
Figure 6: CIFARIO results on ResNet-18. all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
Herr the resulting Kronecker product expressed row-major way.
22

---

0.
land 11p
Sead IT!
120
297
10900
10000 30000 4000 5000
Train Top-1 Acc.
(b) Avg Cos. Margin
(c) BoT
144
12
Foad impiikee
co Imp
cavied
04
(g) Test Top-1 Acc.
(d) NC1
(e) NCC
(1) W H Equinorm
Figure 7: CIFARIO results on VGG-13. In all plots, the x-axis represents the number epochs, except for plot (c), where the x-axis denotes the number of training examples.
00
05
Sund ET
120
207
1090
1530
1530
620E 5966
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PoM BoT
Toad
daaderd
carded
roiie
30
2.-
044
05
(g) Test Top- Acc.
6.3
68
62
61
60
(d) NC1
(c) NC3
(f) W H Equinorm
Figure 8: 100 results on ResNet-50. all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
23

---

1.00
0.2
-036
022
-106
120
293
1900
2000
3000
4000
5000
Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PCM BoT
Foad impiikeee
12
cavied
0.5
(g) Test Top-1 Acc.
02
32
(d) NC1
(e) NC
(1) W H Equinorm
Figure 9: CIFARI00 results on VGG-13. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
18
03
Sunt ET
-10
320
207
1990
2003
2002
1880
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PeM EoT
fead
daader
candedd
molice
05
3.0
(g) Test Top-1 Acc.
67
63
(d) NC1
(c) NC3
(f) W H Equinorm
Figure 10: STL10 results on ResNet-50. plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
24

---

021
TRODeSCHEC Hage
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PCM at EoT
sosced road ETI
Sred FTC impiit.22
candeo Impl.11
(g) Test Top-1 Acc.
(d) NC1
(c) NC3
(1) W Equinorm
Figure 11: STLIO results VGG-13. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
1.
ImplotE
CE lower bound
0.8
TPT Phose Standard
0.8
TPT Phose Foodd
Phose Implet
0.7
Implitt
fiwd
0.
Theoretical ETF Margin
0.6
Throretical Margin
TPT Phuee Standard
0.5
TPT Phose Fied
0.0
TPT Phase implct
0.4
800
1200
1800
2000
1000
2000
3000
4000
5000
(a) CE loss
(b) Avg Cos. Margin
(c) PCM at EoT
impliii ETF
Tieed
1.2
Ried ETF imppcckEE ETF
standard
0.000
Stnndard mppiitt
TPT Phaze Sto Sarddrd
Phase- Find
Implict
0.5
TPT Phane Implct
TPT Phase-Fion Floed
TPT Phase: Standard
0.035
Implict
Shandard
0.6
0.4
0.7
0.0
0.000
(d) NC1
(e) NC3
W H-Equinoor
Figure 12: UFM- 100 results. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x x-axis denotes the number of training examples.
25

---

standard
1.0
Ried ETP
10
implest ETF
0.9
CE
TPT Phose implic
Phose- Floed
0.6
Food ETF
Standard
0.7
standard
Inglit
foed
implist
0.6
0.2
Theoretical
Phate Standard
0.0
Phase
Phase
3000
CF loss
(h) Avg Cos. Margin
(c) PaM FoT
1.4
impidi
1.2
0.0200
Fioed CTT
standbra
Phate Slandard
0.1175
Imskiit ETT
TPT Phase Implist
Phase- Read
1.0
0.8
Phase
Phase Implck
0.0150
IPI Mhase Stantard
Phase Find
Phase
0.6
0.4
0.0255
000
1200
1600
2000
0.0
1200
1000
0.0000
0.02
1200
1002
(d) NC1
(e) NC3
(f) W H EEquioom
Figure 13: UFM-200 results. all plots, x-axis represents the number of epochs, except for plot (c). where the x- x-axis denotes the number of training examples.
1.0
implct ETF
CE lawer bound
0.
0.9
TPT Phose Standard
Phose Feed
Phose
0.4
sandard
0.6
RTd
standard
Thooresical ETF Morgin
Implct
02
Theoretical ETF Margin
0.6
TPT Phose Standard
TPT Phase implet
IPI Mase Fied
0.5
1200
1600
2000
4000
6000
8000
30000
(a) CE loss
(b) Avg Cos. Margin
(c) PcM at EoT
standard
standard
250
find
1.2
Implit ETT
Ried EP
200
150
Phose Foeed
Phane Sandand
Phose Impldt
1.0
TFT Phase Reed
TPT Phase Impick
TFT Phose Standard
---
Instict
TPT Puase Implett
Prase hasd
Phese Standard
0.6
100
0.4
0.0075
0.0000
0.0025
0.0
400
800
1200
1000
2000
800
1200
3600
2000
(d) NC1
(0) N3
H Eqqinorm
Figure 14: UFM-1000 results. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
26

---

1.0
20
22
PolicE TherntoiCT Sure
1EPICEETP No.00
Imslit ITT.
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) Pom at EoT
trplot 61 Ded
implict inplic
Irplenttr Ded
PICKA ETP
25
(g) Test Top-1 Acc.
23
(d) NC1
(c) NC3
(1) W Equinorm
Figure 15: CIFARIO results VGG-13, comparing the implicit ETF method in scenarios: one where the DDN gradient is computed and included in the SGD update, and another where the DDN gradient computation omitted from the update. In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
35
flxed ETF
Implisit
101
impilst
25
median implistt ETF
00 time mmt
20
10
4x100
egochs
(a) Forward pass times in milliseconds.
(b) Forward and backward times in (log) millisecs.
Figure 16: CIFAR10 computational cost results on ResNet-18. In (a), we plot the forward pass time for each method. For the implicit ETF method, which has dynamic computation times, we also include the mean and median time values. In (b), we plot the computational cost for each forward and backward pass across methods. For the implicit ETF forward pass, we have taken its median time. The notation as follows: S/F Standard Forward Pass, S/B = Standard Backward Pass, F/F Fixed ETF Forward Pass, F/B Fixed ETF Backward Pass, I/F Implicit ETF Forward Pass, and I/B Implicit ETF Backward Pass.
27

---

05
Stedad 13
PERCEET
-10
inpict ETT
10900
1000
3900 4000 5000
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PoM EoT
10
Wendord
incilor
conduttiio COW
15
(g) Test Top-1 Acc.
6:38
120
397
203
(d) NC1
(e) NC3
(I) W H Equinorm
Figure 17: CIFARIO results on ResNet-18, where in standard no-norm we do not perform feature and weight normalisation. In all plots, the x-axis represents the number of epochs, except for plot (c). where the x-axis denotes the number of training examples.
05
Dedon
impicc E77
-02
20300 38200 48906 5000
(a) Train Top-1 Acc.
(b) Avg Cos. Margin
(c) PoM at EoT
cenance
2.
Emet nerdon CTT Inctiunakka
65.
food sncon
Fand canoncaal ETT
2.
(g) Test Top-1 Acc.
110
6.1
120
250
A30
80
2.51
292
402
(d) NC1
(c) NC3
(1) W H Equinorm
Figure 18: CIFARIO results on ResNet-18, where in fixed random ETF we choose a random orthogonal direction instead of the canonical In all plots, the x-axis represents the number of epochs, except for plot (c), where the x-axis denotes the number of training examples.
28

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We introduced a novel method leveraging the simplex ETF structure and the neural collapse phenomenon to enhance training convergence and stability in neural networks, and we provided experimental results supporting our claims.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification: We discussed the limitations of our approach in Section of the paper.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
The paper does not propose new theoretical results. We proposed new method and its mathematical formulation as well as the appropriate mathematical background.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All the necessary details required for reproducing the results are presented in the paper. Also, the code will be made publicly available upon acceptance.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Due to anonymity reasons, the code is not included the submission. However, it will be made available GitHub repository upon acceptance.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All the specific implementation details can be found in Section of the paper and the appendix.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have tested our results using different random seeds, and present the median values along with the range.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
29

---

Answer: [Yes]
Justification: We have specified the GPUs that were used for our experiments.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wwrrg...oooo/oo//ccccaaa
Answer: [Yes]
Justification: The paper conforms with the NeurIPS code of ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is societal impact of the work performed in this paper.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not pose such risks.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, every model architecture and dataset used were cited.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
30


Reviews:
[{"review_id": "HwkVpLMWoC", "reviewer": "Reviewer_xmab", "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.", "questions": "See weaknesses section.", "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.\nThe authors claimed that they empirically observed no significant difference in performance for small-scale problems where the DDN gradient can be computed.\nIncluding these results in the supplementary material would also be beneficial.\nMoreover, we agree that future work should verify if this is true for large-scale problems.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality:\nThe 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.\n\nQuality:\nThe paper is well-written, and the proposed approach is carefully supported by theorems and experiments.\n\nClarity:\nThe paper is well-written and easy to follow.\n\nSignificance:\nTheir 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:\nThe 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.\n\nStandard Procedure:\n\"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.\n\nImage Baselines Results are not SOTA:\nThe 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.\n\nFixed ETF Procedure:\nThe 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).\n\nRemarks:\nThe authors should directly clarify in Tables 1 and 2 the ResNet architecture used (18 or 50).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Overhead Cost: Please refer to the tables in our general response and the discussion around computational concerns.\n\nStandard 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.\n\nImage 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.\n\nFixed 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).\n\nRemark: We have updated our table results to reflect the chosen ResNet and VGG architectures.\n\nLimitations:\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6xNVWbkaAL", "reviewer": "Reviewer_m1UJ", "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. \n\nThey 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. \n\nOverall 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. \n\nFrom 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.", "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?\n\nDo you do any, or is it possible to perform a comparison of these two resulting ETFs?\nHow does the test accuracy of your 'encouraged' ETF compare to the one you would have obtained naturally?\n\nIn 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? \n\nOn 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.\nIt 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.\n\nIn 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?\n\nIn 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.\n\n(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. \n\nAre there error bars in Figure 2? I see them for plot (f) but not for the others?\n(clarification) What is depicted in Figure 2(c)? What is \\mathcal{P}_{CM}? I think I somehow missed that.\n\nAre there error bars in Figure 3? Were multiple trials run for these experiments?\n\nTables 1 and 2: The ranges for train and test top-1 accuracy values for STL on VGG seem very large. \n\nIn 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.", "limitations": "N/A The authors address any limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThere 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.\n\nThe 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.\n\nGood work.", "weaknesses": "The main points of concern for me are in regards to the experiments and how the results are reported in the paper.\n\nTable 2 looks good but is a bit misleading particularly when comparing the ranges of the test top-1 accuracy.\nThe 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.\n\nThe 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.\n\nFor Figure 4, also no error bars. Understanding compute restraints, it would be nice to see similar multiple seed runs for ImageNet experiments. \n\nFinally, 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. \nThe authors do mention this in the limitations section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses:\n\nMisleading 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.\n\nNon-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.\n\nFigure 4 no error bars: Fixed and replaced with 5 seed runs (found in the attached pdf).\n\nCost: Please see tables and discussion in the general response. \n\nQuestions:\n\nOptimal 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.\n\nCompare 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.\n\nProximal 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.\n\n$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.\n\n$\\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).\n\n$\\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.\n\nClarification 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.\n\nError 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.\n\nSTL 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3AlJHSCT6i", "reviewer": "Reviewer_oHwR", "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.", "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.\n\n2. 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.\n\n3. 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?", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weakness 1 with question 1 & 2:\n\nEq. 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.\n\nProposition 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).\n\nAs 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.\n\nWeakness 2: Please refer to our general response for computational cost discussions.\n\nWeakness 3: \n\nAblation 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.\n\nQuestion 3:\n\nFixed 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.\nIn 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ANjQlYoi5d", "reviewer": "Reviewer_isxr", "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:\n\n1. The standard training strategy where the final layer weights are updated by backpropagation.\n\n2. The final layer weights are fixed as a simplex ETF (which has been well-studied in previous works).\n\nThe 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.", "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.\n\n2. 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?\n\n3. What is the convergence behavior when employing SGD/Adam instead of the AGD approach?\n\nnit: Where is $U_{init}$ defined?\n\nnit: line 117, below eq (5), is the formulation of $\\widetilde{H}$ correct? shouldn't the denominator be $||\\overline{H}||_F$ ?\n\nReferences\n\n[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\n\n[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).", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n\n2. 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:\n- By what percentage does the step time and memory increase when adding this layer?\n- When should one avoid the backward pass through this layer and consider only the forward pass?\n- 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.\n\nsee more questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses: \n\nRegarding the computational and memory costs of our approach, please refer to our general response. \n\nAvoid 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.\n\nSpecifically, 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.\n\nQuestions:\n\n1. 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. \nIt’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.\n\n2. 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. \nAnother 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.\n\n3. 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.\n\nNit: $U_{init}$ is defined in the paragraph for Hyperparameter Selection and Riemannian Initialisation Schemes in lines 230-235 in our paper.\n\nNit: 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$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their time and thoughtful feedback. Here we address the common question regarding the computational costs of our method, and we will address individual comments for each reviewer separately.


Please refer to Table 1 for computational cost and Table 2 for memory cost measurements. The key insight is that incorporating DDN into the backward pass (in our Python implementation) becomes prohibitively expensive when significantly increasing the number of penultimate layer features and classes. However, as demonstrated in the ImageNet case, where GPU memory constraints prevent us from running this step, omitting the DDN backward pass (similar to a gradient stop operation) still yields strong results as we are still directing the solution to the nearest ETF. Note that there is still an indirect gradient path via the classifier, as shown in Figure 1 of our paper. Despite the strong results, we acknowledge the computational challenges associated with the DDN backward pass. The primary bottleneck in GPU memory usage is during the calculation of the constrained derivatives (Eq. 29), where explicit formulation of certain intermediate matrices occurs. Concerning the time complexity of the DDN backward pass, we use a direct solver to compute the expression in Proposition 1, which can be inefficient for large-scale linear systems. We plan to address these limitations (both time and memory) by using more efficient linear algebra solvers in future work.


                               Table 1: Average time (in milliseconds) of a step update during training

| Model                | Standard Fwd | Standard Bwd | Fixed ETF Fwd | Fixed ETF Bwd | Implicit ETF Fwd | Implicit ETF Bwd (w/o DDN) | Implicit ETF Bwd (w/ DDN) |
|----------------------|--------------|--------------|---------------|---------------|------------------|----------------------------|---------------------------|
| UFM-10               | 0.2          | 0.4          | 0.2           | 0.4           | 7.6              | 1.2                        | 1.8                       |
| UFM-100              | 0.2          | 1.0          | 0.4           | 0.8           | 8.5              | 1.1                        | 608.9                     |
| UFM-200              | 0.3          | 1.1          | 0.3           | 0.9           | 8.7              | 1.2                        | 17,909.9                  |
| UFM-1000             | 0.3          | 1.1          | 0.5           | 1.1           | 23.1             | 0.8                        | N/A                       |
| CIFAR10 / Res18      | 3.4          | 4.3          | 3.7           | 4.5           | 12.8             | 5.1                        | 5.9                       |
| CIFAR10 / Res50      | 7.0          | 9.2          | 6.9           | 7.9           | 19.4             | 9.2                        | 15.0                      |
| CIFAR100 / Res50     | 10.5         | 6.5          | 7.6           | 8.1           | 40.2             | 11.3                       | 1,872.4                   |
| ImageNet / Res50     | 7.4          | 11.4         | 7.3           | 8.1           | 160.2            | 10.0                       | N/A                       |



                               Table 2: GPU memory (in Gigabytes) during training

| Model                | Standard | Fixed ETF | Implicit ETF w/o DDN Bwd | Implicit ETF w/ DDN Bwd |
|----------------------|:---------:|:---------:|:-------------------------:|:------------------------:|
| UFM-10               | 1.5       | 1.5       | 1.5                       | 1.6                      |
| UFM-100              | 1.7       | 1.7       | 1.7                       | 10.7                     |
| UFM-200              | 1.7       | 1.7       | 1.8                       | 70.3                     |
| UFM-1000             | 1.9       | 1.9       | 2.9                       | N/A                      |
| CIFAR10 / Res18      | 2.2       | 2.2       | 2.2                       | 2.3                      |
| CIFAR10 / Res50      | 2.6       | 2.6       | 2.7                       | 2.8                      |
| CIFAR100 / Res50     | 2.6       | 2.6       | 2.7                       | 18.9                     |
| ImageNet / Res50     | 27.5      | 27.2      | 27.8                      | N/A                      |

The attached PDF includes figures of additional results requested by reviewers, which we discuss in the individual responses below.

Meta-review: The paper exploits the neural collapse phenomenon to improve training convergence and stability. Specifically, by leveraging the duality between penultimate layer features and the final layer weights, along with the simplex ETF structure of the final layer weights, the proposed approach gradually guides the weights towards an optimal simplex ETF during training. This approach differs from existing work that uses a fixed simplex ETF during training, which may require the feature mapping to do the heavy lifting to align features with the classifier. Experimental results on various neural network architectures and real-world datasets demonstrate the improvement of the proposed approach in both training and testing performance.

Overall, all the reviewers found the contribution of this work to be solid and a valuable addition to the exploration of neural collapse for improving training.
Decision: Accept (poster)