# On Perfect Bases in Finite Abelian Groups

Béla Bajnok111Department of Mathematics, Gettysburg College, U.S.A. Email:
<EMAIL_ADDRESS>Corresponding author. Connor Berson222Asymmetrik,
U.S.A. Email<EMAIL_ADDRESS>Hoang Anh Just333Department of Electrical
and Computer Engineering, Virginia Polytechnic Institute and State University,
U.S.A. E-mail<EMAIL_ADDRESS>
(September 29, 2021)

###### Abstract

Let $G$ be a finite abelian group and $s$ be a positive integer. A subset $A$
of $G$ is called a perfect $s$-basis of $G$ if each element of $G$ can be
written uniquely as the sum of at most $s$ (not-necessarily-distinct) elements
of $A$; similarly, we say that $A$ is a perfect restricted $s$-basis of $G$ if
each element of $G$ can be written uniquely as the sum of at most $s$ distinct
elements of $A$. We prove that perfect $s$-bases exist only in the trivial
cases of $s=1$ or $|A|=1$. The situation is different with restricted addition
where perfection is more frequent; here we treat the case of $s=2$ and prove
that $G$ has a perfect restricted $2$-basis if, and only if, it is isomorphic
to $\mathbb{Z}_{2}$, $\mathbb{Z}_{4}$, $\mathbb{Z}_{7}$, $\mathbb{Z}_{2}^{2}$,
$\mathbb{Z}_{2}^{4}$, or $\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$.

2020 AMS Mathematics Subject Classification: Primary: 11B13; Secondary: 05B10,
11P70, 11B75, 20K01.

Key words: Abelian group, sumset, restricted addition, basis, $B_{h}$ set.

##  1 Introduction

Throughout this paper we let $G$ be an additively written finite abelian
group. When $G$ is cyclic and of order $n$, we identify it with
$\mathbb{Z}_{n}=\mathbb{Z}/n\mathbb{Z}$; we consider $0,1,\dots,n-1$
interchangeably as integers and as elements of $\mathbb{Z}_{n}$.

For subsets $A$ and $B$ of $G$, we let $A+B$ and $A-B$ denote the set of two-
term sums and differences, respectively, with one term chosen from $A$ and the
other from $B$. (If, say, $A$ consists of a single element $a$, we simply
write $a+B$ and $a-B$ instead of $A+B$ and $A-B$.) For a subset $A$ of $G$ and
a nonnegative integer $h$, $hA$ denotes the $h$-fold sumset of $A$, that is,
the collection of $h$-term sums with (not-necessarily-distinct) elements from
$A$. The study of sumsets is the main theme in much of additive number theory;
Nathanson’s textbook [11] provides a general introduction to the field, while
the monograph [1] discusses and provides the background for over 300 of its
open questions.

Given a positive integer $s$, we say that a subset $A$ of $G$ is an $s$-basis
for $G$ if every element of $G$ can be written as the sum of at most $s$ (not-
necessarily-distinct) elements of $A$; that is, if $\cup_{h=0}^{s}hA=G$. The
concept has enjoyed a rich history since it was first discussed by Erdős and
Turán in [5] eighty years ago; see, for example, [3], [4], [7], [10], [11],
and [12].

A subset $A$ of $G$ is called a perfect $s$-basis if every element of $G$ can
be written uniquely (apart from the order of the terms) as a sum of at most
$s$ elements of $A$. (In somewhat of a contrast to our usual understanding of
bases, this uniqueness property is not part of the definition of $s$-bases
where only the spanning property is required.) One can then ask for each
instance when a group $G$ possesses a perfect $s$-basis. Trivially, the set of
nonzero elements is a perfect $1$-basis in any group $G$, and the 1-element
set consisting of a generator of $G$ is a perfect $s$-basis in the cyclic
group of order $s+1$. We prove that there are no others:

###### Theorem 1

If a subset $A$ of a finite abelian group $G$ is a perfect $s$-basis, then
$s=1$ and $A=G\setminus\\{0\\}$, or $G\cong\mathbb{Z}_{s+1}$ and $|A|=1$.

Our proof of Theorem 1 is based on the fact that if $A$ is a perfect $s$-basis
in $G$, then the $s+1$ subsets

$-A,\;A-A,\;2A-A,\;\dots,\;(s-1)A-A,\;\mbox{and}\;(s-1)A$

are pairwise disjoint. This will allow us to conclude that, when $A$ is a
perfect $s$-basis in $G$ with $|A|\geq 2$ and $s\geq 2$, then the set

$S=\cup_{h=0}^{s-1}(hA-A)\;\cup\;(s-1)A$

has size more than the order of $G$, which is clearly not possible.

We should note that focusing on sums of at most $s$ terms is equivalent to the
case when exactly $s$ terms are considered. A subset $B$ of $G$ is said to be
an additive basis of order $s$ if every element of $G$ can be expressed as the
sum of exactly $s$ terms (that is, if $sA=G$), and it is called a $B_{s}$ set
if all $s$-term sums are distinct (up to the order of the terms). We now prove
that $G$ has a perfect $s$-basis of size $m$ if, and only if, it has a subset
of size $m+1$ that is an additive basis of order $s$ as well as a $B_{s}$ set.

Suppose first that $A$ is a perfect $s$-basis in $G$. Since then $0\not\in A$,
the set $B=A\cup\\{0\\}$ has size $|A|+1$; it is also easy to see that each
element of $G$ can be written uniquely as the sum of exactly $s$ elements of
$B$. Conversely, assume that $B$ is an additive basis of order $s$ as well as
a $B_{s}$ set. Note that this implies that each translate of $B$ has these two
properties as well; in particular, we may assume that $0\in B$. This then
clearly implies that $A=B\setminus\\{0\\}$ is a perfect $s$-basis in $G$. We
can thus restate Theorem 1 as follows:

###### Theorem 2

If a subset $B$ of a finite abelian group $G$ is a $B_{s}$ set and an additive
basis of order $s$, then $s=1$ and $B=G$, or $G\cong\mathbb{Z}_{s+1}$ and
$|B|=2$.

Let us now turn to restricted addition, that is, when the terms in the sums
must be pairwise distinct. The $h$-fold restricted sumset of $A$, denoted by
$h\hat{\;}A$, is the collection of $h$-term sums with pairwise distinct
elements from $A$. We call a subset $A$ of $G$ a restricted $s$-basis if every
element of $G$ can be written as the sum of at most $s$ distinct elements of
$A$, and $A$ is said to be a perfect restricted $s$-basis if these sums are
(apart from the order of the terms) unique for each element of the group. The
instances of $|A|=1$ or $s=1$ are identical for restricted addition and
unrestricted addition, and are as listed above. As a major contrast, however,
there are infinitely many perfect restricted $s$-bases with $|A|\geq 2$ and
$s\geq 2$ (see Section 4).

Our strategy for classifying perfect restricted $s$-bases is similar to the
unrestricted case, but exhibiting sets that are pairwise disjoint seems more
elusive. Here we only treat the case of $s=2$, for which we prove that when
$G$ is not isomorphic to the elementary abelian 2-group, then $G$ has a
perfect restricted $2$-basis in exactly three cases: when $G$ is isomorphic to
$\mathbb{Z}_{4}$, $\mathbb{Z}_{7}$, or
$\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$. Similarly to the unrestricted case,
for a given perfect restricted $2$-basis $A$ we consider the set

$T=(A-A)\cup A,$

and prove that, unless $G$ is isomorphic to the three groups just listed, $T$
has size more than the order of $G$. Elements of order $2$ in $G$ require
additional attention that we are able to handle when $G$ is not isomorphic to
an elementary abelian 2-group. However, if all nonzero elements of $G$ have
order 2, then

$A-A=2A=\\{0\\}\cup 2\hat{\;}A,$

and thus $T=\cup_{h=0}^{2}h\hat{\;}A.$ Therefore, $A$ being a perfect
restricted $2$-basis in $G$ is equivalent to having $T=G$, resulting in no
contradiction. Luckily, a problem of Ramanujan comes to the rescue.

In 1913, Ramanujan asked in [13] whether the quantity $2^{k}-7$ can be a
square number for any integer $k$ besides 3, 4, 5, 7, and 15 (see also
Question 464 in [2]). The negative answer was given by Nagell in 1948 (see
[8]; [9] for the English version). Suppose now that $G$ is the elementary
abelian 2-group of rank $r$ that has a perfect restricted $2$-basis of size
$m$: we then must have

$2^{r}={m\choose 0}+{m\choose 1}+{m\choose 2},$

and therefore

$2^{r+3}-7=(2m+1)^{2}.$

We thus have exactly four choices for $r$: 1, 2, 4, and $12$. We show that
perfect restricted $2$-bases exist in the first three cases but not in
$\mathbb{Z}_{2}^{12}$. In summary, we have the following results:

###### Theorem 3

A finite abelian group $G$ has a perfect restricted $2$-basis if, and only if,
it is isomorphic to one of the following groups:
$\mathbb{Z}_{2},\;\mathbb{Z}_{4},\;\mathbb{Z}_{7},\;\mathbb{Z}_{2}^{2},\;\mathbb{Z}_{2}^{4}$,
or $\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$.

The structure of the paper is as follows: In Section 2 we prove Theorem 1, and
in Section 3 we prove Theorem 3. We then consider some further results and
present some open questions in Section 4.

##  2 The proof of Theorem 1

As mentioned in the Introduction, in order to prove Theorem 1, we argue
indirectly, and show that, when $A$ is a perfect $s$-basis in $G$ with
$|A|\geq 2$ and $s\geq 2$, then the size of

$S=\cup_{h=0}^{s-1}(hA-A)\;\cup\;(s-1)A$

is more than the order of $G$, which is impossible.

Recall that if $A$ is a perfect $s$-basis in $G$, then it is a $B_{h}$ set for
every $0\leq h\leq s$, and $G$ is the disjoint union of sets $hA$ with $0\leq
h\leq s$. Therefore, for each such $h$ we have

$|hA|={m+h-1\choose h}$

and thus

$|G|=\sum_{h=0}^{s}|hA|={m+s\choose s}.$

In order to compute the size of $S$, first note that if $A$ is a perfect
$s$-basis in $G$, then the $s+1$ sets $-A,A-A,2A-A,\dots,(s-1)A-A$, and
$(s-1)A$ are pairwise disjoint. Therefore,

$|S|=\sum_{h=0}^{s-1}|hA-A|+|(s-1)A|=\sum_{h=0}^{s-1}|hA-A|+{m+s-2\choose
s-1}.$

For $h=0$ we have $|hA-A|=|-A|=m$. To find the size of $hA-A$ for a given
$1\leq h\leq s-1$, we first observe that $(h-1)A$ is a subset of $hA-A$, so
there are exactly ${m+h-2\choose h-1}$ elements of $hA-A$ that are in
$(h-1)A$. We will now show that there are precisely $m\cdot{m+h-2\choose h}$
elements in $hA-A$ that are not in $(h-1)A$.

Let $A=\\{a_{1},\dots,a_{m}\\}$. Observe that each element $g$ of
$(hA-A)\setminus(h-1)A$ can be written as

$g=\lambda_{1}a_{1}+\cdots+\lambda_{m}a_{m}-a_{i},$

where $\lambda_{1},\dots,\lambda_{m}$ are nonnegative integers that sum to
$h$, $1\leq i\leq m$, and $\lambda_{i}=0$. As we now show, the expression is
unique for each $g$. Suppose that we also have

$g=\lambda^{\prime}_{1}a_{1}+\cdots+\lambda^{\prime}_{m}a_{m}-a_{j}$

with some $1\leq j\leq m$,
$\lambda^{\prime}_{1}+\cdots+\lambda^{\prime}_{m}=h$, and
$\lambda^{\prime}_{j}=0$. This yields

$\lambda_{1}a_{1}+\cdots+\lambda_{m}a_{m}+a_{j}=\lambda^{\prime}_{1}a_{1}+\cdots+\lambda^{\prime}_{m}a_{m}+a_{i}\in(h+1)A.$

If $i\neq j$, then since $A$ is a $B_{h+1}$ set, we must have
$\lambda_{j}+1=\lambda^{\prime}_{j}$ and $\lambda_{i}=\lambda^{\prime}_{i}+1$,
contradicting $\lambda_{i}=\lambda^{\prime}_{j}=0$. Furthermore, $g+a_{i}\in
hA$, so the coefficients $\lambda_{1},\dots,\lambda_{m}$ are unique as well.

