Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed

Introduction

Explaining the success of deep neural networks in many areas of machine learning remains a key challenge for learning theory. A series of recent theoretical works made progress towards this goal by proving trainability of two-layer neural networks (2LNN) with gradient-based methods (Jacot et al. 2018; Allen-Zhu et al. 2018; Li & Liang 2018; Allen-Zhu et al. 2019; Cao & Gu 2019; Du et al. 2019). These results are based on the observation that strongly over-parameterised 2LNN can achieve good performance even if their first-layer weights remain almost constant throughout training. This is the case if the initial weights are chosen with a particular scaling, which was dubbed the “lazy regime” by Chizat et al. 2019. Going a step further, simply fixing the first-layer weights of a 2LNN at their initial values yields the well-known random features model of Rahimi & Recht 2008; Rahimi & Recht 2009, and can be seen as an approximation of kernel learning (Scholkopf & Smola 2018). This behaviour is to be contrasted with the “feature learning regime”, where the weights of the first layer move significantly during training. Recent empirical studies showed that on some benchmark data sets in computer vision, kernels derived from neural networks achieve comparable performance to neural networks (Matthews et al. 2018; Lee et al. 2018; Garriga-Alonso et al. 2019; Arora et al. 2019; Li et al. 2019; Shankar et al. 2020).

These results raise the question of whether neural networks only learn successfully if random features can also learn successfully, and have led to a renewed interest in the exact conditions under which neural networks trained with gradient descent achieve a better performance than random features (Bach 2017; Yehudai & Shamir 2019; Wei et al. 2019; Li et al. 2020; Daniely & Malach 2020; Geiger et al. 2020; Paccolat et al. 2021; Suzuki & Akiyama 2020). Chizat & Bach 2020 studied the implicit bias of wide two-layer networks trained on data with a low-dimensional structure. They derived strong generalisation bounds, which, when both layers of the network are trained, are independent of ambient dimensions, indicating that the network is able to adapt to the low dimensional structure. In contrast, when only the output layer of the network is trained, the network does not possess such an adaptivity, leading to worse performance. Ghorbani et al. 2019; Ghorbani et al. 2020 analysed in detail how data structure breaks the curse of dimensionality in wide two-layer neural networks, but not in learning with random features, leading to better performance of the former.

We show that even a two-layer neural network with only a few hidden neurons outperforms kernel methods on the classical problem of Gaussian mixture classification. We give a sharp asymptotic analysis of 2LNN and random features on Gaussian mixture classification in the high-dimensional regime where the number of samples $N$ is linearly proportional to the input dimension $D\to\infty$. More precisely:

We analyse 2LNN by deriving a closed set of ordinary differential equations (ODEs), which track the test error of 2LNN with a few hidden neurons $K\!\sim\!O(1)$ trained using one-pass (or online) SGD on Gaussian mixture classification. We thereby extend the classical ODE analysis of Riegler & Biehl 1995; Saad & Solla 1995a where the label $y(x)$ is a function of the input $x$, to a setup where the input is conditional on the label. Solving these equations for their asymptotic fixed point, i.e. taking $t\to\infty$ after $D\to\infty$, yields the final classification error of the 2LNN.

Keeping in mind the high-dimensional limit where the number of samples $N$ is proportional to $D$, We analyse how Gaussian mixtures are transformed under $P$ random features in the regime $P,D\to\infty$ with $\gamma\equiv\nicefrac{{P}}{{D}}$ fixed. In the high-dimensional limit the performance at large $\gamma$ converges to the one of the corresponding kernel (Rahimi & Recht 2008; Rahimi & Recht 2009; El Karoui 2010; Pennington & Worah 2017; Louart et al. 2018; Liao & Couillet 2018), and we can thus recover the performance of kernel learning by taking $\gamma$ large enough.

We compute the asymptotic generalisation of random features on mixtures of Gaussians, which allows us to compare their performance to the performance of 2LNN for various signal-to-noise ratios.

While we do not attempt formal rigorous derivations, to keep the paper readable, our theoretical claims are however amenable to rigorous theorems. In particular the ODEs analysis could be formalised rigorously using the technique of (Wang et al. 2019; Goldt et al. 2019). Our results are valid for generic Gaussian Mixtures with $O(1)$ clusters and we focus on the particular example of the XOR-like mixture in order to make the problematic clear.

2 A paradigmatic example

We compare the performance of a two-layer neural network $\phi_{\theta}$ with parameters $\theta=(K,v,W,g)$,

