# Why Do You Grok?
A Theoretical Analysis of Grokking Modular Addition

Mohamad Amin Mohamadi 22footnotemark: 2 Zhiyuan Li111$\hat{\Theta}_{t}$ is the
NTK of the model output on the training data using the parameter at step $t$.
When the model output $f$ is a vector, we use the NTK of its first coordinate
as an approximation, following [32], where
$\hat{\Theta}_{t}\triangleq[\nabla_{\theta}f_{1}(\theta_{t};\mathcal{X}_{\mathrm{train}})][\nabla_{\theta}f_{1}(\theta_{t};\mathcal{X}_{\mathrm{train}})]^{\top}$.
Lei Wu33footnotemark: 3 Danica J. Sutherland22footnotemark: 2 44footnotemark:
4 Toyota Technological Institute at Chicago ‡University of British Columbia
§Peking University ¶Amii
Correspondence to<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

We present a theoretical explanation of the “grokking” phenomenon [41], where
a model generalizes long after overfitting, for the originally-studied problem
of modular addition. First, we show that early in gradient descent, when the
“kernel regime” approximately holds, no permutation-equivariant model can
achieve small population error on modular addition unless it sees at least a
constant fraction of all possible data points. Eventually, however, models
escape the kernel regime. We show that two-layer quadratic networks that
achieve zero training loss with bounded $\ell_{\infty}$ norm generalize well
with substantially fewer training points, and further show such networks exist
and can be found by gradient descent with small $\ell_{\infty}$
regularization. We further provide empirical evidence that these networks as
well as simple Transformers, leave the kernel regime only after initially
overfitting. Taken together, our results strongly support the case for
grokking as a consequence of the transition from kernel-like behavior to
limiting behavior of gradient descent on deep networks.

##  1 Introduction

Understanding the generalization patterns of modern over-parameterized neural
networks has been a long-standing goal of the deep learning community. [41]
demonstrated an intriguing phenomenon they called “grokking” when learning
transformers on small algorithmic tasks: neural networks can find a
generalizing solution long after overfitting to the training dataset with poor
generalization. This observation has lead to a stream of recent works aimed at
uncovering the mechanisms that can lead a network to “grok,” and properties of
the final solutions, on various algorithmic tasks. Later, it was discovered
that grokking can happen in tasks beyond modular arithmetic: in learning
sparse parities [5, 6], image classifiers [27], greatest common divisors [9],
matrix completion [28], and $k$-sparse linear predictors [28].

Grokking has been variously attributed to difficulty of representation
learning [26], the “slingshot” mechanism [47], weight norm [27, 49],
properties of the loss landscape [40], simplicity of the learned solution [37]
and other feature learning mechanisms [23, 43]. Theoretically, [15] presented
an analytical construction for a two-layer MLP that solves modular addition is
compatible with a grokking pattern.111gromov2023grokking claims this solution
is the one found by gradient descent, but this did not seem to be the case in
our experience. [21] demonstrated grokking when training a two-layer MLP on a
polynomial regression problem, as did [53] for XOR data. The notion of delayed
generalization was perhaps earlier observed by [24] when training diagonal
linear networks with label noise SGD and through sharpness minimization,
before it was known as grokking [41].

[28] present a rigorous theoretical framework in which grokking can be
provably demonstrated through a dichotomy of early and late implicit biases of
the training algorithm. More specifically, they attribute the overfitting
stage of grokking to the initial behavior of gradient descent being similar to
a kernel predictor [20, 4, 22], and the generalization stage is to different
late-phase implicit biases such as sharpness minimization [7, 24, 12, 18],
margin maximization [44, 36, 51, 29], or parameter norm minimization [17, 16,
3]. This transition from kernel to rich regime has been widely observed [10,
34, 14, 46, 28]. Consistent with this framework, [21] hypothesized that
grokking can be explained through the transition from the kernel regime to the
“rich” regime, as long as the size of the training dataset is neither too
small (where generalization would be impossible) nor too large (where
generalization would be easy). They provided empirical support by considering
scaling the model output, which is a rough proxy for the rate of feature
learning in modular addition on two-layer MLPs.

Figure 1: Empirical investigation into grokking modular addition on two-layer
networks in the classification task with cross-entropy loss. Left: Change of
empirical NTK111$\hat{\Theta}_{t}$ is the NTK of the model output on the
training data using the parameter at step $t$. When the model output $f$ is a
vector, we use the NTK of its first coordinate as an approximation, following
[32], where
$\hat{\Theta}_{t}\triangleq[\nabla_{\theta}f_{1}(\theta_{t};\mathcal{X}_{\mathrm{train}})][\nabla_{\theta}f_{1}(\theta_{t};\mathcal{X}_{\mathrm{train}})]^{\top}$.
($\lVert\hat{\Theta}_{t}-\hat{\Theta}_{0}\rVert_{F}$) is negligible before
fitting the training data. NTK changes drastically after overfitting, implying
that the delayed generalization might be caused by delayed a transitioning
from kernel to rich regime. Middle: Reducing initialization scale can mitigate
grokking, to the point of completely eliminating the gap between train and
test curves. $\alpha$ denotes scale multiplied by $\theta_{0}$, the initial
weights according to default PyTorch initialization [19]. The dashed lines
indicate train set statistics, and the solid lines correspond to the test set.
Right: Empirical evaluations support a sample complexity of
$\tilde{\operatorname{\mathcal{O}}}(p^{5/3})$ on the classification task with
cross-entropy loss. More details in Section 4.

This dichotomy between early kernel regime and late feature learning triggered
by weak implicit or explicit regularization (i.e. the transition from lazy to
rich regime) seems to be the most promising theory to explain grokking. Even
so, two fundamental questions as to why grokking occurs on the original
problem of modular addition have remained unanswered:

  1. 1.

Why do a wide variety of architectures all fail to generalize in the initial
phase of training, i.e. in the kernel regime? Is it because their kernels
accidentally share some common property, or it is indeed a property of the
modular addition task itself?

  2. 2.

How does weak regularization encourage the network to learn generalizable
features later in training?

Our Contributions. In this work, we address these questions with rigorous
theoretical analyses of learning modular addition with gradient descent on
two-layer MLPs.

We specifically focus on the problem

$a+b\equiv c\pmod{p}$ (1.1)

where $a,b,c\in{\mathbb{Z}}_{p}$ for a fixed $p\in{\mathbb{N}}$. We use
(regularized) gradient descent on a two-layer MLP with quadratic activations,
the same architecture as [15]. We consider two different tasks: it is somewhat
easier to analyze a “regression” task where we use square loss to learn a
function of $(a,b,c)$ that indicates whether (1.1) holds, but we also study
the “classification” task in which we use cross-entropy loss to learn a
$p$-way classifier to predict which $c$ value satisfies (1.1) for a given
$(a,b)$. This discrete, noiseless setting has only a finite number of possible
distinct data points: $p^{3}$ for regression, $p^{2}$ for classification.

To address the first question, we prove in Sections 3.1 and 4.1 that this task
is _fundamentally hard_ for permutation-equivariant kernel methods, due to
inherent symmetries. Thus, permutation-equivariant networks which are well-
approximated by their neural tangent kernel approximation cannot generalize
well. We also prove that, although no such method can generalize, neural
tangent kernel approaches based on our architecture with realistic widths
_can_ achieve zero training error. This result is highly suggestive of why
drastic overfitting with poor generalization has been empirically observed on
this problem across a wide variety of architectures, losses, and learning
algorithms [41, 26, 15, e.g.].

To prove the generalization lower bounds, we develop a novel general technique
to analyze the population $\ell_{2}$ loss of learning general function classes
with predictors of intrinsic dimension $n$ (for instance kernel predictors
with $n$ training points), presented in Section D.4. This framework allows us
to prove lower bounds on population $\ell_{2}$ loss for the general case of
learning modular addition on $m$ summands (rather than 2 or 3) with kernels,
and might be of independent interest.

To address the second question – _why_ this occurs – we identify
$\ell_{\infty}$ norm as an effective and practically-relevant complexity
measure for generalization, which is closely related to the implicit bias of
Adam [52, 55], and show that networks with small $\ell_{\infty}$ norm in the
regression setting (Section 3.2) and large $\ell_{\infty}$-normalized margin
in the classification setting (Section 4.2) can both provably generalize with
far fewer samples than required in the kernel regime. Thus, models with
corresponding implicit biases can generalize well, answering the second
question. In regression, our proofs are based on “optimistic rates” [45] for
the smooth $\ell_{2}$ loss in terms of the Rademacher complexity of networks
with bounded $\ell_{\infty}$ parameter norm. In classification, our proof
applies the PAC-Bayesian framework of [31] to networks with bounded
$\ell_{\infty}$ parameter norm.

In summary, our main contributions are as follows:

  1. 1.

We prove that networks in the kernel regime can only generalize if trained on
at least a constant fraction of all possible data points, i.e. an $\Omega(1)$
portion, for regression (Section 3.1) and classification (Section 4.1).

  2. 2.

We prove that networks with appropriate regularization can generalize with
many fewer samples: $\bm{\operatorname{\mathcal{O}}(1/p)}$ portion of all
possible data points for square loss generalization on the regression task
(Section 3.2), and $\bm{\tilde{\operatorname{\mathcal{O}}}(1/p^{1/3})}$ for
zero-one loss generalization on classification (Section 4.2).

##  2 Preliminaries and Setup

##### Notations.

We use $[p]$ to denote the set $\\{1,\ldots,p\\}$. We use $e_{i}$ to denote
the $i$ standard basis vector, _i.e._ , the vector with $1$ in its $i$th
component and $0$ elsewhere. For vector $a$, we use $a^{\odot 2}$ denotes the
element-wise square of the vector $a$ For any nonempty set $\mathcal{X}$, a
symmetric function $K:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ is called a
positive semi-definite kernel (p.s.d.) kernel on $\mathcal{X}$ if for all
$n\in\mathbb{N}$, all ${\bm{x}}_{1},\ldots,{\bm{x}}_{n}\in\mathcal{X}$, and
all $\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}$, it holds that
$\sum_{i=1}^{n}\sum_{j=1}^{n}\lambda_{i}\lambda_{j}K({\bm{x}}_{i},{\bm{x}}_{j})\geq
0$. Given two non-negative functions $f,g$, we say that $f(n)=O(g(n))$ (resp.
$f(n)=\Omega(g(n))$) iff. there exists absolute constant $C>0$, such that for
all $n\geq 0$, $f(n)\leq Cg(n)$ (resp. $f(n)\geq Cg(n)$). We use
$f(n)=\omega(g(n))$ to denote that for all $C>0$, there exists $n_{0}>0$ such
that for all $n\geq n_{0}$, $f(n)>Cg(n)$.

##### Setup.

We focus on learning modular addition, (1.1), with a two-layer network with no
biases and quadratic activation, following [15]. More specifically, given
parameters $\theta=(W,V)$, the model $f$ maps the pair of integers $(a,b)$ –
represented as the vector $(e_{a},e_{b})\in\mathbb{R}^{2p}$ – to a vector in
$\mathbb{R}^{p}$. We use the form

$f(\theta;(e_{a},e_{b}))=V\left(W(e_{a},e_{b})\right)^{\odot 2},$ (2.1)

where $W\in\mathbb{R}^{h\times 2p}$, $V\in\mathbb{R}^{p\times h}$ for some
integer $h$. We call $h$ the width of the hidden layer and it will be set
later.

In the classification setting (Section 4), we train $f$ with cross-entropy
loss to identify the $c$ such that $a+b\equiv c\pmod{p}$ in a multi-
classification setting. Letting $\mathcal{Z}=\\{e_{i}:i\in[p]\\}$, the set of
all possible inputs is $\mathcal{X}=\mathcal{Z}\times\mathcal{Z}$ and outputs
is $\mathcal{Y}=\mathcal{Z}$; there are $N=p^{2}$ distinct data points.

In the regression setting (Section 3), we instead aim to learn for each tuple
$(a,b,c)$, if $a+b\equiv c\pmod{p}$, that is, map
${\bm{x}}=(e_{a},e_{b},e_{c})$ to the scalar $y=p\,\mathbf{1}(a+b\equiv
c\pmod{p})$ using $g(\theta;(e_{a},e_{b},e_{c})))\triangleq
e_{c}^{\top}f(\theta;(e_{a},e_{b}))$.222This scaling implies that the
predictor $\Psi_{\bm{0}}(\cdot)=0$ has population square loss $p$; this
scaling allows bounded $\left\lVert\theta\right\rVert_{\infty}$. Here
$\mathcal{X}=\mathcal{Z}^{3}$, $\mathcal{Y}=\\{0,p\\}$, and $N=p^{3}$; we
train the model $g(\theta;{\bm{x}})=\langle
e_{c},f(\theta;(e_{a},e_{b}))\rangle$ to minimize the square loss.

In either setting, we use $\mathcal{D}$ to denote the distribution over
$\mathcal{X}\times\mathcal{Y}$ which is uniform over all $N$ possible input-
output pairs, while
$\mathcal{D}_{\mathrm{train}}=\left\\{({\bm{x}}_{i},y_{i})\right\\}_{i=1}^{n}$
is the training dataset of size $n$. Given parameters $\theta$, we use
$\Psi(\theta;\cdot):\mathcal{X}\to\hat{\mathcal{Y}}$ to denote the predictor
($g(\theta;\cdot)$ for regression, $f(\theta;\cdot)$ for classification), we
train the model with gradient descent with tiny $\ell_{\infty}$
regularization,333More details about the regularization strength can be found
in Appendix A. which is intended to emulate the implicit bias of Adam (AdamW)
which might lead to grokking. In general, we define loss as

$\\!\\!\\!\mathcal{L}^{\lambda}_{\ell}(\Psi,\theta,\mathcal{D}_{\operatorname{train}})\triangleq\frac{1}{n}\\!\\!\sum_{({\bm{x}},y)\in\mathcal{D}_{\operatorname{train}}}\\!\\!\\!\\!\\!\ell(\Psi(\theta;{\bm{x}}),y)+\lambda\left\lVert\theta\right\rVert_{\infty},$
(2.2)

where $\ell:\hat{\mathcal{Y}}\times\mathcal{Y}\to\mathbb{R}_{\geq 0}$ denotes
either square ($\ell_{2}$) or cross-entropy loss and $\lambda\geq 0$ controls
the strength of regularization. Omitting $\lambda$ refers to the case where no
regularization is used and omitting $\mathcal{D}_{\mathrm{train}}$ means the
whole population is used in evaluating the loss. For non-parametric functions
(like $\Psi_{\bm{0}}$ denoted as the predictor which returns zero on any
input) we drop $\theta$ as well.

##### Connection between $\ell_{\infty}$-norm Regularization and Grokking.

[41] shows that grokking phenomenon happens for Adam and gets more prominent
with decoupled weight decay (AdamW). Recent works by [52, 55] show that Adam
implicitly regularizes the $\ell_{\infty}$ norm of the weights. In detail,
[52] shows that AdamW can only converge to KKT points of $\ell_{\infty}$-norm
constrained optimization. [55] shows that Adam converges to max-margin
solutions w.r.t. $\ell_{\infty}$ norm for linear models on separable datasets
and conjecture this result could be generalized to general homogeneous models
(including our model $f$, which is 3-homogeneous). This connection motivates
us to use the explicit $\ell_{\infty}$-norm regularization on top of gradient
descent to understand grokking in the modular addition setting.

######  Definition 2.1.

A deterministic supervised _learning algorithm_ $\mathcal{A}$ is a mapping
from a sequence of training data,
$\mathcal{D}_{\mathrm{train}}\in(\mathcal{X}\times\mathcal{Y})^{n}$, to a
hypothesis
$\mathcal{A}(\mathcal{D}_{\mathrm{train}}):\mathcal{X}\to\hat{\mathcal{Y}}$.
The algorithm $\mathcal{A}$ could also be randomized, in which case the output
$\mathcal{A}(\mathcal{D}_{\mathrm{train}})$ is a distribution over hypotheses.
Two randomized algorithms $\mathcal{A}$ and $\mathcal{A}^{\prime}$ are the
same if for any input, their outputs have the same distribution in function
space, written as
$\mathcal{A}(\mathcal{D}_{\mathrm{train}})\overset{d}{=}\mathcal{A}^{\prime}(\mathcal{D}_{\mathrm{train}})$.

We further define equivariance of learning algorithms which is the main
component of our analysis in deriving lower bounds. This definition is used in
Theorems 3.3 and 4.2 to prove equivariance of GD in learning modular addition.

######  Definition 2.2 (Equivariant Algorithms).

A learning algorithm is _equivariant_ under group $\mathcal{G}_{\mathcal{X}}$
(or $\mathcal{G}_{\mathcal{X}}$-equivariant) if for any dataset
$\mathcal{D}_{\mathrm{train}}\in(\mathcal{X}\times\mathcal{Y})^{n}$ and for
all $T\in\mathcal{G}_{\mathcal{X}}$, it holds that
$\mathcal{A}(\left\\{(T({\bm{x}}_{i}),y_{i})\right\\}_{i=1}^{n})\,\circ\,T\overset{d}{=}\mathcal{A}(\left\\{({\bm{x}}_{i},y_{i})\right\\}_{i=1}^{n})$.
For deterministic learning algorithms, this is equivalent to saying that for
all ${\bm{x}}\in\mathcal{X}$,
$\mathcal{A}(\left\\{(T({\bm{x}}_{i}),y_{i})\right\\}_{i=1}^{n})(T({\bm{x}}))=[\mathcal{A}(\left\\{({\bm{x}}_{i},y_{i})\right\\}_{i=1}^{n})]({\bm{x}})$.

######  Definition 2.3 (Kernel Methods).

For $\hat{\mathcal{Y}}\subseteq\mathbb{R}$, we say a learning algorithm
$\mathcal{A}$ is a _kernel method_ if it first picks a (potentially random)
positive semi-definite kernel $K$ on $\mathcal{X}$ before seeing the
data,444Our definition of kernel methods does not cover learning algorithms
that choose the kernel based on the training data. _Any_ learning algorithm
could be framed as a kernel method with a data-dependent kernel
$K(x,x^{\prime})=\Psi(x)\Psi(x^{\prime})$. and then outputs some hypothesis
$\Psi$ such that there exist $\\{\lambda_{i}\\}_{i=1}^{n}\in\mathbb{R}$ for
which $\Psi(\cdot)=\sum_{i=1}^{n}K(\cdot,{\bm{x}}_{i})\lambda_{i}$, where
$\lambda_{i}$ can depend on the training dataset
$\mathcal{D}_{\mathrm{train}}$.

