1 Introduction↩︎

A major challenge in machine learning is the accurate theoretical understanding of learning trajectories in large predictive models. Typically, model training is performed by some variant of gradient descent, which is not immediately tractable. Theoretical analysis of dynamics commonly involves some simple, solvable models, or reductions of complex models to simpler models under suitable assumptions, or both of these elements [1].

Important examples of solvable models are linear models and multi-layer linear neural networks with suitably aligned initialization [2], [3]. Important examples of reductions to tractable models are model linearizations in “lazy training” scenarios [4], such as the NTK regime [5], [6]. Solvable linear networks can result from random initializations in suitable extreme regimes [7]. Some other notable general methods to achieve tractability include applications of random matrix theory [8], methods of statistical physics such as the replica method [9], and mean-field theory [10]–[12].

Usually, this kind of research assumes a particular extreme learning regime (which, of course, may be relevant for a broad family of models – cf. the NTK regime). We, however, ask a different question: Given a particular family of models, what are all its extreme regimes, and in which of them is the learning dynamics theoretically tractable?

In this paper we show that this question can be reasonably completely answered for weight-tied linear autoencoder [13], [14]. This is a relatively simple model which, however, is nonlinear in the weights and does not have a general theoretical solution.

Theoretical studies typically consider autoencoders in two distinct regimes: undercomplete and overcomplete. The former have the hidden dimension smaller than the input one; this way, it is argued that the model is forced to perform dimensionality reduction; in particular, all global minima of an undercomplete linear autoencoder correspond to performing a PCA [13].

On the other hand, overcomplete autoencoders are always able to fit the identity map perfectly, hence do not have to perform any kind of PCA. Still, if the training dataset is smaller than the input dimension then there are multiple ways of mapping the training data to itself, apart from the perfect identity map. The exact weight configuration our training algorithm chooses directly affects the generalization performance of the learned model. This underlines, in particular, the importance of studying the train and population learning trajectories in this model.

While autoencoders have been analyzed both in the undercomplete [15] and overcomplete [16] regimes, we are not aware of any systematic studies of the full picture of their regimes, except for [17] (see 7 for an overview on training regime classification studies). That work addresses even a broader family of tensor models of different orders, but only considers the population-wide learning, thus not addressing any generalization-related questions. The regime classification proposed in [17] relies on a “diagram expansion” of the loss evolution. Different regimes correspond to different families of diagrams, which can be systematically derived depending on the scaling relations between the hyperparameters of the model.

1.0.0.1 Our contributions.

1. We generalize the diagram-based method of [17] to training on finite training sets, thus allowing to separately study train and population learning trajectories. We achieve this by introducing and analyzing data-related edges and nodes in the diagrams.

2. Using this method, we theoretically derive a hierarchy of extreme learning regimes in the linear autoencoder characterized by the input and latent dimensions, weight initialization magnitude, and training set size. In particular, by examining the formal loss-expansion hierarchy, we argue that this model has five basic theoretical extremes associated with the 2-faces of a triangular prism: (1) large-data, (2) small-data, (3) mean-field, (4) narrow-latent, (5) free. Each regime is characterized by specific scaling relations between the four hyperparameters. Edges and vertices of the prism correspond to more degenerate regimes obtainable by combining the five basic ones.

3. In four of the five basic extremes (except the free regime), we derive explicit limiting descriptions of the train and population loss evolutions: a closed-form formula in the large-data regime, a Marchenko-Pastur integral formula in the mean-field regime, a finite-dimensional ODE characterization in the narrow-latent regime, and a moment hierarchy in the small-data regime. The solutions agree very well with experiments. In the large-data regime (1) the solution was known (e.g., found by another method in [17]), but in the other three regimes (2), (3), (4) they are, to the best of our knowledge, new.

2 Problem statement↩︎

We consider a shallow linear weight-tied autoencoder with input dimension \(p\) and latent dimension \(n\): \(f({\boldsymbol{x}}) = {\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{x}}\), where \({\boldsymbol{x}}\in \mathbb{R}^p\), \({\boldsymbol{U}}\in \mathbb{R}^{n \times p}\). Assuming the data distribution is isotropic Gaussian, consider a training set of size \(m\) and define the train and population square losses for this model: \[\label{eq:loss95def} \widehat L({\boldsymbol{U}}) = \frac{\left\|{\boldsymbol{X}}- {\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{X}}\right\|_F^2}{2 p m}, \qquad L({\boldsymbol{U}}) = \mathbb{E}_{{\boldsymbol{x}}} \left[\frac{\left\|{\boldsymbol{x}}- {\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{x}}\right\|^2}{2 p}\right] = \frac{\left\|{\boldsymbol{I}}- {\boldsymbol{U}}^\top {\boldsymbol{U}}\right\|_F^2}{2 p},\tag{1}\] where \({\boldsymbol{x}}\sim {\mathcal{N}}(0, {\boldsymbol{I}}_p)\), each column of \({\boldsymbol{X}}\in \mathbb{R}^{p \times m}\) is sampled independently from \({\mathcal{N}}(0, {\boldsymbol{I}}_p)\), and \(\|{\boldsymbol{A}}\|_F^2=\mathop{\mathrm{Tr}}[A^\top A].\) We train our model with gradient flow with learning rate \(\eta\): \[\label{eq:gf} \frac{d{\boldsymbol{U}}}{dt} = -\eta \frac{\partial \widehat L({\boldsymbol{U}})}{\partial{\boldsymbol{U}}}, \qquad {\boldsymbol{U}}(0) \sim {\mathcal{N}}(0, \sigma^2 {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p).\tag{2}\] Both \(\eta\) and \(\sigma^2\) may depend on \(p\), \(m\), \(n\). Define the loss values at time \(t\): \[\label{eq:loss95def95t} \widehat L(t) = \widehat L({\boldsymbol{U}}(t)), \qquad L(t) = L({\boldsymbol{U}}(t)).\tag{3}\] We will be interested in the average loss evolution in the limit of large \(p,n,m\): \[\label{eq:average95losses} \widehat{{\mathcal{L}}}(t)= \lim_{p,n,m\to\infty} \mathbb{E}[\widehat L(t)],\quad {{\mathcal{L}}}(t)=\lim_{p,n,m\to\infty} \mathbb{E}[L(t)].\tag{4}\] Here and in the sequel, the expectation is taken w.r.t. \({\boldsymbol{U}}(0)\) and \({\boldsymbol{X}}\). We will see that there are multiple limiting learning regimes depending on the joint scaling of \(p\), \(n\), \(m\), and \(\sigma^2\), and will describe the complete hierarchy of these regimes. Next, we will examine the availability of explicit solutions for \(\widehat{{\mathcal{L}}}(t), {{\mathcal{L}}}(t)\). While they don’t seem to exist in general, we will see that (to a varying degree of explicitness) they can be found in the extreme learning regimes.

3 Classification of learning regimes↩︎

We generally follow the small-\(t\)-expansion approach to classification of learning regimes proposed in [17]. This method treats learning trajectories indirectly, on the level of expansion coefficients, but it is systematic and produces a consistent and clear geometric picture of learning regimes. See additional discussion in 8.3.

3.0.0.1 Loss expansion (8.1).

Our starting point is power series expansions of the losses at \(t=0\):

The polynomials \(Y_s, \widehat Y_s\) are complex, but we sketch how they can be constructively described in terms of suitable diagrams (see 8.1 for details). Observe first that the values \(Y_s, \widehat Y_s\) can be found by time-differentiating the losses, e.g. for the population loss and \(s\ge 1\) we have \(Y_s \eta^s = \mathbb{E}[\tfrac{d^s L}{dt^s}(t=0)].\) These time derivatives can be computed recursively using the gradient flow equation: \[\label{eq:loss95derivatives} \frac{d^s L}{dt^s} = -\eta \Big\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \frac{d^{s-1} L}{dt^{s-1}}, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \widehat L \Big\rangle, \quad \frac{d^s \widehat L}{dt^s} = -\eta \Big\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \frac{d^{s-1} \widehat L}{dt^{s-1}}, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \widehat L \Big\rangle.\tag{5}\] At \(s=0\), \(L\) and \(\widehat L\) can be written as linear combinations of traces of various products of matrices \({\boldsymbol{U}}, {\boldsymbol{U}}^\top, {\boldsymbol{X}},{\boldsymbol{X}}^\top\): \[\label{eq:loss95as95base95diagrams} L = \frac{\frac{1}{2} D - R + \frac{p}{2}}{p},\quad \widehat L = \frac{\frac{1}{2} \widehat D - \widehat R + \frac{1}{2} \widehat F}{p m},\tag{6}\] where \[\begin{align} D ={}& \mathop{\mathrm{Tr}}[({\boldsymbol{U}}^\top {\boldsymbol{U}})^2], \; R = \mathop{\mathrm{Tr}}[{\boldsymbol{U}}^\top {\boldsymbol{U}}], \tag{7}\\ \widehat D ={}& \mathop{\mathrm{Tr}}[({\boldsymbol{U}}^\top {\boldsymbol{U}})^2 {\boldsymbol{X}}{\boldsymbol{X}}^\top], \; \widehat R = \mathop{\mathrm{Tr}}[{\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{X}}{\boldsymbol{X}}^\top], \; \widehat F = \mathop{\mathrm{Tr}}[{\boldsymbol{X}}{\boldsymbol{X}}^\top].\tag{8} \end{align}\] We describe each of these traces by a ring diagram, see Fig. 1, top left. Computing the scalar products of gradients in (5 ) then amounts to merging the diagrams corresponding to \(\widehat L\) and \(\tfrac{d^{s-1} L}{dt^{s-1}}\) or \(\tfrac{d^{s-1} \widehat L}{dt^{s-1}}\) (Fig. 1, top right). The merged diagrams are larger rings, also representing traces of products.

To finally get the polynomials \(Y_s,\widehat Y_s\), it remains to find the expectations of \(\tfrac{d^{s} L}{dt^{s}}\) and \(\tfrac{d^{s} \widehat L}{dt^{s}}\) at \(t=0\). Since \({\boldsymbol{U}}(0)\) and \({\boldsymbol{X}}\) have independent Gaussian entries, this can be done using Wick’s theorem. To find the expectation \(\mathbb{E}[G]\) of a diagram \(G\), we consider all possible pairings between edges of matching types (\({\boldsymbol{U}}\) or \({\boldsymbol{X}}\)) and contract the respective nodes (Fig. 1, bottom). Each pairing then contributes to \(\mathbb{E}[G]\) a monomial \(p^{q_p} n^{q_n} m^{q_m} \sigma^{q_\sigma}\), where \(q_\sigma\) is the number of \({\boldsymbol{U}}\)-edges and \(q_p, q_n, q_m\) are the numbers of respective contracted nodes. The initial term \(\tfrac{1}{2}=\frac{1}{2p}\mathop{\mathrm{Tr}}[{\boldsymbol{I}}]=\frac{1}{2pm}\mathop{\mathrm{Tr}}[{\boldsymbol{X}}{\boldsymbol{X}}^\top]\) is the loss of target identity matrix \({\boldsymbol{I}}\); it is convenient to separate it from the model-dependent terms.

The above diagrammatic picture was developed in [17] only for population-wide learning. By introducing diagrams with \({\boldsymbol{X}}\)-edges we extend it to finite training sets, allowing to separately analyze train and population losses and study model generalization.

3.0.0.2 Classification of learning regimes.

To obtain a general classification of learning regimes, we assume that the hyperparameters \(p,n,m\) and \(\sigma^2\) have a power-law scaling: \[\label{eq:alphascaling} p\asymp a^{\alpha_p}, n\asymp a^{\alpha_n}, m\asymp a^{\alpha_m},\sigma\asymp a^{\alpha_{\sigma}},\quad a\to+\infty,\tag{9}\] with some exponents \(\boldsymbol{\alpha}=(\alpha_p, \alpha_n, \alpha_m, \alpha_{\sigma})\). The condition that \(p,n,\) and \(m\) grow implies that the exponents \(\alpha_p, \alpha_n, \alpha_m\) (but not necessarily \(\alpha_\sigma\)) are positive. As \(a\to\infty\), the scaling of the polynomial \(Y_s\) is determined by its leading monomials maximizing the scalar product \(\mathbf{q}^\top\boldsymbol{\alpha}\): \[\label{eq:yscaling} Y_s\asymp a^{\alpha_{Y_s}},\quad \alpha_{Y_s}=\max_{\mathbf{q}\in Q_s} (\mathbf{q}^\top \boldsymbol{\alpha}),\quad Q_s=\{\mathbf{q}\in \mathbb{N}_0^4: c_{\mathbf{q}; s}\ne 0\}.\tag{10}\] The monomials on which the maximum is attained depend on \(\boldsymbol{\alpha}\). As a general principle, we associate the same learning regime to those scaling vectors \(\boldsymbol{\alpha}\) that have the same sets of leading monomials \[S_{\boldsymbol{\alpha};s}=\operatorname{Argmax}_{\mathbf{q}\in Q_s}(\mathbf{q}^\top \boldsymbol{\alpha})\subset Q_s.\] In general these sets depend on \(s\), but we will see that in our autoencoder model the phase picture will be essentially the same for all \(s\). Since \(Y_s\) are defined through the coefficients in the asymptotic power series (?? ), our learning phases are effectively defined on the level of asymptotic loss expansions.

To simplify description of the sets \(S_{\boldsymbol{\alpha};s}\), it is convenient to introduce the subsets \(P_s\subset Q_s\) of Pareto-optimal monomials. Let us introduce the partial ordering \(\mathbf{q}\preccurlyeq \mathbf{q}'\) by conditions \(q_p\le q'_p, q_n\le q'_n, q_m\le q'_m, q_\sigma=q'_\sigma\) (note the equality rather than inequality for \(q_\sigma\)). Then we define the Pareto set \(P_s\) as the points of \(Q_s\) non-dominated w.r.t. \(\preccurlyeq\): \[P_s=\{\mathbf{q}\in Q_s: \nexists\; \mathbf{q}'\in Q_s \text{ s.t. } \mathbf{q}'\ne \mathbf{q}\text{ and }\mathbf{q}\preccurlyeq \mathbf{q}'\}.\] Clearly, since in our scaling vectors \(\boldsymbol{\alpha}\) the components \(\alpha_p, \alpha_n\) and \(\alpha_m\) are positive, any leading set \(S_{\boldsymbol{\alpha};s}\) is a subset of the respective Pareto set \(P_s\). We define the Pareto set \(\widehat P_s\) for train loss \(\widehat L\) similarly to how \(P_s\) is defined for population loss \(L\). Below we denote by \(\overline{a,b}\) the integers in the interval \([a,b]\).