This means that $(hA-A)\setminus(h-1)A$ is the disjoint union of sets
$h(A\setminus\\{a_{i}\\})-a_{i}$ for $1\leq i\leq m$. Since a subset of a
$B_{h}$ set is also a $B_{h}$ set, we have

$|h(A\setminus\\{a_{i}\\})-a_{i}|=|h(A\setminus\\{a_{i}\\})|={m+h-2\choose
h},$

implying

$|(hA-A)\setminus(h-1)A|=m\cdot{m+h-2\choose h},$

as claimed.

We thus find that

$|hA-A|=m\cdot{m+h-2\choose h}+{m+h-2\choose h-1},$

from which

$\sum_{h=0}^{s-1}|hA-A|=m+\sum_{h=1}^{s-1}m\cdot{m+h-2\choose
h}+\sum_{h=1}^{s-1}{m+h-2\choose h-1}=m{m+s-2\choose s-1}+{m+s-2\choose s-2}.$

This then yields

$|S|=\sum_{h=0}^{s-1}|hA-A|+|(s-1)A|=m{m+s-2\choose s-1}+{m+s-1\choose s-1};$

rewriting this expression, we get

$|S|=m{m+s-2\choose s-1}+{m+s\choose s}-{m+s-1\choose s}={m+s\choose
s}+\frac{(m-1)(s-1)}{s}{m+s-2\choose s-1}.$

Since this quantity is larger than $|G|={m+s\choose s}$ for $m\geq 2$ and
$s\geq 2$, our proof is complete.

##  3 The proof of Theorem 3

We start by treating the cases when $G$ is isomorphic to an elementary abelian
2-group $\mathbb{Z}_{2}^{r}$. As explained in the Introduction, we only need
to examine four cases: when $G$ has rank 1, 2, 4, or 12. It is not hard to
find perfect restricted 2-bases in the first three of these cases; for
example:

  * •

in $\mathbb{Z}_{2}$, the set $\\{1\\}$ is a perfect restricted 2-basis;

  * •

in $\mathbb{Z}_{2}^{2}$, the set $\\{(0,1),(1,0)\\}$ is a perfect restricted
2-basis; and

  * •

in $\mathbb{Z}_{2}^{4}$, the set
$\\{(0,0,0,1),(0,0,1,0),(0,1,0,0),(1,0,0,0),(1,1,1,1)\\}$ is a perfect
restricted 2-basis.

However, we can easily see that $\mathbb{Z}_{2}^{12}$ has no perfect
restricted 2-bases. For the sake of a contradiction, suppose that $A$ is a
perfect restricted 2-basis in $\mathbb{Z}_{2}^{12}$; it then has size
$|A|=90$. If $k$ elements of $A$ have a first component of $1$ for some $0\leq
k\leq 90,$ then $0\hat{\;}A\cup 1\hat{\;}A\cup 2\hat{\;}A$ has $k+k(90-k)$
elements with a first component of $1$. Since there is no $k$ for which
$k+k(90-k)$ equals $|G|/2=2^{11}$, this is a contradiction.

We now turn to the cases when $G$ is not isomorphic to elementary abelian
2-groups: we need to prove that $G$ has a perfect restricted 2-basis if, and
only if, it is isomorphic to $\mathbb{Z}_{4}$, $\mathbb{Z}_{7}$, or
$\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$.

We let $L$ denote the subset consisting of the identity element of $G$ as well
as all involutions in $G$; that is,

$L=\\{g\in G\mid 2g=0\\}.$

Note that $L$ is a subgroup of $G$; in fact, $L$ is isomorphic to the
elementary abelian 2-group whose rank equals the number of even-order terms in
the invariant decomposition of $G$.

Let $A=\\{a_{1},\dots,a_{m}\\}$ be an $m$-subset of $G$, and suppose that $A$
is a perfect restricted 2-basis in $G$. The main idea of the proof is to
estimate the size of $(A-A)\cup A$, namely, to prove that

$|(A-A)\cup A|\geq m^{2}-2m+2-|L|.$

Since, as we show, this bound is more than $|G|$ when $m\geq 6$, we get a
contradiction; the cases of $m\leq 5$ will be examined individually.

For an element $g\in G$, let $R(g)$ denote the number of ordered pairs
$(a_{i},a_{j})\in A^{2}$ for which $g=a_{i}-a_{j}$; we say that $R(g)$ is the
representation number of $g$. Below we prove that $R(g)\in\\{0,1,2,m\\}$ for
each $g\in G$ and, in fact, count the number of elements for each possible
representation number.

Clearly, elements not in $A-A$ have representation number 0, and $R(0)=m$.

Suppose now that $g\in A-A$, $g\neq 0$; we will prove that $R(g)=1$ or
$R(g)=2$. Let’s assume indirectly that $R(g)\geq 3$, and let

$g=a_{1}-a_{2}=a_{3}-a_{4}=a_{5}-a_{6}$

for some $a_{1},\dots,a_{6}\in A$ such that $a_{1},a_{3},a_{5}$ are pairwise
distinct and $a_{2},a_{4},a_{6}$ are pairwise distinct. (Note that if
$a_{1}=a_{3}$, for example, then $a_{2}=a_{4}$ and thus those two
representations of $g$ are identical.) Our equations imply that
$a_{1}+a_{4}=a_{2}+a_{3}$, $a_{1}+a_{6}=a_{2}+a_{5}$, and
$a_{3}+a_{6}=a_{4}+a_{5}$. Since $A$ is a perfect restricted 2-basis, we must
have $a_{1}=a_{4}$ or $a_{2}=a_{3}$; $a_{1}=a_{6}$ or $a_{2}=a_{5}$; and
$a_{3}=a_{6}$ or $a_{4}=a_{5}$. In particular, note that if $a_{1}=a_{4}$,
then we cannot have $a_{1}=a_{6}$ or $a_{4}=a_{5}$, since that would imply
that $a_{4}=a_{6}$ or $a_{1}=a_{5}$, respectively, which we can’t have.
Therefore, there are only two possibilities: $a_{1}=a_{4}$, $a_{2}=a_{5}$, and
$a_{3}=a_{6}$; or $a_{2}=a_{3}$, $a_{1}=a_{6}$, and $a_{4}=a_{5}$. In the
first case, we have

$3g=(a_{1}-a_{2})+(a_{3}-a_{4})+(a_{5}-a_{6})=(a_{1}-a_{2})+(a_{3}-a_{1})+(a_{2}-a_{3})=0,$

and in the second case

$3g=(a_{1}-a_{2})+(a_{3}-a_{4})+(a_{5}-a_{6})=(a_{1}-a_{2})+(a_{2}-a_{4})+(a_{4}-a_{1})=0,$

and, since $g\neq 0$, this implies that the order of $G$ must be divisible by
3. However,

$|G|={m\choose 0}+{m\choose 1}+{m\choose 2}=\frac{m^{2}+m+2}{2},$

which is not divisible by 3 for any $m$. We arrived at a contradiction, and
thus every nonzero element of $A-A$ has representation number equal to either
1 or 2.

We now identify the elements of $A-A$ that have representation number 2. We
call two elements $a_{i},a_{j}\in A$ twins if $2a_{i}=2a_{j}$ or,
equivalently, if $a_{i}-a_{j}\in L\setminus\\{0\\}$. More specifically, we say
that $a_{i}$ and $a_{j}$ are $g$-twins for some $g\in L\setminus\\{0\\}$ when
$a_{i}=a_{j}+g$ (in which case we also have $a_{j}=a_{i}+g$). Clearly, if
$a_{i}$ and $a_{j}$ are twins, then $R(a_{i}-a_{j})=2$ since
$a_{i}-a_{j}=a_{j}-a_{i}$. Now we show that there cannot be more than one
$g$-twin pair for any $g\in L\setminus\\{0\\}$; in other words, if $a_{i}$ and
$a_{j}$ are $g$-twins and $a_{i}^{\prime}$ and $a_{j}^{\prime}$ are $g$-twins,
then $\\{a_{i},a_{j}\\}=\\{a_{i}^{\prime},a_{j}^{\prime}\\}$. Indeed if
$a_{i}-a_{j}=a_{i}^{\prime}-a_{j}^{\prime}=g$, then adding the equations
results in $a_{i}+a_{i}^{\prime}=a_{j}+a_{j}^{\prime}$ which, since $A$ is a
perfect restricted 2-basis, can only happen if $a_{i}=a_{i}^{\prime}$ (but
then $a_{j}=a_{j}^{\prime}$), or $a_{j}=a_{j}^{\prime}$ (but then
$a_{i}=a_{i}^{\prime}$), or $a_{i}=a_{j}^{\prime}$ and $a_{j}=a_{i}^{\prime}$.
In any case, $\\{a_{i},a_{j}\\}=\\{a_{i}^{\prime},a_{j}^{\prime}\\}$, as
claimed, and this implies that there are at most $|L|-1$ pairs of twins in
$A$.

Next, we examine the awesum elements of $A$, that is, elements $a_{i}$ for
which $2a_{i}\in 2\hat{\;}A$. If $a_{i}\in A$ is awesum, then there exists an
unordered pair of distinct elements $a_{i^{\prime}},a_{i^{\prime\prime}}\in A$
with $2a_{i}=a_{i^{\prime}}+a_{i^{\prime\prime}};$ we call $a_{i^{\prime}}$
and $a_{i^{\prime\prime}}$ the summands of $a_{i}$. Note that the summands of
an element $a_{i}$ cannot equal $a_{i}$ since that would imply that the two
summands are equal. Now if $a_{i}$ is an awesum element of $A$ with summands
$a_{i^{\prime}}$ and $a_{i^{\prime\prime}}$, then
$R(a_{i}-a_{i}^{\prime})=R(a_{i}-a_{i}^{\prime\prime})=2$ since
$a_{i}-a_{i}^{\prime}=a_{i}^{\prime\prime}-a_{i}$ and
$a_{i}-a_{i}^{\prime\prime}=a_{i}^{\prime}-a_{i}$.

We are now ready to calculate $|A-A|$ as the number of ordered pairs
$(a_{i},a_{j})$ that yield distinct elements of $A-A$. Since $R(0)=m$ and
$R(g)\in\\{1,2\\}$ for each nonzero element of $A-A$, we have

$|A-A|=m^{2}-(m-1)-\alpha,$

where $\alpha$ is the number of elements of $A-A$ whose representation number
equals 2. Suppose that $g\in A-A$ has $R(g)=2$ and that
$g=a_{1}-a_{2}=a_{3}-a_{4}$ for some $a_{1},a_{2},a_{3},a_{4}\in G$ with
$a_{1}\neq a_{3}$ and $a_{2}\neq a_{4}$. This implies that
$a_{1}+a_{4}=a_{2}+a_{3}$, so we have three possibilities: (i) $a_{1}=a_{4}$
and $a_{2}=a_{3}$, in which case $a_{1}$ and $a_{2}$ are twins; (ii)
$a_{1}=a_{4}$ and $a_{2}\neq a_{3}$, in which case $a_{1}$ is awesum and
$a_{2}$ is one of its summands; or (iii) $a_{1}\neq a_{4}$ and $a_{2}=a_{3}$,
in which case $a_{2}$ is awesum and $a_{1}$ is one of its summands. Thus
$\alpha$ equals the number of twin pairs in $A$ (denoted by $\beta$) plus two
times the number of awesum elements of $A$ (denoted by $\gamma$):

$|A-A|=m^{2}-(m-1)-\beta-2\gamma.$

Finally, we count the number of elements in $(A-A)\cup A$ by subtracting
$|(A-A)\cap A|$ from $|A-A|+|A|$. Suppose that $a=a^{\prime}-a^{\prime\prime}$
for some $a,a^{\prime},a^{\prime\prime}\in A$ or, equivalently, that
$a+a^{\prime\prime}=a^{\prime}$. Since $A$ and $2\hat{\;}A$ are disjoint, this
implies that $a=a^{\prime\prime}$ and so $2a\in A$. But then $2a\not\in
2\hat{\;}A$, so $a$ is not awesum, and thus we have at most $m-\gamma$
elements of $A$ that are in $A-A$. Putting it all together, we get

$\displaystyle|(A-A)\cup A|$ $\displaystyle=$
$\displaystyle|A-A|+|A|-|(A-A)\cap A|$ $\displaystyle\geq$ $\displaystyle
m^{2}-(m-1)-\beta-2\gamma+m-(m-\gamma)$ $\displaystyle=$ $\displaystyle
m^{2}-m+1-\beta-\gamma;$

and since $\beta\leq|L|-1$ and $\gamma\leq m$, we arrive at

