# Neural Collapse with Normalized Features:
A Geometric Analysis over the Riemannian Manifold

Can Yaras111The first two authors contributed to this work equally. Department
of Electrical Engineering & Computer Science, University of Michigan  Peng
Wang11footnotemark: 1 Department of Electrical Engineering & Computer Science,
University of Michigan  Zhihui Zhu  Department of Computer Science &
Engineering, Ohio State University  Laura Balzano  Department of Electrical
Engineering & Computer Science, University of Michigan  Qing Qu  Department of
Electrical Engineering & Computer Science, University of Michigan

###### Abstract

When training overparameterized deep networks for classification tasks, it has
been widely observed that the learned features exhibit a so-called “neural
collapse” phenomenon. More specifically, for the output features of the
penultimate layer, for each class the within-class features converge to their
means, and the means of different classes exhibit a certain tight frame
structure, which is also aligned with the last layer’s classifier. As feature
normalization in the last layer becomes a common practice in modern
representation learning, in this work we theoretically justify the neural
collapse phenomenon for normalized features. Based on an unconstrained feature
model, we simplify the empirical loss function in a multi-class classification
task into a nonconvex optimization problem over the Riemannian manifold by
constraining all features and classifiers over the sphere. In this context, we
analyze the nonconvex landscape of the Riemannian optimization problem over
the product of spheres, showing a benign global landscape in the sense that
the only global minimizers are the neural collapse solutions while all other
critical points are strict saddles with negative curvature. Experimental
results on practical deep networks corroborate our theory and demonstrate that
better representations can be learned faster via feature normalization.

##  1 Introduction

Depsite the tremendous success of deep learning in engineering and scientific
applications over the past decades, the underlying mechanism of deep neural
networks (DNNs) still largely remains mysterious. Towards the goal of
understanding the learned deep representations, a recent seminal line of works
[50, 23, 16, 87, 86] presents an intriguing phenomenon that persists across a
range of canonical classification problems during the terminal phase of
training. Specifically, it has been widely observed that last-layer features
(i.e., the output of the penultimate layer) and last-layer linear classifiers
of a trained DNN exhibit simple but elegant mathematical structures, in the
sense that

  * •

(NC1) Variability Collapse: the individual features of each class concentrate
to their class-means.

  * •

(NC2) Convergence to Simplex ETF: the class-means have the same length and are
maximally distant. In other words, they form a Simplex Equiangular Tight Frame
(ETF).

  * •

(NC3) Convergence to Self-Duality: the last-layer linear classifiers perfectly
match their class- means.

Such a phenomenon is referred to as _Neural Collapse_ ($\mathcal{NC}$) [50],
which has been shown empirically to persist across a broad range of canonical
classification problems, on different loss functions (e.g., cross-entropy (CE)
[50, 87, 16], mean-squared error (MSE) [71, 86], and supervised contrasive
(SC) losses [21]), on different neural network architectures (e.g., VGG [63],
ResNet [24], and DenseNet [28]), and on a variety of standard datasets (such
as MNIST [39], CIFAR [35], and ImageNet [12], etc). Recently, in independent
lines of research, many works are devoted to learning maximally compact and
separated features; see, e.g., [76, 43, 60, 72, 42, 73, 13, 51, 52]. This has
also been widely demonstrated in a number of recent works [9, 46, 14, 48, 61,
18, 25], including state-of-the-art natural language models (such as BERT,
RoBERTa, and GPT) [46].

Figure 1: Comparison of features found with and without normalization. $K=100$
classes with $n=5$ samples per class. Features and classifiers are found
through optimizing the cross-entropy loss with a UFM, where features are
embedded in 2-dimensional space, i.e., $d=2$. (a) No normalization of features
or classifiers. (b) Features and classifiers are constrained to the unit
sphere (features are scaled down for visualization purposes).

| Average CE Loss | Average Accuracy
---|---|---
No Normalization | $1.63\pm 0.03$ | $49.9\%\pm 2.39\%$
Normalization | $3.84\pm 0.00$ | $100.0\%\pm 0.00\%$

Table 1: Average UFM feature loss and accuracy over 10 trials with and without
normalization to sphere, with the same set-up as in Figure 1.

#### Motivations & contributions.

In this work, we further demystify why $\mathcal{NC}$ happens in network
training with a common practice of feature normalization (i.e., normalizing
the last-layer features on the unit hypersphere), mainly motivated by the
following reasons:

  * •

_Feature normalization is a common practice in training deep networks._
Recently, many existing results demonstrated that training with feature
normalization often improves the quality of learned representation with better
class separation [60, 42, 73, 13, 5, 81, 75, 21]. Such a representation is
closely related to the discriminative representation in literature; see, e.g.,
[60, 42, 73, 80]. As illustrated in Figure 1 and Table 1, experimental results
visualized in low-dimensional space show that features learned with
normalization are more uniformly distributed over the sphere and hence are
more linearly separable than those learned without normalization. In
particular, it has been shown that the learned representations with larger
class separation usually lead to improved test performances; see, e.g., [32,
21]. Moreover, it has been demonstrated that discriminative representations
can also improve robustness to mislabeled data [5, 81], and has become a
common practice in recent advances on (self-supervised) pretrained models [8,
75].

  * •

_A common practice of theoretically studying $\mathcal{NC}$ with norm
constraints._ Due to these practical reasons, many existing theoretical
studies on $\mathcal{NC}$ consider formulations with both the norms of
features and classifiers constrained [44, 77, 16, 21, 30]. Based upon
assumptions of unconstrained feature models [47, 87, 16], these works show
that the only global solutions satisfy $\mathcal{NC}$ properties for a variety
of loss functions (e.g., MSE, CE, SC losses, etc). Nonetheless, they only
focused on the global optimality conditions without looking into the nonconvex
landscapes, and therefore failed to explain why these $\mathcal{NC}$ solutions
can be efficiently reached by classical training algorithms such as stochastic
gradient descent (SGD).

In this work we study the global nonconvex landscape of training deep networks
with norm constraints on the features and classifiers. We consider the
commonly used CE loss and formulate the problem as a Riemannian optimization
problem over products of unit spheres (i.e., the oblique manifold). Our study
is also based upon the assumption of the so-called _unconstrained feature
model_ (UFM) [47, 87, 86] or _layer-peeled model_ [17], where the last-layer
features of the deep network are treated as free optimization variables to
simplify the nonlinear interactions across layers. The underlying reasoning is
that modern deep networks are often highly overparameterized with the capacity
of learning any representations [45, 26, 62], so that the last-layer features
can approximate, or interpolate, any point in the feature space.

Assuming the UFM, we show that the Riemannian optimization problem has a
benign global landscape, in the sense that the loss with respect to (w.r.t.)
the features and classifers is a strict saddle function [20, 65] over the
Riemannian manifold. More specifically, we prove that _every_ local minimizer
is a global solution satisfying the $\mathcal{NC}$ properties, and _all_ the
other critical points exhibit directions with negative curvature. Our analysis
for the manifold setting is based upon a nontrivial extension of recent
studies for the $\mathcal{NC}$ with penalized formulations [87, 86, 71, 23],
which could be of independent interest. Our work brings new tools from
Riemannian optimization for analyzing optimization landscapes of training deep
networks with an increasingly common practice of feature normalization. At the
same time, we empirically demonstrate the advantages of the
Riemannian/constrained formulation over its penalized counterpart for training
deep networks – faster training and higher quality representations. Lastly,
under the UFM we believe that the benign landscape over the manifold could
hold for many other popular training losses beyond CE, such as the
(supervised) contrastive loss [32]. We leave this for future exploration.

#### Prior arts and related works on $\mathcal{NC}$.

The empirical $\mathcal{NC}$ phenomenon has inspired a recent line of
theoretical studies on understanding why it occurs [21, 17, 44, 47, 71, 86,
87]. Like ours, most of these works studied the problem under the UFM. In
particular, despite the nonconvexity, recent works showed that the _only_
global solutions are $\mathcal{NC}$ solutions for a variety of nonconvex
training losses (e.g., CE [87, 17, 44], MSE [71, 86], SC losses [21]) and
different problem formulations (e.g., penalized, constrained, and
unconstrained) [87, 86, 71, 23, 44]. Recently, this study has been extended to
deeper models with the MSE training loss [71]. More surprisingly, it has been
further shown that the nonconvex losses under the UFM have benign global
optimization landscapes, in the sense that every local minimizer satisfies
$\mathcal{NC}$ properties and the remaining critical points are strict saddles
with negative curvature. Such results have been established for both CE and
MSE losses [87, 86], where they considered the unconstrained formulations with
regularization on both features and classifiers. We should also mention that
the benign global optimization landscapes of many other problems in neural
networks have been widely found in the literature; see, e.g., [83, 38, 68, 41,
64].

Moreover, there is a line of recent works investigating the benefits of
$\mathcal{NC}$ on generalizations of deep networks. The work [19] shows that
$\mathcal{NC}$ also happens on test data drawn from the same distribution
asymptotically, but less collapse for finite samples [29]. Other works [29,
49] demonstrated that the variability collapse of features is actually
happening progressively from shallow to deep layers, and [2] showed that test
performance can be improved when enforcing variability collapse on features of
intermediate layers. The works [78, 79] showed that fixing the classifier as a
simplex ETF improves test performance on imbalanced training data and long-
tailed classification problems. We refer interested readers to a recent survey
on this emerging topic [34].

#### Notation.

Let $\mathbb{R}^{n}$ be the $n$-dimensional Euclidean space and
$\|\cdot\|_{2}$ be the Euclidean norm. We write matrices in bold capital
letters such as $\mathbf{A}$, vectors in bold lower-case such as $\mathbf{a}$,
and scalars in plain letters such as $a$. Given a matrix
$\mathbf{A}\in\mathbb{R}^{d\times K}$, we denote its $k$-th column by
$\mathbf{a}_{k}$, its $i$-th row by $\mathbf{a}^{i}$, its $(i,j)$-th element
by $a_{ij}$, and let $\|\mathbf{A}\|$ be its spectral norm. We use
$\operatorname{diag}(\mathbf{A})$ to denote a vector that consists of diagonal
elements of $\mathbf{A}$, and we use $\operatorname{ddiag}(\mathbf{A})$ to
denote a diagonal matrix composed by only the diagonal entries of
$\mathbf{A}$. Given a positive integer $n$, we denote the set
$\\{1,\dots,n\\}$ by $[n]$. We denote the unit hypersphere in $\mathbb{R}^{d}$
by $\mathbb{S}^{d-1}:=\\{\mathbf{x}\in\mathbb{R}^{d}:\|\mathbf{x}\|_{2}=1\\}$.

##  2 Nonconvex Formulation with Spherical Constraints

In this section, we review the basic concepts of deep neural networks and
introduce notation that will be used throughout the paper. Based upon this, we
formally introduce the problem formulation over the Riemannian manifold under
the assumption of the UFM.

###  2.1 Basics of Deep Neural Networks

In this work, we focus on the multi-class (e.g., $K$ class) classification
problem. Given input data $\mathbf{x}\in\mathbb{R}^{D}$, the goal of deep
learning is to learn a deep hierarchical representation (or feature)
$\mathbf{h}(\mathbf{x})=\phi_{\mathbf{\theta}}(\mathbf{x})\in\mathbb{R}^{d}$
of the input along with a linear classifier222We write
$\mathbf{W}=\mathbf{W}_{L}^{\top}$ in the transposed form for the simplicity
of analysis. $\mathbf{W}\in\mathbb{R}^{d\times K}$ such that the output
$\psi_{\mathbf{\Theta}}(\mathbf{x})=\mathbf{W}^{\top}\mathbf{h}(\mathbf{x})$
of the network fits the input $\mathbf{x}$ to an one-hot training label
$\mathbf{y}\in\mathbb{R}^{K}$. More precisely, in vanilla form an $L$-layer
fully connected deep neural network can be written as

$\displaystyle\psi_{\mathbf{\Theta}}(\mathbf{x})\;=\;\underbrace{\mathbf{W}_{L}}_{\text{\bf
linear
classifier}\;\mathbf{W}=\mathbf{W}_{L}^{\top}}\;\;\underbrace{\sigma\left(\mathbf{W}_{L-1}\cdots\sigma\left(\mathbf{W}_{1}\mathbf{x}+\mathbf{b}_{1}\right)+\mathbf{b}_{L-1}\right)}_{\text{
\bf
feature}\;\;\mathbf{h}\;=\;\phi_{\mathbf{\theta}}(\mathbf{x})}\;+\;\mathbf{b}_{L},$
(1)

where each layer is composed of an affine transformation, represented by some
weight matrix $\mathbf{W}_{k}$, and bias $\mathbf{b}_{k}$, followed by a
nonlinear activation $\sigma(\cdot)$, and
$\mathbf{\Theta}=\left\\{\mathbf{W}_{k},\mathbf{b}_{k}\right\\}_{k=1}^{L}$ and
$\mathbf{\theta}=\left\\{\mathbf{W}_{k},\mathbf{b}_{k}\right\\}_{k=1}^{L-1}$
denote the weights for _all_ the network parameters and those up to the last
layer, respectively. Given training samples
$\\{(\mathbf{x}_{k,i},\mathbf{y}_{k})\\}\subset\mathbb{R}^{D}\times\mathbb{R}^{K}$
drawn from the same data distribution $\mathcal{D}$, we learn the network
parameters $\mathbf{\Theta}$ via minimizing the empirical risk over these
samples,

$\displaystyle\min_{\mathbf{\Theta}}\;\sum_{k=1}^{K}\sum_{i=1}^{n_{k}}\mathcal{L}_{\mathrm{CE}}\left(\psi_{\mathbf{\Theta}}(\mathbf{x}_{k,i}),\bm{y}_{k}\right),\quad\text{s.t.}\quad\mathbf{\Theta}\in\mathcal{C},$
(2)

where $\bm{y}_{k}\in\mathbb{R}^{K}$ is a one-hot vector with only the $k$th
entry being $1$ and the remaining ones being $0$ for all $k\in[K]$,
$\mathbf{x}_{k,i}\in\mathbb{R}^{D}$ is the $i$-th sample in the $k$-th class,
$\left\\{n_{k}\right\\}_{k=1}^{K}$ denotes the number of training samples in
each class, and the set $\mathcal{C}$ denotes the constraint set of the
network parameters $\mathbf{\Theta}$ that we will specify later. In this work,
we study the most widely used CE loss of the form