This theorem shows that the Pareto sets \(P_s, \widehat P_s\) are discretized triangular prisms (see Fig. 2). The triangular bases of these prisms are parameterized by two parameters \(s_D\in \overline{0, s+1}\), \(q_n\in \overline{1,s_D+ 1}\), while the heights are parameterized by the parameter \(q_m\). In the diagram-based analysis of the polynomials \(Y_s\), the parameter \(s_D\) denotes the number of “free” (i.e., “model-self-interaction”) diagrams appearing in the expansion.

The dependence of \(\mathbf{q}\) on the parameters \(s_D, q_n, q_m\) is linear and independent of \(s\) up to \(s\)-dependent shifts and changing the ranges of \(s_D, q_n, q_m\). As a result, all the prisms \(P_s,\widehat P_s\) have the same orientations of their faces in the \(\mathbf{q}\)-space. It follows that the learning phases – i.e., the sets of scaling vectors \(\boldsymbol{\alpha}\) corresponding to the same subsets \(S_{\boldsymbol{\alpha};s}, \widehat S_{\boldsymbol{\alpha};s}\) of leading terms – will be common for all \(s\).

In general, the sets \(S_{\boldsymbol{\alpha};s}, \widehat S_{\boldsymbol{\alpha};s}\) are various \(d\)-faces of the prisms \(P_s,\widehat P_s\). These faces can have dimensions \(d\) from 0 to 3. Dimension 3 corresponds to the fully balanced regime in which \(S_{\boldsymbol{\alpha};s}=P_s\) and \(\widehat S_{\boldsymbol{\alpha};s}=\widehat P_s\) so that all the Pareto terms contribute to the loss. This regime corresponds to the unique (up to normalization) power law scaling \((\alpha_p, \alpha_n,\alpha_m,\alpha_\sigma)=(1,1,1,-\tfrac{1}{2})\) (or, equivalently, \(p\asymp n\asymp m\asymp \sigma^{-2}\)). Lower-dimensional faces (including vertices and edges) correspond to various extremal regimes. A \(d\)-face corresponds to a dual (orthogonal) \((3-d)\)-dimensional set of normalized scaling vectors \(\boldsymbol{\alpha}\).

In particular, the prisms have five 2-faces corresponding to naturally interpretable specific extremes (see also a summary in Table 2 in the appendix):

1. Large-data regime (\(m\gg p\asymp n\asymp \sigma^{-2}\)) corresponds to the “top” of the prisms, i.e. \(q_m=s\). In particular, this regime describes the population-wide training. One can check that the population and train polynomials \(Y_s\) and \(\widehat Y_s\) coincide if restricted to these monomials, thus implying, as expected, the equality of the losses \(L(t)=\widehat L(t)\) in this regime.

2. Small-data regime (\(m\ll p\asymp n\asymp \sigma^{-2}\)) corresponds to the “bottom” of the prisms, i.e. \(q_m=0\) for the train prisms \(\widehat P_s\) and \(q_m=1\) for the population prisms \(P_s\). This implies, in particular, that in this regime the train loss \(\widehat L(t)\) is independent of data size \(m\), while the decrement \(L(0)-L(t)\) of the population loss scales linearly with \(m\); this is confirmed by experiment (see Figure 3).

3. Mean-field regime (\(n\asymp \sigma^{-2}\gg p\asymp m\)) corresponds to the lateral face \(q_n=s_D+1\). The name reflects the standard mean-field scaling \(n\asymp \sigma^{-2}\) and large latent dimension \(n\).

4. Narrow-latent regime (\(n\ll \sigma^{-2}\asymp p\asymp m\)) corresponds to the lateral face \(q_n=1\). The name refers to small latent dimension \(n\).

5. Free regime (\(\sigma^{-2}\ll n\asymp p\asymp m\)) corresponds to the lateral face \(s_D=s+1\). In this setting, thanks to large \(\sigma^2\), the initialized random model is much larger than the target identity, so gradient flow approximately acts as if no target is present, i.e., the model just “deflates to 0”.

Note that, thanks to cartesian factorization of the prisms, the data-related aspect of learning and the associated parameter \(m\) are effectively decoupled from the layer-size and noise magnitude aspects and the associated parameters \(p,n,\sigma^2\). In particular, the large-data and small-data regimes are essentially the opposite of one another. In contrast, the relations among the parameters \(p,n,\sigma^2\) and among the mean-field, narrow-latent, and free regimes is more subtle because there is no further factorization and these regimes correspond to the three sides of the base triangle. In particular, no two of these three regimes are opposites of one another.

We interpret the above list of five regimes as the exhaustive list of basic theoretical extremes. All other extremes are more degenerate, are associated with lower-dimensional faces, and can be obtained by combining the five basic extremes. For example, a regime combining the mean-field and free regimes can be defined by conditions \(\sigma^{-2}\ll n\gg p\asymp m\) and corresponds to the lateral edge \((q_n, s_D)=(s+2, s+1)\) of the prisms.

Each of the five basic extreme regimes admits a two-parameter refinement by considering three parameters: input-to-sample ratio \(\phi=p/m\), relative initialization magnitude \(\rho=n\sigma^2\), and input-to-hidden ratio \(\psi=p/n\). In each case, two of these parameters (or a combination thereof) can be used to define a specific proportional large-size limit.

4 Limiting solutions in extreme regimes (9)↩︎

We present explicit solutions for both train and population losses in regimes corresponding to four out of five two-dimensional faces of the prism. These regimes are naturally arranged by the complexity of the final solution. For the large-data regime, the solution is given by a closed-form formula. For the mean-field one, the solution is given by an integral. For the narrow-latent regime, the solution follows from a system of two scalar ODEs. Finally, for the small-data one, the evolution could be reformulated to a hierarchical system of scalar ODEs; restricting this system to the first \(K\) ODEs already yields a good match with experiment. We have not found a suitable reformulation of the free regime in terms of finite objects suitable for numerical validation. We leave this case for future work.

The diagrammatic formalism used to derive the regime classification could also be used to obtain explicit solutions in the extreme regimes by summing the respective leading monomials (see [17]). Their method is applicable to tensors of arbitrary order, but is relatively cumbersome and seems less efficient when applied to matrix problems as in our case. We therefore take a more direct approach based on random matrix theory. We believe that the same results could also be reproduced by the diagrammatic method.

We generally choose the learning rate \(\eta\) so that the characteristic time of the loss change is \(\asymp 1\). The choice depends on the regime and can be performed, e.g., by balancing the coefficients \((\frac{-\eta}{pm})^s Y_s\) in eq. (?? ) given scalings (9 ), (10 ). In our basic extreme regimes \(\eta=p, m\), or \(\sigma^{-2}\) (see 2).

The formulas in the large-data and mean-field regimes follow from spectral convergence of Wishart matrices after the corresponding limiting reductions. The narrow-latent and small-data formulas additionally rely on asymptotic row-decoupling and isotropic moment-closure assumptions, respectively. Although our derivations are partly heuristic, they agree very well with direct numerical integration of the gradient flow 2 : see 3 4.

4.1 Mean-field (9.2)↩︎

This regime corresponds to vanishing \(\psi\), with \(\phi\) and \(\rho\) staying finite; the natural learning rate is \(\eta = p\). In this case, the whole dynamics decouple into independent ones for each eigencomponent of the data covariance matrix \({\boldsymbol{X}}{\boldsymbol{X}}^\top\). To be specific, let \(\ell(t)\) be the training loss evolution on a single sample \(x = 1\) of dimensionality one with weights initialized from a sphere of radius \(\sqrt\rho\). Then the loss dynamics on the whole dataset is given as its average over the Marchenko-Pastur distribution: \[{\mathcal{L}}(t) = \int \ell(\lambda t)\,d\mu_\phi(\lambda), \qquad \widehat{\mathcal{L}}(t) = \int \lambda \ell(\lambda t)\,d\mu_\phi(\lambda), \qquad \ell(t) = \frac{(1-\rho)^2 e^{-8 t}}{2 \left(\rho + (1-\rho) e^{-4 t}\right)^2},\] where the Marchenko–Pastur law is \[d\mu_\phi(\lambda) = \left(1-\frac{1}{\phi}\right)_+\delta_0(d\lambda) + \frac{\sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)}}{2\pi\phi\lambda} \boldsymbol{1}_{\lambda\in[\lambda_-,\lambda_+]}\,d\lambda, \qquad \lambda_\pm=(1\pm \sqrt{\phi})^2.\]

4.1.0.1 Large time behavior.

Since \(\ell(t)\) decays exponentially, and \(\mu_\phi\) features a spectral gap for any \(\phi \neq 1\) (i.e. there is a gap between zero and the bulk of the spectrum), \(\widehat{\mathcal{L}}(t)\) also decays exponentially when \(\phi \neq 1\). When \(\phi < 1\), this also applies to \({\mathcal{L}}(t)\). For \(\phi > 1\), there is a point mass at zero, which does not affect \(\widehat{\mathcal{L}}(t)\) due to the factor \(\lambda\), while it results in a remainder term for \({\mathcal{L}}(t)\): \[\label{eq:test-loss-final} {\mathcal{L}}(\infty) = \frac{(1-\rho)^2}{2} \left(1 - \frac{1}{\phi}\right)_+.\tag{11}\] The final population loss is zero when \(\rho = 1\). This is due to the fact that \({\boldsymbol{U}}^\top {\boldsymbol{U}}\to {\boldsymbol{I}}_p\) as \(n \to \infty\) for fixed \(p\) due to LLN, whenever \(\sigma^2 = 1/n\).

When \(\phi = 1\), the spectral gap disappears; because of this both \(\widehat{\mathcal{L}}(t)\) and \({\mathcal{L}}(t)\) decay much slower: \[{\mathcal{L}}(t) - {\mathcal{L}}(\infty) \asymp \begin{cases} \frac{1}{t^{1/2}}, & \phi = 1, \\ \frac{e^{-8 (1 - \sqrt\phi)^2 t}}{t^{3/2}}, & \phi \neq 1, \end{cases} \qquad\qquad \widehat{\mathcal{L}}(t) \asymp \begin{cases} \frac{1}{t^{3/2}}, & \phi = 1, \\ \frac{e^{-8 (1 - \sqrt\phi)^2 t}}{t^{3/2}}, & \phi \neq 1 \end{cases}\] as \(t \to \infty\), where the exact expressions are derived in 9.2.1. As a result, if we fix large enough \(t\), the training loss as function of \(\phi\) features a maximum near \(\phi = 1\), i.e. when the problem becomes critically determined. This is in parallel with double descent phenomena: the population loss peaks at \(m = p\) for linear regression trained on a noisy linear teacher. The nature of the peak is different, however: for a linear regression, it is caused by the fact that the model has to spend all its capacity to fit the teacher, thus becoming very sensitive to noise, while in our case, there is no noise in the target; the peak is caused by vanishing non-zero eigenvalues of the data covariance matrix \({\boldsymbol{X}}{\boldsymbol{X}}^\top\) which appear when \(m = p\) and dramatically slow down the gradient flow.

4.2 Narrow-latent (9.3)↩︎