$|(A-A)\cup A|\geq m^{2}-2m+2-|L|,$

as claimed.

Now since $(A-A)\cup A$ is a subset of $G$, this implies that the order $n$ of
$G$ satisfies

$n=\frac{m^{2}+m+2}{2}\geq m^{2}-2m+2-|L|.$

Recall that $L$ is a subgroup of $G$ and its order is a power of 2. As here we
treat the case when $G$ is not isomorphic to an elementary abelian 2-group, we
must have $|L|\leq n/2$, which then implies that $m\leq 10$. As we mentioned
above, $n$ is not divisible by 3 for any value of $m$, and we also find that
it is never divisible by 5 either, hence $|L|\neq n/3,n/5,$ or $n/6$.
Furthermore, the values of $n=(m^{2}+m+2)/2$ are not powers of 2 for any
$6\leq m\leq 10$. Therefore, we have $|L|\leq n/7$, from which we get $m\leq
5$.

The cases of $m\leq 5$ are easy to handle: since $G$ is not isomorphic to an
elementary abelian 2-group, we only need to consider $G\cong\mathbb{Z}_{4}$
for $m=2$; $G\cong\mathbb{Z}_{7}$ for $m=3$; $G\cong\mathbb{Z}_{11}$ for
$m=4$; and
$G\cong\mathbb{Z}_{16},\;\mathbb{Z}_{2}\times\mathbb{Z}_{8},\;\mathbb{Z}_{4}^{2},$
or $\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$ for $m=5$. One can easily verify
that:

  * •

in $\mathbb{Z}_{4}$, the set $\\{1,2\\}$ is a perfect restricted 2-basis;

  * •

in $G=\mathbb{Z}_{7}$, the set $\\{1,2,4\\}$ is a perfect restricted 2-basis;
and

  * •

in $G=\mathbb{Z}_{2}^{2}\times\mathbb{Z}_{4}$, the set
$\\{(0,0,2),(0,0,3),(0,1,1),(1,0,1),(1,1,1)\\}$ is a perfect restricted
2-basis.

(In Section 4 we will see how each of these three constructions generalizes
for arbitrary $s$.) However, we find – using the computer program [6] – that
$\mathbb{Z}_{11},\;\mathbb{Z}_{16},\;\mathbb{Z}_{2}\times\mathbb{Z}_{8}$ and
$\mathbb{Z}_{4}^{2}$ have no perfect restricted 2-bases. This completes our
proof of Theorem 3.

##  4 Further results and open questions

In this section we generalize the perfect restricted $2$-bases we exhibited
above for arbitrary values of $s$ and present some open questions.

Suppose that a finite abelian group $G$ has a perfect restricted $s$-basis of
size $m$. When $m\leq s$, then for the order of $G$ we have

$|G|=\sum_{h=0}^{s}{m\choose h}=2^{m}.$

Let $G$ be an abelian group of order $2^{m}$; we then may assume that there
exist positive integers $k_{1},\dots,k_{r}$ for which

$G=\mathbb{Z}_{2^{k_{1}}}\times\mathbb{Z}_{2^{k_{2}}}\times\cdots\times\mathbb{Z}_{2^{k_{r}}}$

with $k_{1}+\cdots+k_{r}=m$. For each $i=1,\dots,r$, consider the $k_{i}$
elements of $G$ whose components are all 0, except for their $i$-th
components, which are equal to $1,2,\dots,2^{k_{i}-1}$, respectively. Then it
is easy to see that $m$-subset of $G$ consisting of these elements is a
perfect restricted $m$-basis (and thus a perfect restricted $s$-basis for any
$s\geq m$) in $G$. Therefore, the case of $m\leq s$ is settled:

###### Proposition 4

Suppose that $m$ and $s$ are positive integers with $m\leq s$. Then $G$ has a
perfect restricted $s$-basis of size $m$ if, and only if, $G$ has order
$2^{m}$.

Next, we consider $m=s+1$, in which case

$|G|=\sum_{h=0}^{s}{m\choose h}=2^{s+1}-1.$

As it is easy to see that the set $\\{1,2,,\dots,2^{s}\\}$ is a perfect
restricted $s$-basis in the cyclic group of order $2^{s+1}-1$, we have:

###### Proposition 5

For every positive integer $s$, the cyclic group of order $2^{s+1}-1$ has a
perfect restricted $s$-basis of size $s+1$.

The question then becomes:

###### Question 6

Are there any noncyclic groups that have a perfect restricted $s$-basis of
size $s+1$?

As we have not found any perfect restricted $s$-bases of size more than $s+1$
until reaching size $2s+1$, we ask:

###### Question 7

Are there any groups that have a perfect restricted $s$-basis of size $m$ with
$s+2\leq m\leq 2s$?

This brings us to the case of $m=2s+1$, for which we have

$\sum_{h=0}^{s}{m\choose h}=\frac{1}{2}\cdot\sum_{h=0}^{2s+1}{2s+1\choose
h}=2^{2s}.$

We found two different groups with perfect restricted $s$-bases of size
$2s+1$. One of them is the elementary abelian 2-group of rank $2s$. For
$i=1,\dots,2s$, let ${\bf e}_{i}$ denote the $i$-th unit vector of
$\mathbb{Z}_{2}^{2s}$, that is, the element with a 1 in the $i$-th component
and 0 everywhere else, and let ${\bf e}$ be their sum. Then the $2s$ unit
vectors generate each element of the group with at most $s$ nonzero
components, and the unit vectors added to ${\bf e}$ generate the rest; each
element is generated only once. We thus have:

###### Proposition 8

For every positive integer $s$, the elementary abelian $2$-group of rank $2s$
has a perfect restricted $s$-basis of size $2s+1$.

The other group of order $2^{2s}$ with a perfect restricted $s$-basis that we
found is $\mathbb{Z}_{2}^{2s-2}\times\mathbb{Z}_{4}$. Similarly as above, for
$i=1,\dots,2s-2$ we let ${\bf e}_{i}$ denote the element of $G$ with a 1 in
the $i$-th component and 0 everywhere else, and we let ${\bf e}$ be their sum.
Furthermore, we set ${\bf f}$ equal to the vector with 0s in its first $2s-2$
components and 1 in its last. We now verify that the set

$A=\\{{\bf e}_{i}+{\bf f}\mid i=1,\dots,2s-2\\}\cup\\{2{\bf f},3{\bf f},{\bf
e}+{\bf f}\\}$

is a perfect restricted $s$-basis in
$G=\mathbb{Z}_{2}^{2s-2}\times\mathbb{Z}_{4}$. Since $|A|=2s+1$ and
$|G|=2^{2s}$, it suffices to prove that all elements of $G$ can be written as
the sum of at most $s$ distinct terms from $A$.

For an integer $0\leq k\leq 2s-2$, we let $S_{k}$ denote the collection of
group elements that have exactly $k$ of their first $2s-2$ components equal to
1, and we write

$B=\\{{\bf e}_{i}+{\bf f}\mid i=1,\dots,2s-2\\}$

for short. Now let $g\in S_{k}$ be an arbitrary element. Observe that

$S_{k}=\left(k\hat{\;}B\right)\cup\left(k\hat{\;}B+2{\bf
f}\right)\cup\left(k\hat{\;}B+3{\bf f}\right)\cup\left(k\hat{\;}B+2{\bf
f}+3{\bf f}\right),$

and thus if $k\leq s-2$, then $g\in\cup_{h=0}^{s}h\hat{\;}A.$ Replacing
$k\hat{\;}B$ above by $(2s-2-k)\hat{\;}B+({\bf e}+{\bf f})$ allows us to
conclude that $g\in\cup_{h=0}^{s}h\hat{\;}A$ whenever $2s-2-k$ is at most
$s-3$, that is, when $k\geq s+1$.

This leaves us with the cases of $k=s-1$ and $k=s$, for which we see that

$S_{s-1}=\left((s-1)\hat{\;}B\right)\cup\left((s-1)\hat{\;}B+2{\bf
f}\right)\cup\left((s-1)\hat{\;}B+3{\bf
f}\right)\cup\left((s-1)\hat{\;}B+({\bf e}+{\bf f})\right)$

and

$S_{s}=\left(s\hat{\;}B\right)\cup\left((s-2)\hat{\;}B+({\bf e}+{\bf
f})\right)\cup\left((s-2)\hat{\;}B+({\bf e}+{\bf f})+2{\bf
f}\right)\cup\left((s-2)\hat{\;}B+({\bf e}+{\bf f})+3{\bf f}\right),$

with which all cases are covered.

Therefore, we have the following result:

###### Proposition 9

For every positive integer $s$, the group
$\mathbb{Z}_{2}^{2s-2}\times\mathbb{Z}_{4}$ has a perfect restricted $s$-basis
of size $2s+1$.

###### Question 10

Are there other groups besides those isomorphic to $\mathbb{Z}_{2}^{2s}$ or
$\mathbb{Z}_{2}^{2s-2}\times\mathbb{Z}_{4}$ that are of order $2^{2s}$ and
have a perfect restricted $s$-basis of size $2s+1$?

We also could not find any cases with $m$ greater than $2s+1$, so ask:

###### Question 11

Are there any groups that have a perfect restricted $s$-basis of size $m$ with
$m\geq 2s+2$?

We close by mentioning that while Theorems 1 and 2 are easily seen to be
equivalent, a similar statement does not appear to hold for restricted sums.
In particular, we do not see a correspondence between perfect restricted
$s$-bases, considered above, and subsets that generate every group element
uniquely as the sum of exactly $s$ distinct elements. This then raises the
following general question:

###### Question 12

Given a positive integer $s$, characterize all finite abelian groups $G$ that
have a subset $A$ for which each element of $G$ arises as a unique sum of
exactly $s$ distinct elements of $A$.

## References

  * [1] B. Bajnok, Additive Combinatorics: A Menu of Research Problems. CRC Press, Boca Raton, 2018, xix+390 pp.
  * [2] B. C. Berndt, Y-S. Choi, and S-I. Kang, The problems submitted by Ramanujan to the Journal of the Indian Mathematical Society. Contemp. Math. 236 (1999), 15–56.
  * [3] P. Erdős and R. L. Graham, On bases with an exact order. Acta Arith. 37 (1980), 201–207.
  * [4] P. Erdős and M. B. Nathanson, Problems and results on minimal bases in additive number theory. Lecture Notes in Math. 1240, Springer, Berlin (1987), 87–96.
  * [5] P. Erdős and P. Turán, On a problem of Sidon in additive number theory and some related questions. J. London Math. Soc. 16 (1941), 212–215.
  * [6] I. Ilinkin, Sumset generator. Published at http://addcomb.gettysburg.edu.
  * [7] V. Lambert, T. H. Lê, and A. Plagne, Additive bases in groups. Israel J. Math. 217 (2017), no. 1, 383–411.
  * [8] T. Nagell, Løsning til oppgave nr. 2, 1943, s. 29. Nordisk Mat. Tidskr. 30 (1948), 62–64.
  * [9] T. Nagell, The Diophantine Equation $x^{2}+7=2^{n}$. Ark. Mat. 4 (1961), 185–187.
  * [10] M. B. Nathanson, Minimal bases and maximal nonbases in additive number theory. J. Number Theory 6 (1974), 324–333.
  * [11] M. B. Nathanson, Additive number theory: Inverse problems and the geometry of sumsets. Graduate Texts in Mathematics, 165. Springer–Verlag, New York, 1996. xiv+293 pp.
  * [12] M. B. Nathanson, Paul Erdős and additive bases. arXiv:1401.7598 [math.NT] (2014).
  * [13] S. Ramanujan, Question 464. Journal of the Indian Mathematical Society 5 (1913), 120.


# The Impact of Geometric Complexity on Neural Collapse in Transfer Learning

Michael Munn
Google Research
<EMAIL_ADDRESS>
&Benoit Dherin∗
Google Research
<EMAIL_ADDRESS>
&Javier Gonzalvo
Google Research
<EMAIL_ADDRESS>
Equal contribution.

###### Abstract

Many of the recent remarkable advances in computer vision and language models
can be attributed to the success of transfer learning via the pre-training of
large foundation models. However, a theoretical framework which explains this
empirical success is incomplete and remains an active area of research.
Flatness of the loss surface and neural collapse have recently emerged as
useful pre-training metrics which shed light on the implicit biases underlying
pre-training. In this paper, we explore the geometric complexity of a model’s
learned representations as a fundamental mechanism that relates these two
concepts. We show through experiments and theory that mechanisms which affect
the geometric complexity of the pre-trained network also influence the neural
collapse. Furthermore, we show how this effect of the geometric complexity
generalizes to the neural collapse of new classes as well, thus encouraging
better performance on downstream tasks, particularly in the few-shot setting.

