# Neural collapse with unconstrained features

Dustin G. Mixon, Hans Parshall, Jianzong Pi  Dustin G. Mixon111Department of
Mathematics, The Ohio State University, Columbus, OH 222Translational Data
Analytics Institute, The Ohio State University, Columbus, OH Hans
Parshall333Department of Mathematics, Western Washington University,
Bellingham, WA Jianzong Pi444Department of Electrical and Computer
Engineering, The Ohio State University, Columbus, OH

###### Abstract

Neural collapse is an emergent phenomenon in deep learning that was recently
discovered by Papyan, Han and Donoho. We propose a simple unconstrained
features model in which neural collapse also emerges empirically. By studying
this model, we provide some explanation for the emergence of neural collapse
in terms of the landscape of empirical risk.

##  1 Introduction

Consider the task of learning a function $\mathcal{X}\to[C]$, where
$\mathcal{X}\subseteq\mathbb{R}^{d}$ represents a space of signals and
$[C]:=\\{1,\ldots,C\\}$ represents a set of $C\in\mathbb{N}$ labels. Given a
labeled training set $\\{(x_{i},c_{i})\\}_{i\in\mathcal{I}}$ in
$\mathcal{X}\times[C]$, one might seek both a feature map
$h\colon\mathcal{X}\to\mathbb{R}^{p}$ and a linear classifier
$z\mapsto\arg\max(Wz+b)$ with $W\in\mathbb{R}^{C\times p}$ and
$b\in\mathbb{R}^{C}$ for which empirical risk is small:

$R_{e}(h,W,b):=\sum_{i\in\mathcal{I}}\mathcal{L}(Wh(x_{i})+b,y_{i}).$

Here, $\mathcal{L}\colon\mathbb{R}^{C}\times\mathbb{R}^{C}\to\mathbb{R}^{+}$
denotes some loss function and $y_{i}\in\mathbb{R}^{C}$ denotes the one-hot
vector representation $e_{c_{i}}$ of the label $c_{i}$. To accomplish this, it
is common to take $h$ to reside in some parameterized family
$\mathcal{H}=\\{h_{\theta}:\theta\in\Theta\\}$, such as a family of neural
networks, and then locally minimize empirical risk over $(\theta,W,b)$. Such
methods have revolutionized classification in various domains, notably image
domains [8], but this level of performance is largely under-explained by
existing theory.

Papyan, Han and Donoho [11] recently observed several emergent phenomena in
the terminal phase of training, that is, when the above empirical risk is
minimized even after the classifier $x\mapsto\arg\max(Wh(x)+b)$ interpolates
the training set. They focused on the balanced case in which there exists
$N\in\mathbb{N}$ such that for each $c\in[C]$, it holds that

$|\\{i\in\mathcal{I}:c_{i}=c\\}|=N,$

i.e., $|\mathcal{I}|=CN$. For several common families $\mathcal{H}$ of neural
networks, locally minimizing empirical risk results in $(h,W,b)$ that
satisfies several properties, which are collectively referred to as neural
collapse:

(NC1) Variability collapse. For each $c\in[C]$, there exists
$\mu_{c}\in\mathbb{R}^{p}$ such that for every $i\in\mathcal{I}$ with
$y_{i}=e_{c}$, it holds that $h(x_{i})=\mu_{c}$.

(NC2) Simplex equiangular tight frame structure. Put
$\mu_{G}:=\frac{1}{C}\sum_{c=1}^{C}\mu_{c}$. Then

$\|\mu_{c}-\mu_{G}\|_{2}=\|\mu_{c^{\prime}}-\mu_{G}\|_{2}$

for every $c,c^{\prime}$. Furthermore, let the $c$th column of
$M\in\mathbb{R}^{p\times C}$ be
$\tilde{\mu}_{c}:=\frac{\mu_{c}-\mu_{G}}{\|\mu_{c}-\mu_{G}\|_{2}}$. Then

$M^{\top}M=\frac{C}{C-1}I_{C}-\frac{1}{C-1}1_{C}1_{C}^{\top}.$

(NC3) Self-duality. $\frac{W^{\top}}{\|W\|_{F}}=\frac{M}{\|M\|_{F}}$.

(NC4) Equivalence to nearest class center.
$\arg\max_{c}(Wz+b)_{c}=\arg\min_{c}\|z-\mu_{c}\|_{2}$.

In general, an equiangular tight frame (ETF) is any tuple
$\\{v_{i}\\}_{i=1}^{n}$ of unit vectors in $\mathbb{C}^{d}$ for which there
exist $\alpha,\beta\geq 0$ such that

$\sum_{i=1}^{n}v_{i}v_{i}^{*}=\alpha I_{d},\qquad|\langle
v_{i},v_{j}\rangle|^{2}=\beta\qquad\forall i,j\in[n],~{}i\neq j.$