We emphasise that we study random features in the high-dimensional limit where we let $N,D\to\infty$ with their ratio $t=\nicefrac{{N}}{{D}}\!\sim\!O(1)$ as before, while also letting the number of random features $P\to\infty$ with their ratio $\gamma\equiv\nicefrac{{P}}{{D}}\sim O(1)$ fixed. This regime has been studied in a series of recent works (Lelarge & Miolane 2019; Couillet 2019; Liao & Couillet 2019; Mai & Liao 2019; Deng et al. 2019; Kini & Thrampoulidis 2020; Mignacco et al. 2020a). While we concentrate on random features, we note that we can recover the performance of kernel methods (Rahimi & Recht 2008; Rahimi & Recht 2009) by sending $\gamma\to\infty$. Indeed, as $\gamma=\nicefrac{{P}}{{D}}$ grows, the gram matrix converges to the limiting kernel gram matrix in the high-dimensional regime; detailed studies of the convergence in this regime can be found in (El Karoui 2010; Pennington & Worah 2017; Louart et al. 2018; Liao & Couillet 2018). We can thus recover the performance for any general distance or angle based kernel method, e.g. the NTK of Jacot et al. 2018, by considering $\gamma$ large enough in our computations with random features. Note, however, that this must be done with some care. Our results for random projections are given for $N>P$. As discussed by Ghorbani et al. 2019; Mei et al. 2021, the relevant dimension for random features performances is, rather than $P$, the minimum between $N$ and $P$. Since we focus here in the regime where increasing $P$ beyond $O(D)$, and therefore $\gamma$ beyond $O(1)$, is not allowed. Indeed, we shall see that Lazy training methods such as kernels or random projections require asymptotically $N=O(D^{2})$ samples to beat a random guess, while neural-networks achieves oracle-like performances with only $N=O(D)$ samples.

We provide code to reproduce our plots and solve the equations of Sec. 2.2 at github.com/mariaref/rfvs2lnn_GMM_online.

3 Further related work

Barron 1993 already discussed the limitations of approximating functions with a bounded number of random features within a worst-case analysis. Yehudai & Shamir 2019 construct a data distribution that can be efficiently learnt by a single ReLU neuron, but not by random features. Wei et al. 2019 studied the separation between 2LNN & RF and show the existence of a small ($K\sim O(1)$) network that beats kernels on this data distribution, and study the dynamics of learning in the same mean-field limit as Chizat & Bach 2020 and Ghorbani et al. 2019; Ghorbani et al. 2020. Likewise, Li et al. 2020 show separation between kernels & neural networks in the mean-field limit on the phase retrieval problem. Geiger et al. 2020 investigated numerically the role of architecture and data in determining whether lazy or feature learning perform better. Paccolat et al. 2021 studied how neural networks can compress inputs of effectively low-dimensional data.

Gaussian mixture classification

is a well-studied problem in statistical learning theory, and its supervised version was recently considered in a series of works from the perspective of Bayes-optimal inference (Lelarge & Miolane 2019; Mai & Liao 2019; Deng et al. 2019). Mignacco et al. 2020a; Mignacco et al. 2020b studied the dynamics of stochastic gradient descent on a finite training set using dynamical mean-field theory for the perceptron, which corresponds to the case $K=1,v^{1}=1$ in Eq. (1). Liao & Couillet 2019 and Couillet 2019 studied mixture classification with kernel in an unsupervised setting using random matrix theory.

Dynamics of 2LNN

A classic series of papers by Biehl & Schwarze 1995 and Saad & Solla 1995a studied the dynamics of 2LNN as in Eq. (1) trained using online SGD in the classic teacher-student setup (Gardner & Derrida 1989), where inputs $x$ are element-wise i.i.d. Gaussian variables and labels are obtained from a “teacher” network with random weights. They derived a set of closed ODEs that track the test error of the student (see also Saad & Solla 1995b; Biehl et al. 1996; Saad 2009 for further results and Goldt et al. 2019 for a recent proof of these equations). There have been several extensions of this approach to different data distributions (Yoshida & Okada 2019; Goldt et al. 2020b; Goldt et al. 2020a). All of these works, though, consider the label $y$ as a function of the input $x$, or as a function of a latent variable from which $x$ is generated. Here, we extend this type of analysis to a case where the input is conditional on the label, a point of view taken implicitly by Cohen et al. 2020.

The reduction of the dynamics to a set of low-dimensional ODEs should be contrasted with the “mean-field” approach, where the number of hidden neurons $K$ is sent to infinity while the input dimension $D$ is kept finite. In this limit, the neural networks are still a more expressive function class than the corresponding reproducing kernel Hilbert space (Chizat & Bach 2018; Sirignano & Spiliopoulos 2019; Rotskoff & Vanden-Eijnden 2018; Mei et al. 2018). The evolution of the network parameters in this limit can be described by a high-dimensional partial differential equation. This analysis was used in the aforementioned works by Ghorbani et al. 2019; Ghorbani et al. 2020.

Neural networks for GM classification

where $\mathcal{N}_{\alpha}(x)$ is a multivariate normal distribution with mean $\nicefrac{{\mu^{\alpha}}}{{\sqrt{D}}}$ and covariance $\Omega^{\alpha}$. The index set $\mathcal{S}(y)$ contains all the Gaussians that are associated with the label $y$. We choose the constants $\mathcal{P}_{\alpha}$ such that $q(x,y)$ is correctly normalised. To simplify notation, we focus on binary classification, which can be learnt using a student with a single output unit. Extending our results to $C$-class classification, where the student has $C$ output heads, is straightforward.

