# On Clique Roots of Flat Graphs

Hossein Teimoori Faal
Department of Mathematics and Computer Science,
Allameh Tabataba’i University, Tehran, Iran

(17 December 2021)

###### Abstract

A complete subgraph of a given graph is called a clique. A clique Polynomial
of a graph is a generating function of the number of cliques in $G$. A real
root of the clique polynomial of a graph $G$ is called a _clique root_ of $G$.
Hajiabolhassan and Mehrabadi showed that the clique polynomial of any simple
graph has a clique root in $[-1,0)$. As a generalization of their result, the
author of this paper showed that the class of $K_{4}$-free connected chordal
graphs has also only clique roots.
A given graph $G$ is called flat if each edge of $G$ belongs to at most two
triangles of $G$. In answering the author’s open question about the class of
_non-chordal_ graphs with the same property of having only c;ique roots, we
extend the aforementioned result to the class of $K_{4}$-free flat graphs. In
particular, we prove that the class of $K_{4}$-free flat graphs without
isolated edges has $r=-1$ as one of it’s clique roots. We finally present some
interesting open questions and conjectures regarding clique roots of graphs.

##  1 Introduction

Throughout this paper, we will assume that our graphs are finite, simple and
undirected. For the definitions which are not mentioned here, readers can
refer to the book [1].
For a given graph $G$, a complete subgraph of $G$ with $i$-vertices ($i\geq
1$) is called an $i$-clique of $G$. The number of $i$-cliques of $G$ will be
denoted by $c_{i}(G)$. By convention, we define $c_{0}(G)=1$ for any graph
$G$. The largest clique of $G$ is called the _clique number_ of $G$ and is
denoted by $\omega(G)$. The Clique polynomial of a graph $G$ denoted by
$C(G,x)$ is defined as follows.

$C(G,x)=\sum_{k=0}^{\omega(G)}c_{k}(G)x^{k},$ (1)

where by convention $c_{0}(G)$ is taken to be $1$ for any simple graph $G$. A
real root of $C(G,x)$ is called a _clique root_ of $G$.
Hajiabolhassan and late Mehrabadi [2] showed that the clique polynomial of any
graph $G$ has a clique root in $[-1,0)$. Moreover, they showed that clique
polynomial of triangle-free graphs has only clique roots. In the same line of
research the author this paper [3] showed that clique polynomial of
$K_{4}$-free connected chordal graphs has also only real roots.
An interesting class of graphs which are not (necessarily) chordal is the
class of flat graphs which are those graphs in which each edge belongs to at
most two triangles of them. In this paper, in answering to a question in the
author’s previous paper [3] about the existence of a nonchordal graph which
has the clique polynomial with only real roots, we prove that the class of
flat graphs has only clique roots. In particular, we show that the class of
flat graphs without isolated edges has always a clique root $r=-1$.

##  2 Basic Definitions and Notations

In this section, we quickly review clique polynomials and their algebraic
properties. We also introduce a new class of graphs which are triangulated but
they are not (necessarily) chordal.
For a graph $G=(V,E)$, by $G-v$ we mean a graph obtained from $G$ by deleting
the vertex $v\in V$. We also denote by $G[N_{G}(v)]$ the graph _induced_ by
the (open) neighborhood of the vertex $v$.

###  2.1 Clique Polynomials

######  Definition 2.1 (see [2]).

For a given graph $G$, the clique polynomial of $G$ denoted by $C(G,x)$ is
defined by

$C(G,x)=1+\sum_{i=1}^{\omega(G)}c_{i}(G)x^{i}.$ (2)

######  Example 2.1.

The clique polynomial of the graph $G_{1}$ depicted in Figure1, is the
following

$C(G_{1},x)=1+4x+4x^{2}+x^{3}.$

$4$$3$$1$$2$ Figure 1: The clique polynomial of $G_{1}$ is
$C(G_{1},x)=1+4x+4x^{2}+x^{3}$.

The following vertex-recurrence relation [2] can be obtained for a clique
polynomial of a graph, using a simple counting argument.

$C(G,x)=C(G-v,x)+xC(G[N(v)],x)\hskip 14.22636pt(v\in V(G)).$ (3)

We also recall the following interesting _vertex-deck_ identity for the
cliques of $G$.

######  Lemma 2.1.

Let $G=(V,E)$ be a graph on $n$ vertices. Then for any $i\geq 0$, we have

$(n-i)c_{i}(G)=\sum_{v\in V}c_{i}(G-v)$ (4)

###### Proof.

The proof is the straight-forward double-counting argument and left to the
interested readers as a simple exercise. ∎

Using the vertex-recurrence relation (3) and the vertex-deck identity (4), one
can easily prove the following combinatorial interpretation of the first
derivative of the clique polynomial [4]. For the sake of completeness, we also
include the proof.

######  Lemma 2.2.

For a graph $G=(V,E)$, we have

$\frac{d}{dx}C(G,x)=\sum_{v\in V}C(G[N(v)],x).$ (5)

###### Proof.

By multiplying both sides of (4) by $x^{i}$ and summing over all $i$, we get

$\sum_{i\geq 0}(n-i)c_{i}(G)x^{i}=\sum_{i\geq 0}\sum_{v\in V}c_{i}(G-v)x^{i},$

or equivalently (by changing summation order), we have

$n\sum_{i\geq 0}c_{i}(G)x^{i}-x\sum_{i\geq 1}ic_{i}(G)x^{i-1}=\sum_{v\in
V}\left(\sum_{i\geq 0}c_{i}(G-v)x^{i}\right).$ (6)

On the other hand, by the definition of the first derivative of any polynomial
we obtain

$\frac{d}{dx}C(G,x)=\sum_{i\geq 1}ic_{i}(G)x^{i-1}.$ (7)

Hence, considering the relation (7), we can rewrite the equality (6) as follow

$n.C(G,x)-x\frac{d}{dx}C(G,x)=\sum_{v\in V}C(G-v,x).$

or equivalently,

$\frac{d}{dx}C(G,x)=\sum_{v\in V}\frac{C(G,x)-C(G-v,x)}{x}.$ (8)

Finally based on the relations (8) and vertex-recurrence formula 3, we get the
desired result.

∎

###### Proof.

By multiplying both sides of (4) by $x^{i}$ and summing over all $i$, we get

$\sum_{i\geq 0}(n-i)c_{i}(G)x^{i}=\sum_{i\geq 0}\sum_{v\in V}c_{i}(G-v)x^{i},$

or equivalently (by changing summation order), we have

$n\sum_{i\geq 0}c_{i}(G)x^{i}-x\sum_{i\geq 1}ic_{i}(G)x^{i-1}=\sum_{v\in
V}\left(\sum_{i\geq 0}c_{i}(G-v)x^{i}\right).$ (9)

On the other hand, by the definition of the first derivative of any polynomial
we obtain

$\frac{d}{dx}C(G,x)=\sum_{i\geq 1}ic_{i}(G)x^{i-1}.$ (10)

Hence, considering the relation (10), we can rewrite the equality (9) as
follow

$n.C(G,x)-x\frac{d}{dx}C(G,x)=\sum_{v\in V}C(G-v,x).$

or equivalently,

$\frac{d}{dx}C(G,x)=\sum_{v\in V}\frac{C(G,x)-C(G-v,x)}{x}.$ (11)

Finally based on the relations (11) and vertex-recurrence formula 3, we get
the desired result.

∎

###  2.2 Flat Graphs

Next, we introduce a new class of graphs which can be of special interest for
our purposes. They are triangulated graphs but they are not (necessarily)
chordal. We call them flat graphs [5].

######  Definition 2.2.

A graph $G$ is called _flat_ if every edge of $G$ belongs to at most two
triangles of $G$.

######  Example 2.2.

The wheel $W_{5}$ is a simple example of a flat graph which is clearly not a
chordal graph.

For the sequel purposes, we need the following generalization of the concept
of a _degree_ of a vertex.

######  Definition 2.3 (see [6]).

Let $G=(V,E)$ be a graph and $e=\\{u,v\\}\in E$ be an edge of $G$. The value
of $e$ in $G$ denoted by $val_{G}(e)$ is defined as follows:

$val_{G}(e)=N_{G}(u)\cap N_{G}(v).$

An edge $e\in E$ with $val_{G}(e)=0$ is called an _isolated_ edge of $G$.

##  3 Main Results

In this section, we are going to prove the similar results of the author’s
paper [3]. The basic idea is that if the derivative of a cubic polynomial
$p(x)$ has (at least) one real root, then it has only real roots. Before going
through the proof, we recall the well-known Euler identity for _planar_
graphs.

######  Lemma 3.1.

For a planar graph $G$ with $n$ vertices, $m$ edges and $t$ faces, we have

$n-m+f=2.$ (12)

In particular, for triangulated graph $f=t+1$, where $t$ denotes the number of
triangles.

We first state a simpler version of our main result. From now on, we will call
a real root of $C(G,x)$ a _clique root_ of $G$.

######  Theorem 3.1.

The class of $K_{4}$-free flat graphs without isolated edges has only clique
roots. In particular, this class has always the clique root $r=-1$.

###### Proof.

Since $G$ is $K_{4}$-free connected graph, we have

$C(G,x)=1+nx+mx^{2}+tx^{3}.$

Now, we note that by the definition of flat graphs, the graph $G$ can not
contain any copy of $K_{3,3}$ as a subgraph (or as a minor in general). Hence
by the Kuratowski’s well-known theorem, $G$ is a planar graph. Therefore,
based on the Euler identity3.1, we conclude that $C(G,-1)=0$. Hence, to
complete the proof we only need to show that the polynomial
$\frac{d}{dx}C(G,x)$ has (at least) one real root. Next, we note that for a
flat graph without isolated edges, the neighborhood of any vertex $v\in V(G)$
is a path or a cycle. Let $\\{v_{i}\\}^{k_{1}}_{i=1}$ be the set of vertices
such that each $G[N_{G}(v_{i})]$ is a path of length $l_{i}$ and
$\\{u_{j}\\}^{k_{2}}_{j=1}$ is the set of vertices such that each
$G[N_{G}(u_{j})]$ is a cycle of length $c_{j}$.
Now, considering Proposition 2.2, we get

$\displaystyle\frac{d}{dx}C(G,x)$ $\displaystyle=$ $\displaystyle\sum_{v\in
V}C(G[N_{G}(v)],x),$ $\displaystyle=$
$\displaystyle\sum^{k_{1}}_{i=1}C(G[N_{G}(v_{i})],x)+\sum^{k_{2}}_{j=1}C(G[N_{G}(v_{j}),x]),$
$\displaystyle=$
$\displaystyle\sum^{k_{1}}_{i=1}(1+x)(1+l_{i}x)+\sum^{k_{2}}_{j=1}(1+c_{j}x+c_{j}x^{2}).$

Considering the fact that each $l_{i}\geq 2$ and $c_{j}\geq 4$, we immediately
conclude that

$\frac{d}{dx}C(G,-\frac{1}{2})\leq 0.$ (13)

Finally, the result is immediate considering the fact that $C(G,0)>0$, the
equality 13 and the _intermediate value theorem_ from calculus. ∎

Next, based on the vertex-recurrence formula 3 and some slight modifications
of Theorem 3.1, we can get the following interesting result.

######  Theorem 3.2.

The class of $K_{4}$-free flat graphs in which no isolated edge lies in any
induced cycle of $G$ has only clique roots. In particular, this class has
always a clique root $r=-1$.

Now, we are in position to state and prove the main result of this paper which
is stronger than the previous results.

######  Theorem 3.3.

Let $G$ be a $K_{4}$-free flat graph. Then, the graph $G$ has only clique
roots.

###### Proof.

We partition the vertex set $V(G)$ into three (disjoint) sets $V_{1}$, $V_{2}$
and $V_{3}$. $V_{1}$ is the set vertices $v$ such that $G[N_{G}(v)]$ is a
path. $V_{2}$ is the set of those vertices $u$ such that $G[N_{G}(u)]$ is a
cycle. $V_{3}$ is the set of vertices $w$ such that $G[N_{G}(w)]$ is a
collection of paths and isolated vertices (paths of length zero).
In this case, we do not have necessarily a clique root $r=-1$ but we still
have a real root because $C(G,x)$ is a cubic polynomial. Thus, again we need
to prove that $\frac{d}{dx}C(G,x)$ has (at least) one real root.