It is natural to identify each $v_{i}v_{i}^{*}$ with the corresponding point
in complex projective space. ETFs were introduced in [15] as convenient
maximizers of minimum pairwise distance in this space. As optimal projective
codes, ETFs find applications in multiple description coding [15], digital
fingerprinting [9], compressed sensing [2], and quantum state tomography [13].
These applications have motivated a flurry of recent work to discover various
infinite families of ETFs; see [6] for a living survey. The simplex ETF arises
from one of the simplest constructions: take $\\{v_{i}\\}_{i=1}^{n}$ to be the
$n=d+1$ vertices of an origin-centered regular simplex in $\mathbb{R}^{d}$. We
note that negating any vector in an ETF produces another ETF (as ETFs are
fundamentally a projective objects), whereas property (NC2) is not invariant
to such an operation. The regular simplex also emerges as an optimal spherical
code by virtue of achieving equality in Rankin’s simplex bound [12]. In
addition, $C$ vertices of a regular simplex in $\mathbb{R}^{p}$ induce Voronoi
cells that partition $\mathbb{R}^{p}$ into $C$ isometric cones that are highly
symmetric, and this geometry plays a fundamental role in (NC4).

The rich geometric structure of neural collapse has some advantages in the
context of machine learning. In particular, if the feature map $h$
generalizes, then neural collapse helps the classifier
$x\mapsto\arg\max(Wh(x)+b)$ to also generalize. Indeed, consider a member $x$
of the test set corresponding to class $c$. If the corresponding point $h(x)$
in the feature domain is a perturbation of $\mu_{c}$, then $h(x)$ is most
likely to reside in the Voronoi cell containing $\mu_{c}$ if the class means
form the vertices of a regular simplex. Papyan, Han and Donoho [11] provide
some theoretical justification along these lines.

While the notion of neural collapse appears useful in both theory and
practice, it remains unexplained why neural collapse emerges from empirical
risk minimization. This paper provides some explanation for this emergence. In
the next section, we propose a simple unconstrained features model in which a
strong notion of neural collapse empirically emerges. For this model, Section
3 identifies an invariant subspace of the gradient descent dynamical system
that encourages convergence to this strong notion of neural collapse. We
conclude in Section 4 with a brief discussion.

##  2 Unconstrained features and strong neural collapse

Empirically, neural collapse is a phenomenon that occurs in the terminal phase
of training, in which the trained classifier interpolates the training set.
This behavior is only feasible when the feature maps in $\mathcal{H}$
restricted to the training set $\\{x_{i}\\}_{i\in\mathcal{I}}$ form a high-
dimensional subset of $(\mathbb{R}^{p})^{\mathcal{I}}$. Indeed, the
observations by Papyan, Han and Donoho [11] were made for over-parameterized
families of neural networks. In this paper, we take this to an extreme of
sorts by considering the following unconstrained features model:

$\mathcal{H}:=\\{h\colon\mathcal{X}\to\mathbb{R}^{p}\\},\qquad\mathcal{X}:=\\{x_{i}:i\in\mathcal{I}\\}\subseteq\mathbb{R}^{d}.$

Under the reasonable assumption that $x_{i}\neq x_{j}$ whenever $i\neq j$,
then we may represent $\mathcal{H}$ by the matrix space $\mathbb{R}^{p\times
CN}$; explicitly, we identify $\mathcal{I}$ with $[CN]$, and the $i$th column
of $H\in\mathbb{R}^{p\times CN}$ is given by $h(x_{i})$. Recalling that
$\\{c_{i}\\}_{i\in\mathcal{I}}$ consists of $N$ copies of each member of
$[C]$, we may similarly represent $\\{y_{i}\\}_{i\in\mathcal{I}}$ as the
columns of the matrix $I_{C}\otimes 1_{N}^{\top}$. Note that feature maps in
the unconstrained features model are only defined over the training set, while
feature maps that are typically trained in practice are defined over all of
$\mathbb{R}^{d}$. Of course, such a choice of $\mathcal{H}$ has no hope of
generalizing to a test set, but we will find that it facilitates the study of
neural collapse. For simplicity, we consider the loss defined by
$\mathcal{L}(u,v):=\frac{1}{2}\|u-v\|_{2}^{2}$, in which case empirical risk
reduces to

$R_{e}(H,W,b)=\frac{1}{2}\|WH+b1_{CN}^{\top}-I_{C}\otimes
1_{N}^{\top}\|_{F}^{2},\qquad H\in\mathbb{R}^{p\times
CN},~{}W\in\mathbb{R}^{C\times p},~{}b\in\mathbb{R}^{C}.$

Following common practice in deep learning [10], we initialize gradient
descent close to the origin. As illustrated in Figure 1, the following
phenomena emerges from this local optimization, which we collectively refer to
as strong neural collapse:

$\displaystyle WW^{\top}$
$\displaystyle=\sqrt{N}(I_{C}-\frac{1}{C}1_{C}1_{C}^{\top}),$ (SNC1)
$\displaystyle H$ $\displaystyle=\frac{1}{\sqrt{N}}(W\otimes 1_{N})^{\top},$
(SNC2) $\displaystyle b$ $\displaystyle=\frac{1}{C}1_{C}.$ (SNC3)

The following lemma establishes that the points exhibiting strong neural
collapse form a subset of the global minimizers of empirical risk. In the next
section, we show how the emergence of strong neural collapse is an artifact of
the optimization landscape.

Figure 1:  The emergence of strong neural collapse. Run gradient decent to
minimize $R_{e}(H,W,b)$ for $C=N=3$ and $p=15$, initializing at a random
choice of $H_{0}$ and $W_{0}$ with $\|H_{0}\|_{F}=\|W_{0}\|_{F}=\varepsilon$
and $b_{0}=0$. At each iteration, quantify the error in (SNC1) by
$\|WW^{\top}-\sqrt{N}(I_{C}-\frac{1}{C}1_{C}1_{C}^{\top})\|_{F}$ (plotted on
the left), the relative error in (SNC2) by $\|H-\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top}\|_{F}/\|H\|_{F}$ (plotted in the middle), and the error in
(SNC3) by $\|b-\frac{1}{C}1_{C}\|_{2}$ (plotted on the right). Apparently, the
limit point of gradient descent approaches strong neural collapse as the
initialization approaches the origin.