In particular, when
$\bm{\lambda}(\mathcal{D}_{\mathrm{train}})=\mathbf{K}^{\dagger}\mathbf{y}$,
where $\mathbf{K}=\left(K({\bm{x}}_{i},{\bm{x}}_{j})\right)_{i,j=1}^{n}$,
$\mathbf{y}=\left(y_{i}\right)_{i=1}^{n},$ the kernel method is called a
_(ridgeless) kernel regression_ method.

######  Theorem 2.4.

For any p.s.d. kernel $K:{\mathcal{X}}\times{\mathcal{X}}\to\mathbb{R}$ and
transformation group ${\mathcal{G}}_{g}X$, kernel regression (Definition 2.3)
with respect to kernel $K$ is ${\mathcal{G}}_{{\mathcal{X}}}$-equivariant if
and only if kernel $K$ is equivariant to ${\mathcal{G}}_{g}X$, _i.e._ ,
$K(T(x),T(x^{\prime}))=K(x,x^{\prime})$ for any $T\in{\mathcal{G}}_{g}X$ and
$x,x^{\prime}\in{\mathcal{X}}$.

###### Proof of Theorem 2.4.

For any $\mathcal{D}_{\mathrm{train}}=\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}$
and transformation $T\in{\mathcal{G}}_{\mathcal{X}}$, let
$\mathcal{D}_{\mathrm{train}}^{T}$ be the transformed dataset
$\\{(T({\bm{x}}_{i}),y_{i})\\}_{i=1}^{n}$, and ${\mathcal{A}}_{K}$ be the
kernel regression algorithm w.r.t. $K$. For any ${\bm{x}}$, we have that

$\displaystyle{\mathcal{A}}(T(\mathcal{D}_{\mathrm{train}}))(T({\bm{x}}))=\sum_{i=1}^{n}K(T({\bm{x}}),T({\bm{x}}_{i}))\lambda_{i}(\mathcal{D}_{\mathrm{train}}^{T})=\sum_{i=1}^{n}K({\bm{x}},{\bm{x}}_{i})\lambda_{i}(\mathcal{D}_{\mathrm{train}})={\mathcal{A}}(\mathcal{D}_{\mathrm{train}})({\bm{x}}),$

where the second equality follows from the equivariance of the kernel $K$
w.r.t. ${\mathcal{G}}_{\mathcal{X}}$. ∎

##  3 Regression Task

We will first present our theoretical analysis of the grokking phenomenon on
modular arithmetic with two-layer quadratic networks in the regression
setting, where we learn a function from $(e_{a},e_{b},e_{c})$ to
$p\mathbf{1}(a+b=c\pmod{p})$. Although this is perhaps a less natural way to
model modular arithmetic than the classification task, it admits some useful
theoretical tools.

In this section, we show that networks in the kernel regime will provably fail
to generalize as long as they do not have access to nearly all the points in
the dataset (Theorem 3.4), although they can achieve zero training error
(Theorem 3.1). However, the network eventually leave the kernel regime because
of the weak $\ell_{\infty}$ norm regularization. We further establish that if
the network manages to achieve zero training error with small
$\left\lVert\theta\right\rVert_{\infty}$, it will generalize with only a
$\frac{n}{N}=\tilde{\omega}(1/p)$ portion of the overall dataset, as long as
the width of the network is larger than $4p$ (Theorem 3.6). We also prove that
such networks exist (Proposition 3.8), and demonstrate that gradient descent
can find them with a small amount of explicit regularization (Figure 3).

###  3.1 Kernel Regime

Figure 2: Empirical evidence for kernel regime in early training. Left: Train
(dashed) and test (solid) accuracy while training in the regression setting
with various initialization scales $\alpha$ and a fixed $p=47$. Shrinking the
scale of initialization can mitigate grokking in the regression task,
eventually eliminating of the gap between train and test accuracies (at the
cost of slower improvement in each). Right: eNTK continues to significantly
change after overfitting.

A recent line of work on the _neural tangent kernel_ (NTK) framework [20, 4,
22, 10, 54] has shown that with typical initialization schemes, gradient
descent in over-parameterized neural networks locally approximates the
behavior of a kernel model using the empirical neural tangent kernel (eNTK)
$K_{\theta}({\bm{x}},{\bm{x}}^{\prime})\triangleq\nabla
g(\theta;{\bm{x}})\nabla g(\theta;{\bm{x}}^{\prime})^{\mathsf{T}}$. In the
“kernel regime,” the change in $\theta$ over the course of gradient descent
does not substantially change the eNTK $K_{\theta}$, and hence the neural
network behaves similarly to a kernel predictor trained with $K_{\theta_{0}}$.
With square loss, as here, these kernel predictors follow a particularly
simple optimization path for which a closed form (corresponding to kernel
regression) is available.

For networks of finite width (and in certain infinite-width cases), the eNTK
will stay roughly constant and the network will closely track the kernel model
for the first phase of optimization, but the tiny regularization will
eventually lead it to depart the kernel regime. Thus, bounds on kernel models
are informative about deep networks in the first part of optimization.

In the first phase of grokking, the model overfits to the training data,
achieving very low training loss but retaining high loss on test points.We
establish that both phenomena occur in the kernel regime.

####  3.1.1 Empirical Neural Tangent Kernels Can Achieve Zero Training Error

Our first result shows that kernel regression with the empirical neural
tangent kernel of our quadratic network achieves zero training loss if the
network is mildly wide, for instance, $n=\tilde{\Theta}(p^{2})$ and width
$h=\tilde{\Theta}(p)$. This implies, for example, that in a “lazy training”
setting [10], or alternatively when the first layer is initialized with huge
weights and the second layer with tiny ones, gradient descent can achieve zero
training loss. Informally, this also strongly suggests that networks with more
typical initializations can achieve very small training error without needing
to leave the “kernel regime.”

######  Theorem 3.1.

Initialize the network of Section 2 with any values for $V$, and entries of
$W$ all independently $\operatorname{Uniform}([-s_{W},s_{W}])$ for some
$s_{W}>0$. Let the most frequent value of $c\in[p]$ appear $\rho_{c}$ times.
Then, if $h>36\rho_{c}\log(n/\delta)$, it holds with probability at least
$1-\delta$ over the random initialization of $W$ that kernel regression using
the network’s empirical neural tangent kernel can achieve zero training loss
for _any_ target labels such that if ${\bm{x}}_{i}={\bm{x}}_{j}$ then
$y_{i}=y_{j}$. Conversely, if $h<n/(3p)$, there exist target labels (with
$y_{i}=y_{j}$ when ${\bm{x}}_{i}={\bm{x}}_{j}$) for which this method cannot
achieve zero training loss.

Uniform initialization is the default behavior e.g. in PyTorch [2]. Note that
$\rho_{c}$ is always between $n/p$ and $n$. Moreover, in the usual setting
with randomly selected $\mathcal{D}_{\mathrm{train}}$ and
$n=\Omega\bigl{(}p\log p\bigr{)}$ – recall that full knowledge of a single
$(a,b)$ pair requires $n=p$ in the regression setting – we have that
$\rho_{c}=\Theta(n/p)$ with high probability [42, Theorem 1]. Thus in this
usual setting, the threshold for interpolation is at a network width
$h=\tilde{\Theta}(n/p)$.

###### Proof sketch of Theorem 3.1.

Kernel regression can achieve all possible target labels iff its kernel matrix
is strictly positive definite (Lemmas B.1 and B.2). Evaluating the expected
neural tangent kernel, we can see it is strictly positive definite under only
mild assumptions on the weights (Proposition B.3). With bounded weights, a
matrix Chernoff bound controls the convergence of the lowest eigenvalue of the
kernel matrix to that of the expected kernel matrix (Proposition B.4). The
lower bound follows from the rank of the kernel matrix (Proposition B.6). ∎

####  3.1.2 Permutation-Equivariant Kernel Methods Cannot Generalize

We now show that for any permutation-equivariant kernel method (Definition
2.3) (and hence for networks trained by gradient descent close enough to
initialization), generalization is possible only when training on
$n=\Omega(p^{3})$ samples, i.e. the portion of all possible data points is
$\frac{n}{N}=\Omega(1)$.

The key component of our analysis is the permutation equivariance of learning
modular addition in this setting. We first define the permutation group on the
modular addition data:

######  Definition 3.2.

Let ${\mathbb{S}}_{p}$ denote the set of all permutations on $[p]$. We define
the _permutation group_ $\mathcal{G}_{\mathcal{X}}$ on $\mathcal{X}$ as the
group

$\left\\{(e_{a},e_{b},e_{c})\mapsto(e_{\sigma_{1}(a)},e_{\sigma_{2}(b)},e_{\sigma_{3}(c)}):\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}\right\\},$

with the group operation being composition of each permutation.

The following theorem establishes that our learning algorithm is equivariant
(Definition 2.2) under this permutation group, which we call _permutation-
equivariant_ for short. Note that this result applies to the actual process of
gradient descent on our neural network, not only to its NTK approximation.
(This result is not particularly specific to the architecture defined in
Section 2; it holds broadly.)

######  Theorem 3.3 (Permuation Equivariance in Regression).

Training a two-layer neural network with quadratic activations (such as
$g(\theta;x))$ initialized with random i.i.d. parameters using gradient-based
update rules (such as gradient descent) on the modular addition problem in the
regression setting is an equivariant learning algorithm (as defined in
Definition 2.2) with respect to the permutation group of Definition 3.2
applied on the input data.

More details and the full proof are in Appendix C.

Roughly speaking, this indicates that learning the modular addition task on
this setup is exactly the same difficulty as learning any permuted version of
the dataset. Following the same argument, we can show that the kernel method
corresponding to the neural nets in the early phase of training is also
permutation-equivariant.

Further, as the distribution $\mathcal{D}$ is uniform and thus invariant under
permutation, Theorem 3.3 shows that the difficulty of learning the original
ground-truth is the same as simultaneously learning all the permuted versions
of the ground-truths, which turns out to be difficult for any kernel method.
The following theorem formalizes this idea, establishing that any such kernel
predictor needs at least $n=\Omega(p^{3})$, or equivalently
$\frac{n}{N}=\Omega(1)$, to generalize on the modular addition task; otherwise
it cannot do substantially better than the trivial all-zero predictor.

######  Theorem 3.4 (Lower Bound).

There exists a constant $C>0$ such that for any $p\geq 2$, training data size
$n<Cp^{3}$, and any permutation-equivariant kernel method $\mathcal{A}$, it
holds that

$\\!\\!\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\mathcal{L}_{\ell_{2}}\left(\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)\right)\geq\frac{w}{2}\;\mathcal{L}_{\ell_{2}}\left(\Psi_{\bm{0}}\right)=\frac{p}{2},$

where $\operatorname*{\mathbb{E}}_{\mathcal{A}}$ takes the mean over the
randomness in algorithm $\mathcal{A}$.

A full proof of this result is deferred to Appendix D.

Theorem 3.4 is in fact a special case of a more general theorem that lower
bounds the population $\ell_{2}$ loss of learning modular addition with
$m$-summands using permutation-equivariant kernels, showing that poor
generalization is inevitable. We present an informal version of this result
below and refer the reader to Corollary D.8 for the formal version.

######  Theorem 3.5 (Informal).

Consider the problem of learning modular addition over $m$-summands with a
training set of size $n$ using a permutation equivariant kernel method
$\mathcal{A}$. For any $p\geq 2$ and training data size $n<Cp^{m}$ the
expected population $\ell_{2}$ loss is at least $\Omega(p)$, which is of the
same magnitude as the trivial all-zero predictor.

Poor population loss is thus inevitable for models which are well-approximated
by a permutation-equivariant kernel; since our quadratic network can also
achieve zero training error in the kernel regime, this strongly suggests
drastic overfitting in early training. Fortunately, however, regularized
gradient descent on finite networks will eventually leave the kernel regime
and begin learning features.

###  3.2 Rich Regime

Figure 3: Empirical verification of our theoretical explanation for
generalization. We train a network of width $h=4p$ with gradient descent on
$\ell_{2}$ loss and $10^{-4}$ $\ell_{\infty}$-regularization on $2\times
p^{2.25}$ training samples (out of $p^{3}$) in the regression setting. Left:
Generalization happens when the number of samples are more than
$\Omega(p^{2})$, as predicted by Proposition 3.7. The dashed and solid lines
indicate train and test set statistics respectively. Right: $\ell_{\infty}$
norm of parameters after grokking remains the same for different problem
dimensions $p$, as predicted by

While the transient behavior of gradient descent after leaving the kernel
regime may be complicated, it is often the case that the limiting behavior can
be better-understood based on analyses of implicit bias. Motivated by this, we
present a generalization analysis of networks that achieve zero training error
(as we expect in the long-term optimization limit) with bounded
$\left\lVert\theta\right\rVert_{\infty}$. This assumption is motivated by
recent work of [52], who show that full-batch AdamW can only converge to KKT
points of the $\ell_{\infty}$ norm-constrained optimization problem,
$\min_{\left\lVert\theta\right\rVert_{\infty}\leq\frac{1}{\lambda}}\mathcal{L}(\Psi,\theta,\mathcal{D}_{\mathrm{train}})$,
where $\lambda$ is the weight decay coefficient. We will discuss the
feasibility of this assumption more afterwards.

######  Theorem 3.6 (Upper Bound).

For any width $h\geq 8p$, with probability at least $1-\delta$ over the
randomness of training dataset $\mathcal{D}_{\operatorname{train}}$ of size
$n$, define the set of interpolating solutions as
$\tilde{\Theta}^{*}=\\{\theta\mid\mathcal{L}_{\ell_{2}}(\Psi,\theta,\mathcal{D}_{\operatorname{train}})=0\\}$.
For any interpolating solution $\theta^{*}\in\tilde{\Theta}^{*}$ with small
$\ell_{\infty}$ norm, _i.e._ , satisfying that
$\|\theta^{*}\|_{\infty}=\operatorname{\mathcal{O}}(\min_{\theta\in\tilde{\Theta}^{*}}\left\lVert\theta\right\rVert_{\infty})$,
it holds that

$\displaystyle\mathcal{L}_{\ell_{2}}(g,\theta^{*})$
$\displaystyle=\operatorname{\mathcal{O}}\left(\frac{p^{2}}{n}\left(\log^{3}n+\frac{1}{p}\log\frac{1}{\delta}\right)\right)\cdot\mathcal{L}_{\ell_{2}}\left(\Psi_{\bm{0}}\right).$

Comparing Theorem 3.6 to Theorem 3.4, we see that when
$n=\tilde{\omega}(p^{2})$ and $h=\Omega(p)$, there is a strict separation
between generalization in the kernel and rich regimes. Note that the
classifier
$\phi(e_{a},e_{b})\triangleq\operatorname*{arg\,max}_{c}f_{c}(\theta^{*};(e_{a},e_{b}))$
has population error rate at most $\frac{2}{p}\mathcal{L}_{\ell_{2}}(\phi)$
(as shown in Proposition E.7) and thus, its population error goes to zero when
$n=\tilde{\omega}(p^{2})$. The proof of Theorem 3.6 consists of showing all
networks with small training error and small $\ell_{\infty}$ norms generalize
(Proposition 3.7), and at least one such network exists (Proposition 3.8).

######  Proposition 3.7.

For any $R>0,\delta\in(0,1)$, $\mathcal{D}_{\operatorname{train}}$ of size
$n$, and
$\theta^{*}\in\\{\theta=(W,V):\mathcal{L}_{\ell_{2}}(g,\theta,\mathcal{D}_{\operatorname{train}})=0\wedge\left\lVert\theta\right\rVert_{\infty}\leq
R\\}$, there exists a positive constant $C>0$ such that with probability at
least $1-\delta$ over the randomness of $\mathcal{D}_{\operatorname{train}}$,

$\mathcal{L}_{\ell_{2}}(g,\theta^{*})\leq\frac{CR^{6}h^{2}}{n}\left(p\log^{3}n+\log\frac{1}{\delta}\right).$

###### Proof sketch of Proposition 3.7.

We bound the Rademacher complexity of the set of networks with small
$\ell_{\infty}$ weights, and then apply Theorem E.6, due to [45], which gives
an “optimistic” bound on the excess risk of smooth loss functions. Details in
Appendix E. ∎

######  Proposition 3.8.

Let the set of models with zero population loss be
$\Theta^{*}\triangleq\\{\theta\mid\mathcal{L}_{\ell_{2}}(g,\theta)=0\\}$. For
any $p\geq 2$ and $h\geq 8p$, $\Theta^{*}$ is nonempty and
$\min_{\theta\in\Theta^{*}}\left\lVert\theta\right\rVert_{\infty}\leq\left\lfloor\frac{h}{8p}\right\rfloor^{-\frac{1}{3}}$.

###### Proof sketch of Proposition 3.8.

The main difficulty is a manual Fourier-based construction of a zero-loss
solution with $h=8p$ and $\ell_{\infty}$ norm at most one; this is inspired by
similar constructions of [15] and results of [37]. In a concurrent work, [35]
presented a similar construction. Once we have that, because our model is
3-homogeneous in its parameters, we can easily reduce the $\ell_{\infty}$ norm
by duplicating neurons, without changing the input-output function. Details in
Appendix F. ∎

We know that an interpolating solution with small $\ell_{\infty}$ norm exist,
and that any such solution will generalize. Does gradient descent find such a
solution? In Figure 3, we empirically evaluate the $\ell_{\infty}$ norm of
weights learned by running gradient descent on the regression task, with
varying $p$, and a very small $\ell_{\infty}$ regularizer.555We used a
regularization weight of $10^{-4}$. Without explicit regularization, a small
number of network weights do grow with $p$, but we believe this phenomenon is
not important to the overall behavior of gradient descent on this task.
Consistent with our manual construction of weights, the $\ell_{\infty}$ norm
of the network weights does not grow with the problem dimension $p$. This
supports the applicability of Theorem 3.6 and a far better sample complexity
compared to the kernel regime.