$\displaystyle\frac{d}{dx}C(G,x)$ $\displaystyle=$ $\displaystyle\sum_{v\in
V}C(G[N_{G}(v)],x),$ $\displaystyle=$ $\displaystyle\sum_{v\in
V_{1}}C(G[N_{G}(v)],x)+\sum_{u\in V_{2}}C(G[N_{G}(u)],x)+\sum_{w\in
V_{3}}C(G[N_{G}(w)],x),$ $\displaystyle=$ $\displaystyle\sum_{v\in
V_{1}}(1+x)(1+l_{v}x)+\sum_{u\in V_{2}}(1+c_{u}x+c_{u}x^{2})$ $\displaystyle+$
$\displaystyle\sum_{w\in
V_{3}}\Big{(}(1+x)(1+l_{w}x)+(1+k_{w}x)-k_{w}\Big{)},$

where the last line is based on the clique polynomial formula for the disjoint
union of two graphs. ∎

Once again, considering the fact that each $l_{i}\geq 2$, $c_{j}\geq 4$ and
$k_{w}\geq 1$ it is not hard to prove that

$\frac{d}{dx}C(G,-\frac{1}{2})\leq 0,$

which again by the intermediate value theorem implies that
$\frac{d}{dx}C(G,x)$ has at least one real root. This finishes the proof.

##  4 Open Questions and Conjectures

In this final section, we propose several interesting open questions and
conjectures regarding the clique roots of graphs.
In the paper [3], the author gives a proof of the real rootedness of the
clique polynomial of $K_{4}$-free chordal graphs mainly based on the idea of
the _breadth first search_ algorithm. Therefore, the following question is of
our special interest.

######  Question 4.1.

Can we give an algorithmic proof of Theorem 3.3?

It is also interesting to know that whether Theorem 3.3 is true for the class
of $K_{5}$-free graphs or not. Here is a counter-example. Let $G$ be a the
graph $K^{+}_{4}$ which is a complete graph with an isolated edge attached to
one of its vertices. Hence, we have

$C(K^{+}_{4},x)=(1+x)^{4}+(1+x)^{2}-1.$

Therefore, $\frac{d}{dx}C(K^{+}_{4},x)=4(1+x)^{3}+2(1+x)$ and this implies
that $\frac{d^{2}}{dx^{2}}C(K^{+}_{4},x)=12(1+x)^{2}+2>0.$ Thus, we
immediately conclude that not all roots of $C(K^{+}_{4},x)$ are real. The
following open question naturally arises.

######  Question 4.2.

Which subclasses of flat graphs have only clique roots?

We first note that a generalization of the edge value of an edge $e\in E(G)$
of a graph $G$, which we call it _triangle value_ [6] can be defined as
follows.

######  Definition 4.1.

For a triangle $\delta=\\{u,v,w\\}\subset V(G)$, the value of $\delta$ in $G$
denoted by $val_{G}(\delta)$ is defined by

$val_{G}(\delta)=N_{G}(u)\cap N_{G}(v)\cap N_{G}(w).$

Moreover, a triangle is called _isolated_ if its triangle value is zero.

In analogy with Theorem 3.1, we made the following conjecture.

######  Conjecture 1.

The class of $2$-connected $K_{5}$-free flat graphs without isolated triangles
has only clique roots. In particular, this class has always the clique root
$r=-1$ of multiplicity $2$.

A $i$-clique is called _odd_ clique if $i$ is odd and otherwise is called an
_even_ clique. We can define an interesting parameter in graph theory as the
absolute value of the difference between the total number of odd and even
cliques. More precisely, we define

$D_{clique}(G)=\Big{|}\sum_{i:~{}i~{}is~{}odd}c_{i}(G)-\sum_{i:~{}i~{}is~{}even}c_{i}(G)\Big{|}.$
(14)

As we can see from Theorem 3.1 and Theorem 3.2, for the subclasses of flat
graphs $G$ mentioned in those theorems, we conclude that $D_{clique}(G)=1$ .
Thus, the following question naturally arises.

######  Question 4.3.

For which classes of graphs there exits a positive integer $k$ such that
$D_{clique}(G)\leq k$ ?

We note that when $k=1$, we call the graph $G$ a _balanced_ graph. We believe
that the following conjecture is true.

######  Conjecture 2.

Any flat graph is a balanced graph.

## References

  * [1] J.A. Bondy and U.S.R. Murty, Graph Theory with Applications, North-Holland, NewYork-Amsterdam-Oxford (1982).
  * [2] H. Hajiabolhassan and M. L. Mehrabadi, On clique polynomials, Australasian Journal of Combinatorics, 18 (1998), 313–316.
  * [3] H. Teimoori, Clique roots of $\mathit{K}_{4}$-free chordal graphs, Electronic Journal of Graph Theory and Applications 7 (2018), 87–94.
  * [4] X. Li and I. Gutman, A unified approach to the first derivatives of graph polynomials, Discrete Applied Mathematics, 58 (1995), 293–297.
  * [5] P. Haxell, A. Kostochka and S. Thomassé, Packing and Covering Triangles in $K_{4}$-free Planar Graphs , Graphs and Combinatorics, 28 (2012), 653–662.
  * [6] H. Teimoori, Clique values and Turan-type inequalities, In preparation.


# Gradient descent induces alignment between weights and the empirical NTK for
deep non-linear networks

Daniel Beaglehole  Ioannis Mitliagkas  Atish Agarwala

###### Abstract

Understanding the mechanisms through which neural networks extract statistics
from input-label pairs is one of the most important unsolved problems in
supervised learning. Prior works have identified that the gram matrices of the
weights in trained neural networks of general architectures are proportional
to the average gradient outer product of the model, in a statement known as
the Neural Feature Ansatz (NFA). However, the reason these quantities become
correlated during training is poorly understood. In this work, we explain the
emergence of this correlation. We identify that the NFA is equivalent to
alignment between the left singular structure of the weight matrices and a
significant component of the empirical neural tangent kernels associated with
those weights. We establish that the NFA introduced in prior works is driven
by a centered NFA that isolates this alignment. We show that the speed of NFA
development can be predicted analytically at early training times in terms of
simple statistics of the inputs and labels. Finally, we introduce a simple
intervention to increase NFA correlation at any given layer, which
dramatically improves the quality of features learned.

##  1 Introduction

Neural networks have emerged as the state-of-the-art machine learning methods
for seemingly complex tasks, such as language generation (Brown et al., 2020),
image classification (Krizhevsky et al., 2012), and visual rendering
(Mildenhall et al., 2021). The precise reasons why neural networks generalize
well have been the subject of intensive exploration, beginning with the
observation that standard generalization bounds from statistical learning
theory fall short of explaining their performance (Zhang et al., 2021).

A promising line of work emerged in the form of the neural tangent kernel
(NTK), connecting neural networks to kernels in the wide limit (Jacot et al.,
2018; Chizat et al., 2019). However, subsequent research showed that the
success of neural networks relies critically on aspects of learning which are
absent in kernel approximations (Ghorbani et al., 2019; Allen-Zhu and Li,
2019; Yehudai and Shamir, 2019; Li et al., 2020; Refinetti et al., 2021).
Other work showed that low width suffices for gradient descent to achieve
arbitrarily small test error (Ji and Zhu, 2020), further refuting the idea
that extremely wide networks are necessary.

Subsequently, the success of neural networks has been largely attributed to
_feature learning_ \- the ability of neural networks to learn statistics,
measurements, and representations of data which are useful for downstream
tasks. However, the specific mechanism through which features are learned is
an important unsolved problem in deep learning theory. A number of works have
studied the abilities of neural networks to learn features in structured
settings (Abbe et al., 2022; Ba et al., 2022; Nichani et al., 2023; Barak et
al., 2022; Damian et al., 2022; Moniri et al., 2023; Parkinson et al., 2023).
Some of that work proves strict separation in terms of sample complexity
between neural networks trained with stochastic gradient descent and kernels
(Mousavi-Hosseini et al., 2022).

The work above studies simple structure, such as learning from low-rank data
or functions that are hierarchical compositions of simple elements. Recent
work makes a big step towards generalizing these assumptions by proposing the
_neural feature ansatz_ (NFA) (Radhakrishnan et al., 2022; Beaglehole et al.,
2023), a general structure that emerges in the weights of trained neural
networks. The NFA states that the gram matrix of the weights at a given layer
(known as the neural feature matrix (NFM)) is aligned with the average
gradient outer product (AGOP) of the network with respect to the input to that
layer. In particular, the NFM and AGOP are highly correlated in all layers of
trained neural networks of general architectures, including practical models
such as AlexNet (Krizhevsky et al., 2012) and VGG (Simonyan and Zisserman,
2014).

A major missing element of this theory is the reason the AGOP and NFM become
correlated during training. In this paper, we explain the emergence of this
correlation. We establish that the NFA is equivalent to alignment between the
left singular structure of the weight matrices and the pre-activation neural
tangent kernel (PTK), which is closely related to empirical neural tangent
kernel (ENTK) (Section 2). By introducing a modified NFA, we demonstrate that
the original NFA is driven by alignment of the left singular vectors of the
weight matrices to with the PTK (Section 2). Further, we show that the
alignment speed can be understood analytically in terms of the statistics of
the data and labels at early times, and can be manipulated theoretically
(Section 3).

Finally, we force the strength of the NFA alignment to improve using two
interventions: small initialization and Speed Limited Optimization, a
layerwise gradient normalization scheme (Section 4). This suggests a path
towards rational design of architectures and training procedures that maximize
this notion of feature learning.

##  2 Alignment between the weight matrices and the empirical NTK

###  2.1 Preliminaries

We consider fully-connected neural networks with a single output of depth
$L\geq 1$, where $L$ is the number of hidden layers, written
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}$. We write the input to layer
$\ell\in\\{0,\ldots,L-1\\}$ as $x_{\ell}$, where $x_{0}\equiv x$ is the
original datapoint, and the pre-activation as $h_{\ell}(x)$. Then,

$\displaystyle h_{\ell}(x)=W^{(\ell)}x_{\ell},\qquad
x_{\ell+1}=\phi(h_{\ell}(x))~{},$

where $\phi$ is an element-wise nonlinearity,
$W^{(\ell)}\in\mathbb{R}^{k_{\ell+1}\times k_{\ell}}$ is a weight matrix, and
$k_{\ell}$ is the hidden dimension at layer $\ell$. We restrict $k_{L+1}$ to
be the number of output logits, and set $k_{0}=d$, where $d$ is the input
dimension of the data. Note that $f(x)=h_{L+1}(x)$. We train $f$ by gradient
descent on a loss function $\mathcal{L}(\theta,X)$, where $X$ is an input
dataset, and $\theta$ is the collection of weights.

We consider a supervised learning setup where we are provided $n$ input-label
pairs
$(x^{(1)},y^{(1)}),\ldots,(x^{(n)},y^{(n)})\in\mathbb{R}^{d}\times\mathbb{R}$.
We denote the inputs $X\in\mathbb{R}^{n\times d}$ and the labels
$y\in\mathbb{R}^{n\times 1}$. We train a fully-connected neural network to
learn the mapping from inputs to labels by minimizing a standard loss
function, such as mean-squared-error or cross-entropy, on the dataset.

Figure 1: Various feature learning measures for target function
$y(x)=\sum_{k=1}^{r}x_{k\text{ mod }r}\cdot x_{(k+1)\text{ mod }r}$ with $r=5$
and inputs drawn from standard normal. The EGOP
$\mathbb{E}_{x\sim\mu}\left[\frac{\partial y}{\partial x}\frac{\partial
y}{\partial x}^{\top}\right]$ (first plot) captures the low-rank structure of
the task. The NFM $\left(W^{\top}W\right)$ (second plot) and AGOP
$\left(W^{\top}KW\right)$ (third plot) of a fully-connected network are
similar to each other and the EGOP. Replacing $K$ with a symmetric matrix $Q$
with the same spectrum but independent eigenvectors obscures the low rank
structure (fourth plot), and reduces the correlation from
$\rho\left(F,\bar{G}\right)=0.93$ to $\rho\left(F,W^{\top}QW\right)=0.53$.

One can define two objects associated with neural networks that capture
learned structure. For a given layer $\ell$, the neural feature matrix (NFM)
$F_{\ell}$ is the gram matrix of the columns of the weight matrix
$W^{(\ell)}$, i.e. $F_{\ell}\equiv(W^{(\ell)})^{\top}W^{(\ell)}$. The second
fundamental object we consider is the average gradient outer product (AGOP)
$\bar{G}_{\ell}$, defined as
$\bar{G}_{\ell}\equiv\frac{1}{n}\sum_{\alpha=1}^{n}\frac{\partial
f(x^{(\alpha)}_{\ell})}{\partial x_{\ell}}\frac{\partial
f(x^{(\alpha)}_{\ell})}{\partial x_{\ell}}^{\top}$. To understand the
structure of these objects, consider the following chain-monomial low-rank
task:

$\displaystyle y(x)=\sum_{k=1}^{r}x_{k\text{ mod }r}\cdot x_{(k+1)\text{ mod
}r}~{},$ (1)