######  Lemma 1.

  * (a)

If $(H,W,b)$ exhibits strong neural collapse, then $R_{e}(H,W,b)=0$.

  * (b)

Strong neural collapse implies neural collapse.

###### Proof.

For (a), we have

$\displaystyle WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top}$
$\displaystyle=W(\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top})+\frac{1}{C}1_{C}1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top}$
$\displaystyle=(\frac{1}{\sqrt{N}}WW^{\top}+\frac{1}{C}1_{C}1_{C}^{\top}-I_{C})\otimes
1_{N}^{\top},$

which vanishes by (SNC1). Next, we consider (b). Letting $w_{c}^{\top}$ denote
the $c$th row of $W$, then (NC1) follows from (SNC2) by taking
$\mu_{c}:=\frac{1}{\sqrt{N}}w_{c}$. For (NC2), observe that (SNC1) implies
$\|W^{\top}1_{C}\|_{2}^{2}=1_{C}^{\top}WW^{\top}1_{C}=0$, and so

$\mu_{G}=\frac{1}{C}\sum_{c=1}^{C}\mu_{c}=\frac{1}{C}\sum_{c=1}^{C}\frac{1}{\sqrt{N}}w_{c}=\frac{1}{C\sqrt{N}}W^{\top}1_{C}=0.$

Then for every $c\in[C]$, it holds that

$\|\mu_{c}-\mu_{G}\|_{2}^{2}=\|\mu_{c}\|_{2}^{2}=\|\frac{1}{\sqrt{N}}w_{c}\|_{2}^{2}=\frac{1}{N}(WW^{\top})_{cc}=\frac{1}{\sqrt{N}}(1-\frac{1}{C}).$

Furthermore,
$\tilde{\mu}_{c}=\frac{\mu_{c}}{\|\mu_{c}\|_{2}}=\frac{w_{c}}{\|w_{c}\|_{2}}$,
and so

$(M^{\top}M)_{cc^{\prime}}=\langle\tilde{\mu}_{c},\tilde{\mu}_{c^{\prime}}\rangle=\langle\frac{w_{c}}{\|w_{c}\|_{2}},\frac{w_{c^{\prime}}}{\|w_{c^{\prime}}\|_{2}}\rangle=\frac{(WW^{\top})_{cc^{\prime}}}{\sqrt{N}(1-\frac{1}{C})}=(\frac{C}{C-1}I_{C}-\frac{1}{C-1}1_{C}1_{C}^{\top})_{cc^{\prime}}.$

Since $\|w_{c}\|_{2}^{2}=\sqrt{N}(1-\frac{1}{C})$ is constant over $c\in[C]$,
we also have $M\propto W^{\top}$, from which (NC3) follows. It remains to
verify (NC4). For this, we first apply (SNC3) and the fact that
$\mu_{c}=\frac{1}{\sqrt{N}}w_{c}$ to obtain

$\arg\max_{c}(Wz+b)_{c}=\arg\max_{c}(Wz)_{c}=\arg\max_{c}\langle
z,w_{c}\rangle=\arg\max_{c}\langle z,\mu_{c}\rangle.$

Finally, (SNC1) implies $\|\mu_{c}\|_{2}=\|\mu_{c^{\prime}}\|_{2}$ for every
$c,c^{\prime}$, and so

$\arg\max_{c}\langle z,\mu_{c}\rangle=\arg\min_{c}(\|z\|_{2}^{2}-2\langle
z,\mu_{c}\rangle+\|\mu_{c}\|_{2}^{2})=\arg\min_{c}\|z-\mu_{c}\|_{2}^{2}.\qed$

##  3 The effect of empirical risk minimization

Write $Z:=(H,W,b)$ and consider the gradient flow ordinary differential
equation

$Z^{\prime}(t)=-\nabla R_{e}(Z(t)).$

This serves as a model for gradient descent with small learning rate that has
been used to analyze the training of neural networks [1, 3, 4, 5, 14]. To
analyze gradient flow, we start by computing the gradient:

$\nabla_{H}R_{e}=W^{\top}A,\quad\nabla_{W}R_{e}=AH^{\top},\quad\nabla_{b}R_{e}=A1_{CN},\quad
A:=WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top}.$

Notably, the resulting gradient flow differential equation is nonlinear. In
order to analyze the initial behavior of our trajectory, we consider a
modification in which $A$ is replaced by
$\tilde{A}:=b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top}$. By partially decoupling
$(H,W)$ and $b$, this modification facilitates analysis, while serving as an
good approximation in the regime where $(H,W)$ is small.

######  Theorem 2.

The solution to the ordinary differential equation

$\displaystyle H^{\prime}(t)$ $\displaystyle=-W(t)^{\top}\tilde{A}(t),\qquad
W^{\prime}(t)=-\tilde{A}(t)H(t)^{\top},$ $\displaystyle b^{\prime}(t)$
$\displaystyle=-\tilde{A}(t)1_{CN},\qquad\tilde{A}(t):=b(t)1_{CN}^{\top}-I_{C}\otimes
1_{N}^{\top}$