The network is trained using stochastic gradient descent on the quadratic error for technical reasons related to the analysis. The update equations for the weights at the $\mu$th step of the algorithm, $\operatorname{d}\!w^{k}_{i}\equiv\left(w^{k}_{i}\right)_{\mu+1}-\left(w^{k}_{i}\right)_{\mu}$, read

2 Theory for the learning dynamics of 2LNN

Since we are training on the quadratic error, the first step of our analysis is to rewrite the prediction mean-squared error $\pmse$ as a sum over the error made on inputs from each Gaussian $\alpha$ in the mixture,

Any average over a Gaussian distribution is a function of only the first two moments of that distribution, so the $\pmse$ can be written as a function of the “order parameters” $M$ and $Q$ and of the $K\!\sim\!O(1)$ second-layer weights $v\!=\!(v^{k})$:

Likewise, the classification error $\epsilon_{c}$ (2) can also be written as a function of the order parameters only: $\mathop{\lim}_{D\to\infty}\epsilon_{c}(\theta)\!\to\!\epsilon_{c}(Q,M,v)$. The order parameters have a clear interpretation: $M^{k}_{\alpha}$ encodes the overlap between the $k$th student node and the mean of the $\alpha$ cluster, and plays a similar role to the teacher-student overlap in the vanilla teacher-student scenario. $Q^{kl}_{\alpha}$ instead tracks the overlap between the various student weight vectors, with the input-input covariance $\Omega^{\alpha}$ intervening. The strategy for our analysis is thus to derive equations that describe how the order parameters $(Q,M,v)$ evolve during training, which will in turn allow us to compute the $\pmse$ of the network at all times.

Dynamics

We derived a closed set of ordinary differential equations that describe the evolution of the order parameters in the case where each Gaussian in the mixture has the same covariance matrix $\Omega$. We proceed here with a brief statement of the equations and deffer the detailed derivation to Sec. B.2. The approach is most easily illustrated with the second-layer weights $v^{k}$. The key idea to compute the average change in the weight $v^{k}$ upon an SGD update (7b), $\operatorname{d}\!v^{k}$, which can be decomposed into a contribution from every Gaussian in the mixture,

where the change $\operatorname{d}\!v^{k}_{\alpha^{+}}$ is obtained directly from Eq. (7b),

The averages that remain to be computed only involve the true label and the local fields $\lambda$. The former is a constant within each Gaussian while the latter are jointly Gaussian. It follows, that also these averages can be expressed in terms of only the order parameters and the equation closes. As we discuss in the appendix, in the high-dimensional limit $D\!\to\!\infty$ the normalised number of samples $t\!\equiv\!\nicefrac{{N}}{{D}}$ can be interpreted as a continuous time, which allows the dynamics of $v^{k}$ to be captured by the ODE (B.27).

where $p_{\Omega}(\rho)$ is the spectral density of $\Omega$, and $q^{kl}(\rho)$ is a density whose time evolution can be characterised in the thermodynamic limit. We relegate the full expression of the equation of motion for $q^{kl}(\rho)$ to Eq. (B.26) of the appendix. Crucially, it involves only averages that can be expressed in terms of the order parameters (9), and hence the equation closes. Likewise, the order parameter $M$ can be rewritten in terms of a density as $M^{\alpha k}=\int\operatorname{d}\!\rho\;p_{\Omega}(\rho)m^{\alpha k}(\rho)$. The dynamics of $m^{\alpha k}$ is described by Eq. (B.21).

Solving the equations of motion

Comparing theory and simulation

On the left of Fig. 2, we plot the evolution of the $\pmse$ (8) and the classification error (2) of a 2LNN with $K=8$ neurons trained on the XOR-like mixture of Fig. 1. We plot the test errors obtained from integration of the order parameters with solid lines, and the same quantities computed using a test set during the simulation with crosses. The agreement between ODE predictions and a single run of SGD is good, even at intermediate system size ($D=1000$). In the App. B.2, we give additional plots for the simulated dynamics of the individual order parameters and find very good agreement with predictions obtained from the ODEs (cf. Fig. 7). Note that although we initialise the weights of the student randomly and independently of the means, there is an initial overlap between student weights and the means of order $1/\sqrt{D}$ due to finite-size fluctuations. To capture this with the ODEs, we initialise them in a regime of weak recovery, where $M^{k}_{\alpha}\!\neq\!0$. For a detailed discussion of the early period of learning up to weak recovery, see Arous et al. 2020.

How 2LNNs learn the XOR-like mixture

A closer look at the learning dynamics on the right of Fig. 2 reveals several phases of learning. There we show the first-layer weight vectors of the 2LNN, projected into the plane spanned by the four means of the mixture, at four different times during training. The regions shaded in red and yellow indicate the decision boundaries of the network, which correspond to the line where the network’s output $\phi_{\theta}(x)$ changes its sign. A 2LNN with $K\geq 4$ neurons can approach the classification error of the oracle (3) if its weight vectors approach the four means, with corresponding second-layer weights. Panel (C) shows that network reaches this configuration. However, this configuration does not minimise the mean-squared error used during training (7), so eventually the weights depart slightly from the means to converge to a solution with lower mean squared error (D). This is confirmed by the inset on the left of Fig. 2, where we see that the average angle of the network weights to the means has a maximum around $t=300$, before decaying slightly at the end of training.