where the data inputs are sampled from an isotropic Gaussian distribution
$\mu=\mathcal{N}(0,I)$. In this case, both objects capture the coordinates on
which the target function depends (Figure 1). Prior work has shown that in
trained neural networks, these objects will be proportional to each other to
varying degree of approximation. This notion is formalized in the Neural
Feature Ansatz (NFA):

######  Ansatz 1 (Neural Feature Ansatz (Radhakrishnan et al., 2022)).

The Neural Feature Ansatz states that, for all layers $\ell\in[L]$ of a fully-
connected neural network with $L$ hidden layers trained on input data
$x^{(1)},\ldots,x^{(n)}$,

$\displaystyle(W^{(\ell)})^{\top}W^{(\ell)}\propto\frac{1}{n}\sum_{\alpha=1}^{n}\frac{\partial
f(x^{(\alpha)}_{\ell})}{\partial x_{\ell}}\frac{\partial
f(x^{(\alpha)}_{\ell})}{\partial x_{\ell}}^{\top}~{},$

The left and right hand sides of this equation are the NFM and AGOP with
respect to the input of layer $\ell$. Here, $\frac{\partial f(x)}{\partial
x_{\ell}}\in\mathbb{R}^{k_{\ell}\times 1}$ denotes the gradient of the
function $f$ with respect to the intermediate representation $x_{\ell}$. For
simplicity, we may concatenate these gradients into a single matrix
$\frac{\partial f(X)}{\partial x_{\ell}}\in\mathbb{R}^{n\times k_{\ell}}$.
Note we consider scalar outputs in this work, though the NFA relation is
identical when there are $c\geq 1$ outputs, where in this case $\frac{\partial
f(x)}{\partial x_{\ell}}\in\mathbb{R}^{k_{\ell}\times c}$ is the input-output
Jacobian of the model $f$.

The NFA is equivalent to the statement that, for every layer $\ell$, the NFM
$F_{\ell}$ and AGOP $\bar{G}_{\ell}$, will have cosine similarity, or
correlation,

$\displaystyle\rho\left(F_{\ell},\bar{G}_{\ell}\right)=1~{},$

where the correlation $\rho$, with range $[-1,1]$, is defined as

$\displaystyle\rho\left(A,B\right)=\text{tr}\left(A^{\top}B\right)\cdot\text{tr}\left(A^{\top}A\right)^{-1/2}\cdot\text{tr}\left(B^{\top}B\right)^{-1/2}~{}.$

for any two matrices $A,B\in\mathbb{R}^{d_{1}\times d_{2}}$, for any $d_{1}$
and $d_{2}$.

We note that while the NFA states exact proportionality, in practice, the NFM
and AGOP are highly correlated with correlation less than $1$ (see e.g. Figure
2), where the final correlation depends on many aspects of training and
architecture choice. For example, the final value of the UC-NFA can be
sensitive to the magnitude of the initial weights (e.g. first column of Figure
2). We parameterize this magnitude by the initialization scale, $s_{\ell}$,
for layer $\ell$, where the initial weights are sampled i.i.d. as
$W^{(\ell)}_{ij}\sim\mathcal{N}\left(0,\frac{s_{\ell}}{k_{\ell}}\right)$.

The relation between the NFM and the AGOP is significant, in part, because for
a neural network that has learned enough information about the target
function, the AGOP of this model with respect to the first-layer inputs will
approximate the _expected gradient outer product_ (EGOP) of the target
function (Yuan et al., 2023). In particular, as we will see later with the
example of a low rank polynomial, the EGOP of the target function contains
task-specific structure that is completely independent of the model used to
estimate it. Where the labels are generated from a particular target function
$y(x):\mathbb{R}^{d}\rightarrow\mathbb{R}$ on data sampled from a distribution
$\mu$, the EGOP is equal to,

$\displaystyle\operatorname{EGOP}(y,\mu)=\mathbb{E}_{x\sim\mu}\left[\frac{\partial
y}{\partial x}\frac{\partial y}{\partial x}^{\top}\right]~{}.$

If the NFA holds, the correlation of the EGOP and the AGOP at the end of
training also implies high correlation between the NFM of the first layer and
the EGOP, so that the NFM has encoded this task-specific structure.

To demonstrate the significance of the approximate proportionality between the
NFM and the AGOP in successfully trained networks, we return to the chain-
monomial low rank task. In this case, $\operatorname{EGOP}(y,\mu)$ will be
rank $r$, where $r$ is much less than the ambient dimension (Figure 1).
Similarly, the AGOP of the trained model resembles the EGOP. In this case the
NFA implies that the NFM also resembles the EGOP. Therefore, the neural
network has learned the model-independent and task-specific structure of the
chain-monomial task in the right singular values and vectors of the first
layer weight matrix, as these are determined by the NFM. As a demonstration of
this fact, the NFM of a trained neural network can be recovered up to high
correlation from the AGOP of a fixed kernel method on real datasets
(Radhakrishnan et al., 2022)

This insight has inspired iterative kernel methods which can match the
performance of fully-connected networks (Radhakrishnan et al., 2022, 2024) and
can close the gap with convolutional networks (Beaglehole et al., 2023).
Additional prior works demonstrate the benefit of including the AGOP features
to improve feature-less predictors (Hristache et al., 2001; Trivedi et al.,
2014; Kpotufe et al., 2016). Additionally, because the NFM is correlated with
the AGOP, the AGOP can be used to recover the features from feature-less
methods, such as kernel machines.

###  2.2 Alignment decomposition

In order to understand the NFA, it is useful to decompose the AGOP. Doing so
will allow us to show that the NFA can be interpreted as an alignment between
weight matrices and the pre-activation tangent kernel (PTK), which is closely
related to the neural tangent kernel (NTK).

For any layer $\ell$, we can re-write the AGOP as,

$\displaystyle\bar{G}_{\ell}=(W^{(\ell)})^{\top}K^{(\ell)}W^{(\ell)},~{}K^{(\ell)}\equiv\frac{\partial
f(X)}{\partial h_{\ell}}^{\top}\frac{\partial f(X)}{\partial h_{\ell}}$

This gives us the following proposition:

######  Proposition 2 (Alignment decomposition of NFA).

$\displaystyle\rho\left(F_{\ell},\bar{G}_{\ell}\right)=\rho\left((W^{(\ell)})^{\top}W^{(\ell)},(W^{(\ell)})^{\top}K^{(\ell)}W^{(\ell)}\right)~{}.$

This alignment holds trivially and exactly if $K^{(\ell)}$ is the identity.
However, the correlation can be high in trained networks even with non-trivial
$K^{(\ell)}$. For example, in the chain monomial task (Figure 1), $K^{(0)}$ is
far from identity (standard deviation of its eigenvalues is $5.9$ times its
average eigenvalue), but the NFA correlation is $0.93$ at the end of training.
We also note that if $K^{(\ell)}$ is independent of $W^{(\ell)}$, the
alignment is lower than in trained networks; in the same example, replacing
$K^{(0)}$ with a matrix $Q$ with equal spectrum but random eigenvectors
greatly reduces the correlation to $0.53$ and qualitatively disrupts the
structure relative to the NFM (Figure 1, rightmost column). We show the same
result for the CelebA dataset (see Appendix K). Therefore, the NFA has to do
with the alignment of the left eigenvectors of $W^{(\ell)}$ with $K^{(\ell)}$
in addition to spectral considerations.

The $k_{\ell}\times k_{\ell}$ matrix $K^{(\ell)}$ is the feature covariance of
the pre-activation tangent kernel (PTK) $\mathcal{K}^{(\ell)}$ which we define
as:

$\mathcal{K}^{(\ell)}(x,z)\equiv\frac{\partial f(x)}{\partial
h_{\ell}}\cdot\frac{\partial f(z)}{\partial h_{\ell}}$ (2)

The PTK is related to the layerwise empirical NTK (ENTK) $\hat{\theta}_{\ell}$
by the following identity:

(a)

(b)

Figure 2: Uncentered and centered NFA correlations for a two hidden layer MLP
trained on data drawn from a standard Normal distribution on the chain
monomial task of rank $r=5$. Only, the initialization scale of the first layer
weights is varied, while $s_{\ell}=1$ for $\ell>0$. The top row shows the
values for layer 1, while the bottom row are the values for layer 2. Train
(test) losses are scaled by the maximum train (test) loss achieved so that
they are between $0$ and $1$. Here width is $256$, the input dimension is
$32$, and the dataset contains $256$ points.

######  Proposition 3 (Pre-activation to neural tangent identity).

Consider a depth-$L$ neural network $f(x)$ with inputs
$X_{\ell}\in\mathbb{R}^{n\times k_{\ell}}$ to weight matrices
$W^{(\ell)}\in\mathbb{R}^{k_{\ell}\times k_{\ell}}$ for $\ell\in[L]$. Consider
the empirical NTK where gradients are taken only with respect to $W^{(\ell)}$
evaluated between two points $x$ and $z$ denoted by
$\hat{\Theta}_{\ell}(x_{0},z_{0})=\langle\frac{\partial f(x_{0})}{\partial
W^{(\ell)}},\frac{\partial f(z_{0})}{\partial W^{(\ell)}}\rangle$. Then for
all $\ell$, $\hat{\Theta}_{\ell}$ satisfies,

$\displaystyle\hat{\Theta}_{\ell}(x_{0},z_{0})=\mathcal{K}^{(\ell)}(x_{0},z_{0})\cdot
x_{\ell}^{\top}z_{\ell}~{}.$

Therefore, the NFA is a statement about the alignment of the weight matrices
with the PTK feature covariance - which itself is a factor of the layerwise
NTK.

###  2.3 Centering the NFA isolates weight-PTK alignment

We showed that the NFA is equivalent to PTK-weight alignment (Proposition 2),
and that it emerges during training (Figure 2 and Appendix G). We now ask: is
the development of the NFA due to alignment of the weight matrices to the
current PTK, or the alignment of the PTK to the current weights?

In practice, both effects matter, but numerical evidence suggests that changes
in the PTK don’t drive the early dynamics of the NFA (Appendix B). Instead, we
focus on the alignment of the weights to the PTK at early times. We can
measure this alignment using an alternative, centered version of the NFA:

######  Ansatz 4 (Centered Neural Feature Ansatz).

Let $W_{0}^{(\ell)}$, $W^{(\ell)}$, and $K^{(\ell)}$ be the initial weight
matrix, trained weight matrix, and interior feature matrix, respectively, at
layer $\ell$, and let $\bar{W}^{(\ell)}\equiv W^{(\ell)}-W^{(\ell)}_{0}$. The
centered neural feature ansatz (C-NFA) postulates that the following
quantities will remain proportional to one another,

$\displaystyle(\bar{W}^{(\ell)})^{\top}\bar{W}^{(\ell)}\propto(\bar{W}^{(\ell)})^{\top}K^{(\ell)}\bar{W}^{(\ell)}~{}.$

We refer to the centered NFA as the C-NFA, and the original, uncentered NFA as
the UC-NFA to distinguish them.

Since $\bar{W}=0$ at initialization, the early dynamics of the C-NFA are
dominated by $W$ aligning with the initial $K$:

######  Proposition 5 (Centered NFA dynamics).

At initialization, $\bar{W}^{\top}\bar{W}=\bar{W}^{\top}K\bar{W}=0$, and,

$\begin{split}\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}\bar{W})=0,\quad&\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}\bar{W})=2\dot{W}^{\top}\dot{W},\\\
\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}K\bar{W})=0,\quad&\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}K\bar{W})=2\dot{W}^{\top}K\dot{W}~{}.\end{split}$

The first non-zero derivatives of the two quantities in the C-NFA give us the
change in the gram matrix of $W$, as well as the change in the AGOP for fixed
PTK $K$. This makes it a useful object to study weight-PTK alignment.

We note that the C-NFA correlation remains high throughout training across
initialization scale (Figure 2, second column, and Appendix G), but is
especially high and invariant to scale early on in training. This is in
contrast to the UC-NFA, whose final value can be extremely sensitive to the
initialization scale (e.g. second row, Figure 2).

In the remainder of this work, we first analyze the C-NFA in early training
and show that gradient based training drives the C-NFA, with a strength that
depends on model and datasets. (Section 3). We then discuss initializations
and optimizers which can increase the strength of the C-NFA and induce larger
correlation between $F$ and $\bar{G}$ (Section 4).

##  3 Theoretically predicting the centered NFA

###  3.1 Early time C-NFA dynamics

To understand the dynamics of the weights aligning to the PTK, we compute the
time derivatives of the C-NFA at initialization (Appendix A):

######  Proposition 6 (Centered NFA decomposition).

For a fully-connected neural network at initialization, whose weights are
trained with gradient descent on a loss function $\mathcal{L}$,