This regime corresponds to diverging \(\psi\), with \(\phi\) and \(\psi \rho\) staying finite; the natural learning rate is again \(\eta = p\). Since a vanishingly narrow-latent model (\(\psi \to \infty\)) is not able to fit the data,\[{\mathcal{L}}(t) = \widehat{\mathcal{L}}(t) = 1/2 \qquad \forall t \geq 0,\] where \(1/2\) is the loss of the zero model. The learned parts are given by a \(O(\psi^{-1})\) correction term: \[\label{eq:test95loss95narrow} \lim_{p,n,m\to\infty} \left[\left(\mathbb{E}[L(t)] - {\mathcal{L}}(t)\right) \psi\right] = a(t) \left(\frac{a(t)}{2} - 1\right),\tag{12}\] \[\label{eq:train95loss95narrow} \lim_{p,n,m\to\infty} \left[\left(\mathbb{E}\left[\widehat L(t)\right] - \widehat{\mathcal{L}}(t)\right) \psi\right] = r(t) \left(\frac{a(t)}{2} - 1\right).\tag{13}\] The quantities \(a\) and \(r\) are given in terms of Marchenko-Pastur averages \(M_\phi(q)=\int e^{q\lambda}\, d\mu_\phi(\lambda)\): \[\label{eq:a95r95narrow} a(t) = \psi \rho e^{-2B(t)}M_\phi(2R(t)), \qquad r(t) = \psi \rho e^{-2B(t)}M_\phi'(2R(t)),\tag{14}\] where \(R(t)\) and \(B(t)\) are a solution of the following ODEs. \[\dot{R}(t)=2-a(t), \qquad \dot{B}(t)=r(t), \qquad R(0)=B(0)=0.\]

Comparing to the mean-field regime, the dynamics again decouple, but now wrt hidden neurons rather than data eigencomponents. That is, weights associated with different neurons become orthogonal as their dimension \(p\) grows, and stay so during the training process. For this reason, the above evolution could be derived from that with a single hidden neuron. Suppose its weight vector is \({\boldsymbol{u}}\). The above variables \(a\) and \(r\) are the row norm and the alignment of \({\boldsymbol{u}}\) to the data covariance matrix, respectively: \[a = \lim_{p\to\infty} \|{\boldsymbol{u}}\|^2, \qquad r = \lim_{p,m\to\infty} {\boldsymbol{u}}^\top \frac{{\boldsymbol{X}}{\boldsymbol{X}}^\top}{m} {\boldsymbol{u}}.\]

4.2.0.1 Large time behavior.

The gradient flow evolution yields \(\dot{a} = 4 r (1-a)\). As we shall see shortly, \(r(\infty) \neq 0\), hence \(a(\infty) = 1\). Because of this, \(R(t) \to \infty\) as \(t \to \infty\). 14 then implies \[\frac{r(t)}{a(t)} = \frac{M_\phi'(2R(t))}{M_\phi(2R(t))} \to \lambda_+ = \left(1 + \sqrt\phi\right)^2,\] since for large \(q\), the integrals of \(M_\phi(q)\) and \(M'_\phi(q)\) are dominated by the top of the spectrum. Therefore \(r(\infty) = \lambda_+ \neq 0\). This means that training not only normalizes the row (\(a(t) \to 1\)), but also aligns it with the strongest component of the data covariance matrix. Substituting this to 12 13 for \(t = \infty\) yields \[{\mathcal{L}}(\infty) - \widehat{\mathcal{L}}(\infty) = 0, \qquad \lim_{p,n,m\to\infty} \left[\mathbb{E}\left[L(\infty) - \widehat L(\infty)\right] \psi\right] = \frac{\lambda_+ - 1}{2} = \sqrt\phi + \frac{\phi}{2}.\] Therefore, since \(\phi = p / m\), the more data we introduce, the lower the generalization gap.

As we demonstrate in 9.3.1, \(a(t)\) converges exponentially, while \(r(t)\) does so only algebraically. Hence, training occurs on two timescales: by the time \(\|{\boldsymbol{u}}\|\) almost reaches \(1\), \({\boldsymbol{u}}\) is still aligning itself to the strongest component of the data covariance matrix. The population loss is minimal for any \({\boldsymbol{u}}\) of norm one, while the train one is minimal when \({\boldsymbol{u}}\) is moreover aligned as above. Therefore, the population loss converges much faster than the train one (see 9.3.1 for derivations): \[\lim_{p,n,m\to\infty} \left[\mathbb{E}\left[L(t) - L(\infty)\right] \psi\right] \asymp t^6 e^{-8 \lambda_+ t}, \qquad \lim_{p,n,m\to\infty} \left[\mathbb{E}\left[\widehat L(t) - \widehat L(\infty)\right] \psi\right] \asymp t^{-1}.\]

4.3 Large-data (9.4)↩︎

This regime corresponds to vanishing \(\phi\), with \(\psi\) and \(\rho\) kept finite; the natural learning rate is once again \(\eta = p\). When the data is infinite, the train and the population losses coincide. \[{\mathcal{L}}(t) = \widehat{\mathcal{L}}(t) = \frac{1}{2} - \rho e^{4t} H_\psi\left(\rho (1-e^{4t})\right) + \frac{\rho^2 e^{8t}}{2} H_\psi'\left(\rho (1-e^{4t})\right),\] where \[H_\psi(q) = \frac{1-q\left(1+\psi\right) - \sqrt{1-2q\left(1+\psi\right) + q^2\left(1-\psi\right)^2}}{2 q^2 \psi}.\]

4.3.0.1 Large time behavior.

The model implements a linear map \({\boldsymbol{x}}\to {\boldsymbol{A}}{\boldsymbol{x}}\) for \({\boldsymbol{A}}= {\boldsymbol{U}}^\top {\boldsymbol{U}}\in \mathbb{R}^{p \times p}\). Since the data is infinite, it satisfies \(\dot{\boldsymbol{A}}= 4 {\boldsymbol{A}}({\boldsymbol{I}}- {\boldsymbol{A}})\). Therefore, the eigenvectors of \({\boldsymbol{A}}\) stay unchanged. Every nonzero eigenvalue converges to \(1\), while zero ones stay zero. Since \(\mathop{\mathrm{rank}}{\boldsymbol{A}}(0) = \min(n,p)\), the converged loss becomes \[{\mathcal{L}}(\infty) = \frac{p-\min(n,p)}{2p} = \frac{1}{2} \left(1 - \frac{1}{\psi}\right)_+,\] i.e. the perfect reconstruction happens iff the hidden layer is not a bottleneck.

The spectrum of \({\boldsymbol{A}}(0) / \rho\) follows \(\mu_\psi\). This causes a qualitative difference in loss behavior for large \(t\) between \(\psi = 1\) and \(\psi \neq 1\), similar to that in the mean-field regime. The difference is again caused by the spectral gap present when \(\psi \neq 1\), but in contrast to the mean-field regime, the loss decays exponentially in both scenarios, albeit with different exponents (see 9.4.1 for derivations): \[{\mathcal{L}}(t) - {\mathcal{L}}(\infty) \asymp \begin{cases} e^{-2t}, & \psi = 1,\\ e^{-8t}, & \psi \neq 1, \end{cases} \qquad t \to \infty.\]

4.4 Small-data (9.5)↩︎

This regime corresponds to diverging \(\phi\), with \(\psi\) and \(\rho\) staying finite; the natural learning rate is now \(\eta = m\). The train loss is given as follows. \[\widehat{\mathcal{L}}(t) = \frac{1}{2} \left(1-a(t)\right)^2 + \frac{1}{2} a(t) \mu_1(t),\] where \(a\) and \(\mu_1\) are given as the following hierarchical system of ODEs. \[\label{eq:small95data95a} \dot{a} = 4 a (1-a) - 2a \mu_1, \qquad a(0) = \rho;\tag{15}\] \[\dot{\mu}_k = -2\mu_{k+1} - 2a \sum_{j=0}^{k-1} \mu_j\mu_{k-j} + 2\mu_1\mu_k, \qquad \mu_k(0) = \rho^k \sum_{j=1}^k N_{k,j} \psi^j \left(1 - \frac{1}{\phi}\right)^j \qquad k \geq 1,\] where \(\mu_0 \equiv 1\), and \(N_{k,j}\) are Narayana numbers: \(N_{k,j} = \tfrac{1}{k} \tbinom{k}{j} \tbinom{k}{j-1}\).

As for the population loss, an extremely small data model (\(\phi \to \infty\)) seems to have no chance to generalize. Also because of this, the infinite-dimensional loss \({\mathcal{L}}(t)\) should not depend on time \(t\). However, a weight-tied autoencoder is implicitly biased towards the identity map: if \(p\) was fixed and \(\rho = 1\) then \({\boldsymbol{U}}^\top(0) {\boldsymbol{U}}(0)\) would converge to \({\boldsymbol{I}}\) by LLN. For general \(\rho\) and \(\psi\), \[{\mathcal{L}}(t) = {\mathcal{L}}(0) = \frac{(1-\rho)^2 + \psi \rho^2}{2} \qquad \forall t \geq 0.\] The sub-leading term is \(O(\phi^{-1})\), and is given by \[\begin{align} \lim_{p,n,m\to\infty} \left[\left(\mathbb{E}[L(t)] - {\mathcal{L}}(t)\right) \phi\right] &= \frac{(1-a)^2}{2} + a \mu_1 - \frac{(1-\rho)^2}{2} - \psi \rho^2 + 2 \int_0^t a (\mu_1 - \mu_2) \, ds. \end{align}\]

4.4.0.1 Large time behavior.

Since the dataset is small, the weight matrix naturally divides into two blocks, aligned with the data (active) and orthogonal to it (inactive): \({\boldsymbol{U}}= \left[{\boldsymbol{U}}_\parallel \quad {\boldsymbol{U}}_\perp\right]\). They do not evolve independently. The quantities \(a\) and \(\mu_k\) are the norm of the active block, and its order \(k\) "leakage" to the inactive one, respectively: \[a = \frac{1}{m} \mathop{\mathrm{Tr}}\left[{\boldsymbol{U}}_\parallel^\top {\boldsymbol{U}}_\parallel\right], \qquad \mu_k = \frac{1}{m a} \mathop{\mathrm{Tr}}\left[{\boldsymbol{U}}_\parallel^\top \left({\boldsymbol{U}}_\perp {\boldsymbol{U}}_\perp^\top\right)^k {\boldsymbol{U}}_\parallel\right] \quad \forall k \in {\mathbb{N}}.\] As \(t \to \infty\), \(a(t) \to 1\), while \(\mu_k(t) \to 0\) \(\forall k \in {\mathbb{N}}\). That is, the model fits the training data in the subspace it sees, while the active block gets orthogonal to the inactive one. Therefore \(\widehat{\mathcal{L}}(\infty) = 0\), while \[\lim_{p,n,m\to\infty} \left[\left(\mathbb{E}[L(\infty)] - {\mathcal{L}}(0)\right) \phi\right] = -\frac{(1-\rho)^2}{2} - \psi \rho^2 + 2 \int_0^\infty a(s) (\mu_1(s) - \mu_2(s)) \, ds,\] which depends on the trajectory of the first and second leakage moments.

5 Limitations and future work↩︎

We expect that under a non-isotropic Gaussian dataset (a power-law spectrum is more natural and a more common assumption in the literature), the regime classification does not change, and most of the regimes we resolved could still be solved, but might yield more complicated results. A more challenging generalization would be to introduce an activation function. It is not clear whether the same regime classification holds, or it changes: the diagrammatic technique we used for classification does not directly apply to nonlinear activations. Some prior works analyzed and solved (although in the online SGD setting, not for a gradient flow on a finite dataset) ReLU autoencoders in the mean-field [16] and narrow-latent [15] regimes, therefore they should stay. We leave deriving the solutions for them for future work.

6 Summary of notation and extreme regimes↩︎

See summary of notation in Table 1 and summary of extreme regimes in Table 2

3.5pt

7 Prior works↩︎

7.0.0.1 Training and generalization dynamics of autoencoders.

Following a seminal work of [2] on training dynamics of linear nets, [18] integrated gradient flow for a linear autoencoder trained on a finite dataset. Same as [2], they had to assume data-aligned weight initialization, and do not study generalization. [15], [16] considered a shallow ReLU autoencoder trained with online SGD. [16] obtained exact population loss evolution for a wide hidden layer in the mean-field limit, while [15] obtained a system of equations governing the system evolution for the opposite case of a narrow hidden layer. We emphasize that the setting of online SGD the above studies work with assumes the train dataset to be essentially infinite, in contrast to our case. On the other hand, [19] considered a finite train dataset, same as we do, but worked directly with a regularized empirical loss minimizer, thus completely omitting the training dynamics that we study in the present work.

7.0.0.2 Large-size regime classification.

In the present work, we derive a complete classification of learning regimes when latent and input dimensions, as well as the dataset size, are large. This classification is similar in spirit to that of [17], who considered the gradient flow dynamics on the problem of canonical-polyadic decomposition of the identity tensor. Prior to it, a number of works [20]–[22] proposed different classifications of neural network’s learning regimes in the limit when only hidden dimensions go to infinity, while all others are kept fixed. These regimes include mean-field, NTK, as well as "intermediate" ones. On the one hand, our classification is more general as not only the hidden dimensions may diverge. On the other hand, our classification is specific to shallow weight-tied autoencoders.

8 Diagrammatic analysis of the loss evolution↩︎

8.1 Diagram expansions and proof of Proposition 1↩︎

In this section we provide more details on the diagrammatic loss analysis sketched in Section 3, in particular establishing Proposition 1.

8.1.0.1 Loss expansion.

We start with the Taylor expansions of the population and train losses at \(t=0\): \[\label{eq:elly1} L(t) \sim \sum_{s=0}^\infty \frac{d^s L}{dt^s}(0) \frac{t^s}{s!},\quad \widehat L(t) \sim \sum_{s=0}^\infty \frac{d^s \widehat L}{dt^s}(0) \frac{t^s}{s!}.\tag{16}\] For a finite model with a particular initialization these series converge in a usual sense for sufficiently small \(|t|\), but in general, having in mind the large-system limit, it will be more convenient to view these expansions as asymptotic, well-defined on the level of individual power terms.

The coefficients \(\tfrac{d^s L}{dt^s}(0), \tfrac{d^s \widehat L}{dt^s}(0)\) can be computed using Eqs. (5 ): \[\label{eq:loss95derivatives1} \frac{d^s L}{dt^s} = -\eta \Big\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \frac{d^{s-1} L}{dt^{s-1}}, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \widehat L \Big\rangle, \quad \frac{d^s \widehat L}{dt^s} = -\eta \Big\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \frac{d^{s-1} \widehat L}{dt^{s-1}}, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \widehat L \Big\rangle.\tag{17}\] Note that the loss in the right factor of the scalar products is always the train loss \(\widehat L\), since it determines the evolution of the weights, while the left factor contains the population or train loss, depending on which one we observe.

8.1.0.2 Diagrams.

The recursive computation of \(\tfrac{d^s L}{dt^s}, \tfrac{d^s \widehat L}{dt^s}\) can be conveniently represented by diagrams. First expand the losses \(L, \widehat L\), as defined by (1 ), in linear combination of traces: \[\label{eq:loss95as95base95diagrams1} L = \frac{\frac{1}{2} D - R + \frac{p}{2}}{p},\quad \widehat L = \frac{\frac{1}{2} \widehat D - \widehat R + \frac{1}{2} \widehat F}{p m},\tag{18}\] where \[\begin{align} D ={}& \mathop{\mathrm{Tr}}[({\boldsymbol{U}}^\top {\boldsymbol{U}})^2], \; R = \mathop{\mathrm{Tr}}[{\boldsymbol{U}}^\top {\boldsymbol{U}}], \tag{19}\\ \widehat D ={}& \mathop{\mathrm{Tr}}[({\boldsymbol{U}}^\top {\boldsymbol{U}})^2 {\boldsymbol{X}}{\boldsymbol{X}}^\top], \; \widehat R = \mathop{\mathrm{Tr}}[{\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{X}}{\boldsymbol{X}}^\top], \; \widehat F = \mathop{\mathrm{Tr}}[{\boldsymbol{X}}{\boldsymbol{X}}^\top].\tag{20} \end{align}\] Each of these traces can be written as a sum of products of entries of the matrices \({\boldsymbol{U}},{\boldsymbol{X}}\). For example, \[\widehat D = \sum_{k=1}^n\sum_{k'=1}^n \sum_{l=1}^m \sum_{i=1}^p \sum_{i'=1}^p \sum_{i''=1}^p U_{ki}U_{ki'}U_{k'i'}U_{k'i''}X_{i'l}X_{i''l}.\]

Any expression of this kind can be described by a graph (“diagram”) in which:

1. The vertices correspond to the summation indices and can be one of the three types: \(p\) (input dimension), \(n\) (latent dimension) or \(m\) (dataset size).

2. The edges connecting two vertices correspond to respective entries of the matrices \({\boldsymbol{U}}, {\boldsymbol{X}}\) and accordingly can be of two types (\({\boldsymbol{U}}\) or \({\boldsymbol{X}}\)). \({\boldsymbol{U}}\)-edges connect \(n\)-nodes with \(p\)-nodes, while \({\boldsymbol{X}}\)-edges connect \(p\)-nodes with \(m\)-nodes.

3. The value associated with the diagram is obtained by multiplying the entries of \({\boldsymbol{U}}, {\boldsymbol{X}}\) over the edges and summing resulting products over all configurations of node indices.

Note that:

1. Thanks to the trace-product structure of \(D,R,\widehat D,\widehat R,\widehat F,\) the associated five diagrams are ring diagrams (see Fig. 1). However, in general diagrams defined by above rules may be more general, e.g. the diagrams obtained by contractions of the ring diagrams (see below) are not ring diagrams.

2. In the ring diagrams such as \(D,R,\widehat D,\widehat R,\widehat F,\) the \(n,m\)-nodes alternate with \(p\)-nodes. The number of \({\boldsymbol{X}}\)-edges is twice the number of \(m\)-nodes, while the number of \({\boldsymbol{U}}\)-edges is twice the number of \(n\)-nodes.

8.1.0.3 Diagram merging.

Computation of scalar products (17 ) can be described in terms of diagram merging. Let \(G_1, G_2\) be two functions of the weights represented by diagrams. Then the scalar product \(\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} G_1, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} G_2\rangle\) can be computed by the following rules:

1. Consider all pairs of a \({\boldsymbol{U}}\)-edge \(g_1\) in \(G_1\) and a \({\boldsymbol{U}}\)-edge \(g_2\) in \(G_2\).

2. For each such pair, merge the diagrams \(G_1\) and \(G_2\) by identifying the \(n\)-nodes of \(g_1, g_2\), identifying the \(p\)-nodes of \(g_1,g_2\), and removing the edges \(g_1,g_2\) (see Fig. 1).

3. Add the resulting diagrams.

Note that:

1. Merger of two diagrams produces a linear combination of diagrams (corresponding to different pairs of edges).

2. The diagrams are merged only over \({\boldsymbol{U}}\)-edges and not \({\boldsymbol{X}}\)-edges (since only the \({\boldsymbol{U}}\)-edges contain the trainable model weights).

3. Merger of two ring diagrams \(G_1, G_2\) produces again ring diagrams. Let \(q_p^{(r)}, q_n^{(r)}, q_m^{(r)}, q_\sigma^{(r)}\) denote, respectively, the numbers of \(p\)-, \(n\)-, \(m\)-nodes and edges in \(G_r, r=0,1\). Then in the merged diagrams \[\begin{align} q_p ={}& q_p^{(1)}+q_p^{(2)}-1,\\ q_n ={}& q_n^{(1)}+q_n^{(2)}-1,\\ q_m ={}& q_m^{(1)}+q_m^{(2)}, \\ q_\sigma ={}& q_\sigma^{(1)}+q_\sigma^{(2)}-2. \end{align}\]

We denote the merge operation by \(\star\) (i.e. the merger of \(G_1\) and \(G_2\) is \(G_1\star G_2\)) and naturally bilinearly extend it to linear combinations of diagrams. Note that the operation \(\star\) is commutative, but not generally associative.

By (18 ), for any function of the weights expressed by a diagram \(G\) we have \[\langle {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} G, {\boldsymbol{\nabla}}_{{\boldsymbol{U}}} \widehat L \rangle = \frac{1}{pm} G\star(\tfrac{1}{2}\widehat D - \widehat R + \tfrac{1}{2} \widehat F)=\frac{1}{pm} G\star(\tfrac{1}{2}\widehat D - \widehat R),\] where the \(\widehat F\)-term disappears because the diagram \(\widehat F\) does not include model weights. Applying this identity to the derivatives of the losses \(L\) and \(\widehat L\) and performing iterations (17 ), we obtain for \(s\ge 1\) \[\begin{align} \frac{d^s L}{dt^s}={}&\frac{(-\eta)^s}{p^{s+1}m^s} (\tfrac{1}{2} D- R)[\star(\tfrac{1}{2}\widehat D-\widehat R)]^s,\tag{21}\\ \frac{d^s \widehat L}{dt^s}={}&\frac{(-\eta)^s}{p^{s+1}m^{s+1}} (\tfrac{1}{2}\widehat D- \widehat R)[\star(\tfrac{1}{2}\widehat D-\widehat R)]^s=\frac{(-\eta)^s}{p^{s+1}m^{s+1}} (\tfrac{1}{2}\widehat D-\widehat R)^{\star(s+1)}.\tag{22} \end{align}\] At \(s=0\) the above formulas should additionally include the term corresponding to the target (see Eq. (18 )). This is why we separate the values \(\tfrac{1}{2}\) in Eq. (?? ).

8.1.0.4 Averaging, edge pairings, contractions.

We consider now averaging of the losses w.r.t. random initialization \({\boldsymbol{U}}(0)\) and training set \({\boldsymbol{X}}\). Using loss expansions (16 ), we write \[\label{eq:elly2} \mathbb{E}[L(t)] \sim \sum_{s=0}^\infty \mathbb{E}\Big[\frac{d^s L}{dt^s}(0)\Big] \frac{t^s}{s!},\quad \mathbb{E}[\widehat L(t)] \sim \sum_{s=0}^\infty \mathbb{E}\Big[\frac{d^s \widehat L}{dt^s}(0)\Big] \frac{t^s}{s!},\tag{23}\] reducing the question to computing the expectations \(\mathbb{E}\big[\frac{d^s L}{dt^s}(0)\big], \mathbb{E}\big[\frac{d^s \widehat L}{dt^s}(0)\big]\).

Formulas (21 ), (22 ) show that the derivatives \(\tfrac{d^s L}{dt^s}(0), \tfrac{d^s \widehat L}{dt^s}(0)\) can be expressed as linear combinations of diagrams, so the task can be reduced, in turn, to computing the expectation \(\mathbb{E}[G]\) for an arbitrary diagram \(G\).

Since at \(t=0\) the matrices \({\boldsymbol{U}}\) and \({\boldsymbol{X}}\) are independent and have independent gaussian entries, we can use Wick’s theorem to compute \(\mathbb{E}[G]\) at \(t=0\). Recall that this theorem states that the expectation of a product of an even number of jointly normal centered random variables equals the sum of the products of pair covariances taken over all partitions of these variables into pairs.

In the context of diagrams a pairing of the variables corresponds to an edge pairing. The independence conditions imply that the only nonzero covariances between the entries of \({\boldsymbol{U}},{\boldsymbol{X}}\) are those between the same entries of the same matrix (\({\boldsymbol{U}}\) or \({\boldsymbol{X}}\)). This imposes identity constraints on the indices of the paired edges. Next, in our setting the variance of the entries of \({\boldsymbol{U}}\) is \(\sigma^2\), while the variance of the entries of \({\boldsymbol{X}}\) is \(1\). This implies the following procedure for computing \(\mathbb{E}[G]\) (see Fig. 1):

1. Consider all pairings of the edges of \(G\) between matching edges (i.e., \({\boldsymbol{U}}\)-edges with \({\boldsymbol{U}}\)-edges and \({\boldsymbol{X}}\)-edges with \({\boldsymbol{X}}\)-edges).

2. For each pairing:

   1. For each pair of edges, contract (i,e., identify) their respective \(p\)-, \(n\)- and/or \(m\)-nodes. The resulting contracted nodes correspond to the degrees of freedom left after imposing all the identity constraints.

   2. The resulting contracted diagram contributes to \(\mathbb{E}[G]\) the term \(p^{q_p}n^{q_n}m^{q_m}\sigma^{q_\sigma},\) where \(q_p, q_n,q_m\) are the numbers of respective nodes after contraction, and \(q_\sigma\) is the number of \({\boldsymbol{U}}\)-edges.

8.1.0.5 Proof of Proposition 1.

The initial terms \(\tfrac{1}{2}\) in the expansions (?? ) of \(L(t), \widehat L(t)\) result from \(\tfrac{1}{2p}\mathop{\mathrm{Tr}}[{\boldsymbol{I}}]\) and \(\tfrac{1}{2pm}\mathbb{E}\mathop{\mathrm{Tr}}[{\boldsymbol{X}}{\boldsymbol{X}}^\top]\), respectively.

The values \(Y_s, \widehat Y_s\) can be obtained using expressions (21 ), (22 ) for the order-\(s\) derivatives of the losses: \[\begin{align} Y_s={}&\Big(\frac{(-\eta)^s}{p^{s}m^s}\Big)^{-1}\frac{d^s L}{dt^s}=\frac{1}{p}(\tfrac{1}{2} D- R)[\star(\tfrac{1}{2}\widehat D-\widehat R)]^s,\tag{24}\\ \widehat Y_s={}&\Big(\frac{(-\eta)^s}{p^{s}m^s}\Big)^{-1}\frac{d^s\widehat L}{dt^s}=\frac{1}{pm}(\tfrac{1}{2}\widehat D-\widehat R)^{\star(s+1)}.\tag{25} \end{align}\] As observed above, expectations of diagrams are polynomials in \(p,n,m,\sigma^2\) (\(\sigma\) occurs in even powers since the number of \({\boldsymbol{U}}\)-edges is always even).

Moreover, the diagrams appearing in \((\tfrac{1}{2} D- R)[\star(\tfrac{1}{2}\widehat D-\widehat R)]^s\) have at least one \(p\)-node after contraction, thus always providing a factor \(p\), so \(Y_s\) will still be a polynomial. Also, all the diagrams appearing in \((\tfrac{1}{2}\widehat D-\widehat R)^{\star(s+1)}\) contain \({\boldsymbol{X}}\)-edges and will have at least one \(p\)-node and at least one \(m\)-node after contraction, so \(\widehat Y_s\) will also be a polynomial. This completes the proof.

8.2 Proof of Theorem 2↩︎

8.2.0.1 Preliminaries.

We start from the expressions (24 ), (25 ) for \(Y_s, \widehat Y_s\). When binomially expanded, these expressions contain various sequential mergers of diagrams \(D,\widehat D\) and \(R,\widehat R\). Given such a sequence, denote by \(s_D\) the number of diagrams \(D,\widehat D\) and by \(s_R\) the number of diagrams \(R,\widehat R\), so that \[s_D+s_R=s+1.\] We will use \(s_D\) as one of the independent parameters in the parameterization of Pareto sets (with \(s_R\) found by \(s_R=s+1-s_D\) and the other independent parameters being the numbers \(q_n,q_m\) of \(n\)- and \(m\)- nodes after contraction).

8.2.0.2 Train Pareto sets \(\widehat P_s\).

Part 1: Realizability. First observe that any merger of \(s_D\) diagrams \(\widehat D\) and \(s_R=s+1-s_D\) diagrams \(\widehat R\) produces ring diagrams with \(s_D+1\) \(n\)-nodes, \(s+1\) \(m\)-nodes, \(s+s_D+2\) \(p\)-nodes, \(2(s+1)\) \({\boldsymbol{X}}\)-edges and \(2(s_D+1)\) \({\boldsymbol{U}}\)-edges. Since \(q_\sigma\) equals the number of \({\boldsymbol{U}}\)-edges, we get \(q_\sigma=2(s_D+1)\), as desired. Now, let us show that for any \(q_n\in \overline{1,s_D+ 1}\) and \(q_m\in\overline{0,s}\) there exist merged diagrams and edge pairing that induce contractions leaving \(q_n\) \(n\)-nodes, \(q_m+1\) \(m\)-nodes, and \(s+s_D-q_n-q_m+2\) \(p\)-nodes.

To this end, arrange mergers in such a way that the \({\boldsymbol{U}}\)-edges and \({\boldsymbol{X}}\)-edges form two contiguous arcs on the ring. Pair some \({\boldsymbol{X}}\)-edges connecting nearest \(m\)-nodes, forming \(s-q_m\) pairs — this is possible since \(q_m\in\overline{0,s}\) and there are \(s+1\) \(m\)-nodes. This leaves precisely \(q_m+1\) \(m\)-nodes after contraction induced by this pairing. Each edge pair forms a “petal” ending in \(p\)-node, with the total of \(s-q_m\) \(p\)-nodes.

Analogously, consider the \({\boldsymbol{U}}\)-arc of the ring. Form \(s_D+1-q_n\) pairs of some edges connecting \(n\)-nodes – this is possible since \(q_n\in \overline{1,s_D+ 1}\) and there are \(s_D+1\) \(n\)-nodes. After contraction, this leaves \(q_n\) \(n\)-nodes and \(s_D+1-q_n\) petals each ending in a \(p\)-node.

Now consider the remaining ring (ignoring the already formed petals) and contract all the remaining \(p\)-nodes into one. This contractions is naturally associated with a final pairing of all the remaining edges.

The total number of \(p\)-nodes after all contractions is \((s-q_m)+(s_D+1-q_n)+1=s+s_D-q_n-q_m+2,\) as desired.

Dividing by \(pm\), the contribution from this merger/pairing to \(\widehat Y_s\) is then \(p^{s+s_D-q_n-q_m+1}n^{q_n}m^{q_m}\sigma^{2(s_D+1)}\) with some coefficient, as desired.

Note also that contributions to one monomial in \(\widehat Y_s\) from different mergers/pairings have coefficients of the same sign, so that no cancelling can occur, and presenting one example of a merger/pairing is sufficient to ensure the presence of the monomial. Indeed, the sign of the contribution of a particular merger to \(\widehat Y_s\) is determined by the parity of \(s_R\) and, for given \(s\), \(s_R\) uniquely corresponds to \(q_\sigma\).

We conclude that all monomials listed in the statement of the theorem indeed occur in \(\widehat Y_s\).

Part 2: Optimality. We show now that no monomials Pareto-dominating the listed monomials can be present in \(\widehat Y_s\).

First recall that domination w.r.t. \(\preccurlyeq\) requires the monomials to have the same \(q_\sigma\), which, as already noted, uniquely corresponds to \(s_D\) for given \(s\). Thus, we need to argue that for given \(s,s_D\) we cannot achieve, by a suitable merger/pairing, larger counts of some of the \(p\)-, \(n\)-, \(m\)-nodes after contraction, without decreasing other counts.

To this end, ignore the difference between edge types and observe that any pairing of \(2(s+s_D+2)\) edges in a connected graph induces a node contraction leaving not more than \(s+s_D+3\) nodes. Indeed, the contracted graph is connected and its number of nodes is maximized at \(s+s_D+3\) when it’s a tree. But \(s+s_D+3=(q_m+1)+q_n+(s+s_D-q_n-q_m+2)\) is exactly the total number of nodes in the construction in Part 1 above.

This completes the proof for \(\widehat Y\).

8.2.0.3 Population Pareto sets \(P_s\).

The arguments for \(Y_s\) are analogous, with small differences.

First, at \(s=0\) we have \(Y_0=\tfrac{1}{p}(\tfrac{1}{2}D-R)\), implying by a direct check that \(Y_0\) has the same monomials as \(\widehat Y_0\).

At \(s\ge 1\), due to the omission of \({\boldsymbol{X}}\) edges in the first factor \(\tfrac{1}{2}D-R\), the mergers/pairings contributing to \(Y_s\) have two fewer \({\boldsymbol{X}}\) edges, one fewer \(m\)-nodes, and one fewer \(p\)-nodes compared to those for \(\widehat Y_s\). The realizability construction for \(\widehat Y_s\) extends to \(Y_s\), but with \(q_m\in\overline{1,s}\) \(m\)-nodes after contraction; the numbers of \(n\)- and \(p\)-nodes are still given by \(q_n\) and \(s+s_D-q_n-q_m+2\), respectively.

The proof of optimality is analogous.

8.3 Additional discussion of the diagrammatic approach↩︎

The diagrammatic approach we use to classify the learning regimes following [17] treats the loss evolutions on the level of small-\(t\) asymptotic expansions, by examining the scaling properties of their coefficients. The advantage of this approach is the possibility of a clear and systematic criterion separating different learning regimes. While it may be observed by other means that particular scaling regimes have reasonable nontrivial large-size theoretical limits (e.g., in the mean-field, narrow-latent, etc. settings), it is not obvious what are the concrete boundaries between these regimes and whether there are other regimes that we have possibly missed.

The small-\(t\) expansion approach addresses precisely this question by examining the infinite sequence of expansion coefficients. This sequence encapsulates important information about the dynamics and may display characteristic transitions, signaling changes in the qualitative character of the evolution. This is exactly what Theorem 2 demonstrates: depending on the mutual scaling of the hyper-parameters, the sequence of coefficients displays degeneracies geometrically described in terms of faces of the Pareto prisms.

The connection between the sequence of expansion coefficients and the actual time-dependent loss trajectories \(L(t),\widehat L(t)\) is not a simple question. In [17], it is shown that the degenerate formal small-\(t\) asymptotic expansions resulting in various extreme limiting regimes can often be summed, using appropriate (sometimes non-classical) summation methods, and the results agree well with experiment. We expect that this can also be done for most of our five basic regimes, resulting in the same formulas derived in our Section 9. However, given the large number of regimes, in this paper we restrict ourselves to more direct approaches available for matrix problems, based on random matrix theory and hierarchical systems of ODEs.

A related difficult question is rigorous proofs of convergence of the loss trajectories for finite-size models to their predicted theoretical limits. Here we can draw a parallel with Feynman diagrams in physics, given a long history of related studies [23]. Feynman diagrams also represent coefficients in the asymptotic power series expansions of physical quantities and are widely used, in particular to analyze key features of large systems by formal summation of most relevant parts of associated diagram expansions [24]. However, Feynman diagrams are known to be not so useful for mathematically rigorous convergence proofs, which are typically achieved by other methods [25]. In our case of the autoencoder model the diagrams are ring diagrams, which we expect to simplify the mathematical theory, but still we defer a fully rigorous analysis to future work.

9 Solutions of extreme regimes (4)↩︎

9.1 Notation and preliminary derivations↩︎

Before proceeding with specific cases, let us introduce common notation. Introduce the data covariance matrix \({\boldsymbol{S}}\), weight covariance matrix \({\boldsymbol{A}}\), and model discrepancy \({\boldsymbol{E}}\): \[\boxed{ {\boldsymbol{S}}= \frac{{\boldsymbol{X}}{\boldsymbol{X}}^\top}{m}, \qquad {\boldsymbol{A}}= {\boldsymbol{U}}^\top {\boldsymbol{U}}, \qquad {\boldsymbol{E}}= {\boldsymbol{I}}_p - {\boldsymbol{A}}. }\] Recall 2 , \[\widehat L({\boldsymbol{U}}) = \frac{\left\|{\boldsymbol{X}}- {\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{X}}\right\|_F^2}{2 p m}, \qquad L({\boldsymbol{U}}) = \mathbb{E}_{{\boldsymbol{x}}} \left[\frac{\left\|{\boldsymbol{x}}- {\boldsymbol{U}}^\top {\boldsymbol{U}}{\boldsymbol{x}}\right\|^2}{2 p}\right].\] In new terms, \[\label{eq:losses95as95E} \boxed{ \widehat L({\boldsymbol{E}}) = \frac{1}{2p}\mathop{\mathrm{Tr}}({\boldsymbol{E}}{\boldsymbol{S}}{\boldsymbol{E}}), \qquad L({\boldsymbol{E}}) = \frac{1}{2p}\mathop{\mathrm{Tr}}({\boldsymbol{E}}^2). }\tag{26}\] Since \(d{\boldsymbol{E}}=-d{\boldsymbol{A}}\), \[d\widehat L = -\frac{1}{2p} \mathop{\mathrm{Tr}} \left[ d{\boldsymbol{A}}({\boldsymbol{S}}{\boldsymbol{E}}+ {\boldsymbol{E}}{\boldsymbol{S}}) \right].\] Using \[d{\boldsymbol{A}}=d{\boldsymbol{U}}^\top {\boldsymbol{U}}+{\boldsymbol{U}}^\top d{\boldsymbol{U}},\] and the symmetry of \({\boldsymbol{S}}{\boldsymbol{E}}+ {\boldsymbol{E}}{\boldsymbol{S}}\), \[\frac{\partial \widehat L}{\partial {\boldsymbol{U}}} = -\frac{1}{p} {\boldsymbol{U}}({\boldsymbol{S}}{\boldsymbol{E}}+ {\boldsymbol{E}}{\boldsymbol{S}}).\] Thus the empirical gradient flow is \[\boxed{ \dot{\boldsymbol{U}} = \frac{\eta}{p} {\boldsymbol{U}}\left[{\boldsymbol{S}}({\boldsymbol{I}}_p-{\boldsymbol{A}})+({\boldsymbol{I}}_p-{\boldsymbol{A}}){\boldsymbol{S}}\right]. }\]

9.2 Mean-field (4.1)↩︎

9.2.0.1 Single dim-one sample dynamics.

Consider first the problem of reconstructing a single (\(m = 1\)) data sample \(x = 1\) of dimensionality \(p = 1\): \[\ell({\boldsymbol{u}}) = \frac{1}{2} \left(1 - {\boldsymbol{u}}^\top {\boldsymbol{u}}\right)^2, \qquad \frac{d{\boldsymbol{u}}}{dt} = -\frac{\partial \ell}{\partial {\boldsymbol{u}}}, \quad {\boldsymbol{u}}(0) = \sqrt\rho, \qquad \ell(t) = \ell({\boldsymbol{u}}(t)).\] We shall later use the solution of this auxiliary problem to obtain that for the mean-field regime.

Let \[a(t) = {\boldsymbol{u}}(t)^\top {\boldsymbol{u}}(t).\] The population loss is \[\ell({\boldsymbol{u}}) = \frac{1}{2} (1 - a)^2.\] It satisfies the following ODE. \[\dot{a} = \dot{\boldsymbol{u}}^\top {\boldsymbol{u}}+ {\boldsymbol{u}}^\top \dot{\boldsymbol{u}}.\] Substituting the gradient flow equation gives \[\dot{a} = 4a (1 - a), \qquad a(0) = \rho.\] Its solution is \[a(t) = \frac{\rho}{\rho+(1-\rho)e^{-4t}}.\] Equivalently, \[1-a(t) = \frac{(1-\rho)e^{-4t}}{\rho+(1-\rho)e^{-4t}} = \frac{1-\rho}{\rho e^{4t}+1-\rho}.\] Therefore the loss is \[\label{eq:single95sample95wide95loss} \boxed{ \ell(t) = \frac{1}{2} \left( \frac{1-\rho}{\rho e^{4t}+1-\rho} \right)^2 }.\tag{27}\]

9.2.0.2 Original weight evolution.

The mean-field regime corresponds to vanishing \(\psi\), with \(\phi\) and \(\rho\) staying finite. As the natural choice of learning rate is \(\eta = p\), we get \[\dot{\boldsymbol{U}} = {\boldsymbol{U}}\left[{\boldsymbol{S}}({\boldsymbol{I}}_p-{\boldsymbol{A}})+({\boldsymbol{I}}_p-{\boldsymbol{A}}){\boldsymbol{S}}\right], \qquad {\boldsymbol{U}}(0) \sim {\mathcal{N}}\left(0, \sigma^2 {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right) = {\mathcal{N}}\left(0, \frac{\rho}{n} {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right).\]

Diagonalize the empirical covariance: \[{\boldsymbol{S}}={\boldsymbol{Q}}{\boldsymbol{\Lambda}}{\boldsymbol{Q}}^\top, \qquad {\boldsymbol{\Lambda}}=\operatorname{diag}(\lambda_1,\dots,\lambda_p).\] Rotating the feature coordinates, we may work in the eigenbasis of \({\boldsymbol{S}}\), so that \({\boldsymbol{S}}={\boldsymbol{\Lambda}}\).

In the wide limit, \[{\boldsymbol{A}}(0)={\boldsymbol{U}}(0)^\top {\boldsymbol{U}}(0)\to \rho {\boldsymbol{I}}_p.\] Because the initialization is rotationally invariant and the dynamics are equivariant under rotations that diagonalize \({\boldsymbol{S}}\), the limiting dynamics remain diagonal: \[{\boldsymbol{A}}(t)=\operatorname{diag}(a_1(t),\dots,a_p(t)).\] For a diagonal \({\boldsymbol{A}}(t)\), the empirical flow gives, mode by mode, \[\dot{\boldsymbol{u}}_i = 2\lambda_i(1-a_i){\boldsymbol{u}}_i,\] where \({\boldsymbol{u}}_i\in\mathbb{R}^n\) is the \(i\)-th column of \({\boldsymbol{U}}\), and \[a_i(t)=\|{\boldsymbol{u}}_i(t)\|^2.\] Therefore \[\dot{a}_i = 2{\boldsymbol{u}}_i^\top \dot{\boldsymbol{u}}_i = 4\lambda_i a_i(1-a_i).\] Thus each empirical covariance eigenmode evolves according to \[\dot{a}_i = 4\lambda_i a_i(1-a_i), \qquad a_i(0)=\rho.\] Hence \[a_i(t) = \frac{\rho}{ \rho+(1-\rho)e^{-4\lambda_i t} }.\] Equivalently, \[1-a_i(t) = \frac{1-\rho}{ \rho e^{4\lambda_i t}+1-\rho }.\] Then the loss contribution of mode \(i\) is exactly \(\ell(\lambda_i t)\) with \(\ell\) derived in 27 .

9.2.0.3 Train and population losses.

Since \({\boldsymbol{A}}(t)\) is diagonal in the eigenbasis of \({\boldsymbol{S}}\), \({\boldsymbol{E}}(t)\) also is, and 26 yields \[\widehat L(t) = \frac{1}{2p} \sum_{i=1}^p \lambda_i(1-a_i(t))^2, \qquad L(t) = \frac{1}{2p} \sum_{i=1}^p (1-a_i(t))^2.\] Equivalently, in terms of \(\ell\), \[\widehat L(t) = \frac{1}{p} \sum_{i=1}^p \lambda_i\ell(\lambda_i t), \qquad L(t) = \frac{1}{p} \sum_{i=1}^p \ell(\lambda_i t).\]

As \(p,m\to\infty\) with \(p / m = \phi\), the empirical spectral distribution of \({\boldsymbol{S}}={\boldsymbol{X}}{\boldsymbol{X}}^\top/m\) converges to the Marchenko–Pastur law \(\mu_\phi\): \[d\mu_\phi(\lambda) = \left(1-\frac{1}{\phi}\right)_+\delta_0(d\lambda) + \frac{\sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)}}{2\pi\phi\lambda} \mathbf{1}_{\lambda\in[\lambda_-,\lambda_+]}\,d\lambda,\] where \[\lambda_\pm=(1\pm\sqrt{\phi})^2.\] Therefore the limiting losses are \[\boxed{ {\mathcal{L}}(t) = \int \ell(\lambda t)\,d\mu_\phi(\lambda), \qquad\qquad \widehat{\mathcal{L}}(t) = \int \lambda \ell(\lambda t)\,d\mu_\phi(\lambda) }.\] Explicitly, \[{\mathcal{L}}(t) = \frac{1}{2} \int \left( \frac{1-\rho}{\rho e^{4\lambda t}+1-\rho} \right)^2 d\mu_\phi(\lambda), \qquad \widehat{\mathcal{L}}(t) = \frac{1}{2} \int \lambda \left( \frac{1-\rho}{\rho e^{4\lambda t}+1-\rho} \right)^2 d\mu_\phi(\lambda).\]

9.2.1 Large time behavior↩︎

For every positive empirical covariance eigenvalue \(\lambda>0\), we have \[a_\lambda(t) = \frac{\rho}{ \rho+(1-\rho)e^{-4\lambda t} } \to 1 \qquad \text{as }t\to\infty.\] Therefore \[\ell(\lambda t)\to 0 \qquad \text{for every }\lambda>0.\]

If \(\phi\le 1\), the limiting empirical covariance has no atom at zero, and hence \(\lim_{t\to\infty}L(t)=0\). If \(\phi>1\), then \(\mu_\phi\) has an atom at \(\lambda=0\) of mass \(1-\frac{1}{\phi}\). Those zero-eigenvalue directions are never trained by the empirical loss, so their reconstruction factor remains equal to its initialization value \(\rho\). Hence \[\ell(0) = \frac{1}{2}(1-\rho)^2.\] Therefore \[{\mathcal{L}}(\infty) = \frac{1}{2}(1-\rho)^2 \left(1-\frac{1}{\phi}\right)_+.\]

On the other hand, the empirical train loss contains an additional factor of \(\lambda\). Therefore the zero-eigenvalue directions do not contribute to the empirical loss, and \[\lim_{t\to\infty}\widehat{\mathcal{L}}(t)=0.\]

We now compute the exact train loss asymptotics for large \(t\). For \(\phi = 1\), since \(\lambda_- = 0\) and \(\lambda_+ = 4\), using a Laplace approximation, \[\begin{align} \widehat{\mathcal{L}}(t) &= \frac{1}{2} \int_0^4 \lambda \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \frac{\sqrt{(4-\lambda) \lambda}}{2\pi\lambda} \, d\lambda \\&= \frac{1}{4 \pi} \int_0^4 \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \sqrt{(4-\lambda) \lambda} \, d\lambda \\&\sim \frac{1}{2 \pi} \int_0^\infty \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \sqrt{\lambda} \, d\lambda \end{align}\] as \(t\to\infty\). This integral could be evaluated explicitly: \[\widehat{\mathcal{L}}(t) \sim \frac{\mathrm{Li}_{1/2}\left(\frac{\rho-1}{\rho}\right) - \mathrm{Li}_{3/2}\left(\frac{\rho-1}{\rho}\right)}{32 \sqrt{\pi} t^{3/2}}, \qquad t\to\infty, \qquad \phi = 1.\] In contrast, when \(\phi \neq 1\), \(\lambda_- > 0\), which yields \[\begin{align} \widehat{\mathcal{L}}(t) &= \frac{1}{2} \int_{\lambda_-}^{\lambda_+} \lambda \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \frac{\sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)}}{2\pi\phi\lambda} \, d\lambda \\&= \frac{1}{4 \pi \phi} \int_{\lambda_-}^{\lambda_+} \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)} \, d\lambda \\&\sim \frac{\sqrt{\lambda_+ - \lambda_-}}{4 \pi \phi} \left(\frac{1-\rho}{\rho}\right)^2 \int_{\lambda_-}^{\lambda_+} e^{-8 \lambda t} \sqrt{\lambda-\lambda_-} \, d\lambda \end{align}\] as \(t\to\infty\). This integral could be evaluated explicitly: \[\widehat{\mathcal{L}}(t) \sim \left(\frac{1-\rho}{\rho}\right)^2 \frac{e^{-8 (1 - \sqrt\phi)^2 t}}{64 \sqrt{2\pi} \phi^{3/4} t^{3/2}}, \qquad t\to\infty, \qquad \phi \neq 1.\] Putting all together, \[\boxed{ \widehat{\mathcal{L}}(t) \sim \frac{1}{64 \sqrt{2\pi} \phi^{3/4} t^{3/2}} \begin{cases} 2 \sqrt{2} \left[\mathop{\mathrm{Li}}_{1/2} - \mathop{\mathrm{Li}}_{3/2}\right]\left(\frac{\rho-1}{\rho}\right), & \phi = 1, \\ \left(\frac{1-\rho}{\rho}\right)^2 e^{-8 (1 - \sqrt\phi)^2 t}, & \phi \neq 1 \end{cases} \qquad t\to\infty. }\]

