# Asymptotics of feature learning in two-layer networks
after one gradient-step

Hugo Cui  Statistical Physics Of Computation lab., École Polytechnique
Fédérale de Lausanne (EPFL), 1015 Lausanne, Switzerland Luca Pesce
Information Learning & Physics lab., École Polytechnique Fédérale de Lausanne
(EPFL), 1015 Lausanne, Switzerland  Yatin Dandi  Information Learning &
Physics lab., École Polytechnique Fédérale de Lausanne (EPFL), 1015 Lausanne,
Switzerland  Florent Krzakala  Information Learning & Physics lab., École
Polytechnique Fédérale de Lausanne (EPFL), 1015 Lausanne, Switzerland  Yue M.
Lu  Harvard University, School of Engineering and Applied Sciences
Lenka Zdeborová  Statistical Physics Of Computation lab., École Polytechnique
Fédérale de Lausanne (EPFL), 1015 Lausanne, Switzerland Bruno Loureiro
Département d’Informatique, École Normale Supérieure (ENS) - PSL & CNRS,
F-75230 Paris cedex 05, France

###### Abstract

In this manuscript we investigate the problem of how two-layer neural networks
learn features from data, and improve over the kernel regime, after being
trained with a single gradient descent step. Leveraging a connection from [1]
with a non-linear spiked matrix model and recent progress on Gaussian
universality [2], we provide an exact asymptotic description of the
generalization error in the high-dimensional limit where the number of samples
$n$, the width $p$ and the input dimension $d$ grow at a proportional rate. We
characterize exactly how adapting to the data is crucial for the network to
efficiently learn non-linear functions in the direction of the gradient –
where at initialization it can only express linear functions in this regime.
To our knowledge, our results provides the first tight description of the
impact of feature learning in the generalization of two-layer neural networks
in the large learning rate regime $\eta=\Theta_{d}(d)$, beyond perturbative
finite width corrections of the conjugate and neural tangent kernels.

##  1 Introduction

A common deep learning intuition behind the unreasonable effectiveness of
neural networks is their capacity to effectively adapt to the training data
which makes them superior to kernel methods. While kernel methods and their
finite width approximations are known to be data hungry (see e.g. [3]), neural
networks have proven themselves to be flexible and efficient in practice.
Their adaptivity and capacity to learn features from data is behind the
success in efficiently solving problems from image classification to text
generation. A large part of our current theoretical understanding of neural
networks stems from the investigation of their lazy regime where features are
not learned during training. This includes a set of works that falls under the
umbrella of Gaussian processes [4, 5], the Neural Tangent Kernel (NTK) [6] and
the Lazy regime [7]. A crucial question in the theoretical machine learning
community is thus to characterize the advantages of two-layer neural networks
beyond these convex optimization approaches.

Different theoretical works have offered sharp separation results between
kernel and feature learning regimes (see e.g. [8, 9, 10, 11]). In particular,
[1, 10], and later [2], discussed the advantage of neural networks when
training with only one single step of large-batch gradient descent with a
large learning rate. Specifically, [1] highlighted that the weight matrix
after the first training step can be decomposed in a bulk plus a rank-one
spike, effectively mapping the learned features to a _spiked Random Features
model_ (sRF), defined in eq. (6). This observation has fueled many further
studies on the effect of spiked structure, see e.g. [2, 12, 13, 14].

In this paper, we follow this line of work, and provide an exact high-
dimensional description of the test error achieved by a two-layer network for
after a single, large, gradient step. To our knowledge, our work provides the
first sharp asymptotic treatment of a model where feature learning is
accounted for in the high-dimensional, non-perturbative large learning rate
regime $\eta=\Theta_{d}(d)$, quantitatively illustrating the benefits of
feature learning over the lazy regime.

Our main contributions in this paper are the following:

  * •

Exact asymptotics for two-layer networks after one gradient step – We provide
a sharp asymptotic characterization of the test error, alongside a set of
summary statistics, for two-layers neural networks with first layer weights
trained with a large learning rate gradient step. The derivation leverages the
replica method from statistical physics [15, 16], and provides a set of scalar
self-consistent equations for the generalization error.

  * •

Conditional Gaussian Equivalence – Building upon [2], we show (and provide
strong numerical evidence to support) that the learning properties of the
spiked Random Features model are asymptotically equivalent to a simple
_conditional_ Gaussian model in the high-dimensional proportional regime. The
conditional Gaussian distribution is characterized by the projections of the
input data on the spike in the weight matrix. This mapping constitutes the
extension of related theoretical results that unveiled a similar Gaussian
equivalence property for the training and generalization error for non-spiked
vanilla RFs [17, 18, 19, 20, 21, 22, 23, 24, 25, 26].

  * •

Feature learning – We provide an extensive discussion on how feature learning
leads to a drastic improvement on the generalization performance over random
features in a data limited regime, demonstrating a clear and quantitative
separation with respect to kernel method and random feature models. In
particular, we derive both upper and lower bounds on the generalization error
and discuss under which conditions they are tight.

The code used in the present work is available in this repository.

##  2 Setting, motivation and related work

We study fully-connected two-layer networks

$\displaystyle
f_{W,\bm{a}}(\bm{x})=\frac{1}{\sqrt{p}}\sum\limits_{i=1}^{p}a_{i}\sigma(\bm{w}_{i}^{\top}\bm{x})\,,$
(1)

and their capacity to learn a single-index target function of isotropic
Gaussian covariates:

$\displaystyle
f_{\star}(x)=\sigma_{\star}(\nicefrac{{\bm{\theta}^{\top}\bm{x}}}{{\sqrt{d}}}),\quad\bm{x}\sim\mathcal{N}(0,I_{d})$
(2)

from finite batch ${\mathcal{D}=\\{(\bm{x}^{\mu},y^{\mu})_{\mu=1}^{n}\\}}$ of
$n$ independently drawn training samples. We consider a layer-wise training
procedure where the first layer weights $W\in\mathbb{R}^{p\times d}$ are
trained for a single gradient step:

$\displaystyle\bm{w}^{(1)}_{i}=\bm{w}^{(0)}_{i}-\eta\bm{g}_{i}$ (3)
$\displaystyle\bm{g}_{i}=\frac{1}{\sqrt{p}}\sum\limits_{\mu=1}^{n_{0}}(y^{\mu}-f_{W^{(0)},\bm{a}^{(0)}}(\bm{x}^{\mu}))a_{i}^{(0)}\sigma^{\prime}(\bm{w}^{(0)}_{i}\cdot\bm{x}^{\mu})\bm{x}^{\mu}$

on a subset $\mathcal{D}_{0}\subset\mathcal{D}$ of size $n_{0}$, where
$(W^{(0)},\bm{a}^{(0)})$ denote the initial weights and $\eta\\!>\\!0$ the
learning rate. For simplicity, we assume
$\bm{a}^{(0)}\\!=\\!\nicefrac{{\bm{1}_{p}}}{{\sqrt{p}}}$ (uniform
initialization) and $\bm{w}_{i}^{(0)}$ with unit norm
$\lVert\bm{w}^{(0)}_{i}\lVert=\\!1$ and weak correlation
${\bm{w}^{(0)}_{i}}\cdot\bm{w}^{(0)}_{j}=O_{d}(\nicefrac{{\mathrm{polylog}(d)}}{{\sqrt{d}}})$
for $i\neq j$ (e.g. uniformly drawn from the unit sphere $\mathbb{S}^{d-1}$).
Given the updated weights $W^{(1)}$, we train the read-out layer on the
remaining data $\mathcal{D}_{1}\\!=\\!\mathcal{D}\setminus\mathcal{D}_{0}$:

$\displaystyle\hat{\bm{a}}_{\lambda}\\!=\\!\underset{\bm{a}\in\mathbb{R}^{p}}{\mathrm{argmin}}\frac{1}{2}\sum\limits_{\mu=1}^{n_{1}}\left(y^{\mu}-f_{W^{(1)},\bm{a}}(\bm{x}^{\mu})\right)^{2}\\!+\\!\frac{\lambda}{2}\lVert\bm{a}\lVert^{2}.$
(4)

with $\lambda\in\mathbb{R}_{+}$ being a regularization parameter. Note that
the layer-wise training procedure considered here is commonly studied in the
theoretical machine learning literature [1, 10, 11, 27, 2, 14] due to its
mathematical tractability.

Our main goal in the following is to provide a tight asymptotic
characterization of the generalization error:

$\displaystyle\epsilon_{g}=\mathbb{E}_{\mathcal{D},\bm{x}}\left(f_{\star}(\bm{x})-f_{W^{(1)},\hat{\bm{a}}_{\lambda}}(\bm{x})\right)^{2}.$
(5)

in the high-dimensional proportional limit where
$n_{0},n_{1},d,p,\eta\\!\to\\!\infty$ at fixed ratios
$\alpha_{0}\\!=\\!\nicefrac{{n_{0}}}{{d}}$,
${\alpha\\!=\\!\nicefrac{{n_{1}}}{{d}}},\beta\\!=\\!\nicefrac{{p}}{{d}},\tilde{\eta}\\!=\\!\nicefrac{{\eta}}{{d}}$.

### Motivation

Driven by the lazy-training regime of learning of large-width networks [7], a
large body of literature has been dedicated to the particular case where the
first layer weights are fixed at initialization $W^{(0)}$ ($\eta=0$), also
known as the Random Features (RF) model [28]. In particular, [8, 29, 30] have
shown that with $n_{1}=\Theta_{d}(d)$ samples
$f_{W^{(0)},\bm{\hat{}}{\bm{a}}_{\lambda}}$ can only approximate, at best, a
linear function of $\bm{x}^{\mu}$, with the non-linear part playing a role
akin to additive label noise. This is a strong limitation: RF network requires
polynomials number of data and neurons to fit a simple polynomial [30, 31]. It
is one of our motivation here to discuss how, with a single gradient steps,
these limitations are lifted.

Behind the scenes in this effective linearity of random features is a Gaussian
Equivalent Principle (GEP) [17, 22, 20, 32, 26], which states that in this
regime the random feature map $\varphi=\sigma(W^{(0)}\bm{x})$ is statistically
equivalent to a rescaled stochastic linear map
$\varphi^{g}(\bm{x})\asymp\mu_{0}\bm{1}+\mu_{1}W^{(0)}\bm{x}+\mu_{\star}\bm{z}$,
with $\bm{z}\sim\mathcal{N}(0,I_{p})$. This surprising universality result
allows to go beyond lower bounds for the generalization performance, making
the problem amenable to a tight high-dimensional characterization of all
relevant statistics in these models [22, 19, 21].

Figure 1: Numerical estimation of the function
$f(x_{\theta})=\mathbb{E}_{x}[f_{W^{(1),\hat{\bm{a}}_{\lambda}}}|\nicefrac{{\bm{\theta}^{\top}\bm{x}}}{{\sqrt{d}}}=x_{\theta}]$
implemented by the trained network (1) in the direction spanned by the weights
$\bm{\theta}$ of the target function. The activations are
$\sigma=\sigma_{\star}=\tanh$, and simulations were run in dimensions
$d=p=2000$, for a learning rate $\eta=2.5p$, and a readout regularization
$\lambda=0.01$. The readout was trained with $n_{1}=2d$ samples. Different
colors corresponds to different sample complexities
$\alpha_{0}\equiv\nicefrac{{n_{0}}}{{d}}$ used to implement the gradient step
on the first layer weights, with $\alpha_{0}=0$ corresponding to not
implementing the step.

Figure 1 illustrates this fundamental limitation of RF models contrasted with
the function
$f(x_{\theta})\\!=\\!\mathbb{E}_{\bm{x}}[f_{W^{(1)},\bm{\hat{}}a)\lambda}(\bm{x})|\nicefrac{{\bm{\theta}^{\top}\bm{x}}}{{\sqrt{d}}}=x_{\theta}]$,
along the direction of the target $\bm{\theta}$, implemented by the network
$f_{W^{(1)},\bm{a}}$ trained with a single gradient step (eq. (3)). Varying
the amounts of data $n_{0}\\!=\\!\alpha_{0}d$ used in the first gradient step,
the function $f(x_{\theta})$ moves from a linear approximation of
$\sigma_{\star}$ in the RF limit ($\alpha_{0}\\!=\\!0$), to an accurate non-
linear one ($\alpha_{0}\\!=\\!2.5$).

Learning a non-linear approximation of $\sigma_{\star}$ in the high-
dimensional proportional regime therefore requires learning features. [1] have
proven that with $\eta=\Theta_{d}(1)$, the GEP holds even after a few gradient
steps, corroborating a fact empirically observed by [33]. Indeed, they have
shown that $\eta=\Theta_{d}(d)$ is _sufficient_ to go beyond a linear
approximation of $f_{\star}$ in this regime. [14] considered intermediate
scalings of step-size $\eta=\Theta_{d}(d^{s})$ for
$\nicefrac{{1}}{{2}}\\!<\\!s\\!<\\!1$, which allows the network to fit target
functions along $\bm{\theta}$ having finite degree. In this intermediate
regime, the feature matrix can be approximated through a finite-number of
spikes corresponding to increasing degree of functions along $\bm{\theta}$.
Instead, we consider the full-scaling of $\eta=\Theta_{d}(d)$, where such a
finite-spike approximation is insufficient and the network can fit arbitrary
functions along $\bm{\theta}$. [2] proved that even if the target depends on
multiple directions (multi-index model), only a (non-linear) function of a
single direction $\bm{\theta}$ can be learned with a single gradient step and
$\eta=\Theta_{d}(d)$. This observation justifies the focus on single-index
functions (2) on the regime of interest.

### Further related works

#### Random features —

Random Features (RFs) were first introduced as a computationally efficient
approximation to kernel methods [28]. Recently, they have enjoyed renewed
interest also as models of two-layer neural networks in the lazy regime. Tight
asymptotics for the random features model have been derived by [18, 17, 19,
22, 20, 21] in the two-layer case, and were extended to deep networks in [24,
25] in the deep case. Importantly, with the exception of [19] who considered
rotationally invariant weights and [34] for the case of deep linear random
features, all these work assumed unstructured weights. In sharp contrast,
gradient-trained neural networks have fundamentally structured weights. In the
present manuscript, we consider such a case, when the weights are given by a
bulk random matrix plus a rank-one spike emerging after a single large
gradient step.

#### Feature learning regime –

Perturbative feature learning corrections to the large-width lazy regime have
been extensively studied in the literature [35, 36, 37, 38, 39, 40]. Our work
radically contrasts with this line, since we exactly account for feature
learning in the first step, _non-perturbatively_ (note the gradient in (3) has
a norm comparable with the initial weights). Beyond the lazy regime, a major
recent development has been the understanding that the training dynamics of
two-layer neural networks with small learning rate can be mapped to a
Wasserstein gradient flow, known as the mean-field regime [41, 42, 43, 44].
Over the past few years, this flow was investigated under different classes of
generative data models, such as staircase functions [45, 46, 11], single-index
[27, 47] and multi-index models [48], symmetric targets [49] and Gaussian
mixture models [50, 51].

##  3 Main technical results

Our main technical results are a tight asymptotic characterization of the test
error achieved by two-layer networks trained with a single large gradient step
followed by a ridge regression on the readout weights. These results are
enabled through the mapping to first an equivalent sRF model in subsection
3.1, which can in turn be mapped to a equivalent Gaussian model in subsection
3.2. Subsection 3.3 finally states the tight asymptotic characterization of
the test error.

###  3.1 Asymptotics of the first layer weights after one (large) gradient
step

The first step is to derive an explicit asymptotic expression for the hidden-
layer weights $W^{(1)}$ after one (large) gradient step. In the following, we
show that the learning problem introduced in Section 2 is equivalent to a
_spiked Random Features_ model, which we first define.

######  Definition 3.1 (sRF model).

We define a _spiked Random Features_ (sRF) model with bulk variance $c$, spike
strength $r$ as the two-layer neural network

$\displaystyle
g_{F,\bm{a}}(x)=\frac{1}{\sqrt{p}}\bm{a}^{\top}\sigma\left(Fx\right)$ (6)

with trainable readout $\bm{a}$ and frozen random first layer weights:

$\displaystyle F=W+r\frac{\bm{u}\bm{v}^{\top}}{\sqrt{d}}.$ (7)

where $W$ is a random matrix with rows independently sampled from
$\mathbb{S}^{d-1}(\sqrt{c})$, and
$\bm{u}\in\mathbb{S}^{d-1}(\sqrt{p}),\bm{v}\in\mathbb{S}^{d-1}(\sqrt{d})$. We
further say that a sRF has alignment $\gamma$ with $\bm{\theta}$ when $\bm{v}$
is uniformly sampled uniformly sampled among vectors with norm $\sqrt{d}$
satisfying $\nicefrac{{\bm{v}^{\top}\bm{\theta}}}{{d}}=\gamma$.

The next result shows that after a large gradient step, our problem is
asymptotically equivalent to a particular sRF model with spikes
$\bm{u}=\bm{1}_{p}$ and $\bm{v}$ that correlated to the target weights
$\bm{\theta}$. We give a tight characterization of the parameters
$c,r,\gamma$.

######  Result 3.1 (Equivalence to a sRF model).