$\displaystyle\dot{W}^{\top}\dot{W}=X^{\top}\dot{\mathcal{L}}\mathcal{K}\dot{\mathcal{L}}X,\quad\dot{W}^{\top}K\dot{W}=X^{\top}\dot{\mathcal{L}}\mathcal{K}^{2}\dot{\mathcal{L}}X~{}$
(3)

where $\dot{\mathcal{L}}$ is the $n\times n$ diagonal matrix of logit
derivatives

$\dot{\mathcal{L}}\equiv{{\rm diag}\left(\frac{\partial\mathcal{L}}{\partial
f}\right)}$

We immediately see from these equations how gradient-based training can drive
the NFA. Even though the NFA doesn’t explicitly involve the loss, the gradient
descent dynamics of the NFM and AGOP depend on the labels in a similar way,
and differ by a factor of the kernel $\mathcal{K}$ \- similar to how the AGOP
and NFM differ by a factor of $K$.

When $\mathcal{K}$ is proportional to a projection matrix, i.e.
$\mathcal{K}\propto\mathcal{K}^{2}$, these derivatives will be perfectly
correlated. Even in the general case where $\mathcal{K}$ is dispersed with
moderate eigenvalue decay, and does not commute with
$\dot{\mathcal{L}}XX^{\top}\dot{\mathcal{L}}$, the two quantities have non-
negative covariance. In the next section we will construct a dataset which
interpolates between adversarial and aligned eigenstructure to demonstrate the
range of possible values the derivative correlations can take.

In order to understand the correlation more quantitatively, we will focus on
the case of mean-squared error (MSE) loss in the rest of this work. Here
$\dot{\mathcal{L}}$ corresponds to the diagonal matrix of the residuals
$y-f(x)$. If the outputs of the network is $0$ on the training data, then
$\dot{\mathcal{L}}=Y\equiv{\rm diag}\left(y\right)$, the labels themselves. In
that case, the correlation of the time derivatives at initialization is a
function of the input-label covariances as well as the PTK matrix:

$\begin{split}&\rho\left(\dot{W}^{\top}\dot{W},\dot{W}^{\top}K\dot{W}\right)=\,\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)\\\
&\cdot\text{tr}\left((X^{\top}Y\mathcal{K}YX)^{2}\right)^{-1/2}\cdot\text{tr}\left((X^{\top}Y\mathcal{K}^{2}YX)^{2}\right)^{-1/2}\end{split}$
(4)

In general we expect that these quantities can self-average under the
appropriate high-dimensional limits. It is helpful then to write down the
average (over initializations) of the covariance. Taking expectation of the
first factor,

$\begin{split}&\mathbb{E}\left[\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)\right]=\\\
&\text{tr}\left(X^{\top}Y\mathbb{E}\left[\mathcal{K}\right]YXX^{\top}Y\mathbb{E}\left[\mathcal{K}^{2}\right]YX\right)\\\
&+\text{tr}\left(\text{Cov}\left(X^{\top}Y\mathcal{K}YX,X^{\top}Y\mathcal{K}^{2}YX\right)\right)\end{split}$
(5)

A similar decomposition exist for the denominator terms.

The simplest limit is the NTK regime where the width tends to infinity with
fixed number of data points and input dimension. Here the PTK matrix will
approach its expectation, and the second, covariance term above will tend to
0. This suggests that the first (mean) term encodes much of the interesting
phenomenology near initialization for large networks. We use this fact to
design a dataset that interpolates between small and large initial derivatives
of the C-NFA, even in the non-NTK regime, which will help us demonstrate the
validity of our random matrix theory approach.

We focus most of our random matrix theory analysis on a one-hidden layer
network, as the calculations quickly become complicated with depth. However,
we also show that the mean term can be used to approximately predict the NFA
value for Gaussian data with different spectral decay rates for more
complicated networks (see Appendix C).

###  3.2 Exact predictions with one hidden layer, quadratic activation

Figure 3: Predicted versus observed correlation of the second derivatives of
centered $F$ and $\bar{G}$ on the alignment reversing dataset. Different
shaded color curves correspond to four different seeds for the dataset. The
solid blue curve is the average over all data seeds. The rightmost sub-figure
is a scatter plot of the predicted versus observed correlations of these
second derivatives, with one point for each balance value. We instantiate the
dataset in the proportional regime where width, input dimension, and dataset
size are all equal to $1024$.

We can capture the behavior of both the mean and covariance terms from
Equation (5) in certain high dimensional settings. In particular for one
hidden layer neural networks with quadratic activations, we can exactly
predict the value of Equation (5) in the high dimensional limit. We make the
assumption that $X$ and $Y$ are (asymptotically) freely independent of the
parameters at initialization, and that the resulting average is close to the
value of any specific network initialization (self-averaging). For more
details on the assumptions, see Appendix E. We then show numerically that the
predicted correlation between the second derivatives of the centered NFM and
AGOP matches the observed correlation across individual random instances of
the dataset.

To illustrate our analysis, we derive here the numerator of Equation (4). Let
$M^{(4)}_{X|Y}=(X^{\top}YX)^{2}$, and $M^{(2)}_{X}=X^{\top}X$, and
$F_{a}=W^{\top}{\rm diag}\left(a^{2}\right)W$. Then,

$\displaystyle\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)$
$\displaystyle=\text{tr}\left(M^{(4)}_{X|Y}F_{a}M^{(4)}_{X|Y}F_{a}M^{(2)}_{X}F_{a}\right)~{}.$

We can simplify this alternating product using standard results from random
matrix theory, and factor the parameter and data contributions at the cost of
increasing the number of terms in the expression (Appendix E).

In the simplified case of isotropic data $X^{\top}X=I$, we have:

$\displaystyle\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)$
$\displaystyle=\text{tr}\left(M^{(4)}_{X|Y}\right)^{2}\text{tr}\left(F_{a}^{3}\right)+\text{tr}\left(F_{a}\right)\text{tr}\left(F_{a}^{2}\right)\text{tr}\left((\bar{M}^{(4)}_{X|Y})^{2}\right)~{}.$

From these random matrix calculations, we see that the correlations of the
derivatives are determined by traces of powers of $M^{(4)}_{X|Y}$ (and
$M^{(2)}_{X}$ by the calculations in Section E), which is specific to the
dataset, and $F_{a}$, which is specific to the architecture and
initialization.

#### Manipulating the C-NFA

To numerically explore the validity of the random matrix theory calculations,
we need a way to generate datasets with different values of
$\rho\left(\dot{W}^{\top}\dot{W},\dot{W}^{\top}K\dot{W}\right)$. We construct
a random dataset called the alignment reversing dataset, parameterized by a
balance parameter $\gamma\in(0,1]$ to adversarially disrupt the NFA near
initialization in the regime that width $k$, input dimension $d$, and dataset
size $n$ are all equal ($n=k=d=1024$). By Proposition 12, for the
aforementioned neural architecture, the expected second derivative of the
centered NFM satisfies,
$\mathbb{E}\left[\dot{W}^{\top}\dot{W}\right]=X^{\top}Y\mathbb{E}\left[\mathcal{K}\right]YX=(X^{\top}YX)^{2}$,
while the expected second derivative of the centered AGOP,
$\mathbb{E}\left[\dot{W}^{\top}K\dot{W}\right]=X^{\top}Y\mathbb{E}\left[\mathcal{K}^{2}\right]YX$,
has an additional component $X^{\top}YX\cdot X^{\top}X\cdot X^{\top}YX$. Our
construction exploits this difference in that $X^{\top}X$ becomes
adversarially unaligned to $X^{\top}YX$ as the balance parameter decreases. In
our experiment, we sample multiple random datasets with this construction and
compute the predicted and observed correlation of the second derivatives of
the centered NFA at initialization.

The construction exploits that we can manipulate $X^{\top}YX\cdot
X^{\top}X\cdot X^{\top}YX$ freely of the NFM using a certain choice of $Y$
(see Appendix F for details of the construction). We design the dataset such
that this AGOP-unique term is close to identity, while the NFM second
derivative has many large off-diagonal entries, leading to low correlation
between the second derivatives of the NFM and AGOP.

We observe in Figure 3 that the centered NFA correlations predicted with
random matrix theory closely match the observed values, across individual four
random seeds and for the average of the correlation across them. Crucially, a
single neural network is used across the datasets, confirming the validity of
the self-averaging assumption. The variation in the plot across seeds come
from randomness in the sample of the data, which cause deviations from the
adversarial construction.

##  4 Increasing the centered contribution to the uncentered NFA improves
feature learning

Our theoretical and experimental work has established that gradient based
training leads to alignment of the weight matrices to the PTK feature
covariance. This process is driven by the C-NFA. Therefore, one path towards
improving the correlations defined by the NFA is to increase the contribution
of the C-NFA to the dynamics - as measured, for example, by the ratio
$\text{tr}\left(\bar{W}^{\top}\bar{W}\bar{W}^{\top}K\bar{W}\right)\text{tr}\left(W^{\top}WW^{\top}KW\right)^{-1}$.
When this ratio is large, the C-NFA contributes significantly in magnitude to
the UC-NFA, indicating successful feature learning. We will discuss two
interventions which increase this ratio, with an eye towards improving
generalization.

###  4.1 Feature learning and initialization

(a) Standard training.

(b) Speed Limited Optimization.

Figure 4: The uncentered and centered NFA for the first layer weight matrix as
a function of initialization, with and without SLO. We consider a two hidden
layer network with ReLU activations, where we set $C_{0}=500$, and
$C_{1}=C_{2}=0.002$. The third column shows the ratio of the unnormalized
C-NFA correlation to the UC-NFA:
${\text{tr}\left(\bar{W}^{\top}\bar{W}\bar{W}^{\top}K\bar{W}\right)\cdot{\text{tr}\left(W^{\top}WW^{\top}KW\right)}^{-1}}$.
We show that SLO pushes the ratio of the unnormalized C-NFA to UC-NFA close to
$1$ and improves the UC-NFA across initialization scales.

One factor that modulates the level of feature learning is the scale of the
initialization in each layer. In our experiments training networks with
unmodified gradient descent, we observe that the centered NFA will
increasingly dominate the uncentered NFA with decreasing initialization (third
panel, first row, Figure 4). Further, as the centered NFA increases in
contribution to the uncentered quantity, the strength of the UC-NFA, and to a
lesser extent, the C-NFA increases (Figure 4). The decreases in correspondence
between these quantities is also associated with a decrease in the feature
quality of the NFM (Appendix H) and test error (Figure 2) for the chain
monomial task.

For fully-connected networks with homogeneous activation functions, such as
ReLU, and no normalization layers, decreasing initialization scale is
equivalent to decreasing the scale of the outputs, since we can write
$f(Wx)=a^{-p}f(aWx)$ for any scalar $a$ for any homogenous activation $f$.
This in turn is equivalent to increasing the scale of the labels. Therefore,
decreasing initialization forces the weights to change more in order to fit
the labels, leading to more change in $F$ from its initialization.
Conceptually, this aligns with the substantial line of empirical and
theoretical evidence that increasing initialization scale or output scaling
transitions training between the lazy and feature learning regimes (Chizat et
al., 2019; Woodworth et al., 2020; Agarwala et al., 2020; Lyu et al., 2023).

This relationship suggests that small initialization can be broadly applied to
increase the change in F, and the strength of the UC-NFA. However, this may
not be ideal; for example, if the activation function is differentiable at
$0$, small initialization leads to a network which is approximately linear.
This may lead to low expressivity unless the learning dynamics can increase
the weight magnitude.

###  4.2 Speed Limited Optimization

(a) Standard training. (b) SLO.

Figure 5: The NFM and AGOP from a trained network with and without Speed
Limited Optimization with fixed initialization scale on the chain monomial
task. With standard training and large initialization, the NFM remains close
to identity, and does not learn task structure. SLO increases feature
learning, causing the NFM to become highly correlated with the AGOP and learn
the target features.

We can instead design an intervention which can increase feature learning
without the need to decrease the initialization scale. We do so by fixing the
learning speed layerwise to constant values, which causes the C-NFA to
dominate the UC-NFA dynamics. For weights at layer $\ell$ and learning rate
$\eta>0$, we introduce Speed Limited Optimization (SLO), which is
characterized by the following update rule,

$\displaystyle W^{(\ell)}_{t+1}\leftarrow W^{(\ell)}_{t}+\eta\cdot
C_{\ell}\cdot\frac{\nabla_{W^{(\ell)}_{t}}\mathcal{L}}{\|\nabla_{W^{(\ell)}_{t}}\mathcal{L}\|}$