$\displaystyle\mathcal{L}_{\mathrm{CE}}(\mathbf{z},\mathbf{y}_{k})\;:=\;-\log\left(\frac{\exp(z_{k})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right).$

###  2.2 Riemannian Optimization over the Product of Spheres

For the $K$-class classification problem, let us consider a simple case where
the number of training samples in each class is balanced (i.e.,
$n=n_{1}=n_{2}=\cdots=n_{K}$) and let $N=Kn$, and we assume that all the
biases $\left\\{\mathbf{b}_{k}\right\\}_{k=1}^{L}$ are zero with the last
activation function $\sigma(\cdot)$ before the output to be linear. Analyzing
deep networks $\psi_{\mathbf{\Theta}}(\mathbf{x})$ is a tremendously difficult
task mainly due to the _nonlinear interactions_ across a large number of
layers. To simplify the analysis, we assume the so-called _unconstrained
feature model_ (UFM) following the previous works [30, 47, 21, 87]. More
specifically, we simplify the nonlinear interactions across layers by treating
the last-layer features
$\mathbf{h}_{k,i}=\phi_{\mathbf{\theta}}(\mathbf{x}_{k,i})\in\mathbb{R}^{d}$
as _free_ optimization variables, where the underlying reasoning is that
modern deep networks are often highly overparameterized to approximate any
continuous function [45, 26, 62]. Concisely, we rewrite all the features in a
matrix form as

$\displaystyle\mathbf{H}\;=\;\begin{bmatrix}\mathbf{H}_{1}&\mathbf{H}_{2}&\cdots&\mathbf{H}_{K}\end{bmatrix}\in\mathbb{R}^{d\times
N},\;\mathbf{H}_{k}\;=\;\begin{bmatrix}\mathbf{h}_{k,1}&\mathbf{h}_{k,2}&\cdots&\mathbf{h}_{k,n}\end{bmatrix}\in\mathbb{R}^{d\times
n},\;\forall\ k\in[K],$

and correspondingly denote the classifier $\mathbf{W}$ by

$\displaystyle\mathbf{W}\;=\;\begin{bmatrix}\mathbf{w}_{1}&\mathbf{w}_{2}&\cdots&\mathbf{w}_{K}\end{bmatrix}\in\mathbb{R}^{d\times
K},\quad\mathbf{w}_{k}\in\mathbb{R}^{d},\quad\forall\ k\in[K].$

Based upon the discussion in Section 1, we assume that both the features
$\mathbf{H}$ and the classifiers $\mathbf{W}$ are normalized,333In practice,
it is a common practice to normalize the output feature $\mathbf{h}_{o}$ by
its norm, i.e., $\mathbf{h}=\mathbf{h}_{o}/\left\|\mathbf{h}_{o}\right\|_{2}$,
so that $\left\|\mathbf{h}\right\|_{2}=1$. in the sense that
$\|\mathbf{h}_{k,i}\|_{2}=1$, $\|\mathbf{w}_{k}\|_{2}=\tau$, for all $k\in[K]$
and all $i\in[n]$, where $\tau>0$ is a temperature parameter. As a result, we
obtain a _constrained_ formulation over a Riemannian manifold

$\displaystyle\min_{\mathbf{W},\mathbf{H}}\;\frac{1}{N}\sum_{k=1}^{K}\sum_{i=1}^{n}\mathcal{L}_{\mathrm{CE}}\left(\mathbf{W}^{\top}\mathbf{h}_{k,i},\mathbf{y}_{k}\right)\quad\text{s.t.}\;\|\mathbf{w}_{k}\|_{2}\;=\;\tau,\;\|\mathbf{h}_{k,i}\|_{2}\;=\;1,\
\forall\ i\in[n],\;\forall\ k\in[K].$ (3)

Since the temperature parameter $\tau$ can be absorbed into the loss function,
we replace $\mathbf{w}_{k}$ by $\tau\mathbf{w}_{k}$ and change the original
constraint into $\|\mathbf{w}_{k}\|_{2}=1$ for all $k\in[K]$. In particular,
the product of spherical constraints forms an _oblique manifold_ [3] embedded
in Euclidean space,

$\displaystyle\mathcal{OB}(d,K)\;:=\;\left\\{\mathbf{Z}\in\mathbb{R}^{d\times
K}\mid\mathbf{z}_{k}\in\mathbb{S}^{d-1},\;\forall\ k\in[K]\right\\}.$

Consequently, we can rewrite Problem (3) as a Riemannian optimization problem
over the oblique manifold w.r.t. $\mathbf{W}$ and $\mathbf{H}$:

$\displaystyle\min_{\mathbf{W},\mathbf{H}}$
$\displaystyle\;f(\mathbf{W},\mathbf{H})\;:=\;\frac{1}{N}\sum_{k=1}^{K}\sum_{i=1}^{n}\mathcal{L}_{\mathrm{CE}}\left(\tau\mathbf{W}^{\top}\mathbf{h}_{k,i},\mathbf{y}_{k}\right),$
(4) s.t.
$\displaystyle\quad\mathbf{H}\in\mathcal{OB}(d,N),\;\mathbf{W}\in\;\mathcal{OB}(d,K).$

In Section 3, we will show that all global solutions of Problem (4) satisfy
$\mathcal{NC}$ properties, and its objective function is a strict saddle
function [31, 15] of $(\mathbf{W},\mathbf{H})$ over the oblique manifold so
that the $\mathcal{NC}$ solution can be efficiently achieved.

Figure 2: An illustration of the Riemannian gradient of $f(\mathbf{h})$ on a
simple manifold $\mathcal{OB}(3,1)$.

#### Riemannian derivatives over the oblique manifold.

In Section 3, we will use tools from Riemannian optimization to characterize
the global optimality condition and the geometric properties of the
optimization landscape of Problem (4). Before that, let us first briefly
introduce some basic derivations of the Riemannian gradient and Hessian,
defined on the tangent space of the oblique manifold. For more technical
details, we refer the readers to Section A.1. According to [4, Chapter 3 & 5]
and [27, 1], we can calculate the Riemannian gradients and Hessian of Problem
(4) as follows. Since those quantities are defined on the tangent space,
according to [27, Section 3.1] and the illustration in Figure 2, we first
obtain the tangent space to $\mathcal{OB}(d,K)$ at $\mathbf{W}$ as

$\displaystyle\mathrm{T}_{\mathbf{W}}\mathcal{OB}(d,K)\;=\;\left\\{\mathbf{Z}\in\mathbb{R}^{d\times
K}\mid\operatorname{diag}\left(\mathbf{W}^{\top}\mathbf{Z}\right)=\mathbf{0}\right\\}.$

Note that the tangent space contains all $\mathbf{Z}$ such that
$\mathbf{z}_{k}$ is orthogonal to $\mathbf{w}_{k}$ for all $k$. When $K=1$, it
reduces to the tangent space to the unit sphere $\mathbb{S}^{d-1}$.

Analogously, we can derive the tangent space for $\mathbf{H}$ with a similar
form. Let us define

$\displaystyle\mathbf{M}:=\tau\mathbf{W}^{\top}\mathbf{H},\
g(\mathbf{M}):=f(\mathbf{W},\mathbf{H}).$

First, the Riemannian gradient of $f(\mathbf{W},\mathbf{H})$ of Problem (4) is
basically the projection of the ordinary Euclidean gradient $\nabla
f(\mathbf{W},\mathbf{H})$ onto its tangent space, i.e.,
$\operatorname{grad}_{\mathbf{W}}f(\mathbf{W},\mathbf{H})=\mathcal{P}_{\mathrm{T}_{\mathbf{W}}\mathcal{OB}(d,K)}(\nabla_{\mathbf{W}}f(\mathbf{W},\mathbf{H}))$
and
$\operatorname{grad}_{\mathbf{H}}f(\mathbf{W},\mathbf{H})=\mathcal{P}_{\mathrm{T}_{\mathbf{H}}\mathcal{OB}(d,N)}(\nabla_{\mathbf{H}}f(\mathbf{W},\mathbf{H}))$.
More specifically, we have

$\displaystyle\operatorname{grad}_{\mathbf{W}}f(\mathbf{W},\mathbf{H})\;$
$\displaystyle=\;\tau\mathbf{H}\nabla
g(\mathbf{M})^{\top}-\tau\mathbf{W}\operatorname{ddiag}\left(\mathbf{W}^{\top}\mathbf{H}\nabla
g(\mathbf{W})^{\top}\right),$ (5)
$\displaystyle\operatorname{grad}_{\mathbf{H}}f(\mathbf{W},\mathbf{H})\;$
$\displaystyle=\;\tau\mathbf{W}\nabla
g(\mathbf{M})-\tau\mathbf{H}\operatorname{ddiag}\left(\mathbf{H}^{\top}\mathbf{W}\nabla
g(\mathbf{M})\right).$ (6)

Second, for any
$\mathbf{\Delta}=(\mathbf{\Delta}_{\mathbf{W}},\mathbf{\Delta}_{\mathbf{H}})\in\mathbb{R}^{d\times
K}\times\mathbb{R}^{d\times N}$, we compute the Hessian bilinear form of
$f(\mathbf{W},\mathbf{H})$ along the direction $\mathbf{\Delta}$ by

$\displaystyle\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;=\;$
$\displaystyle\nabla^{2}g(\mathbf{M})\left[\tau\left(\bm{W}^{\top}\bm{\Delta}_{\bm{H}}+\bm{\Delta}_{\bm{W}}^{\top}\bm{H}\right),\tau\left(\bm{W}^{\top}\bm{\Delta}_{\bm{H}}+\bm{\Delta}_{\bm{W}}^{\top}\bm{H}\right)\right]$
$\displaystyle+2\tau\left\langle\nabla
g(\mathbf{M}),\bm{\Delta}_{\bm{W}}^{\top}\bm{\Delta}_{\bm{H}}\right\rangle.$
(7)

We compute the Riemannian Hessian bilinear form of $f(\mathbf{W},\mathbf{H})$
along any direction
$\mathbf{\Delta}\in\mathrm{T}_{\mathbf{W}}\mathcal{OB}(d,K)\times\mathrm{T}_{\mathbf{H}}\mathcal{OB}(d,N)$
by

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;=\;$
$\displaystyle\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]-\langle\mathbf{\Delta}_{\mathbf{W}}\operatorname{ddiag}\left(\mathbf{M}\nabla
g(\mathbf{M})^{\top}\right),\mathbf{\Delta}_{\mathbf{W}}\rangle$
$\displaystyle-\langle\mathbf{\Delta}_{\mathbf{H}}\operatorname{ddiag}\left(\mathbf{M}^{\top}\nabla
g(\mathbf{M})\right),\mathbf{\Delta}_{\mathbf{H}}\rangle,$ (8)

where the extra terms besides
$\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]$
represent the curvatures induced by the oblique manifold. We refer to Section
A.2 for the derivations of (2.2) and (2.2). In the following section, we will
use the Riemannian gradient and Hessian to characterize the optimization
landscape of Problem (4).

##  3 Main Theoretical Analysis

In this section, we first characterize the structure of the global solution
set of Problem (4). Based upon this, we analyze the global landscape of
Problem (4) via characterizing its Riemannian derivatives.

###  3.1 Global Optimality Condition

For the feature matrix $\mathbf{H}$, let us denote the class mean for each
class by

$\displaystyle\overline{\mathbf{h}}_{k}\;:=\;\frac{1}{n}\sum_{i=1}^{n}\mathbf{h}_{k,i},\
\forall\
k\in[K],\quad\text{and}\quad\overline{\mathbf{H}}\;:=\;\begin{bmatrix}\overline{\mathbf{h}}_{1}&\cdots&\overline{\mathbf{h}}_{K}\end{bmatrix}\in\mathbb{R}^{d\times
K}.$ (9)

Based upon this, we show any global solution of Problem (4) exhibits
$\mathcal{NC}$ properties in the sense that it satisfies (NC1) variability
collapse, (NC2) convergence to simplex ETF, and (NC3) convergence to self-
duality.

######  Theorem 3.1 (Global Optimality Condition)

Suppose that the feature dimension is no smaller than the number of classes,
i.e., $d\geq K$, and the training labels are balanced in each class, i.e.,
$n=n_{1}=\cdots=n_{K}$. Then for the CE loss $f(\mathbf{W},\mathbf{H})$ in
Problem (4), it holds that

$\displaystyle f(\mathbf{W},\mathbf{H})\,$
$\displaystyle\;\geq\;\log\left(1+(K-1)\exp\left(-\frac{K\tau}{K-1}\right)\right)$

for all $\mathbf{W}=[\mathbf{w}_{1},\dots,\mathbf{w}_{K}]\in\mathcal{OB}(d,K)$
and
$\mathbf{H}=[\mathbf{h}_{1,1},\dots,\mathbf{h}_{K,n}]\in\mathcal{OB}(d,N)$. In
particular, equality holds if and only if

  * •

(NC1) Variability collapse: $\mathbf{h}_{k,i}=\overline{\mathbf{h}}_{k},\
\forall\ i\in[n]$;

  * •

(NC2) Convergence to simplex ETF: $\\{\overline{\mathbf{h}}_{k}\\}_{k=1}^{K}$
form a sphere-inscribed simplex ETF in the sense that

$\displaystyle\overline{\mathbf{H}}^{\top}\overline{\mathbf{H}}\;=\;\frac{1}{K-1}\left(K\mathbf{I}_{K}-\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right),\quad\overline{\mathbf{H}}\in\mathcal{OB}(d,K).$

  * •

(NC3) Convergence to self-duality: $\mathbf{w}_{k}=\overline{\mathbf{h}}_{k},\
\forall\ k\in[K]$.

Compared to the unconstrained regularized problems in [86, 87], it is worth
noting that the regularization parameters influence the structure of global
solutions, while the temperature parameter $\tau$ only affects the
optimization landscape but not the global solutions. On the other hand, our
result is closely related to [21, Theorem 1] (i.e., spherical constraints vs.
that of ball constraints). In fact, our problem and that in [21] share the
same global solution set. Moreover, our proof follows similar ideas as those
in a line of recent works [21, 17, 44, 71, 86, 87], and we refer the readers
to Appendix B for the proof. It should be noted that we do not claim
originality of this result compared to previous works; instead our major
contribution lies in the following global landscape analysis.

Figure 3: Global optimization of (4) under UFM with $d=100$ and $n=5$.
Theoretical line refers to lower bound (global minimum) from Theorem 3.1.
Empirical values found using gradient descent with random initialization.
Left: Lower bound against number of classes $K$ while fixing $\tau=1$. Right:
Lower bound against temperature $\tau$ while fixing $K=10$. The same empirical
values are achieved over many trials due to the benign global landscape.

###  3.2 Global Landscape Analysis

Due to the nonconvex nature of Problem (4), the characterization of global
optimality in Theorem 3.1 alone is not sufficient for guaranteeing efficient
optimization to those desired global solutions. Thus, we further study the
global landscape of Problem (4) by characterizing all the _Riemannian critical
points_ $(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,K)\times\mathcal{OB}(d,N)$
satisfying

$\displaystyle\operatorname{grad}_{\mathbf{H}}f(\mathbf{W},\mathbf{H})\;=\;\mathbf{0},\quad\operatorname{grad}_{\mathbf{W}}f(\mathbf{W},\mathbf{H})\;=\;\mathbf{0}.$

We now state our major result below.

######  Theorem 3.2 (Global Landscape Analysis)

Assume that the number of training samples in each class is balanced, i.e.,
$n=n_{1}=\cdots=n_{K}$. If the feature dimension is larger than the number of
classes, i.e., $d>K$, and the temperature parameter satisfies
$\tau<2(d-2)(1+(K\mbox{{ mod }}2)/K)^{-1}$, then the function
$f(\mathbf{W},\mathbf{H})$ is a strict saddle function that has no spurious
local minimum, in the sense that

  * •

Any Riemannian critical point $(\mathbf{W},\mathbf{H})$ of Problem (4) that is
not a local minimizer is a Riemannian strict saddle point with negative
curvatures, in the sense that the Riemannian Hessian
$\mathrm{Hess}f(\mathbf{W},\mathbf{H})$ at the critical point
$(\mathbf{W},\mathbf{H})$ is non-degenerate, and there exists a direction
$\mathbf{\Delta}=(\mathbf{\Delta}_{\mathbf{W}},\mathbf{\Delta}_{\mathbf{H}})\in\mathrm{T}_{\mathbf{W}}\mathcal{OB}(d,K)\times\mathrm{T}_{\mathbf{H}}\mathcal{OB}(d,N)$
such that

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;<\;0.$

In other words,
$\lambda_{\mathrm{min}}\left(\mathrm{Hess}f(\mathbf{W},\mathbf{H})\right)<0$
at the corresponding Riemannian critical point.

  * •

Any local minimizer of Problem (4) is a global minimizer of the form shown in
Theorem 3.1.

For the details of the proof, we refer readers to Appendix C. The second
bullet point naturally follows from Theorem 3.1 and the first bullet point.
The major challenge of our analysis is showing the first bullet, i.e., to find
a negative curvature direction $\mathbf{\Delta}$ for
$\operatorname{Hess}f(\mathbf{W},\mathbf{H})$. Our key observation is that the
set of non-global critical points can be partitioned into two separate cases.
In the first case, the last two terms of (2.2) vanish, and we show that the
second term of (2.2) is negative and dominates the first term for an
appropriate direction. We require $\tau$ to not be too large, since the first
term is $O(\tau^{2})$, whereas the second term is $O(\tau)$. In the second
case, using the assumption that $d>K$ we can find a rank-one direction that
makes the first term of (2.2) vanishing. In this case, we similarly show that
the second term of (2.2) is negative but instead dominates the last two terms
of (2.2). In the following, we discuss the implications, relationship, and
limitations of our results in Theorem 3.2.

  * •

_Efficient global optimization to $\mathcal{NC}$ solutions._ Our theorem
implies that the $\mathcal{NC}$ solutions can be efficiently reached by
Riemannian first-order methods (e.g., Riemannian stochastic gradient descent)
with random initialization [31, 11, 69] for solving Problem (4); see Figure 3
for a demonstration. For training practical deep networks, this can be
efficiently implemented by normalizing last-layer features when running SGD.

  * •

_Relation to existing works on $\mathcal{NC}$._ Most existing results have
only studied the global minimizers under the UFM [21, 17, 44, 71], which has
limited implication for optimization. On the other hand, our landscape
analysis is based upon a nontrivial extension of that with the unconstrained
problem formulation [87, 86]. Compared to those works, Problem (4) is much
more challenging for analysis, due to the fact that the set of critical points
of our problem is essentially much larger than that of [87, 86]. Moreover, we
empirically demonstrate the advantages of the manifold formulation over its
regularized counterpart, in terms of representation quality and training
speed.

  * •

_Assumptions on the feature dimension $d$ and temperature parameter $\tau$._
Our current result requires that $d>K$, which is the same requirement in [87,
86]. Furthermore, through numerical simulations we conjecture that the global
landscape also holds even when $d\ll K$, while the global solutions are
uniform over the sphere [44] rather than being simplex ETFs (see Figure 1).
The analysis on $d\ll K$ is left for future work. On the other hand, the
required upper bound on $\tau$ is for the ease of analysis and it holds
generally in practice,444For instance, a standard ResNet-18 [24] model trained
on CIFAR-10 [35] has $d=512$ and $K=10$. In the same setting, we assume
$\tau<1020$, which is far larger than any useful setting of the temperature
parameter (see Section 4.5). but we conjecture that the benign landscape holds
without it.

  * •

_Relation to other Riemannian nonconvex problems._ Our result joins a recent
line of work on the study of global nonconvex landscapes over Riemannian
manifolds, such as orthogonal tensor decomposition [20], dictionary learning
[56, 66, 67, 57, 58], subspace clustering [74], and sparse blind deconvolution
[36, 55, 37, 84]. For all these problems constrained over a Riemannian
manifold, it can be shown that they exhibit “equivalently good” global
minimizers due to symmetries and intrinsic low-dimensional structures, and the
loss functions are usually strict saddles [20, 65, 85]. As we can see, the
global minimizers (i.e., simplex ETFs) of our problem here also exhibit a
similar rotational symmetry, in the sense that
$\mathbf{W}^{\top}\mathbf{H}=(\mathbf{Q}\mathbf{W})^{\top}\left(\mathbf{Q}\mathbf{H}\right)$
for any orthogonal matrix $\mathbf{Q}$. Additionally, our result show that the
tools from Riemmanian optimization can be powerful for the study of deep
learning.

##  4 Experiments

In this section, we support our theoretical results in previous sections and
provide further motivation with experimental results on practical deep network
training. In Section 4.1, we validate the assumption of UFM introduced in
Section 2 for analyzing $\mathcal{NC}$, by demonstrating that $\mathcal{NC}$
occurs for increasingly overparameterized deep networks. In Section 4.2, we
further motivate feature normalization with empirical results showing that
feature normalization can lead to faster training and better collapse than the
unconstrained counterpart with regularization. This occurs not only with the
UFM but also with practical overparameterized networks. In Section 4.3, we
demonstrate that, independent of the algorithm used, the feature normalized
UFM has faster training and collapse than its regularized counterpart. In
Section 4.4, we show that feature normalization leads to better generalization
and test feature collapse on practical deep networks. In Section 4.5, we
investigate the effect of the temperature parameter $\tau$ on training
dynamics for both the UFM and practical deep networks. Finally, in Section
4.6, we empirically explore the global landscape of other commonly used loss
functions for deep learning classification tasks. Before that, we introduce
some basics of the experimental setup and metrics for evaluating
$\mathcal{NC}$.

#### Network architectures, datasets, and training details.