##  1 Introduction

Many of the recent remarkable advances in modern machine learning owe their
success in part to transfer learning [15, 28, 6, 9, 4, 48, 62, 47]. While
there are different approaches to this technique, the standard one involves
two stages. In a first stage, called pre-training, ones trains a deep neural
network on a general, large-scale dataset in the form of a supervised or
unsupervised source task; e.g., ImageNet or CIFAR-100 [14, 32] for image
models or the Common Crawl, C4 or LM1B datasets [8, 11, 47] for language
models. In the second stage, one then leverages portions of the pre-trained
network to use as features map or embeddings that can then be adapted, or
fine-tuned, on a more specific target task. Often these target tasks are
unknown at the time of pre-training and labeled data may be very scarce, such
as in the context of few-shot learning [57, 33, 49]. However, despite these
limitations, this approach often results in a fine-tuned model that achieves
quite impressive performance substantially better than training on the target
task alone and requiring less computational resources [20]. Despite this
empirical success, a comprehensive theoretical understanding of the mechanisms
underlying this effectiveness of transfer learning is not fully understood and
remains an active area of research [60].

One interesting line of research suggests that the effectiveness of transfer
learning is due to the implicit biases encoded in pre-trained models [22, 35,
46]. These implicit biases effectively constrain the hypothesis space, guiding
the model towards solutions that have a preference for smoother functions [12,
39, 51, 37, 17], simpler geometry in the loss surface [3, 52, 18, 13, 55] or
reduced complexity of the internal learned representations [23, 31]. In the
same way that these implicit biases have been used to help explain the success
of deep learning, recent work has shown that the notion of neural collapse
[22, 59, 34] and flatness of the loss surface [35] can also inform the
mechanisms behind transfer learning. This suggests that these preferences
during pre-training also act as a form of prior knowledge which is highly
transferable to downstream tasks, even with limited task-specific data.

In this paper, we present a novel perspective that further sheds light on
these mechanisms and implicit biases hidden within transfer learning. Our
approach analyzes the geometric complexity [16, 17] of the internal
representations learned by the deep neural networks during pre-training and
provides a complementary theoretical framework which unifies previous work
examining the role of neural collapse and loss surface flatness in transfer
learning [22, 35]. In particular, we show through experiments and theory that
the geometric complexity of the pre-trained network directly controls the
neural collapse of the pre-trained model and thus its efficacy in transfer
learning, particularly in the few-shot setting. We argue that geometric
complexity (similarly to flatness of the loss surface and neural collapse) can
be used as hidden progress metrics, cf. [2], for transfer learning, serving as
an informative proxy toward the transfer power of a pre-trained network.

Our primary contributions are the following:

  * •

We uncover relationships between learning-path flatness, neural network
geometric-complexity, and embedding neural-collapse providing a framework to
understand how these implicit biases interact (Section 4).

  * •

We show through theory that the geometric complexity (GC) of a neural network
controls neural collapse and verify this empirically by showing how mechanisms
which regularize the GC in turn put pressure on the neural collapse (Section
4.1 and Fig. 1).

  * •

We demonstrate both theoretically and empirically that pre-trained networks
with lower GC promote lower neural collapse on new unseen target classes, and
thus enable improved transfer accuracy during fine-tuning (Section 5 and Fig.
3).

  * •

We prove a new generalization bound in terms of geometric complexity (Section
4.3).

  * •

We show that the empirical GC can be accurately and efficiently estimated on a
small number of samples, input coordinates, and output features making it
computationally tractable compared to other progress metrics in machine
learning (Section 4.2 and Fig. 2).

##  2 Background and Related Work

The implicit biases introduced by an optimization algorithm play a crucial
role in deep learning and in the generalization ability of the learned models
[43, 52, 61, 42]. They help ensure that the model not only finds a solution
with low error but also one with low complexity which generalizes well [63,
26]. Uncovering the mechanisms of these implicit regularizers is crucial for
understanding how the model learns, both as a means to improve generalization
and as valuable leverage for designing more efficient algorithms and
decreasing costly experiment iteration cycles.

Here we focus on three seemingly different themes behind implicit
regularization in deep learning and explain how they are related: the loss
surface slope measuring the flatness of the learning path [3, 44, 29], the
geometric complexity of the learned model function measuring its variability
with respect to a dataset [16, 41], and the neural collapse [45, 31] measuring
how a neural network efficiently clusters its learned representations of the
input class examples in embedding space.

##### Flatness in parameter space.

In parameter space, the learning dynamics is fully characterized by the
discrete optimization path $\theta_{t}$ where $t\in\mathbb{N}$. At each step,
one can compute the flatness of the learning path as the slope of the tangent
space at $\theta_{t}$ to the loss surface ${\cal
L}=\\{(\theta,L(\theta)):\>\theta\in\mathbb{R}^{n}\\}$. As shown in [3,
Appendix A.2], this slope coincides with the loss-gradient square norm:
$\operatorname{slope}(\theta_{t})=\|\nabla L(\theta_{t})\|^{2}$. This quantity
has been used to bound a generalization gap in [25] and as a beneficial
explicit regularizer in [3, 24, 55] indicating that learning curves with lower
slope values tend to also have better test performance.

Moreover, many standard optimizers and common training heuristics have been
shown to put an implicit pressure on the learning-path flatness, making it an
implicit bias of the training procedure [3, 55, 25, 10, 13, 17, 5, 38].
Related to the learning curve slope is the optima sharpness

$\operatorname{sharpness}(\theta_{*})=\frac{1}{n}\operatorname{trace}H(\theta_{*})$

where $\theta_{*}$ is a global minima toward which the learning dynamics
converges to $\theta_{t}\rightarrow\theta_{*}$ and $H$ is the Hessian of the
loss. Since the sharpness of an optima corresponds to the mean curvature of
the loss surface at that point, flat minima (i.e., minima with low curvature
in all directions) can be reached only through learning paths with shallower
slopes.

##### Neural collapse in embedding space.

Neural collapse [22, 45, 27] refers to a phenomenon observed in deep learning
where the embedding network (i.e., the subnetwork $f$ before the logit layer
$g$ in a neural network $h(x)=\operatorname{softmax}(g(f(x))$) collapses the
input points around their repective class means. Furthermore, these class
means form a simplex creating an equiangular tight frame (ETF) centered around
the global mean with roughly equal distance between its vertices. Intuitively,
such a phenomenon is beneficial to generalization since the embeddings of
different classes are optimally separated, making the job of the classifier
head $\operatorname{softmax}(g(z))$ easier. To measure neural collapse, the
authors in [22] introduce the _class-distance normalized variance_
($\operatorname{CDNV}$) as

$V_{f}(Q_{i},Q_{j}):=\frac{\operatorname{Var}_{f}(Q_{i})+\operatorname{Var}_{f}(Q_{j})}{\|\mu_{f}(Q_{i})-\mu_{f}(Q_{j})\|^{2}},$
(1)

where $Q_{r}=q_{r}(x)dx$ is the input distribution for classes $r\in\\{i,j\\}$
and where

$\mu_{f}(Q_{r})=\mathbb{E}_{x\sim
Q_{r}}[f(x)]\quad\textrm{and}\quad\operatorname{Var}_{f}(Q_{r})=\int\mathbb{\|}\mu_{f}(Q_{r})-f(x)\|^{2}q_{r}(x)dx.$
(2)

Neural collapse between two classes happens when their class variance
decreases while the distance between their class means increases, causing
lower values of their CDNV. Following [22], given a well-balanced input
distribution $Q=\frac{1}{k}(Q_{1}+\cdots+Q_{k})$ with $k$ classes ($Q_{i}$
being the input distribution for class $i$), _neural collapse_ (NC) during
training is characterized by the following limit as the number of training
steps $t$ goes to infinity:

$\lim_{t\to\infty}\operatorname{NC}(f_{t},Q)=0\quad\textrm{where}\quad\operatorname{NC}(f,Q):=\operatorname{Avg}_{i\neq
j}\left(V_{f}(Q_{i},Q_{j})\right),$ (3)

and where $f_{t}$ is the learned embedding network at step $t$. Thus, more
neural collapse (i.e., more clustering around better separated class-means)
happens with lower $\operatorname{NC}$ values. Proposition 5 from [22] proves
that $\operatorname{NC}(f,Q)$ bounds a generalization gap and lower values of
$\operatorname{NC}(Q,f)$ are correlated with better test performance of the
neural network $f$, and [36] shows that explicitly regularizing for neural
collapse is beneficial.

##### Geometric complexity in function space:

The _geometric complexity_ (GC) of a function
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ w.r.t. a data distribution
$Q=q(x)dx$ on $\mathbb{R}^{d}$ is defined as the expectation of the gradient
Frobenius square-norm w.r.t. to the data distribution [17] given by

$\operatorname{GC}(f,Q):=\mathbb{E}_{x\sim
Q}\left[\|\nabla_{x}f(x)\|^{2}\right].$

Intuitively, the GC measures the function complexity or variability and is
closely related to the Dirichlet energy in geometric analysis [21]. Provided
mild Lipschitz smoothness assumptions [41, Proposition 3.4], the GC can be
estimated accurately on batches of training data $D$ by its empirical
counterpart

$\widehat{\operatorname{GC}}(f,D)=\frac{1}{|D|}\sum_{x\in
D}\|\nabla_{x}f(x)\|^{2},\quad\text{where }D\text{ is an i.i.d.~{}sample drawn
from }Q.$ (4)

Previous work has explored the relationship between the GC and model
generalization. When measured on the full neural network, lower GC values
correlate experimentally with higher test accuracy [44] and it has also been
used as a beneficial explicit regularizer [56, 30]. In [16, 17], the authors
ignore the softmax activation and study the GC measured with respect to the
logit network, exploring its connection with various implicit and explicit
regularizers and training heuristics. The logit GC has also be used to prove a
margin based multi-class generalization bound [41], assuming that the input
distribution $Q$ satisfies a mild assumption known as the Poincaré inequality
[1].

In fact, the Poincaré inequality plays a central role in the proof of the
generalization bound and we also make use of this property in the current
work. For completeness, we recall it here, cf. [19]. A distribution $Q$
satisfies a Poincaré inequality if, for all differentiable functions $v$
defined on $\operatorname{supp}(Q)$, there exists a constant $c$ such that

$\operatorname{Var}_{v}(Q)\leq c\mathbb{E}_{x\sim
Q}\left[\|\nabla_{x}v\|^{2}\right]=c\operatorname{GC}(v,Q).$ (5)

Note that the same assumption is used in [7] to prove the law of robustness
for overparameterized neural networks via isoperimetric inequalities. In this
paper, we peel back yet another layer of the network and consider the
embedding geometric complexity measured on a feature map (Section 3).

##### Relationship between flatness, neural collapse, and geometric
complexity.

The learning-path flatness as measured by its slope influences the geometric
complexity of the learned solutions [16, Theorem 5.1]. Namely, for dense
layers a regularizing pressure on the slope of the learning path in parameter
space transfers to a regularizing pressure on the geometric complexity of the
learned solution in function space. A similar result also been shown for
special attention layers as well [50].

In this paper, we will see that imposing a regularization pressure on the
embedding geometric complexity encourages more neural collapse of the model.
This results in the following chain of influences from regularization
pressure:

learning path flatness |  $\rightsquigarrow$ |  function geometric complexity |  $\rightsquigarrow$ |  embedding neural collapse
---|---|---|---|---
Regularizing Pressure |  |  Regularizing Pressure |  |  Lower CDNV

##  3 Problem Formulation

We are interested in understanding the relationship between the geometric
complexity (GC) and neural collapse (NC) in the transfer learning setting and
explore this relationship in two stages. In the first stage, we examine the
general impact of the GC on NC during model training/pre-training. Secondly,
we examine how this relationship provides insight into the mechanisms behind
transfer learning and the advantageous implicit biases of the pre-training
stage. In short, lower GC during pre-training on source classes leads to lower
NC for target classes and improved transfer accuracy.

For the first stage of our inquiry (Section 4), we are concerned with a
$k$-classification task. Let $D:=\\{(x_{i},y_{i})\\}_{i=1}^{m}$ be a dataset
drawn from a distribution $Q$ with $x_{i}\in\mathbb{R}^{d}$ and
$y_{i}\in\\{e_{j}~{}|~{}j=1,\dots,k\\}$ where $e_{j}\in\mathbb{R}^{k}$ are the
canonical basis vectors representing a one-hot encoding of the labels. We aim
to learn this task using a neural network denoted by a function
$h_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{k}$ parameterized by $\theta$ (though
to simplify notation we subsequently drop this dependence on $\theta$). We can
write $h_{\theta}(x)=\operatorname{softmax}(g(f(x))$ where
$g:\mathbb{R}^{p}\to\mathbb{R}^{k}$ is a classifier head mapping from the
feature space $\mathbb{R}^{p}$ to the prediction layer $\mathbb{R}^{k}$ and
$f:\mathbb{R}^{d}\to\mathbb{R}^{p}$ is a feature mapping from the input space
to the penultimate embedding layer of the network before the classification
head. The standard way of training $h_{\theta}$ is to find, via some
stochastic optimization technique, an optimal parameter configuration
$\theta_{*}$ such that
$\theta_{*}=\operatorname{argmin}_{\theta}\sum_{i=1}^{m}\ell(h_{\theta}(x_{i}),y_{i})$
for a given loss function
$\ell:\mathbb{R}^{k}\times\mathbb{R}^{k}\to[0,\infty)$.

To analyze the role of neural collapse and geometric complexity in this
setting, we focus our attention on the feature map
$f:\mathbb{R}^{d}\to\mathbb{R}^{p}$. The NC as described in Section 2 is
defined using this sub-network map $f$ of $h$ and thus we also measure the
geometric complexity with respect to this sub-network feature mapping as well.
Throughout this paper, we use the term _embedding GC_ to refer to the GC of
such a feature map $f$ and unless otherwise specified, when we refer to the
model GC we mean the embedding GC, not the GC of the logit network as in [16,
17].

Our second stage of inquiry (Section 5) is focused on the role of geometric
complexity in transfer learning. For this setting, we assume there is some
$l$-class classification target task we would like to solve and corresponding
target distribution $\mathcal{T}$. We aim to learn a classifier
$h^{\prime}:\mathbb{R}^{d}\to\mathbb{R}^{l}$ given some training data
$D^{\prime}:=\\{(x_{i}^{\prime},y_{i}^{\prime})\\}_{i=1}^{m^{\prime}}$ which
is potentially very limited in size. In order to find a good solution, we can
leverage a pre-trained feature mapping that has been trained on some other
source task and source distribution $\mathcal{S}$ where more data is
available. For example, by leveraging a feature map $f$ pre-trained as in the
above setup, the target task classifier $h^{\prime}$ can instead be trained on
the outputs of our feature map; i.e.,
$\\{f(x_{i}^{\prime}),y_{i}^{\prime})\\}_{i=1}^{m^{\prime}}$. In this way, we
can write $h^{\prime}=\operatorname{softmax}(g^{\prime}\circ f)$ where the
parameters of $f$ are fixed and we only need to learn the parameters of
$g^{\prime}:\mathbb{R}^{p}\to\mathbb{R}^{l}$. Ideally, provided $f$ is a rich
enough feature map trained on a general enough source distribution dataset,
then the task for learning $g^{\prime}$ is much simpler.

##  4 Geometric complexity and neural collapse

In this section, we explore the general relationship between the embedding GC
and neural collapse showing that the geometric complexity can be used to bound
neural collapse. Next, we see how the geometric complexity can be efficiently
and accurately estimated, both as an approximation of the theoretical
$\operatorname{GC}$ as well as through a number of sampling techniques for the
empirical $\operatorname{GC}$ via batch sampling, feature sampling and label
sampling. Lastly, following [22], we derive a new generalization bound for
neural networks expressed in terms of the embedding geometric complexity; cf.
[41].

###  4.1 Geometric complexity controls neural collapse

Previously, it has been shown [16] that the $\operatorname{GC}$ can be
controlled implicitly through choice of learning rate and batch size, as well
as through standard explicit regularization. Practically, this means that
these same beneficial tuning strategies which ensure that models have low
$\operatorname{GC}$ should also work to keep the $\operatorname{NC}$ low as
well. The following proposition (which we prove in Appendix A.1) bounds the
neural collapse by the geometric complexity provided that the input
distribution satisfies a Poincaré inequality also assumed in [41] and [7].

######  Proposition 4.1.

Suppose that we have a balanced multi-class input-distribution $Q$ with $k$
classes satisfying the Poincaré inequality in (5) for some constant $c$, then
the geometric complexity of a network embedding $f$ bounds its neural collapse
as measured by (3); namely, we have the following bound

$\operatorname{NC}(f,Q)\leq\frac{c\cdot\operatorname{GC}(f,Q)}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right),$ (6)