We now perform a similar computation for the population loss. For \(\phi = 1\), again using a Laplace approximation, \[\begin{align} {\mathcal{L}}(t) - {\mathcal{L}}(\infty) &= \frac{1}{2} \int_0^4 \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \frac{\sqrt{(4-\lambda) \lambda}}{2\pi\lambda} \, d\lambda \\&= \frac{1}{4 \pi} \int_0^4 \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \sqrt{\frac{4-\lambda}{\lambda}} \, d\lambda \\&\sim \frac{1}{2 \pi} \int_0^\infty \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \, \frac{d\lambda}{\sqrt{\lambda}} \\&= \frac{\mathrm{Li}_{-1/2}\left(\frac{\rho-1}{\rho}\right) - \mathrm{Li}_{1/2}\left(\frac{\rho-1}{\rho}\right)}{4 \sqrt{\pi} t^{1/2}} \end{align}\] as \(t\to\infty\). When \(\phi \neq 1\), \(\lambda_- > 0\), which yields similarly to the train loss case, \[\begin{align} {\mathcal{L}}(t) - {\mathcal{L}}(\infty) &= \frac{1}{2} \int_{\lambda_-}^{\lambda_+} \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \frac{\sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)}}{2\pi\phi\lambda} \, d\lambda \\&= \frac{1}{4 \pi \phi} \int_{\lambda_-}^{\lambda_+} \left(\frac{1-\rho}{\rho e^{4\lambda t}+1-\rho}\right)^2 \sqrt{(\lambda_+-\lambda)(\lambda-\lambda_-)} \, \frac{d\lambda}{\lambda} \\&\sim \frac{\sqrt{\lambda_+ - \lambda_-}}{4 \pi \phi} \left(\frac{1-\rho}{\rho}\right)^2 \int_{\lambda_-}^{\lambda_+} e^{-8 \lambda t} \sqrt{\lambda-\lambda_-} \, \frac{d\lambda}{\lambda} \\&= \left(\frac{1-\rho}{\rho}\right)^2 \frac{e^{-8 (1 - \sqrt\phi)^2 t}}{64 \sqrt{2\pi} \phi^{3/4} \left(1 - \sqrt\phi\right)^2 t^{3/2}} \end{align}\] as \(t\to\infty\). Putting all together, \[\boxed{ {\mathcal{L}}(t) - {\mathcal{L}}(\infty) \sim \frac{1}{32 \sqrt\pi \phi^{3/4}} \begin{cases} \frac{8}{t^{1/2}} \left[\mathop{\mathrm{Li}}_{-1/2} - \mathop{\mathrm{Li}}_{1/2}\right]\left(\frac{\rho-1}{\rho}\right), & \phi = 1, \\ \left(\frac{1-\rho}{\rho}\right)^2 \frac{e^{-8 \left(1 - \sqrt\phi\right)^2 t}}{\left(1 - \sqrt\phi\right)^2 \left(2 t\right)^{3/2}}, & \phi \neq 1 \end{cases} \qquad t\to\infty. }\]