Consider two-layer networks with first-layer weights trained with a single
gradient step of learning rate $\eta$ from initial conditions
$\bm{a}^{(0)}=\nicefrac{{\bm{1}_{p}}}{{\sqrt{p}}}$ and $\bm{w}_{i}^{(0)}$ with
unit norm $\lVert\bm{w}^{(0)}_{i}\lVert=1$ and weak correlation
${\bm{w}^{(0)}_{i}}\cdot\bm{w}^{(0)}_{j}=O_{d}(\nicefrac{{1}}{{\sqrt{d}}})$
for $i\neq j$ (eq. (3)). In the asymptotic limit $n_{0},n,d,p\to\infty$, with
$\alpha_{0}=\nicefrac{{n_{0}}}{{d}},\alpha=\nicefrac{{n}}{{d}},\beta=\nicefrac{{p}}{{d}},\tilde{\eta}=\nicefrac{{\eta}}{{d}}=\Theta_{d}(1)$,
the test error achieved by performing ridge regression on the readout weights
$\bm{a}$ of the network after the gradient step (4) is identical to that
achieved by an equivalent sRF model with parameters

$\displaystyle
c=1+\frac{\tilde{\eta}^{2}h_{1}^{2}\check{h}_{1}^{2}h_{2}^{\star}}{\alpha_{0}\beta^{2}}$
(8) $\displaystyle
r=\frac{\tilde{\eta}h_{1}}{\beta}\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}}$ (9)
$\displaystyle\gamma={\frac{h_{1}^{\star}}{\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}}}}$ (10)

where:

$\displaystyle h_{1}=\mathbb{E}_{z}[z\sigma(z)],$ $\displaystyle
h^{\star}_{1}=\mathbb{E}_{z}[z\sigma_{\star}(z)],$
$\displaystyle\check{h}_{1}^{2}=\mathbb{E}_{z}[(\sigma^{\prime}(z)-h_{1})^{2}],$
$\displaystyle h^{\star}_{2}=\mathbb{E}_{z}[\sigma_{\star}(z)^{2}],$ (11)

with $z\sim\mathcal{N}(0,1)$.

The derivation of Result (3.1), detailed in Appendix A, leverages the
decomposition from [1] of the gradient
$\bm{g}_{i}=\nicefrac{{ru_{i}\bm{v}}}{{d}}+\Delta$ into a rank-one term and a
correction term. Note that above we have assumed a uniform initialization for
the readout layer. This can be relaxed in the equivalence above, for instance
for $\bm{a}^{(0)}$ taking a finite number of values, leading to finite-rank
term instead of a single spike.

###  3.2 Conditional Gaussian equivalence

The sharp characterization of the test performance, which we state in Result
3.3 in the following subsection, is enabled by further mapping the sRF model
to an exactly solvable (conditional) Gaussian model. We adapt the rigorous
result in [2] (Theorem 4) by constructing explicitly the equivalent stochastic
feature map, that we believe is of independent interest.

######  Result 3.2 (Conditional Gaussian Equivalence).

Consider the sRF model with weights
$F=W+r\nicefrac{{\bm{u}\bm{v}^{\top}}}{{\sqrt{d}}}$, with $\bm{u}=\bm{1}_{p}$
and parameters $c,r,\gamma$, and the corresponding feature map given by

$\displaystyle\bm{\varphi}(\bm{x})=\sigma\left(F\bm{x}\right)$ (12)

Define the equivalent stochastic feature map

$\displaystyle\bm{\varphi}^{g}(\bm{x})\overset{d}{=}\mu_{0}(\kappa)\bm{1}_{p}+\mu_{1}(\kappa)W\bm{x}+\mu_{2}(\kappa)\mathcal{N}(0,\mathbb{I}_{p}),$
(13)

where $\kappa\equiv\nicefrac{{\bm{v}^{\top}\bm{x}}}{{\sqrt{d}}}$. We
introduced the coefficients $\mu_{0}(\kappa),\mu_{1}(\kappa),\mu_{2}(\kappa)$
defined as

$\begin{split}\mu_{0}(\kappa)&=\operatorname{\mathbb{E}}_{z}{\sigma(z+r\kappa)}\\\
\mu_{1}(\kappa)&=\frac{1}{c}\operatorname{\mathbb{E}}_{z}{z\sigma(z+r\kappa)}\\\
\mu_{2}(\kappa)&=\sqrt{\operatorname{\mathbb{E}}_{z}{\sigma^{2}(z+r\kappa)-c(\mu_{1}(\kappa))^{2}-(\mu_{0}(\kappa))^{2}}},\end{split}$
(14)

with expectations bearing over $z\sim\mathcal{N}(0,c)$. The test error
$\epsilon_{g}$ achieved by ridge regression

$\displaystyle\hat{\bm{a}}_{\lambda}=\underset{\bm{a}\in\mathbb{R}^{p}}{\mathrm{argmin}}\frac{1}{2}\sum\limits_{\mu=1}^{n_{1}}\left(y^{\mu}-\frac{1}{\sqrt{p}}\bm{a}^{\top}\phi(\bm{x})\right)^{2}\\!\\!\\!\\!+\\!\\!\frac{\lambda}{2}\lVert\bm{a}\lVert^{2}.$
(15)

is asymptotically identical for $\phi=\varphi$ and $\phi=\varphi^{g}$.

Result 3.2 extends the similar linearizations provided e.g. in [18, 20, 23]
for unstructured RFs to sRFs (7), and we refer to App. B for a detailed
discussion on the derivation. Informally, the quantity $\kappa$ in the
stochastic feature map (13) represents the projection of the input on the
spike defining the sRF $\kappa=\nicefrac{{\bm{x}^{\top}\bm{v}}}{{\sqrt{d}}}$.
The equivalent network
$\nicefrac{{1}}{{\sqrt{d}}}\bm{a}^{\top}\bm{\varphi}^{g}(x)$ obtained by
replacing $\bm{\varphi}$ by the equivalent $\bm{\varphi}^{g}$ is a linear
combination of terms such as $\mu_{0}(\kappa)$, $\mu_{1}(\kappa)\kappa$,
$\mu_{1}(\kappa)\bm{a}^{\top}W(\Pi^{\perp}\bm{x})$, plus noise. On an
intuitive level, this makes it apparent that sRFs can thus express non-linear
functions of the component $\kappa$ along the spike $\bm{v}$, but only linear
functions of the component $\Pi^{\perp}\bm{x}$ orthogonal thereto. This
feature learning correction to the linear regime is visually exemplified in
Fig. 1. In the next subsection, we make this discussion more quantitative by
providing a tight asymptotic characterization of the test error achieved by
two-layer networks trained with a single large gradient step followed by a
ridge regression on the readout weights.

###  3.3 Tight asymptotic characterization of the test error

Finally, we leverage on the sequential mappings of Result 3.1 and Result 3.2
to offer sharp asymptotic guarantees on the test error achieved after training
the readout weights using eq. (4).

######  Assumption 3.2.

Denote by $\\{\bm{e}_{i}\\}_{i=1}^{p}$ ($\\{\bm{f}_{i}\\}_{i=1}^{d}$) the left
(resp. right) singular vectors of $W$. We further note
$\\{\lambda^{\ell}_{i}\\}_{i=1}^{p}$ the squared singular values of $W$. The
squared singular values $\\{\lambda^{\ell}_{i}\\}_{i=1}^{p}$ and the
projection of the teacher vector $\bm{\theta}$ and the spike $v$ on the
eigenvectors
$\\{\bm{f}_{i}^{\top}\bm{v}\\}_{i,\ell}$,$\\{\bm{f}_{i}^{\top}\Pi^{\perp}\bm{\theta}\\}_{i,\ell}$
are assumed to admit a well-defined joint distribution $\nu$ as $d\to\infty$.

$\displaystyle\frac{1}{p}\sum\limits_{i=1}^{\min(p,d)}\delta\left(\lambda_{i}-\varrho\right)\delta\left(\bm{f}_{i}^{\top}v-\tau\right)\delta\left(\bm{f}_{i}^{\top}\Pi^{\perp}\bm{\theta}-\pi\right)\xrightarrow[]{d\to\infty}\nu(\varrho,\tau,\pi).$

######  Result 3.3 (Test error asymptotics).

Consider the ERM problem (4) associated with the training of the readout
weights $\bm{a}$, and assume 3.2 to hold. Define
$\Pi^{\perp}\equiv\mathbb{I}_{d}-\nicefrac{{vv^{\top}}}{{d}}$ the projection
to the subspace orthogonal to the spike $\bm{v}$. In the asymptotic limit
$d,p,n\to\infty$ with
$\alpha=\nicefrac{{n}}{{d}},\beta=\nicefrac{{p}}{{d}}=\Theta_{d}(1)$, the
summary statistics

$\displaystyle
q_{1}=\frac{\hat{\bm{a}}^{\top}W\Pi^{\perp}W^{\top}\hat{\bm{a}}}{p},$
$\displaystyle q_{2}=\frac{\hat{\bm{a}}^{\top}\hat{\bm{a}}}{p},$
$\displaystyle m=\frac{\bm{1}_{p}^{\top}\hat{\bm{a}}}{\sqrt{p}},$
$\displaystyle\zeta=\frac{\hat{\bm{a}}^{\top}W\bm{v}}{\sqrt{dp}},$
$\displaystyle\psi=\frac{\hat{\bm{a}}^{\top}W\Pi^{\perp}\bm{\theta}}{\sqrt{dp}},$
$\displaystyle\rho^{2}=\frac{\bm{\theta}^{\top}\Pi^{\perp}\bm{\theta}}{d}$
(16)

concentrate in probability to the solutions of the system of equations

$\displaystyle\begin{cases}q_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}\\\
\qquad-\beta\hat{\zeta}^{2}\frac{I(\hat{V}_{1},\hat{V}_{2})^{2}}{\left(1-\beta\hat{V}_{1}I(\hat{V}_{1},\hat{V}_{2})\right)^{2}}\\\
\qquad-\hat{\zeta}^{2}\frac{\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho^{2}}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\left[\left(1-\beta\hat{V}_{1}I(\hat{V}_{1},\hat{V}_{2})\right)^{2}-1\right]}{\left(1-\beta\hat{V}_{1}I(\hat{V}_{1},\hat{V}_{2})\right)^{2}}\\\
q_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}\\\
\qquad-\hat{\zeta}^{2}\\!\\!\int\\!\\!d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\\!\\!\left[\\!1\\!\\!-\\!\\!\frac{1}{\left(1-\beta\hat{V}_{1}I(\hat{V}_{1},\hat{V}_{2})\right)^{2}}\right]\\\
V_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
V_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
m=\frac{1}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)^{2}}{1+V(\kappa)}\right]}\mathbb{E}_{\kappa,y}\left[\frac{\mu_{0}(\kappa)(\sigma_{\star}(\kappa,y)-\mu_{1}(\kappa)\kappa\zeta)}{1+V(\kappa)}\right]\\\
\zeta=\hat{\zeta}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\tau^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
\qquad+\beta^{\nicefrac{{3}}{{2}}}\hat{\zeta}\hat{V}_{1}\frac{I(\hat{V}_{1},\hat{V}_{2})^{2}}{1-\beta\hat{V}_{1}I(\hat{V}_{1},\hat{V}_{2})}\\\
\psi=\hat{\psi}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\pi^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\end{cases},$
$\displaystyle\begin{cases}\hat{V}_{1}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa}\frac{\rho\mu_{1}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}_{1}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y}\mu_{1}(\kappa)^{2}\frac{b(\kappa,y)^{2}+\rho
q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{V}_{2}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa}\frac{\rho\mu_{2}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}_{2}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y}\mu_{2}(\kappa)^{2}\frac{b(\kappa,y)^{2}+\rho
q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{\zeta}=\frac{\alpha}{\sqrt{\beta}}\mathbb{E}_{\kappa,y}\kappa\mu_{1}(\kappa)\frac{b(\kappa,y)}{1+V(\kappa)}\\\
\hat{\psi}=\frac{\alpha}{\sqrt{\beta}}\mathbb{E}_{\kappa,y}\frac{y\mu_{1}(\kappa)b(\kappa,y)+\psi\mu_{1}(\kappa)^{2}}{1+V(\kappa)}\end{cases}$
(17)

and

$\displaystyle\rho^{2}=1-\gamma^{2}.$ (18)

We introduced the shorthands

$\displaystyle
b(y,\kappa)\equiv\sigma_{\star}(\gamma\kappa+\sqrt{1-\gamma^{2}}y)-\mu_{0}(\kappa)m-\kappa\mu_{1}(\kappa)\zeta\\!\\!-\\!\\!\mu_{1}(\kappa)\psi
y$ (19) $\displaystyle
V(\kappa)\equiv\mu_{1}(\kappa)^{2}V_{1}+\mu_{2}(\kappa)^{2}V_{2}$ (20)
$\displaystyle
q(\kappa)\equiv\mu_{1}(\kappa)^{2}q_{1}+\mu_{2}(\kappa)^{2}q_{2},$ (21)
$\displaystyle I(\hat{V}_{1},\hat{V}_{2})=\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}$
(22)

In (17), the expectations over $\kappa,y$ bear over standard Gaussian
variables. We remind that the quantities $r,c,\gamma$ are characterized in
Result 3.1.
Finally, the test error (5) admits the sharp characterization

$\displaystyle\epsilon_{g}=\mathbb{E}_{\kappa,y}\Bigg{[}$
$\displaystyle\Bigg{(}\sigma_{\star}(\gamma\kappa+\sqrt{1-\gamma^{2}}y)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\kappa\zeta-\frac{\mu_{1}(\kappa)\psi}{\sqrt{\rho}}y\Bigg{)}^{2}{+q(\kappa)-\frac{\mu_{1}(\kappa)^{2}\psi^{2}}{\rho}}\Bigg{]}$
(23)

where expectations bear over standard Gaussian variables.

Result 3.3 thus rephrases the high-dimensional learning problem (4) in terms
of a finite set of scalar summary statistics which can be efficiently
evaluated, yielding excellent agreement with finite $d$ numerical simulations,
see Figs. 2-4. While we state Result 3.3 for the square loss and $\ell_{2}$
regularization for clarity, we provide a sharp characterization for any
generic convex loss $\ell$ in Appendix C. While we have used the (non-
rigorous) replica method to derive these equations, an interesting avenue of
future research is to provide a rigorous probabilistic proof. A possible way
to address the problem is to apply Gordon’s Gaussian min-max inequalities [52,
53, 54, 55]), and generalized the approach used for pure random features [20,
56]. This is, however, beyond the scope of this manuscript.

##  4 Discussion of main results

While the self-consistent equations in Result 3.3 might appear cumbersome,
they offer valuable insight into the mechanism behind feature learning in two-
layer neural networks trained with gradient descent. In this section, we
discuss and highlight some of these insights.

Figure 2: Test error achieved by a two-layer network with activation $\sigma$
whose first layer has been trained following the protocol detailed in section
2 on a single-index target with activation $\sigma_{\star}$. (Left)
Activations $\sigma=\sigma_{\star}=\tanh$, learning rate $\tilde{\eta}=1$,
readout regularization $\lambda=0.01$. (Right) Activations
$\sigma=\tanh,\sigma_{\star}=\sin$, learning rate $\tilde{\eta}=3$, readout
regularization $\lambda=0.1$. The dashed black line represents the lowest
achievable MSE for kernel/linear methods, namely
$h^{\star}_{2}-(h^{\star}_{1})^{2}$ [1]. In both plots, solid lines correspond
to the theoretical Result 3.3, and crosses correspond numerical experiments in
$d=3000$ (left) and $d=2000$ (right). All points were averaged over $5$
instances. Different colours represent different initial sample complexities
$\alpha_{0}=\nicefrac{{n_{0}}}{{d}}$ used for the first gradient step.

###  4.1 Spiked Random Features vs Random Features

The asymptotic characterization 3.3 encompasses, as a special case, usual RFs
(when setting the spike strength $r$ to zero). More precisely, for zero spike
strength $r=0$ in (7), sRFs coincide with RFs, as the coefficients (14) lose
their $\kappa$ dependence, reducing to usual Hermite coefficients. The
equivalent feature map $\bm{\varphi}^{g}$ then reduces to the Gaussian
equivalent feature map employed in e.g. [18, 20, 24]. Importantly, while the
equivalent feature map for unspiked RFs is linear in the input $\bm{x}$, the
equivalent feature map $\bm{\varphi}^{g}$ (13) of Result 3.2 is non-linear in
the component $\kappa$ of the input $\bm{x}$.

sRF models, which are equivalent to two-layer networks after a single large
gradient step through Result 3.1, offer an ideal playground to test the
intertwined influence of the spike/target correlation and the test performance
of the model. Fig. 4 (left) presents the learning curves of (s)RFs for varying
spike strengths $r$. As is intuitive, larger spikes allow the sRF to more
easily express non-linear function in the direction of the spike, and thus
lead to relatively smaller test errors. Furthermore, note that even rather
small spike strengths $r=0.2$ already yield test errors which are sizably
lower than vanilla (unspiked) RFs ($r=0$), hinting at the qualitative
difference between sRF and RF models discussed in section 2 and Fig. 1. This
is in qualitative agreement with the result of [14], who studied the
intermediate scaling regime where the learning rate $\eta=\Theta_{d}(d^{s})$
with $s\in(\nicefrac{{1}}{{2}},1)$, and have shown that in this regime a
polynomial approximation of $\sigma_{\star}$ is learned.

The plot shows a compelling agreement between the theoretical predictions
(continuous line) and numerical simulations. Fig. 4 (right) represents the
theoretical closed-form expression for the test error (23) for different
values of the spike/target alignment $\gamma$, for a single-index target
$\mathrm{sign}(\bm{\theta}^{\top}\bm{x})$, with good agreement with numerical
experiments. Higher alignments $\gamma$ lead to overall lower test errors
escaping the linear curse of Gaussian models.