where $k$ is the number of classes, and $d_{ij}$ is the distance between the
mean of class $i$ and class $j$.

We call the RHS of the bound in Eq. (6), excepting the Poincaré constant $c$,
the geometric collapse:

$\frac{\operatorname{GC}(f,Q)}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right).$ (7)

This quantity can be seen as another proxy metric measuring neural collapse.
It is made of two main parts decoupling the variances and mean differences
present in the neural collapse framework. The variances are consolidated into
a single factor through the GC, while the mean differences are averaged across
classes in a separate factor. This separation allows the overall intra-class
variability to be influenced through the GC term, ensuring sufficient between-
class separation through the mechanisms identified in neural collapse.

This bound provides a powerful method to control the overall within-class
variability via the L2-norm of the model gradient. As a result, the geometric
collapse provides a simplified and more refined approach to managing class
separability and variance in deep learning models.

We verify the relationship posed in Eq. (6) through experiments on VGG-13
trained on CIFAR-10 (see Figure 1). Through both implicit and explicit
regularizers the pressure on the embedding $\operatorname{GC}$ translates to a
direct pressure on the neural collapse of the model via the geometric
collapse. In Appendix A.5 we see the same relationship holds across various
datasets; e.g., MNIST, Fashion-MNIST, CIFAR-100, and architectures; e.g.,
ResNet-18, ResNet-50.

Figure 1: Controlling the neural collapse via the model geometric complexity
for VGG-13 trained on CIFAR-10. Lower embedding $\operatorname{GC}$ produces
lower geometric collapse (Eq. 7) and more neural collapse (i.e., lower
$\operatorname{NC}$) for Top row: increased learning rates, Middle row:
decreased batch sizes, and Bottom row: increased L2 regularization.

###  4.2 The geometric complexity is a reliable and robust measure

One of the key advantages of the GC as a complexity measure is its
computational efficiency. Under mild regularity constraints on the model
function the theoretical geometric complexity of a map can be efficiently
estimated by its empirical counterpart, as stated in the proposition below,
already proven in [41] for logit networks (see Appendix A.2 for a proof).
Additionally and crucially, we verify that the empirical $\operatorname{GC}$
is a robust and reliable measure with respect to that data sample.

######  Proposition 4.2.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{p}$ be a map with Lipschitz
constant $L$ and let $D_{X}$ be a sample of $m$ elements drawn independently
from an input distribution $Q$. Then, for any $\delta>0$, we have with
probability $1-\delta/2$ the following bound

$\operatorname{GC}(f,Q)\leq\widehat{\operatorname{GC}}(f,D_{X})+L\sqrt{\frac{\log\frac{2}{\delta}}{2m}}.$
(8)

For a function $f$ and data sample $D_{X}$ as in Proposition 4.2, to measure
$\widehat{\operatorname{GC}}(f,D_{X})$ requires computing the Frobenius norm
of a potentially very large Jacobian matrix, particularly when $p$ represents
the embedding dimension of a feature map. Note that,

$\widehat{\operatorname{GC}}(f,D_{X})=\frac{1}{m}\sum_{i=1}^{m}\sum_{j=1}^{p}\sum_{s=1}^{d}\left(\frac{\partial
f^{j}(x^{(i)})}{\partial x_{s}}\right)^{2}.$ (9)

Although the computational complexity of the GC is impervious to increasing
complexity of the network architecture, e.g., in terms of parameter count or
number of layers, one may run into computation bottlenecks as the sample size
$m$ increases or when increasing the dimensionality $d$ (resp. $p$) of the
inputs (resp. outputs). In these scenarios, it is necessary to find efficient
means to accurately approximate or sample the Jacobian; cf. [58].

We explore the robustness of the GC when measured via samplings along these
three axes; i.e., decreasing the number of examples $m$ in the batch, randomly
sampling the full Jacobian matrix which has order $d\times p$, and randomly
sampling the number of model outputs $p$. As shown Figure 2, we find that the
value of the sampled $\widehat{\operatorname{GC}}$ remains stable and
consistent to its true value through these fairly simple and naive sampling
tricks.

Figure 2: The $\operatorname{GC}$ computation is robust and consistent to
sampling via Left: number of examples in the batch, Middle: number of elements
in the Jacobian, or Right: by sampling the embedding dimension of the model.
Here the $\operatorname{GC}$ and subnet $\operatorname{GC}$ have been computed
over 20 trials, plotting the mean and standard deviations for a ResNet-18
model that has been trained to convergence on CIFAR-100. The true value of the
$\operatorname{GC}$ for each setting is indicated by dotted line.

###  4.3 A new generalization bound with GC through NC

In [41], the authors derive a margin-based multi-class generalization bound
for neural networks which depends on the geometric complexity
$\operatorname{GC}(h,Q)$ of the full logit network. In this section, we extend
this result. With inspiration from [22, Proposition 5], we show that the
geometric complexity measured on the sub-network feature map
$\operatorname{GC}(f,Q)$ bounds the classification error of the nearest-mean
classifier defined by the feature map $f$.

As in [22], given a balanced $k$-class classification problem, let
$S=\bigcup_{c\in[k]}S_{c}$ denote all class samples and define the nearest-
mean classifier by
$h_{f,S}(x):=\operatorname{argmin}_{c\in[k]}\|f(x)-\mu_{f}(S_{c})\|$ given by
the feature map $f$ where $\mu_{f}(S_{c})$ denotes the sample mean of the set
$S_{c}$ under the map $f$. Define the generalization error

$\operatorname{Error}:=\mathbb{E}_{(x,y)\sim P}\left[\mathbb{I}[h_{f,S}(x)\neq
y]\right]$

where $\mathbb{I}[h_{f,S}(x)\neq y]$ is the indicator function of the error
set $[h_{f,S}(x)\neq y]$ and $P$ is the full data distribution from which
samples $S$ are drawn.

######  Proposition 4.3.

Suppose that we have a balanced sample $S$ from a $k$-class input-distribution
$(x,y)\sim P$ with $m_{c}$ samples per class. Assume further that the induced
input distribution $x\sim Q$ satisfies a Poincaré inequality as in (5) for
some constant $c$. Then, for any $\delta>0$, with probability $1-\delta/2$ we
have the following bound for the generalization error

$\mathbb{E}\left[\operatorname{Error}\right]\leq
16c\left(\frac{1}{p}+\frac{1}{m_{c}}\right)\left(\widehat{\operatorname{GC}}(f,S)+L\sqrt{\frac{\log\frac{1}{\delta}}{2m_{c}k}}\right)\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right),$ (10)

where $p$ is the embedding dimension of the feature map $f$ and
$d_{ij}=\|\mu_{f}(S_{i})-\mu_{f}(S_{j})\|$.

###### Proof.

This follows immediately from [22, Proposition 5] which gives the bound in
terms of the neural collapse instead of the geometric complexity. By using
Proposition A.1 we can replace the neural collapse with the geometric
complexity and using Proposition 4.2 we then replace the theoretical GC with
the empirical GC, yielding the additional term with the Poincaré constant. ∎

##  5 The impact of geometric complexity on transfer learning

Many implicit regularization mechanisms in gradient descent exert a pressure
on geometric complexity which in turn constrain the neural collapse. In this
section, we discuss the affect of this implicit bias in transfer learning.
Specifically, we show with theory how this bias during pre-training can help
explain the mechanisms behind transfer learning. Namely, that lower GC in a
pre-trained embedding network promotes neural collapse on new target classes,
simplifying the fine-tuning process.

However, our bound makes it clear that certain compatibility conditions
between the pre-training (source) distribution and the fine-tuning (target)
distribution are needed, which we argue are satisfied in image foundational
models and large language models. At last, we verify empirically that known
methods to reduce geometric complexity for the pre-trained embedding result in
better performance during fine-tuning on unseen tasks, in agreement with our
neural collapse bound for transfer learning.

###  5.1 Lower pre-trained GC leads to improved fine-tuning