with initial condition $H(0)=H_{0}$, $W(0)=W_{0}$, $b(0)=0$ satisfies

$\|(H(t),W(t))-e^{\sqrt{N}t}\cdot\Pi_{T}(H_{0},W_{0})\|_{E}\leq
e^{1/(C\sqrt{N})}\cdot\|\Pi_{T^{\perp}}(H_{0},W_{0})\|_{E},\quad
b(t)=(\frac{1-e^{-CNt}}{C})1_{C}$

for all $t\geq 0$, where $\|(H,W)\|_{E}^{2}:=\|H\|_{F}^{2}+\|W\|_{F}^{2}$ and
$\Pi_{T}$ denotes orthogonal projection onto the subspace

$T:=\\{(H,W):H=\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top},~{}1_{C}^{\top}W=0\\}.$

###### Proof.

We start by solving for $b(\cdot)$, which is governed by the differential
equation

$b^{\prime}(t)=-\tilde{A}(t)1_{CN}=(I_{C}\otimes
1_{N}^{\top}-b(t)1_{CN}^{\top})1_{CN}=N(1_{C}-Cb(t)).$

Since $b(0)=0$, we may write $b(t)=\beta(t)1_{C}$, in which case
$\beta^{\prime}(t)=N(1-C\beta(t))$, and so $\beta(t)=\frac{1-e^{-CNt}}{C}$, as
claimed. Next, writing $U=(H,W)$, then the unsolved portion of our system is
given by

$U^{\prime}(t)=L_{t}(U(t)),\quad
L_{t}(H,W):=(W^{\top}M_{t},M_{t}H^{\top}),\quad M_{t}:=(I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top}.$

First, we observe that each $L_{t}$ is self-adjoint:

$\displaystyle\langle L_{t}(H,W),(\tilde{H},\tilde{W})\rangle_{E}$
$\displaystyle=\langle(W^{\top}M_{t},M_{t}H^{\top}),(\tilde{H},\tilde{W})\rangle_{E}$
$\displaystyle=\operatorname{tr}((W^{\top}M_{t})^{\top}\tilde{H})+\operatorname{tr}((M_{t}H^{\top})^{\top}\tilde{W})$
$\displaystyle=\operatorname{tr}(M_{t}^{\top}W\tilde{H})+\operatorname{tr}(HM_{t}^{\top}\tilde{W})$
$\displaystyle=\operatorname{tr}(\tilde{H}M_{t}^{\top}W)+\operatorname{tr}(M_{t}^{\top}\tilde{W}H)$
$\displaystyle=\operatorname{tr}((M_{t}\tilde{H}^{\top})^{\top}W)+\operatorname{tr}((\tilde{W}^{\top}M_{t})^{\top}H)$
$\displaystyle=\langle(\tilde{W}^{\top}M_{t},M_{t}\tilde{H}^{\top}),(H,W)\rangle_{E}=\langle(H,W),L_{t}(\tilde{H},\tilde{W})\rangle_{E}.$

Next, we claim that $\\{L_{t}\\}_{t\geq 0}$ are simultaneously diagonalizable
over five eigenspaces:

$\displaystyle E_{1}^{\epsilon}$
$\displaystyle:=\\{(H,W):H=\epsilon\cdot\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top},~{}1_{C}^{\top}W=0\\},$ $\displaystyle E_{2}^{\epsilon}$
$\displaystyle:=\\{(H,W):H=\epsilon\cdot
z1_{CN}^{\top},~{}W=\sqrt{N}1_{C}z^{\top},~{}z\in\mathbb{R}^{p}\\},$
$\displaystyle E_{3}$ $\displaystyle:=\\{(H,W):(I_{C}\otimes
1_{N}^{\top})H^{\top}=0,~{}W=0\\},$

where $\epsilon\in\\{\pm\\}$. To see this, first note that

$\operatorname{dim}E_{1}^{\epsilon}=p(C-1),\qquad\operatorname{dim}E_{2}^{\epsilon}=p,\qquad\operatorname{dim}E_{3}=pC(N-1).$

Since these dimensions sum to $pC(N+1)=\operatorname{dim}(\mathbb{R}^{p\times
CN}\oplus\mathbb{R}^{C\times p})$, it suffices to show that each nonzero
member of each claimed eigenspace is an eigenvector, and that the claimed
eigenspaces have distinct eigenvalues. First, suppose $(H,W)\in
E_{1}^{\epsilon}$. Then