###  4.2 Beating kernels in a single step

A key parameter in our formulas is
$\gamma=\nicefrac{{\bm{v}^{\top}\bm{\theta}}}{{d}}\in[0,1]$, the correlation
between the effective gradient spike and the target weights. From its
asymptotic expression eq. 8 in Result 3.1, this is an increasing function of
$\alpha_{0}=\nicefrac{{n_{0}}}{{d}}$, the number of samples in the first step.
As shown in Fig. 2 (left) for $\sigma=\sigma_{\star}=\tanh$, larger sample
complexities in the representation learning step $\alpha_{0}$ allow for better
feature learning when implementing a gradient step on the first layer,
enabling a lower test error after the readout layer is retrained. As expected,
the lowest error is achieved when $\alpha_{0}\to\infty$, in which case the
spike $\bm{v}$ is perfectly aligned with the target weights $\bm{\theta}$ and
$\gamma=1$, as can be seen from (8).

Figure 2 (right) presents similar curves for another target activation
$\sigma^{\star}=\sin$, including the test error achieved by the network at
initialization ($\alpha_{0}=0$), which corresponds to a usual RF model. The
latter is well above the lowest Mean Squared Error (MSE) achievable by a
linear estimator (plotted as a dashed black line), namely the projection of
the teacher function on Hermite polynomials
$\lVert\mathbb{P}_{>1}f_{\star}\lVert_{2}^{2}$. This performance corresponds
to the kernel one when the number of samples scales proportionally with the
input dimension [29]; using the notations of Result 3.1 the best linear MSE is
readily written as $h^{\star}_{2}-(h^{\star}_{1})^{2}$. In sharp contrast,
networks with trained first layers ($\alpha_{0}>0$) can learn non-linear
functions of the inputs and outperform this baseline.

###  4.3 What can be learned with a single step?

Figure 3: ($c=1,r=0.9$ and $\gamma=1$) Illustration of the functions realizing
the upper bound (24) (orange) and lower bound (25) (blue), for $\sigma=\tanh$,
for a target $\sigma_{\star}=\sin$ (dashed black).

While training a two-layer network with large gradient steps allows it to
escape the linear limitations of, e.g. RF models, it generically only learns
the target up to small test error, yet not perfectly. In fact, a closer
examination of the sharp asymptotic expression (23) for $\epsilon_{g}$ in
Result 3.3 reveals that even for large initial batches $\alpha_{0}\to\infty$
(and thus perfect spike/target alignments $\gamma=1$), at fixed learning rate
strength $\tilde{\eta}$ and sample complexity $\alpha>0$, the test error
optimized over the regularization $\lambda$ is upper bounded as

$\displaystyle\underset{\lambda\geq
0}{\inf}\epsilon_{g}\leq\underset{b_{1}}{\inf}~{}\mathbb{E}_{\kappa}\left[\sigma_{\star}(\kappa)-b_{1}\mu_{0}(\kappa)\right]^{2},$
(24)

and lower bounded by:

$\displaystyle\underset{\lambda\geq
0}{\inf}~{}\epsilon_{g}\\!\geq\\!\underset{b_{1},b_{2}}{\inf}~{}\mathbb{E}_{\kappa}\left[\sigma_{\star}(\kappa)-b_{1}\mu_{0}(\kappa)-b_{2}\mu_{1}(\kappa)\kappa\right]^{2}.$
(25)

The upper bound (24) is the equivalent of Lemma 6 of [1] in our case of
uniform readout initialization
$\bm{a}^{(0)}=\nicefrac{{\bm{1}_{p}}}{{\sqrt{p}}}$, and is achieved for
$\lambda\to\infty$, as we show in App. E. The lower bound (25), however, shows
that the test error cannot be lower than the Gaussian-weighted $L^{2}$
distance between the target function $\sigma_{\star}$ and
$\mathrm{span}(\mu_{0},\tilde{\mu}_{1})$, where
$\tilde{\mu}_{1}(\kappa)\equiv\kappa\mu_{1}(\kappa)$, and the best
approximation would be reached for the projection of the target thereon. Fig.
3 illustrates what functions realize the upper (24) and lower bounds (25), and
how they compare with the target $\sigma_{\star}$. Finally, note that in the
vanilla RF limit $r=0$, the functions $\mu_{0}(\kappa),\mu_{1}(\kappa)$ reduce
to constants independent of $\kappa$, constraining the class of functions that
can be learned to that of linear functions.

Finally, let us mention that while this discussion was made at fixed learning
rate $\tilde{\eta}$ for clarity, the latter is in practice a tunable hyper-
parameter, and the functions $\mu_{0},\mu_{1}$ depend thereupon via the spike
strength $r$ (14). One can thus refine – and lower– the bounds (24) and (25)
over the functions $\mu_{0},\mu_{1}$, which can take values in the realizable
set $\mathcal{M}=\\{(\mu_{0}(r),\mu_{1}(r)\\}_{r\geq 0}$ (emphasizing the
dependence on $r$ via equation (75)) as $\tilde{\eta}$ is varied:

$\displaystyle\underset{\lambda\geq 0,\tilde{\eta}\geq
0}{\inf}\epsilon_{g}\leq\underset{b_{1}\in\mathbb{R},\nu_{0}\in\mathcal{M}}{\inf}\mathbb{E}_{\kappa}\left[\sigma_{\star}(\kappa)-b_{1}\nu_{0}(\kappa)\right]^{2},$
(26) $\displaystyle\underset{\lambda\geq 0,\tilde{\eta}\geq
0}{\inf}\epsilon_{g}\geq\\!\\!\\!\\!\\!\\!\\!\\!\\!\underset{\begin{array}[]{cc}{\scriptstyle
b_{1},b_{2}\in\mathbb{R}},\\\
{\scriptstyle(\nu_{0},\nu_{1})\in\mathcal{M}}\end{array}}{\inf}\\!\\!\\!\\!\\!\\!\\!\\!\\!\mathbb{E}_{\kappa}\left[\sigma_{\star}(\kappa)-b_{1}\nu_{0}(\kappa)-b_{2}\nu_{1}(\kappa)\kappa\right]^{2}.$
(29)

In other words, by tuning the learning rate $\tilde{\eta}$ – and thus the
spike strength $r$–, one gains the freedom to choose the "best" subspace
$\mathrm{span}(\mu_{0}(r),\mu_{1}(r))$, i.e. the one allowing to approximate
the target $\sigma_{\star}$ best.

Finally, as discussed in [1], observe that for
$\sigma=\sigma_{\star}=\mathrm{erf}$, the upper bound in (26) is zero provided
one tunes the learning rate to
$\tilde{\eta}=\nicefrac{{\sqrt{3\beta}}}{{h_{1}}}$, and perfect learning is
therefore achievable.

Figure 4: Test error for a sRF with activation $\sigma$ learning from a
single-index model
$\sigma_{\star}(\nicefrac{{\bm{\theta}^{\top}\bm{x}}}{{\sqrt{d}}})$, with
regularization $\lambda=0.1$. (Left) $\sigma=\sin,\sigma_{\star}={\rm sign}$.
Different colours correspond to different spike strengths $r$ (7), with $r=0$
corresponding to the vanilla RF model. (Right) $\sigma=\sigma_{\star}=\tanh$.
Different colours correspond to different overlaps
$\gamma\equiv\nicefrac{{\bm{\theta}^{\top}\bm{v}}}{{d}}$ between the target
weights $\bm{\theta}$ and the spike $\bm{v}$. Solid lines: theoretical
characterization of Result 3.3. Crosses : numerical simulations in dimensions
$d\\!=\\!p\\!=\\!2000$. Each point is averaged over $10$ instances of the
problem.

###  4.4 More variability means better feature learning

The discussion of subsection 4.3 thus affords an insightful perspective on the
learning of two-layer neural networks in terms of approximating the target
activation $\sigma_{\star}$ in a two-dimensional functional space.
Interestingly, introducing variability in the readout layer at initialization
leads to an even richer functional basis, and hence greater expressivity of
the network. As we further discuss in Appendix D, when $\bm{a}^{(0)}$ is no
longer proportional to $\bm{1}_{p}$, but is rather initialized from a
distribution over a finite vocabulary $V$ of size $|V|>1$–e.g.
$V=\\{-1,0,+1\\}$– the equivalent feature map takes the form:

$\displaystyle\bm{\varphi}^{g}(\bm{x})=\begin{pmatrix}\scriptstyle\scriptstyle\mu_{0}(u_{1}\kappa)\\\
\scriptstyle\mu_{0}(u_{2}\kappa)\\\ \vdots\\\
\scriptstyle\mu_{0}(u_{p}\kappa)\end{pmatrix}+\begin{pmatrix}\scriptstyle\mu_{1}(u_{1}\kappa)\\\
\scriptstyle\mu_{1}(u_{2}\kappa)\\\ \vdots\\\
\scriptstyle\mu_{1}(u_{p}\kappa)\end{pmatrix}\odot
W\bm{x}+\begin{pmatrix}\scriptstyle\mu_{2}(u_{1}\kappa)\\\
\scriptstyle\mu_{2}(u_{2}\kappa)\\\ \vdots\\\
\scriptstyle\mu_{2}(u_{p}\kappa)\end{pmatrix}\odot\bm{\xi}$ (30)

where $\odot$ denotes element-wise multiplication,
$\bm{\xi}\sim\mathcal{N}(0,\mathbb{I}_{p})$, and $\bm{u}\not\propto\bm{1}_{p}$
has entries which can now take a finite number of different values. The
coefficients of the equivalent map (30) are thus neuron-dependent and thus
afford a richer functional basis
$\\{\mu_{0}(\omega\cdot),\tilde{\mu}_{1}(\omega\cdot)\\}_{\omega\in V}$,
thereby allowing the network to express a larger class of functions. As a
matter of fact, the functional space spanned by these functions is generically
of dimension $2|V|$ for non-uniform readout initializations $\bm{a}^{(0)}$,
compared to just $2$ in the uniform readout case. A sharp asymptotic
characterization of the test error for the case of non-uniform readout
initialization with finite vocabulary is further provided in Appendix D.

We briefly discuss the limiting case of interest
$\lambda,\alpha_{0},\tilde{\eta}\rightarrow\infty$, for which the equivalent
feature map $\bm{\varphi}^{g}(\bm{x})$ (30) reduces to its first term
$(\mu_{0}(u_{i}\kappa))_{i=1}^{p}$. Further observe that
$\mu_{0}(u_{i}\kappa)$ can be viewed as a one-dimensional neuron acting on the
one-dimensional input variable $\kappa$ with a random weight $u_{i}$. Standard
results on approximation errors for random feature mappings of finite-
dimensional inputs [28] imply that a large class of functions can be
approximated from the network features $\bm{\varphi}^{g}(\bm{x})$, provided
the vocabulary size $|V|$ is large enough. Similar random features
approximations have been leveraged in [1, 12, 10]. The equivalent feature map
(30) provides an intuitive picture on how such random feature mappings of low-
dimensional inputs can naturally emerge in the setting of the learning of a
two-layer network.

Finally, note that the bounds (24) and (25) can be readily generalised to non-
uniform readout initializations, provided one replaces in the discussion of
subsection 4.3 the two-dimensional functional basis
$\mathrm{span}(\mu_{0},\tilde{\mu}_{1})$ in the case of uniform initialization
by the richer functional space
$\mathrm{span}(\\{\mu_{0}(\omega\cdot),\tilde{\mu}_{1}(\omega\cdot)\\}_{\omega\in
V})$ for non-uniform initialization. For instance, the lower bound (25) thus
involves in the non-uniform case the distance between the target
$\sigma_{\star}$ and
$\mathrm{span}(\\{\mu_{0}(\omega\cdot),\tilde{\mu}_{1}(\omega\cdot)\\}_{\omega\in
V})$.

## Conclusion

We provided a tight asymptotic description of the learning of a two-layer
neural network after training its first layer with a large, single, gradient
step, in the limit where the number of samples, the hidden layer width, and
the input dimension are proportionally large. Our results sharply characterize
the feature maps learned from the data, and how it achieves a test error which
non-perturbatively improves over the kernel regime. Crucially, the trained
network can efficiently approximate non-linear functions in the direction of
the gradient – sizeably improving upon the network at initialization, which
can only express linear functions. We further discuss bounds for the test
error and which functions are learnable after a single gradient step. Finally,
extensive numerical support is provided to illustrate our findings. To our
knowledge, our work provides the first exact asymptotic result in the non-
perturbative $\eta=\Theta_{d}(d)$ regime for feature learning.

We believe the present work opens exciting research avenues, paving the way
towards a tight theoretical understanding of feature learning in gradient-
trained networks. Prominent among these research directions is the extension
of our results to readout initialization with generic (not necessarily finite)
support. to a finite number of gradient steps, and ultimately fully trained
networks.

##  5 Acknowledgements

We thank Denny Wu, Gabriel Arpino, Weiyu Li, Marco Mondelli and Theodor
Misiakiewicz for discussion during the course of this project. Early
discussions on this project started as an open problem posed by BL during the
2022 Les Houches Summer School on _Statistical Physics and Machine Learning_
organized by LZ and FK. We acknowledge funding from the Swiss National Science
Foundation grant SNFS OperaGOST (grant number $200390$), and SMArtNet (grant
number $212049$), and the Choose France - CNRS AI Rising Talents program.

## References

  * [1] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang.  High-dimensional asymptotics of feature learning: How one gradient step improves the representation.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 37932–37946. Curran Associates, Inc., 2022.
  * [2] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan.  How two-layer neural networks learn, one (giant) step at a time, 2023.
  * [3] Ben Adlam, Jaehoon Lee, Shreyas Padhy, Zachary Nado, and Jasper Snoek.  Kernel regression with infinite-width neural networks on millions of examples.  arXiv preprint arXiv:2303.05420, 2023.
  * [4] R.M. Neal.  Bayesian Learning for Neural Networks.  Lecture Notes in Statistics. Springer New York, 1996.
  * [5] Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri.  Deep neural networks as gaussian processes.  In International Conference on Learning Representations, 2018.
  * [6] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  Advances in neural information processing systems, 31, 2018.
  * [7] Lénaïc Chizat, Edouard Oyallon, and Francis Bach.  On Lazy Training in Differentiable Programming.  In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [8] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Limitations of lazy training of two-layers neural network.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [9] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová.  Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed.  In International Conference on Machine Learning, pages 8936–8947. PMLR, 2021.
  * [10] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi.  Neural networks can learn representations with gradient descent.  In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 5413–5452. PMLR, 02–05 Jul 2022.
  * [11] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz.  Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, 2023.
  * [12] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu.  Learning in the presence of low-dimensional structure: a spiked random matrix perspective.  In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
  * [13] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A Erdogdu.  Gradient-based feature learning under structured data.  arXiv preprint arXiv:2309.03843, 2023.
  * [14] Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban.  A theory of non-linear feature learning with one gradient step in two-layer neural networks, 2023.
  * [15] Giorgio Parisi.  Toward a mean field theory for spin glasses.  Physics Letters A, 73(3):203–205, 1979.
  * [16] Giorgio Parisi.  Order parameter for spin-glasses.  Physical Review Letters, 50(24):1946, 1983.
  * [17] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mezard, and Lenka Zdeborova.  The gaussian equivalence of generative models for learning with shallow neural networks.  In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference, volume 145 of Proceedings of Machine Learning Research, pages 426–471. PMLR, 16–19 Aug 2022.
  * [18] Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová.  Modeling the influence of data structure on learning in neural networks: The hidden manifold model.  Physical Review X, 10(4):041044, 2020.
  * [19] Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mezard, and Lenka Zdeborova.  Generalisation error in learning with random features and the hidden manifold model.  In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3452–3462. PMLR, 13–18 Jul 2020.
  * [20] Hong Hu and Yue M Lu.  Universality laws for high-dimensional learning with random features.  IEEE Transactions on Information Theory, 2022.
  * [21] Oussama Dhifallah and Yue M. Lu.  A precise performance analysis of learning with random features, 2022.
  * [22] Song Mei and Andrea Montanari.  The generalization error of random features regression: Precise asymptotics and the double descent curve.  Communications on Pure and Applied Mathematics, 75(4):667–766, 2022.
  * [23] Hugo Cui, Florent Krzakala, and Lenka Zdeborova.  Bayes-optimal learning of deep random networks of extensive-width.  In International Conference on Machine Learning, pages 6468–6521. PMLR, 2023.
  * [24] Dominik Schröder, Hugo Cui, Daniil Dmitriev, and Bruno Loureiro.  Deterministic equivalent and error universality of deep random features learning.  arXiv preprint arXiv:2302.00401, 2023.
  * [25] David Bosch, Ashkan Panahi, and Babak Hassibi.  Precise asymptotic analysis of deep random feature models.  In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 4132–4179. PMLR, 12–15 Jul 2023.
  * [26] Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová.  Universality laws for gaussian mixtures in generalized linear models, 2023.
  * [27] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou.  Learning time-scales in two-layers neural networks, 2023.
  * [28] Ali Rahimi and Benjamin Recht.  Random features for large-scale kernel machines.  Advances in neural information processing systems, 20, 2007.
  * [29] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  When do neural networks outperform kernel methods?  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 14820–14830. Curran Associates, Inc., 2020.
  * [30] Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration.  Applied and Computational Harmonic Analysis, 59:3–84, 2022.  Special Issue on Harmonic Analysis and Machine Learning.
  * [31] Lechao Xiao, Hong Hu, Theodor Misiakiewicz, Yue Lu, and Jeffrey Pennington.  Precise learning curves and higher-order scalings for dot-product kernel regression.  Advances in Neural Information Processing Systems, 35:4558–4570, 2022.
  * [32] Andrea Montanari and Basil N Saeed.  Universality of empirical risk minimization.  In Conference on Learning Theory, pages 4310–4312. PMLR, 2022.
  * [33] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová.  Learning curves of generic features maps for realistic datasets with a teacher-student model.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 18137–18151. Curran Associates, Inc., 2021.
  * [34] Jacob A Zavatone-Veth and Cengiz Pehlevan.  Learning curves for deep structured gaussian feature models.  arXiv preprint arXiv:2303.00564, 2023.
  * [35] Sho Yaida.  Non-Gaussian processes and neural networks at finite widths.  In Jianfeng Lu and Rachel Ward, editors, Proceedings of The First Mathematical and Scientific Machine Learning Conference, volume 107 of Proceedings of Machine Learning Research, pages 165–192. PMLR, 20–24 Jul 2020.
  * [36] Boris Hanin and Mihai Nica.  Finite depth and width corrections to the neural tangent kernel.  In International Conference on Learning Representations, 2020.
  * [37] Ethan Dyer and Guy Gur-Ari.  Asymptotics of wide networks from feynman diagrams.  In International Conference on Learning Representations, 2020.
  * [38] Inbar Seroussi, Gadi Naveh, and Zohar Ringel.  Separation of scales and a thermodynamic description of feature learning in some cnns.  Nature Communications, 14(1):908, Feb 2023.
  * [39] Gadi Naveh and Zohar Ringel.  A self consistent theory of gaussian processes captures feature learning effects in finite cnns.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 21352–21364. Curran Associates, Inc., 2021.
  * [40] Blake Bordelon and Cengiz Pehlevan.  Dynamics of finite width kernel and prediction fluctuations in mean field neural networks, 2023.
  * [41] Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A mean field view of the landscape of two-layer neural networks.  Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
  * [42] Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  Advances in neural information processing systems, 31, 2018.
  * [43] Grant Rotskoff and Eric Vanden-Eijnden.  Trainability and accuracy of artificial neural networks: An interacting particle system approach.  Communications on Pure and Applied Mathematics, 75(9):1889–1935, 2022.
  * [44] Justin Sirignano and Konstantinos Spiliopoulos.  Mean field analysis of neural networks: A central limit theorem.  Stochastic Processes and their Applications, 130(3):1820–1852, 2020.
  * [45] Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj.  The staircase property: How hierarchical structure can guide deep learning.  Advances in Neural Information Processing Systems, 34:26989–27002, 2021.
  * [46] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz.  The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks.  In Conference on Learning Theory, pages 4782–4887. PMLR, 2022.
  * [47] Luca Arnaboldi, Florent Krzakala, Bruno Loureiro, and Ludovic Stephan.  Escaping mediocrity: how two-layer networks learn hard single-index models with sgd, 2023.
  * [48] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro.  From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks, February 2023.  arXiv:2302.05882 [cond-mat, stat] type: article.
  * [49] Karl Hajjar and Lenaic Chizat.  On the symmetries in the dynamics of wide two-layer neural networks, 2023.
  * [50] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborova.  Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8936–8947. PMLR, 18–24 Jul 2021.
  * [51] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath.  High-dimensional limit theorems for sgd: Effective dynamics and critical scaling.  Advances in Neural Information Processing Systems, 35:25349–25362, 2022.
  * [52] Yehoram Gordon.  On milman’s inequality and random subspaces which escape through a mesh in rn.  In Geometric Aspects of Functional Analysis: Israel Seminar (GAFA) 1986–87, pages 84–106. Springer, 1988.
  * [53] Christos Thrampoulidis, Samet Oymak, and Babak Hassibi.  The gaussian min-max theorem in the presence of convexity.  arXiv preprint arXiv:1408.4837, 2014.
  * [54] Mihailo Stojnic.  Meshes that trap random subspaces.  arXiv preprint arXiv:1304.0003, 2013.
  * [55] Mihailo Stojnic.  Upper-bounding l1-optimization weak thresholds.  arXiv preprint arXiv:1303.7289, 2013.
  * [56] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová.  Learning curves of generic features maps for realistic datasets with a teacher-student model.  Advances in Neural Information Processing Systems, 34:18137–18151, 2021.
  * [57] Urte Adomaityte, Leonardo Defilippis, Bruno Loureiro, and Gabriele Sicuro.  High-dimensional robust regression under heavy-tailed data: Asymptotics and universality.  arXiv preprint arXiv:2309.16476, 2023.
  * [58] Urte Adomaityte, Gabriele Sicuro, and Pierpaolo Vivo.  Classification of heavy-tailed features in high dimensions: a superstatistical approach.  In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  * [59] Benjamin Aubin, Florent Krzakala, Yue Lu, and Lenka Zdeborová.  Generalization error in high-dimensional perceptrons: Approaching bayes error with convex optimization.  Advances in Neural Information Processing Systems, 33:12199–12210, 2020.
  * [60] Francesca Mignacco, Florent Krzakala, Yue Lu, Pierfrancesco Urbani, and Lenka Zdeborova.  The role of regularization in classification of high-dimensional noisy gaussian mixture.  In International conference on machine learning, pages 6874–6883. PMLR, 2020.
  * [61] Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová.  Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions.  Advances in Neural Information Processing Systems, 34:10144–10157, 2021.
  * [62] Luca Pesce, Florent Krzakala, Bruno Loureiro, and Ludovic Stephan.  Are Gaussian data all you need? The extents and limits of universality in high-dimensional generalized linear estimation.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 27680–27708. PMLR, 23–29 Jul 2023.
  * [63] Hugo Cui and Lenka Zdeborová.  High-dimensional asymptotics of denoising autoencoders.  arXiv preprint arXiv:2305.11041, 2023.

##  Appendix A Derivation of Result 3.1

This Appendix details the derivation of Result 3.1, which provides tight
asymptotic formulae (8) for the parameters of the sRF associated to a two-
layer neural network trained with a large gradient step.

###  A.1 First layer weights after one gradient step –

After one large gradient step, the weights $W^{(1)}$ can be decomposed as a
bulk plus a spike in the following way:

$\displaystyle
W^{(1)}=W^{(0)}+\Delta+\frac{\check{\bm{u}}\check{\bm{v}}^{\top}}{\sqrt{p}},$
(31)

where

$\displaystyle\check{\bm{u}}=\eta h_{1}\bm{a}^{(0)},$ (32)
$\displaystyle\check{\bm{v}}=\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}f_{\star}(\bm{x}^{\mu})\bm{x}^{\mu},$
(33)