To analyze the impact of GC in the context of transfer learning, we consider
the same formal setup as in [22, Proposition 2]. Assume a class $c$ in a set
$\mathcal{C}$ of classes is represented by a class-conditional distribution
$Q_{c}(x)=P(x|y=c)$. Both the pre-training/source classes
$\widetilde{Q}_{1},\dots,\widetilde{Q}_{k}$ and the fine-tuning/target classes
$Q_{1},\dots,Q_{l}$ are assumed to be drawn from a distribution over all
classes in $\mathcal{C}$. The induced input distribution on all the classes
(i.e., the combined source and target input distribution) is denoted $Q$,
while the source-only input distribution is denoted by $\widetilde{Q}$. Let
$\mathcal{F}^{*}$ denote the set of pre-trained feature maps
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{p}$ selected by the training procedure
(e.g., the set of trained embeddings with different initialization seeds but
the same training protocol) and consider

$\Delta(\mathcal{F}^{*})=\inf_{f\in\mathcal{F}^{*}}\inf_{c\neq
c^{\prime}}\|\mu_{f}(Q_{c})-\mu_{f}(Q_{c^{\prime}})\|.$ (11)

With this, we can state the transfer learning bound

######  Proposition 5.1.

Suppose that the source and target input distributions satisfy a Poincaré
inequality with constant $c_{\tilde{Q}}$ and $c_{Q}$ (resp.). Then, with
probability $1-\delta$, over the selection $Q_{c}$ and $Q_{c^{\prime}}$ of
target class distributions, we have that the CDNV expectation for two target
classes is bounded by the geometric complexity of the embedding network $f$ in
the following way

$\displaystyle\mathbb{E}_{Q_{c}\neq
Q_{c^{\prime}}}\left[V_{f}(Q_{c},Q_{c^{\prime}})\right]$ $\displaystyle\leq$
$\displaystyle\frac{c_{\tilde{Q}}\operatorname{GC}(f,\widetilde{Q})}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right)$ $\displaystyle+\left(8+\frac{16c_{Q}\sup_{c\in
C}\operatorname{GC}(\mathcal{F}^{*},Q_{c})}{\Delta(\mathcal{F}^{*})}\right)\left(\frac{\sqrt{2\pi\log(k)}\mathfrak{H}(\mathcal{F}^{*},\widetilde{Q})}{(k-1)\Delta(\mathcal{F}^{*})}\right)$
$\displaystyle+\left(1+\dfrac{4\sup_{\begin{subarray}{c}x\in\operatorname{supp}(Q)\\\
f\in\mathcal{F}^{*}\end{subarray}}\|f(x)\|}{\Delta(\mathcal{F}^{*})^{2}}\right)\left(\frac{2\sqrt{\log\frac{1}{\delta}}c_{Q}\sup_{\begin{subarray}{c}c\in\mathcal{C}\end{subarray}}\operatorname{GC}(\mathcal{F}^{*},Q_{c})}{\sqrt{k}\Delta(\mathcal{F}^{*})^{2}}\right)$

where $k$ is the number of source classes,
$\operatorname{GC}(\mathcal{F}^{*},Q_{c})=\sup_{f\in\mathcal{F}^{*}}\operatorname{GC}(f,Q_{c})$,
$d_{ij}=\|\mu_{f}(Q_{c_{i}})-\mu_{f}(Q_{c_{j}})\|$, and
$\mathfrak{H}(\mathcal{F}^{*},\widetilde{Q})$ is a complexity measure for
$\mathcal{F}^{*}$ over $\widetilde{Q}$ (see Appendix A.3).

###### Proof.

This follows immediately from [22, Proposition 2] which proves this bound for
the NC and distribution variances. Using our Proposition 4.1, we can replace
the NC by the GC and swap the variances by the corresponding geometric
complexities via the respective Poincaré inequality. ∎

##### Interpretation.

In the bound above, the first term depends on the geometric complexity
measured over the source distribution $\tilde{Q}$ and decreases as the number
of source classes $k$ increases. Thus, we can predict that lower
$\operatorname{GC}(f,\tilde{Q})$ encourages lower NC values (i.e. more neural
collapse) on the new target classes provided that the second and third term of
the bound are also small.

However these other two terms involve both the source and target classes,
making the full bound no longer dependent only on the geometric complexity
over $\tilde{Q}$. Moreover these two terms can apriori explode since the
infimum over the class-mean distances $\Delta(\mathcal{F}^{*})$ no longer only
involves source classes (where this distance is bounded away from zero due to
neural collapse into an equidistant simplex) but also target classes (for
which we no longer have theoretical guarantees).

However, if the source labels are granular enough in the sense the target
labels can be represented as combinations of the source labels (as for
instance a target label of “dog” can be subsumed as all the breeds of dogs in
source classes on ImageNet), then the issues above are mitigated. Formally, we
can summarize this granularity compatibility condition between the source and
target classes as follows: For each label $c$ of the target class there exists
labels $c_{1},\dots,c_{k}$ in the source classes such that
$Q_{c}=1/k(Q_{c_{1}}+\cdots+Q_{c_{k}})$. Because $\operatorname{GC}(f,\cdot)$
is linear, we have

$\sup_{c\in\mathcal{C}}\operatorname{GC}(\mathcal{F}^{*},Q_{c})\leq\sup_{c\in[k]}\operatorname{GC}(\mathcal{F}^{*},\tilde{Q}_{c})$

and thus recover dependence only on source classes.

To address the $\Delta(\mathcal{F}^{*})$ term, observe that the compatibility
condition above implies in terms of class-means that
$\mu_{f}(Q_{c})=\frac{1}{k}(\mu_{f}(Q_{c_{1}})+\cdots+\mu_{f}(Q_{c_{k}}))$.
Because of neural collapse during pre-training, these source class-means form
the vertices of a face in the neural-collapse class-mean simplex making the
target-class mean $\mu_{f}(Q_{c})$ the barycenter of this face. Since the
distance between any two points taken in the vertices of a simplex plus its
barycenters is bounded away from zero, so is $\Delta(\mathcal{F}^{*})$. Note,
albeit intuitive and natural, this granularity condition is hard to verify in
practice.

###  5.2 Improve fine-tuning by controlling pre-trained GC

When the compatibility conditions above are satisfied, Proposition 5.1
indicates that increasing the amount of neural collapse for the target classes
can be achieved by lowering the embedding GC during pre-training. We verify
this indeed occurs experimentally using the same regularization techniques
exploited in Section 4. Furthermore, we verify that these implicit methods of
controlling pre-training GC produce feature maps that perform better on fine-
tuning tasks on CIFAR-FS with a ResNet-18 in Figure 3 and on mini-ImageNet
with VGG-16 in Appendix A.5.5.

Figure 3: Controlling target neural collapse through source GC on CIFAR-FS
with ResNet-18. Lower Source GC produces more neural collapse on target
classes (i.e. lower target NC), and higher 5-shot transfer accuracy for Top
row: increased learning rates, Middle row: decreased batch sizes, and Bottom
row: increased L2 regularization.

##  6 Conclusion

Uncovering and understanding the implicit biases that enable successful
transfer learning is a critical area of research in modern machine learning.
Here we provide a framework connecting three different themes behind implicit
regularization: flatness of the loss surface, geometric complexity, and neural
collapse. We show that the embedding geometric complexity directly controls
the neural collapse during training and, moreover, plays a role in the success
of transfer learning. Extensive experiments on different image classification
and fine tuning tasks across different hyperparameters verify our hypothesis.
We believe this opens up an intriguing direction of research further exploring
the role of geometric collapse and geometric complexity in deep learning and
provides valuable insight for designing more efficient and effective
techniques for model training and fine-tuning.