$\displaystyle W^{\top}M_{t}$ $\displaystyle=W^{\top}((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})=W^{\top}(I_{C}\otimes
1_{N}^{\top})=W^{\top}\otimes 1_{N}^{\top}=\epsilon\sqrt{N}\cdot H,$
$\displaystyle M_{t}H^{\top}$ $\displaystyle=((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})(\epsilon\cdot\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top})^{\top}$
$\displaystyle=\epsilon\cdot\frac{1}{\sqrt{N}}\cdot((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{C}^{\top}\otimes 1_{N}^{\top})(W\otimes
1_{N})=\epsilon\sqrt{N}\cdot W,$

i.e., $(H,W)$ is an eigenvector of $L_{t}$ with eigenvalue $\epsilon\sqrt{N}$.
Next, suppose $(H,W)\in E_{2}^{\epsilon}$. Then

$\displaystyle W^{\top}M_{t}$
$\displaystyle=(\sqrt{N}1_{C}z^{\top})^{\top}((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})$
$\displaystyle=\sqrt{N}z1_{C}^{\top}((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})$
$\displaystyle=\sqrt{N}z(1_{C}^{\top}\otimes
1_{N}^{\top}-C\beta(t)1_{CN}^{\top})=\sqrt{N}(1-C\beta(t))\cdot
z1_{CN}^{\top}=\epsilon\sqrt{N}(1-C\beta(t))\cdot H,$ $\displaystyle
M_{t}H^{\top}$ $\displaystyle=((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})(\epsilon\cdot
z1_{CN}^{\top})^{\top}$ $\displaystyle=\epsilon\cdot((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})1_{CN}z^{\top}$
$\displaystyle=\epsilon\cdot((I_{C}\otimes 1_{N}^{\top})(1_{C}\otimes
1_{N})-\beta(t)1_{C}1_{CN}^{\top}1_{CN})z^{\top}$
$\displaystyle=\epsilon\cdot(N1_{C}-CN\beta(t)1_{C})z^{\top}=\epsilon
N(1-C\beta(t))\cdot 1_{C}z^{\top}=\epsilon\sqrt{N}(1-C\beta(t))\cdot W,$

i.e., $(H,W)$ is an eigenvector of $L_{t}$ with eigenvalue
$\epsilon\sqrt{N}(1-C\beta(t))$. Finally, suppose $(H,W)\in E_{3}$. Then
$W^{\top}M_{t}=0$ and

$M_{t}H^{\top}=((I_{C}\otimes
1_{N}^{\top})-\beta(t)1_{C}1_{CN}^{\top})H^{\top}=-\beta(t)1_{C}1_{CN}^{\top}H^{\top}=-\beta(t)1_{C}1_{C}^{\top}(I_{C}\otimes
1_{N}^{\top})H^{\top}=0,$

i.e., $(H,W)$ is an eigenvector of $L_{t}$ with eigenvalue $0$. Overall,
letting $\Pi_{i}^{\epsilon}$ denote orthogonal projection onto
$E_{i}^{\epsilon}$, we have the spectral decomposition

$L_{t}=\sqrt{N}\Big{(}\Pi_{1}^{+}-\Pi_{1}^{-}+(1-C\beta(t))\Pi_{2}^{+}-(1-C\beta(t))\Pi_{2}^{-}\Big{)}.$

Finally, we solve the differential equation $U^{\prime}(t)=L_{t}(U(t))$ by
finding the orthogonal projection of $U(t)$ onto each eigenspace of $L_{t}$.
First,
$\Pi_{1}^{\epsilon}U^{\prime}(t)=\epsilon\sqrt{N}\Pi_{1}^{\epsilon}U(t)$, and
so

$\Pi_{1}^{\epsilon}U(t)=e^{\epsilon\sqrt{N}t}\Pi_{1}^{\epsilon}U(0).$

Next,
$\Pi_{2}^{\epsilon}U^{\prime}(t)=\epsilon\sqrt{N}(1-C\beta(t))\Pi_{2}^{\epsilon}U(t)$,
and so

$\Pi_{2}^{\epsilon}U(t)=f_{\epsilon}(t)\Pi_{2}^{\epsilon}U(0),\qquad
f_{\epsilon}(t):=\operatorname{exp}(\frac{\epsilon}{C\sqrt{N}}(1-e^{-CNt})).$

An application of the Pythagorean theorem then gives

$\displaystyle\|U(t)-e^{\sqrt{N}t}\Pi_{1}^{+}U(0)\|_{E}^{2}$
$\displaystyle=\|e^{-\sqrt{N}t}\Pi_{1}^{-}U(0)+f_{+}(t)\Pi_{2}^{+}U(0)+f_{-}(t)\Pi_{2}^{-}U(0)\|_{E}^{2}$
$\displaystyle=e^{-2\sqrt{N}t}\|\Pi_{1}^{-}U(0)\|_{E}^{2}+f_{+}(t)^{2}\|\Pi_{2}^{+}U(0)\|_{E}^{2}+f_{-}(t)^{2}\|\Pi_{2}^{-}U(0)\|_{E}^{2}$
$\displaystyle\leq\|\Pi_{1}^{-}U(0)\|_{E}^{2}+e^{2/(C\sqrt{N})}\|\Pi_{2}^{+}U(0)\|_{E}^{2}+\|\Pi_{2}^{-}U(0)\|_{E}^{2}$
$\displaystyle\leq e^{2/(C\sqrt{N})}\|(I-\Pi_{1}^{+})U(0)\|_{E}^{2}.$

The result then follows by observing that $T=E_{1}^{+}$. ∎

Theorem 2 indicates that gradient flow initially magnifies the portion of
$(H_{0},W_{0})$ that resides in $T$ while sending $b$ to $\frac{1}{C}1_{C}$
along $\operatorname{span}\\{1_{C}\\}$. As a consequence, the initial
trajectory of $(H,W,b)$ approximately travels along the subspace

$S:=\Big{\\{}(H,W,b):H=\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top},~{}1_{C}^{\top}W=0,~{}b\in\operatorname{span}\\{1_{C}\\}\Big{\\}}.$

In fact, Figure 2 illustrates that (empirically) the full trajectory
approximately travels along this subspace. As the following result
demonstrates, $S$ is an invariant subspace of our differential equation that
encourages convergence to points that exhibit strong neural collapse.