and

$\displaystyle\Delta=W^{(0)}+\frac{\eta
h_{1}}{n_{0}\sqrt{p}}\mathrm{diag}(\bm{a}^{(0)})\left(\sum\limits_{\mu=1}^{n_{0}}\check{\sigma}(W^{(0)}\bm{x}^{\mu})f_{\star}(\bm{x}^{\mu})\bm{x}^{\mu}\right).$
(34)

We noted $h_{1}$ the first Hermite coefficient of the student activation
$\sigma$ and $\check{\sigma}(\cdot)=\sigma^{\prime}(\cdot)-h_{1}$. As
discussed in [1, 26], the bulk $\Delta$ behaves as an effective noise which
rescales the initial weights $W^{(0)}$; we refer to Sec. B.6 of [2] for
additional rigorous characterization of the bulk $\Delta$. To connect to the
convention used below (7) of normalized $\bm{u}=\bm{1}_{p}$,
$\lVert\bm{v}\lVert=\sqrt{d}$, and a single bulk $W$, one thus needs to assess
the asymptotic limits of
$r=\nicefrac{{\lVert\check{\bm{u}}\lVert\lVert\check{\bm{v}}\lVert}}{{\sqrt{pd}}}$,
$c=\nicefrac{{1}}{{\sqrt{d}}}\lVert W^{(0)}+\Delta\lVert$ and
$\gamma=\nicefrac{{\check{\bm{v}}^{\top}\bm{\theta}}}{{\lVert\check{\bm{v}}\lVert}}$.
This program is carried out in the following subsection.

###  A.2 Spike scale

We first focus on the scale of the spike $r$. We remind here the convention to
assume normalized spike vectors, i.e.
$\lVert\bm{u}\lVert=\lVert\bm{v}\lVert=\sqrt{d}$. Comparing to (32), it
follows that

$\displaystyle r=\frac{1}{\sqrt{pd}}\left\lVert\eta
h_{1}a^{(0)}\right\lVert\times\left\lVert\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}f_{\star}(\bm{x}^{\mu})\bm{x}^{\mu}\right\lVert.$
(35)

Remembering that $\bm{a}^{(0)}=\nicefrac{{\bm{1}_{p}}}{{\sqrt{p}}}$, it
follows that

$\displaystyle\left\lVert\eta
h_{1}\bm{a}^{(0)}\right\lVert=d\tilde{\eta}h_{1}.$ (36)

We now turn to the second term. As a mean of independent vectors, it
concentrates in probability. We thus compute the expectation of the squared
norm:

$\displaystyle\mathbb{E}_{\mathcal{D}_{0}}\left\lVert\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}f_{\star}(\bm{x}^{\mu})\bm{x}^{\mu}\right\lVert^{2}$
$\displaystyle=\frac{1}{n_{0}}\sum\limits_{i=1}^{d}\mathbb{E}_{\bm{x}}[f_{\star}(\bm{x})^{2}x_{i}^{2}]+\frac{(n_{0}-1)}{n_{0}}\sum\limits_{i=1}^{d}\mathbb{E}_{x}[f_{\star}(\bm{x})x_{i}]^{2}$
(37)

Note that $\nicefrac{{\bm{\theta}^{\top}\bm{x}}}{{\sqrt{d}}}$ and $x_{i}$ have
asymptotically vanishing correlation. To leading order, therefore

$\displaystyle\mathbb{E}_{\bm{x}}[f_{\star}(\bm{x})^{2}x_{i}^{2}]=\mathbb{E}_{z}[\sigma_{\star}(z)^{2}]+O_{d}(\nicefrac{{1}}{{\sqrt{d}}})=h_{2}^{\star}$
(38)
$\displaystyle\mathbb{E}_{\bm{x}}[f_{\star}(\bm{x})x_{i}]=\mathbb{E}_{z}[\sigma_{\star}^{\prime}(z)]\frac{\theta_{i}}{\sqrt{d}}=h_{1}^{\star}\frac{\theta_{i}}{\sqrt{d}}$
(39)

Therfore

$\displaystyle\mathbb{E}_{\mathcal{D}_{0}}\left\lVert\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}f_{\star}(\bm{x}^{\mu})\bm{x}^{\mu}\right\lVert^{2}$
$\displaystyle=\frac{dh_{2}^{\star}}{n_{0}}+\frac{(n_{0}-1)h_{1}^{\star
2}}{n_{0}}=\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star 2}.$ (40)

Thus

$\displaystyle
r=\frac{\tilde{\eta}h_{1}}{\sqrt{\beta}}\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}}.$ (41)

###  A.3 Spike/target alignment

We now turn to the spike/target alignment, defined as the overlap between the
(normalized) target weights $\theta$ and spike $v$.

$\displaystyle\gamma=\frac{\bm{\theta}^{\top}\bm{v}\sqrt{d}}{d\lVert\bm{v}\lVert}=\frac{\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}\sigma_{\star}(z^{\mu})z^{\mu}}{\left\lVert\frac{1}{n_{0}}\sum\limits_{\mu=1}^{n_{0}}f_{\star}(x^{\mu})x^{\mu}\right\lVert}.$
(42)

We noted
$z^{\mu}\equiv\nicefrac{{\bm{\theta}^{\top}\bm{x}^{\mu}}}{{\sqrt{d}}}$. The
numerator converges by the law of large numbers to $h^{\star}_{1}$, while the
numerator was computed in the previous subsection. All in all, we reach

$\displaystyle\gamma=\frac{h_{1}^{\star}}{\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}}}$ (43)

###  A.4 Bulk variance

By symmetry, all rows of the bulk $\Delta$ (34) asymptotically share the same
norm $c_{\Delta}$. The norm of a row $w$ of $W=W^{(0)}+\Delta$ thus has
asymptotic norm

$\displaystyle\lVert w_{i}\lVert=\sqrt{1+c_{\Delta}^{2}+2c_{W\Delta}}$ (44)

where

$\displaystyle c_{\Delta}=\left\lVert\frac{\eta
h_{1}}{n_{0}p}\left(\sum\limits_{\mu=1}^{n_{0}}\check{\sigma}(w^{(0)\top}x^{\mu})f_{\star}(x^{\mu})x^{\mu}\right)\right\lVert$
(45) $\displaystyle c_{W\Delta}=\frac{\eta
h_{1}}{n_{0}p}\left(\sum\limits_{\mu=1}^{n_{0}}\check{\sigma}(w^{(0)\top}x^{\mu})f_{\star}(x^{\mu})w^{(0)\top}x^{\mu}\right),$
(46)

where with a slight abuse of notation we dropped the row index $i$, since all
rows are equivalent.

$\displaystyle\mathbb{E}_{\mathcal{D}_{0}}[c_{\Delta}^{2}]=\frac{\eta^{2}h_{1}^{2}}{n_{0}p^{2}}\sum\limits_{i=1}^{d}\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)^{2}\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})^{2}x_{i}^{2}]+\frac{(n_{0}-1)\eta^{2}h_{1}^{2}}{n_{0}p^{2}}\sum\limits_{i=1}^{d}\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})x_{i}]^{2}$
(47)

Because $w^{(0)\top}x,\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}},x_{i}$ all have
asymptotically vanishing $\Theta_{d}(\nicefrac{{1}}{{\sqrt{d}}})$
correlations, the first summand reads to leading order

$\displaystyle\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)^{2}\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})^{2}x_{i}^{2}]=\check{h}_{1}^{2}h_{2}^{\star}$
(48)

The second summand reads, using Stein’s lemma

$\displaystyle\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})x_{i}]$
$\displaystyle=w^{(0)}_{i}\mathbb{E}_{x}[\check{\sigma}^{\prime}(w^{(0)\top}x)\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})]+\frac{\theta_{i}}{\sqrt{d}}\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)\sigma_{\star}^{\prime}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})]$
$\displaystyle=\Theta_{d}(\nicefrac{{1}}{{d}})$ (49)

We used

$\displaystyle\mathbb{E}_{z}[\check{\sigma}(z)]=\mathbb{E}_{z}[\check{\sigma}^{\prime}(z)]=0$
(50)

respectively by definition of $\check{\sigma}$ and because $\sigma$ is odd.
Thus

$\displaystyle\mathbb{E}_{\mathcal{D}_{0}}[c_{\Delta}^{2}]=\frac{\eta^{2}h_{1}^{2}\check{h}_{1}^{2}h_{2}d}{n_{0}p^{2}}+\Theta_{d}(\nicefrac{{1}}{{d}})=\frac{\tilde{\eta}^{2}h_{1}^{2}\check{h}_{1}^{2}h_{2}^{\star}}{\alpha_{0}\beta^{2}}+o_{d}(1).$
(51)

The cross-term $c_{\Delta W}$ can similarly be computed as

$\displaystyle c_{W\Delta}=\frac{\eta
h_{1}}{p}\mathbb{E}_{x}[\check{\sigma}(w^{(0)\top}x)\sigma_{\star}(\nicefrac{{\theta^{\top}x}}{{\sqrt{d}}})w^{(0)\top}x]$
(52)

Again because all variables in the expectation are weakly correlated, and
because the average of $\check{\sigma}$ is zero, we have that $c_{W\Delta}$ is
vanishing to leading order.

$\displaystyle c_{W\Delta}=o_{d}(1)$ (53)

Finally, we thus have

$\displaystyle
c=1+\frac{\tilde{\eta}^{2}h_{1}^{2}\check{h}_{1}^{2}h_{2}^{\star}}{\alpha_{0}\beta^{2}}$
(54)

###  A.5 Summary of the mapping

In summary, we have derived the asymptotic formulae for the equivalent sRF:

$\displaystyle
c=1+\frac{\tilde{\eta}^{2}h_{1}^{2}\check{h}_{1}^{2}h_{2}^{\star}}{\alpha_{0}\beta^{2}}$
(55) $\displaystyle
r=\frac{\tilde{\eta}h_{1}}{\sqrt{\beta}}\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}},$ (56)
$\displaystyle\gamma=\frac{h_{1}^{\star}}{\left(\frac{h_{2}^{\star}}{\alpha_{0}}+h_{1}^{\star
2}\right)^{\nicefrac{{1}}{{2}}}},$ (57)

which recovers equation (8) of Result 3.1.

##  Appendix B Derivation of Result 3.2

In this section we detail the heuristic derivation of Result 3.2, which shows
the asymptotic equivalence of the sRF feature map (12) with the conditional
feature map (13). This mapping is at the core of the exact asymptotic
description detailed in the next appendix.  The equivalent features
$\varphi^{g}(x)$ first appeared in [2] (Theorem 4), where they are formulated
in terms of second-order population statistics of the features
$\varphi(\bm{x})$. In the same asymptotic limit as Result 3.1, these
statistics admit compact asymptotically equivalent closed-form expressions,
which we explicitly derive in the next subsection, thereby yielding the
compact formulation of Result 3.2.

###  B.1 Building the equivalent feature map $\varphi^{g}$

We first remark that, since the data is Gaussian,
$x\sim\mathcal{N}(0,\mathbb{I}_{d})$, we can decompose it as