color=blue!30,noinline]Zhiyuan: check later Empirical Evidence that Kernel
Lower Bounds Lead to Overfitting: One way to mitigate the grokking effect in
learning modular addition is by preventing the network from overfitting the
training set in the kernel regime. To do so, one can reduce the scale of
initialization (e.g. smaller variance scales for weight initializations in the
scheme of [19]). Roughly speaking, networks initialized with a very small
scale of parameters need to undergo a significant growth in the norm of
parameters to be able to fit the training data. This norm growth will
eventually lead the network to leave the kernel regime and begin learning
features before overiftting the data. Figure 2 confirms that in our setting:
using a very small weight initialization can substantially mitigate the
grokking effect by preventing the network from overfitting in the kernel
regime, incidacting that grokking is indeed caused by a delayed transition
from kernel (overiftting) to rich (generalizing) regime. However, this comes
at a cost: training becomes increasingly more difficult as the scale of
initialization decreases [10, 34, 46].

In Figure 2, we also evaluate the change of the empirical NTK during training,
by computing the difference in eNTK matrices through training. To make this
empirical investigation computationally feasible, we evaluated the eNTK
approximation of [32] on 20,000 random data points, similar to previous
schemes [13, 51, 32]. We see that the change in empirical NTK is orders of
magnitude larger after overfitting the training set, implying that most
feature learning happens only later, supporting our hypothesis that the
initial overfitting occurs roughly in the kernel regime.

##  4 Classification Task

We now move onto the multi-class classification setting established in Section
2, where we train with cross-entropy loss. Similar to the regression problem,
we first analyze the early stage of training where the network operates like a
kernel machine, and then move onto the rich regime and focus on the weights
learned through margin-maximization implicit bias of gradient descent. We
prove that a sample complexity gap between kernel (Theorem 4.3) and rich
(Theorem 4.6) regime exists when learning a two-layer neural network with
gradient descent on modular addition modeled as a multi-class classification
task. Our results imply that as long as the max-margin implicit bias of
gradient descent on exponential-type loss functions drives the net to leave
the kernel regime and $\hat{\Omega}(p^{5/3})$ samples are used for training,
generalization to the whole population is guaranteed.

###  4.1 Kernel Regime

In the multi-class setting, the output is $p$-dimensional, and thus the eNTK
for each pair of points is a $p\times p$ psd matrix [56, see]. Our notion of
kernel methods (Definition 4.1) slightly changes due to account for multi-
output functions, and the main lower bound result is similar to that in the
regression case: kernel methods must see a constant fraction of data before
learning.

######  Definition 4.1.

For $\mathcal{Y}\subseteq\mathbb{R}^{p}$, we say a learning algorithm
$\mathcal{A}$ is a _kernel method_ if before seeing the data it first picks a
(potentially random) kernel $K$ on $\mathcal{X}$ such that for any
$x,x^{\prime}\in\mathcal{X}$, $K(x,x^{\prime})$ is a $p$ by $p$ positive semi-
definite matrix, and then outputs some hypothesis $h$ based on
$\mathcal{D}_{\mathrm{train}}$ such that there exist
$\\{\lambda_{i}\\}_{i=1}^{n}\in\mathbb{R}^{p}$ for which
$h(\cdot)=\sum_{i=1}^{n}K(\cdot,{\bm{x}}_{i})\lambda_{i}$.

We next establish the permutation-equivariance of gradient-based learning
algorithms in the classification setting. They key difference is that in the
multi-output setting, gradient-based learning algorithms are equivariant to
permutations on _both input data and output labels_. More details are
available in Appendix C.

######  Proposition 4.2 (Permuation Equivariance in Classification).

We define the _permutation group_ $\mathcal{G}_{\mathcal{X},\mathcal{Y}}$ on
$\mathcal{X}\times\mathcal{Y}$ (defined in Section 2) as the group

$\left\\{(e_{a},e_{b}),e_{c}\mapsto(e_{\sigma_{1}(a)},e_{\sigma_{2}(b)}),e_{\sigma_{3}(c)}:\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}\right\\}.$

Training a two-layer neural network with quadratic activations (such as
$f(\theta;x))$ initialized with random i.i.d. parameters using gradient-based
update rules (as defined in Definition C.1)666GD is an example of such
algorithm. on the modular addition problem in the classification setting is an
equivariant learning algorithm (as defined in Definition 2.2) with respect to
$\mathcal{G}_{\mathcal{X},\mathcal{Y}}$.777A similar result holds for adaptive
algorithms like Adam as well, as discussed in Corollary C.7.

We say a learning algorithm is _input-output permutation-equivariant_ it is
equivariant with respect to $\mathcal{G}_{\mathcal{X},\mathcal{Y}}$. The
following theorem establishes that any such kernel learning algorithm needs at
least $n=\Omega(p^{2})$, or equivalently $\frac{n}{N}=\Omega(1)$, to
generalize on the modular addition task in the classification setting.

######  Theorem 4.3 (Kernel Lower Bound).

There exists a constant $C>0$ such that for any training data size $n<Cp^{2}$
and any kernel method $\mathcal{A}$ which is input-output permutation-
equivariant (with respect to $\mathcal{G}_{\mathcal{X},\mathcal{Y}}$), it
holds that

$\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\mathcal{L}_{\ell_{2}}\left(\mathcal{A}\left(\\{{\bm{x}}_{i},y_{i}\\}_{i=1}^{n}\right)\right)\geq\frac{1}{2}\;\mathcal{L}_{\ell_{2}}\left(\Psi_{\bm{0}}\right),$

where $\operatorname*{\mathbb{E}}_{\mathcal{A}}$ takes expectation over the
randomness in algorithm $\mathcal{A}$.

Figure 4:  Empirical investigation of grokking in the classification setting
with tiny $\ell_{\infty}$ regularization on different problem dimensions $p$.
Networks are trained with normalized gradient descent888We refer the reader to
Appendix A for the definition of normalized GD. on cross-entropy loss and an
$\ell_{\infty}$ regularization of $10^{-20}$ strength. The dashed lines in the
indicate train set statistics, and the solid lines correspond to the test set.

Intuitively, the classification setting is not very different from the
regression setting. They share the same goal and the amount of knowledge
contained in each data in the classification setting is exactly equal to $p$
data in the regression setting, where these $p$ data share the same first two
coordinates and only differ in the last coordinate.

More formally, we can define the following correspondence: given one data
point in classification $({\bm{x}},y)\in\mathbb{R}^{2p}\times\mathbb{R}^{p}$,
we can view it as $p$ data points in regression,
$\\{(({\bm{x}},e_{i}),y_{i})\\}_{i=1}^{p}$, denoted by $F({\bm{x}},y)$.
Similarly, any function $\Psi$ mapping from ${\bm{x}}=\\{(e_{i},e_{j})\\}$ to
$\mathbb{R}^{p}$ can be viewed as a function mapping
$\\{(e_{i},e_{j},e_{k})\\}$ to $\mathbb{R}$ by defining
$\Psi^{\prime}({\bm{x}},e_{k})=[\Psi({\bm{x}})]_{k}$. Moreover, these two
functions $\Psi$ and $\Psi^{\prime}$ share the same population $\ell_{2}$
loss. Under this view, matrix-valued kernel learning with $n$ classification
data points is exactly the same as scalar-valued kernel learning with $np$
regression data points.

The only obstacle to directly applying Theorem 3.4 is that the data
distribution is different. For the regression setting, each data is sampled
independently and uniformly, and the number of data points can be any integer;
in the classification setting, the data points are sampled in independent
groups and the number of data must be a multiple of $p$. However, it is easy
to see this new distribution is still invariant under permutations as in
Definition 3.2. Thus, we can directly apply Theorem 3.4 on this new
distribution to get Theorem 4.3. A formal version of this argument is
available in Section D.3.

Therefore, for networks operating in the kernel regime, generalization to
unseen data in $\ell_{2}$ loss is impossible unless $n/N=\Omega(1)$, i.e. we
have observed a constant fraction of all the $N=p^{2}$ possible data points.
It is worth emphasizing, however, that a large $\ell_{2}$ loss does not
necessarily imply a large classification error. It remains open to prove a
classification error lower bound for permutation-equivariant kernel methods.

###  4.2 Rich Regime

To analyze the behaviour of the network in the rich regime, we first introduce
the notion of margin and elaborate on the margin-maximization implicit bias of
gradient descent. For a multi-class classification problem with $p$ classes
and a fixed network $f$, the margin of a data point $({\bm{x}},y)$ is

$q(\theta;{\bm{x}},y)\triangleq f_{y}(\theta;{\bm{x}})-\max_{y^{\prime}\neq
y}f_{y^{\prime}}(\theta;{\bm{x}}).$

The margin for a dataset $\mathcal{D}$ is defined as the minimum margin of all
points on the dataset:

$q_{\min}(\theta;\mathcal{D})\triangleq\min_{({\bm{x}},y)\in\mathcal{D}}q(\theta;{\bm{x}},y).$

When the network $f$ is homogeneous with respect to its parameters (as is the
case in our setup), one can observe that as long as $\theta$ linearly
separates the dataset $\mathcal{D}$ , _i.e._ ,
$q_{\min}(\theta;\mathcal{D})>0$, it is possible to arbitrary scale the
minimum margin, through scaling the parameters of the network. Hence, in such
homogenous networks, one is usually concerned with a normalized margin
$q_{\min}(\theta/\left\lVert\theta\right\rVert;\mathcal{D})$ according to some
norm.

[29] proved that gradient descent on homogeneous models with the cross-entropy
(or similar) losses on dataset $\mathcal{D}$, in the absence of explicit
regularization, maximizes the normalized margin. Specifically, although
$\lVert\theta\rVert_{2}\to\infty$ as $t\to\infty$,
$\theta/\lVert\theta\rVert_{2}$ converges to a solution (or more generally, a
KKT point) of the following problem when one exists:

$\min\;\frac{1}{2}\left\lVert\theta\right\rVert_{2}^{2}\quad\text{s.t.}\quad
q_{\min}(\theta;\mathcal{D})\geq 1.$ (4.1)

To establish our results in the classification setting, we borrow the
following theorem from [51], who prove that when the strength of the
regularization used in (2.2) is small enough, the maximum normalized margin of
the regularized loss converges to that of the unregularized loss.

######  Proposition 4.4 ([51], Theorem 4.1).

Consider a positively homogeneous function $f$ with respect to parameters
$\theta$ and a dataset $\mathcal{D}_{\mathrm{train}}$ separable with $f$. Let
$\lVert\cdot\rVert$ be any norm. Let $\gamma^{*}$ be the maximum normalized
margin of the unregularized loss and $\gamma^{\lambda}$ be normalized margin
of the minimizer of the regularized loss with strength $\lambda$. As
$\lambda\to 0$, $\gamma^{\lambda}\to\gamma^{*}$.

That is, if we use a small enough regularization weight $\lambda$ with any
norm in (2.2), we will obtain approximately the same solution as the
unregularized problem. This confirms that training our two-layer network using
gradient descent with a small enough $\ell_{\infty}$ regularizer can lead to
weights close to the solution of the max $\ell_{\infty}$-normalized margin
problem. Next we present the main result of this subsection, a upper bound on
test error, showing that all parameters whose $\ell_{\infty}$-normalized
margin is close enough to the max $\ell_{\infty}$-normalized margin solution
will generalize with a sample complexity of
$\tilde{\operatorname{\mathcal{O}}}(p^{5/3})$

######  Theorem 4.5 (Upper Bound).

Let $\delta\in(0,1)$ be positive constants such that $8p\leq h=O(p)$
independent of $p$ and $n$. For any $n$ representing the size of the training
set $\mathcal{D}_{\mathrm{train}}$ it holds with probability at least
$1-\delta$ over randomness of $\mathcal{D}_{\mathrm{train}}$, for all
interpolating solutions $\theta$ with normalized $\ell_{\infty}$-margin
$\Omega(p)$, it holds that

$\displaystyle
L_{0}(f,\theta,\mathcal{D})\leq\tilde{\operatorname{\mathcal{O}}}\left(\sqrt{\frac{p^{5/3}}{n}}\right).$
(4.2)

In other words, any interpolating solution $\theta$ with approximately maximal
normalized $\ell_{\infty}$-margin, _i.e._ ,
$q_{\min}(\theta/\left\lVert\theta\right\rVert_{\infty};\mathcal{D}_{\mathrm{train}})\geq\Omega(\max_{\|\theta^{\prime}\|_{\infty}\leq
1}q_{\min}(\theta^{\prime};\mathcal{D}_{\mathrm{train}}))$, has a sample
complexity of $\tilde{\operatorname{\mathcal{O}}}(p^{5/3})$, since Proposition
3.8 shows that the maximal normalized $\ell_{\infty}$-margin is at least
$\Omega(p)$.

Our proof of Theorem 4.5 is based on the PAC-Bayesian framework [31],
specifically using Lemma 1 of [38]. This result provides a margin-based high
probability generalization bound for any predictor based on the _margin loss_
, which counts a prediction as correct only if it predicts with a margin at
least $\gamma$:

$L_{\gamma}\left(f,\theta,\mathcal{D}\right)\triangleq\Pr_{({\bm{x}},y)\sim\mathcal{D}}\Big{[}q(\theta,{\bm{x}},y)>\gamma\Big{]}.\vspace*{-1mm}$
(4.3)

$L_{0}$ is simply the misclassification rate, or 0-1 error. The following
statement uses this result to prove an upper bound on the population 0-1 error
of networks using margin loss on the training dataset.

######  Theorem 4.6.

For any $p\geq 2$, training set size $n\geq 1,\delta\in(0,1)$, norm $r>0$,
width $h^{\prime}>4\log\frac{2}{\delta}$ and normalized margin
$\gamma/r^{3}=\tilde{\Omega}(\sqrt{h})$ it holds with probability at least
$1-\delta$ over the randomness of the training set
$\mathcal{D}_{\operatorname{train}}$ that for any $\theta^{\prime}$ of width
$h^{\prime}$ with $\left\lVert\theta^{\prime}\right\rVert_{\infty}\leq r$

$\displaystyle L_{0}(f,\theta^{\prime},\mathcal{D})\leq
L_{\gamma}(f,\theta^{\prime},\mathcal{D}_{\mathrm{train}})+\tilde{\operatorname{\mathcal{O}}}\left(\sqrt{\frac{p}{n}}\cdot\sqrt[3]{\frac{h^{2}}{\gamma/r^{3}}}\right)$
(4.4)

where $\mathcal{D}$ denotes the population and $f,L$ are defined in Equations
2.1 and 4.3 respectively.

###### Proof Sketch of Theorem 4.6.

Through a series of concentration inequalities, we show that as long as the
assumptions of Theorem 4.6 are met, the output logits under gaussian
perturbation on parameters,
$f(\theta^{\prime}+\tilde{\theta}^{\prime};\cdot)$, where
$\tilde{\theta}^{\prime}$ is entry-wise Gaussian noise of mean zero and
variance $\sigma^{2}$, are close to that of $f(\theta^{\prime};\cdot)$ up to
an absolute difference of
$\tilde{\operatorname{\mathcal{O}}}(\sqrt{h}\sigma^{3})$. The proof is
completed by applying a PAC-bayesian generalization bound, e.g., using Lemma 1
of [38] and setting $\sigma$ based on the value of $\ell_{\infty}$-normalized
margin $\gamma/r^{3}$ and width $h$. The full proof is available in Appendix
G. ∎

Now we show how to derive Theorem 4.5 using Theorem 4.6. Because Theorem 4.6
only relies on the normalized margin and our model is 3-homogeneous, without
loss of generality, we can fix norm bound $r=1$. Then it suffices to set
$8p<h=O(p)$ and $\gamma=\Omega(p)$ in Equation 4.4, where the first term
becomes $0$ because Theorem 4.5 assumes $\theta$ has normalized margin
$\Omega(p)$ and the second term becomes $\tilde{O}(\sqrt{\frac{p^{5/3}}{n}})$.

Theorem 4.5 confirms and explains previous observations [41, 15, 37, 27, e.g.]
on the minimum threshold for the fraction of data used to achieve
generalization. In combination with Theorem 4.2 of [29], this shows that with
enough training data, gradient descent will eventually find a generalizing
solution for this setting of the modular addition problem.

Figure 5: Grokking in transformers happens after a delayed transition from
kernel to rich regime. A one-layer transformer is trained with gradient
descent using cross-entropy loss and a tiny $\ell_{\infty}$ regularization of
$10^{-20}$ strength on $2\times p^{5/3}$ training samples from the modular
addition problem with various $p$s. Change of eNTK up to the point of fitting
the training set is negligible. The eNTK has a drastic change only after
fitting the whole training set, implying minimal feature learning until past
overfitting. The dashed lines in the middle and left figures indicate train
set statistics, and the solid lines correspond to the test set.

Empirical Verification of Our Theory: In terms of whether gradient descent
finds a model satisfying these conditions: Figure 4 empirically evaluates the
$\ell_{\infty}$ norm of the learned weights through gradient descent on the
classification task, with tiny $\ell_{\infty}$ regularization (weight of
$10^{-20}$), across different values of $p$. Again, it can be seen that the
$\ell_{\infty}$ norm of the weights of the network do not grow with the
problem dimension. These experiment use $n=2\,p^{5/3}$ data points for each
$p$, with a learning rate of 10.

Similar to the regression setting, Footnote 1 presents empirical evidence on
the impact of kernel regime on the poor generalization capabilities in the
early phase of training by showing the minimal change of empirical NTK until
after overfitting in the classification setting.

Through changing the _scale of initialization_ , in Footnote 1 we demonstrate
that by decreasing the scale of initialization in the classification task one
can mitigate grokking such that the gap between overfitting and generalization
becomes larger or smaller. This further supports that as long as the network
lies in the kernel regime, generalization without having access to the a
constant fraction of the dataset is impossible.

Grokking Modular Addition in Transformers: Figure 5 suggests that, similarly
to the two-layer network, grokking in the original Transformer studied by [41]
might be explainable through the same mechanism. In fact, Theorem 3.3, albeit
with small modifications to incorporate the shared embeddings in the
transformer architecture, applies to transformers as well.

##  5 Additional Related Works

Limitations imposed by Equivariance of Learning Algorithms. [1] study the
impact of equivariance of the training algorithms on the efficiency of
learning different functions on different networks. In particular, they
consider two main setups: a) learning with FCNs using noisy GD with clipped
gradients throughout the training, and b) learning a specific instance of the
modular addition task ($p=2$ with noisy inputs) with FCNs using SGD. Although
their approach in studying lower bounds for efficient learning shares some
high level similarity with ours in using equivariance of the training
algorithm, the settings considered are significantly differs from ours.
Moreover, in Section D.4 we present a novel abstract framework for analyzing
lower bounds on population $\ell_{2}$ loss for general function classes. [30]
present another technique in analysis of population loss lower bounds, which
shares some high-level similarity with our framework, albeit their analysis is
more restrictive on function classes. [39] also discusses rotational
equivariance of many learning algorithms and presents a general lower bound on
the 0-1 population error of such algorithms in the general case.