In our experiments, we use ResNet [24] architectures for the feature encoder.
For the normalized network, we project the output of the encoder onto the
sphere of radius $\tau$ (as done in [21]) and also project the weight
classifiers to the unit sphere after each optimization step to maintain
constraints. In all experiments, we set $\tau=1$. For the regularized UFM and
network, we use a weight and feature decay of $10^{-4}$ (using the loss in
[87]). We do not use a bias term for the classifier for either architecture.
For all experiments, we use the CIFAR dataset555Both CIFAR10 and CIFAR100 are
publicly available and are licensed under the MIT license. [35], where we use
CIFAR100 for all experiments except for the experiment in Section 4.1, where
we use CIFAR10. In all experiments, we train the networks using SGD with a
batch size of $128$ and momentum $0.9$ with an initial learning rate of
$0.05$, and we decay the learning rate by a factor of $0.1$ after every $40$
epochs - these hyperparameters are chosen to be the same as those in [87] for
fair comparisons. All networks are trained on Nvidia Tesla V100 GPUs with 16G
of memory.

#### Neural collapse metrics.

For measuring different aspects of neural collapse as introduced in Section 1,
we adopt similar $\mathcal{NC}$ metrics from [50, 87, 86], given by

$\displaystyle\mathcal{NC}_{1}$
$\displaystyle:=\frac{1}{K}\mbox{trace}(\mathbf{\Sigma}_{W}\mathbf{\Sigma}_{B}^{\dagger})$
$\displaystyle\mathcal{NC}_{2}$
$\displaystyle:=\left\|\frac{\mathbf{W}^{\top}\mathbf{W}}{\|\mathbf{W}^{\top}\mathbf{W}\|_{F}}-\frac{1}{\sqrt{K-1}}\left(\mathbf{I}_{K}-\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right)\right\|_{F}$
$\displaystyle\mathcal{NC}_{3}$
$\displaystyle:=\left\|\frac{\mathbf{W}^{\top}\overline{\mathbf{H}}}{\|\mathbf{W}^{\top}\overline{\mathbf{H}}\|_{F}}-\frac{1}{\sqrt{K-1}}\left(\mathbf{I}_{K}-\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right)\right\|_{F},$

where $\mathbf{\Sigma}_{W}$ and $\mathbf{\Sigma}_{B}$ are the within-class and
between-class covariance matrices (see [50, 87] for more details),
$\mathbf{\Sigma}_{B}^{\dagger}$ denotes pseudo inverse of
$\mathbf{\Sigma}_{B}$, and $\overline{\mathbf{H}}$ is the centered class mean
matrix in (9). More specifically, $\mathcal{NC}_{1}$ measures NC1 (i.e.,
within class variability collapse), $\mathcal{NC}_{2}$ measures NC2 (i.e., the
convergence to the simplex ETF), and $\mathcal{NC}_{3}$ measures NC3 (i.e.,
the duality collapse).

###  4.1 Validation of the UFM for training networks with feature
normalization

In Section 2, our study of the Riemannian optimization problem (4) is based
upon the UFM, where we assume $\mathbf{H}$ is a free optimization variable.
Here, we justify this assumption by showing that $\mathcal{NC}$ happens for
training overparameterized networks even when the training labels are
completely random. By using random labels, we disassociate the input from
their class labels, by which we can characterize the approximation power of
the features of overparmeterized models. To show this, we train ResNet-18 with
varying widths (i.e., the number of feature maps resulting from the first
convolutional layer) on CIFAR10 with random labels, with normalized features
and classifiers.

As shown in Figure 4, we observe that increasing the width of the network
allows for perfect classification on the training data even when the labels
are random. Furthermore, increasing the network width also leads to better
$\mathcal{NC}$, measured by the decrease in each $\mathcal{NC}$ metric. This
corroborates that (_i_) our assumption of UFM is reasonable given that
$\mathcal{NC}$ seems to be independent of the input data, and (_ii_)
$\mathcal{NC}$ happens under the our constraint formulation (4) on practical
networks.

Figure 4: Validation of UFM on ResNet with varying network width.
$\mathcal{NC}$ metrics and training accuracy of ResNet-18 networks of various
widths on CIFAR10 with $n=200$ over 200 epochs.

###  4.2 Feature normalization for improved training speed and representation
quality

We now investigate the benefits of using feature normalization, more
specifically for improved training speed and better representation quality.
First, we consider the UFM formulation, where we optimize Problem (4) and
compare to the regularized UFM in [87]. The results are shown in Figure 5. We
can see that normalizing features over the sphere consistently results in
reaching perfect classification and greater feature collapse (i.e., smaller
$\mathcal{NC}_{1}$) quicker than penalizing the features. To demonstrate that
these behaviors are reflected in training practical deep networks, we train
both ResNet-18 and ResNet-50 architectures on a reduced CIFAR100 [35] dataset
with $N=3000$ total samples, comparing the training accuracy and metrics of
$\mathcal{NC}$ with and without feature normalization. The results are shown
in Figure 6.

Figure 5: Faster training/feature collapse of UFM with feature normalization.
Average (deviation denoted by shaded area) training accuracy and
$\mathcal{NC}_{1}$ of UFM over 10 trials of (Riemmanian) gradient descent with
exact line search. We set $K=100$ classes, $n=30$ samples per class. Figure 6:
Faster training/feature collapse with ResNet on CIFAR100 with feature
normalization. Training accuracy and $\mathcal{NC}_{1}$ of ResNet-18 and
ResNet-50 on CIFAR100 with $n=30$ over 100 epochs.

From Figure 6 (left), we can see that for the ResNet-18 network, we reach
perfect classification of the training data about 10-20 epochs sooner by using
feature normalization compared to that of the unconstrained formulation. From
Figure 6 (right), training the ResNet-50 network without feature normalization
for 100 epochs shows slow convergence with poor training accuracy, whereas
using feature normalization arrives at above 90% training accuracy in the same
number of epochs. By keeping the size of the dataset the same and increasing
the number of parameters, it is reasonable that optimizing the ResNet-50
network is more challenging due to the higher degree of overparameterization,
yet this effect is mitigated by using feature normalization.

At the same time, for both architectures, using feature normalization leads to
greater feature collapse (i.e., smaller $\mathcal{NC}_{1}$) compared to that
of the unconstrained counterpart. As shown in recent work [50, 86, 19] ,
better $\mathcal{NC}$ often leads to better generalization performance, as
corroborated by Section 4.4. Last but not least, we believe the benefits of
feature normalization are not limited to the evidence that we showed here, as
it could also lead to better robustness [5, 81] that is worth further
exploration.

###  4.3 Neural collapse occurs independently of training algorithms under UFM

In Section 4.2, we demonstrated that optimizing Problem (4), which corresponds
to the feature normalized UFM, results in quicker training and feature
collapse as opposed to the regularized UFM formulation, as shown in Figure 5.
To show that this phenomenon is independent of the algorithm used, we
additionally test the conjugate gradient (CG) method [1] as well as the trust-
region method (TRM) [1] to solve (4) with the same set-up as in Figure 5.
While the Riemannian conjugate gradient method is also a first order method
like gradient descent, the Riemannian trust-region method is a second order
method, so the convergence speed is much faster compared to Riemannian
gradient descent or conjugate gradient method. The results are shown in Figure
7 and Figure 8 for the CG method and TRM respectively.

We see that optimizing the feature normalized UFM with CG gives similar
results to using GD, whereas optimizing the feature normalized UFM using TRM
results in an even greater gap in convergence speed to the global solutions,
when compared with optimizing the regularized counterpart using TRM. These
results suggest that the benefits of feature normalization are not limited to
vanilla gradient descent or even first order methods.

Figure 7: Faster training/feature collapse of UFM with feature normalization
with CG. Average (deviation denoted by shaded area) training accuracy and
$\mathcal{NC}_{1}$ of UFM over 10 trials of (Riemmanian) conjugate gradient
method. We set $K=100$ classes, $n=30$ samples per class. Figure 8: Faster
training/feature collapse of UFM with feature normalization with TRM. Average
(deviation denoted by shaded area) training accuracy and $\mathcal{NC}_{1}$ of
UFM over 10 trials of (Riemmanian) trust-region method. We set $K=100$
classes, $n=30$ samples per class.

###  4.4 Feature normalization generalizes better than regularization

In Section 4.2, we showed that using feature normalization over regularization
improves training speed and feature collapse when training increasingly
overparameterized ResNet models on a small subset of CIFAR100. We now
demonstrate that feature normalization leads to better generalization than
regularization.

We train a ResNet-18 and ResNet-50 model on the _entirety_ of the CIFAR100
training split without any data augmentation for 100 epochs, and test the
accuracy and $\mathcal{NC}_{1}$ metric on the standard test split. The results
are shown in Table 2. We immediately see that using feature normalization
gives both better test accuracy and test feature collapse than using
regularization. Furthermore, the test generalization performance is coupled
with the degree of feature collapse, supporting the claim that better
$\mathcal{NC}$ often leads to better generalization performance. Finally, as
we have trained both ResNet architectures with the same set-up and number of
epochs, there is a substantial drop in performance (both test accuracy and
test $\mathcal{NC}_{1}$) of the regularized ResNet-50 model compared to the
ResNet-18. However, using feature normalization, this effect is mostly
mitigated, suggesting that feature normalization is more robust compared to
regularization and effective for generalizing highly overparameterized models
on fixed-size datasets.

| ResNet-18 | ResNet-50
---|---|---
| Test Accuracy | Test $\mathcal{NC}_{1}$ | Test Accuracy | Test $\mathcal{NC}_{1}$
Regularization | 55.3% | 3.838 | 48.9% | 4.486
Normalization | 58.6% | 3.143 | 56.4% | 3.127

Table 2: Better generalization and test feature collapse with ResNet on
CIFAR100 with feature normalization. Test accuracy and test $\mathcal{NC}_{1}$
of ResNet-18 and ResNet-50 on CIFAR100.

###  4.5 Investigating the effect of the temperature parameter $\tau$

In Section 2, although the temperature parameter $\tau>0$ does not affect the
global optimality and critical points, it does affect the training speed of
specific learning algorithms and hence test performance. In all experiments in
Section 4, we set $\tau=1$ and have not discussed in detail the effects of
$\tau$ in practice.

However, as mentioned in [21], $\tau$ has important side-effects on
optimization dynamics and must be carefully tuned in practice. Hence, we now
present a brief study of the temperature parameter $\tau$ when optimizing the
problem (4) under the UFM and training a deep network in practice. To begin,
we consider the UFM formulation. We first note that $\tau$ does not affect the
theoretical global solution or benign landscape of the UFM (although it does
affect the attained theoretical lower bound, see Theorem 3.1 and Figure 3).
However, it does impact the rate of convergence to neural collapse as well as
the attained numerical values of the $\mathcal{NC}$ metrics. To see this, we
apply (Riemannian) gradient descent with exact line search to Problem (4) for
various settings of $\tau$. These results are shown in Figure 9.

Figure 9: Effect of temperature parameter on collapse and training accuracy of
UFM. $K=10$ classes, $n=5$ samples per class, $d=32$. Top row: Average
$\mathcal{NC}$ metrics and training accuracy of UFM over 20 trials for various
settings of $\tau$ with respect to each iteration of (Riemannian) gradient
descent. Bottom row: Final average $\mathcal{NC}$ metrics and training
accuracy of UFM over 20 trials for various settings of $\tau$.

First, it is evident that for all tested $\tau$ values, we achieve perfect
classification in a similar number of iterations. Furthermore, the rate of
convergence of $\mathcal{NC}_{1}$ is somewhat the same for most settings of
$\tau$, and we essentially have feature collapse for most settings of $\tau$.
On the other hand, it appears that the rate of convergence of
$\mathcal{NC}_{2}$ and $\mathcal{NC}_{3}$ are dramatically affected by $\tau$,
with values in the range of 1 to 10 yielding the greatest collapse. This
aligns with the choice of the temperature parameter in the experimental
section of [21], where the equivalent parameter is set
$\rho=1/\sqrt{0.1}\approx 3.16$.

We now look to the setting of training practical deep networks. We train a
feature normalized ResNet-18 architecture on CIFAR-10 for various settings of
$\tau$. The results are shown in Figure 10.

Figure 10: Effect of temperature parameter on collapse and training accuracy
of ResNet. Top row: Average $\mathcal{NC}$ metrics and training accuracy of
ResNet-18 on CIFAR-10 with $n=100$ for various settings of $\tau$ with respect
to each epoch over 200 epochs. Bottom row: Final average $\mathcal{NC}$
metrics and training accuracy of ResNet-18 on CIFAR-10 with $n=100$ for
various settings of $\tau$.

One immediate difference from the UFM formulation is that we arrive at perfect
classification of the training data for a particular range of values for
$\tau$ (from about $0.1$ to $10$) but not for all settings of $\tau$. Within
this range, we can see that values of $\tau$ around 1 to 10 lead to the
fastest training, and values close to $\tau=5$ lead to the greatest collapse
in all $\mathcal{NC}$ metrics, as was the case with the UFM. All this evidence
suggests that $\tau=1$ is not the optimal setting of the temperature parameter
for either the UFM or ResNet, particularly when measuring $\mathcal{NC}_{2}$
and $\mathcal{NC}_{3}$, and instead $\tau=5$ may perform better. In the
practical experiments of the main text, however, we mainly focused on training
speed and feature collapse, and for these purposes it appears that the $\tau$
parameter can be set in a fairly nonstringent manner.

###  4.6 Exploration of benign global landscapes of other commonly-used losses

Although the cross-entropy (CE) loss studied in this work is arguably the most
common loss function for deep classification tasks, it is not the only one.
Some other commonly used loss functions include focal loss (FL) [40], label
smoothing (LS) [70], and supervised contrastive (SC) loss [32], each of which
has demonstrated various benefits over vanilla CE. In this section, we briefly
explore the empirical global landscape of these losses under the UFM with
normalized features (and classifiers). Specifically, we consider the problem

$\displaystyle\min_{\mathbf{W},\mathbf{H}}$
$\displaystyle\;f(\mathbf{W},\mathbf{H})\;:=\;\frac{1}{N}\sum_{k=1}^{K}\sum_{i=1}^{n}\mathcal{L}\left(\tau\mathbf{W}^{\top}\mathbf{h}_{k,i},\mathbf{y}_{k}\right),$
(10) s.t.
$\displaystyle\quad\mathbf{H}\in\mathcal{OB}(d,N),\;\mathbf{W}\in\;\mathcal{OB}(d,K).$

where $\mathcal{L}$ is either the focal loss or label smoothing loss.

First, we consider the focal loss, defined as