3 Predicting the long-time performance of 2LNN

4 The impact of over-parametrisation

We also studied the effect of over-parametrisation, which we define as the number of additional neurons a student has on top of the $K=4$ neurons that it needs to reach the oracle’s performance on the XOR mixture. We show in the inset of Fig. 4 that over-parametrisation does not improve final performance, since the remaining error of the student is dominated by “spill-over” of points from one mixture into adjacent quadrants. However, over-parametrisation leads to an “implicit acceleration” effect: over-parametrised networks are much more likely to converge to a solution that approaches the oracle’s performance, as we show in the main of Fig. 4. The term “implicit acceleration” was coined by Arora et al. 2018 for similar effects in deep neural networks, and analysed for two-layer networks in the teacher-student setup by Livni et al. 2014; Safran & Shamir 2018. A complete understanding of the phenomenon remains an open problem, which we leave for future work.

Random features on GM classification

To understand the performance of random features on Gaussian Mixtures classification, we analyse the performances of the linear model (5) trained with online SGD with the squared loss on the random features $z$ (4) (Steinwart et al. 2009; Caponnetto & De Vito 2007).

First, we assume that we have enough samples, so that $N\gg P$, and discuss the situation when $N~\ll~P$ later. For any finite $D,P$, running the algorithm up to convergence then corresponds to taking the limit $t\to\infty$. The random features’ weights converge to an estimate $\hat{W}$ which can be computed analytically, see Eq. D.6, and allows to precisely characterise the test error:

As discussed in App. C.3, in the case of ReLU activation function, the feature distribution $p(z_{i})$ is a truncated Gaussian. Hence, at large $D$, $P$, both the mean of $z_{i}$ and the population covariance $\cov(z_{i},z_{j})$ can be obtained analytically in terms of the matrix $F$ and means $\mu$, see Eq. (C.16) and (C.19) for the full result.

We used this formula to obtain precisely the error (14), and the results are shown in Fig. 5. We see that the RF error is a function of $\sigma D^{1/2}/P^{1/4}=\sigma(\nicefrac{{D}}{{\gamma}})^{1/4}$ leading to the conclusion that – as discussed in Fig. 1 – the “transition” from the high to low $\snr$ regime happens when $\sigma^{-1}~\approx~\nicefrac{{{D}^{1/2}}}{{P^{1/4}}}$. This scaling further reveals that $P\approx D^{2}$ features are required in order to obtain good performance. The validity of Eq. 16 is verified in Fig. 9. Reaching this performance, however, requires the number of samples $N$ to be larger than $P$, so $N>O(D^{2})$. In the so-called high-dimensional regime analysed in this paper, where $N\approx D$, such performances remain out of reach. The scaling analysis can be easily generalised; as discussed by Ghorbani et al. 2019; Mei et al. 2021, the relevant dimension for RF performances is, rather than $P$, the minimum between $N$ and $P$.

The classification error is thus a function of $\sigma\nicefrac{{D^{1/2}}}{{\min(N,P)}}^{1/4}$. If $N$ is $O(D)$, then even in the kernel limit when $P\to\infty$, the performance degrades to no more than a random guess as soon as

and therefore for any value of $\sigma$ when $D,N\to\infty$ with fixed $N/D$. In a nutshell, for any fixed $\sigma$, lazy training methods such as random features or kernels will fail to beat a random guess in the high-dimensional limit. This, and the requirement of at least $N=O(D^{2})$ samples to learn, are to be contrasted with the the oracle-like performance achieved by a simple neural net with only $N=O(D)$ samples.

Since the mixture remains a mixture after the application of random features, our main task is to compute the new means and variances of the distribution in the transformed space. We thus focus on transformation of a random variable drawn from a single Gaussian $x_{r}=\mu_{r}/\sqrt{D}+\sigma w_{r}$, where $w_{r}$ is a standard Gaussian, in the kernel, or random feature, space.

For generic activation function the first two moments of the features can be obtained in the well studied low signal-to-noise regime $\snr\sim O(1)$. Key to do so, is the observation that $\nicefrac{{F_{ir}\mu_{r}}}{{D}}\sim O(\nicefrac{{1}}{{\sqrt{D}}})$. The activation function can thus be expanded in orders of $\nicefrac{{1}}{{D}}$ and its action is essentially linear. We define the constants

with the expectation taken over the standard Gaussian random variable $\zeta$. To leading order, the mean and covariance of the features are given by (cf. Sec. C):