Margin Maximization as the Late Phase Implicit Bias. [35] present an
analytical solution for the max-margin solution of learning modular addition
in a classification setting similar to Section 4 when using all of the dataset
in training the network. Similar to our analysis of the rich regime in this
setting, they face difficulties in proving results under the assumption of
bounded $\ell_{2}$ norm for weights of the network, and assume an $\ell_{2,3}$
bound instead.

##  6 Conclusion

In this work, we studied the phenomenon of grokking in learning modular
addition with gradient descent on two-layer networks, modeled as regression or
classification. We showed that learning modular addition as presented is
fundamentally a difficult task for kernel models (for example neural networks
in kernel regime) due to the inherent symmetry and permutation-equivariance of
the task. We theoretically established this difficulty by presenting sample
complexity lower bounds of order of constant fraction of the whole dataset. We
further showed that networks satisfying certain conditions generalize far
better than those in the kernel regime, that such networks exist, and showed
empirical evidence that simple regularized gradient descent can eventually
find them, once it escapes the kernel regime. These results, in combination,
attempt to address why grokking is observed when learning modular addition. We
provide strong evidence to the hypothesis [28, 21] that, on this important
problem, it is indeed due to a separation between kernel and non-kernel
behavior of gradient descent.

Future Work. We have guaranteed large $\ell_{2}$ population loss for networks
in kernel regime on both regression and classification settings. It is
possible, however, for networks to have arbitrarily high $\ell_{2}$ loss while
having perfect classification accuracy. We conjecture that impossibility
results for accuracy-based generalization may also be possible based on
permutation equivariance in the early phase of training, but leave it as
future work. Moreover, we only study the cause of grokking in these settings,
but do not analyze possible training techniques to enable quick generalization
on this task. Although we are able to eliminate grokking through changing the
scale of initialization, doing so actually slows down the time to final
generalization; finding practical methods to enable quick generalization would
be more useful.

## Acknowledgments

We would like to thank Kaifeng Lyu and Wei Hu for helpful discussions. This
work was enabled in part by support provided by the Natural Sciences and
Engineering Research Council of Canada, the Canada CIFAR AI Chairs program,
Advanced Research Computing at the University of British Columbia, Calcul
Québec, the BC DRI Group, and the Digital Research Alliance of Canada.

## References

  * [1] Emmanuel Abbe and Enric Boix-Adsera  “On the non-universality of deep learning: quantifying the cost of symmetry”  In _NeurIPS_ , 2022  arXiv:2208.03113
  * [2] Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Shunting Zhang, Michael Suo, Phil Tillet, Xu Zhao, Eikan Wang, Keren Zhou, Richard Zou, Xiaodong Wang, Ajit Mathews, William Wen, Gregory Chanan, Peng Wu and Soumith Chintala  “PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation”  In _Architectural Support for Programming Languages and Operating Systems_ , 2024
  * [3] Sanjeev Arora, Nadav Cohen, Wei Hu and Yuping Luo  “Implicit Regularization in Deep Matrix Factorization”  In _NeurIPS_ , 2019  arXiv:1905.13655
  * [4] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov and Ruosong Wang  “On exact computation with an infinitely wide neural net”  In _NeurIPS_ , 2019  arXiv:1904.11955
  * [5] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach and Cyril Zhang  “Hidden progress in deep learning: SGD learns parities near the computational limit”  In _NeurIPS_ , 2022  arXiv:2207.08799
  * [6] Satwik Bhattamishra, Arkil Patel, Varun Kanade and Phil Blunsom  “Simplicity Bias in Transformers and their Ability to Learn Sparse Boolean Functions”  In _ACL_ , 2023  arXiv:2211.12316
  * [7] Guy Blanc, Neha Gupta, Gregory Valiant and Paul Valiant  “Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process”  In _COLT_ , 2020  arXiv:1904.09080
  * [8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang  “JAX: composable transformations of Python+NumPy programs”, 2018  URL: http://github.com/google/jax
  * [9] François Charton  “Learning the greatest common divisor: explaining transformer predictions”  In _ICLR_ , 2024  arXiv:2308.15594
  * [10] Lenaic Chizat, Edouard Oyallon and Francis Bach  “On Lazy Training in Differentiable Programming”  In _NeurIPS_ , 2019  arXiv:1812.07956
  * [11] Richard Courant and David Hilbert  “Methods of Mathematical Physics”  Interscience Publishers, 1953
  * [12] Alex Damian, Tengyu Ma and Jason D. Lee  “Label Noise SGD Provably Prefers Flat Global Minimizers”  In _NeurIPS_ , 2021  arXiv:2106.06530
  * [13] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy and Surya Ganguli  “Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel”  In _NeurIPS_ , 2020  arXiv:2010.15110
  * [14] Mario Geiger, Stefano Spigler, Arthur Jacot and Matthieu Wyart  “Disentangling feature and lazy training in deep neural networks”  In _Journal of Statistical Mechanics: Theory and Experiment_ 2020.11 IOP Publishing, 2020, pp. 113301  arXiv:1906.08034
  * [15] Andrey Gromov  “Grokking modular arithmetic”, 2023  arXiv:2301.02679
  * [16] Suriya Gunasekar, Jason Lee, Daniel Soudry and Nathan Srebro  “Characterizing Implicit Bias in Terms of Optimization Geometry”  In _ICML_ , 2018  arXiv:1802.08246
  * [17] Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur and Nathan Srebro  “Implicit Regularization in Matrix Factorization”  In _NeurIPS_ , 2017  arXiv:1705.09280
  * [18] Jeff Z. HaoChen, Colin Wei, Jason D. Lee and Tengyu Ma  “Shape Matters: Understanding the Implicit Bias of the Noise Covariance”  In _COLT_ , 2020  arXiv:2006.08680
  * [19] Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun  “Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification”  In _ICCV_ , 2015  arXiv:1502.01852
  * [20] Arthur Jacot, Franck Gabriel and Clément Hongler  “Neural tangent kernel: Convergence and generalization in neural networks”  In _NeurIPS_ , 2018  arXiv:1806.07572
  * [21] Tanishq Kumar, Blake Bordelon, Samuel J. Gershman and Cengiz Pehlevan  “Grokking as the Transition from Lazy to Rich Training Dynamics”  In _ICLR_ , 2024  arXiv:2310.06110
  * [22] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein and Jeffrey Pennington  “Wide neural networks of any depth evolve as linear models under gradient descent”  In _NeurIPS_ , 2019  arXiv:1902.06720
  * [23] Noam Levi, Alon Beck and Yohai Bar-Sinai  “Grokking in Linear Estimators – A Solvable Model that Groks without Understanding”  In _ICLR_ , 2024  arXiv:2310.16441
  * [24] Zhiyuan Li, Tianhao Wang and Sanjeev Arora  “What Happens after SGD Reaches Zero Loss? –A Mathematical Framework”  In _ICLR_ , 2022  arXiv:2110.06914
  * [25] Zhiyuan Li, Yi Zhang and Sanjeev Arora  “Why are convolutional nets more sample-efficient than fully-connected nets?”  In _ICLR_ , 2021  arXiv:2010.08515
  * [26] Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark and Mike Williams  “Towards understanding grokking: An effective theory of representation learning”  In _NeurIPS_ , 2022  arXiv:2205.10343
  * [27] Ziming Liu, Eric J Michaud and Max Tegmark  “Omnigrok: Grokking beyond algorithmic data”  In _ICLR_ , 2023  arXiv:2210.01117
  * [28] 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”  In _ICLR_ , 2024  arXiv:2311.18817
  * [29] Kaifeng Lyu and Jian Li  “Gradient descent maximizes the margin of homogeneous neural networks”  In _ICLR_ , 2020  arXiv:1906.05890
  * [30] Eran Malach and Shai Shalev-Shwartz  “When Hardness of Approximation Meets Hardness of Learning”  In _Journal of Machine Learning Research_ 23.91, 2022, pp. 1–24  arXiv:2008.08059
  * [31] David A. McAllester  “Simplified PAC-Bayesian Margin Bounds”  In _COLT_ , 2003
  * [32] Mohamad Amin Mohamadi, Wonho Bae and Danica J Sutherland  “A fast, well-founded approximation to the empirical neural tangent kernel”  In _International Conference on Machine Learning_ , 2023, pp. 25061–25081  PMLR
  * [33] Mehryar Mohri, Afshin Rostamizadeh and Ameet Talkwalkar  “Foundations of Machine Learning”  MIT Press, 2018  URL: https://cs.nyu.edu/~mohri/mlbook/
  * [34] Edward Moroshko, Suriya Gunasekar, Blake Woodworth, Jason D. Lee, Nathan Srebro and Daniel Soudry  “Implicit Bias in Deep Linear Classification: Initialization Scale vs Training Accuracy”  In _NeurIPS_ , 2020  arXiv:2007.06738
  * [35] Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu, Rosie Zhao and Sham Kakade  “Feature emergence via margin maximization: case studies in algebraic tasks”  In _ICLR_ , 2024  arXiv:2311.07568
  * [36] Mor Shpigel Nacson, Suriya Gunasekar, Jason D. Lee, Nathan Srebro and Daniel Soudry  “Lexicographic and Depth-Sensitive Margins in Homogeneous and Non-Homogeneous Deep Models”  In _ICML_ , 2019  arXiv:1905.07325
  * [37] Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith and Jacob Steinhardt  “Progress measures for grokking via mechanistic interpretability”  In _ICLR_ , 2023  arXiv:2301.05217
  * [38] Behnam Neyshabur, Srinadh Bhojanapalli and Nathan Srebro  “A PAC-Bayesian Approach to Spectrally-Normalized Margin Bounds for Neural Networks”  In _ICLR_ , 2018  arXiv:1707.09564
  * [39] Andrew Y Ng  “Feature selection, $L_{1}$ vs. $L_{2}$ regularization, and rotational invariance”  In _ICML_ , 2004
  * [40] Pascal Jr. Tikeng Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish and Guillaume Dumas  “Predicting Grokking Long Before it Happens: A look into the loss landscape of models which grok”, 2023  arXiv:2306.13253
  * [41] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin and Vedant Misra  “Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets”, 2022  arXiv:2201.02177
  * [42] Martin Raab and Angelika Steger  “‘Balls into Bins’ — A Simple and Tight Analysis”  In _Randomization and Approximation Techniques in Computer Science_ , 1998
  * [43] Noa Rubin, Inbar Seroussi and Zohar Ringel  “Grokking as a First Order Phase Transition in Two Layer Networks”  In _ICLR_ , 2024  arXiv:2310.03789
  * [44] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar and Nathan Srebro  “The Implicit Bias of Gradient Descent on Separable Data”  In _JMLR_ 19.70, 2018  arXiv:1710.10345
  * [45] Nathan Srebro, Karthik Sridharan and Ambuj Tewari  “Smoothness, low noise and fast rates”  In _NeurIPS_ , 2010  arXiv:1009.3896
  * [46] Matus Telgarsky  “Feature selection with gradient descent on two-layer networks in low-rotation regimes”, 2022  arXiv:2208.02789
  * [47] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss and Joshua Susskind  “The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon”, 2022  arXiv:2206.04817
  * [48] Joel A. Tropp  “An Introduction to Matrix Concentration Inequalities”  In _Foundations and Trends® in Machine Learning_ 8.1-2, 2015, pp. 1–230  arXiv:1501.01571
  * [49] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár and Ramana Kumar  “Explaining grokking through circuit efficiency”, 2023  arXiv:2309.02390
  * [50] Martin J. Wainwright  “High-Dimensional Statistics: A Non-Asymptotic Viewpoint”  Cambridge University Press, 2019
  * [51] Colin Wei, Jason D. Lee, Qiang Liu and Tengyu Ma  “Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel”  In _NeurIPS_ , 2019  arXiv:1810.05369
  * [52] Shuo Xie and Zhiyuan Li  “Implicit Bias of AdamW: $\ell_{\infty}$ Norm Constrained Optimization”  In _ICML_ , 2024  arXiv:2404.04454
  * [53] Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi and Wei Hu  “Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data”  In _ICLR_ , 2024  arXiv:2310.02541
  * [54] Greg Yang and Edward J. Hu  “Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks”  In _ICML_ , 2021  arXiv:2011.14522
  * [55] Chenyang Zhang, Difan Zou and Yuan Cao  “The Implicit Bias of Adam on Separable Data”, 2024  arXiv:2406.10650
  * [56] Mauricio A Álvarez, Lorenzo Rosasco and Neil D Lawrence  “Kernels for vector-valued functions: A review”  In _Foundations and Trends® in Machine Learning_ 4.3, 2012, pp. 195–266  arXiv:1106.6251

##  Appendix A Experimental Setup

In this section, we briefly explain the setup used for our experimental
evaluations.

###  A.1 Regression

We use vanilla gradient descent with squared loss and tiny $\ell_{\infty}$
regularization for 50,000, 100,000 or 200,000 steps for each experiment. In
regression, our learning rate has been fixed to 1, and the regularization
strength has been set to $10^{-4}$. The network has been initialized according
to [19]. The amount of data used for training in regression task has been set
to $2\times p^{2.25}$.

###  A.2 Classification

We use vanilla gradient descent with cross-entropy loss and tiny
$\ell_{\infty}$ regularization, for up to 100,000 steps. The learning rate in
the presented experiments was set to 10 and was kept constant during the
training. The regularization strength of $\ell_{\infty}$ regularizer has been
set to $10^{-20}$. To accelerate the training with cross-entropy loss, we use
the "normalized" GD trick, where the learning rate of each step is scaled by
the inverse of the norm of the gradient:

$\theta_{t+1}=\theta_{t}-\eta\,\frac{\nabla_{\theta}\ell(\theta_{t})}{\left\lVert\nabla_{\theta}\ell(\theta_{t})\right\rVert_{2}}$
(A.1)

where $\ell$ denotes the loss function and $\eta$ denotes the learning rate.
The network has been initialized according to [19]. The amount of data used
for training in regression task has been set to $2\times p^{5/3}$.

###  A.3 Transformer

To train the one-layer transformer we have used full-batch gradient descent
with a learning rate of $\eta=0.25$ and a tiny $\ell_{\infty}$ regularization
with the strength of $1.0e-20$. The network has been initialized according to
default PyTorch initialziation (migrated ot JAX). We have used $2\times
p^{5/3}$ of the data for training.

###  A.4 Logistics

We used the JAX framework [8] to implement and run the experiments on machines
using NVIDIA V100 or A100 GPUs.

##  Appendix B Ability of Neural Tangent Kernel Models to Interpolate

In this section, we prove Theorem 3.1, that neural tangent kernel models for
our one-hidden-layer quadratic network are able to exactly interpolate their
training data, in the regression setting.

If the training set contains duplicate ${\bm{x}}$, say
${\bm{x}}_{i}={\bm{x}}_{j}$ for $i\neq j$, then $\mathbf{K}$ cannot be full
rank – the $i$th and $j$th rows are necessarily identical. Of course, there
also exist unattainable targets; just set $y_{i}\neq y_{j}$. Call a labeling
_consistent_ if for all $i$ and $j$ such that ${\bm{x}}_{i}={\bm{x}}_{j}$,
$y_{i}=y_{j}$. A learning algorithm can achieve any consistent labeling if and
only if that algorithm applied to the largest distinct subset (e.g. removing
$j$ but keeping $i$) can achieve any labeling. Thus, we assume without loss of
generality in the remainder of this section that the training set contains no
duplicates.

###  B.1 Full-Rank Kernels are Expressive

Recalling that neural tangent kernel models correspond to “ridgeless” kernel
regression [20, 22, e.g.], we first notice that this method can interpolate
_any_ set of training labels when the kernel is strictly positive definite.

Specifically, empirical neural tangent kernel regression corresponds to
ridgeless regression with a prior mean corresponding to the network at
initialization $f_{0}$,

$\operatornamewithlimits{argmin}_{f:\forall i,\;f({\bm{x}}_{i})=y_{i}}\lVert
f-f_{0}\rVert_{K}=f_{0}+\operatornamewithlimits{argmin}_{f:\forall
i,\;f({\bm{x}}_{i})=y_{i}-f_{0}}\lVert
f\rVert_{K}=f_{0}+\sum_{i=1}^{n}K(\cdot,{\bm{x}}_{i})\lambda_{i}\;\;\text{ for
}\lambda=\mathbf{K}^{\dagger}\mathbf{y^{\prime}},$

where $\mathbf{K}^{\dagger}$ is the Moore-Penrose pseudoinverse of the
$n\times n$ kernel matrix $[K({\bm{x}}_{i},{\bm{x}}_{j})]_{ij}$, and
$\mathbf{y^{\prime}}$ has $i$th entry $y_{i}-f_{0}({\bm{x}}_{i})$. By changing
variables between $\mathbf{y^{\prime}}$ and $\mathbf{y}$, we need not
explicitly consider $f_{0}$ in the following two results.

######  Lemma B.1.

Let $\mathcal{D}_{\mathrm{train}}=\\{(x_{i},y_{i})\\}_{i\in[n]}$, and let $K$
be a kernel such that the kernel matrix $\mathbf{K}=[K(x_{i},x_{j})]_{ij}$ is
strictly positive definite. Then kernel ridgeless regression achieves zero
training error on $\mathcal{D}_{\mathrm{train}}$.

###### Proof.

This well-known result follows from the fact that
$[\hat{f}({\bm{x}}_{i})]_{i}=\mathbf{K}\mathbf{K}^{\dagger}\mathbf{y}$. If
$\mathbf{K}$ is strictly positive definite,
$\mathbf{K}^{\dagger}=\mathbf{K}^{-1}$, and so
$[\hat{f}(x_{i})]_{i}=\mathbf{y}$. ∎

The following result is a partial converse.

######  Lemma B.2.

Let $\\{x_{i}\\}_{i\in[n]}$ and $K$ be a kernel such that the kernel matrix
$\mathbf{K}=[K(x_{i},x_{j})]_{ij}$ is singular. Then there exists an
assignment of $y_{i}\in\mathbb{R}$ such that kernel ridgeless regression
achieves nonzero training error on
$\mathcal{D}_{\mathrm{train}}=\\{(x_{i},y_{i})\\}_{i\in[n]}$.

###### Proof.

Let $\mathbf{y}$ be any nonzero vector in the null space of $\mathbf{K}$; such
vectors exist since $\mathbf{K}$ is singular. Then $\mathbf{y}$ is also in the
null space of $\mathbf{K}^{\dagger}$, and the training set predictions are
$\mathbf{K}\mathbf{K}^{\dagger}\mathbf{y}=0\neq\mathbf{y}$. ∎

###  B.2 Expected NTK is Full-Rank

We next show that, in the regression setting, the _expected_ neural tangent
kernel is strictly positive definite. While perhaps of interest of its own
accord, this will be a key component in our analysis of finite-width networks
that follows.

######  Proposition B.3.

Let the entries of $W$ follow a distribution $\mathcal{P}_{W}$, all mutually
independent. Assume $\mathcal{P}_{W}$ has mean zero, variance
$\sigma_{W}^{2}>0$, skewness zero, and kurtosis $\kappa_{W}$ such that
$\operatorname*{\mathbb{E}}_{w\sim\mathcal{P}_{W}}w^{4}=\kappa_{W}\sigma_{W}^{4}$.
The entries of $V$ can be arbitrary. Then the expected neural tangent kernel
for the regression network with any finite $h\geq 1$ is strictly positive
definite on any set of distinct inputs $\\{x_{i}:i\in[n]\\}$, with minimum
eigenvalue at least $4h\sigma_{W}^{4}$.

We first note that if $\mathcal{P}_{W}=\mathcal{N}(0,\sigma_{W}^{2})$,
$\kappa_{V}=3$. If $\mathcal{P}_{W}=\mathrm{Unif}([-s_{W},s_{W}])$,
$\sigma_{W}^{2}=\frac{1}{3}s_{W}^{2}$ and $\kappa_{W}=\frac{9}{5}$. These two
distributions cover the vast majority of initialization schemes used in
practice.

###### Proof.

It will be more convenient in this proof to give different names to the sub-
matrix of the parameter vector $W$ that act on the $a$ inputs and the $b$
inputs; we will write $W=\begin{bmatrix}Q&R\end{bmatrix}$, where $Q,R$ are
each $h\times p$ matrices. Then we can write the full model as

$g\left(\theta;(e_{a},e_{b},e_{c})\right)=e_{c}^{\mathsf{T}}V(Qe_{a}+Re_{b})^{\odot
2}=\sum_{k=1}^{h}V_{ck}(Q_{ka}+R_{kb})^{2}.$

The empirical neural tangent kernel between two inputs $x=(e_{a},e_{b},e_{c})$
and $x^{\prime}=(e_{a^{\prime}},e_{b^{\prime}},e_{c^{\prime}})$ is then given
by

$\begin{split}K_{\theta}(x,x^{\prime})=\sum_{k=1}^{h}\Biggl{[}\sum_{c^{\prime\prime}=1}^{p}&\frac{\partial
g(\theta;x)}{\partial V_{c^{\prime\prime}k}}\frac{\partial
g(\theta;x^{\prime})}{\partial V_{c^{\prime\prime}k}}\\\
&+\sum_{a^{\prime\prime}=1}^{p}\frac{\partial g(\theta;x)}{\partial
Q_{ka^{\prime\prime}}}\frac{\partial g(\theta;x^{\prime})}{\partial
Q_{ka^{\prime\prime}}}+\sum_{b^{\prime\prime}=1}^{p}\frac{\partial
g(\theta;x)}{\partial R_{kb^{\prime\prime}}}\frac{\partial
g(\theta;x^{\prime})}{\partial R_{kb^{\prime\prime}}}\Biggr{]}.\end{split}$
(B.1)

Evaluating the $V$ derivatives,

$\displaystyle\frac{\partial g(\theta;x)}{\partial V_{c^{\prime\prime}k}}$
$\displaystyle=\begin{cases}(Q_{ka}+R_{kb})^{2}&\text{if
}c=c^{\prime\prime}\\\ 0&\text{otherwise}\end{cases}$
$\displaystyle\sum_{c^{\prime\prime}=1}^{p}\frac{\partial
g(\theta;x)}{\partial V_{c^{\prime\prime}k}}\frac{\partial
g(\theta;x^{\prime})}{\partial V_{c^{\prime\prime}k}}$
$\displaystyle=\begin{cases}(Q_{ka}+R_{kb})^{2}(Q_{ka^{\prime}}+R_{kb^{\prime}})^{2}&\text{if
}c=c^{\prime}\\\ 0&\text{otherwise}.\end{cases}$

While it would not be difficult to analyze the $Q$ and $R$ derivatives as
well, their exact form will not be important. Instead, we only need to write

$\displaystyle K_{\theta}(x,x^{\prime})$
$\displaystyle=J_{\theta}(x,x^{\prime})+\sum_{k=1}^{h}L_{\theta_{k}}(x,x^{\prime})$
(B.2) $\displaystyle J_{\theta}(x,x^{\prime})$
$\displaystyle=\sum_{k=1}^{h}\sum_{a^{\prime\prime}=1}^{p}\frac{\partial
g(\theta;x)}{\partial Q_{ka^{\prime\prime}}}\frac{\partial
g(\theta;x^{\prime})}{\partial
Q_{ka^{\prime\prime}}}+\sum_{b^{\prime\prime}=1}^{p}\frac{\partial
g(\theta;x)}{\partial R_{kb^{\prime\prime}}}\frac{\partial
g(\theta;x^{\prime})}{\partial R_{kb^{\prime\prime}}}$ $\displaystyle
L_{\theta_{k}}(x,x^{\prime})$
$\displaystyle=(Q_{ka}+R_{kb})^{2}(Q_{ka^{\prime}}+R_{kb^{\prime}})^{2}\mathbf{1}(c=c^{\prime}).$
(B.3)

The function $J_{\theta}$ has an explicit feature map corresponding to the
relevant gradients; thus, for any set of inputs $\\{x_{i}:i\in[n]\\}$ and any
value of $\theta$, the kernel matrix
$\mathbf{J}_{\theta}=[J_{\theta}(x_{i},x_{j})]_{ij}$ is positive semi-
definite. This implies, e.g. via Weyl’s inequality, that
$\operatorname*{\mathbb{E}}_{\theta}\mathbf{J}_{\theta}$ is also positive
semi-definite.

For any fixed set of inputs $\\{x_{i}:i\in[n]\\}$, the kernel matrices
$\mathbf{L}_{\theta_{k}}=[L_{\theta_{k}}({\bm{x}}_{i},{\bm{x}}_{j})]_{ij}$ are
independent and identically distributed, and in particular have the same mean
$\operatorname*{\mathbb{E}}\mathbf{L}_{\theta_{k}}$ for any arbitrary choice
of $k\in[p]$. Thus the expected neural tangent kernel matrix can be written as

$[\operatorname*{\mathbb{E}}K_{\theta}(x_{i},x_{j})]_{ij}=\operatorname*{\mathbb{E}}\mathbf{J}_{\theta}+h\operatorname*{\mathbb{E}}\mathbf{L}_{\theta_{k}}.$

We will show that $\operatorname*{\mathbb{E}}\mathbf{L}_{\theta_{k}}$ has
minimum eigenvalue at least $4\sigma_{W}^{4}$, from which the result follows.

To do so, we will evaluate

$\operatorname*{\mathbb{E}}L_{\theta_{k}}(x,x^{\prime})=\operatorname*{\mathbb{E}}\Bigl{[}(Q_{ka}+R_{kb})^{2}(Q_{ka^{\prime}}+R_{kb^{\prime}})^{2}\Bigr{]}\mathbf{1}(c=c^{\prime})$

for arbitrary nonrandom ${\bm{x}}=(e_{a},e_{b},e_{c})$ and
${\bm{x}}^{\prime}=(e_{a^{\prime}},e_{b^{\prime}},e_{c^{\prime}})$, where all
expectations will be over the relevant parameters
$\\{Q_{ka}:a\in[p]\\}\cup\\{R_{kb}:b\in[p]\\}$.

If $a=a^{\prime}$ and $b=b^{\prime}$,

$\displaystyle\operatorname*{\mathbb{E}}\left[(Q_{ka}+R_{kb})^{4}\right]$
$\displaystyle=\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{4}}_{\kappa_{W}\sigma_{W}^{4}}+4\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{3}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb}}_{0}+6\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{2}}_{\sigma_{W}^{2}}+4\underbrace{\operatorname*{\mathbb{E}}Q_{ka}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{3}}_{0}+\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{4}}_{\kappa_{W}\sigma_{W}^{4}}$
$\displaystyle=(2\kappa_{W}+6)\sigma_{W}^{4}.$