$\displaystyle\mathcal{L}_{\mathrm{FL}}(\mathbf{z},\mathbf{y}_{k})=-\left(1-\frac{\exp(z_{k})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right)^{\gamma}\log\left(\frac{\exp(z_{k})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right)$

where $\gamma\geq 0$ is the _focusing_ parameter (with $\gamma=0$, we recover
ordinary CE). As seen in Figure 11, using gradient descent with random
initialization on the focal loss, we achieve neural collapse over a range of
settings of $K$ and $\tau$. Characterizing the global solutions of (10) for
the focal loss and a general landscape analysis are left as future work.

Figure 11: Global optimization of focal loss (FL) with $\gamma=3$ under UFM
with $d=16$ and $n=3$. Black line refers to theoretical value of (10) at
$\mathcal{NC}$ solutions. Empirical values found using gradient descent with
random initialization. Left: Lower bound against number of classes $K$ while
fixing $\tau=1$. Right: Lower bound against temperature $\tau$ while fixing
$K=3$. The same empirical values are achieved over many trials.

Next, we consider the label smoothing loss, defined as

$\displaystyle\mathcal{L}_{\mathrm{LS}}(\mathbf{z},\mathbf{y}_{k})=-\left(1-\frac{K-1}{K}\alpha\right)\log\left(\frac{\exp(z_{k})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right)-\frac{\alpha}{K}\sum_{j\neq
k}\log\left(\frac{\exp(z_{j})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right)$

where $\alpha\geq 0$ is the _smoothing_ parameter (with $\alpha=0$, we recover
ordinary CE). As seen in Figure 12, for small enough $\tau$ we achieve neural
collapse, but for larger $\tau$, we do not. In fact, the global solutions of
the label smoothing loss are not neural collapse for large enough $\tau$. To
see this, let $(\mathbf{W}_{1},\mathbf{H}_{1})$ denote a $\mathcal{NC}$
solution, and let $(\mathbf{W}_{2},\mathbf{H}_{2})$ denote a solution where
$\mathbf{W}_{2}=\mathbf{a}\mathbf{1}_{K}^{\top}$ and
$\mathbf{H}_{2}=\mathbf{a}\mathbf{1}_{N}^{\top}$, where $\mathbf{a}$ is unit-
norm. It is easy to compute that

$\displaystyle
f(\mathbf{W}_{1},\mathbf{H}_{1})=\log\left(1+(K-1)\exp\left(-\frac{K\tau}{K-1}\right)\right)+\alpha\tau$

so $f(\mathbf{W}_{1},\mathbf{H}_{1})\rightarrow\infty$ as
$\tau\rightarrow\infty$, whereas $f(\mathbf{W}_{2},\mathbf{H}_{2})=\log(K)$ is
independent of $\tau$. Again, characterizing the global solutions of (10) for
label smoothing and a general landscape analysis are left as future work.

Figure 12: Global optimization of label smoothing (LS) with $\alpha=0.1$ under
UFM with $d=16$ and $n=3$. Black line refers to value of (10) at
$\mathcal{NC}$ solutions. Empirical values found using gradient descent with
random initialization. Left: Lower bound against number of classes $K$ while
fixing $\tau=1$. Right: Lower bound against temperature $\tau$ while fixing
$K=3$. The same empirical values are achieved over many trials.

Finally, we look to the supervised contrastive loss. Unlike the other losses,
we do not have classifier $\mathbf{W}$ when training, so we instead have the
problem

$\displaystyle\min_{\mathbf{H}}f(\mathbf{H}):=$
$\displaystyle-\frac{1}{N(n-1)}\sum_{i=1}^{N}\sum_{\begin{subarray}{c}j\neq
i\\\
y_{j}=y_{i}\end{subarray}}\log\left(\frac{\exp(\tau^{2}\mathbf{h}_{i}^{\top}\mathbf{h}_{j})}{\sum_{\ell\neq
i}\exp(\tau^{2}\mathbf{h}_{i}^{\top}\mathbf{h}_{\ell})}\right)$ (11) s.t.
$\displaystyle\quad\mathbf{H}\in\mathcal{OB}(d,N).$

We note that the loss as written above computes the loss over the entire
dataset, as opposed to computing over all minibatches of a fixed size as in
[21]. As seen in Figure 13, using gradient descent with random initialization
on the supervised contrastive loss, we achieve neural collapse over a range of
settings of $K$ and $\tau$. In fact, it is proven in [21] that the global
minimizers of (11) are $\mathcal{NC}$ solutions. However, an understanding of
the global landscape requires further exploration and is left as future work.

Figure 13: Global optimization of supervised contrastive (SC) loss under UFM
with $d=16$ and $n=3$. Black line refers to value of (11) at $\mathcal{NC}$
solutions. Empirical values found using gradient descent with random
initialization. Left: Lower bound against number of classes $K$ while fixing
$\tau=1$. Right: Lower bound against temperature $\tau$ while fixing $K=3$.
The same empirical values are achieved over many trials.

##  5 Conclusion & Discussion

In this work, motivated by the common practice of feature normalization in
modern deep learning, we study the prevalence of the $\mathcal{NC}$ phenomenon
when last-layer features and classifiers are constrained on the sphere. Based
upon the assumption of the UFM, we formulate the problem as a Riemannian
optimization problem over the product of sphere (i.e., oblique manifold). We
showed that the loss function is a strict saddle function over the manifold
with respect to the last-layer features and classifiers, with no other
spurious local minimizers. We demonstrated that this phenomenon occurs for
overparameterized deep network training, and show the benefits of feature
normalization in terms of training speed, learned representation quality, and
generalization, both for the UFM and for practical deep networks on
classification tasks. We conclude by placing our work in the context of
existing literature, and we briefly discuss several exciting future research
directions, motivated both by previous work as well as the work presented
here.

#### Global optimality of $\mathcal{NC}$ under UFM.

The seminal works [50, 23] inspired many recent theoretical studies of the
$\mathcal{NC}$ phenomenon. Because the training loss of a deep neural network
is highly nonlinear, most works simplify the analysis by assuming
unconstrained feature models (UFM) [47, 87] or layer peeled models [17]. It
basically assumes that the network has infinite expression power so that the
features can be reviewed as free optimization variables. Based upon the UFM,
[44] is the first work justifying the global optimality of $\mathcal{NC}$ and
uniformity based upon a CE loss with normalized features, although their study
is quite simplified in the sense that they assume each class has only one
training sample. The work [17] provided global optimality analysis for the CE
loss with constrained features under more generic settings, and they also
studied the case when the training samples are imbalanced in each class.
However, this work constrains the sum of feature norms, whereas in our work,
we constrain the features independently. The follow-up work [30] extended the
analysis to the unconstrained setting without any penalty. Additionally,
motivated by the commonly used weight decay on network parameters, the work
[87] justifies the global optimality of $\mathcal{NC}$ for the CE loss under
the unconstrained formulation, with penalization on both the features and
classifiers. Its companion work [86] extended the analysis to the MSE loss.
Under the same assumption, other work [21] studied both the CE and SC losses
with features lying in the ball, proving that the only global solutions
satisfy $\mathcal{NC}$ properties. Our setting is most closely related to that
in [21], except we explicitly constrain the features on the sphere. Moreover,
the work [71] studied the setting beyond the simple UFM, showing that, even
for a three-layer nonlinear network, the $\mathcal{NC}$ solutions are the only
global solutions with the MSE training loss. Motivated by feature
normalization, as future work one could study a three-layer nonlinear network
where the penultimate output is projected onto the sphere.

#### Benign global landscape and learning dynamics under UFM.

Since the training loss is highly nonconvex even under the UFM, merely
studying global optimality is not sufficient for guaranteeing efficient global
optimization. More recent works address this issue by investigating the global
landscape properties and learning dynamics of specific training algorithms.
More specifically, under the UFM, [87, 86] showed that the optimization
landscapes of CE and MSE losses have benign global optimization landscapes, in
the sense that every local minimizer satisfies $\mathcal{NC}$ properties and
the remaining critical points are strict saddles with negative curvatures.
These works considered the unconstrained formulation with regularization on
both features and classifiers. In comparison, our work studies the benign
global landscape with features and classifiers constrained over the product of
spheres - to our knowledge, this is the first work to study the global
landscape of the constrained formulation. We have also empirically
demonstrated benign landscapes for other losses such as focal loss and
supervised constrastive loss, but further development is required for a
theoretical analysis of the benign landscape. On the other hand, there is
another line of works studying the implicit bias of learning dynamics under
UFM [47, 59, 22, 54, 53, 30], showing that the convergent direction is along
the direction of the minimum-norm separation problem for both CE and MSE
losses.

#### The empirical phenomena of $\mathcal{NC}$ and feature engineering

Although the seminal works [50] and [23] are the first to summarize the
empirical prevalence of $\mathcal{NC}$ for the commonly used CE and MSE losses
respectively, the idea of designing features with intra-class compactness and
inter-class separability has a richer history. More specifically, in the past
many loss functions, such as center loss [76], large-margin softmax
(L-Softmax) loss [43], and its variants [60, 42, 73, 13] were designed with
similar goals for the task of visual face recognition. Additionally, related
works [51, 52] introduced similar ideas of learning maximal separable features
by fixing the linear classifiers with a simplex-shaped structure. Furthermore,
the work [86] demonstrated that better collapse could potentially lead to
better generalization, which is corroborated by our experiments with
normalized features. However, learning neural collapsed features could also
easily lead to overfitting [82] and vulnerability to data corruptions [87].
Additionally, the collapse of the feature dimension could cause the loss of
intrinsic structure of input data, making the learned features less
transferable [33]. In self-supervised learning, recent works promote feature
diversity and uniformity via contrastive learning [8, 75]. In contrast, a line
of recent work proposed to learn diverse while discriminative representation
by designing a loss that maximizes the coding rate reduction [81, 6]. Instead
of collapsing the features to a single dimension, these works promote within-
class diversity while maintaining maximum between-class separability. As such,
it leads to better robustness and transferability. Since the coding rate is
monotonic in the scale of the features, these works also normalize features on
the sphere, making it a suitable problem to be studied over the manifold
similar to our work. A geometric analysis of the coding rate reduction
landscape is a topic of future research.

#### Benefits of feature normalization

The works [60, 42, 73, 13] first introduced feature normalization and
demonstrated its advantages for learning more separable/discriminative
features, which well motivates our study in this work. In fact, we
demonstrated in Figure 1 and Table 1 that features learned on the sphere are
more separable in low dimensions. However, as seen in our experimental
results, the benefits of feature normalization extend beyond better
separability. We have empirically demonstrated that using feature
normalization leads to faster training and better generalization in some
settings. Further investigation is necessary to gain a better understanding of
the role of feature normalization in these phenomena.

#### Large number of classes $K\gg d$

Lastly, in this work we have presented both theoretical and experimental
results either in the case $d\geq K$ or $d>K$. In many applications, this is
quite reasonable, since the number of classes isn’t too large and we can pick
$d$ accordingly in the network architecture. However, in the case of self-
supervised constrastive learning, we do not have label information to assign
training samples to common class labels, and therefore the number of classes
$K$ can grow very large. In other applications such as recommendation systems
[10] and document retrieval [7] the number of classes can also grow very
large. In Figure 1 and Table 1, we can see that feature normalization under
the UFM gives better separability and classification accuracy when $K\gg d$,
making feature normalization a good candidate for improving performance on
these problems. However, our theory does not readily extend to this case.
Although feature collapse is still a reasonable notion, we can not form a
simplex ETF when $d<K+1$. Instead, we can measure the _uniformity_ of features
over the sphere in place, and in fact the work [75] shows that feature
alignment (collapse) and uniformity are asymptotically minimized for
contrastive loss. A finite sample analysis of the $K\gg d$ using the
uniformity metric would be a meaningful extension to our work and of great
interest for the aforementioned applications.

## Acknowledgement

Can Yaras and Qing Qu acknowledge support from U-M START & PODS grants, NSF
CAREER CCF 2143904, NSF CCF 2212066, NSF CCF 2212326, and ONR N00014-22-1-2529
grants. Peng Wang and Laura Balzano acknowledge support from ARO YIP
W911NF1910027, AFOSR YIP FA9550-19-1-0026, and NSF CAREER CCF-1845076. Zhihui
Zhu acknowledges support from NSF grants CCF-2008460 and CCF-2106881. We would
like to thank Zhexin Wu (ETH), Yan Wen (Tsinghua), and Pengru Huang (UMich)
for fruitful discussion through various stages of the work.

## References

  * [1] P.-A. Absil, R. Mahony, and R. Sepulchre.  Optimization algorithms on matrix manifolds.  Princeton University Press, 2009.
  * [2] I. Ben-Shaul and S. Dekel.  Nearest class-center simplification through intermediate layers.  arXiv preprint arXiv:2201.08924, 2022.
  * [3] N. Boumal.  An introduction to optimization on smooth manifolds.  Available online, May, 3, 2020.
  * [4] N. Boumal.  An introduction to optimization on smooth manifolds.  To appear with Cambridge University Press, Jan 2022.
  * [5] K. H. R. Chan, Y. Yu, C. You, H. Qi, J. Wright, and Y. Ma.  Deep networks from the principle of rate reduction.  arXiv preprint arXiv:2010.14765, 2020.
  * [6] K. H. R. Chan, Y. Yu, C. You, H. Qi, J. Wright, and Y. Ma.  Redunet: A white-box deep network from the principle of maximizing rate reduction.  ArXiv, abs/2105.10446, 2021.
  * [7] W.-C. Chang, X. Y. Felix, Y.-W. Chang, Y. Yang, and S. Kumar.  Pre-training tasks for embedding-based large-scale retrieval.  In International Conference on Learning Representations, 2019.
  * [8] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton.  A simple framework for contrastive learning of visual representations.  arXiv preprint arXiv:2002.05709, 2020.
  * [9] U. Cohen, S. Chung, D. D. Lee, and H. Sompolinsky.  Separability and geometry of object manifolds in deep neural networks.  Nature communications, 11(1):1–13, 2020.
  * [10] P. Covington, J. Adams, and E. Sargin.  Deep neural networks for youtube recommendations.  In Proceedings of the 10th ACM conference on recommender systems, pages 191–198, 2016.
  * [11] C. Criscitiello and N. Boumal.  Efficiently escaping saddle points on manifolds.  Advances in Neural Information Processing Systems, 32, 2019.
  * [12] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [13] J. Deng, J. Guo, N. Xue, and S. Zafeiriou.  Arcface: Additive angular margin loss for deep face recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4690–4699, 2019.
  * [14] D. Doimo, A. Glielmo, A. Ansuini, and A. Laio.  Hierarchical nucleation in deep neural networks.  Advances in Neural Information Processing Systems, 33:7526–7536, 2020.
  * [15] S. S. Du, C. Jin, J. D. Lee, M. I. Jordan, A. Singh, and B. Poczos.  Gradient descent can take exponential time to escape saddle points.  In Advances in neural information processing systems, pages 1067–1077, 2017.
  * [16] C. Fang, H. He, Q. Long, and W. J. Su.  Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training.  Proceedings of the National Academy of Sciences, 118(43), 2021.
  * [17] C. Fang, H. He, Q. Long, and W. J. Su.  Layer-peeled model: Toward understanding well-trained deep neural networks.  arXiv preprint arXiv:2101.12699, 2021.
  * [18] N. Frosst, N. Papernot, and G. Hinton.  Analyzing and improving representations with the soft nearest neighbor loss.  In International conference on machine learning, pages 2012–2020. PMLR, 2019.
  * [19] T. Galanti, A. György, and M. Hutter.  On the role of neural collapse in transfer learning.  In International Conference on Learning Representations, 2022.
  * [20] R. Ge, F. Huang, C. Jin, and Y. Yuan.  Escaping from saddle points—online stochastic gradient for tensor decomposition.  In Proceedings of The 28th Conference on Learning Theory, pages 797–842, 2015.
  * [21] F. Graf, C. Hofer, M. Niethammer, and R. Kwitt.  Dissecting supervised constrastive learning.  In International Conference on Machine Learning, pages 3821–3830. PMLR, 2021.
  * [22] F. Graf, C. Hofer, M. Niethammer, and R. Kwitt.  Neural collapse in deep homogeneous classifiers and the role of weight decay.  In International Conference on Machine Learning, pages 3821–3830. PMLR, 2021.
  * [23] X. Han, V. Papyan, and D. L. Donoho.  Neural collapse under mse loss: Proximity to and dynamics on the central path.  arXiv preprint arXiv:2106.02073, 2021.
  * [24] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [25] C. Hofer, F. Graf, M. Niethammer, and R. Kwitt.  Topologically densified distributions.  In International Conference on Machine Learning, pages 4304–4313. PMLR, 2020.
  * [26] K. Hornik.  Approximation capabilities of multilayer feedforward networks.  Neural networks, 4(2):251–257, 1991.
  * [27] J. Hu, X. Liu, Z.-W. Wen, and Y.-X. Yuan.  A brief introduction to manifold optimization.  Journal of the Operations Research Society of China, 8(2):199–248, 2020.
  * [28] G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger.  Densely connected convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700–4708, 2017.
  * [29] L. Hui, M. Belkin, and P. Nakkiran.  Limitations of neural collapse for understanding generalization in deep learning.  arXiv preprint arXiv:2202.08384, 2022.
  * [30] W. Ji, Y. Lu, Y. Zhang, Z. Deng, and W. J. Su.  An unconstrained layer-peeled perspective on neural collapse.  arXiv preprint arXiv:2110.02796, 2021.
  * [31] C. Jin, M. Jordan, R. Ge, P. Netrapalli, and S. Kakade.  How to escape saddle points efficiently.  In International Conference on Machine Learning, pages 2727–2752, 2017.
  * [32] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan.  Supervised contrastive learning.  Advances in Neural Information Processing Systems, 33:18661–18673, 2020.
  * [33] S. Kornblith, T. Chen, H. Lee, and M. Norouzi.  Why do better loss functions lead to less transferable features?  In NeurIPS, 2021.
  * [34] V. Kothapalli, E. Rasromani, and V. Awatramani.  Neural collapse: A review on modelling principles and generalization.  arXiv preprint arXiv:2206.04041, 2022.
  * [35] A. Krizhevsky, G. Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [36] H.-W. Kuo, Y. Lau, Y. Zhang, and J. Wright.  Geometry and symmetry in short-and-sparse deconvolution.  In International Conference on Machine Learning, pages 3570–3580. PMLR, 2019.
  * [37] Y. Lau, Q. Qu, H.-W. Kuo, P. Zhou, Y. Zhang, and J. Wright.  Short and sparse deconvolution — a geometric approach.  In International Conference on Learning Representations, 2020.
  * [38] T. Laurent and J. Brecht.  Deep linear networks with arbitrary loss: All local minima are global.  In International conference on machine learning, pages 2902–2907. PMLR, 2018.
  * [39] Y. LeCun, C. Cortes, and C. Burges.  Mnist handwritten digit database. at&t labs, 2010.
  * [40] T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Dollár.  Focal loss for dense object detection, 2017.
  * [41] C. Liu, L. Zhu, and M. Belkin.  Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.  Applied and Computational Harmonic Analysis, 59:85–116, 2022.
  * [42] W. Liu, Y. Wen, Z. Yu, M. Li, B. Raj, and L. Song.  Sphereface: Deep hypersphere embedding for face recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 212–220, 2017.
  * [43] W. Liu, Y. Wen, Z. Yu, and M. Yang.  Large-margin softmax loss for convolutional neural networks.  arXiv preprint arXiv:1612.02295, 2016.
  * [44] J. Lu and S. Steinerberger.  Neural collapse with cross-entropy loss.  arXiv preprint arXiv:2012.08465, 2020.
  * [45] Z. Lu, H. Pu, F. Wang, Z. Hu, and L. Wang.  The expressive power of neural networks: a view from the width.  In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 6232–6240, 2017.
  * [46] J. Mamou, H. Le, M. Del Rio, C. Stephenson, H. Tang, Y. Kim, and S. Chung.  Emergence of separable manifolds in deep language representations.  In International Conference on Machine Learning, pages 6713–6723. PMLR, 2020.
  * [47] D. G. Mixon, H. Parshall, and J. Pi.  Neural collapse with unconstrained features.  arXiv preprint arXiv:2011.11619, 2020.
  * [48] G. Naitzat, A. Zhitnikov, and L.-H. Lim.  Topology of deep neural networks.  J. Mach. Learn. Res., 21(184):1–40, 2020.
  * [49] V. Papyan.  Traces of class/cross-class structure pervade deep learning spectra.  Journal of Machine Learning Research, 21(252):1–64, 2020.
  * [50] V. Papyan, X. Han, and D. L. Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  Proceedings of the National Academy of Sciences, 117(40):24652–24663, 2020.
  * [51] F. Pernici, M. Bruni, C. Baecchi, and A. Del Bimbo.  Maximally compact and separated features with regular polytope networks.  In CVPR Workshops, pages 46–53, 2019.
  * [52] F. Pernici, M. Bruni, C. Baecchi, and A. Del Bimbo.  Regular polytope networks.  IEEE Transactions on Neural Networks and Learning Systems, 2021\.
  * [53] T. Poggio and Q. Liao.  Explicit regularization and implicit bias in deep network classifiers trained with the square loss.  arXiv preprint arXiv:2101.00072, 2020.
  * [54] T. Poggio and Q. Liao.  Implicit dynamic regularization in deep networks.  Technical report, Center for Brains, Minds and Machines (CBMM), 2020.
  * [55] Q. Qu, X. Li, and Z. Zhu.  Exact recovery of multichannel sparse blind deconvolution via gradient descent.  SIAM Journal on Imaging Sciences, 13(3):1630–1652, 2020.
  * [56] Q. Qu, J. Sun, and J. Wright.  Finding a sparse vector in a subspace: Linear sparsity using alternating directions.  In Advances in Neural Information Processing Systems, pages 3401–3409, 2014.
  * [57] Q. Qu, Y. Zhai, X. Li, Y. Zhang, and Z. Zhu.  Geometric analysis of nonconvex optimization landscapes for overcomplete learning.  In International Conference on Learning Representations, 2020.
  * [58] Q. Qu, Z. Zhu, X. Li, M. C. Tsakiris, J. Wright, and R. Vidal.  Finding the sparsest vectors in a subspace: Theory, algorithms, and applications.  arXiv preprint arXiv:2001.06970, 2020.
  * [59] A. Rangamani, M. Xu, A. Banburski, Q. Liao, and T. Poggio.  Dynamics and neural collapse in deep classifiers trained with the square loss.  Technical report, Center for Brains, Minds and Machines (CBMM), 2021\.
  * [60] R. Ranjan, C. D. Castillo, and R. Chellappa.  L2-constrained softmax loss for discriminative face verification.  arXiv preprint arXiv:1703.09507, 2017.
  * [61] S. Recanatesi, M. Farrell, M. Advani, T. Moore, G. Lajoie, and E. Shea-Brown.  Dimensionality compression and expansion in deep neural networks.  arXiv preprint arXiv:1906.00443, 2019.
  * [62] U. Shaham, A. Cloninger, and R. R. Coifman.  Provable approximation properties for deep neural networks.  Applied and Computational Harmonic Analysis, 44(3):537–557, 2018\.
  * [63] K. Simonyan and A. Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [64] M. Soltanolkotabi, A. Javanmard, and J. D. Lee.  Theoretical insights into the optimization landscape of over-parameterized shallow neural networks.  IEEE Transactions on Information Theory, 65(2):742–769, 2018.
  * [65] J. Sun, Q. Qu, and J. Wright.  When are nonconvex problems not scary?  arXiv preprint arXiv:1510.06096, 2015.
  * [66] J. Sun, Q. Qu, and J. Wright.  Complete dictionary recovery over the sphere I: Overview and the geometric picture.  IEEE Transactions on Information Theory, 63(2):853–884, 2016.
  * [67] J. Sun, Q. Qu, and J. Wright.  Complete dictionary recovery over the sphere II: Recovery by Riemannian trust$-$region method.  IEEE Transactions on Information Theory, 63(2):853–884, 2017.
  * [68] R. Sun, D. Li, S. Liang, T. Ding, and R. Srikant.  The global landscape of neural networks: An overview.  IEEE Signal Processing Magazine, 37(5):95–108, 2020.
  * [69] Y. Sun, N. Flammarion, and M. Fazel.  Escaping from saddle points on riemannian manifolds.  Advances in Neural Information Processing Systems, 32, 2019.
  * [70] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna.  Rethinking the inception architecture for computer vision, 2015.
  * [71] T. Tirer and J. Bruna.  Extended unconstrained features model for exploring deep neural collapse.  arXiv preprint arXiv:2202.08087, 2022.
  * [72] F. Wang, X. Xiang, J. Cheng, and A. L. Yuille.  Normface: L2 hypersphere embedding for face verification.  In Proceedings of the 25th ACM international conference on Multimedia, pages 1041–1049, 2017.
  * [73] H. Wang, Y. Wang, Z. Zhou, X. Ji, D. Gong, J. Zhou, Z. Li, and W. Liu.  Cosface: Large margin cosine loss for deep face recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5265–5274, 2018.
  * [74] P. Wang, H. Liu, A. M.-C. So, and L. Balzano.  Convergence and recovery guarantees of the k-subspaces method for subspace clustering.  In International Conference on Machine Learning, pages 22884–22918. PMLR, 2022.
  * [75] T. Wang and P. Isola.  Understanding contrastive representation learning through alignment and uniformity on the hypersphere.  In International Conference on Machine Learning, pages 9929–9939. PMLR, 2020.
  * [76] Y. Wen, K. Zhang, Z. Li, and Y. Qiao.  A discriminative feature learning approach for deep face recognition.  In European conference on computer vision, pages 499–515. Springer, 2016.
  * [77] S. Wojtowytsch et al.  On the emergence of simplex symmetry in the final and penultimate layers of neural network classifiers.  arXiv preprint arXiv:2012.05420, 2020.
  * [78] L. Xie, Y. Yang, D. Cai, D. Tao, and X. He.  Neural collapse inspired attraction-repulsion-balanced loss for imbalanced learning.  arXiv preprint arXiv:2204.08735, 2022.
  * [79] Y. Yang, L. Xie, S. Chen, X. Li, Z. Lin, and D. Tao.  Do we really need a learnable classifier at the end of deep neural network?  arXiv preprint arXiv:2203.09081, 2022.
  * [80] E. Yu, Z. Li, and S. Han.  Towards discriminative representation: Multi-view trajectory contrastive learning for online multi-object tracking.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8834–8843, 2022.
  * [81] Y. Yu, K. H. R. Chan, C. You, C. Song, and Y. Ma.  Learning diverse and discriminative representations via the principle of maximal coding rate reduction.  arXiv preprint arXiv:2006.08558, 2020.
  * [82] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals.  Understanding deep learning requires rethinking generalization.  arXiv preprint arXiv:1611.03530, 2016.
  * [83] J. Zhang, Y. Zhang, M. Hong, R. Sun, and Z.-Q. Luo.  When expressivity meets trainability: Fewer than $n$ neurons can work.  Advances in Neural Information Processing Systems, 34:9167–9180, 2021.
  * [84] Y. Zhang, H.-W. Kuo, and J. Wright.  Structured local optima in sparse blind deconvolution.  IEEE Transactions on Information Theory, 66(1):419–452, 2019.
  * [85] Y. Zhang, Q. Qu, and J. Wright.  From symmetry to geometry: Tractable nonconvex problems.  arXiv preprint arXiv:2007.06753, 2020.
  * [86] J. Zhou, X. Li, T. Ding, C. You, Q. Qu, and Z. Zhu.  On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features.  arXiv preprint arXiv:2203.01238, 2022.
  * [87] Z. Zhu, T. Ding, J. Zhou, X. Li, C. You, J. Sulam, and Q. Qu.  A geometric analysis of neural collapse with unconstrained features.  Advances in Neural Information Processing Systems, 34, 2021.

Appendix

#### Organization of the appendices.

The appendix is organized as follows. In Appendix A, we provide some
preliminary tools for analyzing our manifold optimization problem. Based upon
this, the proof of Theorem 3.1 and the proof of Theorem 3.2 are provided in
Appendix B and Appendix C, respectively.

#### Notations.

Before we proceed, let us first introduce the notations that will be used
throughout the appendix. Let $\mathbb{R}^{n}$ denote $n$-dimensional Euclidean
space and $\|\cdot\|_{2}$ be the Euclidean norm. We write matrices in bold
capital letters such as $\mathbf{A}$, vectors in bold lower-case such as
$\mathbf{a}$, and scalars in plain letters such as $a$. Given a matrix
$\mathbf{A}\in\mathbb{R}^{d\times K}$, we denote by $\mathbf{a}_{k}$ its
$k$-th column, $\mathbf{a}^{i}$ its $i$-th row, $a_{ij}$ its $(i,j)$-th
element, and $\|\mathbf{A}\|$ its spectral norm. We use
$\operatorname{diag}(\mathbf{A})$ to denote a vector that consists of diagonal
elements of $\mathbf{A}$ and $\operatorname{ddiag}(\mathbf{A})$ to denote a
diagonal matrix whose diagonal elements are the diagonal ones of $\mathbf{A}$.
We use $\operatorname{diag}(\mathbf{a})$ to denote a diagonal matrix whose
diagonal is $\mathbf{a}$. Given a positive integer $n$, we denote by $[n]$ the
set $\\{1,\dots,n\\}$. We denote the unit sphere in $\mathbb{R}^{d}$ by
$\mathbb{S}^{d-1}:=\\{\mathbf{x}\in\mathbb{R}^{d}:\|\mathbf{x}\|_{2}=1\\}$.

##  Appendix A Preliminaries

In this section, we first review some basic aspects of the Riemannian
optimization and then compute the derivative of the CE loss.

###  A.1 Riemannian Derivatives

According to [4, Chapter 3 & 5] and [27, 1], the tangent space of a general
manifold $\mathcal{M}\subseteq\mathbb{R}^{d}$ at $\mathbf{x}$, denoted by
$\mathrm{T}_{\mathbf{x}}\mathcal{M}$, is defined as the set of all vectors
tangent to $\mathcal{M}$ at $\mathbf{x}$. Based on this, the Riemannian
gradient $\operatorname{grad}f$ of a function $f$ at $\mathbf{x}$ is a unique
vector in $\mathrm{T}_{\mathbf{x}}\mathcal{M}$ satisfying

$\displaystyle\left\langle\operatorname{grad}f,\mathbf{\xi}\right\rangle\;=\;Df(\mathbf{x})[\mathbf{\xi}],\quad\forall\
\mathbf{\xi}\in\mathrm{T}_{\mathbf{x}}\mathcal{M}.$

where $Df(\mathbf{x})[\mathbf{\xi}]$ is the derivative of $f(\gamma(t))$ at
$t=0$, $\gamma(t)$ is any curve on the manifold that satisfies
$\gamma(0)=\mathbf{x}$ and $\dot{\gamma}(0)=\mathbf{\xi}$. The Riemannian
Hessian $\operatorname{Hess}f(\mathbf{x})$ is a mapping from the tangent space
$\mathrm{T}_{\mathbf{x}}\mathcal{M}$ to the tangent space
$\mathrm{T}_{\mathbf{x}}\mathcal{M}$ with

$\displaystyle\operatorname{Hess}f(\mathbf{x})[\mathbf{\xi}]\;=\;\widetilde{\nabla}_{\mathbf{\xi}}\operatorname{grad}f(\mathbf{x}),$

where $\widetilde{\nabla}$ is the Riemannian connection. For a function $f$
defined on the manifold $\mathcal{M}$, if it can be extended smoothly to the
ambient Euclidean space, we have

$\displaystyle\operatorname{grad}f(\mathbf{x})\;$
$\displaystyle=\;\mathcal{P}_{\mathrm{T}_{\mathbf{x}}\mathcal{M}}\left(\nabla
f(\mathbf{x})\right),$
$\displaystyle\operatorname{Hess}f(\mathbf{x})[\mathbf{\xi}]\;$
$\displaystyle=\;\mathcal{P}_{\mathrm{T}_{\mathbf{x}}\mathcal{M}}\left(D\operatorname{grad}f(\mathbf{x})[\mathbf{\xi}]\right).$

where $D$ is the Euclidean differential, and
$\mathcal{P}_{\mathrm{T}_{\mathbf{x}}\mathcal{M}}$ is the projection on the
tangent space $\mathrm{T}_{\mathbf{x}}\mathcal{M}$. According to [4, Example
3.18], if $\mathcal{M}=\mathbb{S}^{p-1}$, then the tangent space and
projection are

$\displaystyle\mathrm{T}_{\mathbf{x}}\mathbb{S}^{p-1}\;=\;\left\\{\mathbf{z}\in\mathbb{R}^{p}\mid\mathbf{x}^{\top}\mathbf{z}=0\right\\},\quad\mathcal{P}_{\mathrm{T}_{\mathbf{x}}\mathbb{S}^{p-1}}\mathbf{z}\;=\;(\mathbf{I}-\mathbf{x}\mathbf{x}^{\top})\mathbf{z}.$

Moreover, the oblique manifold $\mathcal{M}=\mathcal{OB}(p,q)$ is a product of
$q$ unit spheres, and it is also a smooth manifold embedded in
$\mathbb{R}^{p\times q}$, where

$\displaystyle\mathcal{M}=\mathcal{OB}(p,q)\;=\;\underbrace{\mathbb{S}^{p-1}\times\mathbb{S}^{p-1}\times\cdots\times\mathbb{S}^{p-1}}_{q\;\text{times}}\;=\;\left\\{\mathbf{Z}\in\mathbb{R}^{p\times
q}\mid\operatorname{diag}\left(\mathbf{Z}^{\top}\mathbf{Z}\right)=\mathbf{1}\right\\}.$

Correspondingly, the tangent space
$\mathcal{P}_{\mathrm{T}_{\mathbf{X}}\mathcal{OB}(p,q)}$ is

$\displaystyle\mathrm{T}_{\mathbf{X}}\mathcal{OB}(p,q)\;=\;\mathrm{T}_{\mathbf{x}_{1}}\mathbb{S}^{p-1}\times\cdots\times\mathrm{T}_{\mathbf{x}_{q}}\mathbb{S}^{p-1}\;$
$\displaystyle=\;\left\\{\mathbf{Z}\in\mathbb{R}^{p\times
q}\mid\mathbf{x}_{i}^{\top}\mathbf{z}_{i}\;=\;0,\;1\leq i\leq q\right\\},$
$\displaystyle=\;\left\\{\mathbf{Z}\in\mathbb{R}^{p\times
q}\mid\operatorname{diag}\left(\mathbf{X}^{\top}\mathbf{Z}\right)=\mathbf{0}\right\\}$

and the projection operator $\mathrm{T}_{\mathbf{x}_{1}}\mathbb{S}^{p-1}$ is

$\displaystyle\mathcal{P}_{\mathrm{T}_{\mathbf{X}}\mathcal{OB}(p,q)}(\mathbf{Z})\;$
$\displaystyle=\;\begin{bmatrix}\left(\mathbf{I}-\mathbf{x}_{1}\mathbf{x}_{1}^{\top}\right)\mathbf{z}_{1}&\cdots&\left(\mathbf{I}-\mathbf{x}_{q}\mathbf{x}_{q}^{\top}\right)\mathbf{z}_{q}\end{bmatrix}$
$\displaystyle=\;\mathbf{Z}-\mathbf{X}\operatorname{ddiag}(\mathbf{X}^{\top}\mathbf{Z}).$

###  A.2 Derivation of (2.2) and (2.2)

We first derive (2.2). Define the curve

$\displaystyle\phi(t):$
$\displaystyle=f(\mathbf{W}+t\mathbf{\Delta}_{\mathbf{W}},\mathbf{H}+t\mathbf{\Delta}_{\mathbf{H}})$
$\displaystyle=g(\tau(\mathbf{W}+t\mathbf{\Delta}_{\mathbf{W}})^{\top}(\mathbf{H}+t\mathbf{\Delta}_{\mathbf{H}}))$
$\displaystyle=g(\tau\mathbf{W}^{\top}\mathbf{H}+\tau(\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}+\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}})t+\tau\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}}t^{2})$
$\displaystyle=g(\mathbf{M}+\mathbf{\delta}(t))$