This computation immediately reveals the reason random features cannot hope to learn in the low $\snr$ regime: the transformation of the means is only linear; hence a Gaussian mixture that is not linearly separable in input space will remain so even after random features. In other words, if the centres of the Gaussian are too close, the kernel fails to map the data non-linearly to a large dimensional space. In contrast, in the high-$\snr$ regime, where the centres are separated enough, the non-linearity kicks-in and the data becomes separable in feature space.

Relation to kernel methods

The same argument explains the failure of kernel methods: if two centres $\bf x$ and $\bf y$ are close, the kernel function $K({\bf x},{\bf y})$ can be expanded to low order and the kernel is essentially linear, leading to bad performances. The connection can be made explicit using the convergence of random features to a kernel (Rahimi & Recht 2008; Rahimi & Recht 2009):

At low SNR, the constants $a,b,c$ can be obtained from the kernel via

Neural networks vs random features

We now collect our results for a comparison of the performance of 2LNN and RF on the XOR-like mixture from Fig. 1. We look at three different regimes for the $\snr$, illustrated in the first column of Fig. 6. The second column visualises the mixture after the Gaussian random features transformation with $\psi(x)=\max(0,x)$ (4). The third and fourth columns show the evolution of the $\pmse$ of 2LNN and RF, respectively, during training with online SGD. Since overparametrisation does not impact the 2LNN’s performance in these tasks, Sec. 2.3, we train a $K=10$ network to increase the number of runs that converge.

At low $\snr$ (a) the distance of each Gaussian to the origin is $O(1)$ and the standard deviation $\sigma\sim O(1)$ as well. The two-layer neural network learns to predict the correct labels almost as well as the oracle (3). Its performance does not depend on $D$, and using the long-time solution of Sec. 2.3, we can predict its asymptotic error (black line) which agrees well with simulations (crosses). In contrast, random features display an asymptotic error that approaches random guessing as the input dimensions increases (inset). This is clear from Eq. 20: in the large $D$ limit, random features only produce a linear transformation of their input. The XOR therefore remains a XOR in RF space, leading linear regression’s failure to do better than chance.

At high $\snr$ (b), the distance between the clusters scales as $\sqrt{D}$ while the $\sigma$ remains fixed. The asymptotic error of the 2LNN thus decreases with $D$ and the network is able to learn perfectly in the $D\to\infty$ limit (black line). The error of random features also approaches 0 as $D\to\infty$, since the mixture is now well separated in random feature space, too.

We finally consider a regime of mixed $\snr$ (c) where the mixture is well-separated in one dimension, but very close in the other dimension. We achieve this by setting $\mu^{0}_{1}\!~\sim~\!\sqrt{D},\mu^{0}_{2}~\!\sim\!~D$ for the mean of the first mixture, etc. Random features then achieve a non-trivial generalisation error, which can be understood by considering the means of the features $z_{i}$. The large component $\mu_{2}$, induces the activation function to perform a non-linear transformation of the centres and allows for opposite sign centroids to be separated by a hyper-plane in feature space. The small component $\mu_{1}$, causes the distance between opposite sign centroids, which is of order $O(1)$ in input space, to remain of order one in feature space, for all $D$. This leads to a finite generalisation error of RF which remains invariant with increasing input dimension. In this regime, the 2LNN still achieve better performance than the random features, thereby completing the picture we developed in Fig. 1.

Acknowledgements

We acknowledge funding from the ERC under the European Union’s Horizon 2020 Research and Innovation Programme Grant Agreement 714608-SMiLe, from “Chaire de recherche sur les modèles et sciences des données”, and from the French National Research Agency grants ANR-17-CE23-0023-01 PAIL and ANR-19-P3IA-0001 PRAIRIE.

References

Appendix A Summary of Notations

$\text{ }q(x|y)\!=\!\!\sum\limits_{\alpha\in\mathcal{S}(y)}\!P_{\alpha^{\pm}}\mathcal{N}_{\alpha}(x)$ conditional probability of $x$ given the true label $y$

Random Features (RF)

Appendix B Derivation of the dynamical equations

In this appendix, we derive the dynamical equations that describe the dynamics of two-layer neural networks trained on the Gaussian mixture from Sec. 2.2. We first derive a useful Lemma for the averages of weakly correlated random variables B.1, which we we then use in the derivation of the dynamical equations B.2

Here, we show how to compute expectation of functions of weakly correlated variables with non zero mean. The derivation follows the ones of (Goldt et al. 2020b) (see App. A). We extend their computations to include variables with non-zero means.

where we defined the mean of $x$, respectively $y$, as $\bar{x}$, respectively $\bar{y}$ and the covariance matrix:

The weak correlation between $x$ and $y$ is encapsulated in the parameter $\epsilon\ll 1$ while $M_{12}\!\sim\!O(1)$.

Using the above, one can compute the expectations:

The expectations are now taken over the 1-dimensional distributions of $x\!\sim\!\mathcal{N}(0,C_{x})$ and $y\!\sim\!\mathcal{N}(0,C_{x})$.

Similarly, consider the case of three weakly correlated real random variables $x_{i}$, $i=1,2,3$ with mean $\bar{x}_{i}$ and covariance matrix $M_{3}$ such that