If instead $a=a^{\prime}$ but $b\neq b^{\prime}$, then

$\displaystyle\operatorname*{\mathbb{E}}\left[(Q_{ka}+R_{kb})^{2}(Q_{ka}+R_{kb^{\prime}})^{2}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}\left[(Q_{ka}^{2}+2Q_{ka}R_{kb}+R_{kb}^{2})(Q_{ka}^{2}+2Q_{ka}R_{kb^{\prime}}+R_{kb^{\prime}}^{2})\right]$
$\displaystyle=\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{4}}_{\kappa_{W}\sigma_{W}^{4}}+2\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{3}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}}_{0}+\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}^{2}}_{\sigma_{W}^{2}}$
$\displaystyle\qquad+2\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{3}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb}}_{0}+4\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}R_{kb}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}}_{0}+2\underbrace{\operatorname*{\mathbb{E}}Q_{ka}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}^{2}}_{\sigma_{W}^{2}}$
$\displaystyle\qquad+\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{2}}_{\sigma_{W}^{2}}+2\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}Q_{ka}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}}_{0}+\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{2}}_{\sigma_{W}^{2}}\underbrace{\operatorname*{\mathbb{E}}R_{kb^{\prime}}^{2}}_{\sigma_{W}^{2}}$
$\displaystyle=\left(\kappa_{W}+3\right)\sigma_{W}^{4};$

the case where $a\neq a^{\prime}$ but $b=b^{\prime}$ is the same by symmetry.

Finally, when $a\neq a^{\prime}$ and $b\neq b^{\prime}$, we have by
independence that

$\displaystyle\operatorname*{\mathbb{E}}\bigl{[}(Q_{ka}+R_{kb})^{2}(Q_{ka^{\prime}}+R_{kb^{\prime}})^{2}\bigr{]}$
$\displaystyle=\operatorname*{\mathbb{E}}\bigl{[}(Q_{ka}+R_{kb})^{2}\bigr{]}\operatorname*{\mathbb{E}}\bigl{[}(Q_{ka^{\prime}}+R_{kb^{\prime}})^{2}\bigr{]}=\Bigl{(}\operatorname*{\mathbb{E}}\bigl{[}(Q_{ka}+R_{kb})^{2}\bigr{]}\Bigr{)}^{2}$
$\displaystyle=\Bigl{(}\underbrace{\operatorname*{\mathbb{E}}Q_{ka}^{2}}_{\sigma_{W}^{2}}+2\underbrace{\operatorname*{\mathbb{E}}Q_{ka}}_{0}\underbrace{\operatorname*{\mathbb{E}}R_{kb}}_{0}+\underbrace{\operatorname*{\mathbb{E}}R_{kb}^{2}}_{\sigma_{W}^{2}}\Bigr{)}^{2}=(2\sigma_{W}^{2})^{2}=4\sigma_{W}^{4}.$

Combining the cases, it holds in general that

$\displaystyle\operatorname*{\mathbb{E}}L_{\theta_{k}}(x,x^{\prime})$
$\displaystyle=\sigma_{W}^{4}\mathbf{1}(c=c^{\prime})\begin{cases}4&\text{if
}a\neq a^{\prime},b\neq b^{\prime}\\\ \kappa_{W}+3&\text{if
}a=a^{\prime},b\neq b^{\prime}\\\ \kappa_{W}+3&\text{if }a\neq
a^{\prime},b=b^{\prime}\\\ 2\kappa_{W}+6&\text{if
}a=a^{\prime},b=b^{\prime}\\\ \end{cases}$
$\displaystyle=4\sigma_{W}^{4}\mathbf{1}(c=c^{\prime})$
$\displaystyle\quad+(\kappa_{W}-1)\sigma_{W}^{4}\mathbf{1}(a=a^{\prime},c=c^{\prime})$
$\displaystyle\quad+(\kappa_{W}-1)\sigma_{W}^{4}\mathbf{1}(b=b^{\prime},c=c^{\prime})$
$\displaystyle\quad+4\sigma_{W}^{4}\mathbf{1}(a=a^{\prime},b=b^{\prime},c=c^{\prime}).$

The kurtosis of any probability distribution is at least 1 by Jensen’s
inequality, so all the coefficients in this last form are nonnegative. We can
then use this to construct an explicit feature map for each term in
$\operatorname*{\mathbb{E}}L_{\theta_{k}}$, because if $\gamma\geq 0$,

$\gamma\mathbf{1}(c=c^{\prime})=\begin{bmatrix}\sqrt{\gamma}\mathbf{1}(c=1)\\\
\sqrt{\gamma}\mathbf{1}(c=2)\\\ \vdots\\\
\sqrt{\gamma}\mathbf{1}(c=p)\end{bmatrix}^{\mathsf{T}}\begin{bmatrix}\sqrt{\gamma}\mathbf{1}(c^{\prime}=1)\\\
\sqrt{\gamma}\mathbf{1}(c^{\prime}=2)\\\ \vdots\\\
\sqrt{\gamma}\mathbf{1}(c^{\prime}=p)\end{bmatrix}$ (B.4)

corresponds to explicit features in $\mathbb{R}^{p}$. The other indicator
functions can be implemented in the same way, in $\mathbb{R}^{p^{2}}$ or
$\mathbb{R}^{p^{3}}$; their sum can then be obtained by concatenating the
individual features together. This construction makes it clear that the
function

$M(x,x^{\prime})=4\sigma_{W}^{4}\mathbf{1}(c=c^{\prime})+(\kappa_{W}-1)\sigma_{W}^{4}\mathbf{1}(a=a^{\prime},c=c^{\prime})+(\kappa_{W}-1)\sigma_{W}^{4}\mathbf{1}(b=b^{\prime},c=c^{\prime})$

is a positive semi-definite kernel, so $\mathbf{M}=[M(x_{i},x_{j})]_{ij}$ is a
positive semi-definite matrix for any set of inputs $\\{x_{i}\\}$.

It remains to show that the final component of the kernel is strictly positive
definite. Noticing that
$\mathbf{1}(a=a^{\prime},b=b^{\prime},c=c^{\prime})=\mathbf{1}(x=x^{\prime})$,
if the $\\{{\bm{x}}_{i}:i\in[n]\\}$ are distinct, the expected kernel matrix
is
$\operatorname*{\mathbb{E}}\mathbf{L}_{\theta_{k}}=\mathbf{M}+4\sigma_{W}^{4}\mathbf{I}$.
Since $\mathbf{M}$ is positive semi-definite, this has minimum eigenvalue at
least $4\sigma_{W}^{4}>0$, as desired. ∎

###  B.3 Empirical NTKs are Likely Full-Rank

Now, for _bounded_ initialization schemes, we use matrix concentration
inequalities to show that the empirical neural tangent kernel is also likely
to be full-rank when $h$ is large enough.

######  Proposition B.4.

In the setting of Proposition B.3, further assume that
$\Pr_{w\sim\mathcal{P}_{W}}(\lvert w\rvert\leq s_{W})=1$. Let the set of
inputs $\\{x_{i}:i\in[n]\\}$ be distinct and such that the most common $c$
value is seen $\rho_{c}$ times. Then the empirical neural tangent kernel for
the regression network is strictly positive definite with probability at least
$1-\delta$ over the choice of random parameters $\theta$ as long as

$h>\frac{4s_{W}^{4}}{\sigma_{W}^{4}}\rho_{c}\log\frac{n}{\delta}.$

If $\mathcal{P}_{W}$ is uniform on $[-s_{W},s_{W}]$ for $s_{W}>0$, this
condition is equivalent to

$h>36\rho_{c}\log\frac{n}{\delta}.$

###### Proof.

Recall the decomposition of
$K_{\theta}({\bm{x}},{\bm{x}}^{\prime})=J_{\theta}({\bm{x}},{\bm{x}}^{\prime})+\sum_{k=1}^{h}L_{\theta_{k}}({\bm{x}},{\bm{x}}^{\prime})$
from (B.2), and the corresponding $n\times n$ matrices
$\mathbf{K}_{\theta}=[K_{\theta}({\bm{x}}_{i},{\bm{x}}_{j})]_{ij}$,
$\mathbf{J}_{\theta}=[J_{\theta}({\bm{x}}_{i},{\bm{x}}_{j})]_{ij}$, and
$\mathbf{L}_{\theta_{k}}=[L_{\theta_{k}}({\bm{x}}_{i},{\bm{x}}_{j})]_{ij}$, so
that
$\mathbf{K}_{\theta}=\mathbf{J}_{\theta}+\sum_{k=1}^{h}\mathbf{L}_{\theta_{k}}$.
As shown in the proof of Proposition B.3, $\mathbf{J}_{\theta}$ is positive
semi-definite; we now wish to show that the sum of $h$ iid matrices
$\mathbf{L}_{\theta_{k}}$, whose mean
$h\operatorname*{\mathbb{E}}\mathbf{L}_{\theta_{k}}$ has minimum eigenvalue
$\mu_{\min}:=4h\sigma_{W}^{4}$, is likely to be full-rank. We can do so by
applying Lemma B.5, a direct corollary of matrix Chernoff bounds, if we
additionally have an almost sure upper bound on the operator norm
$\lVert\mathbf{L}_{\theta_{k}}\rVert$.

Notate the input ${\bm{x}}_{i}$ as $(e_{a_{i}},e_{b_{i}},e_{c_{i}})$ for
$a_{i},b_{i},c_{i}\in[p]$. Using (B.3), we can write

$\mathbf{L}_{\theta_{k}}=\operatorname{diag}(\mathbf{w}_{\theta_{k}})\mathbf{C}\operatorname{diag}(\mathbf{w}_{\theta_{k}}),$
(B.5)

where $\operatorname{diag}:\mathbb{R}^{n}\to\mathbb{R}^{n\times n}$ constructs
a diagonal matrix with the given vector on its diagonal and