where the hyperparameter $C_{\ell}\geq 0$ controls the amount of learning in
layer $\ell$. We expect this rule to increase the strength of the UC-NFA in
layers where $C_{\ell}$ is large relative to $C_{m}$ for $m\neq\ell$, as
$W^{(\ell)}$ will be forced to change significantly from initialization. As a
result,
$\left\|W_{0}^{(\ell)}\right\|\left\|W_{t}^{(\ell)}\right\|^{-1}\rightarrow 0$
for large $t$, causing the centered and uncentered NFA to coincide for this
layer.

We demonstrate the effects of SLO on the chain-monomial task (Equation (1)).
We found that fixing the learning speed the be high in the first layer and low
in the remaining layers causes the ratio of the unnormalized C-NFA to the UC-
NFA to become close to $1$ across initialization scales (Figure 4b). The same
result holds for the SVHN dataset (see Appendix K). We note that this
intervention can be applied to target underperforming layers and improve
generalization in deeper networks (see Appendix J for details).

We observe that the correlations from both the C-NFA and the UC-NFA become
close to $1$ after training with SLO, independent of initialization scale
(Appendix H). Further, the quality of the features learned, measured by the
similarity of $F$ and $\bar{G}$ to the true EGOP, significantly improve and
are more similar to each other with SLO (Figure 5), even with large
initialization. In contrast, in standard training the UC-NFA fails to develop
with large initialization, as $F$ resembles identity (no feature learning),
while $\bar{G}$ only slightly captures the relevant features.

##  5 Discussion

#### Centered NFA through training

Our theoretical analysis in this work shows that the PTK feature covariance at
initialization has a relatively simple structure in terms of the weights and
the data. To predict the correlation of the NFA later in training, we will
likely need to account for the change in this matrix. One should be able to
predict this development at short times by taking advantage of the fact that
eigenvectors of the Hessian change slowly during training (Bao et al., 2023).
An alternative approach would be to use a quadratic model for neural network
dynamics (Agarwala et al., 2022; Zhu et al., 2022) which can capture dynamics
of the NTK (and therefore the PTK).

#### Depth and general activations

We provide an analytic prediction for the C-NFA at initialization for one
hidden layer networks with quadratic activation functions, and provide first-
order predictions for two hidden layer networks with quadratic and ReLU
activations. A followup direction would be to use random matrix analysis to
more accurately predict the NFA for deep architectures with general
activations. A promising approach in this direction is to approximate non-
linear activations with a linear expansion and independent correction terms
(Adlam et al., 2019; Adlam and Pennington, 2020; El Karoui, 2008; Hu and Lu,
2022).

More generally it’s an open question as to how the NFA develops with depth in
more complicated architectures like CNNs and transformers. More research is
needed to understand how PTK-weight alignment relates to generalization and
trainability in deep networks.

#### Speed Limited Optimization

There are prior works that implement differential learning rates across layers
(Howard and Ruder, 2018; Singh et al., 2015), though these differ from
differential learning speeds as in SLO, in which the L2-norm of the gradients
are fixed, and not the scaling of the gradients. Our intervention is more
similar to the LARS (You et al., 2017) and LAMB (You et al., 2019) optimizers
that fix learning speeds to decrease training time in ResNet and BERT
architectures. However, these optimizers fix learning speeds to the norm of
the weight matrices, while our intervention sets learning speeds to free
hyperparameters in order to increase feature learning, as measured by the NFA.
We also note as an additional explanation for the success of SLO in our
experiments, the PTK feature covariance may change slowly (or just in a single
direction) with this optimizer, because the learning rate is small for later
layers, allowing the first layer weights to align with the PTK.

#### Adaptive and stochastic optimizers

We demonstrate in this work that the dataset (by our construction in Section
3) and optimizer (with SLO in Section 4) play a significant role in the
strength of feature learning. Important future work would be to understand
other settings where significant empirical differences exist between
optimization choices. In particular, analyzing the NFA may clarify the role of
gradient batch size and adaptive gradient methods in generalization (Zhu et
al., 2023).

##  6 Acknowledgements

We thank Lechao Xiao for detailed feedback on the manuscript. We also thank
Jeffrey Pennington for helpful discussions.

## References

  * Abbe et al. [2022] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz.  The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks.  In _Conference on Learning Theory_ , 2022.
  * Adlam and Pennington [2020] Ben Adlam and Jeffrey Pennington.  The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization.  In _International Conference on Machine Learning_ , pages 74–84. PMLR, 2020.
  * Adlam et al. [2019] Ben Adlam, Jake Levinson, and Jeffrey Pennington.  A random matrix perspective on mixtures of nonlinearities for deep learning.  _arXiv preprint arXiv:1912.00827_ , 2019.
  * Agarwala et al. [2020] Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schoenholz.  Temperature check: theory and practice for training models with softmax-cross-entropy losses.  _arXiv preprint arXiv:2010.07344_ , 2020.
  * Agarwala et al. [2022] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington.  Second-order regression models exhibit progressive sharpening to the edge of stability.  _arXiv preprint arXiv:2210.04860_ , 2022.
  * Allen-Zhu and Li [2019] Zeyuan Allen-Zhu and Yuanzhi Li.  What Can ResNet Learn Efficiently, Going Beyond Kernels?  In _Advances in Neural Information Processing Systems_ , 2019.
  * Ba et al. [2022] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang.  High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation.  _arXiv preprint arXiv:2205.01445_ , 2022.
  * Bao et al. [2023] Xuchan Bao, Bietti Alberto, Aaron Defazio, and Vivien Cabannes.  Hessian inertia in neural networks.  _1st Workshop on High-dimensional Learning Dynamics_ , 2023.
  * Barak et al. [2022] Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang.  Hidden progress in deep learning: Sgd learns parities near the computational limit.  _arXiv preprint arXiv:2207.08799_ , 2022.
  * Beaglehole et al. [2023] Daniel Beaglehole, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin.  Mechanism of feature learning in convolutional neural networks.  _arXiv preprint arXiv:2309.00570_ , 2023.
  * 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.
  * Chizat et al. [2019] Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  _Advances in neural information processing systems_ , 32, 2019.
  * Damian et al. [2022] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi.  Neural networks can learn representations with gradient descent.  In _Conference on Learning Theory_ , pages 5413–5452. PMLR, 2022\.
  * El Karoui [2008] Noureddine El Karoui.  Spectrum estimation for large dimensional covariance matrices using random matrix theory.  2008\.
  * Ghorbani et al. [2019] B. Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Limitations of Lazy Training of Two-layers Neural Networks.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Howard and Ruder [2018] Jeremy Howard and Sebastian Ruder.  Universal language model fine-tuning for text classification.  _arXiv preprint arXiv:1801.06146_ , 2018.
  * Hristache et al. [2001] Marian Hristache, Anatoli Juditsky, Jorg Polzehl, and Vladimir Spokoiny.  Structure adaptive approach for dimension reduction.  _Annals of Statistics_ , pages 1537–1566, 2001.
  * Hu and Lu [2022] Hong Hu and Yue M Lu.  Universality laws for high-dimensional learning with random features.  _IEEE Transactions on Information Theory_ , 69(3):1932–1964, 2022.
  * Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clement Hongler.  Neural Tangent Kernel: Convergence and Generalization in Neural Networks.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Ji and Zhu [2020] Guangda Ji and Zhanxing Zhu.  Knowledge distillation in wide neural networks: Risk bound, data efficiency and imperfect teacher.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Kpotufe et al. [2016] Samory Kpotufe, Abdeslam Boularias, Thomas Schultz, and Kyoungok Kim.  Gradients weights improve regression and classification.  _Journal of Machine Learning Research_ , 2016.
  * Krizhevsky et al. [2012] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ , 25, 2012.
  * Li et al. [2020] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang.  Learning over-parametrized two-layer neural networks beyond NTK.  In _Conference on Learning Theory_ , 2020.
  * Lyu et al. [2023] Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S Du, Jason D Lee, and Wei Hu.  Dichotomy of early and late phase implicit biases can provably induce grokking.  _arXiv preprint arXiv:2311.18817_ , 2023.
  * Mildenhall et al. [2021] Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng.  Nerf: Representing scenes as neural radiance fields for view synthesis.  _Communications of the ACM_ , 65(1):99–106, 2021\.
  * Moniri et al. [2023] Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban.  A theory of non-linear feature learning with one gradient step in two-layer neural networks.  _arXiv preprint arXiv:2310.07891_ , 2023.
  * Mousavi-Hosseini et al. [2022] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu.  Neural networks efficiently learn low-dimensional representations with sgd.  _arXiv preprint arXiv:2209.14863_ , 2022.
  * Nichani et al. [2023] Eshaan Nichani, Alex Damian, and Jason D Lee.  Provable guarantees for nonlinear feature learning in three-layer neural networks.  _arXiv preprint arXiv:2305.06986_ , 2023.
  * Parkinson et al. [2023] Suzanna Parkinson, Greg Ongie, and Rebecca Willett.  Linear neural network layers promote learning single-and multiple-index models.  _arXiv preprint arXiv:2305.15598_ , 2023.
  * Radhakrishnan et al. [2022] Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin.  Feature learning in neural networks and kernel machines that recursively learn features.  _arXiv preprint arXiv:2212.13881_ , 2022.
  * Radhakrishnan et al. [2024] Adityanarayanan Radhakrishnan, Mikhail Belkin, and Dmitriy Drusvyatskiy.  Linear recursive feature machines provably recover low-rank matrices, 2024\.
  * Refinetti et al. [2021] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová.  Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed.  In _International Conference on Machine Learning_ , 2021.
  * Simonyan and Zisserman [2014] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Singh et al. [2015] Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor.  Layer-specific adaptive learning rates for deep networks.  In _2015 IEEE 14th International Conference on Machine Learning and Applications (ICMLA)_ , pages 364–368. IEEE, 2015.
  * Trivedi et al. [2014] Shubhendu Trivedi, Jialei Wang, Samory Kpotufe, and Gregory Shakhnarovich.  A consistent estimator of the expected gradient outerproduct.  In _UAI_ , pages 819–828, 2014.
  * Woodworth et al. [2020] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models.  In _Conference on Learning Theory_ , pages 3635–3673. PMLR, 2020\.
  * Yehudai and Shamir [2019] Gilad Yehudai and Ohad Shamir.  On the Power and Limitations of Random Features for Understanding Neural Networks.  In _Advances in Neural Information Processing Systems_ , 2019.
  * You et al. [2017] Yang You, Igor Gitman, and Boris Ginsburg.  Large batch training of convolutional networks.  _arXiv preprint arXiv:1708.03888_ , 2017.
  * You et al. [2019] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training bert in 76 minutes.  _arXiv preprint arXiv:1904.00962_ , 2019.
  * Yuan et al. [2023] Gan Yuan, Mingyue Xu, Samory Kpotufe, and Daniel Hsu.  Efficient estimation of the central mean subspace via smoothed gradient outer products.  _arXiv preprint arXiv:2312.15469_ , 2023.
  * 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_ , 2021.
  * Zhu et al. [2022] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin.  Quadratic models for understanding neural network dynamics.  _arXiv preprint arXiv:2205.11787_ , 2022.
  * Zhu et al. [2023] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin.  Catapults in sgd: spikes in the training loss and their impact on generalization through feature learning.  _arXiv preprint arXiv:2306.04815_ , 2023.

##  Appendix A Omitted proofs of propositions in the main text

###### Proof of Proposition 3.

Note that $\frac{\partial f(x_{0})}{\partial W^{(\ell)}}=\frac{\partial
f(x_{0})}{\partial h_{\ell}}x_{0}^{\top}\in\mathbb{R}^{k_{\ell}\times
k_{\ell}}$. Then,

$\displaystyle\hat{\Theta}_{\ell}(x_{0},z_{0})$
$\displaystyle=\text{tr}\left(\frac{\partial f(x_{0})}{\partial
h_{\ell}}x_{\ell}^{\top}z_{\ell}\frac{\partial f(z_{0})}{\partial
h_{\ell}}^{\top}\right)$ $\displaystyle=\frac{\partial f(x_{0})}{\partial
h_{\ell}}^{\top}\frac{\partial f(z_{0})}{\partial h_{\ell}}\cdot
x_{\ell}^{\top}z_{\ell}$ $\displaystyle=\mathcal{K}^{(\ell)}(x_{0},z_{0})\cdot
x_{\ell}^{\top}z_{\ell}~{}.$

∎

###### Proof of Proposition 6.

At initialization, we have, $\dot{W}=\frac{\partial f(X)}{\partial
h_{1}}^{\top}\dot{\mathcal{L}}X$. Therefore, using that
$\mathcal{K}\equiv\frac{\partial f(X)}{\partial h_{1}}\frac{\partial
f(X)}{\partial h_{1}}^{\top}$,

