Extended Unconstrained Features Model for Exploring Deep Neural Collapse

Abstract

The modern strategy for training deep neural networks for classification tasks includes optimizing the network's weights even after the training error vanishes to further push the training loss toward zero. Recently, a phenomenon termed "neural collapse" (NC) has been empirically observed in this training procedure. Specifically, it has been shown that the learned features (the output of the penultimate layer) of within-class samples converge to their mean, and the means of different classes exhibit a certain tight frame structure, which is also aligned with the last layer's weights. Recent papers have shown that minimizers with this structure emerge when optimizing a simplified "unconstrained features model" (UFM) with a regularized cross-entropy loss. In this paper, we further analyze and extend the UFM. First, we study the UFM for the regularized MSE loss, and show that the minimizers' features can be more structured than in the cross-entropy case. This affects also the structure of the weights. Then, we extend the UFM by adding another layer of weights as well as ReLU nonlinearity to the model and generalize our previous results. Finally, we empirically demonstrate the usefulness of our nonlinear extended UFM in modeling the NC phenomenon that occurs with practical networks.

On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features

“…The underlying reasoning is that modern deep networks are often highly overparameterized with the capacity of learning any representations [19][20][21][22], so that the last-layer features can approximate, or interpolate, any point in the feature space. Under the unconstrained feature model, the work [14][15][16][23][24][25][26] showed that the N C solutions are the only global optimal solution for nonconvex training losses under different settings. However, given the nonconvexity of the problem, even under the unconstrained feature model these global optimality results do guarantee that the N C solutions can be efficiently achieved.…”

“…In the meanwhile, the MSE loss is not only appealing for its algebraic simplicity, but it also demonstrates on-par or even better generalization performances compared to the CE loss, as reported by recent line of work [10]. However, the theoretical study of MSE loss for N C is still limited [12,15,26]. Under the unconstrained feature model, their work proved that the continuous gradient flow of the MSE loss converges to N C solutions.…”

“…The work [28][29][30][31] study deep homogeneous classification networks (without bias terms but beyond the unconstrained features model) trained with MSE loss, stochastic gradient descent, and weight decay. In particular, the solutions satisfying the so-called symmetric quasi-interpolation assumption are proved to obey N C properties, but the properties of other solutions are not investigated [30,31] As far as we know, the work closest to ours is the concurrent work [26]. Under similar unconstrained feature models, the work studies the global optimality condition of N C for the MSE loss for both two-layer and three layer networks, but not the global optimization landscape.…”

Inducing Neural Collapse in Imbalanced Learning: Do We Really Need a Learnable Classifier at the End of Deep Neural Network?

“…They prove that neural collapse emerges under the CE loss with proper constraint or regularization. Other studies prove that the MSE loss also leads to the neural collapse solution [6,12,13,14]. [31] proposes a convex formulation for a norm-regularized ReLU network and explains neural collapse accordingly.…”

“…It has been shown that the optimality for LPM under the CE loss leads to the neural collapse solution with feature norm constraints [8,5,7,9], regularization [10], or even no explicit constraint [11]. Other studies turn to analyze the unconstrained LPM under the mean squared error (MSE) loss and also derive the neural collapse solution [6,12,13,14].…”