$\displaystyle(\mathbf{C})_{ij}=\mathbf{1}(c_{i}=c_{j})$
$\displaystyle(\mathbf{w}_{\theta_{k}})_{i}=(Q_{ka_{i}}+R_{kb_{i}})^{2}.$

Using $\lVert AB\rVert\leq\lVert A\rVert\lVert B\rVert$ and
$\lVert\operatorname{diag}(v)\rVert=\max_{i}\lvert v_{i}\rvert$, we obtain
that

$\lVert\mathbf{L}_{\theta_{k}}\rVert\leq(2s_{W})^{2}\lVert\mathbf{C}\rVert(2s_{W})^{2}=16s_{W}^{4}\lVert\mathbf{C}\rVert.$

Using (B.4), we can write $\mathbf{C}=\Phi_{c}\Phi_{c}^{\mathsf{T}}$ so that
$\lVert\mathbf{C}\rVert=\lVert\Phi_{c}\rVert^{2}$, where
$\Phi_{c}\in\mathbb{R}^{n\times p}$ is given by

$\Phi_{c}=\begin{bmatrix}\mathbf{1}(c_{1}=1)&\cdots&\mathbf{1}(c_{1}=p)\\\
\vdots&\ddots&\vdots\\\
\mathbf{1}(c_{n}=1)&\cdots&\mathbf{1}(c_{n}=p)\end{bmatrix}.$

We can evaluate this operator norm with

$\displaystyle(\Phi_{c}x)_{i}$
$\displaystyle=\sum_{\ell=1}^{p}\mathbf{1}(c_{i}=\ell)x_{\ell}=x_{c_{i}}$
$\displaystyle\lVert\Phi_{c}x\rVert^{2}$
$\displaystyle=\sum_{i=1}^{n}x_{c_{i}}^{2}=\sum_{\ell=1}^{p}x_{\ell}^{2}\left(\sum_{i=1}^{n}\mathbf{1}(c_{i}=\ell)\right)$
$\displaystyle\lVert\Phi_{c}\rVert^{2}$ $\displaystyle=\sup_{\lVert
x\rVert=1}\sum_{\ell=1}^{p}x_{\ell}^{2}\left(\sum_{i=1}^{n}\mathbf{1}(c_{i}=\ell)\right)=\max_{\ell\in[p]}\sum_{i=1}^{n}\mathbf{1}(c_{i}=\ell)=\rho_{c}.$

Thus it holds almost surely that

$\lVert\mathbf{L}_{\theta_{k}}\rVert\leq 16s_{W}^{4}\rho_{c}=:L.$

Plugging in Lemma B.5, we have shown that

$\Pr\left(\lambda_{\min}\left(\mathbf{K}_{\theta}\right)>0\right)\geq
1-n\exp\left(-\frac{4h\sigma_{W}^{4}}{16s_{W}^{4}\rho_{c}}\right),$

and so $\mathbf{K}_{\theta}$ is full rank with probability at least $1-\delta$
as long as

$h>\frac{4s_{W}^{4}}{\sigma_{W}^{4}}\rho_{c}\log\frac{n}{\delta}.$

If the entries of $W$ are iid $\mathrm{Unif}([-s_{w},s_{w}])$ for $s_{w}>0$,
$\sigma_{W}^{2}=\frac{1}{3}s_{W}^{2}$, giving the condition

$h>36\rho_{c}\log\frac{n}{\delta}.\qed$

######  Lemma B.5.

Consider a finite sequence of independent positive semi-definite $d\times d$
real random matrices $\mathbf{X}_{1},\dots,\mathbf{X}_{m}$. Assume that
$\lVert\mathbf{X}_{k}\rVert\leq L$ almost surely, and that
$\mu_{\min}=\lambda_{\min}\left(\sum_{k=1}^{m}\operatorname*{\mathbb{E}}\mathbf{X}_{k}\right)>0$.
Then

$\Pr\left(\lambda_{\min}\left(\sum_{k=1}^{m}\mathbf{X}_{k}\right)>0\right)\geq
1-d\exp\left(-\frac{\mu_{\min}}{L}\right).$

###### Proof.

A matrix Chernoff bound [48, Theorem 5.1.1] in gives that for each
${\epsilon}\in[0,1)$,

$\Pr\left(\lambda_{\min}\left(\sum_{k=1}^{h}\mathbf{X}_{k}\right)>(1-{\epsilon})\mu_{\min}\right)\geq
1-d\left(\frac{e^{-{\epsilon}}}{(1-{\epsilon})^{1-{\epsilon}}}\right)^{{\mu_{\min}}/{L}}.$
(B.6)

If the smallest eigenvalue is strictly positive, it must exceed some
$\delta>0$, which corresponds to some ${\epsilon}<1$. Thus the event of being
full-rank is the limit of the events of exceeding each ${\epsilon}$, and since
the right-hand side of (B.6) is continuous for ${\epsilon}<1$, we can take the
limit as ${\epsilon}\nearrow 1$; since
${e^{-{\epsilon}}}/{(1-{\epsilon})^{1-{\epsilon}}}\to 1/e$, this gives the
desired result. ∎

This threshold is indeed tight up to logarithmic factors, in the sense that
there are some labels which kernel ridgeless regression cannot achieve when
$h=o(n/p)$.

######  Proposition B.6.

In the setting of Proposition B.3, the empirical neural tangent kernel of the
regression network $\mathbf{K}_{\theta}$ is guaranteed to be singular when
$h<n/(3p)$.

###### Proof.

Expanding on (B.1), we will additionally need to evaluate the $Q$ and $R$
derivatives:

$\displaystyle\frac{\partial g(\theta;{\bm{x}})}{\partial
Q_{ka^{\prime\prime}}}$
$\displaystyle=\begin{cases}2V_{ck}(Q_{ka}+R_{kb})&\text{if
}a=a^{\prime\prime}\\\ 0&\text{otherwise}\end{cases}$
$\displaystyle\sum_{a^{\prime\prime}=1}^{p}\frac{\partial
g(\theta;{\bm{x}})}{\partial Q_{ka^{\prime\prime}}}\frac{\partial
g({\bm{x}}^{\prime})}{\partial Q_{ka^{\prime\prime}}}$
$\displaystyle=\begin{cases}4V_{ck}V_{c^{\prime}k}(Q_{ka}+R_{kb})(Q_{ka}+R_{kb^{\prime}})&\text{if
}a=a^{\prime}\\\ 0&\text{otherwise}\end{cases}$ $\displaystyle\frac{\partial
g(\theta;{\bm{x}})}{\partial R_{kb^{\prime\prime}}}$
$\displaystyle=\begin{cases}2V_{ck}(Q_{ka}+R_{kb})&\text{if
}b=b^{\prime\prime}\\\ 0&\text{otherwise}\end{cases}$
$\displaystyle\sum_{b^{\prime\prime}=1}^{p}\frac{\partial
g({\bm{x}})}{\partial R_{kb^{\prime\prime}}}\frac{\partial
g({\bm{x}}^{\prime})}{\partial R_{kb^{\prime\prime}}}$
$\displaystyle=\begin{cases}4V_{ck}V_{c^{\prime}k}(Q_{ka}+R_{kb})(Q_{ka^{\prime}}+R_{kb})&\text{if
}b=b^{\prime}\\\ 0&\text{otherwise}.\end{cases}$

Writing as in (B.2) and (B.5), we have that

$\mathbf{K}_{\theta}=\sum_{k=1}^{h}\operatorname{diag}(\mathbf{w}_{\theta_{k}})\mathbf{C}\operatorname{diag}(\mathbf{w}_{\theta_{k}})+\operatorname{diag}(\mathbf{v}_{\theta_{k}})\mathbf{A}\operatorname{diag}(\mathbf{v}_{\theta_{k}})+\operatorname{diag}(\mathbf{v}_{\theta_{k}})\mathbf{B}\operatorname{diag}(\mathbf{v}_{\theta_{k}}),$

where $(\mathbf{v}_{\theta_{k}})_{i}=2V_{ck}(Q_{ka}+R_{kb})$,
$\mathbf{A}_{ij}=\mathbf{1}(a_{i}=a_{j})$, and
$\mathbf{B}_{ij}=\mathbf{1}(b_{i}=b_{j})$. By (B.4), each of the matrices
$\mathbf{A}$, $\mathbf{B}$, and $\mathbf{C}$ have rank at most $p$; thus
$\operatorname{rank}(\mathbf{K}_{\theta})\leq 3ph$. If $3ph<n$, this $n\times
n$ matrix cannot be full-rank. ∎

##  Appendix C Permutation-Equivariance of Gradient-Based Training

We first define the notion of permutation equivariance. Towards this, we
borrow the following proof from Appendix C of [25]:

######  Definition C.1 (Gradient-Based Algorithm $\mathcal{A}$).

We borrow the definition of Algorithm 1color=violet!20,noinline]Danica: Should
write out the actual definition to be self-contained, it’s not very long in
[25] with the slight modification of restricting the update rule
$F(U,\mathcal{M},\mathcal{D}_{\operatorname{train}})$ where $U$ denotes the
parameters and $\mathcal{M}:U\to(\mathcal{X}\to\mathbb{R})$ denotes the model
mapping parameters to a function. We restrict the update rule to only allow
gradient-based update rules, such that

$F(U,\mathcal{M},\mathcal{D}_{\operatorname{train}})=U-\eta\nabla_{U}\mathcal{L}(\mathcal{M}(U),\mathcal{D}_{\operatorname{train}})$

for some $\eta\in\mathbb{R}$ and loss function $\mathcal{L}$ mapping a
function and a dataset to a loss
$(\mathcal{X}\to\mathbb{R})\times\mathcal{D}\to\mathbb{R}$.

######  Theorem C.2.

Suppose $\mathcal{G}_{\mathcal{X}}$ is a group acting on $\mathcal{X}$. The
gradient-based iterative algorithm $\mathcal{A}$ (defined in Definition C.1)
is $\mathcal{G}_{\mathcal{X}}$-equivariant if: color=blue!30,noinline]Zhiyuan:
you didn’t copy the definition correct. 1. do not use $\theta$, give parameter
space some name, e.g., $\mathcal{\Theta}$. 2. the defined algorithm $A$ is not
used. the transformation should be on the updated parameters, but not grad. 3.
We need a theorem saying Adam is permutation equivariant.

  1. 1.

There exists a group $\mathcal{G}_{\theta}$ acting on parameters $U$ and a
group isomorphism $\tau:\mathcal{G}_{\mathcal{X}}\to\mathcal{G}_{\theta}$ such
that for all $x\in\mathcal{X},T\in\mathcal{G}_{\mathcal{X}},U$ we have that
$g(U;x)=g(\tau(T)(U);T(x))$.

  2. 2.

The gradient update rule is invariant under joint group action $(T,\tau(T))$
for all $T\in\mathcal{G}_{\mathcal{X}}$:
$\tau(T)(\nabla_{U}g(U;x))=\nabla_{U}g(\tau(T)(U);T(x))$.

  3. 3.

The initialization distribution $P_{\operatorname{init}}$ is invariant under
$\mathcal{G}_{\mathcal{X}}$.

We now present our permutation-equivariance results.

######  Definition C.3.

We define $\mathcal{G}_{\theta}$ as a group of actions to be applied on
$U=(W,V)$ as

$\mathcal{G}_{\theta}\triangleq\left\\{\pi_{\sigma_{1},\sigma_{2},\sigma_{3}}\mid\sigma_{1},\sigma_{2},\sigma_{3}\in\mathbb{S}_{p}\right\\}$

where $\pi_{\sigma_{1},\sigma_{2},\sigma_{3}}$ is defined as
$\pi_{\sigma_{1},\sigma_{2},\sigma_{3}}(W,V)\triangleq(W^{\prime},V^{\prime})$,
where $\forall
i\in[h],j\in[p],W^{\prime}_{ij}=W_{i,\sigma_{1}(j)},W_{i(j+p)}^{\prime}=W_{i,\sigma_{2}(j)},V^{\prime}_{\sigma_{3}(j),i}=V_{j,i}$.
This means the concatenation of the row is permuted in the same way as the
data does. Furthermore, the rows of $V$ also get permuted according to
$\sigma_{3}$.

######  Remark C.4.

There is a one-to-one mapping $\tau$ between $\mathcal{G}_{\theta}$ and
$\mathcal{G}_{\mathcal{X}}$, which is
$\tau(\sigma_{1},\sigma_{2},\sigma_{3})=\pi_{\sigma_{1},\sigma_{2},\sigma_{3}}$.

######  Lemma C.5.

For every $x=(i,j,k)$ where $i,j,k\in[p]$, $T\in\mathcal{G}_{\mathcal{X}}$ we
have

$g(U;x)=g(\tau(T)(U);T(x)).$

###### Proof.

For each
$x=(i,j,k)\in\mathcal{X};\sigma_{1},\sigma_{2},\sigma_{3}\in\mathbb{S}_{p};U=(W,V)$
we have that

$\displaystyle
g(\tau(\sigma_{1},\sigma_{2},\sigma_{3})(U);(\sigma_{1},\sigma_{2},\sigma_{3})(x))$
$\displaystyle=\left\langle
e_{\sigma_{3}(k)},V^{\prime}_{\sigma_{3}}\left(\begin{bmatrix}W_{1}^{\prime\top}(\sigma_{1}(i),\sigma_{2}(j))\\\
W_{2}^{\prime\top}(\sigma_{1}(i),\sigma_{2}(j))\\\ \vdots\\\
W_{h}^{\prime\top}(\sigma_{1}(i),\sigma_{2}(j))\end{bmatrix}\right)^{\odot
2}\right\rangle$ $\displaystyle=\left\langle
e_{k},V\left(\begin{bmatrix}W_{1,\sigma_{1}(i)}^{\prime}+W_{1,\sigma_{2}(j)+p}^{\prime}\\\
W_{2,\sigma_{1}(i)}^{\prime}+W_{2,\sigma_{2}(j)+p}^{\prime}\\\ \vdots\\\
W_{h,\sigma_{1}(i)}^{\prime}+W_{h,\sigma_{2}(j)+p}^{\prime}\end{bmatrix}\right)^{\odot
2}\right\rangle$ $\displaystyle=\left\langle
e_{k},V\left(\begin{bmatrix}W_{1,i}+W_{1,j+p}\\\ W_{2,i}+W_{2,j+p}\\\
\vdots\\\ W_{h,i}+W_{h,j+p}\end{bmatrix}\right)^{\odot 2}\right\rangle$
$\displaystyle=\left\langle e_{k},V(Wx)^{\odot 2}\right\rangle$
$\displaystyle=g(U;x).$

∎

######  Lemma C.6.

For every $x=(i,j,k)$ where $i,j,k\in[p]$, $T\in\mathcal{G}_{\mathcal{X}}$ we
have

$\tau(T)(\nabla_{U}g(U;x))=\nabla_{U}g(\tau(T)(U);T(x)).$

###### Proof.

We first consider the gradient of the second layer. For all $a,b\in[p]$ and
$\sigma_{1},\sigma_{2},\sigma_{3}\in\mathbb{S}_{p}$:

$\displaystyle\tau^{-1}(\pi_{\sigma_{1},\sigma_{2},\sigma_{3}})\left(\nabla_{V_{b}}g(\tau(\sigma_{1},\sigma_{2},\sigma_{3})(U);(\sigma_{1},\sigma_{2},\sigma_{3})(x))\right)$
$\displaystyle=I(\sigma_{3}(k)=\sigma_{3}(b))\nabla_{V_{b}}g\left(\tau(\sigma_{1},\sigma_{2},\sigma_{3})(U);(\sigma_{1},\sigma_{2},\sigma_{3})(x)\right)$
$\displaystyle=I(k=b)(W^{\prime}(\sigma_{1},\sigma_{2})(x))^{\odot 2}$
$\displaystyle=I(k=b)(Wx)^{\odot 2}$ $\displaystyle=\nabla_{V_{b}}g(U;x).$

For the gradients of the first layer we have:

$\displaystyle\tau^{-1}(\pi_{\sigma_{1},\sigma_{2},\sigma_{3}})\left(\nabla_{W}g(\tau(\sigma_{1},\sigma_{2},\sigma_{3})(U);(\sigma_{1},\sigma_{2},\sigma_{3})(x))\right)$
$\displaystyle=\tau^{-1}(\pi_{\sigma_{1},\sigma_{2},\sigma_{3}})\left(2V_{k}\odot(W^{\prime}_{(\sigma_{1},\sigma_{2})}(e_{\sigma_{1}(i)},e_{\sigma_{2}(j)})^{\top})\;(e_{\sigma_{1}(i)},e_{\sigma_{2}(j)})\right)$
$\displaystyle=\tau^{-1}(\sigma_{1},\sigma_{2},\sigma_{3})\left(2V_{k}\odot(W(e_{i},e_{j}))\;(e_{\sigma_{1}(i)},e_{\sigma_{2}(j)})^{\top}\right)$
$\displaystyle=2V_{k}\odot(Wx)\;(\sigma_{1},\sigma_{2})^{-1}\left((e_{\sigma_{1}(i)},e_{\sigma_{2}(j)})\right)$
$\displaystyle=2V_{k}\odot(W(e_{i},e_{j})^{\top})\;(e_{i},e_{j})$
$\displaystyle=\nabla_{W}g(U;x).$

∎

See 3.3

###### Proof.

We show that training a neural network with gradient descent in our setting
satisfies the three conditions proposed in Theorem C.2 and thus, this theorem
applies to the training process. Equivariance of the forward pass and the
backward pass (update rule) are settled through Lemmas C.5 and C.6. Finally,
it is straightforward to see that the initialization is invariant under
$\mathcal{G}_{\theta}$ defined in Definition C.3. Since the distribution of
initialization is symmetric and each parameter is initialized independently
and identically from the same distribution, the action of swapping rows or
columns doesn’t change the distribution. ∎

See 4.2

###### Proof Sketch.