where $\mathbf{M}=\tau\mathbf{W}^{\top}\mathbf{H}$ and
$\mathbf{\delta}(t)=\tau(\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}+\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}})t+\tau\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}}t^{2}$
satisfies

$\displaystyle\dot{\mathbf{\delta}}(t)$
$\displaystyle=\tau(\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}+\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}})+2\tau\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}}t$
$\displaystyle\ddot{\mathbf{\delta}}(t)$
$\displaystyle=2\tau\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}}$

so by chain rule and product rule we have

$\displaystyle\dot{\phi}(t)=\left<\dot{\mathbf{\delta}}(t),\nabla
g(\mathbf{M}+\mathbf{\delta}(t))\right>$

and

$\displaystyle\ddot{\phi}(t)=\left<\ddot{\mathbf{\delta}}(t),\nabla
g(\mathbf{M}+\mathbf{\delta}(t))\right>+\nabla^{2}g(\mathbf{M}+\mathbf{\delta}(t))[\dot{\mathbf{\delta}}(t),\dot{\mathbf{\delta}}(t)].$

Then since
$\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]=\ddot{\phi}(0)$,
we have

$\displaystyle\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]$
$\displaystyle=$ $\displaystyle\left<\ddot{\mathbf{\delta}}(0),\nabla
g(\mathbf{M}+\mathbf{\delta}(0))\right>+\nabla^{2}g(\mathbf{M}+\mathbf{\delta}(0))[\dot{\mathbf{\delta}}(0),\dot{\mathbf{\delta}}(0)]$
$\displaystyle=$ $\displaystyle
2\tau\left<\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}},\nabla
g(\mathbf{M})\right>+\nabla^{2}g(\mathbf{M})[\tau(\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}+\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}}),\tau(\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}+\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}})]$

giving the result.

Now we derive (2.2). First, we consider the general case of a function $f$
defined on the oblique manifold $\mathcal{M}=\mathcal{OB}(p,q)$, where $f$ can
be smoothly extended to the ambient Euclidean space. We have

$\displaystyle\operatorname{grad}f(\mathbf{X})$ $\displaystyle=\nabla
f(\mathbf{X})-\mathbf{X}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X})).$

Then

$\displaystyle D\operatorname{grad}f(\mathbf{X})[\mathbf{U}]$
$\displaystyle=\lim_{t\rightarrow 0}\mathbf{\Delta}(t)$

where

$\displaystyle\mathbf{\Delta}(t)=$
$\displaystyle\frac{\operatorname{grad}f(\mathbf{X}+t\mathbf{U})-\operatorname{grad}f(\mathbf{X})}{t}$
$\displaystyle=$ $\displaystyle\frac{\nabla
f(\mathbf{X}+t\mathbf{U})-(\mathbf{X}+t\mathbf{U})\operatorname{ddiag}((\mathbf{X}+t\mathbf{U})^{\top}\nabla
f(\mathbf{X}+t\mathbf{U}))-\nabla
f(\mathbf{X})+\mathbf{X}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X}))}{t}$ $\displaystyle=$ $\displaystyle\frac{\nabla
f(\mathbf{X}+t\mathbf{U})-\nabla
f(\mathbf{X})}{t}-\mathbf{U}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X}+t\mathbf{U}))-\mathbf{X}\operatorname{ddiag}(\mathbf{U}^{\top}\nabla
f(\mathbf{X}+t\mathbf{U}))$
$\displaystyle-\mathbf{X}\operatorname{ddiag}\left(\mathbf{X}^{\top}\frac{\nabla
f(\mathbf{X}+t\mathbf{U})-\nabla f(\mathbf{X})}{t}\right)-t\
\mathbf{U}\operatorname{ddiag}(\mathbf{U}^{\top}\nabla
f(\mathbf{X}+t\mathbf{U}))$