## References

  * Bakry et al. [2008] Dominique Bakry, Franck Barthe, Patrick Cattiaux, and Arnaud Guillin.  A simple proof of the poincaré inequality for a large class of probability measures.  2008\.
  * Barak et al. [2022] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang.  Hidden progress in deep learning: Sgd learns parities near the computational limit.  _Advances in Neural Information Processing Systems_ , 35:21750–21764, 2022.
  * Barrett and Dherin [2021] David G.T. Barrett and Benoit Dherin.  Implicit gradient regularization.  In _International Conference on Learning Representations_ , 2021.
  * Bengio [2012] Yoshua Bengio.  Deep learning of representations for unsupervised and transfer learning.  In _Proceedings of ICML workshop on unsupervised and transfer learning_ , pages 17–36. JMLR Workshop and Conference Proceedings, 2012.
  * Blanc et al. [2019] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In _Annual Conference Computational Learning Theory_ , 2019.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Bubeck and Sellke [2023] Sébastien Bubeck and Mark Sellke.  A universal law of robustness via isoperimetry.  _Journal of the ACM_ , 70(2):1–18, 2023.
  * Carlini et al. [2021] Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, et al.  Extracting training data from large language models.  In _30th USENIX Security Symposium (USENIX Security 21)_ , pages 2633–2650, 2021.
  * Caruana [1994] Rich Caruana.  Learning many related tasks at the same time with backpropagation.  _Advances in neural information processing systems_ , 7, 1994.
  * Cattaneo et al. [2023] Matias Cattaneo, Jason Klusowski, and Boris Shigida.  On the implicit bias of adam.  _arXiv:2309.00079_ , 2023.
  * Chelba et al. [2013] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson.  One billion word benchmark for measuring progress in statistical language modeling.  _arXiv preprint arXiv:1312.3005_ , 2013.
  * Choromanska et al. [2015] Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun.  The loss surfaces of multilayer networks.  In _Artificial intelligence and statistics_ , pages 192–204. PMLR, 2015.
  * Damian et al. [2021] Alex Damian, Tengyu Ma, and Jason D. Lee.  Label noise SGD provably prefers flat global minimizers.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, _NeurIPS 2021_ , 2021.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pages 248–255. Ieee, 2009.
  * Devlin et al. [2018] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dherin et al. [2021] Benoit Dherin, Micheal Munn, and David GT Barrett.  The geometric occam’s razor implicit in deep learning.  _NeurIPS, OPT2021_ , 2021.
  * Dherin et al. [2022] Benoit Dherin, Michael Munn, Mihaela Rosca, and David Barrett.  Why neural networks find simple solutions: the many regularizers of geometric complexity.  _Advances in Neural Information Processing Systems_ , 35:2333–2349, 2022.
  * Dinh et al. [2017] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.  Sharp minima can generalize for deep nets.  In _International Conference on Machine Learning_ , pages 1019–1028. PMLR, 2017.
  * Evans [2022] Lawrence C Evans.  _Partial differential equations_ , volume 19.  American Mathematical Society, 2022.
  * Farahani et al. [2021] Abolfazl Farahani, Behrouz Pourshojae, Khaled Rasheed, and Hamid R. Arabnia.  A concise review of transfer learning.  _CoRR_ , abs/2104.02144, 2021.  URL https://arxiv.org/abs/2104.02144.
  * Federer [2014] Herbert Federer.  _Geometric measure theory_.  Springer, 2014.
  * Galanti et al. [2022] Tomer Galanti, András György, and Marcus Hutter.  On the role of neural collapse in transfer learning.  In _ICLR_ , 2022.
  * Gao et al. [2023] Peifeng Gao, Qianqian Xu, Yibo Yang, Peisong Wen, Huiyang Shao, Zhiyong Yang, Bernard Ghanem, and Qingming Huang.  Towards demystifying the generalization behaviors when neural collapse emerges.  _arXiv preprint arXiv:2310.08358_ , 2023.
  * [24] Jonas Geiping, Micah Goldblum, Phil Pope, Michael Moeller, and Tom Goldstein.  Stochastic training is not necessary for generalization.  In _ICLR 2022_.
  * Ghosh et al. [2023] Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang.  Implicit regularization in heavy-ball momentum accelerated stochastic gradient descent.  _ICLR_ , 2023.
  * Gunasekar et al. [2018] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro.  Implicit bias of gradient descent on linear convolutional networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Han et al. [2022] X.Y. Han, Vardan Papyan, and David L. Donoho.  Neural collapse under MSE loss: Proximity to and dynamics on the central path.  In _ICLR_ , 2022.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber.  Flat minima.  _Neural computation_ , 9(1):1–42, 1997.
  * Hoffman et al. [2020] Judy Hoffman, Daniel A. Roberts, and Sho Yaida.  Robust learning with jacobian regularization.  2020\.  URL https://arxiv.org/abs/1908.02729.
  * Kothapalli [2022] Vignesh Kothapalli.  Neural collapse: A review on modelling principles and generalization.  _arXiv preprint arXiv:2206.04041_ , 2022.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Lee et al. [2019] Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto.  Meta-learning with differentiable convex optimization.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10657–10665, 2019.
  * Li et al. [2022] Xiao Li, Sheng Liu, Jinxin Zhou, Xinyu Lu, Carlos Fernandez-Granda, Zhihui Zhu, and Qing Qu.  Principled and efficient transfer learning of deep models via neural collapse.  _arXiv preprint arXiv:2212.12206_ , 2022.
  * Liu et al. [2023a] Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma.  Same pre-training loss, better downstream: Implicit bias matters for language models.  In _ICML_ , 2023a.
  * Liu et al. [2023b] Xuantong Liu, Jianfeng Zhang, Tianyang Hu, He Cao, Yuan Yao, and Lujia Pan.  Inducing neural collapse in deep long-tailed learning.  In _AISTATS_ , 2023b.
  * Ma and Ying [2021a] Chao Ma and Lexing Ying.  The sobolev regularization effect of stochastic gradient descent.  2021a.  URL https://arxiv.org/abs/2105.13462.
  * Ma and Ying [2021b] Chao Ma and Lexing Ying.  On linear stability of sgd and input-smoothness of neural networks.  In _NeurIPS_ , 2021b.
  * Martin and Mahoney [2021] Charles H Martin and Michael W Mahoney.  Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning.  _Journal of Machine Learning Research_ , 22(165):1–73, 2021.
  * Mohri et al. [2018] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar.  _Foundations of machine learning_.  MIT press, 2018.
  * Munn et al. [2023] Michael Munn, Benoit Dherin, and Javier Gonzalvo.  A margin-based multiclass generalization bound via geometric complexity.  _ICML, TAGML Workshop_ , 2023.
  * Neyshabur [2017] Behnam Neyshabur.  Implicit regularization in deep learning.  _arXiv preprint arXiv:1709.01953_ , 2017.
  * Neyshabur et al. [2014] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro.  In search of the real inductive bias: On the role of implicit regularization in deep learning.  _arXiv preprint arXiv:1412.6614_ , 2014.
  * Novak et al. [2018] Roman Novak, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Sensitivity and generalization in neural networks: an empirical study.  In _ICLR_ , 2018.
  * Papyan et al. [2020] Vardan Papyan, X. Y. Han, and David L. Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  _Proceedings of the National Academy of Sciences_ , 117(40), 2020.
  * Power et al. [2022] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra.  Grokking: Generalization beyond overfitting on small algorithmic datasets.  _arXiv preprint arXiv:2201.02177_ , 2022.
  * Raffel et al. [2020] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of machine learning research_ , 21(140):1–67, 2020.
  * Ramesh et al. [2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In _International conference on machine learning_ , pages 8821–8831. Pmlr, 2021.
  * Ravi and Larochelle [2016] Sachin Ravi and Hugo Larochelle.  Optimization as a model for few-shot learning.  In _International conference on learning representations_ , 2016.
  * Ren et al. [2023] Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai.  Sparse modular activation for efficient sequence modeling.  In _NeurIPS_ , 2023.
  * Rosca et al. [2020] Mihaela Rosca, Theophane Weber, Arthur Gretton, and Shakir Mohamed.  A case for new neural network smoothness constraints.  In _Proceedings on "I Can’t Believe It’s Not Better!" at NeurIPS Workshops_ , volume 137 of _Proceedings of Machine Learning Research_ , 2020\.
  * Shirish et al. [2016] Keskar Nitish Shirish, Mudigere Dheevatsa, Nocedal Jorge, Smelyanskiy Mikhail, and Tang Ping Tak Peter.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_ , 2016.
  * Simonyan and Zisserman [2014] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Simonyan and Zisserman [2015] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  In _ICLR_ , 2015.
  * Smith et al. [2021] Samuel L Smith, Benoit Dherin, David G.T. Barrett, and Soham De.  On the origin of implicit regularization in stochastic gradient descent.  In _International Conference on Learning Representations_ , 2021.
  * Sokolic et al. [2017] Jure Sokolic, Raja Giryes, Guillermo Sapiro, and Miguel R. D. Rodrigues.  Robust large margin deep neural networks.  _IEEE Transactions on Signal Processing_ , 65:4265–4280, 2017.
  * Vinyals et al. [2016] Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al.  Matching networks for one shot learning.  _Advances in neural information processing systems_ , 29, 2016.
  * Wang and Fessler [2023] Guanhua Wang and Jeffrey A Fessler.  Efficient approximation of jacobian matrices involving a non-uniform fast fourier transform (nufft).  _IEEE transactions on computational imaging_ , 9:43–54, 2023.
  * Wang et al. [2023] Zijian Wang, Yadan Luo, Liang Zheng, Zi Huang, and Mahsa Baktashmotlagh.  How far pre-trained models are from neural collapse on the target dataset informs their transferability.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 5549–5558, 2023.
  * Weiss et al. [2016] Karl Weiss, Taghi Khoshgoftaar, and DingDing Wang.  A survey of transfer learning.  _Journal of Big Data_ , 3, 05 2016.  doi: 10.1186/s40537-016-0043-6.
  * Wilson et al. [2017] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht.  The marginal value of adaptive gradient methods in machine learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Yosinski et al. [2014] Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson.  How transferable are features in deep neural networks?  _Advances in neural information processing systems_ , 27, 2014.
  * Zhang et al. [2021] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning (still) requires rethinking generalization.  _Communications of the ACM_ , 64(3):107–115, 2021\.

##  Appendix A Appendix

###  A.1 Proof of Proposition 3.2 (Neural collapse bound)

###### Proof.

Suppose that we have an input probability distribution $Q(x)=q(x)dx$ coming
from a data distribution $P(X,Y)$ (i.e. $q(x)dx$ is the marginal probability
distribution of $p(x,y)dxdy$) where $Y$ represents the $k$ possible classes.
Suppose further that $P(X,Y)$ is a balanced multi-class distribution. In terms
of the input distribution $Q(x)$ this means that

$Q=\frac{1}{k}(Q_{1}+\cdots+Q_{k})$

where $Q_{i}(x)=q_{i}(x)dx$ is the input distribution of class $i$. Suppose
moreover that $Q(x)$ satisfies the Poincaré inequality in (5), that is,

$\operatorname{Var}_{f}(Q)\leq
c\mathbb{E}_{Q}(\|\nabla_{x}f\|^{2})=c\operatorname{GC}(f,Q)$

for some constant $c$. Then the statement we want to prove is that the
geometric complexity of $f$, that is,

$\operatorname{GC}(f,Q):=\int\|\nabla_{x}f(x)\|^{2}q(x)dx$

bounds its neural collapse as measured by

$\operatorname{NC}(f,Q):=\frac{1}{\\#\\{i\neq j\\}}\sum_{i\neq
j}\left(\frac{\operatorname{Var}_{f}(Q_{i})+\operatorname{Var}_{f}(Q_{j})}{2\|\mu_{f}(Q_{i})-\mu_{f}(Q_{j})\|^{2}}\right)$
(12)

More precisely, we want to prove that we have the following inequality:

$\operatorname{NC}(f,Q)\leq\frac{c\cdot\operatorname{GC}(f,Q)}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right),$ (13)

where $d_{ij}=\|\mu_{f}(Q_{i})-\mu_{f}(Q_{j})\|$ is the distance between the
mean of class $i$ and class $j$. Let us now prove this statement.

First of all, since the geometric complexity respects convex sums of data
densities, we have that for a distribution $Q=\frac{1}{k}(Q_{1}+\cdots Q_{k})$
we can write

$\operatorname{GC}(f,Q)=\frac{1}{k}(\operatorname{GC}(f,Q_{1})+\cdots+\operatorname{GC}(f,Q_{k})).$
(14)

In particular, this means that $\operatorname{GC}(f,Q_{i})\leq
k\operatorname{GC}(f,Q)$. Furthermore, the Poincare inequality ensures that
$\operatorname{Var}_{Q_{i}}(f)\leq c\operatorname{GC}(f,Q_{i})$. Using these
two properties and the definition of the neural collapse, we obtain that:

$\displaystyle\operatorname{NC}(f,Q)$ $\displaystyle=$
$\displaystyle\frac{1}{\\#\\{i\neq j\\}}\sum_{i\neq
j}\left(\frac{\operatorname{Var}_{f}(Q_{i})+\operatorname{Var}_{f}(Q_{j})}{2\|\mu_{f}(Q_{i})-\mu_{f}(Q_{j})\|^{2}}\right)$
$\displaystyle\leq$ $\displaystyle\frac{c}{\\#\\{i\neq j\\}}\sum_{i\neq
j}\frac{\operatorname{GC}(f,Q_{i})+\operatorname{GC}(f,Q_{j})}{2d_{ij}^{2}}$
$\displaystyle\leq$
$\displaystyle\frac{ck\operatorname{GC}(f,Q)}{k(k-1)}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right),$

which simplifies to the desired result. ∎

######  Remark A.1.

Excepting the Poincaré constant, the quantity on the RHS above given by

$\frac{\operatorname{GC}(f,Q)}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right)$ (15)

could be taken as an alternative measure of neural collapse. We call this the
geometric collapse.

###  A.2 Proof of Proposition 4.2 (Estimating the Theoretical
$\operatorname{GC}$ by $\widehat{\operatorname{GC}}$)

###### Proof.

Let $Q$ be an (input) probability distribution defined over $\mathbb{R}^{d}$
and let $D=\\{x_{i}\\}_{i=1}^{m}$ of $m\geq 1$ points drawn as i.i.d. samples
from $Q$. Given the map $f:\mathbb{R}^{d}\to\mathbb{R}^{p}$, we denote the
empirical geometric complexity of $f$ over $D$ by
$\widehat{\operatorname{GC}}(f,D)$ which is defined as

$\widehat{\operatorname{GC}}(f,D)=\dfrac{1}{m}\sum_{i=1}^{m}\|\nabla_{x}f(x_{i})\|^{2}_{F}.$

We begin by showing that

$\mathbb{E}_{D\sim
Q^{m}}\left[\widehat{\operatorname{GC}}(f,D)\right]=\operatorname{GC}(f,Q).$

This follows by direct computation, keeping in mind that $Q$ is a probability
distribution and that the points are independently sampled. We can write $Q$
as $Q(x)=q(x)dx$ and note that