Consider the permutation group $\mathcal{G}_{\theta}$ defined in Definition
C.3. Note that there is a one-to-one mapping $\kappa$ mapping
$\mathcal{G}_{\mathcal{X},\mathcal{Y}}$ defined in Definition 3.2 to
$\mathcal{G}_{\mathcal{X}}$ and another one-to-one mapping $\tau$ mapping
$\mathcal{G}_{\mathcal{X}}$ to $\mathcal{G}_{\theta}$ where
$\tau(\sigma_{1},\sigma_{2},\sigma_{3})=\pi_{\sigma_{1},\sigma_{2},\sigma_{3}}$.
Since $f$ and $g$ share the parameters, it’s clear that the distribution of
initialization for $f$ is equivariant under $\mathcal{G}_{\theta}$. To see the
equivariance of forward and backward passes on $f(U;\cdot)$ it suffices to see
that Lemmas C.5 and C.6 hold for
($\mathcal{G}_{\mathcal{X},\mathcal{Y}},\mathcal{G}_{\theta}$) and function
$f$ under the isomorphism $\tau\circ\kappa$ since for any parameters $U$ and
inputs $i,j,k\in[p]$ it holds that $f_{k}(U;(i,j))=g(U;(i,j,k))$. ∎

######  Corollary C.7 (Equivariance of Adam).

Theorem 3.3 (and similarly Proposition 4.2) applies to other gradient-based
training algorithms that need memory, such as Adam.

###### Proof Sketch for Corollary C.7.

To prove that other gradient-based algorithms, particularly Adam, are also
equivariant to the permutation group $\mathcal{G}_{\mathcal{X}}$ we just need
to ensure that the update rule of these algorithms is equivariant under the
joint group action $(T,\tau(T))$. First, note that linear operations (such as
weight decay) on gradients and parameters equivariant. To track the momentum
at different steps of the algorithm we can apply an induction on equivariance
of these variables on the step number. At $t=0$ they’re both zero. $m_{t}$ is
a linear combination $m_{t-1}$ and $g_{t}$ which both are equivarint under the
joint action $(T,\tau(T))$. Since the gradient is equivariant, the coordinate-
wise squared gradient is also equivariant and the linear combination of it
with $v_{t_{1}}$ is also equivariant. This settles the equivariance of the
update rule of Adam and other similar gradient-based algorithms that need
memory and buffers. ∎

##  Appendix D Lower Bound of Population Loss for Kernel Methods

In this section we present the formal version of Theorem 3.4 alongside a proof
of it. Note that a kernel-based predictor $h$ on a training data
$\left\\{({\bm{x}}_{i},y_{i})\right\\}_{i=1}^{n}$ can be expressed as
$h(x)=\sum_{i=1}^{n}\lambda_{i}K({\bm{x}}_{i},x)$ where
$\lambda_{i};i\\!\in\\![n]$ are constants. Assuming that the kernel’s feature
maps are of dimension $d$, the predictions are linear combinations of
$d$-dimensional feature maps. We first present a general proof that every
permutation invariant kernel requires $\Omega(p^{2})$ training points to
outperform the null predictor in terms of $\ell_{2}$ loss, and then show that
this theorem applies to the distribution of empirical NTKs at initialization.

###  D.1 Notation

We use $[p]$ to denote the set $\\{1,\ldots,p\\}$. We use $\mathbb{S}_{p}$ to
denote the permutation groups over $p$ elements, ${\mathbb{S}}_{m}$ as
permutation group over $m$ elements and $\mathsf{id}$ is the identity mapping.
For any nonempty set $\mathcal{X}$, a symmetric function
$K:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ is called a positive semi-
definite kernel (p.s.d) kernel on $\mathcal{X}$ if for any $n\in\mathbb{N}$,
any ${\bm{x}}_{1},\ldots,{\bm{x}}_{n}$ and
$\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}$, it holds that
$\sum_{i=1}^{n}\sum_{j=1}^{n}\lambda_{i}\lambda_{j}K({\bm{x}}_{i},{\bm{x}}_{j})\geq
0$. For a subspace $V$ of $\mathbb{R}^{n}$ and vector $x\in\mathbb{R}^{n}$, we
define $\mathsf{dist}(x,V)\triangleq\min_{v\in V}\left\lVert
x-v\right\rVert_{2}$.

For any $p^{m}$-dimensional vector $v\in\mathbb{R}^{p^{m}}$ we denote by
$v(x)$ the vector $v$ indexed by an $m$-dimensional index vector
$x\in[p]^{m}$. We define the vector $s_{i,a}\in\mathbb{R}^{p^{m}}$ whose
entries are

$\displaystyle s_{i,a}(x)\triangleq\begin{cases}1,&x[i]=a;\\\
0,&\text{otherwise.}\end{cases}$ (D.1)

which helps us denote the all-once slices in this vector space,
$V_{s}\triangleq\operatorname{span}\\{s_{i,a}\\}_{i\in[m],a\in[p]}$. We
further define $\Delta(x,x^{\prime})$ where $x,x^{\prime}$ are two index
vectors of size $m$ as the number of equal indices between them, formally:

$\Delta(x,x^{\prime})=\sum_{i=1}^{m}\mathbf{1}_{x[i]=x^{\prime}[i]}.$

We also define $\mathcal{X}_{a,b}$ as the set of all $x,x^{\prime}\in[p]^{b}$
index vector pairs such that $\mathsf{dist}(x,x^{\prime})=a$, formally:

$\mathcal{X}_{a,b}\triangleq\\{(x,x^{\prime}):x,x^{\prime}\in[p]^{b}\wedge\Delta(x,x^{\prime})=a\\}.$

When $b=m$, we drop the second index and write $\mathcal{X}_{a}$ (instead of
$\mathcal{X}_{a,m}$) for simplicity. It’s clear that the collection of all
$\mathcal{X}_{d}$ for $0\leq d\leq m$ is a partitioning of the set of all
pairs of index vectors. Unless stated otherwise, $x$ refers to the
$m$-dimensional index vector. Finally, we define
$\mathbb{U}_{m}\triangleq[\operatorname{Unif}({\mathbb{S}}_{p})]^{m}$ as the
product of $m$ Uniform distribution on ${\mathbb{S}}_{p}$.

###  D.2 Loss Lower Bound: Regression Setting

For convenience of notation, we will use the $(i,j,k)$ and
$e_{i}+e_{j+p}+e_{k+2p}$ interchangebly for $i,j,k\in[p]$. We denote the
function $K({\bm{x}}_{t},\cdot):[p]\times[p]\times[p]\to\mathbb{R}$ as a
tensor on $\mathbb{R}^{p\times p\times p}$ by $v_{t}(\cdot)$ for each
$t\in[n]$. We also define function
$\Psi_{\sigma_{1},\sigma_{2},\sigma_{3}}(i,j,k)\triangleq{\bm{1}}\bigg{(}\sigma_{1}(i)+\sigma_{2}(j)\equiv\sigma_{3}(k)\
(\\!\\!\\!\\!\mod p)\bigg{)}$. We can view a function mapping from
$[p]\times[p]\times[p]\to\mathbb{R}$ as a vector of size $p^{3}$ and define
inner products and $\mathsf{dist}$ on the function space, _i.e._ ,
$\left\langle\Psi,\Psi^{\prime}\right\rangle\triangleq\sum_{i,j,k\in[p]}\Psi(i,j,k)\Psi^{\prime}(i,j,k)$
and $\left\lVert h\right\rVert_{2}^{2}\triangleq\left\langle
h,h\right\rangle$.

######  Theorem D.1.

For any integers $n\geq 1$, $p\geq 2$ and kernel
$K:\left([p]\times[p]\times[p]\right)\times\left([p]\times[p]\times[p]\right)\to\mathbb{R}$,
for any ${\bm{x}}_{t}=(i_{t},j_{t},k_{t})\in[p]^{3}$ for each $t\in[n]$, it
holds that

$\displaystyle\min_{{\bm{x}}_{1},\ldots,{\bm{x}}_{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}(\mathbb{S}_{p})}\inf_{\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}}\left\lVert\sum_{t=1}^{n}\lambda_{t}K({\bm{x}}_{t},\cdot)-\Psi_{\sigma_{1},\sigma_{2},\sigma_{3}}(\cdot)\right\rVert_{2}^{2}\geq
p^{2}\bigg{(}1-\frac{1}{p}-\frac{n}{p^{3}}\exp\left(\frac{2}{p-1}\right)\bigg{)}$
(D.2)

In other words, if $n\leq(1-\Omega(1))p^{3}$, then the expected population
$\ell_{2}$ loss is at least $\Omega(p^{2})$, which is of the same magnitude as
the trivial all-zero predictor.

###### Proof of Theorem D.1.

This Theorem is a direct result of Corollary D.8 for the case where $m=3$. ∎

See 3.4

###### Proof of Theorem 3.4.

Because $\mathcal{A}$ is permutation-equivariant, we have that

$\displaystyle\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\mathcal{L}_{\ell_{2}}\left(\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)\right)$
$\displaystyle=$
$\displaystyle\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\|\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)-p\cdot\Psi_{\mathsf{id},\mathsf{id},\mathsf{id}}\|^{2}/p^{3}$
$\displaystyle=$
$\displaystyle\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\|\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)/p-\Psi_{\mathsf{id},\mathsf{id},\mathsf{id}}\|^{2}/p$
$\displaystyle=$
$\displaystyle\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}(\mathbb{S}_{p})}\operatorname*{\mathbb{E}}_{\mathcal{A}}\|\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)/p-\Psi_{\sigma_{1},\sigma_{2},\sigma_{3}}\|^{2}/p.$

Because $\mathcal{A}$ is a kernel method, we have for any
$({\bm{x}}_{i},y_{i})_{i=1}^{n}$ and
$\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}$

$\|\mathcal{A}\left(\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\right)/p-\Psi_{\sigma_{1},\sigma_{2},\sigma_{3}}\|^{2}\geq\inf_{\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}}\left\lVert\sum_{t=1}^{n}\lambda_{t}K({\bm{x}}_{t},\cdot)-\Psi_{\sigma_{1},\sigma_{2},\sigma_{3}}(\cdot)\right\rVert_{2}^{2}.$

Applying Theorem D.1 completes the proof. ∎

###  D.3 Loss Lower Bound: Classification Setting

For this subsection, we define the modular addition function
$\Psi:[p]\times[p]\to[p]$ as
$\big{[}\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(i,j)\big{]}_{k}\triangleq\mathbf{1}\Big{(}\sigma_{1}(i)+\sigma_{2}(j)\equiv\sigma_{3}(k)\pmod{p}\Big{)}$
for all $k\in[p],\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}$.

######  Theorem D.2.

For any integers $n>1,p\geq 2$ and input-output permutation equivariant kernel
$K:([p]\times[p])\times([p]\times[p])\to\mathbb{R}^{p\times p}$ (according to
Proposition 4.2), suppose
$x_{t}=(i_{t},j_{t})\overset{i.i.d.}{\sim}\mathrm{Unif}([p]\times[p])$ for
each $t\in[n]$ it holds that

$\operatorname*{\mathbb{E}}_{x_{1},\ldots,x_{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2}\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})}\inf_{\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}^{p}}\sum_{x\in[p]\times[p]}\left\lVert\sum_{t=1}^{n}K(x_{t},x)\lambda_{t}-\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right\rVert_{2}^{2}\geq
p^{2}\bigg{(}1-\frac{1}{p}-\frac{n}{p^{2}}\exp\left(\frac{2}{p-1}\right)\bigg{)}.$

In other words, if $n\leq(1-\Omega(1))p^{2}$, then the expected population
$\ell_{2}$ loss is at least $\Omega(p^{2})$, which is of the same magnitude as
the trivial all-zero predictor.

###### Proof.

Note that

$\displaystyle\operatorname*{\mathbb{E}}_{x_{1},\ldots,x_{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})}\inf_{\lambda\in\mathbb{R}^{n\times
p}}\sum_{x\in[p]\times[p]}\left\lVert\sum_{i=1}^{n}K(x_{i},x)\lambda_{i}-\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right\rVert_{2}^{2}$
$\displaystyle\qquad=\operatorname*{\mathbb{E}}_{x_{1},\ldots,x_{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2}\sim\mathrm{Unif}({\mathbb{S}}_{p})}\inf_{\lambda\in\mathbb{R}^{n\times
p}}\sum_{x\in[p]\times[p]}\sum_{j=1}^{p}\left(\sum_{i=1}^{n}\left\langle\left[K(x_{i},x)\right]_{j,:},\lambda_{i}\right\rangle-\left[\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right]_{j}\right)^{2}$
$\displaystyle\qquad=\operatorname*{\mathbb{E}}_{x_{1},\ldots,x_{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})}\inf_{\lambda\in\mathbb{R}^{n\times
p}}\sum_{x\in[p]\times[p]}\sum_{j=1}^{p}\left(\sum_{i=1}^{n}\sum_{k=1}^{p}\lambda_{i,k}\left[K(x_{i},x)\right]_{j,k}-\left[\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right]_{j}\right)^{2}.$
(D.3)

We define

$\displaystyle
K^{\prime}\big{(}(i,j,k),({i^{\prime}},{j^{\prime}},{k^{\prime}})\big{)}\triangleq
K\big{(}(i,j),({i^{\prime}},{j^{\prime}})\big{)}_{k,k^{\prime}},$
$\displaystyle\Psi^{\prime}_{\sigma_{1},\sigma_{2},\sigma_{3}}\big{(}(i,j,k)\big{)}\triangleq\left[\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}\big{(}(i,j)\big{)}\right]_{k},$
$\displaystyle\lambda^{\prime}_{t}\triangleq\lambda_{\lfloor t/p\rfloor,t\mod
p},$ $\displaystyle x^{\prime}_{t}\triangleq(i_{\lfloor t/p\rfloor},j_{\lfloor
t/p\rfloor},t\mod p)$ (D.4)

for any $i,j,k,i^{\prime},j^{\prime},k^{\prime}\in[p]$,
$\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}$, $t\in[np]$,
$\lambda\in\mathbb{R}^{n\times p}$ and $\lambda^{\prime}\in\mathbb{R}^{np}$.
It can be seen that

$\displaystyle\operatorname*{\mathbb{E}}_{\begin{subarray}{c}x_{1},\ldots,x_{n}\\\
\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})\end{subarray}}\inf_{\lambda\in\mathbb{R}^{n\times
p}}\sum_{x\in[p]\times[p]}\sum_{j=1}^{p}\left(\sum_{i=1}^{n}\sum_{k=1}^{p}\lambda_{i,k}\left[K(x_{i},x)\right]_{j,k}-\left[\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right]_{j}\right)^{2}$
$\displaystyle\operatorname*{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}_{1},\ldots,x^{\prime}_{np}\\\
\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})\end{subarray}}\inf_{\lambda^{\prime}\in\mathbb{R}^{np}}\sum_{x\in[p]\times[p]}\sum_{j=1}^{p}\left(\sum_{i=1}^{np}\lambda^{\prime}_{i}K^{\prime}\big{(}x^{\prime}_{i},(x[0],x[1],i\mod
p)\big{)}-\Psi^{\prime}_{\sigma_{1},\sigma_{2},\sigma_{3}}(x)\right)^{2}.$
(D.5)

This is similar to Theorem D.1 except that the underlying data is generated
through sampling from independent groups each having a size of $p$ (we sample
i.i.d from $\mathrm{Unif}([p]\times[p])$ and for each sample we consider the
set of all possible responses). Applying the result of Theorem D.1 yields

$\displaystyle\operatorname*{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}_{1},\ldots,x^{\prime}_{np}\\\
\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}({\mathbb{S}}_{p})\end{subarray}}\inf_{\lambda^{\prime}\in\mathbb{R}^{np}}\sum_{x\in[p]\times[p]}\sum_{j=1}^{p}\left(\sum_{i=1}^{np}\lambda^{\prime}_{i}K^{\prime}\big{(}x^{\prime}_{i},(x[0],x[1],i\mod
p)\big{)}-\Psi^{\prime}_{\sigma_{1},\sigma_{2},\sigma_{3}}(x)\right)^{2}$
$\displaystyle\qquad\qquad\geq
p^{2}\left(1-\frac{1}{p}-\frac{n}{p^{2}}\exp\left(\frac{2}{p-1}\right)\right)$
(D.6)

which completes the proof. ∎

See 4.3

###### Proof of Theorem 4.3.

Because $\mathcal{A}$ is input-output permutation-equivariant, we have that

$\displaystyle\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\mathcal{L}_{\ell_{2}}\left(\mathcal{A}\big{(}\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\big{)}\right)$
$\displaystyle\qquad=\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\sum_{x\in[p]\times[p]}\left\|\mathcal{A}\big{(}\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\big{)}(x)-p\cdot\Psi^{\mathsf{id}}_{\mathsf{id},\mathsf{id}}(x)\right\|_{2}^{2}/p^{2}$
$\displaystyle\qquad=\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\mathcal{A}}\sum_{x\in[p]\times[p]}\left\|\mathcal{A}\big{(}\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\big{)}(x)/p-\Psi^{\mathsf{id}}_{\mathsf{id},\mathsf{id}}(x)\right\|_{2}^{2}$
$\displaystyle\qquad=\operatorname*{\mathbb{E}}_{({\bm{x}}_{i},y_{i})_{i=1}^{n}\sim\mathcal{D}^{n}}\operatorname*{\mathbb{E}}_{\sigma_{1},\sigma_{2},\sigma_{3}\sim\mathrm{Unif}(\mathbb{S}_{p})}\operatorname*{\mathbb{E}}_{\mathcal{A}}\sum_{x\in[p]\times[p]}\left\|\mathcal{A}\big{(}\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\big{)}(x)/p-\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right\|_{2}^{2}.$

Because $\mathcal{A}$ is a kernel method, we have for any
$({\bm{x}}_{i},y_{i})_{i=1}^{n}$ and
$\sigma_{1},\sigma_{2},\sigma_{3}\in{\mathbb{S}}_{p}$

$\sum_{x\in[p]\times[p]}\left\|\mathcal{A}\big{(}\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}\big{)}(x)/p-\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right\|_{2}^{2}\geq\inf_{\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}}\sum_{x\in[p]\times[p]}\left\lVert\sum_{t=1}^{n}\lambda_{t}K({\bm{x}}_{t},x)-\Psi^{\sigma_{3}}_{\sigma_{1},\sigma_{2}}(x)\right\rVert_{2}^{2}.$

Applying Theorem D.2 completes the proof. ∎

###  D.4 Loss Lower Bound: General Theorem For Arbitrary Functions

######  Lemma D.3.

For any subspace $V$ of $\mathbb{R}^{n}$ and vector $x\in\mathbb{R}^{n}$, let
$\\{v_{i}\\}_{i=1}^{m}$ be an orthonormal basis of $V$. It holds that
$\mathsf{dist}^{2}(x,V)=\left\lVert
x\right\rVert_{2}^{2}-\sum_{i=1}^{m}\left\langle x,v_{i}\right\rangle^{2}$.