$\displaystyle\dot{W}^{\top}\dot{W}=X^{\top}\dot{\mathcal{L}}\frac{\partial
f(X)}{\partial h_{1}}\frac{\partial f(X)}{\partial
h_{1}}^{\top}\dot{\mathcal{L}}X=X^{\top}\dot{\mathcal{L}}\mathcal{K}\dot{\mathcal{L}}X~{},$

and,

$\displaystyle\dot{W}^{\top}K\dot{W}$
$\displaystyle=X^{\top}\dot{\mathcal{L}}\frac{\partial f(X)}{\partial
h_{1}}\,K\,\frac{\partial f(X)}{\partial h_{1}}^{\top}\dot{\mathcal{L}}X$
$\displaystyle=X^{\top}\dot{\mathcal{L}}\frac{\partial f(X)}{\partial
h_{1}}\frac{\partial f(X)}{\partial h_{1}}^{\top}\frac{\partial f(X)}{\partial
h_{1}}\frac{\partial f(X)}{\partial h_{1}}^{\top}\dot{\mathcal{L}}X$
$\displaystyle=X^{\top}\dot{\mathcal{L}}\mathcal{K}^{2}\dot{\mathcal{L}}X~{}.$

∎

##  Appendix B Additional centerings of the NFA

#### Double-centered NFA

One may additionally center the PTK feature map to understand the co-evolution
of the PTK feature covariance and the weight matrices. In this work, we
consider such a centering that we refer to as the double-centered NFA.

######  Ansatz 7 (Double-centered NFA).

$\displaystyle(\bar{W}^{(\ell)})^{\top}\bar{W}^{(\ell)}\propto(\bar{W}^{(\ell)})^{\top}\bar{K}^{(\ell)}\bar{W}^{(\ell)}~{},$

where $\bar{K}^{(\ell)}=\left(\frac{\partial f(X)}{\partial
h_{\ell}}-\frac{\partial f_{0}(X)}{\partial
h_{\ell}}\right)^{\top}\left(\frac{\partial f(X)}{\partial
h_{\ell}}-\frac{\partial f_{0}(X)}{\partial h_{\ell}}\right)$, and $f_{0}$ is
the neural network at initialization.

However, the double-centered NFA term corresponds to higher-order dynamics
that do not significantly contribute the centered and uncentered NFA (Figure
6) when initialization is large or for early periods of training. Note however
this term becomes relevant over longer periods of training.

Figure 6: Ratio of the unnormalized double-centered NFA correlation to the
centered NFA correlation throughout neural network training. In particular, we
plot
$\text{tr}\left(\bar{W}^{\top}\bar{W}\bar{W}^{\top}\bar{K}\bar{W}\right)\cdot{\text{tr}\left(\bar{W}^{\top}\bar{W}\bar{W}^{\top}K\bar{W}\right)}^{-1}$
throughout training for both layers of a two-hidden layer MLP with ReLU
activations.

#### Isolating alignment of the PTK to the initial weight matrix

One may also center just the PTK feature map, while substituting the initial
weights for $W$ to isolate how the PTK feature covariance aligns to the weight
matrices. To measure this alignment, we consider the PTK-centered NFA.

######  Ansatz 8 (PTK-centered NFA).

$\displaystyle(W_{0}^{(\ell)})^{\top}W_{0}^{(\ell)}\propto(W_{0}^{(\ell)})^{\top}\bar{K}^{(\ell)}W_{0}^{(\ell)}~{},$

where $W_{0}^{(\ell)}$ is the initial weight matrix at layer $\ell$.

However, this correlation decreases through training, indicating that the
correlation of these quantities does not drive alignment between the
uncentered NFM and AGOP (Figure 7).

(a) Isotropic data.

(b) Data spectrum with decay $\lambda_{k}\sim\frac{1}{1+k^{2}}$.

Figure 7: PTK-centered NFA correlation throughout training for both layers of
a two-hidden layer MLP with ReLU activations on Gaussian data with two
different spectra.

##  Appendix C Extending our theoretical predictions to depth and general
activations

Precise predictions of the C-NFA become more complicated with additional depth
and general activation functions. However, we note that the deep C-NFA will
remain sensitive to a first-order approximation in which $K$ is replaced by
its expectation. We demonstrate that this term qualitatively captures the
behavior of the C-NFA for 2 hidden layer architectures with quadratic and, to
a lesser extent, ReLU activation functions in Figure 8. In this experiment, we
sample Gaussian data with mean $0$ and covariance with a random eigenbasis. We
parameterize the eigenvalue decay of the covariance matrix by a parameter
$\alpha$, called the data decay rate, so that the eigenvalues have values
$\lambda_{k}=\frac{1}{1+k^{\alpha}}$. As $\alpha$ approaches $0$ or $\infty$
the data covariance approaches a projector matrix.

In this experiment, we see that the data covariance spectrum will also
parameterize the eigenvalue decay of $\mathbb{E}\left[K\right]$, allowing us
to vary how close the expected PTK matrix (and its dual, the PTK feature
covariance) is from a projector, where the NFA holds exactly. We see that for
intermediate values of $\alpha$, both the observed and the predicted
derivatives of the C-NFA decreases in value.

We plot the observed values in two settings corresponding to different
asymptotic regimes. One setting is the proportional regime, where $n=k=d=128$.
The other is the NTK regime where $n=d=128$ and $k=1024$. For the quadratic
case, as the network approaches infinite width, the prediction more closely
matches the observed values. Additional terms corresponding to the nonlinear
part of $\phi^{\prime}$ in ReLU networks, the derivative of the activation
function, are required to capture the correlation more accurately in this
case.

(a) Quadratic.

(b) ReLU.

Figure 8: Observed versus the first-order predicted C-NFA for the input to the
first layer of a two hidden layer MLP. The dashed line is neural network width
$k=n=d=128$, where $n$ and $d$ are the number of data point and data
dimension, respectively, while the solid line uses $n=d=128$ and $k=1024$.

##  Appendix D Proofs and derivations

###### Proof of Proposition 5.

Consider the second derivatives of the NFM and the AGOP.

$\displaystyle\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}\bar{W})$
$\displaystyle=\dot{W}^{\top}\bar{W}+\bar{W}^{\top}\dot{W}~{}.$

Then,

$\displaystyle\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}\bar{W})$
$\displaystyle=2\dot{W}^{\top}\dot{W}+\ddot{W}^{\top}\bar{W}+\bar{W}^{\top}\ddot{W}~{}.$

At initialization $\bar{W}=0$, therefore,

$\displaystyle\bar{W}^{\top}\bar{W}=0,\quad\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}\bar{W})=0,\quad\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}\bar{W})=2\dot{W}^{\top}\dot{W}~{}.$
(6)

Meanwhile, for the (centered) AGOP,

$\displaystyle\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}K\bar{W})$
$\displaystyle=\dot{W}^{\top}K\bar{W}+\bar{W}^{\top}\dot{K}_{z}\bar{W}+\bar{W}^{\top}K\dot{W}~{}.$

Then,

$\displaystyle\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}K\bar{W})$
$\displaystyle=\ddot{W}^{\top}K\bar{W}+\dot{W}^{\top}\dot{K}\bar{W}+\dot{W}^{\top}K_{z}\dot{W}$
$\displaystyle+\dot{W}^{\top}K\dot{W}+\bar{W}^{\top}\dot{K}\dot{W}+\bar{W}^{\top}K\ddot{W}$
$\displaystyle+\dot{W}^{\top}\dot{K}\bar{W}+\bar{W}^{\top}\ddot{K}\bar{W}+\bar{W}^{\top}\dot{K}\dot{W}~{}.$

At initialization,

$\displaystyle\bar{W}^{\top}K\bar{W}=0,$
$\displaystyle\quad\frac{\operatorname{d}}{\operatorname{d}t}(\bar{W}^{\top}K\bar{W})=0,$
(7)
$\displaystyle\frac{\operatorname{d}^{2}}{\operatorname{d}t^{2}}(\bar{W}^{\top}K\bar{W})$
$\displaystyle=2\dot{W}^{\top}K\dot{W}~{}.$ (8)

Therefore, the correlation between the centered neural feature matrix and the
centered AGOP at initialization is determined by,

$\displaystyle\rho\left(\dot{W}^{\top}\dot{W},\dot{W}^{\top}K\dot{W}\right)~{}.$

Both sides simplify at initialization:

$\displaystyle\dot{W}^{\top}\dot{W}=X^{\top}YKYX,\quad\dot{W}^{\top}K\dot{W}=X^{\top}YK^{2}YX~{},$

where $K$ is its value at initialization. ∎

##  Appendix E Free probability calculations of C-NFA

In order to understand the development of the NFA, we analyze the centered NFA
in the limit that learning rate is much smaller than the initialization for a
one hidden layer MLP with quadratic activations. We write this particular
network as,

$\displaystyle f(x)=a^{\top}(Wx)^{2}~{},$

where $a\in\mathbb{R}^{1\times k}$ and $W\in\mathbb{R}^{k\times d}$, where $d$
is the input dimension and $k$ is the width. In this case, the NFA has the
following form,

$\displaystyle\rho\left(F,\bar{G}\right)=\frac{\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)}{\text{tr}\left((X^{\top}Y\mathcal{K}YX)^{2}\right)^{-1/2}\text{tr}\left((X^{\top}Y\mathcal{K}^{2}YX)^{2}\right)^{-1/2}}~{},$
(9)

where $\mathcal{K}=XW^{\top}{\rm diag}\left(a\right)^{2}WX^{\top}$.

We assume three properties hold in the finite dimensional case we consider,
that will hold asymptotically in the infinite dimensional limit.

######  Assumption 9 (Self-averaging).

We assume that computing the average of the NFA quantities across
initializations is equal to the quantities themselves in the high-dimensional
limit.

######  Assumption 10 (Asymptotic freeness).

We assume that the collections $\\{X,Y\\}$ and $\\{W,a\\}$ are asymptotically
free with respect to the operator $\mathbb{E}\left[{\rm tr}(\cdot)\right]$,
where ${\rm tr}[M]=\frac{1}{n}\sum_{i=1}^{n}M_{ii}$.

######  Assumption 11 (Commutativity of expectation).

We will also make the approximation that the expectation commutes with ratio
and square root.

$\mathbb{E}\left[\rho(A,B)\right]=\mathbb{E}\left[\frac{{\rm
tr}(A^{\top}B)}{\sqrt{{\rm tr}(A^{\top}A){\rm
tr}(B^{\top}B)}}\right]\approx\frac{\mathbb{E}\left[{\rm
tr}(A^{\top}B)\right]}{\sqrt{\mathbb{E}\left[{\rm tr}(A^{\top}A)]{\rm E}[{\rm
tr}(B^{\top}B)\right]}}$ (10)

We will compute the expected values of the centered NFA under these
assumptions. In the remainder of the section we will drop the
$\mathbb{E}\left[\cdot\right]$ in the trace for ease of notation.

###  E.1 Free probability identities

The following lemmas will be useful: let $\\{\bar{C}_{i}\\}$ and $\\{R_{i}\\}$
be freely independent of each other with respect to ${\rm tr}$, with ${\rm
tr}[\bar{C}_{i}]=0$. Alternating words have the following products:

${\rm tr}[\bar{C}_{1}R_{1}]=0$ (11) ${\rm
tr}[\bar{C}_{1}R_{1}\bar{C}_{2}R_{2}]={\rm tr}[R_{1}]{\rm tr}[R_{2}]{\rm
tr}[\bar{C}_{1}\bar{C}_{2}]$ (12) ${\rm
tr}[\bar{C}_{1}R_{1}\bar{C}_{2}R_{2}\bar{C}_{3}R_{3}]={\rm tr}[R_{1}]{\rm
tr}[R_{2}]{\rm tr}[R_{3}]{\rm tr}[\bar{C}_{1}\bar{C}_{2}\bar{C}_{3}]$ (13)
$\begin{split}&{\rm
tr}[\bar{C}_{1}R_{1}\bar{C}_{2}R_{2}\bar{C}_{3}R_{3}\bar{C}_{4}R_{4}]={\rm
tr}[R_{1}]{\rm tr}[R_{2}]{\rm tr}[R_{3}]{\rm tr}[R_{4}]{\rm
tr}[\bar{C}_{1}\bar{C}_{2}\bar{C}_{3}\bar{C}_{4}]+\\\ &{\rm tr}[R_{1}]{\rm
tr}[R_{3}]{\rm tr}[\bar{R}_{2}\bar{R}_{4}]{\rm tr}[\bar{C}_{1}\bar{C}_{2}]{\rm
tr}[\bar{C}_{3}\bar{C}_{4}]+{\rm tr}[R_{2}]{\rm tr}[R_{4}]{\rm
tr}[\bar{R}_{1}\bar{R}_{3}]{\rm tr}[\bar{C}_{2}\bar{C}_{3}]{\rm
tr}[\bar{C}_{1}\bar{C}_{4}]\end{split}$ (14)