so

$\displaystyle D\operatorname{grad}f(\mathbf{X})[\mathbf{U}]$
$\displaystyle=\nabla^{2}f(\mathbf{X})[\mathbf{U}]-\mathbf{U}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X}))$
$\displaystyle-\mathbf{X}\operatorname{ddiag}(\mathbf{U}^{\top}\nabla
f(\mathbf{X}))-\mathbf{X}\operatorname{ddiag}(\mathbf{U}^{\top}\nabla^{2}f(\mathbf{X})[\mathbf{U}]).$

Now, for $\mathbf{U}\in\mathrm{T}_{\mathbf{X}}\mathcal{M}$, we have
$\operatorname{diag}(\mathbf{U}^{\top}\mathbf{X})=\mathbf{0}$ so

$\displaystyle\operatorname{Hess}f(\mathbf{X})[\mathbf{U},\mathbf{U}]$
$\displaystyle=\left<\mathbf{U},\operatorname{Hess}f(\mathbf{X})[\mathbf{U}]\right>$
$\displaystyle=\left<\mathbf{U},D\operatorname{grad}f(\mathbf{X})-\mathbf{X}\operatorname{ddiag}(\mathbf{X}^{\top}D\operatorname{grad}f(\mathbf{X}))\right>$
$\displaystyle=\left<\mathbf{U},D\operatorname{grad}f(\mathbf{X})\right>$
$\displaystyle=\left<\mathbf{U},\nabla^{2}f(\mathbf{X})[\mathbf{U}]\right>-\left<\mathbf{U},\mathbf{U}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X}))\right>$
$\displaystyle=\nabla^{2}f(\mathbf{X})[\mathbf{U},\mathbf{U}]-\left<\mathbf{U}\operatorname{ddiag}(\mathbf{X}^{\top}\nabla
f(\mathbf{X})),\mathbf{U}\right>.$

Now let $f$ be defined as in (4). Since $(\mathbf{W},\mathbf{H})$ lies on the
product manifold
$\mathcal{OB}(d,K)\times\mathcal{OB}(d,N)=\mathcal{OB}(d,K+N)$ which is also
an oblique manifold, we can simply use the general result above, i.e., for
$\mathbf{\Delta}\in\mathrm{T}_{(\mathbf{W},\mathbf{H})}\mathcal{\mathcal{missing}}{OB}(d,N+K)$,

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]=\;$
$\displaystyle\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]-\left<\mathbf{\Delta}_{\mathbf{W}}\operatorname{ddiag}(\mathbf{W}^{\top}\nabla_{\mathbf{W}}f(\mathbf{W},\mathbf{H})),\mathbf{\Delta}_{\mathbf{W}}\right>$
$\displaystyle-\left<\mathbf{\Delta}_{\mathbf{H}}\operatorname{ddiag}(\mathbf{H}^{\top}\nabla_{\mathbf{H}}f(\mathbf{W},\mathbf{H})),\mathbf{\Delta}_{\mathbf{H}}\right>$

which gives (2.2) after substituting the ordinary Euclidean gradient of $f$.

###  A.3 Derivatives of CE Loss

Note that the CE loss is of the form

$\displaystyle\mathcal{L}_{\mathrm{CE}}\left(\mathbf{z},\mathbf{y}_{k}\right)\;=\;-\log\left(\frac{\exp(z_{k})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}\right)\;=\;\log\left(\sum_{\ell=1}^{K}\exp(z_{\ell})\right)-z_{k}.$

Then, one can verify

$\displaystyle\frac{\partial\mathcal{L}_{\mathrm{CE}}\left(\mathbf{z},\mathbf{y}_{k}\right)}{\partial
z_{j}}\;=\;\begin{cases}\frac{\exp(z_{j})}{\sum_{\ell=1}^{K}\exp(z_{\ell})},&j\neq
k,\\\ \frac{\exp(z_{j})}{\sum_{\ell=1}^{K}\exp(z_{\ell})}-1,&j=k,\end{cases}$

for all $j\in[K]$. Thus, we have

$\displaystyle\nabla\mathcal{L}_{\mathrm{CE}}(\mathbf{z},\mathbf{y}_{k})\;=\;\frac{\exp\left(\mathbf{z}\right)}{\sum_{\ell=1}^{K}\exp(z_{\ell})}-\mathbf{e}_{k}\;=\;\eta(\mathbf{z})-\mathbf{e}_{k},$

where $\eta(\mathbf{z})$ is a softmax function, with

$\displaystyle\eta(z_{j})\;:=\;\frac{\exp\left(z_{j}\right)}{\sum_{\ell=1}^{K}\exp(z_{\ell})}.$

Furthermore, we have

$\displaystyle\nabla^{2}\mathcal{L}_{\mathrm{CE}}(\mathbf{z},\mathbf{y}_{k})\;=\;\operatorname{diag}(\eta(\mathbf{z}))-\eta(\mathbf{z})\eta(\mathbf{z})^{\top}.$

##  Appendix B Proof of Theorem 3.1

In this section, we first simplify Problem (4) by utilizing its structure,
then characterize the structure of global solutions of the simplified problem,
and finally deduce the struture of global solutions of Problem (4) based on
their relationship. Before we proceed, we can first reformulate Problem (4) as
follows. Let

$\displaystyle\mathbf{H}\;=\;\begin{bmatrix}\mathbf{H}^{1}&\mathbf{H}^{2}&\cdots&\mathbf{H}^{n}\end{bmatrix}\in\mathbb{R}^{d\times
N},\;\mathbf{H}^{i}\;=\;\begin{bmatrix}\mathbf{h}_{1,i}&\mathbf{h}_{2,i}&\cdots&\mathbf{h}_{K,i}\end{bmatrix}\in\mathbb{R}^{d\times
K},\;\forall\ i\in[N],$

and $\bar{f}:\mathbb{R}^{d\times K}\times\mathbb{R}^{d\times
K}\rightarrow\mathbb{R}$ be such that

$\displaystyle\bar{f}(\mathbf{W},\mathbf{Q})=\frac{1}{K}\sum_{k=1}^{K}\mathcal{L}_{\mathrm{CE}}\left(\tau\mathbf{W}^{\top}\mathbf{q}_{k},\mathbf{y}_{k}\right).$
(12)

Then, we can rewrite the objective function of Problem (4) as

$\displaystyle
f(\mathbf{W},\mathbf{H})=\frac{1}{n}\sum_{i=1}^{n}\bar{f}(\mathbf{W},\mathbf{H}^{i}).$
(13)

###### Lemma B.1

Suppose that $(\mathbf{W}^{*},\mathbf{Q}^{*})$ is an optimal solution of

$\displaystyle\min_{\mathbf{W}\in\mathbb{R}^{d\times
K},\mathbf{Q}\in\mathbb{R}^{d\times
K}}\bar{f}(\mathbf{W},\mathbf{Q})\quad\operatorname*{s.t.\ }\
\mathbf{Q}\in\mathcal{OB}(d,K),\;\mathbf{W}\in\;\mathcal{OB}(d,K).$ (14)

Then, $(\mathbf{W}^{*},\mathbf{H}^{*})$ with
$\mathbf{H}^{*}=\begin{bmatrix}\mathbf{Q}^{*}&\mathbf{Q}^{*}&\cdots&\mathbf{Q}^{*}\end{bmatrix}$
is an optimal solution of Problem (4).

Proof According to (13), we note that

$\displaystyle\min\left\\{f(\mathbf{W},\mathbf{H}):\
\mathbf{H}\in\mathcal{OB}(d,N),\;\mathbf{W}\in\;\mathcal{OB}(d,K)\right\\}$
$\displaystyle\geq\ $
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\min\left\\{\bar{f}(\mathbf{W}^{i},\mathbf{H}^{i}):\
\mathbf{H}^{i}\in\mathcal{OB}(d,K),\;\mathbf{W}^{i}\in\;\mathcal{OB}(d,K)\right\\},$

where equality holds if
$(\mathbf{W}^{i},\mathbf{H}^{i})=(\mathbf{W}^{*},\mathbf{Q}^{*})$ for all
$i\in[n]$ and $(\mathbf{W},\mathbf{H})=(\mathbf{W}^{*},\mathbf{Q}^{*})$. This,
together with the fact that $(\mathbf{W}^{*},\mathbf{Q}^{*})$ is an optimal
solution of Problem (14), implies the desired result.

Based on the above lemma, it suffices to consider the global optimality
condition of Problem (14).

###### Proposition 1

Suppose that the feature dimension is no smaller than the number of classes
(i.e., $d\geq K$) and the training labels are balanced in each class (i.e.,
$n=n_{1}=\cdots=n_{K}$). Then, any global minimizer
$(\mathbf{W},\mathbf{Q})\in\mathcal{OB}(d,K)\times\mathcal{OB}(d,K)$ of
Problem (14) satisfies

$\displaystyle\mathbf{Q}=\mathbf{W},\quad\mathbf{Q}^{T}\mathbf{Q}=\frac{K}{K-1}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right).$
(15)

Proof According to [87, Lemma D.5], it holds for all $k\in[K]$ and any
$c_{1}>0$ that

$\displaystyle(1+c_{1})(K-1)\left(\mathcal{L}_{\mathrm{CE}}\left(\tau\mathbf{W}^{\top}\mathbf{q}_{k},\mathbf{y}_{k}\right)-c_{2}\right)$
$\displaystyle\geq\tau\left(\sum_{\ell=1}^{K}\mathbf{w}_{\ell}^{\top}\mathbf{q}_{k}-K\mathbf{w}_{k}^{\top}\mathbf{q}_{k}\right),$

where

$\displaystyle
c_{2}=\frac{1}{1+c_{1}}\log\left((1+c_{1})(K-1)\right)+\frac{c_{1}}{1+c_{1}}\log\left(\frac{1+c_{1}}{c_{1}}\right)$

and the equality holds when
$\mathbf{w}_{i}^{\top}\mathbf{q}_{k}=\mathbf{w}_{j}^{\top}\mathbf{q}_{k}$ for
all $i,j\neq k$ and

$\displaystyle
c_{1}=\left((K-1)\exp\left(\frac{\sum_{\ell=1}^{K}\mathbf{w}_{\ell}^{\top}\mathbf{q}_{k}-K\mathbf{w}_{k}^{\top}\mathbf{q}_{k}}{K-1}\right)\right)^{-1}.$

This, together with (12), implies

$\displaystyle(1+c_{1})(K-1)\left(\bar{f}(\mathbf{W},\mathbf{Q})-c_{2}\right)$
$\displaystyle\geq\frac{\tau}{K}\sum_{k=1}^{K}\left(\sum_{\ell=1}^{K}\mathbf{w}_{\ell}^{\top}\mathbf{q}_{k}-K\mathbf{w}_{k}^{\top}\mathbf{q}_{k}\right)$
$\displaystyle=\frac{\tau}{K}\left(\sum_{k=1}^{K}\sum_{\ell=1}^{K}\mathbf{w}_{k}^{\top}\mathbf{q}_{\ell}-K\sum_{k=1}^{K}\mathbf{w}_{k}^{\top}\mathbf{q}_{k}\right)$
$\displaystyle=\tau\sum_{k=1}^{K}\mathbf{w}_{k}^{\top}\left(\bar{\mathbf{q}}-\mathbf{q}_{k}\right),$

where the first inequality becomes equality when
$\mathbf{w}_{i}^{\top}\mathbf{q}_{k}=\mathbf{w}_{j}^{\top}\mathbf{q}_{k}$ for
all $i,j\neq k$ and all $k\in[K]$ and
$\bar{\mathbf{q}}=\frac{1}{K}\sum_{\ell=1}^{K}\mathbf{q}_{\ell}$ in the last
equality. Note that that
$\mathbf{u}^{\top}\mathbf{v}\geq-\frac{c_{3}}{2}\|\mathbf{u}\|_{2}^{2}-\frac{1}{2c_{3}}\|\mathbf{v}\|_{2}^{2}$
for any $c_{3}>0$, where the equality holds when
$c_{3}\mathbf{u}=-\mathbf{v}$. Consequently, it holds for any $c_{3}>0$ that

$\displaystyle(1+c_{1})(K-1)\left(\bar{f}(\mathbf{W},\mathbf{Q})-c_{2}\right)$
$\displaystyle\geq-\tau\sum_{k=1}^{K}\left(\frac{c_{3}}{2}\|\mathbf{w}_{k}\|_{2}^{2}+\frac{1}{2c_{3}}\|\bar{\mathbf{q}}-\mathbf{q}_{k}\|_{2}^{2}\right)$
$\displaystyle=-\frac{\tau}{2}\left(c_{3}\sum_{k=1}^{K}\|\mathbf{w}_{k}\|_{2}^{2}+\frac{1}{c_{3}}\sum_{k=1}^{K}\|\mathbf{q}_{k}\|_{2}^{2}-\frac{K}{c_{3}}\|\bar{\mathbf{q}}\|_{2}^{2}\right)$
$\displaystyle\geq-\frac{\tau}{2}\left(c_{3}\sum_{k=1}^{K}\|\mathbf{w}_{k}\|_{2}^{2}+\frac{1}{c_{3}}\sum_{k=1}^{K}\|\mathbf{q}_{k}\|_{2}^{2}\right)=-\frac{\tau}{2}\left(c_{3}K+\frac{K}{c_{3}}\right),$

where the first inequality becomes equality when
$c_{3}\mathbf{w}_{k}=\mathbf{q}_{k}-\bar{\mathbf{q}}$ for all $k\in[K]$, the
second inequality becomes equality when $\bar{\mathbf{q}}=\mathbf{0}$, and the
last equality is due to $\mathbf{Q}\in\mathcal{OB}(d,K)$ and
$\mathbf{W}\in\;\mathcal{OB}(d,K)$. Thus, we have

$\displaystyle(1+c_{1})(K-1)\left(\bar{f}(\mathbf{W},\mathbf{Q})-c_{2}\right)\geq-\frac{\tau
K}{2}\left(c_{3}+\frac{1}{c_{3}}\right),$

where the equality holds when
$\mathbf{w}_{i}^{\top}\mathbf{q}_{k}=\mathbf{w}_{j}^{\top}\mathbf{q}_{k}$ for
all $i,j\neq k$ and all $k\in[K]$, $c_{3}\mathbf{w}_{k}=\mathbf{q}_{k}$ for
all $k\in[K]$, and $\sum_{k=1}^{K}\mathbf{q}_{k}=\mathbf{0}$. This, together
with $\mathbf{Q}\in\mathcal{OB}(d,K)$ and $\mathbf{W}\in\;\mathcal{OB}(d,K)$,
implies $c_{3}=1$. Thus, we have $\mathbf{q}_{k}=\mathbf{w}_{k}$ for all
$k\in[K]$ and

$\displaystyle\bar{f}(\mathbf{W},\mathbf{Q})\geq-\frac{\tau
K}{(1+c_{1})(K-1)}+c_{2}.$

This further implies that $\sum_{k=1}^{K}\mathbf{w}_{k}=0$,
$\mathbf{w}_{i}^{\top}\mathbf{w}_{k}=\mathbf{w}_{j}^{\top}\mathbf{w}_{k}$ for
all $i,j\neq k$ and all $k\in[K]$. Then, it holds that for all $1\leq
k\neq\ell\leq K$ that

$\displaystyle\langle\mathbf{w}_{\ell},\mathbf{w}_{k}\rangle=-\frac{1}{K-1}.$

These, together with $\mathbf{Q}\in\mathcal{OB}(d,K)$ and
$\mathbf{W}\in\;\mathcal{OB}(d,K)$, imply (15).

Proof [Proof of Theorem 3.1] According to (13), Lemma B.1, and 1, the global
solutions of Problem (4) take the form of

$\displaystyle\mathbf{h}_{k,i}=\mathbf{q}_{k},\
\mathbf{w}_{k}=\mathbf{q}_{k},\ \forall\ k\in[K],\ i\in[N],$

and

$\displaystyle\mathbf{Q}^{T}\mathbf{Q}=\frac{K}{K-1}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right).$

Based on this and the objective function in Problem (4), the value at an
optimal solution $(\mathbf{W}^{*},\mathbf{H}^{*})$ is

$\displaystyle
f(\mathbf{W}^{*},\mathbf{H}^{*})=\log\left(1+\frac{(K-1)\exp\left(-\frac{\tau}{K-1}\right)}{\exp(\tau)}\right)=\log\left(1+(K-1)\exp\left(-\frac{K\tau}{K-1}\right)\right).$

Then, we complete the proof.

##  Appendix C Proof of Theorem 3.2

In this section, we first analyze the first-order optimality condition of
Problem (4), then characterize the global optimality condition of Problem (4),
and finally prove no spurious local minima and strict saddle point property
based on the previous optimality conditions. For ease of exposition, let us
denote

$\displaystyle\mathbf{M}:=\tau\mathbf{W}^{\top}\mathbf{H},\
g(\mathbf{M}):=f(\mathbf{W},\mathbf{H})=\frac{1}{N}\sum_{i=1}^{n}\sum_{k=1}^{K}\mathcal{L}_{\mathrm{CE}}(\mathbf{m}_{k,i},\mathbf{y}_{k}).$
(16)

Then we have the gradient

$\displaystyle\nabla
f(\mathbf{W},\mathbf{H})\;=\;\left(\nabla_{\mathbf{W}}f(\mathbf{W},\mathbf{H}),\nabla_{\mathbf{H}}f(\mathbf{W},\mathbf{H})\right)$

with

$\displaystyle\nabla_{\mathbf{W}}f(\mathbf{W},\mathbf{H})\;=\;\tau\mathbf{H}\nabla
g(\mathbf{M})^{\top},\
\nabla_{\mathbf{H}}f(\mathbf{W},\mathbf{H})\;=\;\tau\mathbf{W}\nabla
g(\mathbf{M}),$ (17)

and

$\displaystyle\nabla
g(\mathbf{M})\;=\;\begin{bmatrix}\eta(\mathbf{m}_{1,1})&\cdots&\eta(\mathbf{m}_{K,n})\end{bmatrix}-\mathbf{I}_{K}\otimes\mathbf{1}_{n}^{\top},\quad\eta(\mathbf{m})\;=\;\frac{\exp\left(\mathbf{m}\right)}{\sum_{i=1}^{K}\exp\left(m_{i}\right)}.$
(18)