$\displaystyle x=\kappa+x^{\perp},$ (58)

where $\kappa\equiv\nicefrac{{v^{\top}x}}{{\sqrt{d}}}\sim\mathcal{N}(0,1)$ and
we noted $x^{\perp}\equiv\Pi^{\perp}x$ the component orthogonal to the spike,
with Gaussian statistics $x^{\perp}\sim\mathcal{N}(0,\Pi^{\perp}).$ The sRF
feature map (12) thus reads

$\displaystyle\varphi(x)=\sigma\left(\kappa
1_{p}+\kappa\frac{Wv}{\sqrt{d}}+Wx^{\perp}\right)$ (59)

Conditioning on the projection $\kappa$ along the spike, this corresponds to a
RF feature map with random biases $\kappa
1_{p}+\kappa\nicefrac{{Wv}}{{\sqrt{d}}}$. One can compute the population mean
and covariance of $\varphi_{x}$ along the exact same lines as e.g. [23],
leading to

$\displaystyle\mathbb{E}_{x}[\varphi(x)_{i}]=\tilde{\kappa}^{0}_{i}.$ (60)
$\displaystyle\mathrm{Cov}_{x}[\varphi(x)_{i},\varphi(x)_{j}]=\tilde{\kappa}^{1}_{i}\tilde{\kappa}^{1}_{j}(W)_{i}^{\top}\Pi^{\perp}(W)_{j}+\delta_{ij}(\kappa^{*}_{i})^{2}.$
(61)

where

$\displaystyle\tilde{\kappa}^{0}_{i}\equiv\mathbb{E}_{z}[\sigma(\kappa+\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}+z)]$
(62)
$\displaystyle\tilde{\kappa}^{1}_{i}\equiv\mathbb{E}_{z}[z\sigma(\kappa+\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}+z)]$
(63)
$\displaystyle\tilde{\kappa}^{*}_{i}=\sqrt{\mathbb{E}_{z}[\sigma(\kappa+\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}+z)^{2}]-(\tilde{\kappa}^{1}_{i})^{2}-(\tilde{\kappa}^{0}_{i})^{2}}$
(64)

In other words, the features $\varphi(x)$ share the same second order
statistics as the noisy features

$\displaystyle\varphi^{g}(x)_{i}\equiv\tilde{\kappa}^{0}_{i}+\tilde{\kappa}^{1}_{i}(W)_{i}^{\top}x^{\perp}+\tilde{\kappa}^{*}_{i}\xi_{i}$
(65)

where $\xi_{i}\sim\mathcal{N}(0,1)$ is a stochastic noise. Note that in the
definition of the $\kappa^{0,1,*}_{i}$ coefficients (62),
$\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}=\Theta_{d}(\nicefrac{{1}}{{\sqrt{d}}})$,
so one can expand them as

$\displaystyle\tilde{\kappa}^{0}_{i}\equiv\underbrace{\mathbb{E}_{z}[\sigma(\kappa+z)]}_{\mu_{0}(\kappa)}+\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}\mu_{1}(\kappa)$
(66)
$\displaystyle\tilde{\kappa}^{1}_{i}\equiv\underbrace{\mathbb{E}_{z}[z\sigma(\kappa+z)]}_{\mu_{1}(\kappa)}+\nicefrac{{(W)_{i}^{\top}v}}{{\sqrt{d}}}\underbrace{\mathbb{E}_{z}[z\sigma^{\prime}(\kappa+z)}_{\equiv\mu_{3}(\kappa)}]$
(67)
$\displaystyle\tilde{\kappa}^{*}_{i}=\underbrace{\sqrt{\mathbb{E}_{z}[\sigma(\kappa+z)^{2}]-(\kappa^{1})^{2}-(\kappa^{0})^{2}}}_{\mu_{2}(\kappa)}+O_{d}(\nicefrac{{1}}{{\sqrt{d}}})$
(68)

So

$\displaystyle\varphi^{g}(x)_{i}$
$\displaystyle\equiv\mu_{0}(\kappa)+\mu_{1}(\kappa)(W)_{i}^{\top}x+\mu_{3}(\kappa)\times\frac{(W)_{i}^{\top}v}{\sqrt{d}}\times(W)_{i}x^{\perp}+\mu_{2}(\kappa)\xi_{i}+\xi_{i}\Theta_{d}(\nicefrac{{1}}{{\sqrt{d}}})$
(69)

Finally, note that the $\mu_{3}(\kappa)$ term is subleading compared to the
second term, and should be asymptotically unconsequential for the learning,
and we assume that it can be safely neglected. By the same token, the last
term is a subleading correction to the noise term and is assumed negligible.
These approximations allow to reach the asymptotically equivalent compact
expression

$\displaystyle\varphi^{g}(x)\equiv\mu_{0}(\kappa)1_{p}+\mu_{1}(\kappa)W^{\top}x+\mu_{2}(\kappa)\xi.$
(70)

Putting these assumptions on fully rigorous grounds warrants a more careful
control on how the neglected term impact the learning, and is left for future
work. Finally, from the rigorous result of [2], only the two first
(conditional) moments of the features matter for the asymptotic test error.
Therefore, the features $\varphi(x),\varphi^{g}(x)$ share the same asymptotic
test error, which finishes the derivation of Result 3.2.

##  Appendix C Derivation of Result 3.3

In this Appendix, we detail the derivation of Result 3.3. The derivation
leverages the replica method from statistical physics [15, 16] in its replica-
symmetric formulation, in conjunction with Result 3.2. On a technical level,
the derivation builds upon the formal equivalence between the distribution of
the equivalent features $\bm{\varphi}^{g}(x)$ (13) with a mixture of Gaussians
with infinitely many $\kappa-$indexed clusters, similarly to the
superstatistical approach of [57, 58] in another setting.

First observe that for any test function $\phi(\hat{\bm{a}})$ of the minimizer
$\hat{\bm{a}}$ of (4),

$\displaystyle\phi(\hat{\bm{a}})=\underset{\omega\to\infty}{\lim}\mathbb{E}_{\mathcal{D}}\frac{1}{Z}\int
d\bm{a}\phi(\bm{a})e^{-\omega\mathcal{R}_{\lambda}[\bm{a}]},$ (71)

where we remind that $R[\bm{a}]$ denotes the empirical risk (4), and we
denoted

$\displaystyle Z\equiv\int d\bm{a}e^{-\omega\mathcal{R}_{\lambda}[\bm{a}]}$
(72)

the normalization factor, also known as the partition function in statistical
physics. It is therefore important to the generating function associated to
the measure (71), namely $\mathbb{E}\ln Z$. Such computations can be addressed
using the replica method from statistical physics [15, 16], leveraging the
identity

$\displaystyle\ln Z=\underset{s\to 0}{\lim}\frac{Z^{s}-1}{s}.$ (73)

The backbone of the derivation thus lies in the computation of
$\mathbb{E}Z^{s}$.

###  C.1 Reminder of setting and notations

Because of Result 3.2, we take as a starting point that one can assume the sRF
features $\bm{z}=\bm{\varphi}(x)$ (12) are distributed as

$\displaystyle\bm{z}$
$\displaystyle=\mu_{0}(\kappa)\bm{1}_{p}+\mu_{1}(\kappa)W\bm{x}+\mu_{2}(\kappa)\bm{\xi}$
$\displaystyle=\mu_{0}(\kappa)\bm{1}_{p}+\nicefrac{{1}}{{\sqrt{d}}}\mu_{1}(\kappa)\kappa
W\bm{v}+\mu_{1}(\kappa)W\Pi^{\perp}\bm{x}+\mu_{2}(\kappa)\bm{\xi}$ (74)

where
$\kappa\equiv\frac{\bm{v}^{\top}\bm{z}}{\sqrt{d}}\sim\mathcal{N}(0,1),\bm{u}=\bm{1}_{d},\bm{x}\sim\mathcal{N}(0,\mathbb{I}_{d}),\bm{\xi}\sim\mathcal{N}(0,\mathbb{I}_{p})$.
We remind that

$\displaystyle\mu_{0}(\kappa)$
$\displaystyle=\operatorname{\mathbb{E}}_{w}{\sigma(w+r\kappa)}$ (75)
$\displaystyle\mu_{1}(\kappa)$
$\displaystyle=\frac{1}{c}\operatorname{\mathbb{E}}_{w}{w\sigma(w+r\kappa)}$
(76) $\displaystyle\mu_{2}(\kappa)$
$\displaystyle=\sqrt{\operatorname{\mathbb{E}}_{w}{\sigma^{2}(w+r\kappa)-c(\mu_{1}(\kappa))^{2}-(\mu_{0}(\kappa))^{2}}},$
(77)

where expectations bear over $w\sim\mathcal{N}(0,c)$. Finally, we consider a
generic target function

$\displaystyle
f_{\star}(\bm{x})=\sigma_{\star}(\kappa,\nicefrac{{\bm{\theta}^{\top}\Pi^{\perp}\bm{x}}}{{\sqrt{d}}}).$
(78)

Note that single-index models considered in the main text constitute a special
case of this family of functions, which we consider here for generality. We
look at the empirical risk minimization over $\bm{a}$ of

$\displaystyle
R[\bm{a}]=\sum\limits_{\mu=1}^{n}\ell(f_{\star}(\bm{x}^{\mu}),\nicefrac{{\bm{a}^{\top}\bm{z}^{\mu}}}{{\sqrt{p}}})+g(\bm{a}),$
(79)

for generic convex loss $\ell$ and $\ell_{2}$ regularizer
$g(\cdot)=\nicefrac{{\lambda}}{{2}}\lVert\cdot\lVert^{2}$. We shall provide a
detailed derivation in this generic case, and only specialize to the
particular case of square loss addressed in Result 3.3 at the end of the
derivation.

###  C.2 Replica computation

The replicated partition function reads

$\displaystyle Z^{s}=\int\prod\limits_{r=1}^{s}d\bm{a}_{r}e^{-\omega
g(\bm{a}_{r})}\left[\underbrace{\mathbb{E}_{x}\prod\limits_{r=1}^{s}e^{-\omega\ell(f_{\star}(\bm{x}),\nicefrac{{\bm{a}_{r}^{\top}\bm{z}}}{{\sqrt{p}}})}}_{(\star)}\right]^{n}$
(80)

We now expand $(\star)$.

$\displaystyle(\star)$
$\displaystyle=\mathbb{E}_{\kappa,\bm{x}^{\perp},\bm{\xi}}\prod\limits_{r=1}^{s}e^{-\omega\ell(\sigma_{\star}(\kappa,\lambda_{\star}),\mu_{0}(\kappa)m_{r}+\mu_{1}(\kappa)\kappa\zeta_{r}+\mu_{1}(\kappa)g_{r}+\mu_{2}(\kappa)h_{r})}$
(81)

We defined the summary statistics

$\displaystyle m_{r}=\frac{\bm{a}_{r}^{\top}\bm{1}_{p}}{\sqrt{p}},$
$\displaystyle\zeta_{r}=\frac{\bm{a}_{r}^{\top}W\bm{v}}{\sqrt{dp}},$
$\displaystyle
q^{1}_{rr^{\prime}}=\frac{\bm{a}_{r}^{\top}\Omega\bm{a}_{r^{\prime}}}{p},$
$\displaystyle
q^{2}_{rr^{\prime}}=\frac{\bm{a}_{r}^{\top}\bm{a}_{r^{\prime}}}{p},$ (82)

and local fields

$\displaystyle g_{r}=\frac{\bm{a}_{r}^{\top}W\Pi^{\perp}\bm{x}}{\sqrt{p}},$
$\displaystyle h_{r}=\frac{\bm{a}_{r}^{\top}\bm{\xi}}{\sqrt{p}},$
$\displaystyle\lambda_{\star}=\frac{\bm{\theta}^{\top}\Pi^{\perp}\bm{x}}{\sqrt{d}}$
(83)

where

$\displaystyle\Omega=W\Pi^{\perp}W^{\top}.$ (84)

One can then rewrite $(\star)$ more compactly as

$\displaystyle(\star)=\mathbb{E}_{\kappa}\left[\mathbb{E}_{\lambda_{\star},\\{\lambda_{r}\\}_{r=1}^{s}}e^{-\omega\sum\limits_{r=1}^{s}\ell(\sigma_{\star}(\kappa,\lambda_{\star}),\mu_{0}(\kappa)m_{r}+\mu_{1}(\kappa)\kappa\zeta_{r}+\lambda_{r})}\right]$
(85)

where the fields $\lambda_{\star},\lambda_{r}$ are Gaussian with statistics

$\displaystyle\langle\lambda_{r}\lambda_{r^{\prime}}\rangle=\mu_{1}(\kappa)^{2}q^{1}_{rr^{\prime}}+\mu_{2}(\kappa)^{2}q^{2}_{rr^{\prime}}$
$\displaystyle\langle\lambda_{\star}^{2}\rangle=\frac{\bm{\theta}^{\top}\Pi^{\perp}\bm{\theta}}{d}\equiv\rho$
$\displaystyle\langle\lambda_{r}\lambda_{\star}\rangle=\mu_{1}(\kappa)\frac{\bm{a}_{r}^{\top}W\Pi^{\perp}\bm{\theta}}{\sqrt{pd}}\equiv\mu_{1}(\kappa)\psi_{r}$
(86)

Introducing Dirac’s deltas in their Fourier form to enforce the definitions of
the summary statistics $m_{r},\zeta_{r},q^{1,2}_{rr^{\prime}},\psi_{r}$, the
replicated partition function then reads

$\displaystyle Z^{s}=$
$\displaystyle\int\prod\limits_{r=1}^{s}dm_{r}d\hat{m}_{r}d\zeta_{r}d\hat{\zeta}_{r}d\psi_{r}d\hat{\psi}_{r}\prod\limits_{r,r^{\prime}}dq^{1}_{rr^{\prime}}d\hat{q}^{1}_{rr^{\prime}}dq^{2}_{rr^{\prime}}d\hat{q}^{2}_{rr^{\prime}}\underbrace{e^{-\sqrt{p}\sum\limits_{r}m_{r}\hat{m}_{r}-\sqrt{dp}\sum\limits_{r}(\zeta_{r}\hat{\zeta}_{r}+\psi_{r}\hat{\psi}_{r})-p\sum\limits_{a\leq
b}(q^{1}_{rr^{\prime}}\hat{q}^{1}_{rr^{\prime}}+q^{2}_{rr^{\prime}}\hat{q}^{2}_{rr^{\prime}})}}_{e^{s\omega
p\Psi_{t}}}$
$\displaystyle\underbrace{\int\prod\limits_{r}d\bm{a}_{r}e^{-\omega\sum\limits_{r}g(\bm{a}_{r})}e^{+\sum\limits_{r=1}^{s}\bm{a}_{r}^{\top}(\hat{m}_{r}\bm{1}_{p}+\hat{\zeta}_{r}W\bm{v}+\hat{\psi}_{r}W\Pi^{\perp}\bm{\theta})+\sum\limits_{r\leq
r^{\prime}}(\hat{q}^{1}_{rr^{\prime}}\bm{w}_{r}^{\top}\Omega\bm{w}_{r^{\prime}}+\hat{q}^{2}_{rr^{\prime}}\bm{w}_{r}^{\top}\bm{w}_{b})}}_{e^{s\omega
p\Psi_{w}}}$
$\displaystyle\underbrace{\left[\mathbb{E}_{\kappa}\left(\mathbb{E}_{\lambda_{\star},\\{\lambda_{r}\\}_{r=1}^{s}}e^{-\omega\sum\limits_{r=1}^{s}\ell(\sigma_{\star}(\kappa,\lambda_{\star}),\mu_{0}(\kappa)m^{a}+\mu_{1}(\kappa)\kappa\zeta^{a}+\lambda_{r})}\right)\right]^{n}}_{e^{\nicefrac{{\alpha}}{{\beta}}s\omega
p\Psi_{y}}}$ (87)

We’ll refer to these three terms respectively as the trace, entropic and
energetic terms. Note that all terms in the exponents are a priori
$\Theta_{d}(d)$, thus large. It is therefore possible to compute the integral
over the summary statistics $m_{r},\zeta_{r},q^{1,2}_{rr^{\prime}},\psi_{r}$
and respective conjugate variables using Laplace’s method, thereby reducing
the problem to finding the extremum of $s\Psi_{t}+s\Psi_{w}+\alpha s\Psi_{y}$.

###  C.3 Replica symmetric ansatz

The extremization remains challenging on several levels, notably because it
involves $s^{2}+5s$ variables, and the replica method requires to take the
limit $s\to 0$. To make the problem tractable, we will seek the extremum
assuming it has the replica-symmetric form

$\displaystyle\forall 1\leq r,r^{\prime}\leq
s,~{}~{}q^{\iota}_{rr^{\prime}}=\delta_{rr^{\prime}}(r^{\iota}-q^{\iota})+q^{\iota}$
$\displaystyle\forall 1\leq r\leq s,~{}~{}m_{r}=m$ $\displaystyle\forall 1\leq
r\leq s,~{}~{}\zeta_{r}=\zeta$ $\displaystyle\forall 1\leq r\leq
s,~{}~{}\psi_{r}=\psi$ (88)

for $\iota=1,2$, and similarly for the hat variables:

$\displaystyle\forall 1\leq r,r^{\prime}\leq
s,~{}~{}\hat{q}^{\iota}_{ab}=\delta_{ab}(-\frac{1}{2}\hat{r}^{\iota}-\hat{q}^{\iota})+\hat{q}^{\iota}$
$\displaystyle\forall 1\leq r\leq s,~{}~{}\hat{m}_{r}=\hat{m}$
$\displaystyle\forall 1\leq r\leq s,~{}~{}\hat{\zeta}_{r}=\hat{\zeta}$
$\displaystyle\forall 1\leq r\leq s,~{}~{}\hat{\psi}_{r}=\hat{\psi}$ (89)

Like in e.g. [56], we also introduce the variance order parameters

$\displaystyle V_{\iota}=r_{\iota}-q_{\iota},$
$\displaystyle\hat{V}_{\iota}=\hat{r}_{\iota}+\hat{q}_{\iota}.$ (90)

###  C.4 Computing the free energy

#### Entropic term

The computation of the entropic term follows from the derivation of [59], with
but minor changes. The entropic term thus reads,

$\displaystyle p\omega\Psi_{w}=\mathbb{E}_{\bm{\xi}}\ln\int d\bm{a}e^{-\omega
g(\bm{a})-\frac{1}{2}\bm{a}^{\top}(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{p})\bm{a}+\bm{a}^{\top}\left(\hat{m}1_{p}+\hat{\zeta}W\bm{v}+\hat{\psi}W\Pi^{\perp}\bm{\zeta}+\left(\hat{q}_{1}\Omega+\hat{q}_{2}\mathbb{I}_{p}\right)^{\nicefrac{{1}}{{2}}}\bm{\xi}\right)}$
(91)

In the $\omega\to\infty$ limit, rescaling the conjugate statistics as

$\displaystyle\hat{q}_{\iota}\to\omega^{2}\hat{q}_{\iota},$
$\displaystyle\hat{m},\hat{V}_{\iota},\hat{\zeta},\hat{\psi}\to\omega\hat{m},\omega\hat{V}_{\iota},\omega\hat{\zeta},\omega\hat{\psi},$
(92)

the entropic potential can finally be rewritten in the $\omega\to\infty$ limit
as

$\displaystyle\Psi_{w}=-\mathbb{E}_{\bm{\xi}}\mathcal{M}_{g}(\bm{\xi})$ (93)

where we define the Moreau envelope

$\displaystyle\mathcal{M}_{g}(\bm{\xi})=\underset{\bm{a}}{\inf}\left\\{g(\bm{a})+\frac{1}{2}\bm{a}^{\top}(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{p})\bm{a}-a^{\top}\left(\hat{m}\bm{1}_{p}+\hat{\zeta}W\bm{v}+\hat{\psi}W\Pi^{\perp}\bm{\theta}+\left(\hat{q}_{1}\Omega+\hat{q}_{2}\mathbb{I}_{p}\right)^{\nicefrac{{1}}{{2}}}\bm{\xi}\right)\right\\}$
(94)

For the case $g=\nicefrac{{\lambda}}{{2}}\lVert\cdot\lVert^{2}$, this
simplifies to

$\displaystyle\Psi_{w}=\frac{1}{2p}\Tr\left[\left(\hat{m}^{2}\bm{1}_{p}\bm{1}_{p}^{\top}+\hat{\zeta}^{2}W\bm{v}\bm{v}^{\top}W^{\top}+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}+\hat{q}_{1}\Omega+\hat{q}_{2}\mathbb{I}_{d}\right)\left(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}\right)^{-1}\right]$
(95)

#### Energetic term

The energetic term is the same as for a Gaussian mixture computation as in
e.g. [60, 61, 62, 63] for Gaussian mixture distributions, with the
$\mathbb{E}_{\kappa}$ playing the role of the sum over clusters. In other
words, the equivalent feature map (13) is formally equivalent to an infinite
mixture of Gaussian clusters indexed by the projection $\kappa$ on the spike
$\bm{v}$. Note that the formal equivalence to infinite Gaussian mixtures was
already studied in [57, 58], in the context of learning from heavy-tailed
data. Exploiting this formal equivalence, the expectation over
$\\{\lambda_{a}\\}$ can be carried out and give a Moreau envelope in the
$\omega\to\infty$ limit, yielding

$\displaystyle\Psi_{y}=-\mathbb{E}_{\kappa,\xi,y}\mathcal{M}(\xi,\kappa,y)$
(96)

where $\xi\sim\mathcal{N}(0,1)$ and

$\displaystyle\mathcal{M}(\xi,\kappa,y)=\underset{x}{\inf}\left\\{\frac{1}{2V(\kappa)}(x-\mu_{1}(\kappa)\psi
y-\sqrt{\rho
q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}\xi)^{2}+\ell(\sigma_{\star}(\kappa,y),\mu_{0}(\kappa)m+\mu_{1}(\kappa)\kappa\zeta+x)\right\\}$
(97)

we noted

$\displaystyle q(\kappa)=\mu_{1}(\kappa)^{2}q_{1}+\mu_{2}(\kappa)^{2}q_{2},$
$\displaystyle V(\kappa)=\mu_{1}(\kappa)^{2}V_{1}+\mu_{2}(\kappa)^{2}V_{2}.$
(98)

In (96), all expectations bear over zero-mean, unit variance Gaussian
variables. Note that $q(\kappa),V(\kappa)$ are the exact equivalents of
$q_{k},V_{k}$ in e.g. [61] or [63], i.e. the overlaps associated to cluster
index $\kappa$. As previously mentionned, the equivalent map (13) formally
coincides with an infinite Gaussian mixture, with $\kappa$ indexing the
cluster.

#### Trace term

Note that in contrast to e.g. [60, 61, 63], in the trace term $\Psi_{t}$, the
$\hat{m}m$ term is multiplied by $\sqrt{d}$ instead of $d$. This is due to the
fact that the cluster centroids, whose role is here played by $u=1_{p}$, is of
norm $\sqrt{d}$, whereas [60][61] consider Gaussian clusters separated by
centroids of norm $\Theta_{d}(1)$. The computation of the trace term however
follows identical steps as e.g. [56], yielding

$\displaystyle\Psi_{t}$
$\displaystyle=-\frac{1}{2}(\hat{V}_{1}q_{1}-\hat{q}_{1}V_{1})-\frac{1}{2}(\hat{V}_{2}q_{2}-\hat{q}_{2}V_{2})+\frac{1}{\sqrt{d}}m\hat{m}+\frac{1}{\sqrt{\beta}}\psi\hat{\psi}+\frac{1}{\sqrt{\beta}}\zeta\hat{\zeta}$
$\displaystyle=-\frac{1}{2}(\hat{V}_{1}q_{1}-\hat{q}_{1}V_{1})-\frac{1}{2}(\hat{V}_{2}q_{2}-\hat{q}_{2}V_{2})+\frac{1}{\sqrt{\beta}}\psi\hat{\psi}+\frac{1}{\sqrt{\beta}}\zeta\hat{\zeta}$
(99)

Comparing to [61], the $\hat{m}m$ term in the free energy is thus vanishing
asymptotically.

#### Free energy

Putting everything together, the free energy is

$\displaystyle f=$
$\displaystyle\underset{q_{\iota},V_{\iota},m,\psi,\zeta,\hat{q}_{\iota},\hat{V}_{\iota},\hat{m},\hat{\psi},\hat{\zeta}}{\mathrm{extr}}\Bigg{\\{}-\frac{1}{2}(\hat{V}_{1}q_{1}-\hat{q}_{1}V_{1})-\frac{1}{2}(\hat{V}_{2}q_{2}-\hat{q}_{2}V_{2})+\frac{1}{\sqrt{\beta}}\psi\hat{\psi}+\frac{1}{\sqrt{\beta}}\zeta\hat{\zeta}$
$\displaystyle-\frac{1}{2p}\Tr[\left(\hat{m}^{2}1_{p}1_{p}^{\top}+\hat{\zeta}^{2}Wvv^{\top}W^{\top}+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}+\hat{q}_{1}\Omega+\hat{q}_{2}\mathbb{I}_{d}\right)\left(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}\right)^{-1}]$
$\displaystyle+\frac{\alpha}{\beta}\mathbb{E}_{\kappa,\xi}\mathcal{M}(\xi,\kappa)\Bigg{\\}}$
(100)

The free energy can be simplified, noting that the zero-gradient (saddle-
point) condition on $\hat{m}$ can be written as

$\displaystyle
0\overset{!}{=}\hat{m}\Tr[1_{p}1_{p}^{\top}\left(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}\right)^{-1}],$
(101)

which imposes $\hat{m}=0$. The free energy now simplifies to

$\displaystyle f=$
$\displaystyle\underset{q_{\iota},V_{\iota},m,\psi,\zeta,\hat{q}_{\iota},\hat{V}_{\iota},\hat{\psi},\hat{\zeta}}{\mathrm{extr}}\Bigg{\\{}-\frac{1}{2}(\hat{V}_{1}q_{1}-\hat{q}_{1}V_{1})-\frac{1}{2}(\hat{V}_{2}q_{2}-\hat{q}_{2}V_{2})+\frac{1}{\sqrt{\beta}}\psi\hat{\psi}+\frac{1}{\sqrt{\beta}}\zeta\hat{\zeta}$
$\displaystyle-\frac{1}{2p}\Tr[\left(\hat{\zeta}^{2}Wvv^{\top}W^{\top}+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}+\hat{q}_{1}\Omega+\hat{q}_{2}\mathbb{I}_{d}\right)\left(\hat{V}_{1}\Omega+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}\right)^{-1}]$
$\displaystyle+\frac{\alpha}{\beta}\mathbb{E}_{\kappa,\xi}\mathcal{M}(\xi,\kappa)\Bigg{\\}}$
(102)

Let us massage the expression in the trace further, by explicitly replacing
$\Omega$ by its definition, namely
$\Omega=WW^{\top}-\nicefrac{{1}}{{d}}Wvv^{\top}W^{\top}$. The term in the
trace then becomes

$\displaystyle(a)\\!\\!\equiv\\!\\!\frac{1}{p}\\!\\!\Tr[\\!\\!\left((\hat{\zeta}^{2}\\!\\!-\\!\\!\nicefrac{{\hat{q}_{1}}}{{d}})Wvv^{\top}W^{\top}\\!\\!+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}\\!\\!+\\!\\!\hat{q}_{1}WW^{\top}\\!\\!+\\!\\!\hat{q}_{2}\mathbb{I}_{d}\right)\\!\\!\left(\hat{V}_{1}WW^{\top}\\!\\!-\\!\\!\nicefrac{{\hat{V}_{1}}}{{d}}Wvv^{\top}W+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}\right)^{-1}]$
(103)

The $\nicefrac{{\hat{q}_{1}}}{{d}}$ in the numerator can be safely neglected
as a small correction to $\hat{\zeta}^{2}$. On the other hand, one has to take
care of the spike in the denominator using the Sherman-Morrison lemma. Using
the shorthand

$\displaystyle
S\equiv\hat{V}_{1}WW^{\top}+\hat{V}_{2}\mathbb{I}_{d}+\lambda\mathbb{I}_{d}$
(104)

one can write the denominator as

$\displaystyle(S-\nicefrac{{\hat{V}_{1}}}{{d}}Wvv^{\top}W^{\top})^{-1}=S^{-1}+\frac{\hat{V}_{1}}{d}\frac{S^{-1}Wvv^{\top}W^{\top}S^{-1}}{1-\frac{\hat{V}_{1}}{d}\Tr[Wvv^{\top}W^{\top}S^{-1}]}.$
(105)

The trace term $(a)$ thus decomposes as

$\displaystyle(a)=$
$\displaystyle\frac{1}{p}\Tr[\left(\hat{\zeta}^{2}Wvv^{\top}W^{\top}+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}+\hat{q}_{1}WW^{\top}+\hat{q}_{2}\mathbb{I}_{d}\right)S^{-1}]$
$\displaystyle+\underbrace{\frac{\hat{V}_{1}}{pd}\frac{\Tr[\left(\hat{\zeta}^{2}Wvv^{\top}W^{\top}+\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}+\hat{q}_{1}WW^{\top}+\hat{q}_{2}\mathbb{I}_{d}\right)S^{-1}Wvv^{\top}W^{\top}S^{-1}]}{1-\frac{\hat{V}_{1}}{d}\Tr[Wvv^{\top}W^{\top}S^{-1}]}}_{(b)}$
(106)

We turn to the simplification of the numerator of $(b)$, sequentially
examining each of the three terms composing it. First,

$\displaystyle\hat{\psi}^{2}W\Pi^{\perp}\bm{\theta}\bm{\theta}^{\top}\Pi^{\perp}W^{\top}S^{-1}Wvv^{\top}W^{\top}S^{-1}$
$\displaystyle=\hat{\psi}^{2}\left((\Pi^{\perp}\theta)^{\top}W^{\top}S^{-1}Wv\right)^{2}=O_{d}(d),$
(107)

where we used that $\Pi^{\perp}\theta,v$ are orthogonal, so the squared scalar
product of these vectors, with the uncorrelated matrix in between, should be
asymptotically at most $O_{d}(\sqrt{d})$. Turning to the second term:

$\displaystyle\hat{q}_{1}WW^{\top}S^{-1}Wvv^{\top}W^{\top}S^{-1}$
$\displaystyle=\hat{q}_{1}\lVert W^{\top}S^{-1}Wv\lVert^{2}=O_{d}(d).$ (108)

Finally

$\displaystyle\hat{\zeta}^{2}Wvv^{\top}W^{\top}S^{-1}Wvv^{\top}W^{\top}S^{-1}=\hat{\zeta}^{2}\left(v^{\top}W^{\top}S^{-1}Wv\right)^{2}=\Theta_{d}(d^{2}).$
(109)

In short, only the third term yields a $\Theta_{d}(1)$ contribution in $(b)$,
i.e.

$\displaystyle(b)=o_{d}(1)+\frac{\hat{\zeta}^{2}\hat{V}_{1}}{pd}\frac{\Tr[Wvv^{\top}W^{\top}S^{-1}]^{2}}{1-\frac{\hat{V}_{1}}{d}\Tr[Wvv^{\top}W^{\top}S^{-1}]}$
(110)

Finally, the $(a)$ term can be written be written in fully asymptotic form, by
introducing the limiting distribution $\nu(\varrho,\tau\pi)$ of Assumption
3.2:

$\displaystyle(a)=$ $\displaystyle\int
d\nu(\varrho,\tau,\pi)\frac{\hat{\zeta}^{2}\tau^{2}\varrho+\hat{\psi}^{2}\varrho\pi^{2}+\hat{q}_{1}\varrho+\hat{q}_{2}}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}+\beta\hat{\zeta}^{2}\hat{V}_{1}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}{1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}},$
(111)

thereby yielding the expression for the free energy

$\displaystyle f=$
$\displaystyle\underset{q_{\iota},V_{\iota},m,\psi,\zeta,\hat{q}_{\iota},\hat{V}_{\iota},\hat{\psi},\hat{\zeta}}{\mathrm{extr}}\Bigg{\\{}-\frac{1}{2}(\hat{V}_{1}q_{1}-\hat{q}_{1}V_{1})-\frac{1}{2}(\hat{V}_{2}q_{2}-\hat{q}_{2}V_{2})+\frac{1}{\sqrt{\beta}}\psi\hat{\psi}+\frac{1}{\sqrt{\beta}}\zeta\hat{\zeta}+\frac{\alpha}{\beta}\mathbb{E}_{\kappa,\xi}\mathcal{M}(\xi,\kappa)$
$\displaystyle-\frac{1}{2}\int
d\nu(\varrho,\tau,\pi)\frac{\hat{\zeta}^{2}\tau^{2}\varrho+\hat{\psi}^{2}\varrho\pi^{2}+\hat{q}_{1}\varrho+\hat{q}_{2}}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}-\beta\frac{\hat{\zeta}^{2}\hat{V}_{1}}{2}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}{1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}}\Bigg{\\}}$
(112)

#### Comments on the limiting density $\nu$

Let us briefly comment on the form of the limiting distribution $\nu$. To
build intuition, let us further relax in this paragraph only the strict
normalization on the norm of the rows of $W$, and assume the rows are
independently sampled from $\mathcal{N}(0,c\mathbb{I}_{d})$. Finally, by the
same token, let us assume $\bm{v}\sim\mathcal{N}(0,\mathbb{I}_{d})$ and
$\bm{\theta}\overset{d}{=}\gamma\bm{v}+\mathcal{N}(0,(1-\gamma^{2})\mathbb{I}_{d})$,
so that they are Gaussian vectors with standard deviation $\sqrt{d}$ and
average overlap $\gamma$. Note that relaxing the spherical constraint on all
vectors, and allowing them to be sampled from Gaussian distributions with
matching statistics, is expected to lead to an asymptotically equivalent
problem, and is only considered in this paragraph to ease the discussion. In
this setting, it is straightfoward to see that the measure $\nu$ decomposes as
the simple product measure

$\displaystyle\nu=c\otimes\mu^{\beta}_{\rm MP}\times\check{\nu}$ (113)

where $\mu^{\beta}_{\rm MP}$ denotes the Marcenko-Pastur density with aspect
ratio $\beta$, and the notation $a\otimes\mu$ denotes the pushforward
$(a\otimes\mu)(\lambda)=\mu(a\lambda)$. We denoted $\check{\nu}$ the joint
distribution of the random variables $x,y-\gamma x$ for
$x,y\sim\mathcal{N}(0,1)$ independently.

###  C.5 Saddle-point equations

We explicit the saddle-point equations associated to the optimization problem
(C.4). Define $x^{\star}$ as the minimizer of the Moreau enveloppe, and the
proximal