where $\bar{R}_{i}\equiv R_{i}-{\rm tr}[R_{i}]$.

Applying these identities to the one hidden layer quadratic case, we use the
following definitions:

$R=W^{\top}{\rm diag}\left(a^{2}\right)W,~{}A=(X^{\top}YX)^{2},~{}B=X^{\top}X$
(15)

Crucially, $R$ is freely independent of the set $\\{A,B\\}$. We will also use
the notation $\bar{M}$ to indicated the centered version of $M$,
$\bar{M}=M-{\rm tr}[M]$.

###  E.2 Numerator term of NFA

The numerator in Equation (9) is

$\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)=\text{tr}\left(ARARBR\right)$
(16)

Re-writing $A=\bar{A}+{\rm tr}[A]$ and $B=\bar{B}+{\rm tr}[B]$ we have:

$\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)=\text{tr}\left((\bar{A}+{\rm
tr}[A])R(\bar{A}+{\rm tr}[A])R(\bar{B}+{\rm tr}[B])R\right)$ (17)

This expands to

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)&=\text{tr}\left(\bar{A}R\bar{A}R\bar{B}R\right)+2\text{tr}\left(A\right)\text{tr}\left(\bar{A}R\bar{B}R^{2}\right)+\text{tr}\left(B\right)\text{tr}\left(\bar{A}R\bar{A}R^{2}\right)\\\
&\text{tr}\left(A\right)^{2}\text{tr}\left(\bar{B}R^{3}\right)+2\text{tr}\left(A\right)\text{tr}\left(B\right)\text{tr}\left(\bar{A}R^{3}\right)+\text{tr}\left(A\right)^{2}\text{tr}\left(B\right)\text{tr}\left(R\right)^{3}\end{split}$
(18)

Using the identities we arrive at:

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}^{2}YX\right)&={\rm
tr}[A]^{2}{\rm tr}[B]\text{tr}\left(R^{3}\right)\\\ &+2{\rm tr}[A]{\rm
tr}[R^{2}]{\rm tr}[R]\text{tr}\left(\bar{A}\bar{B}\right)\\\ &+{\rm tr}[B]{\rm
tr}[R]{\rm tr}[R^{2}]\text{tr}\left(\bar{A}^{2}\right)\\\ &+{\rm
tr}[R]^{3}\text{tr}\left(\bar{A}^{2}\bar{B}\right)\end{split}$ (19)

###  E.3 First denominator term of NFA

The first denominator term in Equation (9) is

$\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}YX\right)=\text{tr}\left(ARAR\right)$
(20)

This is a classic free probability product:

$\text{tr}\left(X^{\top}Y\mathcal{K}YXX^{\top}Y\mathcal{K}YX\right)={\rm
tr}[A^{2}]{\rm tr}[R]^{2}+{\rm tr}[A]^{2}\text{tr}\left(R^{2}\right)-{\rm
tr}[A]^{2}{\rm tr}[R]^{2}$ (21)

which can be derived from the lemmas.

###  E.4 Second denominator term of NFA

For the second denominator term of Equation (9) we have

$\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)=\text{tr}\left(ARBRARBR\right)$
(22)

Expanding the first $A$ we have

$\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)=\text{tr}\left(\bar{A}RBRARBR\right)+{\rm
tr}[A]\text{tr}\left(R^{2}BRARB\right)$ (23)

Next we expand the first $B$:

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)&=\text{tr}\left(\bar{A}R\bar{B}RARBR\right)+{\rm
tr}[B]\text{tr}\left(\bar{A}R^{2}ARBR\right)\\\ &+{\rm tr}[A]{\rm
tr}[B]\text{tr}\left(R^{3}ARB\right)+{\rm
tr}[A]\text{tr}\left(R^{2}\bar{B}RARB\right)\end{split}$ (24)

The next $A$ gives us

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)&=\text{tr}\left(\bar{A}R\bar{B}R\bar{A}RBR\right)+2{\rm
tr}[A]\text{tr}\left(\bar{A}R\bar{B}R^{2}BR\right)+{\rm
tr}[B]\text{tr}\left(\bar{A}R^{2}\bar{A}RBR\right)\\\ &+2{\rm tr}[A]{\rm
tr}[B]\text{tr}\left(R^{3}\bar{A}RB\right)+{\rm tr}[A]^{2}{\rm
tr}[B]\text{tr}\left(R^{4}B\right)+{\rm
tr}[A]^{2}\text{tr}\left(R^{2}\bar{B}R^{2}B\right)\end{split}$ (25)

Expanding the final $B$ we have

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)&=\text{tr}\left(\bar{A}R\bar{B}R\bar{A}R\bar{B}R\right)+2{\rm
tr}[B]\text{tr}\left(\bar{A}R\bar{B}R\bar{A}R^{2}\right)+2{\rm
tr}[A]\text{tr}\left(\bar{A}R\bar{B}R^{2}\bar{B}R\right)\\\ &+4{\rm tr}[A]{\rm
tr}[B]\text{tr}\left(R^{3}\bar{A}R\bar{B}\right)+2{\rm tr}[A]{\rm
tr}[B]^{2}\text{tr}\left(R^{4}\bar{A}\right)+2{\rm tr}[A]^{2}{\rm
tr}[B]\text{tr}\left(R^{4}\bar{B}\right)\\\ &+{\rm tr}[A]^{2}{\rm
tr}[B]^{2}{\rm tr}[R^{4}]+{\rm tr}[A]^{2}{\rm
tr}[R^{2}\bar{B}R^{2}\bar{B}]+{\rm tr}[B]^{2}{\rm
tr}[R^{2}\bar{A}R^{2}\bar{A}]\end{split}$ (26)

Now all terms are in the form of alternating products from the lemma. This
means we can factor out the non-zero traces of the other terms. Simplifying we
have:

$\begin{split}\text{tr}\left(X^{\top}Y\mathcal{K}^{2}YXX^{\top}Y\mathcal{K}^{2}YX\right)&={\rm
tr}[R]^{4}\text{tr}\left((\bar{A}\bar{B})^{2}\right)+2{\rm tr}[R]^{2}({\rm
tr}[R^{2}]-{\rm tr}[R]^{2}){\rm tr}[\bar{A}\bar{B}]^{2}\\\ &+2{\rm
tr}[R]^{2}{\rm tr}[R^{2}]\left({\rm
tr}[B]\text{tr}\left(\bar{A}^{2}\bar{B}\right)+{\rm
tr}[A]\text{tr}\left(\bar{A}\bar{B}^{2}\right)\right)\\\ &+4{\rm tr}[A]{\rm
tr}[B]{\rm tr}[R^{3}]{\rm tr}[R]\text{tr}\left(\bar{A}\bar{B}\right)+{\rm
tr}[A]^{2}{\rm tr}[B]^{2}{\rm tr}[R^{4}]\\\ &+{\rm tr}[A]^{2}{\rm
tr}[\bar{B}^{2}]{\rm tr}[R^{2}]^{2}+{\rm tr}[\bar{A}]^{2}{\rm tr}[B^{2}]{\rm
tr}[R^{2}]^{2}\end{split}$ (27)

All terms of the NFA are now in terms of traces of the matrices $A$, $B$, and
$R$ and functions on each term separately. The matrices $A$ and $B$ are
determined by the data, while the moments of the eigenvalues of $R$ are
determined by the initialization distribution of the weights in the neural
network, and neither training nor the data.

##  Appendix F Balance dataset

The data consists of a mixture of two distributions from which two subsets of
the data $X_{1}$ and $X_{2}$ are sampled from, and is parametrized by a
balance parameter $\gamma\in(0,1]$ and two variance parameters
$\epsilon_{1},\epsilon_{2}>0$. The subset $X_{1}$ which has label $y_{1}=1$
and constitutes a $\gamma$ fraction of the entire dataset, is sampled from a
multivariate Gaussian distribution with mean $0$ and covariance

$\displaystyle\Sigma=\bm{1}\bm{1}^{\top}+\epsilon_{1}\cdot I~{}.$

Then the second subset, $X_{2}$, is constructed such that
$X_{2}^{\top}X_{2}\approx(X_{1}^{\top}X_{1})^{-2}$, and has labels $y_{2}=0$.
Then, for balance parameter $\gamma$ sufficiently small, the AGOP second
derivative approximately satisfies,

$\displaystyle\mathbb{E}\left[\dot{W}^{\top}K\dot{W}\right]$
$\displaystyle\sim X^{\top}YXX^{\top}XX^{\top}YX$
$\displaystyle=X_{1}^{\top}X_{1}X^{\top}XX_{1}^{\top}X_{1}$
$\displaystyle\approx X_{1}^{\top}X_{1}X_{2}^{\top}X_{2}X_{1}^{\top}X_{1}$
$\displaystyle\approx I~{},$

In contrast, the NFM second derivative,
$\mathbb{E}\left[\dot{W}^{\top}\dot{W}\right]=(X^{\top}YX)^{2}=(X_{1}^{\top}X_{1})^{2}\approx\Sigma^{2}$,
will be significantly far from identity.

Motivated by this derivation, we construct $X_{2}$ by the following procedure:

  1. 1.

Extract singular values $S_{1}$ and right singular vectors $U_{1}$ from a
singular-value decomposition (SVD) of $X_{1}^{\top}X_{1}$.

  2. 2.

Extract the left singular vectors $V_{2}$ from a sample $\tilde{X}_{2}$ that
is sampled from the same distribution as $X_{1}$.

  3. 3.

Construct $X_{2}=V_{2}S_{1}^{-1}U_{1}^{\top}$.

  4. 4.

Where $X=X_{1}\oplus X_{2}$, Set $X\leftarrow X+\epsilon_{2}Z$, where
$Z\sim\mathcal{N}(0,I)$.

  5. 5.

Set $y\leftarrow y+10^{-5}\cdot\bm{1}$.

Note that
$U_{1}S_{1}^{-1}V_{2}^{\top}V_{2}S_{1}^{-1}U_{1}^{\top}=U_{1}S_{1}^{-2}U_{1}^{\top}=(X_{1}^{\top}X_{1})^{-2}$,
therefore, we should set $X_{2}=V_{2}S_{1}^{-1}U_{1}^{\top}$ to get
$X_{2}^{\top}X_{2}=(X_{1}^{\top}X_{1})^{-2}$. Regarding the variance
parameters, in practice we set $\epsilon_{1}=0.5$ and $\epsilon_{2}=10^{-2}$.

######  Proposition 12 (Expected NFM and AGOP).

For a one hidden layer quadratic network, $f(x)=a^{\top}(Wx)^{2}$, with
$a\sim\mathcal{N}(0,I)$ and $W\sim\frac{1}{\sqrt{k}}\cdot\mathcal{N}(0,I)$,

$\displaystyle\mathbb{E}_{a,W}\left[\dot{W}^{\top}\dot{W}\right]=(X^{\top}YX)^{2}~{},$

and,

$\displaystyle\mathbb{E}_{a,W}\left[\dot{W}^{\top}K\dot{W}\right]$
$\displaystyle=3\cdot\text{tr}\left(X^{\top}X\right)\cdot(X^{\top}YX)^{2}$
$\displaystyle+6X^{\top}YXX^{\top}XX^{\top}YX$

###### Proof of Proposition 12.

$\displaystyle\mathbb{E}\left[\dot{W}^{\top}\dot{W}\right]$
$\displaystyle=X^{\top}YX\mathbb{E}\left[W_{0}^{\top}{\rm
diag}\left(a\right)^{2}W_{0}\right]X^{\top}YX$
$\displaystyle=(X^{\top}YX)^{2}~{}.$

Further,

$\displaystyle\mathbb{E}\left[\dot{W}^{\top}K\dot{W}\right]$
$\displaystyle=X^{\top}Y\mathbb{E}\left[K^{2}\right]YX~{}.$

We note that,

$\displaystyle K^{2}$ $\displaystyle=W_{0}^{\top}{\rm
diag}\left(a\right)^{2}W_{0}X^{\top}XW_{0}^{\top}{\rm
diag}\left(a\right)^{2}W_{0}$ (28)
$\displaystyle=\sum_{s_{1},s_{2}}^{k}\sum_{\alpha}^{n}\sum_{p_{1},p_{2}}^{d}$
(29) $\displaystyle
a_{s_{1}}^{2}a_{s_{2}}^{2}W_{s_{1},p_{1}}X_{\alpha,p_{1}}X_{\alpha,p_{2}}W_{s_{2},p_{2}}XW_{s_{1}}W_{s_{2}}^{\top}X^{\top}~{}.$
(30)