Figure 2:  Gradient descent maintains small distance from the invariant
subspace $S$. Run gradient decent to minimize $R_{e}(H,W,b)$ for $C=N=3$ and
$p=15$, initializing at a random choice of $H_{0}$ and $W_{0}$ with
$\|H_{0}\|_{F}=\|W_{0}\|_{F}=\varepsilon$ and $b_{0}=0$. At each iteration,
quantify the relative distance from $S$ by $\|Z-\Pi_{S}Z\|_{E}/\|Z\|_{E}$,
where $\|(H,W,b)\|_{E}^{2}:=\|H\|_{F}^{2}+\|W\|_{F}^{2}+\|b\|_{2}^{2}$.

######  Theorem 3.

Select $Z_{0}:=(H_{0},W_{0},b_{0})\in S$ and consider $Z(\cdot)$ such that

$Z^{\prime}(t)=-\nabla R_{e}(Z(t)),\qquad Z(0)=Z_{0}.$

  * (a)

$Z(t)\in S$ for every $t\geq 0$.

  * (b)

Select $\alpha\in\mathbb{R}$ such that $b_{0}=\frac{1-\alpha}{C}1_{C}$. Then
$(H(\cdot),W(\cdot),b(\cdot))=Z(\cdot)$ is given by

$\displaystyle W^{\prime}(t)$
$\displaystyle=\sqrt{N}W(t)-W(t)W(t)^{\top}W(t),\qquad W(0)=W_{0},$
$\displaystyle H(t)$ $\displaystyle=\frac{1}{\sqrt{N}}(W(t)\otimes
1_{N})^{\top},\qquad b(t)=\frac{1-\alpha e^{-CNt}}{C}1_{C}.$

  * (c)

Let $\Pi$ denote orthogonal projection onto $\operatorname{ker}W_{0}^{\top}$.
Then

$\lim_{t\to\infty}W(t)W(t)^{\top}=\sqrt{N}(I-\Pi).$

In particular, if $\operatorname{rank}W_{0}\geq C-1$, then
$\lim_{t\to\infty}Z(t)$ exhibits strong neural collapse.

###### Proof.

We start by verifying (a). Since $S$ is a subspace, it suffices to show that
$Z:=(H,W,b)\in S$ implies $\nabla R_{e}(Z)\in S$. First, the constraint
$1_{C}^{\top}W=0$ implies

$H1_{CN}=\frac{1}{\sqrt{N}}(W\otimes 1_{N})^{\top}(1_{C}\otimes
1_{N})=\frac{1}{\sqrt{N}}(W^{\top}1_{C}\otimes 1_{N}^{\top}1_{N})=0.$

This then implies

$\displaystyle\nabla_{b}R_{e}(Z)$
$\displaystyle=(WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top})1_{CN}$
$\displaystyle=CNb-(I_{C}\otimes 1_{N}^{\top})(1_{C}\otimes
1_{N})=CNb-N1_{C}\in\operatorname{span}\\{1_{C}\\}.$ (1)

In addition, we have

$\displaystyle 1_{C}^{\top}\nabla_{W}R_{e}(Z)$
$\displaystyle=1_{C}^{\top}(WH+b1_{CN}^{\top}-I_{C}\otimes
1_{N}^{\top})H^{\top}$ $\displaystyle=-(1_{C}^{\top}\otimes 1)(I_{C}\otimes
1_{N}^{\top})H^{\top}=-1_{CN}^{\top}H^{\top}=0.$

It remains to verify that
$\nabla_{H}R_{e}(Z)=\frac{1}{\sqrt{N}}(\nabla_{W}R_{e}(Z)\otimes
1_{N})^{\top}$. To this end, writing $b=\beta 1_{C}$ with $\beta\in\mathbb{R}$
gives $W^{\top}b=\beta W^{\top}1_{C}=0$. As such, we have

$\displaystyle\nabla_{H}R_{e}(Z)$
$\displaystyle=W^{\top}(WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top})$
$\displaystyle=W^{\top}WH-W^{\top}(I_{C}\otimes 1_{N}^{\top})$
$\displaystyle=W^{\top}W\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top}-W^{\top}(I_{C}\otimes
1_{N}^{\top})=W^{\top}(\frac{1}{\sqrt{N}}WW^{\top}-I_{C})\otimes
1_{N}^{\top}.$

Next,

$\displaystyle\nabla_{W}R_{e}(Z)$
$\displaystyle=(WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top})H^{\top}$
$\displaystyle=WHH^{\top}-(I_{C}\otimes 1_{N}^{\top})H^{\top}$
$\displaystyle=\frac{1}{N}W(W\otimes 1_{N})^{\top}(W\otimes
1_{N})-\frac{1}{\sqrt{N}}(I_{C}\otimes 1_{N}^{\top})(W\otimes 1_{N})$
$\displaystyle=WW^{\top}W-\sqrt{N}W,$ (2)

and so $\frac{1}{\sqrt{N}}(\nabla_{W}R_{e}(Z)\otimes
1_{N})^{\top}=\nabla_{H}R_{e}(Z)$, as desired.

For (b), we see from (1) that

$b^{\prime}(t)=-\nabla_{b}R_{e}(Z(t))=N(1_{C}-Cb(t)).$