The proof of Lemma D.3 is straightforward and thus omitted.

######  Lemma D.4.

For any distribution $\mathcal{D}$ over functions mapping from
$\mathcal{X}\to\mathbb{R}$ and $n$ functions $\\{k_{i}\\}_{i=1}^{n}$ where
$k_{i}:\mathcal{X}\to\mathbb{R}$ for each $i\in[n]$ it holds that

$\operatorname*{\mathbb{E}}_{h\sim\mathcal{D}}\left[\min_{\alpha\in\mathbb{R}^{n}}\frac{1}{\lvert\mathcal{X}\rvert}\sum_{x\in\mathcal{X}}\left(\Psi(x)-\sum_{i=1}^{n}\alpha_{i}k_{i}(x)\right)^{2}\right]\geq\frac{1}{\lvert\mathcal{X}\rvert}\sum_{i=n+1}^{\lvert\mathcal{X}\rvert}\lambda_{i}(\Sigma).$

where
$\Sigma_{\mathcal{D}}(x,x^{\prime})\triangleq\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\Psi(x)\Psi(x^{\prime})\right]$
and $\lambda_{i}$ denotes the $i$’th largest eigenvalue function. For
notational convenience, we also view $\Sigma_{\mathcal{D}}$ as a
$\mathbb{R}^{\lvert\mathcal{X}\rvert\times\lvert\mathcal{X}\rvert}$ matrix.

###### Proof of Lemma D.4.

For each $\Psi$, we define $r_{\Psi}\in\mathbb{R}^{\lvert\mathcal{X}\rvert}$
whose entries are realization of the function $\Psi$ on inputs
$x\in\mathcal{X}$. It suffices to show that

$\operatorname*{\mathbb{E}}_{h\sim\mathcal{D}}\left[\mathsf{dist}^{2}(r_{\Psi},V)\right]\geq\sum_{i=n+1}^{\lvert\mathcal{X}\rvert}\lambda_{i}(\Sigma_{\mathcal{D}})$

for any subspace
$V=\operatorname{span}\\{v_{1},v_{2},\cdots,v_{n}\\}\subset\mathbb{R}^{\lvert\mathcal{X}\rvert\times\lvert\mathcal{X}\rvert}$
where $v_{i}$ for $i\in[n]$ are orthonormal vectors. Note that

$\displaystyle\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\mathsf{dist}^{2}(r_{\Psi},V)\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\left\lVert
r_{\Psi}\right\rVert_{2}^{2}-\sum_{t=0}^{n}\left\langle
v_{t},r_{\Psi}\right\rangle^{2}\right]$
$\displaystyle=\operatorname{Tr}\left(\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\Psi\Psi^{\top}\right]\right)-\sum_{t=1}^{n}v_{t}^{\top}\left(\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\Psi\Psi^{\top}\right]\right)v_{t}$
$\displaystyle\geq\operatorname{Tr}\left(\Sigma_{\mathcal{D}}\right)-\sum_{i=1}^{n}\lambda_{i}\left(\Sigma_{\mathcal{D}}\right)$
$\displaystyle=\sum_{i=n+1}^{\lvert\mathcal{X}\rvert}\lambda_{i}(\Sigma_{\mathcal{D}})$
(D.7)

where the inequality in the second to last line is due to the min-max theorem
(also called Courant–Fischer–Weyl min-max principle) [11]. This completes the
proof. ∎

The following Corollary D.5 is a direct consequence of Lemma D.4, noting that
$\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\frac{1}{\lvert\mathcal{X}\rvert}\sum_{x\in\mathcal{X}}\Psi(x)^{2}\right]=\operatorname{Tr}(\Sigma_{\mathcal{D}})$.

######  Corollary D.5.

For any distribution $\mathcal{D}$ over functions mapping from
$\mathcal{X}\to\mathbb{R}$ and $n$ functions $\\{k_{i}\\}_{i=1}^{n}$ where
$k_{i}:\mathcal{X}\to\mathbb{R}$ for each $i\in[n]$, if
$\sum_{i=1}^{n}\lambda_{i}(\Sigma_{\mathcal{D}})\leq\frac{1}{2}\operatorname{Tr}(\Sigma_{\mathcal{D}})$
then it is guaranteed that

$\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\min_{\alpha\in\mathbb{R}^{n}}\frac{1}{\lvert\mathcal{X}\rvert}\sum_{x\in\mathcal{X}}\left(\Psi(x)-\sum_{i=1}^{n}\alpha_{i}k_{i}(x)\right)^{2}\right]\geq\frac{1}{2}\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\frac{1}{\lvert\mathcal{X}\rvert}\sum_{x\in\mathcal{X}}\Psi(x)^{2}\right]$

where the right-hand side denotes the expected loss of the all-0 predictor.

######  Lemma D.6.

For any matrix $\Sigma\in\mathbb{R}^{d\times d}$, subspace $V$ and projection
matrix $P_{V}\in\mathbb{R}^{d\times d}$ corresponding to $V$, it holds that

$\sum_{i=1}^{n}\lambda_{i}(\Sigma)\leq
n\cdot\lambda_{1}\bigg{(}(I-P_{V})\Sigma(I-P_{V})\bigg{)}+\operatorname{Tr}\left(P_{V}\Sigma
P_{V}\right).$

###### Proof.

Let $V_{n}=\operatorname{span}\\{\alpha_{i}\\}_{i=1}^{n}$ where
$\alpha_{i}\in\mathbb{R}^{d}$ is the $i$’th eigenvector of $\Sigma$ and let
$P_{V_{n}}$ be its corresponding projection matrix. Let $P_{V_{n}+V}$ be the
projection onto the sum of two subspaces $V_{n}+V$, it holds that

$\displaystyle\sum_{i=1}^{n}\lambda_{i}(\Sigma)$
$\displaystyle=\operatorname{Tr}(P_{V_{n}}\Sigma P_{V_{n}})$
$\displaystyle\leq\operatorname{Tr}(P_{V_{n}+V}\Sigma P_{V_{n}+V})$
$\displaystyle=\operatorname{Tr}(P_{V}\Sigma
P_{V})+\operatorname{Tr}\big{(}(P_{V_{n}+V}-P_{V})\Sigma(P_{V_{n}+V}-P_{V})\big{)}$
$\displaystyle\leq\operatorname{Tr}(P_{V}\Sigma
P_{V})+n\cdot\lambda_{1}\big{(}(P_{V_{n}+V}-P_{V})\Sigma(P_{V_{n}+V}-P_{V})\big{)}$
$\displaystyle\leq\operatorname{Tr}(P_{V}\Sigma
P_{V})+n\cdot\lambda_{1}\big{(}(I-P_{V})\Sigma(I-P_{V})\big{)}.$ (D.8)

Here the second to last inequality is because $P_{V_{n}+V}-P_{V}$ is at most
rank-$n$ and so is $(P_{V_{n}+V}-P_{V})\Sigma(P_{V_{n}+V}-P_{V})$. The last
inequality is because $P_{V_{n}+V}\leq I$ and thus
$(P_{V_{n}+V}-P_{V})\Sigma(P_{V_{n}+V}-P_{V})\preceq(I-P_{V})\Sigma(I-P_{V})$.
∎

###  D.5 Loss Lower Bound: Modular Addition with $m$ Summands

######  Lemma D.7.

Let $\mathcal{D}$ be the uniform distribution over

$\displaystyle\mathcal{H}\triangleq\left\\{\Psi(x)=\mathbf{1}\bigg{(}\sum_{i=1}^{m}\sigma_{i}(x_{i})\equiv
0\;(\mod p)\bigg{)}\,\Bigg{|}\;\sigma_{i}\in{\mathbb{S}}_{p}\text{ for all
}i\in[m]\right\\}$ (D.9)

and
$\Sigma_{\mathcal{D}}(x,x^{\prime})=\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\Psi(x)\Psi(x^{\prime})\right]$.
It holds that

$\displaystyle\sum_{i=1}^{n}\lambda_{i}(\Sigma)\leq
p^{m-2}+\frac{n}{p}\exp\left(\frac{m-1}{p-1}\right).$

###### Proof of Lemma D.7.

Consider the vector space $V_{s}$ defined in Equation D.1 and the projection
matrix
$P_{V_{s}}\in\mathbb{R}^{\lvert\mathcal{X}\rvert\times\lvert\mathcal{X}\rvert}$
onto $V_{S}$. To find $\operatorname{Tr}\big{(}P_{V_{s}}\Sigma
P_{V_{s}}\big{)}$ we can first derive
$\operatorname{Tr}\big{(}P_{V_{s}}\Psi\Psi^{\top}P_{V_{s}}\big{)}$ where
$\Psi\in\mathbb{R}^{\lvert\mathcal{X}\rvert}$ is a sample from $\mathcal{D}$.
Note that since $(I-P_{V_{s}})h$ should be orthogonal to $V_{s}$ (the sum of
each slice of the projected vector should be zero) we have that

$\Big{[}(I-P_{V_{s}})\Psi\Big{]}(x)=\begin{cases}-\frac{1}{p}&\Psi(x)=0\\\
\frac{p-1}{p}&\Psi(x)=1\end{cases}.$

Hence, $P_{V_{s}}\Psi(x)=\frac{1}{p}$ for all $x\in\mathcal{X}$. Thus,

$\displaystyle\operatorname{Tr}\bigg{(}P_{V_{s}}\Sigma P_{V_{s}}\bigg{)}$
$\displaystyle=\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\operatorname{Tr}\bigg{(}P_{V_{s}}\Psi\Psi^{\top}P_{V_{s}}\bigg{)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\left[\operatorname{Tr}\bigg{(}\frac{1}{p^{2}}\mathbf{1}_{\lvert\mathcal{X}\rvert\times\lvert\mathcal{X}\rvert}\bigg{)}\right]$
$\displaystyle=\frac{\lvert\mathcal{X}\rvert}{p^{2}}$ $\displaystyle=p^{m-2}$
(D.10)

where $\mathbf{1}_{\lvert\mathcal{X}\rvert\times\lvert\mathcal{X}\rvert}$
denotes the all-one square matrix of size $\lvert\mathcal{X}\rvert$. Moreover,
by Lemma D.9 we have that

$\displaystyle
n\cdot\lambda_{1}\bigg{(}(I-P_{V_{s}})\Sigma(I-P_{V_{s}})\bigg{)}=\sup_{\left\lVert
v\right\rVert_{2}\leq 1,v\perp
V_{s}}v^{\top}\Sigma_{\mathcal{D}}v\leq\frac{n}{p}\exp\left(\frac{m-1}{p-1}\right).$
(D.11)

Combining the two equations above, we can see that

$\sum_{i=1}^{n}\lambda_{i}(\Sigma)\leq
p^{m-2}+\frac{n}{p}\exp\left(\frac{m-1}{p-1}\right).$

This concludes the proof. ∎

######  Corollary D.8.

Consider the function class $\mathcal{H}$ defined in Equation D.9 and the
uniform distribution over it denoted by $\mathcal{D}$. For any integers $n\geq
1,p\geq 2$, $1\leq m<p$, permutation-equivariant kernel
$K:[p]^{m}\times[p]^{m}\to\mathbb{R}$ it holds that

$\min_{x_{1},x_{2},\cdots
x_{n}}\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\inf_{\alpha\in\mathbb{R}^{n}}\left\lVert\sum_{t=1}^{n}\lambda_{t}K(x_{t},\cdot)-\Psi\right\rVert_{2}^{2}\geq
p^{m-1}\bigg{(}1-\frac{1}{p}-\frac{n}{p^{m}}\exp\left(\frac{m-1}{p-1}\right)\bigg{)}.$

for any $x_{1},x_{2},\cdots,x_{n}\in[p]^{m}$. In other words, if
$n<(1-\Omega(1))p^{m}$, then the expected population $\ell_{2}$ loss is at
least $\Omega(p^{m-1})$, which is of the same magnitude as the trivial all-
zero predictor.

###### Proof of Corollary D.8.

It suffices to show that for any $n$-dimensional subspace
$V\subset\mathbb{R}^{p^{m}}$ it holds that

$\operatorname*{\mathbb{E}}_{\Psi\sim\mathcal{D}}\mathsf{dist}^{2}(V,\Psi)\geq
p^{m-1}\bigg{(}1-\frac{1}{p}-\frac{n}{p^{m}-1}\exp\left(\frac{m-1}{p-1}\right)\bigg{)}.$

Combining Lemma D.7 and Corollary D.5 yields this statement. ∎

######  Lemma D.9.

For any $v\in{\mathbb{R}}^{p^{m}}$ such that $\left\lVert v\right\rVert=1$ and
$v\perp V_{s}$ (defined in Equation D.1), it holds that

$\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\left\langle\Psi_{\sigma},v\right\rangle^{2}\right]\leq\frac{1}{p}\exp\left(\frac{m-1}{p-1}\right).$

The main implication of this Lemma D.9 is that for any $n$-dimensional space
$V\subset\mathbb{R}^{p^{m}}$ can not "cover" the vector space of all
$\Psi_{\sigma}$ functions for different permutations
$\sigma\in{\mathbb{S}}_{p}$. To prove this Lemma, we decompose the inner
product $\left\langle v,h_{\sigma}\right\rangle$ to $d+1$ sums using the
Multinomial Theorem. This decomposition is enabled through observing the fact
that there are $d+1$ equivalence groups in the possible set of indices of $v$.
Based on Lemma D.10, we can use the Binomial Theorem to decompose the
expectation of the inner product as follows

$\displaystyle\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\left\langle
v,\Psi_{\sigma}\right\rangle^{2}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\left(\sum_{x}v(x)\Psi_{\sigma}(x)\right)^{2}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\sum_{d=0}^{m}\sum_{\begin{subarray}{c}x,x^{\prime}\\\
\mathsf{dist}(x,x^{\prime})=d\end{subarray}}\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})v(x)v(x^{\prime})\right]$
$\displaystyle=\sum_{d=0}^{m}C_{d}\sum_{\begin{subarray}{c}x,x^{\prime}\\\
\mathsf{dist}(x,x^{\prime})=d\end{subarray}}v(x)v(x^{\prime})$ (D.12)

where
$C_{d}\triangleq\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})\right]$
for any $(x,x^{\prime})\in\mathcal{X}_{d}$.

To complete the proof, we now have to bound two terms. First, we need to show
that the expectation
$\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})\right]$
for each $(x,x^{\prime})$ in the same equivalence group is bounded. Next, we
need to show that for each set $\mathcal{X}_{d}$, the sum
$\sum_{\begin{subarray}{c}(x,x^{\prime})\in\mathcal{X}_{d}\end{subarray}}v(x)v(x^{\prime})$
is also bounded. These are correspondingly shown in Lemmas D.10 and D.11.
Based on these two Lemmas, we can now present the proof of Lemma D.9.

###### Proof of Lemma D.9.

$\displaystyle\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\left\langle
v,\Psi_{\sigma}\right\rangle^{2}\right]$
$\displaystyle=\sum_{d=0}^{m}C_{d}\sum_{\begin{subarray}{c}x,x^{\prime}\\\
\mathsf{dist}(x,x^{\prime})=d\end{subarray}}v(x)v(x^{\prime})$
$\displaystyle=\sum_{d=0}^{m}\frac{1}{p^{2}}\left(1-\frac{1}{(1-p)^{d-1}}\right)(-1)^{d}{m\choose
d}$ $\displaystyle=\frac{1}{p}\left(1+\frac{1}{p-1}\right)^{m-1}$
$\displaystyle\leq\frac{1}{p}\exp\left(\frac{m-1}{p-1}\right).$ (D.13)

where the second to last step is due to the binomial Theorem and the last step
is due to the fact that for all $x\in\mathbb{R}$, $1+x\leq\exp(x)$.

∎

######  Lemma D.10.

For any index vector pair $(x,x^{\prime})\in\mathcal{X}_{d}$ it holds that

$\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})\right]=\frac{1}{p^{2}}\left(1-\frac{1}{(1-p)^{d-1}}\right).$

###### Proof of Lemma D.10.

Let us re-iterate the definition of $h$:

$\Psi_{\sigma}(x)=\mathbf{1}\bigg{(}\sum_{i=1}^{m}\sigma_{i}(x_{i})\equiv
0\;(\mod p)\bigg{)}.$

For each pair $(x^{\prime},x^{\prime})$ we define
$I(x,x^{\prime})\triangleq\\{i:x[i]\neq x^{\prime}[i]\\}$ and
$E(x,x^{\prime})\triangleq\\{i:x[i]=x^{\prime}[i]\\}$. We also define

$C_{d}\triangleq\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})\right]$

where for each $(x,x^{\prime})\in\mathcal{X}_{d}$. Note that for $d=0$, we
have that

$\displaystyle C_{0}$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\Psi_{\sigma}(x)\Psi_{\sigma}(x^{\prime})\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim{\mathbb{U}}_{m}}\left[\Psi_{\sigma}(x)\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim{\mathbb{U}}_{m}}\left[\mathbf{1}\bigg{(}\sum_{i=1}^{m}\sigma_{i}(x_{i})\equiv
0\pmod{p}\bigg{)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{a\sim\mathrm{Unif}([p])}\left[a\equiv
0\pmod{p}\right]$ $\displaystyle=1/p.$ (D.14)

For any $2\leq d\leq m$ it holds that

$\displaystyle C_{d}$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\mathbf{1}\bigg{(}\sum_{i=1}^{m}\sigma_{i}(x_{i})\equiv\sum_{i=1}^{m}\sigma_{i}(x^{\prime}_{i})\equiv
0\;(\mod p)\bigg{)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\mathbf{1}\bigg{(}\sum_{i\in
I(x,x^{\prime})}\sigma_{i}(x_{i})+\sum_{i\in
E(x,x^{\prime})}\sigma_{i}(x_{i})\equiv\sum_{i\in
I(x,x^{\prime})}\sigma_{i}(x^{\prime}_{i})+\sum_{i\in
E(x,x^{\prime})}\sigma_{i}(x_{i})\equiv 0\;(\mod p)\bigg{)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{\sigma\sim\mathbb{U}_{m}}\left[\mathbf{1}\bigg{(}\sum_{i\in
I(x,x^{\prime})}\sigma_{i}(x_{i})\equiv\sum_{i\in