9.3 Narrow-latent (4.2)↩︎

The narrow-latent regime corresponds to diverging \(\psi\), with \(\phi\) and \(\psi \rho\) staying finite. As the natural choice of learning rate is \(\eta = p\), we get \[\dot{\boldsymbol{U}} = {\boldsymbol{U}}\left[{\boldsymbol{S}}({\boldsymbol{I}}_p-{\boldsymbol{A}})+({\boldsymbol{I}}_p-{\boldsymbol{A}}){\boldsymbol{S}}\right], \qquad {\boldsymbol{U}}(0) \sim {\mathcal{N}}\left(0, \sigma^2 {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right) = {\mathcal{N}}\left(0, \frac{\psi \rho}{p} {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right).\] Equivalently, \[\dot{\boldsymbol{U}} = {\boldsymbol{U}}(2 {\boldsymbol{S}}-{\boldsymbol{S}}{\boldsymbol{A}}-{\boldsymbol{A}}{\boldsymbol{S}}).\] Write the rows of \({\boldsymbol{U}}\) as \[{\boldsymbol{U}}= \begin{pmatrix} {\boldsymbol{u}}_1^\top\\ \vdots\\ {\boldsymbol{u}}_n^\top \end{pmatrix}, \qquad {\boldsymbol{u}}_\alpha \in \mathbb{R}^p.\] Then \[\|{\boldsymbol{u}}_\alpha(0)\|^2\to \psi \rho,\] whereas for \(\alpha\neq \beta\), \[{\boldsymbol{u}}_\alpha^\top(0) {\boldsymbol{u}}_\beta(0)\to 0.\] Thus the rows are asymptotically orthogonal in the narrow-latent limit.

Because of this, the dynamics decouples over rows, and we consider the row dynamics individually. For a single row \({\boldsymbol{u}}\), after dropping cross-row overlaps in the narrow-latent limit, \[{\boldsymbol{A}}\approx {\boldsymbol{u}}{\boldsymbol{u}}^\top.\] Define \[a(t)={\boldsymbol{u}}(t)^\top {\boldsymbol{u}}(t), \qquad r(t)={\boldsymbol{u}}(t)^\top {\boldsymbol{S}}{\boldsymbol{u}}(t).\] Then \[{\boldsymbol{S}}{\boldsymbol{A}}{\boldsymbol{u}} = {\boldsymbol{S}}{\boldsymbol{u}}{\boldsymbol{u}}^\top {\boldsymbol{u}} = a {\boldsymbol{S}}{\boldsymbol{u}}, \qquad {\boldsymbol{A}}{\boldsymbol{S}}{\boldsymbol{u}} ={\boldsymbol{u}}{\boldsymbol{u}}^\top {\boldsymbol{S}}{\boldsymbol{u}} = r {\boldsymbol{u}}.\] Therefore the leading narrow-latent row dynamics are \[\dot{\boldsymbol{u}} = (2-a){\boldsymbol{S}}{\boldsymbol{u}}-r {\boldsymbol{u}}.\]

9.3.0.1 Solving the row dynamics.

Introduce the two scalar clocks \[R(t)=\int_0^t (2-a(s))\,d s, \qquad B(t)=\int_0^t r(s)\,d s.\] Then \[\dot{R}(t)=2-a(t), \qquad \dot{B}(t)=r(t), \qquad R(0)=B(0)=0.\] The row dynamics have the solution \[{\boldsymbol{u}}(t) = e^{-B(t)}e^{R(t){\boldsymbol{S}}} {\boldsymbol{u}}(0).\]

Let \(\mu_\phi\) be the Marchenko–Pastur law associated with \({\boldsymbol{S}}\). Define \[M_\phi(q) = \int e^{q\lambda}\,d\mu_\phi(\lambda), \qquad M_\phi'(q) = \int \lambda e^{q\lambda}\,d\mu_\phi(\lambda).\]

At initialization, \[{\boldsymbol{u}}(0) \sim \mathcal{N}\left(0,\frac{\psi\rho}{p}{\boldsymbol{I}}_p\right).\] Therefore, by self-averaging, \[a(t) = {\boldsymbol{u}}(t)^\top {\boldsymbol{u}}(t) \to \psi \rho e^{-2B(t)}M_\phi(2R(t)),\] and \[r(t) = {\boldsymbol{u}}(t)^\top {\boldsymbol{S}}{\boldsymbol{u}}(t) \to \psi \rho e^{-2B(t)}M_\phi'(2R(t)).\] Hence the finite-data narrow-latent dynamics are \[\boxed{ a(t) = \psi \rho e^{-2B(t)}M_\phi(2R(t)), \qquad r(t) = \psi \rho e^{-2B(t)}M_\phi'(2R(t)), }\] with \[\boxed{ \dot{R}(t)=2-a(t), \qquad \dot{B}(t)=r(t), \qquad R(0)=B(0)=0. }\]

9.3.0.2 Train and population losses.

Expanding 26 , \[\widehat L(t) = \frac{1}{2p}\mathop{\mathrm{Tr}}{\boldsymbol{S}}- \frac{1}{p}\mathop{\mathrm{Tr}}({\boldsymbol{A}}{\boldsymbol{S}}) + \frac{1}{2p}\mathop{\mathrm{Tr}}({\boldsymbol{A}}{\boldsymbol{S}}{\boldsymbol{A}}).\] Since \(\tfrac1p\mathop{\mathrm{Tr}}{\boldsymbol{S}}\to 1\), the leading term is \(1/2\). The other two terms vanish as \(n / p = 1 / \psi\): \[\mathop{\mathrm{Tr}}({\boldsymbol{A}}{\boldsymbol{S}}) = \sum_{\alpha,\beta=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{S}}{\boldsymbol{u}}_\beta) \sim \sum_{\alpha=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{S}}{\boldsymbol{u}}_\alpha) = n r,\] since in the narrow-latent limit, the cross terms are negligible at order \(1/\psi\). Similarly, \[\mathop{\mathrm{Tr}}({\boldsymbol{A}}{\boldsymbol{S}}{\boldsymbol{A}}) = \sum_{\alpha,\beta=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{S}}{\boldsymbol{u}}_\beta) ({\boldsymbol{u}}_\beta^\top {\boldsymbol{u}}_\alpha) \sim \sum_{\alpha=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{S}}{\boldsymbol{u}}_\alpha) ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{u}}_\alpha) = n a r.\] Therefore \[\boxed{ \widehat L(t) = \frac{1}{2} + \frac{r(t)}{\psi} \left[ \frac{1}{2}a(t)-1 \right] + o_{\psi\to\infty}(1/\psi). }\]