One can use an expansion of the joint probability distribution of $\{x_{i}\}$ to linear order in $\epsilon$ to compute three point moments of real valued functions $f,g,h$ as:

In the case in which $x_{1}$ and $x_{2}$ are weakly correlated with $x_{3}$ but not between each other, i.e. $\operatorname{Cov}(x_{1},x_{2})=M_{12}\!\sim\!O(1)$, one has:

B.2 Derivation of the ODEs

In order to track the training dynamics, we analyse the evolution of the macroscopic operators defined in Eq. (9) allowing to compute the performances of the network at all training times.

At the $s$th step of training, the SGD update for the networks parameter is given by Eq. (7):

In order to guarantee that the dynamics can be described by a set of ordinary equations in the $D\to\infty$ limit, we choose different scalings for the first and second layer learning rates:

To make progress, consider the eigen-decomposition of the covariance matrix:

where we denote the eigenvalues as $\rho_{\tau}$, their corresponding eigenvector as $\Gamma_{\tau}$ and the eigenvalue distribution as $p_{\Omega}$. We further define the projection of the weights into the projected basis as

The expectation of this update over the distribution Eq. (6) is given by:

where we decomposed the expectation into the different clusters and introduced:

with the expectations $\mathcal{A}^{kj}_{\tau},\mathcal{B}^{k}_{\tau}$ and $\mathcal{C}^{k}_{\tau}$ defined as:

Thus, we can compute the expressions Eq. (B.15) using the proposition for weakly correlated variables derived in App. B.1. This gives:

Update of the Order parameters

In order to derive the update equations for the order parameters, we introduce the densities $m(\rho,t)$ and $q(\rho,t)$. These depend on $\rho$ and on the normalised number of steps $t=\nicefrac{{\mu}}{{D}}$, which we interpret as a continuous time variable.

The equation of motion of $m$ can can be easily computed using the update (B.14) and is given by:

Note how, in order to close the equation, we introduced an additional order parameter $T^{\alpha\beta}=\sum_{r=1}^{D}\frac{\mu^{\alpha}_{r}\mu^{\beta}_{r}}{D}$, which is entirely defined by the overlap of the means of the mixture under consideration and is therefore a constant of motion. For compactness, we defined the multidimensional integrals $I$ of the activation function over the local fields as:

The update of $q$ can similarly be decomposed as a sum over the different Gaussian clusters:

Let us define the constant of motion $\chi^{\alpha}=\frac{1}{D}\sum_{\tau}\rho_{\tau}^{2}$, then the quadratic term in the update for $q_{\alpha^{\pm}}$ is given by:

The multidimensional integrals $I$ are given by:

Finally, the full equation of motion of $q$ is written:

Update for the second layer weights

The update of the second layer weights is also decomposed into the contribution of the different Gaussian clusters and follows from taking the expectation of Eq. (7b) on the GM distribution (6):

Equations (B.21),(B.26) and (B.27) suffice to fully characterise the training dynamics, in the limit of high dimensions and online-learning, of a 2LNN trained on an arbitrary Gaussian mixture with $O(1)$ clusters each having mean $\mu^{\alpha}$ and same covariance matrix $\Omega$.

Agreement with Numerical Simulations

Here, we verify the agreement of the ODEs derived above with simulation of 2LNN trained via online SGD.

Note that the equations of motion describe the evolution of the densities $m$ and $q$ averaged over the input distribution. The agreement between this evolution and simulations justifies, at posteriori, the implicit assumption that the stochastic part of the SGD increment (7) can be neglected in the $D\to\infty$ limit. We can thus conjecture that in the $D\to\infty$ limit, the stochastic process defined by the SGD updates converges to a deterministic process parametrised by the continuous time variables $t\equiv\nicefrac{{N}}{{D}}$. We further add that the proof of this conjecture is not a straight-forward extension of the one of Goldt et al. 2019 for i.i.d. inputs since here, one must take into account the density of the covariance matrix.

The ODEs are valid for generic covariance matrix and means. Thus, they can be used to analyse the role of data structure in training 2LNNs. Although we leave a detailed analysis for future work, Fig. 8 gives an example of how this could be done in the case where a $K=3$ 2LNN is trained on a GM obtained from the FashionMnist dataset. The GM is obtained by computing the means $\bar{x}^{\alpha}$ and covariance matrix $\cov_{\alpha}(x,x)$ of each class in the dataset and assigning a label $+1$ or $-1$ to the different classes, as is commonly done in binary classification tasks. One could, for example, assign label $y=+1$ to the sneakers, boots, sandals, trousers and shorts categories and $y=-1$ to all others. Extending our analysis to $C$-class classification is straight forward and follows the analysis of Yoshida et al. 2019. The inputs are then sampled from a GM where the cluster’s mean are given by $\bar{x}^{\alpha}$ and the covariance matrix $\Omega$ is the mean covariance of all classes: $\Omega=\nicefrac{{1}}{{\text{n}_{\text{classes}}}}\sum\limits_{\alpha}\cov_{\alpha}(x,x)$. Note the similarity between this procedure and linear discriminant analysis commonly used in statistics. The agreement between simulations and analytical predictions is again very good, both at the level of the test error and of the order parameters.