Writing $b(t)=\beta(t)1_{C}$, then $\beta^{\prime}(t)=N(1-C\beta(t))$, and so
$\beta(t)=\frac{1-\alpha e^{-CNt}}{C}$, as claimed. Also, we see from (2) that

$W^{\prime}(t)=-\nabla_{W}R_{e}(Z(t))=\sqrt{N}W(t)-W(t)W(t)^{\top}W(t).$

The expression for $H(t)$ follows from the fact that $Z(t)\in S$.

For (c), consider $G(t):=W(t)W(t)^{\top}$. Then

$\displaystyle G^{\prime}(t)$
$\displaystyle=W^{\prime}(t)W(t)^{\top}+W(t)W^{\prime}(t)^{\top}$
$\displaystyle=(\sqrt{N}W(t)-W(t)W(t)^{\top}W(t))W(t)^{\top}+W(t)(\sqrt{N}W(t)-W(t)W(t)^{\top}W(t))^{\top}$
$\displaystyle=2\sqrt{N}W(t)W(t)^{\top}-2W(t)W(t)^{\top}W(t)W(t)^{\top}$
$\displaystyle=2\sqrt{N}G(t)-2G(t)^{2}.$

Since $G^{\prime}(t)$ and $G(t)$ are simultaneously diagonalizable, it follows
that $G(\cdot)$ takes the form $G(t)=\sum_{i}\lambda_{i}(t)\Pi_{i}$, where
$\\{\Pi_{i}\\}$ denote orthogonal projections onto the eigenspaces of $G(0)$.
Furthermore, we have

$\lambda_{i}^{\prime}(t)=2\lambda_{i}(t)(\sqrt{N}-\lambda_{i}(t)).$

Since $G(0)$ is positive semidefinite, we have $\lambda_{i}(0)\geq 0$ for
every $i$. If $\lambda_{i}(0)=0$, then $\lambda_{i}(t)=0$ for all $t\geq 0$.
If $\lambda_{i}(0)>0$, then $\lim_{t\to\infty}\lambda_{i}(t)=\sqrt{N}$. It
follows that

$\lim_{t\to\infty}G(t)=\lim_{t\to\infty}\sum_{i}\lambda_{i}(t)\Pi_{i}=\sum_{\lambda_{i}(0)>0}\sqrt{N}\Pi_{i}=\sqrt{N}(I-\Pi).$

For the last claim, we first show that $\lim_{t\to\infty}Z(t)$ exists. Since
$\lim_{t\to\infty}b(t)$ exists and $H(t)$ is determined by $W(t)$, it suffices
to show that $\lim_{t\to\infty}W(t)$ exists. Observe that the singular vectors
of $W(t)$ are also singular vectors of $W^{\prime}(t)$, and so
$W(t)=\sum_{j}\sigma_{j}(t)u_{j}v_{j}^{\top}$, where
$\sum_{j}\sigma_{j}(0)u_{j}v_{j}^{\top}$ denotes any singular value
decomposition of $W_{0}$. For each $j$, select $i(j)$ such that
$\Pi_{i(j)}u_{j}=u_{j}$. Considering the above analysis of
$\lambda_{i(j)}(t)=\sigma_{j}(t)^{2}$, it follows that
$\lim_{t\to\infty}W(t)=N^{1/4}\sum_{\sigma_{i}(0)>0}u_{i}v_{i}^{\top}$. Next,
note that
$\operatorname{ker}(W_{0}^{\top})\supseteq\operatorname{span}\\{1_{C}\\}$
since $(H_{0},W_{0},b_{0})\in S$. Thus, $\operatorname{rank}W_{0}\geq C-1$
implies $\operatorname{ker}(W_{0}^{\top})=\operatorname{span}\\{1_{C}\\}$, and
so $\Pi=\frac{1}{C}1_{C}1_{C}^{\top}$. As such, $\lim_{t\to\infty}G(t)$
satisfies (SNC1), while (SNC2) and (SNC3) follow from taking limits of $H(t)$
and $b(t)$, respectively. ∎

At this point, there appears to be an “unreasonable effectiveness” of gradient
descent in which locally minimizing empirical risk implicitly encourages
strong neural collapse. We conclude with the following result, which explains
why this behavior is actually reasonable:

######  Lemma 4.

For every $(H,W,b)\in S$, it holds that

$R_{e}(H,W,b)=\frac{1}{2}\|WW^{\top}-\sqrt{N}(I-\frac{1}{C}1_{C}1_{C}^{\top})\|_{F}^{2}+\frac{CN}{2}\|b-\frac{1}{C}1_{C}\|_{2}^{2}.$
(3)

In particular, the global minimizers of

$\text{minimize}\qquad R_{e}(H,W,b)\qquad\text{subject to}\qquad(H,W,b)\in S$

are precisely the $(H,W,b)\in\mathbb{R}^{p\times CN}\oplus\mathbb{R}^{C\times
p}\oplus\mathbb{R}^{C}$ which exhibit strong neural collapse.

###### Proof.

Suppose $(H,W,b)\in S$. The implied form of $H$ then gives

$\displaystyle WH+b1_{CN}^{\top}-I_{C}\otimes 1_{N}^{\top}$
$\displaystyle=W(\frac{1}{\sqrt{N}}(W\otimes
1_{N})^{\top})+b1_{CN}^{\top}-(I_{C}-\frac{1}{C}1_{C}1_{C}^{\top}+\frac{1}{C}1_{C}1_{C}^{\top})\otimes
1_{N}^{\top}$
$\displaystyle=(\frac{1}{\sqrt{N}}WW^{\top}-(I_{C}-\frac{1}{C}1_{C}1_{C}^{\top}))\otimes
1_{N}^{\top}+(b-\frac{1}{C}1_{C})1_{CN}^{\top}.$