$\displaystyle\mathrm{prox}(y,\kappa,\xi)=\frac{1}{V(\kappa)}(x^{\star}-\mu_{1}(\kappa)\psi
y-\sqrt{\rho q(\kappa)-\psi^{2}\mu_{1}(\kappa)^{2}}\xi).$ (114)

The saddle-point equations read

$\displaystyle\begin{cases}q_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}\\\
\qquad-\beta\hat{\zeta}^{2}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}+\frac{1}{\beta}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho^{2}}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\left[\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}-1\right]}{\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}\\\
q_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}-\hat{\zeta}^{2}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\left[1-\frac{1}{\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}\right]\\\
V_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
V_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
m=\underset{m}{\inf}\mathbb{E}_{\kappa,y,\xi}\mathcal{M}\\\
\zeta=\hat{\zeta}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\tau^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}+\beta^{\nicefrac{{3}}{{2}}}\hat{\zeta}\hat{V}_{1}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}{1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}}\\\
\psi=\hat{\psi}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\pi^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\end{cases}$
(115)
$\displaystyle\begin{cases}\hat{V}_{1}=-\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y,\xi}\frac{\rho\mu_{1}(\kappa)^{2}}{\sqrt{\rho
q(\kappa)-\psi^{2}\mu_{1}(\kappa)^{2}}}\mathrm{prox}(y,\kappa,\xi)\\\
\hat{q}_{1}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y.\xi}\mu_{1}(\kappa)^{2}\mathrm{prox}(y,\kappa,\xi)^{2}\\\
\hat{V}_{2}=-\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y,\xi}\frac{\rho\mu_{2}(\kappa)^{2}}{\sqrt{\rho
q(\kappa)-\psi^{2}\mu_{1}(\kappa)^{2}}}\mathrm{prox}(y,\kappa,\xi)\\\
\hat{q}_{2}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y.\xi}\mu_{2}(\kappa)^{2}\mathrm{prox}(y,\kappa,\xi)^{2}\\\
\hat{\zeta}=-\frac{\alpha}{\sqrt{\beta}}\mathbb{E}_{\kappa,y,\xi}\kappa\mu_{1}(\kappa)\ell^{\prime}(\sigma_{\star}(\kappa,y),\mu_{0}(\kappa)m+\mu_{1}(\kappa)\kappa\zeta+x^{\star})\\\
\hat{\psi}=\frac{\alpha}{\sqrt{\beta}}\mathbb{E}_{\kappa,y,\xi}\left(\mu_{1}(\kappa)y-\frac{\psi\mu_{1}(\kappa)^{2}}{\rho
q(\kappa)-\psi^{2}\mu_{1}(\kappa)^{2}}\xi\right)\mathrm{prox}(y,\kappa,\xi)\end{cases}$
(116)

###  C.6 Saddle point equations for the square loss

These expressions simplify for the square loss. Define the minimizer
$x^{\star}$ of the Moreau envelope for
$\ell(y,z)=\nicefrac{{1}}{{2}}(y-z)^{2}$:

$\displaystyle
x^{\star}(\kappa,\xi)=\frac{V(\kappa)\left(\sigma_{\star}(\kappa,y)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\kappa\zeta-\mu_{1}(\kappa)\psi\right)+\sqrt{\rho
q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}\xi+\mu_{1}(\kappa)\psi
y}{1+V(\kappa)}.$ (117)

The saddle point for $m$ admits a compact closed form expression. Writing the
saddle point equation for $m$ indeed imposes that

$\displaystyle 0$
$\displaystyle=\partial_{m}\mathbb{E}_{\kappa,\xi}\mathcal{M}(\xi,\kappa)$
$\displaystyle=\mathbb{E}_{\kappa,\xi}\left[\mu_{0}(\kappa)\left(f_{\star}(\kappa)-\mu_{0}(\kappa)m-\frac{V(\kappa)f^{\star}(\kappa)-V(\kappa)\mu_{0}(\kappa)m+\sqrt{q(\kappa)}\xi}{1+V(\kappa)}\right)\right],$
(118)

from which it follows that

$\displaystyle
m\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)^{2}}{1+V(\kappa)}\right]=\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)f_{\star}(\kappa)}{1+V(\kappa)}\right],$
(119)

i.e.

$\displaystyle
m=\frac{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)f_{\star}(\kappa)}{1+V(\kappa)}\right]}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)^{2}}{1+V(\kappa)}\right]}$
(120)

The saddle-point equations thus read

$\displaystyle\begin{cases}q_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}\\\
\qquad-\beta\hat{\zeta}^{2}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}+\frac{1}{\beta}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho^{2}}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\left[\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}-1\right]}{\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}\\\
q_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{\left(\hat{q}_{1}\varrho+\hat{q}_{2}+\hat{\zeta}^{2}\varrho\tau^{2}+\hat{\psi}^{2}\varrho\pi^{2}\right)}{\left(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}\right)^{2}}-\hat{\zeta}^{2}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{(\lambda+\hat{V}_{1}\varrho+\hat{V}_{2})^{2}}\left[1-\frac{1}{\left(1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}\right]\\\
V_{1}=\int
d\nu(\varrho,\tau,\pi)\varrho\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
V_{2}=\int
d\nu(\varrho,\tau,\pi)\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\\\
m=\frac{1}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(\kappa)^{2}}{1+V(\kappa)}\right]}\mathbb{E}_{\kappa,y}\left[\frac{\mu_{0}(\kappa)(\sigma_{\star}(\kappa,y)-\mu_{1}(\kappa)\kappa\zeta)}{1+V(\kappa)}\right]\\\
\zeta=\hat{\zeta}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\tau^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}+\beta^{\nicefrac{{3}}{{2}}}\hat{\zeta}\hat{V}_{1}\frac{\left(\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}\right)^{2}}{1-\beta\hat{V}_{1}\int
d\nu(\varrho,\tau,\pi)\frac{\tau^{2}\varrho}{\hat{V}_{1}\varrho+\hat{V}_{2}+\lambda}}\\\
\psi=\hat{\psi}\sqrt{\beta}\int
d\nu(\varrho,\tau,\pi)\varrho\pi^{2}\frac{1}{\lambda+\hat{V}_{1}\varrho+\hat{V}_{2}}\end{cases}$
(121)
$\displaystyle\begin{cases}\hat{V}_{1}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa}\frac{\rho\mu_{1}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}_{1}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y}\mu_{1}(\kappa)^{2}\frac{b(\kappa,y)^{2}+q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{V}_{2}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa}\frac{\rho\mu_{2}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}_{2}=\frac{\alpha}{\beta}\mathbb{E}_{\kappa,y}\mu_{2}(\kappa)^{2}\frac{b(\kappa,y)^{2}+q(\kappa)-\mu_{1}(\kappa)^{2}\psi^{2}}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{\zeta}=\frac{\alpha}{\sqrt{\beta}}\mathbb{E}_{\kappa,y}\kappa\mu_{1}(\kappa)\frac{b(\kappa,y)}{1+V(\kappa)}\\\
\hat{\psi}=\frac{\alpha}{\sqrt{\beta}}\sqrt{\beta}\mathbb{E}_{\kappa,y}\frac{y\mu_{1}(\kappa)b(\kappa,y)+\psi\mu_{1}(\kappa)^{2}}{1+V(\kappa)}\end{cases},$
(122)

which recovers equations (115) from Result 3.3.

###  C.7 Test error

We conclude this Appendix by deriving the sharp asymptotic characterization
for the test error. Using once more Result 3.2, the test error reads

$\displaystyle\epsilon_{g}$
$\displaystyle=\mathbb{E}_{x}\left(\sigma_{\star}(\kappa,\nicefrac{{\bm{\theta}^{\top}\Pi^{\perp}x}}{{\sqrt{d}}})-\nicefrac{{\hat{a}^{\top}\varphi^{g}(x)}}{{\sqrt{p}}}\right)^{2}$
$\displaystyle=\mathbb{E}_{x,\xi}\left(\sigma_{\star}(\kappa,\nicefrac{{\bm{\theta}^{\top}\Pi^{\perp}x}}{{\sqrt{d}}})-\mu_{0}(\kappa)\frac{\hat{a}^{\top}1_{p}}{\sqrt{p}}-\mu_{1}(\kappa)\kappa\frac{\hat{a}^{\top}Wv}{\sqrt{dp}}-\mu_{1}(\kappa)\frac{\hat{a}^{\top}W\Pi^{\perp}x}{\sqrt{p}}-\mu_{2}(\kappa)\xi\right)^{2}$
$\displaystyle=\mathbb{E}_{\kappa,y,z}\left(\sigma_{\star}(\kappa,y)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\zeta-z\right)^{2}$
(123)

with

$\displaystyle
y,z\sim\mathcal{N}\left[0,\begin{pmatrix}\rho&\mu_{1}(\kappa)\psi\\\
\mu_{1}(\kappa)\psi&q(\kappa)\end{pmatrix}\right].$ (124)

The expression thus simplifies to

$\displaystyle\epsilon_{g}$
$\displaystyle=\mathbb{E}_{\kappa,y}\left[\left(\sigma_{\star}(\kappa,y)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\kappa\zeta-\frac{\mu_{1}(\kappa)\psi}{\sqrt{\rho}}y\right)^{2}+q(\kappa)-\frac{\mu_{1}(\kappa)^{2}\psi^{2}}{\rho}\right]$
(125)

where all expectations now bear on unit variance zero mean Gaussian variables.

##  Appendix D Non-uniform readout initialization

###  D.1 Setting

In this Appendix, we briefly discuss the case where the readout layer is not
initialized as a vector with equal components,
$a^{(0)}=\nicefrac{{1_{p}}}{{\sqrt{p}}}$, but the case where the components of
$a^{(0)}$ can take values in a finite vocabulary $V$, for instance
$\\{-1,0,+1\\}$. We further assume the fraction of all letters are
asymptotically finite (i.e. the number of occurences of a letter $\sigma\in V$
in $a^{(0)}$, divided by $d$, tends to a well-defined finite value as
$d\to\infty$). For simplicity, we provide a sketch of the derivation directly
for the sRF model (7), and do not explicit for conciseness the mapping from
the gradient-trained network thereto, which follows identical steps from the
mapping detailed in Appendix A. More precisely, we consider directly a sRF
with weights

$\displaystyle W+\frac{uv^{\top}}{\sqrt{d}}.$ (126)

In contrast to the main text, we thus do not assume $u\propto 1_{p}$. We
further also allow the norm of the rows of $W$ to take different values
$c(u_{i})^{\nicefrac{{1}}{{2}}}$ depending on the value taken by the component
of $u$ with the corresponding index. This assumption originates from the form
of the $\Delta$ matrix (34), which exhibits such row-wise norm variations
matching that of $a^{(0)}$ (thus $u$ in the equivalent sRF). Finally, for
simplicity and not to complicate the formulae, we assume perfect alignment
with the target $\gamma=1$, and assume aspect ratio $p=d$, i.e. $\beta=1$. The
generic case $0\leq\gamma\leq 1$ can be derived following identical steps to
Appendix C.

###  D.2 Sketch of derivation

For non-uniform $u$, an equivalent feature map can be obtained along the same
lines as for 3.2, as detailed in Appendix B. The sRF feature map has
asymptotically equivalent test error as the equivalent feature map

$\displaystyle\varphi^{g}(x)=\sum\limits_{\sigma\in
V}\mu_{0}^{\sigma}(\kappa)1_{\sigma}+\mu_{1}^{\sigma}(\kappa)\Pi_{\sigma}Wx+\mu_{2}^{\sigma}(\kappa)\Pi_{\sigma}\xi$
(127)

where

$\displaystyle 1_{\sigma}=(\delta_{u_{i},\sigma})_{i=1}^{d}$
$\displaystyle\Pi_{\sigma}=\mathrm{diag}(1_{\sigma})\in\mathbb{R}^{d\times d}$
(128)

and ($z\sim\mathcal{N}(0,1)$ in the following)

$\displaystyle\mu_{0}^{\tau}(\kappa)=\mathbb{E}_{z\sim\mathcal{N}(0,c(\tau))}[\sigma(z+\tau\kappa)]$
(129)
$\displaystyle\mu_{1}^{\tau}(\kappa)=\frac{1}{c(\tau)}\mathbb{E}_{z\sim\mathcal{N}(0,c(\tau))}[z\sigma(z+\tau\kappa)]$
(130)
$\displaystyle\mu_{0}^{\tau}(\kappa)=\sqrt{\mathbb{E}_{z}[\sigma(z+\tau\kappa)^{2}]-\mu_{0}^{\tau}(\kappa)^{2}-c(\tau)\mu_{1}^{\tau}(\kappa)^{2}}$
(131)

Note $\Pi_{\sigma}\Pi_{\tau}=\delta_{\sigma\tau}\Pi_{\sigma}$. Building on
this equivalent map, the test error can by asymptotically characterized using
the replica method, following the same lines as Appendix C. Like in Appendix
C, we need to compute

$\displaystyle(a)$
$\displaystyle=\mathbb{E}_{z,\xi}\prod\limits_{a=1}^{s}e^{-\omega\ell(\sigma_{\star}(\kappa),\sum\limits_{\sigma}\mu_{0}^{\sigma}(\kappa)m^{a}_{\sigma}+\mu_{1}^{\sigma}(\kappa)\kappa\zeta_{\sigma}^{a}+\mu_{1}^{\sigma}(\kappa)g^{a}_{\sigma}+\mu_{2}^{\sigma}(\kappa)h^{a}_{\sigma})}$
(132)

We defined the order parameters and local fields

$\displaystyle m^{a}_{\sigma}=\frac{a_{a}^{\top}1_{\sigma}}{\sqrt{d}}$
$\displaystyle\zeta^{a}_{\sigma}=\frac{a_{a}^{\top}\Pi_{\sigma}Wv}{d}$
$\displaystyle
g^{a}_{\sigma}=\frac{a_{a}^{\top}\Pi_{\sigma}W\Pi^{\perp}x}{\sqrt{d}}$
$\displaystyle h^{a}_{\sigma}=\frac{a_{a}^{\top}\Pi_{\sigma}\xi}{\sqrt{d}}$
(133)

Again, let us introduce the self-overlaps

$\displaystyle(q^{1}_{\sigma\tau})_{ab}=\frac{a_{a}^{\top}\Omega_{\sigma\tau}a_{b}}{d}$
$\displaystyle(q^{2}_{\sigma})_{ab}=\frac{a_{a}^{\top}\Pi_{\sigma}a_{b}}{d}$
(134)

where we denoted

$\displaystyle\Omega_{\sigma\tau}=\Pi_{\sigma}W\Pi W^{\top}\Pi_{\tau}$ (135)

One can then rewrite $(a)$ more compactly as

$\displaystyle(a)=\mathbb{E}_{\kappa}\left[\mathbb{E}_{\\{\lambda_{a}\\}_{a=1}^{s}}e^{-\omega\sum\limits_{a=1}^{s}\ell(\sigma_{\star}(\kappa),m^{a}(\kappa)+\lambda_{a})}\right]$
(136)

where the fields $\lambda_{a}$ are Gaussian with covariance

$\displaystyle\langle\lambda_{a}\lambda_{b}\rangle=q_{ab}(\kappa)$ (137)

where

$\displaystyle
q(\kappa)_{ab}=\sum\limits_{\sigma\tau}\mu_{1}^{\sigma}(\kappa)\mu_{1}^{\tau}(\kappa)q_{\sigma\tau,ab}^{1}+\sum\limits_{\sigma}\mu_{2}^{\sigma}(\kappa)^{2}q^{2}_{\sigma,ab}$
(138) $\displaystyle m^{a}(\kappa)=\sum\limits_{\sigma\in
V}(\mu_{0}^{\sigma}(\kappa)m^{a}_{\sigma}(\kappa)+\mu_{1}^{\sigma}(\kappa)\kappa\zeta^{a}_{\sigma}(\kappa))$
(139)

We again use the RS ansatz form

$\displaystyle\forall\sigma,\tau\in V,~{}~{}\forall 1\leq r,r^{\prime}\leq
s,~{}~{}(q^{\iota}_{\sigma\tau})_{rr^{\prime}}=\delta_{rr^{\prime}}(r^{\iota}_{\sigma\tau}-q^{\iota}_{\sigma\tau})+q^{\iota}_{\sigma\tau}$
$\displaystyle\forall\sigma\in V,~{}~{}\forall 1\leq r\leq
s,~{}~{}(m_{\sigma})_{r}=m_{\sigma}$ $\displaystyle\forall\sigma\in
V,~{}~{}\forall 1\leq r\leq s,~{}~{}(\zeta_{\sigma})_{r}=\zeta_{\sigma}$ (140)

for $\iota=1,2$, and similarly for the hat variables:

$\displaystyle\forall\sigma,\tau\in V,~{}~{}\forall 1\leq r,r^{\prime}\leq
s,~{}~{}(\hat{q}^{\iota}_{\sigma\tau})_{rr^{\prime}}=\delta_{ab}(-\frac{1}{2}\hat{r}_{\sigma\tau}^{\iota}-\hat{q}_{\sigma\tau}^{\iota})+\hat{q}_{\sigma\tau}^{\iota}$
$\displaystyle\forall\sigma\in V,~{}~{}\forall 1\leq r\leq
s,~{}~{}(\hat{m}_{\sigma})_{r}=\hat{m}_{\sigma}$
$\displaystyle\forall\sigma\in V,~{}~{}\forall 1\leq r\leq
s,~{}~{}(\hat{\zeta}_{\sigma})_{r}=\hat{\zeta}_{\sigma}$ (141)