As for the population loss, 26 yields \[L(t) = \frac{1}{2} - \frac{1}{p}\mathop{\mathrm{Tr}}({\boldsymbol{A}}) + \frac{1}{2p}\mathop{\mathrm{Tr}}({\boldsymbol{A}}^2).\] The same row-decoupling argument gives \[\mathop{\mathrm{Tr}}({\boldsymbol{A}}) = \sum_{\alpha,\beta=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{u}}_\beta) \sim \sum_{\alpha=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{u}}_\alpha) = n a,\] \[\mathop{\mathrm{Tr}}({\boldsymbol{A}}^2) = \sum_{\alpha,\beta=1}^n ({\boldsymbol{u}}_\beta^\top {\boldsymbol{u}}_\alpha)^2 \sim \sum_{\alpha=1}^n ({\boldsymbol{u}}_\alpha^\top {\boldsymbol{u}}_\alpha)^2 = n a^2.\] Thus \[\boxed{ L(t) = \frac{1}{2} + \frac{a(t)}{\psi} \left[ \frac{1}{2}a(t)-1 \right] + o_{\psi\to\infty}(1/\psi). }\]

9.3.1 Large time behavior↩︎

Introduce \[q(t)=2R(t),\] so that \[\dot{q}(t)=2(2-a(t)).\] The solution obtained above gives \[a(t) = \psi \rho e^{-2B(t)}M_\phi(q(t)), \qquad r(t) = \psi \rho e^{-2B(t)}M_\phi'(q(t)).\] Hence \[\frac{r(t)}{a(t)} = \frac{M_\phi'(q(t))}{M_\phi(q(t))}.\] It is convenient to define \[m_\phi(q) = \frac{M_\phi'(q)}{M_\phi(q)}.\] Then \[r(t)=a(t)m_\phi(q(t)).\]

The upper edge of the Marchenko–Pastur spectrum is \[\lambda_+ = (1+\sqrt{\phi})^2.\] As \(q\to\infty\), the moment-generating function of the Marchenko–Pastur law satisfies \[M_\phi(q) \sim A_\phi e^{\lambda_+q}q^{-3/2},\] where \[A_\phi = \frac{1}{2\sqrt{\pi}\,\phi^{3/4}\lambda_+}.\] Therefore \[m_\phi(q) = \frac{M_\phi'(q)}{M_\phi(q)} = \lambda_+ - \frac{3}{2q} + O(q^{-2}).\]

Since \(a(t)\to 1\), we have \[\dot{q}(t)=2(2-a(t))\to 2,\] and therefore \[q(t)=2t+q_\infty+o(1)\] for some constant \(q_\infty\) depending on the initialization and on \(\phi\). Consequently, \[m_\phi(q(t)) = \lambda_+ - \frac{3}{4t} + O(t^{-2}).\] Since \(a(t)-1\) is exponentially small, this gives \[\boxed{ r(t) = \lambda_+ - \frac{3}{4t} + O(t^{-2}) }\] up to exponentially small corrections.

9.3.1.1 Asymptotics of the row norm.

One can also obtain the large-time behavior of \(a(t)\). From \[\dot{a} = 4r(1-a)\] and \[r=a m_\phi(q), \qquad \dot{q}=2(2-a),\] we get \[\frac{da}{dq} = \frac{2a(1-a)m_\phi(q)}{2-a}.\] Equivalently, \[\frac{2-a}{a(1-a)}\,da = 2m_\phi(q)\,dq.\] Integrating gives the first integral \[2\log a-\log|1-a| = 2\log M_\phi(q)+C.\] Using \(a(0)=\psi \rho\) and \(M_\phi(0)=1\), we obtain \[\frac{a(t)^2}{|1-a(t)|} = \frac{(\psi \rho)^2}{|1-\psi \rho|} M_\phi(q(t))^2, \qquad \psi \rho\neq 1.\] Therefore \[|1-a(t)| = \frac{|1-\psi \rho|}{(\psi \rho)^2} \frac{a(t)^2}{M_\phi(q(t))^2}.\] Since \(a(t)\to 1\), and using the large-\(q\) asymptotic for \(M_\phi\), we find \[|1-a(t)| \sim \frac{|1-\psi \rho|}{(\psi \rho)^2} A_\phi^{-2} e^{-2\lambda_+q(t)}q(t)^3.\] Because \[q(t)=2t+q_\infty+o(1),\] this can be written as \[\boxed{ |1-a(t)| \sim K_a t^3 e^{-4\lambda_+ t}, }\] where \(K_a>0\) is a constant depending on \(\psi \rho\) and \(\phi\). Moreover, the sign of \(a(t)-1\) is the same as the sign of \(\psi \rho-1\). Thus, for \(\psi \rho<1\), \(a(t)\) approaches \(1\) from below, while for \(\psi \rho>1\), it approaches \(1\) from above.

The special case \(\psi \rho=1\) is degenerate at the level of the limiting scalar dynamics: then \[a(t)\equiv 1.\] In this case the row norm is already at its limiting value, but the row direction may still evolve by aligning with high-eigenvalue directions of the empirical covariance.

9.3.1.2 Clock asymptotics.

The clock \(R(t)\) satisfies \[R(t)=\frac{1}{2} q(t),\] hence \[R(t) = t+R_\infty+o(1)\] for some constant \(R_\infty\). The second clock satisfies \[\dot{B}(t)=r(t).\] Using \[r(t)=\lambda_+-\frac{3}{4t}+O(t^{-2}),\] we obtain \[\boxed{ B(t) = \lambda_+ t - \frac{3}{4} \log t + B_\infty + O(t^{-1}). }\] Thus the dominant growth of \(B(t)\) is linear with slope \(\lambda_+\), but there is also a universal logarithmic correction coming from the spectral edge.

9.3.1.3 Large-time population loss .

The population loss along empirical gradient flow is \[\mathbb{E}[L(t)] = \frac{1}{2}-\frac{1}{2 \psi} + \frac{1}{2 \psi}(a(t)-1)^2 + o(1/\psi).\] Therefore \[\mathbb{E}[L(\infty)] = \frac{1}{2}-\frac{1}{2 \psi} + o(1/\psi).\] Using \[|1-a(t)| \sim K_a t^3 e^{-4\lambda_+t},\] we get \[\boxed{ \psi \times \mathbb{E}[L(t) - L(\infty)] \sim \frac{1}{2} K_a^2 t^6 e^{-8\lambda_+t} \qquad t \to \infty. }\] Thus the population loss converges exponentially fast, with rate \(8\lambda_+\), up to the polynomial prefactor \(t^6\).

9.3.1.4 Large-time train loss.

The empirical train loss is \[\mathbb{E}\left[\widehat L(t)\right] = \frac{1}{2} + \frac{r(t)}{\psi} \left( \frac{1}{2} a(t)-1 \right) + o(1/\psi).\] Its limiting value is \[\mathbb{E}\left[\widehat L(\infty)\right] = \frac{1}{2} - \frac{1}{2\psi}\lambda_+ + o(1/\psi).\] Therefore \[\mathbb{E}\left[\widehat L(t) - \widehat L(\infty)\right] = \frac{1}{\psi} \left[ r(t) \left( \frac{1}{2} a(t)-1 \right) + \frac{1}{2}\lambda_+ \right] + o(1/\psi).\] Since \(a(t)-1\) is exponentially small, the leading correction comes from the algebraic convergence of \(r(t)\) to \(\lambda_+\). Using \[r(t) = \lambda_+ - \frac{3}{4t} + O(t^{-2}),\] we get \[\boxed{ \psi \times \mathbb{E}[\widehat L(t) - \widehat L(\infty)] = \frac{3}{8t} + O\left(\frac{1}{t^2}\right) + O\left(t^3 e^{-4\lambda_+t}\right) + o_{\psi\to\infty}(1). }\] Thus, unlike the population loss, the train loss approaches its limiting value only algebraically, with a leading \(1/t\) tail.

9.3.1.5 Qualitative interpretation.

These asymptotic formulas separate two effects of empirical training.

First, the row norm equilibrates exponentially fast: \[a(t)\to 1.\] Since the population loss depends, to leading order in \(1/\psi\), only on \(a(t)\), the population loss also converges exponentially fast to \[L_\phi(\infty) = \frac{1}{2}-\frac{1}{2\psi} + o(1/\psi).\] Thus the population performance saturates quickly once the row norms have reached order one.

Second, the empirical Rayleigh quotient \[r(t)=u(t)^\top S u(t)\] converges much more slowly: \[r(t) = \lambda_+ - \frac{3}{4t} + O(t^{-2}).\] This slow convergence reflects the fact that empirical gradient flow keeps aligning the learned row with increasingly extreme high-variance directions of the sample covariance. The limiting value is the upper spectral edge \[\lambda_+=(1+\sqrt{\phi})^2,\] rather than the population value \(1\).

Consequently, the train loss continues to decrease on a \(1/t\) scale even after the population loss has essentially saturated. In this sense, the late-time part of empirical training is mostly “spectral overfitting”: it improves the empirical loss by exploiting the top edge of the sample covariance spectrum, but it produces no corresponding improvement in the isotropic population loss at order \(1/\psi\).

This also clarifies the role of the sample ratio \(\phi\). Larger \(\phi=p/m\) pushes the upper edge \[\lambda_+=(1+\sqrt{\phi})^2\] farther above \(1\), increasing the final train-test gap. By contrast, the initialization scale \(\psi \rho\) affects the constants and the transient delay, but not the final asymptotic losses, provided \(\rho>0\).

9.4 Large-data (4.3)↩︎

The large data regime corresponds to vanishing \(\phi\), with \(\psi\) and \(\rho\) staying finite. As the natural choice of learning rate is \(\eta = p\), we get \[\dot{\boldsymbol{U}} = {\boldsymbol{U}}\left[{\boldsymbol{S}}({\boldsymbol{I}}_p-{\boldsymbol{A}})+({\boldsymbol{I}}_p-{\boldsymbol{A}}){\boldsymbol{S}}\right], \qquad {\boldsymbol{U}}(0) \sim {\mathcal{N}}(0, \sigma^2 {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p) = {\mathcal{N}}\left(0, \frac{\rho}{n} {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right).\] Since \(m\) is large compared to other dimensions, \({\boldsymbol{S}}\sim {\boldsymbol{I}}_p\), and \[\dot{\boldsymbol{U}} = 2 {\boldsymbol{U}}\left({\boldsymbol{I}}_p-{\boldsymbol{A}}\right).\]

Therefore \[\dot{\boldsymbol{A}} = \dot{\boldsymbol{U}}^\top {\boldsymbol{U}}+{\boldsymbol{U}}^\top \dot{\boldsymbol{U}} = 4 {\boldsymbol{A}}({\boldsymbol{I}}_p-{\boldsymbol{A}}).\] Since the right-hand side is a polynomial in \({\boldsymbol{A}}\), the eigenvectors of \({\boldsymbol{A}}(t)\) are fixed. Each eigenvalue evolves independently according to the scalar logistic equation \[\dot{a}_i(t)=4a_i(t)(1-a_i(t)).\] Hence \[a_i(t) = \frac{a_i(0)e^{4t}}{1+a_i(0)(e^{4t}-1)} = \frac{a_i(0)}{a_i(0)+(1-a_i(0))e^{-4t}}.\] Equivalently, at the matrix level, \[{\boldsymbol{A}}(t) = {\boldsymbol{A}}(0) \left[ {\boldsymbol{A}}(0) + ({\boldsymbol{I}}_p-{\boldsymbol{A}}(0))e^{-4t} \right]^{-1},\] with the convention that zero eigenvalues of \({\boldsymbol{A}}(0)\) remain zero. Therefore the population loss is exactly \[L(t) = \frac{1}{2p} \sum_{i=1}^p \left( \frac{a_i(0)}{a_i(0)+(1-a_i(0))e^{-4t}} -1 \right)^2.\]

9.4.0.1 Closed-form population loss in the proportional limit.

We now average the exact eigenvalue solution over the random initialization. Write \[{\boldsymbol{U}}(0)=\sqrt{\frac{\rho}{n}} {\boldsymbol{Z}},\] where \({\boldsymbol{Z}}\in\mathbb{R}^{n\times p}\) has i.i.d. standard Gaussian entries. Then \[{\boldsymbol{A}}(0) = {\boldsymbol{U}}^\top(0) {\boldsymbol{U}}(0) = \rho\, {\boldsymbol{W}}, \qquad {\boldsymbol{W}}:=\frac{1}{n} {\boldsymbol{Z}}^\top {\boldsymbol{Z}}.\] In the proportional limit \(p = n \psi \to \infty\), the empirical spectral distribution of \({\boldsymbol{W}}\) converges to the Marchenko–Pastur law \(\mu_\psi\).

Define the scalar flow map \[\Theta_t(s) := \frac{s}{s+(1-s)e^{-4t}}.\] Since the initial eigenvalues are \(a_i(0) = \rho\lambda_i\), the limiting population loss is \[{\mathcal{L}}(t) = \lim_{p,n\to\infty} \mathbb{E}[L(t)] = \frac{1}{2} \int \left(\Theta_t(\rho\lambda)-1\right)^2 \, d\mu_\psi(\lambda).\]

This integral has a closed form. Let \[q_t := \rho(1-e^{4t}).\] Then \[\Theta_t(\rho\lambda) = \frac{\rho\lambda e^{4t}}{\rho\lambda e^{4t} - \rho\lambda + 1} = \frac{\rho\lambda e^{4t}}{1-q_t\lambda}.\] Therefore \[\left(\Theta_t(\rho\lambda)-1\right)^2 = 1 - 2 \rho e^{4t} \frac{\lambda}{1-q_t\lambda} + \rho^2 e^{8t} \frac{\lambda^2}{(1-q_t\lambda)^2}.\]

Introduce the Marchenko–Pastur generating function \[H_\psi(q) := \int \frac{\lambda}{1-q\lambda} \,d\mu_\psi(\lambda).\] Then \[H_\psi'(q) = \int \frac{\lambda^2}{(1-q\lambda)^2} \,d\mu_\psi(\lambda).\] Hence \[\boxed{ {\mathcal{L}}(t) = \frac{1}{2} - \rho e^{4t} H_\psi(q_t) + \frac{\rho^2 e^{8t}}{2} H_\psi'(q_t), \qquad q_t = \rho (1-e^{4t}). }\] It remains to write \(H_\psi\) explicitly: \[\boxed{ H_\psi(q) = \frac{1-q\left(1+\psi\right) - \sqrt{1-2q\left(1+\psi\right) + q^2\left(1-\psi\right)^2}}{2 q^2 \psi}. }\] The branch of the square root is chosen so that \(H_\psi(q)\) is regular at \(q=0\). In particular, \[H_\psi(0) = 1, \qquad H_\psi'(0) = 1 + \psi.\] Thus the formula for \({\mathcal{L}}(t)\) is understood by continuous extension at \(t=0\).

9.4.1 Large time behavior↩︎

As \(t\to\infty\), every nonzero eigenvalue of \({\boldsymbol{A}}(0)\) converges to \(1\), while zero eigenvalues remain zero. Hence the limiting final loss is exactly the contribution of the null space of \({\boldsymbol{A}}(0)\): \[{\mathcal{L}}(\infty) = \frac{1}{2} \left(1-\frac{1}{\psi}\right)_+.\] In particular, if \(n\ge p\), equivalently \(\psi\le 1\), then the population loss converges to zero. If \(n<p\), equivalently \(\psi>1\), the rank constraint prevents perfect reconstruction.

These conclusions match the basic linear-algebraic intuition. In the population problem the data covariance is isotropic, so there are no preferred input directions to discover. Gradient flow simply turns the initial row space of \(U\) into an isometric subspace. If the hidden dimension is large enough, this subspace can cover all of \(\mathbb{R}^p\), and the final loss is zero. If the hidden dimension is too small, the best possible solution is a rank-\(n\) projector, leaving a fraction \(1-n/p\) of the input space unreconstructed.

Let us study the precise asymptotics for large \(t\). Recall that for every mode with \(a_i(0)>0\), \[a_i(t)-1 = \frac{(a_i(0)-1)e^{-4t}}{a_i(0)+(1-a_i(0))e^{-4t}},\] and hence, at fixed finite dimension, \[L(t)-L(\infty) = \frac{e^{-8t}}{2p} \sum_{i:a_i(0)>0} \left( \frac{1-a_i(0)}{a_i(0)} \right)^2 + O(e^{-12t}).\] Thus the loss itself relaxes at rate \(e^{-8t}\), while individual eigenvalue errors relax at rate \(e^{-4t}\). The prefactor, however, depends strongly on small initial eigenvalues through inverse powers of \(a_i(0)\). Therefore modes that are initialized very close to zero are learned very slowly.

In the proportional limit, with \(p/n = \psi\) and \(a_i(0)=\rho\lambda_i\), the same conclusion gives \[{\mathcal{L}}(t) - {\mathcal{L}}(\infty) \sim \frac{e^{-8t}}{2} \int_{\lambda > 0} \left(\frac{1-\rho\lambda}{\rho\lambda}\right)^2 \, d\mu_\psi(\lambda) \qquad t \to \infty, \qquad \psi \neq 1.\] At the critical aspect ratio \(\psi=1\), however, the MP density touches zero with a square-root singularity. In that case the above coefficient diverges, and the large-time decay in the proportional limit is slower: \[{\mathcal{L}}(t) - {\mathcal{L}}(\infty) \sim \frac{e^{-2t}}{4\sqrt{\rho}}, \qquad t \to \infty, \qquad \psi = 1.\] This is a useful non-commutation of limits: at any fixed finite dimension with full rank, the eventual decay is exponential with rate \(e^{-8t}\), but after taking the proportional limit at the square point \(n=p\), the hard edge of the MP spectrum produces the slower decay \(e^{-2t}\).

9.5 Small-data (4.4)↩︎

The small data regime corresponds to diverging \(\phi\), with \(\psi\) and \(\rho\) staying finite. As the natural choice of learning rate is now \(\eta = m\), we get \[\dot{\boldsymbol{U}} = \frac{1}{\phi} {\boldsymbol{U}}\left[{\boldsymbol{S}}{\boldsymbol{E}}+ {\boldsymbol{E}}{\boldsymbol{S}}\right], \qquad {\boldsymbol{U}}(0) \sim {\mathcal{N}}(0, \sigma^2 {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p) = {\mathcal{N}}\left(0, \frac{\rho}{n} {\boldsymbol{I}}_n \otimes {\boldsymbol{I}}_p\right).\]

We rotate coordinates so that the data span is the first \(m\)-dimensional coordinate subspace. In the small-sample regime, and exactly for orthogonalized data \({\boldsymbol{X}}=\sqrt p\,{\boldsymbol{Q}}\), \({\boldsymbol{Q}}^\top {\boldsymbol{Q}}= {\boldsymbol{I}}_m\), we have \[{\boldsymbol{S}}= \phi {\boldsymbol{P}}_m, \qquad {\boldsymbol{P}}_m = \begin{pmatrix} {\boldsymbol{I}}_m & 0\\ 0 & 0 \end{pmatrix}.\] Substituting this into the gradient flow equation above, the factors \(\phi\) and \(1/\phi\) cancel: \[\dot{\boldsymbol{U}} = {\boldsymbol{U}}({\boldsymbol{P}}_m {\boldsymbol{E}}+ {\boldsymbol{E}}{\boldsymbol{P}}_m). \label{eq:tied95lr95m95projected95gf}\tag{28}\]

Decompose \[{\boldsymbol{U}}= \begin{bmatrix} {\boldsymbol{U}}_\parallel & {\boldsymbol{U}}_\perp \end{bmatrix}, \qquad {\boldsymbol{U}}_\parallel\in\mathbb{R}^{n\times m}, \qquad {\boldsymbol{U}}_\perp\in\mathbb{R}^{n\times(p-m)}.\] For notational simplicity, define \[{\boldsymbol{A}}={\boldsymbol{U}}_\parallel, \qquad {\boldsymbol{B}}={\boldsymbol{U}}_\perp, \qquad {\boldsymbol{M}}={\boldsymbol{B}}{\boldsymbol{B}}^\top.\] Since \[{\boldsymbol{E}} = {\boldsymbol{I}}_p-{\boldsymbol{U}}^\top {\boldsymbol{U}} = \begin{pmatrix} {\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}}& -{\boldsymbol{A}}^\top {\boldsymbol{B}}\\ -{\boldsymbol{B}}^\top {\boldsymbol{A}}& {\boldsymbol{I}}_{p-m}-{\boldsymbol{B}}^\top {\boldsymbol{B}} \end{pmatrix},\] 28 gives the reduced dynamics \[\begin{align} \dot{\boldsymbol{A}} &= 2 {\boldsymbol{A}}({\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}})-{\boldsymbol{M}}{\boldsymbol{A}}, \tag{29}\\ \dot{\boldsymbol{B}} &= -{\boldsymbol{A}}{\boldsymbol{A}}^\top {\boldsymbol{B}}. \tag{30} \end{align}\] Since \({\boldsymbol{M}}={\boldsymbol{B}}{\boldsymbol{B}}^\top\), we also have \[\dot{\boldsymbol{M}} = \dot{\boldsymbol{B}}{\boldsymbol{B}}^\top+{\boldsymbol{B}}\dot{\boldsymbol{B}}^\top = -{\boldsymbol{A}}{\boldsymbol{A}}^\top {\boldsymbol{M}}-{\boldsymbol{M}}{\boldsymbol{A}}{\boldsymbol{A}}^\top. \label{eq:tied95lr95m95M95dynamics}\tag{31}\]

9.5.0.1 Reduced empirical loss.

From 26 , \[\widehat L({\boldsymbol{U}}) = \frac{1}{2pm} \left\|{\boldsymbol{E}}{\boldsymbol{X}}\right\|_F^2.\] Using \({\boldsymbol{X}}=\sqrt p\,{\boldsymbol{Q}}\), where the columns of \({\boldsymbol{Q}}\) span the first \(m\)-dimensional coordinate subspace, this becomes \[\widehat L({\boldsymbol{U}}) = \frac{1}{2m} \left\|{\boldsymbol{E}}{\boldsymbol{P}}_m\right\|_F^2.\] Since \[{\boldsymbol{E}}{\boldsymbol{P}}_m = \begin{pmatrix} {\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}}\\ -{\boldsymbol{B}}^\top {\boldsymbol{A}} \end{pmatrix},\] we obtain \[\widehat L({\boldsymbol{U}}) = \frac{1}{2m} \left\|{\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}}\right\|_F^2 + \frac{1}{2m} \left\|{\boldsymbol{B}}^\top {\boldsymbol{A}}\right\|_F^2. \label{eq:tied95lr95m95train95loss95exact}\tag{32}\] Equivalently, \[\widehat L({\boldsymbol{U}}) = \frac{1}{2m} \left\|{\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}}\right\|_F^2 + \frac{1}{2m} \mathop{\mathrm{Tr}}({\boldsymbol{A}}^\top {\boldsymbol{M}}{\boldsymbol{A}}).\]

9.5.0.2 Order parameters.

Define \[a = \frac{1}{m} \mathop{\mathrm{Tr}}({\boldsymbol{A}}^\top {\boldsymbol{A}}), \qquad y=1-a.\]

Introduce also the active leakage moments \[\mu_k = \frac{1}{ma} \mathop{\mathrm{Tr}}({\boldsymbol{A}}^\top {\boldsymbol{M}}^k {\boldsymbol{A}}), \qquad k\ge0.\] For brevity, denote \(q = \mu_1\) and \(r = \mu_2\). The closure assumes \[{\boldsymbol{A}}^\top {\boldsymbol{M}}^k {\boldsymbol{A}} \approx a\mu_k {\boldsymbol{I}}_m.\] The empirical loss 32 becomes \[\widehat L(t) = \frac{1}{2}y(t)^2 + \frac{1}{2}a(t)q(t). \label{eq:tied95lr95m95train95loss95scalar}\tag{33}\]

Using this in 29 , we obtain \[\dot{a} = 4a(1-a)-2a\mu_1. \label{eq:tied95lr95m95adot95closed}\tag{34}\] Equivalently, since \(y=1-a\), \[\dot{y} = -4a y+2a\mu_1. \label{eq:tied95lr95m95ydot95closed}\tag{35}\]

Also from 29 , \[\dot{q} = -2\mu_2 - 2a q + 2q^2.\]

9.5.0.3 Moment hierarchy.

Using 29 and 31 one obtains, under the same scalar isotropic closure, \[\boxed{ \dot{\mu}_k = -2\mu_{k+1} - 2a \sum_{j=0}^{k-1} \mu_j\mu_{k-j} + 2\mu_1\mu_k, \qquad k\ge1. } \label{eq:tied95lr95m95moment95hierarchy}\tag{36}\] The correct object is the active spectral measure \(\nu_t\) of \({\boldsymbol{M}}\) seen by the columns of \({\boldsymbol{A}}\), defined by \[\mu_k(t) = \int \lambda^k\,d\nu_t(\lambda) = \frac{1}{ma(t)} \mathop{\mathrm{Tr}}\left[{\boldsymbol{A}}(t)^\top {\boldsymbol{M}}(t)^k {\boldsymbol{A}}(t)\right].\]

9.5.0.4 Initial conditions.

At initialization, the active and inactive blocks are independent: \[{\boldsymbol{A}}(0)={\boldsymbol{U}}_\parallel(0), \qquad {\boldsymbol{B}}(0)={\boldsymbol{U}}_\perp(0).\] Since \(m / n = \psi / \phi \to 0\), \[{\boldsymbol{A}}(0)^\top {\boldsymbol{A}}(0)\approx \rho {\boldsymbol{I}}_m,\] and therefore \[a(0)=\rho, \qquad y(0)=1-\rho.\] The leakage matrix \[{\boldsymbol{M}}(0)={\boldsymbol{B}}(0){\boldsymbol{B}}(0)^\top\] is Wishart-like with effective inactive aspect ratio \[\psi_\perp = \frac{p-m}{n}.\] Thus \[\mu_1(0) = q(0) = \rho\psi_\perp,\] and, in the limit \(p / m = \phi \to \infty\), \[q(0) \to \psi \rho.\] The second moment is \[\mu_2(0) = \rho^2(\psi_\perp^2+\psi_\perp).\] More generally, the initial moments are the moments of the Marchenko–Pastur law: \[\mu_k(0) = \rho^k \sum_{j=1}^k N(k,j)\psi_\perp^j,\] where \[N(k,j) = \frac{1}{k} \binom{k}{j} \binom{k}{j-1}\] are the Narayana numbers.

9.5.0.5 Train loss evolution.

Summing up, \[\boxed{ \mathbb{E}\left[\widehat L(t)\right] = \frac{1}{2} \left(1-a(t)\right)^2 + \frac{1}{2}a(t)\mu_1(t). } \label{eq:tied95lr95m95train95loss95final}\tag{37}\] The dynamical system is \[\begin{align} \dot{a} &= 4a(1-a)-2a\mu_1, \tag{38}\\ \dot{\mu}_k &= -2\mu_{k+1} - 2a \sum_{j=0}^{k-1} \mu_j\mu_{k-j} + 2\mu_1\mu_k, \qquad k\ge1. \tag{39} \end{align}\]

9.5.0.6 Population loss evolution.

We now compute the population loss correction. We have \[{\boldsymbol{U}}^\top {\boldsymbol{U}} = \begin{pmatrix} {\boldsymbol{A}}^\top {\boldsymbol{A}}& {\boldsymbol{A}}^\top {\boldsymbol{B}}\\ {\boldsymbol{B}}^\top {\boldsymbol{A}}& {\boldsymbol{B}}^\top {\boldsymbol{B}} \end{pmatrix}.\] Therefore \[\begin{align} L(t) &= \frac{1}{2p} \left\|{\boldsymbol{I}}_p-{\boldsymbol{U}}^\top {\boldsymbol{U}}\right\|_F^2 \\&= \frac{1}{2p} \left\|{\boldsymbol{I}}_m-{\boldsymbol{A}}^\top {\boldsymbol{A}}\right\|_F^2 + \frac{1}{p} \left\|{\boldsymbol{B}}^\top {\boldsymbol{A}}\right\|_F^2 + \frac{1}{2p} \left\|{\boldsymbol{I}}_{p-m}-{\boldsymbol{B}}^\top {\boldsymbol{B}}\right\|_F^2. \end{align} \label{eq:tied95lr95m95distribution95block95decomp}\tag{40}\] The first two terms give \[\frac{1}{\phi} \left[ \frac{1}{2}(1-a)^2 + aq\right]. \label{eq:tied95lr95m95distribution95active95terms}\tag{41}\] The factor in front of \(aq\) is \(1\), not \(1/2\), because the population loss contains both off-diagonal blocks \({\boldsymbol{A}}^\top {\boldsymbol{B}}\) and \({\boldsymbol{B}}^\top {\boldsymbol{A}}\), while the empirical loss sees only \({\boldsymbol{B}}^\top {\boldsymbol{A}}\).

Define the inactive block contribution \[F_\perp(t) = \frac{1}{2p} \left\|{\boldsymbol{I}}_{p-m}-{\boldsymbol{B}}^\top(t) {\boldsymbol{B}}(t)\right\|_F^2.\] Differentiating and using \(\dot{\boldsymbol{B}}=-{\boldsymbol{A}}{\boldsymbol{A}}^\top {\boldsymbol{B}}\), we get \[\dot{F}_\perp = \frac{2}{p} \mathop{\mathrm{Tr}}\left[({\boldsymbol{I}}_{p-m}-{\boldsymbol{B}}^\top {\boldsymbol{B}}){\boldsymbol{B}}^\top {\boldsymbol{A}}{\boldsymbol{A}}^\top {\boldsymbol{B}}\right] = \frac{2}{p} \left[\mathop{\mathrm{Tr}}({\boldsymbol{A}}^\top {\boldsymbol{M}}{\boldsymbol{A}}) - \mathop{\mathrm{Tr}}({\boldsymbol{A}}^\top {\boldsymbol{M}}^2 {\boldsymbol{A}})\right] = \frac{2}{\phi} (aq-r). \label{eq:tied95lr95m95Fperp95dot}\tag{42}\] Therefore \[F_\perp(t) = F_\perp(0) + \frac{2}{\phi} \int_0^t \left[a(s)q(s)-r(s)\right] \, ds + o_{\phi\to\infty}\left(\frac{1}{\phi}\right). \label{eq:tied95lr95m95Fperp95integrated}\tag{43}\]

At initialization, \[\begin{align} F_\perp(0) &= \frac{1}{2p} \left\|{\boldsymbol{I}}_{p-m}-{\boldsymbol{B}}(0)^\top {\boldsymbol{B}}(0)\right\|_F^2 \\&= \frac{1}{2} \left[1-2\rho+(1+\psi)\rho^2\right] - \frac{1}{2 \phi} \left[1-2\rho+(1+2\psi)\rho^2\right] + o_{\phi\to\infty}\left(\frac{1}{\phi}\right). \end{align} \label{eq:tied95lr95m95Fperp95initial}\tag{44}\] Combining equation 40 , equation 41 , equation 43 , and equation 44 , we obtain \[\boxed{ \begin{align} \mathbb{E}[L(t)] &= \frac{1}{2} \left[ 1-2\rho+(1+\psi)\rho^2 \right] \\ &\quad + \frac{1}{\phi} \Bigg\{ \frac{1}{2} \left[ 1-a(t) \right]^2 + a(t) \mu_1(t) - \frac{1}{2} \left[ 1-2\rho+(1+2\psi)\rho^2 \right] \\ & + 2\int_0^t a(s) \left[\mu_1(s)-\mu_2(s)\right] \, ds \Bigg\} + o_{\phi\to\infty}\left(\frac{1}{\phi}\right). \end{align} } \label{eq:tied95lr95m95distribution95loss95final}\tag{45}\]

10 Experimental details↩︎

All experimental results in Fig. 3 and Fig. 4 can be reproduced from a single Jupyter notebook provided in the supplementary material. The experiments were run on a single NVIDIA Tesla P100 GPU and took approximately 40 minutes in total. The train datasets were generated synthetically as Gaussian matrices \({\boldsymbol{X}}\in\mathbb{R}^{p\times m}\) with iid \({\mathcal{N}}(0,1)\) entries, as in 2. For each seed, both the dataset and initialization were resampled.

10.1 Discrete-time implementation↩︎

Our theoretical results are stated for gradient flow, see 2 . In the experiments, we use full-batch gradient descent, interpreted as an explicit Euler discretization of the gradient-flow dynamics: \[{\boldsymbol{U}}_{k+1} = {\boldsymbol{U}}_k - \tau \eta \frac{\partial \widehat L({\boldsymbol{U}}_k)}{\partial {\boldsymbol{U}}_k}.\] Thus the gradient descent learning rate is \(\eta_{\rm GD}=\tau\eta\). As \(\tau\to 0\), the discrete-time dynamics converge to gradient flow; hence, for sufficiently small \(\tau\), gradient flow accurately describes full-batch gradient descent.

10.2 Experimental parameters↩︎

For each regime, we fixed a final time \(t_{\max}\) and a number of gradient descent steps \(N_{\rm step}\). The Euler step size was then \(\tau=t_{\max}/N_{\rm step}\), and the implemented gradient descent learning rate was \(\eta_{\rm GD}=\tau\eta\), where \(\eta\) denotes the corresponding gradient-flow learning rate. The values of \(t_{\max}\) and \(N_{\rm step}\) were chosen so that decreasing \(\tau\) further did not visibly change the curves.

Table 3 lists the parameters used in all our experiments. For each setting, the plotted curve is the empirical mean over independent random seeds. Error bars in our figures indicate \(\pm 2\) standard errors of the mean, computed across seeds.

For the small-data regime, the analytical solution is expressed through the moment hierarchy in 15 36 . To evaluate this solution numerically, we truncate the hierarchy at order \(K\) by imposing the closure \(\mu_{K+1}=0\). In all comparisons with gradient descent in Fig. 3, we use \(K=128\). We also checked that reducing the truncation to \(K=32\) gives no visible change on the plotted time intervals.

References↩︎