###  C.1 First-Order Optimality Condition

By using the tools in Section A.1, we can calculate the Riemannian gradient at
a given point
$(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,N)\times\mathcal{OB}(d,K)$ as in (6)
and (5). Thus, for a point
$(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,N)\times\mathcal{OB}(d,K)$, the
first-order optimality condition of Problem (4) is

$\displaystyle\operatorname{grad}_{\mathbf{W}}f(\mathbf{W},\mathbf{H})\;$
$\displaystyle=\;\tau\mathbf{W}\nabla
g(\mathbf{M})-\tau\mathbf{H}\operatorname{ddiag}\left(\mathbf{H}^{\top}\mathbf{W}\nabla
g(\mathbf{M})\right)\;=\;\mathbf{0},$ (19)
$\displaystyle\operatorname{grad}_{\mathbf{H}}f(\mathbf{W},\mathbf{H})\;$
$\displaystyle=\;\tau\mathbf{H}\nabla
g(\mathbf{M})^{\top}-\tau\mathbf{W}\operatorname{ddiag}\left(\mathbf{W}^{\top}\mathbf{H}\nabla
g(\mathbf{W})^{\top}\right)\;=\;\mathbf{0}.$ (20)

We denote the set of all critical points by

$\displaystyle\mathcal{C}\;:=\;\left\\{(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,K)\times\mathcal{OB}(d,N)\;\mid\;\operatorname{grad}_{\mathbf{H}}f(\mathbf{W},\mathbf{H})=\mathbf{0},\;\operatorname{grad}_{\mathbf{W}}f(\mathbf{W},\mathbf{H})=\mathbf{0}\right\\}.$

###### Lemma C.1

Suppose that $\mathbf{g}_{i}\in\mathbb{R}^{K}$ and
$\mathbf{g}^{k}\in\mathbb{R}^{N}$ denote the $i$-th column and $k$-th row
vectors of the matrix

$\displaystyle\mathbf{G}:=\nabla g(\mathbf{M})\in\mathbb{R}^{K\times N},$

respectively. Let $\mathbf{\alpha}\in\mathbb{R}^{K}$ and
$\mathbf{\beta}\in\mathbb{R}^{N}$ be such that

$\displaystyle\alpha_{k}\;=\;\left\langle\mathbf{w}_{k},\mathbf{H}\mathbf{g}^{k}\right\rangle,\forall\
k\in[K],\quad\beta_{i}\;=\;\left\langle\mathbf{h}_{i},\mathbf{W}\mathbf{g}_{i}\right\rangle,\forall\
i\in[N].$ (21)

Then it holds for any $(\mathbf{W},\mathbf{H})\in\mathcal{C}$ that

$\displaystyle\mathbf{H}\mathbf{g}^{k}\;=\;\alpha_{k}\mathbf{w}_{k},\ \forall\
k\in[K],\quad\mathbf{W}\mathbf{g}_{i}\;=\;\beta_{i}\mathbf{h}_{i},\ \forall\
i\in[N].$ (22)

and

$\displaystyle|\alpha_{k}|=\|\mathbf{H}\mathbf{g}^{k}\|_{2},\ k=1,\dots
K,\quad|\beta_{i}|=\|\mathbf{W}\mathbf{g}_{i}\|_{2},\ i=1,\dots,N.$ (23)

Proof According to (17), we have

$\displaystyle\mathbf{H}\mathbf{G}^{\top}=\begin{bmatrix}\mathbf{H}\mathbf{g}^{1}&\dots&\mathbf{H}\mathbf{g}^{K}\end{bmatrix},\quad\mathbf{W}\mathbf{G}=\begin{bmatrix}\mathbf{W}\mathbf{g}_{1}&\dots&\mathbf{W}\mathbf{g}_{K}\end{bmatrix}$

Using (21), we can compute

$\displaystyle\operatorname{ddiag}\left(\mathbf{W}^{\top}\mathbf{H}\mathbf{G}^{\top}\right)=\operatorname{diag}(\mathbf{\alpha}),\quad\operatorname{ddiag}\left(\mathbf{H}^{\top}\mathbf{W}\mathbf{G}\right)=\operatorname{diag}(\mathbf{\beta})$

This, together with (19) and (20), implies (22). Since
$\|\mathbf{w}_{k}\|_{2}=1$ for all $k\in[K]$ and $\|\mathbf{h}_{i}\|_{2}=1$
for all $i\in[N]$, by

$\displaystyle\alpha_{k}^{2}=\langle\alpha_{k}\mathbf{w}_{k},{\mathbf{H}\mathbf{g}^{k}}\rangle=\left\|\mathbf{H}\mathbf{g}^{k}\right\|_{2}^{2},\quad\beta_{i}^{2}=\langle\beta_{i}\mathbf{h}_{i},\mathbf{W}\mathbf{g}_{i}\rangle=\left\|\mathbf{W}\mathbf{g}_{i}\right\|_{2}^{2}$

which implies (23).

###  C.2 Characterization of Global Optimality

According to Theorem 3.1, it holds that for any global solution
$(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,N)\times\mathcal{OB}(d,K)$ that

$\displaystyle\mathbf{H}=\mathbf{W}\otimes\mathbf{1}_{n}^{\top},\
\mathbf{W}^{\top}\mathbf{W}=\frac{K}{K-1}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right),$
(24)

where $\otimes$ denotes the Kronecker product.

###### Lemma C.2

Given any critical point $(\mathbf{W},\mathbf{H})\in\mathcal{C}$, let
$\mathbf{\alpha}\in\mathbb{R}^{K}$ and $\mathbf{\beta}\in\mathbb{R}^{N}$ be
defined as in (21). Then, $(\mathbf{W},\mathbf{H})$ is a global solution of
Problem (4) if and only if the corresponding
$(\mathbf{\alpha},\mathbf{\beta})$ satisfies

$\displaystyle\alpha_{k}\leq-\sqrt{n}\|\nabla g(\mathbf{M})\|,\ \forall\
k\in[K],\quad\beta_{i}\leq-\frac{\|\nabla g(\mathbf{M})\|}{\sqrt{n}},\
\forall\ i\in[N],$ (25)

where $\mathbf{M}=\tau\mathbf{W}^{\top}\mathbf{H}$.

Proof Suppose that $(\mathbf{W},\mathbf{H})\in\mathcal{C}$ is an optimal
solution. According to (24), one can verify that

$\displaystyle\mathbf{W}^{\top}\mathbf{H}=\mathbf{W}^{\top}\left(\mathbf{W}\otimes\mathbf{1}_{n}^{\top}\right)=\frac{K}{K-1}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right)\otimes\mathbf{1}_{n}^{\top}.$

According to this and (16), we can compute

$\displaystyle\nabla
g(\mathbf{M})=\frac{-K\exp\left(-\frac{1}{K-1}\right)}{\exp(1)+(K-1)\exp\left(-\frac{1}{K-1}\right)}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right)\otimes\mathbf{1}_{n}^{\top}.$
(26)

This, together with
$\alpha_{k}=\langle\mathbf{w}_{k},\mathbf{H}\mathbf{g}^{k}\rangle$, yields for
all $k\in K$,

$\displaystyle\alpha_{k}=\langle\mathbf{H}^{\top}\mathbf{w}_{k},\mathbf{g}^{k}\rangle\;=\;\frac{-nK\exp\left(-\frac{1}{K-1}\right)}{\exp(1)+(K-1)\exp\left(-\frac{1}{K-1}\right)}.$
(27)

By the same argument, we can compute for all $i\in[N]$,

$\displaystyle\beta_{i}=\frac{-K\exp\left(-\frac{1}{K-1}\right)}{\exp(1)+(K-1)\exp\left(-\frac{1}{K-1}\right)}.$
(28)

According to (26), one can verify

$\displaystyle\|\nabla
g(\mathbf{M})\|=\frac{\sqrt{n}K\exp\left(-\frac{1}{K-1}\right)}{\exp(1)+(K-1)\exp\left(-\frac{1}{K-1}\right)}.$

This, together with (27) and (28), implies (25)

Suppose that a critical point $(\mathbf{W}^{*},\mathbf{H}^{*})\in\mathcal{C}$
satisfies (25). Let $\mathbf{M}^{*}=\tau\mathbf{W}^{*^{\top}}\mathbf{H}^{*}$
and $\lambda=\|\nabla g(\mathbf{M}^{*})\|$. According to (23) and the fact
that $\|\mathbf{w}_{k}^{*}\|=1$ and $\|\mathbf{h}_{k}^{*}\|=1$ for all
$k=1,\dots K$, we have

$\displaystyle\sum_{k=1}^{K}\alpha_{k}^{*^{2}}=\|\mathbf{H}^{*}\nabla
g(\mathbf{M}^{*})^{\top}\|_{F}^{2}\leq\|\nabla
g(\mathbf{M}^{*})\|^{2}\|\mathbf{H}^{*}\|_{F}^{2}=\lambda^{2}N,$
$\displaystyle\sum_{i=1}^{N}\beta_{i}^{*^{2}}=\|\mathbf{W}^{*}\nabla
g(\mathbf{M}^{*})\|_{F}^{2}\leq\|\nabla
g(\mathbf{M}^{*})\|^{2}\|\mathbf{W}^{*}\|_{F}^{2}=\lambda^{2}K.$

This, together with (25), implies

$\displaystyle\alpha_{k}^{*}=-\sqrt{n}\lambda,\ \forall\
k\in[K],\quad\beta_{i}^{*}=-\frac{\lambda}{\sqrt{n}},\ \forall\ i\in[N].$ (29)

Then, we consider the following regularized problem:

$\displaystyle\min_{\mathbf{W}\in\mathbb{R}^{d\times
K},\mathbf{H}\in\mathbb{R}^{d\times
N}}f(\mathbf{W},\mathbf{H})+\frac{\lambda\sqrt{n}}{2}\|\mathbf{W}\|_{F}^{2}+\frac{\lambda}{2\sqrt{n}}\|\mathbf{H}\|_{F}^{2}.$
(30)

According to the fact that $(\mathbf{W}^{*},\mathbf{H}^{*})$ is a critical
point of Problem (4) and satisfies (29), (19), and (20), we have

$\displaystyle\begin{cases}\mathbf{H}^{*}\nabla
g(\mathbf{M}^{*})^{\top}+\lambda{\sqrt{n}}\mathbf{W}^{*}=\mathbf{0},\\\
\mathbf{W}^{*}\nabla
g(\mathbf{M}^{*})+\lambda\mathbf{H}^{*}/\sqrt{n}=\mathbf{0}.\end{cases}$ (31)

This, together with the first-order optimality condition of Problem (30),
yields that $(\mathbf{W}^{*},\mathbf{H}^{*})$ is a critical point of Problem
(30). According to [57, Lemma C.4] and $\|\nabla g(\mathbf{M}^{*})\|=\lambda$,
it holds that $(\mathbf{W}^{*},\mathbf{H}^{*})$ is an optimal solution of
Problem (30). This, together with [57, Theorem 3.1], yields that
$(\mathbf{W}^{*},\mathbf{H}^{*})\in\mathcal{C}$ satisfies

$\displaystyle\mathbf{H}^{*}=\mathbf{W}^{*}\otimes\mathbf{1}_{n}^{\top},\
\mathbf{W}^{*^{\top}}\mathbf{W}^{*}=\frac{K}{K-1}\left(\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\right).$

According to Theorem 3.1, we conclude that $(\mathbf{W}^{*},\mathbf{H}^{*})$
is an optimal solution of Problem (4). Then, we complete the proof.

###  C.3 Negative Curvature at Saddle Points

###### Lemma C.3

Let $\mathbf{\alpha}$ and $\mathbf{\beta}$ be defined as in Lemma C.1. Then
$\sum_{k=1}^{K}\alpha_{k}=\sum_{i=1}^{N}\beta_{i}$.

Proof Given the definition of $\mathbf{\alpha}$ and $\mathbf{\beta}$ in (21),
this follows directly from cyclic property of trace:

$\displaystyle\sum_{k=1}^{K}\alpha_{k}=\operatorname{trace}(\mathbf{W}^{\top}\mathbf{H}\mathbf{G}^{\top})=\operatorname{trace}(\mathbf{G}\mathbf{H}^{\top}\mathbf{W})=\operatorname{trace}(\mathbf{H}^{\top}\mathbf{W}\mathbf{G})=\sum_{i=1}^{N}\beta_{i},$

as desired.

###### Lemma C.4

Suppose $(\mathbf{W},\mathbf{H})$ is a critical point and there exists
$i\in[N]$ such that $\beta_{i}=0$. Then there exists
$\mathbf{w}\in\mathbb{S}^{d-1}$ such that
$\mathbf{W}=\mathbf{w}\mathbf{1}_{K}^{\top}$. Furthermore, we have
$\beta_{1}=\ldots=\beta_{N}=0$.

Proof Suppose $nk\leq i<n(k+1)$ for $k\in[K]$ (i.e., $\mathbf{h}_{i}$ has
label $y_{k}$). Thus, we can write each entry of the gradient $\mathbf{g}_{i}$
of the CE loss as

$\displaystyle g_{i\ell}=\begin{cases}p_{ik}-1&\ell=k\\\ p_{i\ell}&\ell\neq
k\end{cases}\quad\mbox{ where }\quad
p_{i\ell}=\frac{\exp(\tau\mathbf{w}_{\ell}^{\top}\mathbf{h}_{i})}{\sum_{j=1}^{K}\exp(\tau\mathbf{w}_{j}^{\top}\mathbf{h}_{i})}.$

Since $\exp(\cdot)>0$ and $K\geq 2$, we have $0<p_{i\ell}<1$. Given that
$\beta_{i}=0$ and $\left\|\mathbf{h}_{i}\right\|_{2}=1$, from (22) we know
that we must have $\mathbf{W}\mathbf{g}_{i}=\mathbf{0}$, which further gives

$\displaystyle g_{ik}\mathbf{w}_{k}+\sum_{\ell\neq
k}g_{i\ell}\mathbf{w}_{\ell}=0.$

Given $1-p_{ik}>0$, equivalently we have

$\displaystyle\mathbf{w}_{k}=\sum_{\ell\neq
k}\frac{p_{i\ell}}{1-p_{ik}}\mathbf{w}_{\ell},$

where $\sum_{\ell\neq k}\frac{p_{i\ell}}{1-p_{ik}}=1$ and $p_{i\ell}>0$ so
$\mathbf{w}_{k}$ is a strict convex combination of points
$\\{\mathbf{w}_{\ell}\\}_{\ell\neq k}$ on the unit sphere. But since
$\mathbf{w}_{k}$ also lies on the unit sphere, and the convex hull of points
on the sphere only intersects with the sphere at
$\\{\mathbf{w}_{\ell}\\}_{\ell\neq k}$, we must have all $\mathbf{w}_{\ell}$
be identical, i.e., $\mathbf{w}_{1}=\ldots=\mathbf{w}_{K}$. Therefore, we can
write $\mathbf{W}=\mathbf{w}_{1}\mathbf{1}_{K}^{\top}$, and consequently

$\displaystyle\mathbf{W}\mathbf{G}=\mathbf{w}_{1}\mathbf{1}_{K}^{\top}\mathbf{G}=\mathbf{0},$

where the last equality follows from the fact that
$\mathbf{1}_{K}^{\top}\mathbf{G}=\mathbf{1}_{K}^{\top}\nabla\mathbf{g}(\mathbf{M})=\mathbf{0}$.
Thus, given
$\beta_{i}=\left\langle\mathbf{h}_{i},\mathbf{W}\mathbf{g}_{i}\right\rangle$,
from the above we have $\beta_{1}=\ldots=\beta_{N}=0$.

###### Lemma C.5

For any $\mathbf{H}\in\mathcal{OB}(d,N)$ and $\mathbf{w}\in\mathbb{S}^{d-1}$,
there exists at least one $\mathbf{a}\in\mathbb{S}^{d-1}$ such that for any
$0<\tau<2(d-2)(1+(K\mbox{{ mod }}2)/K)^{-1}$, we have

$\displaystyle\mathbf{a}^{\top}\mathbf{w}=\mathbf{0}\quad\mbox{ and
}\quad\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}<\Gamma:=\frac{2N}{\tau(1+(K\mbox{{
mod }}2)/K)+2}.$ (32)

Proof To establish the result, we need to show that there exists a linear
subspace $\mathcal{S}\subset\mathbb{R}^{d}$ with $\dim(\mathcal{S})\geq 2$
such that for any nonzero $\mathbf{z}\in\mathcal{S}$ we have
$\|\mathbf{H}^{\top}\mathbf{z}\|_{2}^{2}<\Gamma\|\mathbf{z}\|_{2}^{2}$. Then

$\displaystyle\dim(\mathcal{S}\cap\mathcal{N}(\mathbf{w}))>0,$

where $\mathcal{N}(\mathbf{w})$ denotes the null space of $\mathbf{w}$, so if
we choose unit-norm $\mathbf{a}\in\mathcal{S}\cap\mathcal{N}(\mathbf{w})$, we
can obtain the desired results. Let $(\sigma_{\ell}^{2}(\mathbf{H}),v_{\ell})$
denote the $\ell$-th eigenvalue-eigenvector pair of
$\mathbf{H}\mathbf{H}^{\top}\in\mathbb{R}^{d\times d}$ for $\ell\in[d]$. Given
the fact $\mathbf{H}\in\mathcal{OB}(d,N)$, it is obvious that

$\displaystyle\|\mathbf{H}\|_{F}^{2}=\sum_{\ell=1}^{d}\sigma_{\ell}^{2}(\mathbf{H})=\sum_{j=1}^{N}\left\|\mathbf{h}_{j}\right\|_{2}^{2}=N.$

Now suppose that $\sigma^{2}_{d-1}(\mathbf{H})\geq\Gamma$. Then we must have

$\displaystyle
N=\sum_{i=1}^{d}\sigma_{i}^{2}(\mathbf{H})\geq(d-1)\Gamma=(d-1)\frac{2N}{\tau(1+(K\mbox{{
mod }}2)/K)+2}$