Finally, denoting

$\displaystyle
q(\kappa)=\sum\limits_{\sigma\tau}\mu_{1}^{\sigma}(\kappa)\mu_{1}^{\tau}(\kappa)q_{\sigma\tau}^{1}+\sum\limits_{\sigma}\mu_{2}^{\sigma}(\kappa)^{2}q^{2}_{\sigma}$
(142) $\displaystyle
V(\kappa)=\sum\limits_{\sigma\tau}\mu_{1}^{\sigma}(\kappa)\mu_{1}^{\tau}(\kappa)V_{\sigma\tau}^{1}+\sum\limits_{\sigma}\mu_{2}^{\sigma}(\kappa)^{2}V^{2}_{\sigma}$
(143)

one reaches along identical lines as Appendix C the free energy

$\displaystyle f=$
$\displaystyle-\frac{1}{2}\sum\limits_{\sigma\tau}(\hat{V}^{1}_{\sigma\tau}q^{1}_{\sigma\tau}-\hat{q}^{1}_{\sigma\tau}V^{1}_{\sigma\tau})-\frac{1}{2}\sum\limits_{\sigma}(\hat{V}^{2}_{\sigma}q^{2}_{\sigma}-\hat{q}^{2}_{\sigma}V^{2}_{\sigma})+\sum\limits_{\sigma}\zeta_{\sigma}\hat{\zeta}_{\sigma}$
$\displaystyle-\frac{1}{2d}\Tr[\left(\sum\limits_{\sigma\leq\tau}\hat{q}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{q}^{2}_{\sigma}\Pi_{\sigma}+\sum\limits_{\sigma}\hat{\zeta}^{2}_{\sigma}\Pi_{\sigma}Wvv^{\top}W^{\top}\Pi_{\sigma}\right)\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-1}]$
$\displaystyle+\alpha\mathbb{E}_{\kappa,\xi}\mathcal{M}(\xi,\kappa)$ (144)

where

$\displaystyle\Omega_{\sigma\tau}=\frac{1}{2}\Pi_{\sigma}W\Pi
W^{\top}\Pi_{\tau}+\frac{1}{2}\Pi_{\tau}W\Pi W^{\top}\Pi_{\sigma}.$ (145)

The expression of the Moreau is almost unchanged compared to Appendix C:

$\displaystyle\mathcal{M}(\xi,\kappa)=\underset{x}{\inf}\frac{1}{2V(\kappa)}(x-\sqrt{q(\kappa)}\xi)^{2}+\ell(f^{\star}(\kappa),m(\kappa)+x),$
(146)

where we noted

$\displaystyle m(\kappa)=\sum\limits_{\sigma\in
V}\mu_{0}^{\sigma}(\kappa)m_{\sigma}+\mu_{1}^{\sigma}(\kappa)\kappa\zeta_{\sigma}.$
(147)

The saddle point equations read, for the square loss
$\ell(y,z)=\nicefrac{{1}}{{2}}(y-z)^{2}$:

$\displaystyle\begin{cases}q^{1}_{\sigma\tau}=\frac{1}{d}\Tr[\Omega_{\sigma\tau}\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-2}\left(\sum\limits_{\sigma\leq\tau}\hat{q}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{q}^{2}_{\sigma}\Pi_{\sigma}+\sum\limits_{\sigma}\hat{\zeta}^{2}_{\sigma}\Pi_{\sigma}Wvv^{\top}W^{\top}\Pi_{\sigma}\right)]\\\
q^{2}_{\sigma}=\frac{1}{d}\Tr[\Pi_{\sigma}\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-2}\left(\sum\limits_{\sigma\leq\tau}\hat{q}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{q}^{2}_{\sigma}\Pi_{\sigma}+\sum\limits_{\sigma}\hat{\zeta}^{2}_{\sigma}\Pi_{\sigma}Wvv^{\top}W^{\top}\Pi_{\sigma}\right)]\\\
V^{1}_{\sigma\tau}=\frac{1}{d}\Tr[\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-1}\Omega_{\sigma\tau}]\\\
V^{2}_{\sigma}=\frac{1}{d}\Tr[\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-1}\Pi_{\sigma}]\\\
\zeta_{\sigma}=\frac{\hat{\zeta}_{\sigma}}{d}\Tr[\left(\sum\limits_{\sigma\leq\tau}\hat{V}_{\sigma\tau}^{1}\Omega_{\sigma\tau}+\sum\limits_{\sigma\in
V}\hat{V}^{2}_{\sigma}\Pi_{\sigma}+\lambda\mathbb{I}_{d}\right)^{-1}\Pi_{\sigma}Wvv^{\top}W^{\top}\Pi_{\sigma}]\\\
m_{\sigma}=\frac{1}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}^{\sigma}(\kappa)^{2}}{1+V(\kappa)}\right]}\mathbb{E}_{\kappa}\left[\frac{\mu_{0}^{\sigma}(\kappa)(f^{\star}(\kappa)-\sum\limits_{\tau\neq\sigma}\mu_{0}^{\tau}(\kappa)m_{\tau}-\sum\limits_{\tau}\mu_{1}^{\tau}(\kappa)\kappa\zeta_{\tau})}{1+V(\kappa)}\right]\end{cases}$
(148)
$\displaystyle\begin{cases}\hat{V}^{1}_{\sigma\tau}=\alpha\mathbb{E}_{\kappa}\frac{\mu_{1}^{\sigma}(\kappa)\mu_{1}^{\tau}(\kappa)}{1+V(\kappa)}\\\
\hat{q}^{1}_{\sigma\tau}=\alpha\mathbb{E}_{\kappa}\mu_{1}^{\sigma}(\kappa)\mu_{1}^{\tau}(\kappa)\frac{(f^{\star}(\kappa)-m(\kappa))^{2}+q(\kappa)}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{V}^{2}_{\sigma}=\alpha\mathbb{E}_{\kappa}\frac{\mu_{2}^{\sigma}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}^{2}_{\sigma}=\alpha\mathbb{E}_{\kappa}\mu_{2}^{\sigma}(\kappa)^{2}\frac{(f^{\star}(\kappa)-m(\kappa))^{2}+q(\kappa)}{\left(1+V(\kappa)\right)^{2}}\\\
\hat{\zeta}_{\sigma}=\alpha\mathbb{E}_{\kappa}\mu_{1}^{\sigma}(\kappa)\kappa\frac{(f^{\star}(\kappa)-m(\kappa))^{2}+q(\kappa)}{\left(1+V(\kappa)\right)^{2}}\end{cases}$
(149)

Importantly, these equations are conjectured to be asymptotically exact, but
note that they are not fully asymptotic, in the sense that they still involve
traces of large-dimensional matrices. Finally, the test error admits the
compact asymptotic expression

$\displaystyle\epsilon_{g}$
$\displaystyle=\mathbb{E}_{\kappa}\left[\left(\sigma_{\star}(\kappa)-\sum\limits_{\sigma\in
V}\left[\mu_{0}^{\sigma}(\kappa)m_{\sigma}+\mu_{1}^{\sigma}(\kappa)\kappa\zeta_{\sigma}\right]\right)^{2}+q(\kappa)\right]$
$\displaystyle=\mathbb{E}_{\kappa}\left[\left(\sigma_{\star}(\kappa)-m(\kappa)\right)^{2}+q(\kappa)\right]$
(150)

###  D.3 The special case $W=\Pi^{\perp}$

We briefly discuss a special case of technical interest, for which the saddle-
point equations admit a fully asymptotic formulation, namely when
$W=\Pi^{\perp}$. Then

$\displaystyle\Omega_{\sigma\tau}\approx\delta_{\sigma\tau}\Pi_{\sigma}$ (151)

and the saddle point equations reduce to

$\displaystyle\begin{cases}q_{\sigma}=n_{\sigma}\frac{\hat{q}_{\sigma}}{\left(\hat{V}_{\sigma}+\lambda\right)^{2}}\\\
V_{\sigma}=n_{\sigma}\frac{1}{\hat{V}_{\sigma}+\lambda}\\\
m_{\sigma}=\frac{1}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}^{\sigma}(\kappa)^{2}}{1+V(\kappa)}\right]}\mathbb{E}_{\kappa}\left[\frac{\mu_{0}^{\sigma}(\kappa)(f^{\star}(\kappa)-\sum\limits_{\tau\neq\sigma}\mu_{0}^{\tau}(\kappa)m_{\tau})}{1+V(\kappa)}\right]\end{cases}$
(152)
$\displaystyle\begin{cases}\hat{V}_{\sigma}=\alpha\mathbb{E}_{\kappa}\frac{\mu_{1}^{\sigma}(\kappa)^{2}+\mu_{2}^{\sigma}(\kappa)^{2}}{1+V(\kappa)}\\\
\hat{q}_{\sigma}=\alpha\mathbb{E}_{\kappa}(\mu_{1}^{\sigma}(\kappa)^{2}+\mu_{2}^{\sigma}(\kappa)^{2})\frac{(f^{\star}(\kappa)-m(\kappa))^{2}+q(\kappa)}{\left(1+V(\kappa)\right)^{2}}\end{cases}$
(153)

Where $n_{\sigma}$ is the fraction of components of $u$ taking value $\sigma$.
Fig. 5 represents the theoretical prediction of (152), and displays good
agreement with the corresponding numerical experiments.

Figure 5: Test error for a sRF with $c=1,r=1,\gamma=1$, and ReLU activation
for non-uniform $u$ in the special case $W=\Pi^{\perp}$, when learning from a
single-index model with $\sigma_{\star}=$ReLU. A regularization $\lambda=0.1$
is used to train the readout. The components of $u$ were drawn as independent
Rademacher variables,
$u_{i}\sim\mathrm{Rad}(\nicefrac{{1}}{{2}})=\nicefrac{{1}}{{2}}\delta_{\cdot,1}+\nicefrac{{1}}{{2}}\delta_{\cdot,-1}$.
Solid line: theoretical prediction of (152), for $W=\Pi^{\perp}$. Crosses
represent numerical experiments in dimension $d=2000$, averaged over $10$
runs.

Note that the formulae (152) are further amenable to being rewritten in
functional form, denoting $\nu(x)$ the probability density of the components
of $u$. One can then conjecture in this limit the following limiting equations

$\displaystyle\begin{cases}q(x)=\nu(x)\frac{\hat{q}(x)}{\left(\hat{V}(x)+\lambda\right)^{2}}\\\
V(x)=\nu(x)\frac{1}{\hat{V}(x)+\lambda}\\\
m(x)=\frac{1}{\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(x,\kappa)^{2}}{1+V(\kappa)}\right]}\mathbb{E}_{\kappa}\left[\frac{\mu_{0}(x,\kappa)(f^{\star}(\kappa)-\int
dz\nu(z)\mu_{0}(z,\kappa)m(z)}{1+V(\kappa)}\right]\end{cases}$ (154)
$\displaystyle\begin{cases}\hat{V}(x)=\alpha\mathbb{E}_{\kappa}\frac{(\mu_{2}(x,\kappa)^{2}+\mu_{1}(x,\kappa)^{2})}{1+V(\kappa)}\\\
\hat{q}(x)=\alpha\mathbb{E}_{\kappa}(\mu_{2}(x,\kappa)^{2}+\mu_{1}(x,\kappa)^{2})\frac{(f^{\star}(\kappa)-m(\kappa))^{2}}{\left(1+V(\kappa)\right)^{2}}\end{cases}$
(155)

With the notations

$\displaystyle q(\kappa)=\int
dx\nu(x)(\mu_{1}(x,\kappa)^{2}+\mu_{2}(x,\kappa)^{2})q(x)$ (156)
$\displaystyle V(\kappa)=\int
dx\nu(x)(\mu_{1}(x,\kappa)^{2}+\mu_{2}(x,\kappa)^{2})V(x).$ (157)

As discussed in the main text, the exploration of the case of non-uniform $u$
is of interest, as it affords a richer functional basis to learn the target
$\sigma_{\star}$ (see also discussion below (25)). This study falls well
beyond the scope of the present work, and is left for future work.

##  Appendix E Derivation of the bounds (24) and (25)

Figure 6: Vizualization of the functions $\mu_{0}(r),\tilde{\mu}_{1}(r)$ for
$\sigma=\tanh$ (bottom) and $\sigma=\sin$ (top), for various parameters $r$,
in the perfectly aligned case $\gamma=1,c=1$, as can be achieved for
$\alpha_{0}\to\infty$, see also discussion below equation (25) in the main
text. We remind that the spike strength $r$ cna be tuned via the learning rate
$\tilde{\eta}$. The functions $\mu_{0}(r),\tilde{\mu}_{1}(r)$ provide the
basis of functions that can be used to approximate the target $\sigma_{\star}$
.

We detail – and illustrate– in this Appendix the derivation of the bounds (24)
and (25) on the test error. Again, for definiteness, we consider the case of
perfect spike/target alignment $\gamma=1$ (as can be achieved e.g. by using a
large amount of dta $\alpha_{0}\to\infty$ in the first step).

###  E.1 Derivation of the upper bound (24)

We first focus on the upper bound (24), showing that it can be attained for
$\lambda\to\infty$. In this limit, equations (115) imply that all order
parameters but $m$ are vanishing $q_{1}=q_{2}=V_{1}=V_{2}=\zeta=\psi=0$. Then
the expression for $m$ simplifies as

$\displaystyle
m=\frac{\mathbb{E}_{\kappa}[\mu_{0}(\kappa)\sigma_{\star}(\kappa)]}{\mathbb{E}_{\kappa}[\mu_{o}(\kappa)^{2}]}=\frac{\langle\mu_{0},\sigma_{\star}\rangle}{\langle\mu_{0},\mu_{0}\rangle},$
(158)

where for two univariate functions $f,g$ we denoted the scalar product with
respect to the Gaussian measure:

$\displaystyle\langle
f,g\rangle\equiv\int\frac{1}{\sqrt{2\pi}}e^{-\nicefrac{{x^{2}}}{{2}}}f(x)g(x).$
(159)

In other terms, $m$ is proportional to the projection of the target function
$\sigma_{\star}$ on $\mathrm{span}(\mu_{0})$. Note that $m$ is can also be
rewritten as

$\displaystyle
m=\underset{b_{1}}{\mathrm{arginf}}\lVert\sigma_{\star}-b_{1}\mu_{0}\lVert^{2},$
(160)

where the norm is with respect to the scalar product (159). Finally, from
(23), the test error in the $\lambda\to\infty$ limit admits the compact
expression

$\displaystyle\underset{\lambda\to\infty}{\lim}\epsilon_{g}=\mathbb{E}_{\kappa}[(\sigma_{\star}(\kappa)-m\mu_{0}(\kappa))^{2}]=\underset{b_{1}}{\inf}\lVert\sigma_{\star}-b_{1}\mu_{0}\lVert^{2}.$
(161)

Finally, it follows that

$\displaystyle\underset{\lambda\leq
0}{\epsilon_{g}}\leq\underset{\lambda\to\infty}{\lim}\epsilon_{g}=\underset{b_{1}}{\inf}\lVert\sigma_{\star}-b_{1}\mu_{0}\lVert^{2},$
(162)

which recovers the upper bound (24).

###  E.2 Derivation of the lower bound (25)

We now turn to the lower bound (25). First observe that by the definitions of
$q_{1},\psi,\rho$ (3.3), it follows from the Cauchy-Schwarz inequality that

$\displaystyle\psi\leq\sqrt{q_{1}\rho}.$ (163)

Therefore, for all $\kappa$

$\displaystyle q(\kappa)-\frac{\mu_{1}(\kappa)^{2}\psi^{2}}{\rho}$
$\displaystyle=\mu_{1}(\kappa)^{2}\frac{\rho
q_{1}-\psi^{2}}{\rho}+\mu_{2}(\kappa)^{2}q_{2}\geq 0$ (164)

as a consequence, the test error (23) is lower bounded as

$\displaystyle\epsilon_{g}$
$\displaystyle=\mathbb{E}_{\kappa}\Bigg{[}\Bigg{(}\sigma_{\star}(\kappa)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\kappa\zeta\Bigg{)}^{2}{+q(\kappa)-\frac{\mu_{1}(\kappa)^{2}\psi^{2}}{\rho}}\Bigg{]}$
$\displaystyle\geq\mathbb{E}_{\kappa}\Bigg{[}\Bigg{(}\sigma_{\star}(\kappa)-\mu_{0}(\kappa)m-\mu_{1}(\kappa)\kappa\zeta\Bigg{)}^{2}\Bigg{]}$
$\displaystyle\geq\underset{b_{1},b_{2}}{\inf}\mathbb{E}_{\kappa}\Bigg{[}\Bigg{(}\sigma_{\star}(\kappa)-\mu_{0}(\kappa)b_{1}-\mu_{1}(\kappa)\kappa
b_{2}\Bigg{)}^{2}\Bigg{]},$ (165)

which recovers the lower bound (25). The functions corresponding to the upper
(resp. lower) bound (see (24), (25)) are represented for a
$\sigma_{\star}=\sin$ target in Fig. 3.

###  E.3 Derivation of bounds (26)

Since the bounds (24) and (25) hold for all learning rates $\tilde{\eta}$, the
bounds (26) follow immediately from taking the minimum over $\tilde{\eta}$ on
the right hand side and left hand side of (24) and (25). The functions
$\\{\mu_{0}(r),\tilde{\mu}_{1}(r)\\}_{r\geq 0}$, which form the pool of
functions which can be used to approximate the target function, are
represented for $\sigma=\sin,\tanh$ in Fig. 6.