Since $1_{C}^{\top}W=0$ and $b\in\operatorname{span}\\{1_{C}\\}$, it follows
that the two terms on the right-hand side are orthogonal to each other. The
Pythagorean theorem then gives

$\displaystyle R_{e}(H,W,b)$
$\displaystyle=\frac{1}{2}\|(\frac{1}{\sqrt{N}}WW^{\top}-(I_{C}-\frac{1}{C}1_{C}1_{C}^{\top}))\otimes
1_{N}^{\top}\|_{F}^{2}+\frac{1}{2}\|(b-\frac{1}{C}1_{C})1_{CN}^{\top}\|_{F}^{2}$
$\displaystyle=\frac{1}{2}\|WW^{\top}-\sqrt{N}(I-\frac{1}{C}1_{C}1_{C}^{\top})\|_{F}^{2}+\frac{CN}{2}\|b-\frac{1}{C}1_{C}\|_{2}^{2}.$

For the second part of the result, first observe that the constraint
$(H,W,b)\in S$ implies (SNC2). Subject to this constraint, (3) gives that
equality in $R_{e}(H,W,b)\geq 0$ holds precisely when (SNC1) and (SNC3) both
hold. In addition, $(H,W,b)$ exhibits strong neural collapse only if
$(H,W,b)\in S$. The result follows. ∎

##  4 Discussion

In this paper, we introduced the unconstrained features model that captures
the neural collapse phenomena observed by Papyan, Han and Donoho [11], and
then we identified an invariant subspace of the gradient descent dynamical
system that encourages neural collapse. While we have explained much about the
emergence of neural collapse in the unconstrained features model, there are
several directions for future investigation. First, it would be nice to fully
characterize the dynamics of gradient flow in the unconstrained features
model. How does the distance from $S$ behave over the full gradient flow
trajectory? To what extent does strong neural collapse hold for the limit
point of this trajectory when the initialization is at most $\epsilon>0$ away
from the origin? What behaviors emerge from other popular local optimization
methods such as Adaptive Moment Estimation [7]? What alternatives to the
unconstrained features model are amenable to theoretical treatment?

## Acknowledgments

DGM thanks Arje Nachman and Soledad Villar for (independently) bringing the
article [11] to his attention. DGM was partially supported by AFOSR
FA9550-18-1-0107 and NSF DMS 1829955. HP was partially supported by an AMS-
Simons Travel Grant.

## References

  * [1] S. Arora, N. Cohen, E. Hazan, On the optimization of deep networks: Implicit acceleration by overparameterization, ICML 2018, 372–389.
  * [2] A. S. Bandeira, M. Fickus, D. G. Mixon, P. Wong, The road to deterministic matrices with the restricted isometry property, J. Fourier Anal. Appl. 19 (2013) 1123–1149.
  * [3] L. Chizat, F. Bach, On the global convergence of gradient descent for over-parameterized models using optimal transport, NeurIPS 2018, 3036–3046.
  * [4] S. S. Du, W. Hu, J. D. Lee, Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced, NeurIPS 2018, 384–395.
  * [5] S. S. Du, X. Zhai, B. Poczos, A. Singh, Gradient descent provably optimizes over-parameterized neural networks, ICLR 2018.
  * [6] M. Fickus, D. G. Mixon, Tables of the existence of equiangular tight frames, arXiv:1504.00253
  * [7] D. P. Kingma, J. Ba, Adam: A method for stochastic optimization, arXiv:1412.6980
  * [8] A. Krizhevsky, I. Sutskever, G. Hinton, ImageNet Classification with Deep Convolutional Neural Networks, NIPS 2012, 1097–1105.
  * [9] D. G. Mixon, C. J. Quinn, N. Kiyavash, M. Fickus, Fingerprinting with equiangular tight frames, IEEE Trans. Inf. Theory 59 (2013) 1855–1865.
  * [10] Multi-Layer Neural Network, UFLDL Tutorial, http://ufldl.stanford.edu/tutorial/supervised/MultiLayerNeuralNetworks/
  * [11] V. Papyan, X. Y. Han, D. L. Donoho, Prevalence of neural collapse during the terminal phase of deep learning training, Proc. Natl. Acad. Sci. U.S.A. 117 (2020) 24652–24663.
  * [12] R. A. Rankin, The closest packing of spherical caps in $n$ dimensions, In: Proceedings of the Glasgow Mathematical Association, vol. 2, Cambridge University Press, 1955, pp. 139–144.
  * [13] J. M. Renes, R. Blume-Kohout, A. J. Scott, C. M. Caves, Symmetric informationally complete quantum measurements, J. Math. Phys. 45 (2004) 2171–2180.
  * [14] M. Song, A. Montanari, P. Nguyen, A mean field view of the landscape of two-layers neural networks.” Proc. Natl. Acad. Sci. U.S.A. 115 (2018) E7665–E7671.
  * [15] T. Strohmer, R. W. Heath, Grassmannian frames with applications to coding and communication, Appl. Comput. Harmon. Anal. 14 (2003) 257–275.