B.3 Simplified ansatz to solve the ODEs for the XOR-like mixture

Here, we detail the procedure, introduced in Sec. 2.3, used to find the long time $t\to\infty$ performance of 2LNN by making an ansatz on the form of the order parameters that solve the fix point equations. The motivation for doing so, as argued in of the main text, is that integrating the ODEs is numerically expensive as it requires evaluating various multidimensional integrals and the number of equations to integrate scales as $K^{2}$. In order to extract information about the asymptotic performances of the network, one can look for a fix point of the ODEs. However, the number of coupled equations to be solved, also scales quadraticaly with $K$ and is already $26$ for a $K=4$ student. The trick is to make an ansatz, with fewer degrees of freedom, on the order parameters that solve the equations. Used in this way, the ODEs have generated a wealth of analytical insights into the dynamics and the performance of 2LNN in the classical teacher-student setup (Biehl & Schwarze 1995; Saad & Solla 1995b; Saad & Solla 1995b; Biehl et al. 1996; Saad 2009; Yoshida & Okada 2019; Yoshida et al. 2019; Goldt et al. 2020b). In all these works though, an important simplification occured because the means of the local fields were all zero by construction. This simplification allowed the fixed points to be found analytically in some cases. Here, the means of the local fields evaluated over individual Gaussians in the mixture are not zero, so we have to resort to numerical means to find the fixed points of the ODEs.

This decomposition, fully constrains the overlap matrix $Q^{kl}$ in terms of $M^{\alpha k}$:

where we used that in the XOR-like mixture, $|\mu^{\alpha}|=\sqrt{D}$ and $\mu^{0+}\cdot\mu^{0-}=0$. From the symmetry between the positive and negative sign clusters of the mixture, in the fix point configuration, for every weight having norm $|w|$ and at an angle $\alpha$ with the mean of a positive cluster, there is a corresponding weight of the same norm, at an angle $\alpha$ with a negative mean. I.e. the angles of the weight vectors $w^{k}$ to the means $\mu^{\alpha}$, as well as the norms of the weights, are $2\times 2$ equal (one for the positive sign cluster and the other for the negative sign one). This constrains further half the number of free parameters in the overlap matrix $M$, which are down to $4K/(2\times 2)=K$. The second layer weights $v$ are fully constrained by requiring the output of the student to be $\pm 1$ when evaluated on the means. Putting everything together, one is left with $K$ equations to solve for the $\nicefrac{{K}}{{2}}$ angles and the $\nicefrac{{K}}{{2}}$ norms, or equivalently, for the $K$ free parameters in the overlap matrix $M$. The agreement between the solution found by solving this reduced set of equations and simulations is displayed both in Fig. 3, where we use it to predict the evolution of the test error with the $L2-$regularisation constant.

Appendix C Transforming a Gaussian mixture with random features

Crucially, the distribution of the features $z$ is still a mixture of distributions. We can thus restrict to studying the transformation of a Gaussian random variable

where $w_{r}$ is a standard Gaussian. The scaling of $\mu_{r}$ and $\sigma$ is chosen according to which regime (low or high $\snr$) one chooses to study. We aim at computing the distribution, in particular the two first moments, of the feature $z$ defined in Eq. (C.1). By construction, the random variables $u_{i}$ are Gaussian with first two moments:

C.2 Low signal-to-noise ratio

Here we compute the statistics of the features, for general activation function, in the low signal to noise regime, for which $\nicefrac{{|\mu|}}{{\sqrt{D}}}\sim O(1)$ and $\sigma\sim O(1)$ so that the Gaussian clusters are a distance of order 1 away from the origin.

For the covariance matrix, we separate the computation of the diagonal from the off-diagonal. Starting with the diagonal elements:

In order to compute the off-diagonal elements, we note that different components of $u$ are weakly correlated since $\operatorname{Cov}(u_{i},u_{j})=\sigma^{2}\sum_{r}\nicefrac{{F_{ir}F_{jr}}}{{D}}\!\sim\!O(1/\sqrt{D})$. We can therefore apply formula Eq. (B.1) for weakly correlated variables:

We define the constants $a$, $b$ and $c$ as in Eq. (18) and $d$ as:

These definitions together with Eq. (C.8), Eq. (C.9) and Eq. (C.12) lead to the statistics of Eq. (19) and Eq. (20):

C.3 ReLU features

In the case of Relu activation function i.e. $\psi(x)=\max(0,x)$, the mean and the covariance of the features can be evaluated analytically for all $\snr$ regimes. The distribution of the features within each cluster is given by a modified Gaussian: the probability mass of the Gaussian on the negative real axis is concentrated at the origin while the distribution on the positive axis is unchanged.

In particular, the integral to obtain the mean of $z$ can be computed analytically and is given by:

The covariance is once again computed by separating the diagonal terms from the off-diagonal ones. The integral to obtain the diagonal terms has an analytical expression found to be:

where the expectations above are over one dimensional distributions $u_{i}\sim\mathcal{N}\left(\sum_{r}\frac{F_{ir\mu_{r}}}{D},\sigma^{2}\right)$. The integrals have an analytical closed form expression, which yields the final result for the covariance:

C.4 Relation with the kernel

As discussed in Sec. 3 of the main text , the performances of kernel methods can be studied by using the convergence of RF to a kernel in the $\gamma\to\infty$ limit taken after the $D,P\to\infty$ limit (Rahimi & Recht 2008):

Finally, for $b$ one has to perform a linear expansion of the kernel around the noise variable $\sigma\omega_{\begin{subarray}{c}1\\ 2\end{subarray}}$:

These expressions allow to express the statistical properties of the features $z$, and to asses the performance of RF and kernel methods, directly in terms of the kernel without requiring the explicit form of the activation function.

For completeness, we give the analytical expression of the kernel corresponding to ReLU random features, i.e. $\psi(x)=\max(0,x)$:

From Eq. (C.23), one sees that in case of ReLU activation function, the kernel is an angular kernel, i.e. it depends on the angle between $x$ and $y$.

Appendix D Final test error of random features

This section details the computations leading to Eq. (14) and Eq. (16) allowing to obtain the asymptotic performances of RF trained via online SGD on a mixture of Gaussian distribution.

The expectation of the SGD update over the distribution of $z$ is thus:

Importantly, both the $\pmse$ and the average update only depend on the distribution of the features through the covariance matrix of the features $\Omega^{z}$ and the input label covariance $\Psi$.

To make progress, consider the eigen-decomposition of the covariance matrix $\Omega^{z}$:

where $\rho_{\tau}$ are the eigenvalues and $\Gamma_{\tau}$ their corresponding eigenvectors.

Define the rotation of $W$ and $\Phi$ into this eigenbasis:

Rotating back in the original basis one finds the asymptotic solution for $W$ as:

The asymptotic test error is thus given by:

From the solution of Eq. (D.6) for the asymptotic solution found by linear regression, one can obtain the asymptotic classification error performed by random features as:

These moments can be computed analytically from the statistics of the features computed in Sec. C and from the optimal weights $W$ obtained in Eq. (D.6). The classification error, Eq. (D.8), can thus be evaluated by means of a one dimensional integral over the distribution of $\lambda$.

Appendix E The three-cluster model

Similar to the analysis of the XOR-like mixture of Fig. 6, we analyse a data model with three clusters that was the subject of several recent works (Deng et al. 2019; Mai & Liao 2019; Lelarge & Miolane 2019; Mignacco et al. 2020a; Mignacco et al. 2020b). The Gaussian mixture in input space can be seen in the first column of Fig. 10. The means of both positive clusters are set to $0$ while the means of the negative sign clusters have first component $\pm\mu_{0}$ and all other $D-1$ components $0$. The mixture after random feature transformation is displayed in the second column and the third and fourth column show the performance of a 2LNN, respectively, a random feature network, trained via online SGD, on this problem. Here again, we build on the observation that overparametrisation does not impact performances and train a $K=10$ 2LNN in order to increase the number of runs that converged. The three rows, are as before, three different $\snr$ regimes, they are in order the low, high and mixed $\snr$ regime.

The phenomenology observed in the XOR-like mixture carries through here. In the low $\snr$ regime, $\mu_{0}=\sqrt{D}$ (top row), the 2LNN can learn the problem and its performance remains constant with increasing input dimension. On the other hand, in this regime, the transformation performed by the random features is only linear in the large $D$ limit. Consequently, the RF performances degrade with increasing $D$ and are as bad as random chance in the limit of infinite input dimension. In the high $\snr$ regime instead (second row), where $\mu_{0}=D$, the mixtures becomes well separated in feature space allowing RF to perform well. Both the performance of 2LNN and RF improve as the clusters in the mixture become more separated. The mixed $\snr$ regime (bottom row) is obtained by setting one of the negative sign clusters a distance $\sqrt{D}$ from the origin while maintaining the other one at a distance $1$. Here, the random feature perform a non trivial transformation of the far away cluster while its action on the nearby cluster is linear. Hence, in feature space, one of the negative clusters remains close to the positive clusters while the other is well separated. The RF thus achieve a test error which is better than random but still worse that that of the 2LNN. Its error is constant with increasing $D$ since it is dominated by the “spill-over” of the negative cluster into the positive cluster at the origin.

Lastly, let us comment that in all our work, we did not add a bias to the model. Adding a bias, does not change the conclusion that small 2LNN considerably outperform RF. In fact, the learning curves are only slightly modified. This is due to to our minimisation of the $\pmse$ when training the network, which, unlike classification loss that only cares about the sign of the estimate, penalises large differences between label $y=\pm 1$ and the output. For simplicity, we thus chose to remove the bias in our analysis, although including it is a straight forward operation.