Therefore, applying Wick’s theorem, element $i,j$ of $K^{2}$ satisfies,

$\displaystyle\mathbb{E}\left[K^{2}_{ij}\right]$
$\displaystyle=\sum_{s}^{k}\sum_{\alpha}^{n}\sum_{p_{1},p_{2}}^{d}\mathbb{E}\left[a_{s}^{4}W_{s,p_{1}}X_{\alpha,p_{1}}X_{\alpha,p_{2}}W_{s,p_{2}}X_{i}^{\top}W_{s}W_{s}^{\top}X_{j}\right]$
$\displaystyle=\sum_{s}^{k}\sum_{\alpha}^{n}\sum_{p_{1},p_{2},q_{1},q_{2}}^{d}$
$\displaystyle\mathbb{E}\left[a_{s}^{4}W_{s,p_{1}}W_{s,p_{2}}W_{s,q_{1}}W_{s,q_{2}}X_{\alpha,p_{1}}X_{\alpha,p_{2}}X_{i,q_{1}}X_{j,q_{2}}\right]$
$\displaystyle=3\sum_{s}^{k}\sum_{\alpha}^{n}\sum_{p_{1},p_{2},q_{1},q_{2}}^{d}$
$\displaystyle\Big{(}\mathbb{E}\left[W_{s,p_{1}}W_{s,p_{2}}\right]\mathbb{E}\left[W_{s,q_{1}}W_{s,q_{2}}\right]+$
$\displaystyle\mathbb{E}\left[W_{s,p_{1}}W_{s,q_{1}}\right]\mathbb{E}\left[W_{s,p_{2}}W_{s,p_{2}}\right]+$
$\displaystyle\mathbb{E}\left[W_{s,p_{1}}W_{s,q_{2}}\right]\mathbb{E}\left[W_{s,p_{2}}W_{s,q_{1}}\right]\Big{)}$
$\displaystyle\cdot X_{\alpha,p_{1}}X_{\alpha,p_{2}}X_{i,q_{1}}X_{j,q_{2}}$
$\displaystyle=3\sum_{\alpha}^{n}\sum_{p_{1},p_{2},q_{1},q_{2}}^{d}$
$\displaystyle\Big{(}\delta_{p_{1}p_{2}}\delta_{q_{1}q_{2}}+\delta_{p_{1}q_{1}}\delta_{p_{2}q_{2}}+\delta_{p_{1}q_{2}}\delta_{p_{2}q_{1}}\Big{)}$
$\displaystyle\cdot X_{\alpha,p_{1}}X_{\alpha,p_{2}}X_{i,q_{1}}X_{j,q_{2}}$
$\displaystyle=3\sum_{\alpha}^{n}\Big{(}\sum_{p_{1},q_{1}}^{d}X_{\alpha,p_{1}}X_{\alpha,p_{1}}X_{i,q_{1}}X_{j,q_{1}}$
$\displaystyle+\sum_{p_{1},p_{2}}^{d}X_{\alpha,p_{1}}X_{\alpha,p_{2}}X_{i,p_{1}}X_{j,p_{2}}$
$\displaystyle+\sum_{p_{1},p_{2}}^{d}X_{\alpha,p_{1}}X_{\alpha,p_{2}}X_{i,p_{2}}X_{j,p_{1}}\Big{)}$
$\displaystyle=3\cdot\text{tr}\left(X^{\top}X\right)\cdot
X_{i}^{\top}X_{j}+3\sum_{\alpha}^{n}X_{\alpha}^{\top}X_{i}X_{\alpha}^{\top}X_{j}$
$\displaystyle+3\sum_{\alpha}^{n}X_{\alpha}^{\top}X_{j}X_{\alpha}^{\top}X_{i}$
$\displaystyle=3\cdot\text{tr}\left(X^{\top}X\right)\cdot
X_{i}^{\top}X_{j}+3X_{i}X^{\top}XX_{j}+3X_{j}X^{\top}XX_{i}~{}.$

Finally, we conclude,

$\displaystyle\mathbb{E}\left[K^{2}\right]=3\left(\text{tr}\left(X^{\top}X\right)XX^{\top}+2XX^{\top}XX^{\top}\right)~{},$

giving the second statement of the proposition. ∎

##  Appendix G Varying the data distribution

We verify that our observations for isotropic Gaussian data hold even when the
data covariance has a significant spectral decay. (Figures 9 and 10). We again
consider Gaussian data that is mean $0$ and where the covariance is
constructed from a random eigenbasis. In Figure 9, we substitute the
eignevalue decay as $\lambda_{k}\sim\frac{1}{1+k}$, while in Figure 10, we use
$\lambda_{k}\sim\frac{1}{1+k^{2}}$. We plot the values of the UC-NFA, C-NFA,
train loss, and test loss throughout training for the first and second layer
of a two hidden layer network with ReLU activations, while additionally
varying initialization scale. Similar to Figure 2, we observe that the C-NFA
is more robust to the initialization scale than the UC-NFA, and UC-NFA value
become high through training, while being small at initialization. We see that
the test loss improves for smaller initializations, where the value of the
C-NFA and UC-NFA are higher.

(a)

(b)

Figure 9: Centered NFA correlations. Data covariance decay rate
$\lambda_{k}\sim\frac{1}{1+k}$. Top row is layer 1, bottom row is layer 2.
Train (test) losses are scaled by the maximum train (test) loss achieved so
that they are between $0$ and $1$.

(a)

(b)

Figure 10: Centered NFA correlations. Data covariance decay rate
$\lambda_{k}\sim\frac{1}{1+k^{2}}$. Top row is layer 1, bottom row is layer 2.
Train (test) losses are scaled by the maximum train (test) loss achieved so
that they are between $0$ and $1$.

##  Appendix H Effect of initialization on feature learning

We see that when initialization is small, the C-NFA and UC-NFA are high at the
end of training with and without fixing the learning speed (Figure 4). This is
reflected by the quality of the features learned by the NFM and the
qualitative similarity of the NFM and AGOP at small initialization scale
(Figure 11). Further, we notice that as we increase initialization, without
fixing speeds, the correspondence between the NFM and decreases and the
quality of the NFM features decreases (at a faster rate than the AGOP).
Strikingly, when learning speeds are fixed, the quality of the features in the
AGOP and NFM becomes invariant to the initialization scale.

(a) NFM, speeds not fixed.

(b) AGOP, speeds not fixed.

(c) NFM, speeds fixed.

(d) AGOP, speeds fixed.

Figure 11: The NFM at the end of training as a function of initialization
scale in the first layer weights, with and without fixing learning speeds. The
task again is the chain monomial with rank $r=5$. The title of each plot is
the initialization scale of the first layer $s_{0}$.

##  Appendix I Experimental details

We describe the neural network training and architectural hyperparameters in
the experiments of this paper. Biases were not used for any networks. Further,
in all polynomial tasks, we scaled the label vector to have standard deviation
$1$.

#### Corrupted AGOP

For the experiments in Figure 1, we used $n=384$ data points, $d=32$, $k=128$
as the width in all layers, isotropic Gaussian data, initialization scale
$0.01$ in the first layer and default scale in the second. We used ReLU
activations and two hidden layers. For the experiments in Figure 2,9,10,6, and
7, we used a two hidden layer network with ReLU activations, learning rate
$0.05$, 800 steps of gradient decent, and took correlation/covariance
measurements every 5 steps.

#### Alignment reversing dataset

For the experiments in Figure 3, we used $k=n=d=1024$ for the width, dataset
size, and input dimension, respectively. Further, the traces of powers of
$F_{a}$ are averaged over 30 neural net seeds to decouple these calculated
values from the individual neural net seeds. The mean value plotted in the
first two squares of figure is computed over 10 data seeds.

#### SLO experiments

For the SLO figures (Figures 5, 11, 4), we use isotropic Gaussian data, 600
steps of gradient descent. The learning rates are chosen based on
initialization scale in the first layer. For initialization scales 1, 0.1,
0.01, and 0.001, we used learning rates $0.03$, $0.1$, $0.2$, $0.4$,
respectively. We again used two hidden layers with ReLU activations. We chose
$n=256$, $d=32$, and $k=256$ as the width. We divided the linear readout
weights by $0.01$ at initialization to promote feature learning, and modified
SLO to scale gradients by $(\epsilon+\|\nabla\mathcal{L}\|)^{-1}$, rather than
just the inverse of the norm of the gradient, for $\epsilon=0.1$. This
technique smooths the training dynamics as the parameters approach a loss
minimum, allowing the network to interpolate the labels.

#### Predictions with depth

For the Deep C-NFA predictions (Figure 8), we used $n=128$, $d=128$,
initialization scale of $1$. The low rank task is just the chain monomial of
rank $r=5$. The high rank polynomial task is
$y(x)=\sum_{i=1}^{d}(Qx)_{i}^{2}$, where $Q\in\mathbb{R}^{d\times d}$ is a
matrix with standard normal entries.

#### Real datasets

For the experiments on the SVHN dataset, we train a four hidden layer neural
network with ReLU activations, initialization scale $1.0$ in all layers and
width $256$. For SVHN, we subset the dataset to 4000 points. We train for 3000
epochs with learning rate 0.2 for standard training, and 0.3 for SLO, and take
NFA measurements every 50 epochs. For SLO, we set $C_{0}=2.5$,
$C_{1}=C_{2}=0.4$, and relaxation parameter $\epsilon=0.2$. We pre-process the
dataset so that each pixel is mean $0$ and standard deviation $1$. For the
experiments on CelebA, we train a two hidden layer network on a balanced
subset of 7500 points with Adam with learning rate $0.0001$ and no weight
decay. We use initialization scale 0.02 in the first layer, and width 128. We
train for 500 epochs. We pre-process the dataset by scaling the pixel values
to be between 0 and 1.

##  Appendix J Additional SLO experiments

We demonstrate the SLO can be applied adaptively to increase the strength of
the UC-NFA in all layers of a deep network on the chain monomial task of rank
$r=3$. We train a three hidden layer MLP with ReLU activations and an
initialization scale of $0.1$ by SLO, and find that all layers finish at the
same high UC-NFA (Figure 12). Further, this final UC-NFA value is higher than
the highest UC-NFA achieved by any layer with standard training. The
generalization loss is also lower with SLO on this example, corresponding to
better feature learning (through the UC-NFA).

(a)  Figure 12: Training with SLO where the learning speeds are chosen
adaptively based on the UC-NFA values of all layers. The dashed lines
correspond to training with standard GD.

At every time step we choose $C_{i}=s$ for the layer $i$ with the smallest UC-
NFA correlation value, while setting $C_{j}=s^{-1}$ for all other layers, with
$s=20$. We again modify SLO by dividing the gradients by
$\epsilon\|\nabla\mathcal{L}\|$ for $\epsilon=0.01$. The learning rate is set
to $0.05$ in SLO and $0.25$ for the standard training (gradient descent), and
the networks are trained for 500 epochs. We sample $n=256$ points with $d=32$,
and use width $k=256$.

##  Appendix K Experiments on real datasets

We replicate Figures 1, 4, and 5 on celebrity faces (CelebA) and Street View
House Numbers (SVHN). We begin by showing that one can disrupt the NFA
correspondence by replacing the PTK feature covariance with a random matrix of
the same spectral decay. For this example, we measure the Pearson correlation,
which subtracts the mean of the image. I.e.
$\bar{\rho}(A,B)\equiv\rho\left(A-m(A),B-m(B)\right)$, where $m(A),m(B)$ are
the average of the elements of $A$ and $B$.

Figure 13: Various feature learning measures for the CelebA binary subtask of
predicting glasses. The diagonals of the NFM $\left(W^{\top}W\right)$ (first
plot) and AGOP $\left(W^{\top}KW\right)$ (second plot) of a fully-connected
network are similar to each other. Replacing $K$ with a symmetric matrix $Q$
with the same spectrum but independent eigenvectors obscures the low rank
structure (third plot), and reduces the Pearson correlation of the diagonal
from $\bar{\rho}\left({\rm diag}\left(F\right),{\rm
diag}\left(\bar{G}\right)\right)=0.91$ to $\bar{\rho}\left({\rm
diag}\left(F\right),{\rm diag}\left(W^{\top}QW\right)\right)=0.04$.

(a)

(b)  (c)

Figure 14: We demonstrate on the SVHN dataset, with a 4 hidden layer neural
network with large initialization scale, how SLO can improve the strength of
the UC-NFA, the C-NFA, the ratio of the unnormalized C-NFA to UC-NFA (plot
(a)) and the feature quality (plots (b) and (c)). In plots (b) and (c), we
visualize the diagonal of the NFM and AGOP for the first layer of the trained
network, where SLO was applied with $C_{0}=2.5$, $C_{1}=C_{2}=0.4$.