which implies $\tau\geq 2(d-2)(1+(K\mbox{{ mod }}2)/K)^{-1}$, but this
contradicts the assumption on $\tau$. Therefore
$\sigma^{2}_{d-1}(\mathbf{H})<\Gamma$, so we can choose
$\mathcal{S}=\mbox{span}(\\{v_{d-1},v_{d}\\})$, which suffices to give the
result by the above argument.

We are now ready to show that at any critical point that is not globally
optimal, we can find a direction along which the Riemannian Hessian has a
strictly negative curvature at this point.

Recall $\mathbf{M}:=\tau\mathbf{W}^{\top}\mathbf{H}$ and $\mathbf{G}:=\nabla
g(\mathbf{M})$, as well as the definition of
$\mathbf{\alpha}\in\mathbb{R}^{K},\ \mathbf{\beta}\in\mathbb{R}^{N}$ in (21).
As mentioned at the beginning of Appendix B, we can write $\mathbf{H}$ as

$\displaystyle\mathbf{H}\;=\;\begin{bmatrix}\mathbf{H}^{1}&\mathbf{H}^{2}&\cdots&\mathbf{H}^{n}\end{bmatrix}\in\mathbb{R}^{d\times
N},\;\mathbf{H}^{i}\;=\;\begin{bmatrix}\mathbf{h}_{1,i}&\mathbf{h}_{2,i}&\cdots&\mathbf{h}_{K,i}\end{bmatrix}\in\mathbb{R}^{d\times
K},\;\forall\ i\in[N].$

As a final remark, the bilinear form of the Riemannian Hessian in (2.2) can be
written as

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]=\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]-\tau\sum_{k=1}^{K}\alpha_{k}\|\mathbf{\delta}_{W_{k}}\|_{2}^{2}-\tau\sum_{i=1}^{N}\beta_{i}\|\mathbf{\delta}_{H_{i}}\|_{2}^{2}$
(33)

where $\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]$ is
given in (2.2), and $\mathbf{\delta}_{W_{k}}$, $\mathbf{\delta}_{H_{i}}$ are
the $k$-th and $i$-th columns of $\mathbf{\Delta}_{\mathbf{W}}$ and
$\mathbf{\Delta}_{\mathbf{H}}$ respectively.

###### Proposition 2

Suppose $d>K$ and $\tau<2(d-2)(1+(K\mbox{{ mod }}2)/K)^{-1}$. For any critical
point $(\mathbf{W},\mathbf{H})\in\mathcal{C}$ that is not globally optimal,
there exists
$\mathbf{\Delta}=\left(\mathbf{\Delta}_{\mathbf{W}},\mathbf{\Delta}_{\mathbf{H}}\right)\in\mathrm{T}_{\mathbf{W}}\mathcal{OB}(d,K)\times\mathrm{T}_{\mathbf{H}}\mathcal{OB}(d,N)$
such that

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;<\;0.$
(34)

Proof We proceed by considering two separate cases for the value of
$\mathbf{\beta}$: $\beta_{i}=0$ for some $i\in[N]$, and $\beta_{i}\neq 0$ for
all $i\in[N]$.

Case 1: Suppose $\beta_{i}=0$ for some $i\in[N]$. In this case, by Lemma Lemma
C.4, we know that $\mathbf{W}=\mathbf{w}\mathbf{1}_{K}^{\top}$ for some
$\mathbf{w}\in\mathbb{S}^{d-1}$ and $\mathbf{\beta}=\mathbf{0}$. We have that
$\mathbf{M}=\tau\mathbf{1}_{K}\mathbf{w}^{\top}\mathbf{H}$, and so

$\displaystyle\mathbf{G}\;=\;-\frac{1}{N}\begin{bmatrix}\mathbf{A}&\cdots&\mathbf{A}\end{bmatrix}\in\mathbb{R}^{K\times
N},\quad\mathbf{A}\;=\;\mathbf{I}_{K}-\frac{1}{K}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}\in\mathbb{R}^{K\times
K}.$ (35)

For the the $i$-th column of $\mathbf{M}$, i.e. $\mathbf{m}_{i}$, we have the
Hessian

$\displaystyle\nabla^{2}\mathcal{L}_{\mathrm{CE}}(\mathbf{m}_{i},\mathbf{y}_{k})=\frac{1}{K}\mathbf{I}_{K}-\frac{1}{K^{2}}\mathbf{1}_{K}\mathbf{1}_{K}^{\top}=\frac{1}{K}\mathbf{A}$
(36)

Using Lemma C.5, choose $\mathbf{a}\in\mathbb{S}^{d-1}$ satisfying (32).
Additionally, choose a vector $\mathbf{u}\in\mathbb{R}^{K}$ with each entry
$u_{k}=(-1)^{k+1}$ (noting that $\sum_{k}u_{k}=K\mbox{ mod }2$). Now, we
construct the negative curvature direction
$\mathbf{\Delta}=(\mathbf{\Delta}_{\mathbf{W}},\mathbf{\Delta}_{\mathbf{H}})$
as

$\displaystyle\mathbf{\Delta}_{\mathbf{W}}=\mathbf{a}\mathbf{u}^{\top},\;\mathbf{\Delta}_{\mathbf{H}}=\begin{bmatrix}\mathbf{\Delta}_{\mathbf{H}^{1}}&\cdots\mathbf{\Delta}_{\mathbf{H}^{n}}\end{bmatrix}$

where

$\displaystyle\mathbf{\Delta}_{\mathbf{H}^{i}}=\mathbf{a}\mathbf{u}^{\top}-\mathbf{H}^{i}\operatorname{ddiag}(\mathbf{H}^{i\top}\mathbf{a}\mathbf{u}^{\top}),\;\forall
i\in[n].$

First, let $\mathbf{\delta}_{M_{i}}$ denote the $i$-th column of
$\mathbf{\Delta}_{\mathbf{M}}:=\bm{W}^{\top}\bm{\Delta}_{\bm{H}}+\bm{\Delta}_{\bm{W}}^{\top}\bm{H}$,
so that

$\displaystyle\mathbf{\delta}_{M_{i}}=(\mathbf{w}^{\top}\mathbf{\delta}_{H_{i}})\mathbf{1}_{K}+(\mathbf{h}_{i}^{\top}\mathbf{a})\mathbf{u}.$
(37)

Then from (16) and (36), we know that

$\displaystyle\nabla^{2}g(\mathbf{W}^{\top}\mathbf{H})\left[\tau\mathbf{\Delta}_{\mathbf{M}},\tau\mathbf{\Delta}_{\mathbf{M}}\right]=\frac{\tau^{2}}{NK}\sum_{i=1}^{N}\mathbf{\delta}_{M_{i}}^{\top}\mathbf{A}\mathbf{\delta}_{M_{i}}.$

Since $\mathbf{A}\mathbf{1}_{K}=\mathbf{0}$ and
$\mathbf{u}^{\top}\mathbf{A}\mathbf{u}=K-(K\mbox{{ mod }}2)/K$, by (37) we
have

$\displaystyle\nabla^{2}g(\mathbf{W}^{\top}\mathbf{H})\left[\tau\mathbf{\Delta}_{\mathbf{M}},\tau\mathbf{\Delta}_{\mathbf{M}}\right]$
$\displaystyle=\frac{\tau^{2}}{NK}\left(K-\frac{K\mbox{{ mod
}}2}{K}\right)\sum_{i=1}^{N}(\mathbf{h}_{i}^{\top}\mathbf{a})^{2}$
$\displaystyle=\frac{\tau^{2}}{NK}\left(K-\frac{K\mbox{{ mod
}}2}{K}\right)\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}.$

On the other hand, by (35) we have

$\displaystyle
2\tau\left<\mathbf{G},\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}}\right>$
$\displaystyle=-\frac{2\tau}{N}\sum_{i=1}^{n}\operatorname{trace}(\mathbf{A}\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{\Delta}_{\mathbf{H}^{i}})$
$\displaystyle=-\frac{2\tau}{N}\sum_{i=1}^{n}\operatorname{trace}\left(\mathbf{A}\mathbf{u}\mathbf{u}^{\top}\operatorname{diag}\left(1-(\mathbf{h}_{1,i}^{\top}\mathbf{a})^{2},\ldots,1-(\mathbf{h}_{K,i}^{\top}\mathbf{a})^{2}\right)\right)$
$\displaystyle=-\frac{2\tau}{N}\sum_{i=1}^{n}\mathbf{u}^{\top}\operatorname{diag}\left(1-(\mathbf{h}_{1,i}^{\top}\mathbf{a})^{2},\ldots,1-(\mathbf{h}_{K,i}^{\top}\mathbf{a})^{2}\right)\left(\mathbf{u}-\frac{K\mbox{{
mod }}2}{K}\mathbf{1}_{K}\right)$
$\displaystyle=-\frac{2\tau}{N}\sum_{i=1}^{n}\sum_{k=1}^{K}(1-(\mathbf{h}_{k,i}^{\top}\mathbf{a})^{2})u_{k}^{2}+(K\mbox{{
mod
}}2)\frac{2\tau}{NK}\sum_{i=1}^{n}\sum_{k=1}^{K}(1-(\mathbf{h}_{k,i}^{\top}\mathbf{a})^{2})u_{k}$
$\displaystyle\leq-\frac{2\tau}{N}\left(N-\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}\right)+(K\mbox{{
mod
}}2)\frac{2\tau}{NK}\left(N-\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}\right)$
$\displaystyle=-\frac{2\tau}{NK}\left(N-\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}\right)\left(K-(K\mbox{{
mod }}2)\right).$

Finally, the remaining term
$-\tau\sum_{k=1}^{K}\alpha_{k}\|\mathbf{\delta}_{W_{k}}\|_{2}^{2}-\tau\sum_{i=1}^{N}\beta_{i}\|\mathbf{\delta}_{H_{i}}\|_{2}^{2}$
in (33) vanishes, which is due to the fact that $\mathbf{\beta}=\mathbf{0}$
and

$\displaystyle\sum_{k=1}^{K}\alpha_{k}\|\mathbf{\delta}_{W_{k}}\|_{2}^{2}=\sum_{k=1}^{K}\alpha_{k}u_{k}^{2}=\sum_{k=1}^{K}\alpha_{k}=0,$

where the last equality follows by Lemma C.3 that
$\sum_{k=1}^{K}\alpha_{k}=\sum_{i=1}^{N}\beta_{i}=0$. Therefore, plugging both
bounds above into (33), we obtain

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]$
$\displaystyle\leq\;$ $\displaystyle\frac{\tau^{2}}{NK}\left(K-\frac{K\mbox{{
mod
}}2}{K}\right)\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}-\frac{2\tau}{NK}(N-\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2})(K-(K\mbox{{
mod }}2))$ $\displaystyle=\;$ $\displaystyle\frac{\tau(K-(K\mbox{{ mod
}}2))}{NK}\left(\tau\left[\frac{K^{2}-(K\mbox{{ mod }}2)}{K(K-(K\mbox{{ mod
}}2))}\right]\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}-2(N-\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2})\right)$
$\displaystyle=\;$ $\displaystyle\frac{\tau(K-(K\mbox{{ mod
}}2))}{NK}\left[\left(\tau[1+(K\mbox{{ mod
}}2)/K)]+2\right)\|\mathbf{H}^{\top}\mathbf{a}\|_{2}^{2}-2N\right]<0,$

where the last inequality follows by our choice of
$\mathbf{a}\in\mathbb{S}^{d-1}$ in Lemma C.5. Thus we obtain the desired
result in (34) for this case.

Case 2: Suppose $\beta_{i}\neq 0$ for all $i\in[N]$. Using the fact that
$d>K$, choose $\mathbf{a}\in\mathbb{S}^{d-1}$ such that
$\mathbf{W}^{\top}\mathbf{a}=\mathbf{0}$. By Lemma C.1, given that
$\mathbf{W}\mathbf{g}_{i}=\beta_{i}\mathbf{h}_{i}$ for all $i\in[N]$, we have

$\displaystyle\mathbf{a}^{\top}\mathbf{W}\mathbf{g}_{i}=\beta_{i}\mathbf{a}^{\top}\mathbf{h}_{i}=0,\quad\forall\;i\in[N].$

Thus, as $\beta_{i}\neq 0$ for all $i\in[N]$, this simply implies that
$\mathbf{H}^{\top}\mathbf{a}=\mathbf{0}$. Now using Lemma C.2, for any non-
optimal critical point $(\mathbf{W},\mathbf{H})$, there exists at least one
$k\in[K]$ or $i\in[N]$ such that either

$\displaystyle\alpha_{k}>-\sqrt{n}\|\mathbf{G}\|,\quad\text{or}\quad\beta_{i}>-{\|\mathbf{G}\|}/{\sqrt{n}}.$
(38)

Let $\mathbf{u}_{1}\in\mathbb{R}^{K}$ and $\mathbf{v}_{1}\in\mathbb{R}^{N}$ be
the left and right unit singular vectors associated with the leading singular
values of $\mathbf{G}$, respectively. In other words, we have

$\displaystyle\mathbf{u}_{1}^{\top}\mathbf{G}\mathbf{v}_{1}=\|\mathbf{G}\|.$
(39)

By letting $\mathbf{u}=-\mathbf{u}_{1}/\sqrt[4]{n},\
\mathbf{v}=\sqrt[4]{n}\mathbf{v}_{1}$, we construct the negative curvature
direction as

$\displaystyle\mathbf{\Delta}=\left(\mathbf{\Delta}_{\mathbf{W}},\mathbf{\Delta}_{\mathbf{H}}\right)=\left(\mathbf{a}\mathbf{u}^{\top},\
\mathbf{a}\mathbf{v}^{\top}\right).$ (40)

Since
$\mathbf{W}^{\top}\mathbf{a}=\mathbf{0},\mathbf{H}^{\top}\mathbf{a}=\mathbf{0}$,
we have

$\displaystyle\mathbf{W}^{\top}\mathbf{\Delta}_{\mathbf{H}}+\mathbf{\Delta}_{\mathbf{W}}^{\top}\mathbf{H}\;=\;\mathbf{W}^{\top}\mathbf{a}\mathbf{v}^{\top}+\mathbf{u}\mathbf{a}^{\top}\mathbf{H}\;=\;\mathbf{0},$

so that from (2.2) we have

$\displaystyle\nabla^{2}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;=\;$
$\displaystyle\nabla^{2}g(\mathbf{M})\left[\tau\left(\bm{W}^{\top}\bm{\Delta}_{\bm{H}}+\bm{\Delta}_{\bm{W}}^{\top}\bm{H}\right),\tau\left(\bm{W}^{\top}\bm{\Delta}_{\bm{H}}+\bm{\Delta}_{\bm{W}}^{\top}\bm{H}\right)\right]$
$\displaystyle+2\tau\left\langle\mathbf{G},\bm{\Delta}_{\bm{W}}^{\top}\bm{\Delta}_{\bm{H}}\right\rangle.$

Thus, from (33), combining all the above derivations we obtain

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]$
$\displaystyle\;=\;2\tau\left\langle\mathbf{G},\bm{\Delta}_{\bm{W}}^{\top}\bm{\Delta}_{\bm{H}}\right\rangle-\tau\sum_{k=1}^{K}\alpha_{k}\left\|\mathbf{\delta}_{W_{k}}\right\|_{2}^{2}-\tau\sum_{i=1}^{N}\beta_{i}\left\|\mathbf{\delta}_{H_{i}}\right\|_{2}^{2}.$
$\displaystyle\;=\;-2\tau\left\langle\mathbf{G},\mathbf{u}_{1}\mathbf{v}_{1}^{\top}\right\rangle-\tau\left(\sum_{k=1}^{K}\frac{\alpha_{k}u_{1,k}^{2}}{\sqrt{n}}+\sum_{i=1}^{N}\sqrt{n}\beta_{i}v^{2}_{1,i}\right)$
$\displaystyle\;=\;\tau\left(-2\|\mathbf{G}\|-\sum_{k=1}^{K}\frac{\alpha_{k}u_{1,k}^{2}}{\sqrt{n}}-\sum_{i=1}^{N}\sqrt{n}\beta_{i}v^{2}_{1,i}\right)$

where the last equality follows from (39). On the other hand, by Lemma C.2,
the fact we derived in (38) that there exists $k\in[K]$ such that
$\alpha_{k}>-\sqrt{n}\|\mathbf{G}\|$ or there exists $i\in[N]$ such that
$\beta_{i}>-{\|\mathbf{G}\|}/{\sqrt{n}}$, and that
$\|\mathbf{u}_{1}\|_{2}=\|\mathbf{v}_{1}\|_{2}=1$, we obtain

$\displaystyle-\sum_{k=1}^{K}\frac{\alpha_{k}u_{1,k}^{2}}{\sqrt{n}}-\sum_{i=1}^{N}\sqrt{n}\beta_{i}v^{2}_{1,i}\;<\;\left\|\mathbf{G}\right\|\left(\sum_{k=1}^{K}u_{1,k}^{2}+\sum_{i=1}^{N}v_{1,i}^{2}\right)\;=\;2\left\|\mathbf{G}\right\|.$

Therefore, we have

$\displaystyle\operatorname{Hess}f(\mathbf{W},\mathbf{H})[\mathbf{\Delta},\mathbf{\Delta}]\;<\;\tau\left(-2\|\mathbf{G}\|+2\|\mathbf{G}\|\right)\;=\;0,$

as desired.

Proof [Proof of Theorem 3.2] Let
$(\mathbf{W},\mathbf{H})\in\mathcal{OB}(d,K)\times\mathcal{OB}(d,N)$ be a
local minimizer of Problem (4). Suppose that it is not a global minimizer.
This implies $(\mathbf{W},\mathbf{H})$ is a critical point that is not a
global minimizer. According to 2, the Riemannian Hessian at
$(\mathbf{W},\mathbf{H})$ has negative curvature. This contradicts with the
fact that $(\mathbf{W},\mathbf{H})$ is a local minimizer. Thus, we concludes
that any local minimizer of Problem (4) is a global minimizer in Theorem 3.1.
Moreover, according to 2, any critical point of Problem (4) that is not a
local minimizer is a Riemmannian strict saddle point with negative curvature.