$\displaystyle\mathbb{E}_{D\sim
Q^{m}}\left[\widehat{\operatorname{GC}}(f,D)\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{x_{1},\dots,x_{m}\sim
Q^{m}}\left[\dfrac{1}{m}\sum_{i=1}^{m}\|\nabla_{x}f(x_{i})\|^{2}_{F}\right]$
$\displaystyle=$ $\displaystyle\dfrac{1}{m}\int_{\mathbb{R}^{m\times
d}}\sum_{i=1}^{m}\|\nabla_{x}f(x_{i})\|^{2}_{F}dQ^{m}(x_{1},\dots,x_{m})$
$\displaystyle=$ $\displaystyle\dfrac{1}{m}\int_{\mathbb{R}^{m\times
d}}\sum_{i=1}^{m}\|\nabla_{x}f(x_{i})\|^{2}_{F}q(x_{1})\cdots
q(x_{m})dx_{1}\cdots dx_{m}$ $\displaystyle=$
$\displaystyle\dfrac{1}{m}\sum_{i=1}^{m}\int_{\mathbb{R}^{(m-1)\times
d}}\left[\int_{\mathbb{R}^{d}}\|\nabla_{x}f(x_{i})\|^{2}_{F}q(x_{i})dx_{i}\right]q(x_{1})\cdots\widehat{q(x_{i})}\cdots
q(x_{m})dx_{1}\cdots\widehat{dx_{i}}\cdots dx_{m}$ $\displaystyle=$
$\displaystyle\dfrac{1}{m}\sum_{i=1}^{m}\left[\int_{\mathbb{R}^{d}}\|\nabla_{x}f(x_{i})\|^{2}_{F}q(x_{i})dx_{i}\right]$
$\displaystyle=$
$\displaystyle\dfrac{1}{m}\sum_{i=1}^{m}\left[\int_{\mathbb{R}^{d}}\|\nabla_{x}f(x_{i})\|^{2}_{F}dQ(x_{i})\right]$
$\displaystyle=$
$\displaystyle\dfrac{1}{m}\sum_{i=1}^{m}\operatorname{GC}(f,Q)$
$\displaystyle=$ $\displaystyle\operatorname{GC}(f,Q).$

Now suppose we have two separate independent samples $D$ and $D^{\prime}$ each
of of size $m\geq 1$ which differ by exactly one point, say $x_{i}$ in $D$ and
$x_{i}^{\prime}$ in $D^{\prime}$. By assumption, the map $f$ is $L$-Lipschitz
for some Lipschitz constnt $L>0$. Therefore, we have

$\widehat{\operatorname{GC}}(f,D)-\widehat{\operatorname{GC}}(f,D^{\prime})=\dfrac{1}{m}\left(\|\nabla_{x}f(x_{i})\|^{2}_{F}-\|\nabla_{x}f(x_{i}^{\prime})\|^{2}_{F}\right)\leq
L^{2}/m,$

and similarly,
$\widehat{\operatorname{GC}}(f,D^{\prime})-\widehat{\operatorname{GC}}(f,D)\leq
L^{2}/m$. It follows that
$|\widehat{\operatorname{GC}}(f,D)-\widehat{\operatorname{GC}}(f,D^{\prime})|\leq
L^{2}/m$ and by applying McDiarmind’s inequality (e.g., [40]), we have that
for any $\epsilon>0$,

$\mathbb{P}\left[\widehat{\operatorname{GC}}(f,D)-\mathbb{E}_{D\sim\mu^{m}}[\widehat{\operatorname{GC}}(f,D)]\leq\epsilon\right]\geq
1-\exp(-2m\epsilon^{2}/L^{2}).$ (16)

Since, from the computation above,
$\mathbb{E}_{D\sim\mu^{m}}[\widehat{\operatorname{GC}}(f,D)]=\operatorname{GC}(f,Q)$
and setting $\delta/2=\exp(-2m\epsilon^{2}/L^{2})$ and substituting for
$\epsilon$ in (16), we get that for any $\delta>0$ with probability as least
$1-\delta/2$ the following holds:

$\operatorname{GC}(f,Q)\leq\widehat{\operatorname{GC}}(f,D)+L\sqrt{\dfrac{\log\frac{2}{\delta}}{2m}}.$

This completes the proof. ∎

###  A.3 Definition of the complexity measure $\mathfrak{H}(\mathcal{F},Q)$ in
Proposition 4.1

Let us restate here for the sake of completeness the definition given in
Galanti et al. [22]. First consider the Rademacher complexity $R(A)$ of a set
$A\subset\mathbb{R}^{n}$ which is given as the expectation
$\mathbb{E}_{\epsilon}(X)$ of the random variable $X=\sup_{a\in
A}\langle\epsilon,a\rangle$ where $\epsilon=(\epsilon_{1},\dots,\epsilon_{n})$
is a random vector with components uniformly distributed among the two values
-1 and 1. Now consider the input distribution $Q$ coming from a multiclass
distribution with labels in $c\in\mathcal{C}$ with label input distribution
$Q_{c}$. We define the complexity measure $\mathfrak{H}(\mathcal{F},Q)$ for
the candidate functions in $\mathcal{F}$ over the source distribution $Q$ as
the Rademacher complexity of the following set:

$\left\\{\left(\mu_{f}(Q_{c}),\,\operatorname{Var}_{f}(Q_{c})):\>c\in\mathcal{C},\,f\in\mathcal{F}\right)\right\\}.$

###  A.4 Experiment details

####  A.4.1 Figure 1: Geometric Complexity Controls Neural Collapse on
CIFAR-10.

We trained a VGG-13 neural network on the full CIFAR-10 dataset with the
provided architecture [53] and using the standard train/test split. Throughout
training, we reported the following metrics measured and averaged over
multiple batches of the training dataset: 1) the geometric complexity of the
model embedding layer; i.e., the layer on which the neural collapse is
measured, 2) the neural collapse as measured by (3) measured on the model
embedding layer, and 3) the geometric collapse of the model also measured on
the penultimate embedding layer of the model and given by

$\frac{\operatorname{GC}(f,Q)}{k-1}\left(\sum_{i\neq
j}\frac{1}{d_{ij}^{2}}\right)$ (17)

where $Q$ is the training distribution, $k=10$ as this was CIFAR-10, and (as
in the paper) where $d_{ij}=\|\mu_{f}(Q_{i})-\mu_{f}(Q_{j})\|$ is the distance
between the sample means of class $i$ and class $j$.

These quantities were measured every 1000 steps of a very long pre-training of
$100,000$ steps. The optimizer was plain SGD (note we obtained similar results
with momentum) without any regularization nor schedule to avoid masking
effects. We used random crop and random flip for data augmentation. The
results were all averaged over 5 random seeds for the neural network parameter
default initialization. The plots have no smoothing applied to the learning
curves. Top row: We swept over a learning rate range of
$\\{0.001,0.0025,0.005,0.01,0.025,0.1\\}$ with a constant batch size of 512.
Middle row: We swept over a batch size range of $\\{8,16,32,64,128,256\\}$
with a constant learning rate of 0.01. Bottom row: We swept over a L2
regularization rate range of $\\{0.0,0.00025,0.0005,0.001,0.0025\\}$ with
learning rate 0.01 and batch size 256. Each sweep took roughly 10h of training
on a single Google Cloud TPU V3 accessed via a Google colab.

####  A.4.2 Figure 3: Controlling target performance through source GC on
CIFAR-FS.

We trained a RestNet-18 neural network with width 1 implemented in Flax
https://github.com/google/flax/blob/main/examples/imagenet/models.py on CIFAR-
FS with its initial convolution modified to stride 1 and kernel size of 3 to
adapt to CIFAR-FS instead of ImageNet. We used only 10 classes for the source
datasets with 600 images per class, which further we split in a 10/90 split
for the test and train splits. The metrics were computed on an average of 100
random samples of 5 classes in the remaining classes of CIFAR-FS. We reported
1) the geometric complexity of the embedding layer measured on the source
dataset (Source GC), 2) the neural collapse as measured by (3) measured and
averaged on the 100 random samples of 5-label samples from the target dataset,
3) the test accuracy obtained by solving the normal equation directly for a
ridge regression with only 5 examples per class obtained from the target set.
These three quantities were measured every 1000 steps of a very long pre-
training of $100_{0}00$ steps. The optimizer was plain SGD (note we obtained
similar results with momentum) without any regularization nor schedule to
avoid masking effects. We used random crop and random flip for data
augmentation. The results were all averaged over 5 random seeds for the neural
network parameter default initialization. The plots have no smoothing applied
to the learning curves. Top row: We swept over a learning rate range of
$\\{0.001,0.005,0.01,0.05,0.1,0.25,0.5\\}$ with a constant batch size of 256.
Middle row: We swept over a batch size range of $\\{8,16,32,,64,128\\}$ with a
constant learning rate of 0.005. Bottom row: We swept over a L2 regularization
rate range of $\\{0,0.0001,0.0005\\}$ Each sweep took roughly 10h of training
on a single Google Cloud TPU V3 accessed via a Google colab.

###  A.5 Additional Experiments

####  A.5.1 Cifar-10 on ResNet-18

In Figure 4, we trained a ResNet-18 neural network on CIFAR-10. The
experimental setup is the same as described in Appendix A.4.1. The results
were all averaged over 5 random seeds for the neural network parameter default
initialization. The plots have no smoothing applied to the learning curves.

Top row: We swept over a learning rate range of
$\\{0.005,0.01,0.025,0.05,0.1,0.2\\}$ with a constant batch size of 256.
Middle row: We swept over a batch size range of $\\{8,16,32,64,128,256\\}$
with a constant learning rate of 0.2. Bottom row: We swept over a L2
regularization rate range of $\\{0.0,0.0001,0.00025,0.0005,0.00075\\}$ with
learning rate 0.02 and batch size 512.

Figure 4: Controlling the neural collapse via the model geometric complexity
for ResNet-18 trained on CIFAR-10. Lower $\operatorname{GC}$ produces lower
geometric collapse and more neural collapse (i.e., lower $\operatorname{NC}$)
for Top row: increased learning rates, Middle row: decreased batch sizes, and
Bottom row: increased L2 regularization.

####  A.5.2 MNIST on VGG-11

In Figure 5, we trained a VGG-11 neural network [53] on MNIST. The
experimental setup is the same as described in Appendix A.4.1. The results
were all averaged over 5 random seeds for the neural network parameter default
initialization. The plots have no smoothing applied to the learning curves.

Top row: We swept over a learning rate range of
$\\{0.0025,0.005,0.01,0.025,0.05,0.1\\}$ with a constant batch size of 512.
Middle row: We swept over a batch size range of
$\\{8,16,32,64,128,256,512,1024\\}$ with a constant learning rate of 0.005.
Bottom row: We swept over a L2 regularization rate range of
$\\{0,0.00001,0.0001,0.00025,0.0005\\}$ with learning rate 0.02 and batch size
512.

Figure 5: Controlling the neural collapse via the model geometric complexity
for VGG-11 trained on MNIST. Lower $\operatorname{GC}$ produces lower
geometric collapse and more neural collapse (i.e., lower $\operatorname{NC}$)
for Top row: increased learning rates, Middle row: decreased batch sizes, and
Bottom row: increased L2 regularization.

####  A.5.3 Fashion-MNIST on VGG-11

In Figure 6, we trained a VGG-11 neural network [53] on Fashion-MNIST. The
experimental setup is the same as described in Appendix A.4.1. The results
were all averaged over 5 random seeds for the neural network parameter default
initialization. The plots have no smoothing applied to the learning curves.

Top row: We swept over a learning rate range of
$\\{0.005,0.01,0.025,0.05,0.1,0.2\\}$ with a constant batch size of 256.
Bottom row: We swept over a L2 regularization rate range of
$\\{0,0.00001,0.0001,0.00025,0.0005\\}$ with learning rate 0.02 and batch size
512.

Figure 6: Controlling the neural collapse via the model geometric complexity
for VGG-11 trained on Fashion-MNIST. Lower $\operatorname{GC}$ produces lower
geometric collapse and more neural collapse (i.e., lower $\operatorname{NC}$)
for Top row: increased learning rates, Bottom row: increased L2
regularization.

####  A.5.4 Cifar-100 on ResNet50

In Figure 7, we trained a ResNet-50 neural network on CIFAR-100. The
experimental setup is the same as described in Appendix A.4.1. The results
were all averaged over 5 random seeds for the neural network parameter default
initialization. The plots have no smoothing applied to the learning curves.

Top row: We swept over a learning rate range of
$\\{0.005,0.01,0.05,0.1,0.2\\}$ with a constant batch size of 256. Bottom row:
We swept over a L2 regularization rate range of
$\\{0.0,0.0001,0.00025,0.0005,0.00075\\}$ with learning rate 0.02 and batch
size 256.

Figure 7: Controlling the neural collapse via the model geometric complexity
for ResNet-50 trained on CIFAR-100. Lower $\operatorname{GC}$ produces lower
geometric collapse and more neural collapse (i.e., lower $\operatorname{NC}$)
for Top row: increased learning rates, Bottom row: increased L2
regularization.

####  A.5.5 Lower Pre-trained GC Leads to Improved Fine-tuning: mini-ImageNet
dataset on VGG architecture

In Figure 8, we trained a VGG16 neural network (as described in [54])
implemented in Flax on Mini-ImageNet (as described in [57]). We used 10
classes with 600 examples per class for training, with a 10/90 train/test
split. To evaluate the downstream performance we used 100 downstream tasks
consisting of 5 labels for few shot learning randomly chosen from a pool of 20
separate from the training labels. The experimental setup was the same as
described in section A.4.2. Namely, we reported 1) the geometric complexity of
the embedding layer measured on the source dataset (Source GC), 2) the neural
collapse as measured by (3) measured and averaged on the 100 random samples of
5-label samples from the target dataset, 3) the test accuracy obtained by
solving the normal equation directly for a ridge regression with only 5
examples per class obtained from the target set. These three quantities were
measured every 1000 steps of a very long pre-training of 100,000 steps. The
optimizer was plain SGD (note we obtained similar results with momentum)
without any regularization nor schedule to avoid masking effects. We used
random crop and random flip for data augmentation. The results were all
averaged over 5 random seeds for the neural network parameter default
initialization. We swept over a learning rate range of
$\\{0.001,0.005,0.01\\}$ with a constant batch size of 256. Each sweep took
roughly 10h of training on a single Google Cloud TPU V3 accessed via a Google
colab.

Figure 8: Controlling target NC through source GC on mini-imagenet with
VGG-16: Increased learning rates produces lower $\operatorname{GC}$ and more
neural collapse (i.e. lower $\operatorname{NC}$) resulting in higher 5-shot
transfer accuracy.