The Impact of Anisotropic Covariance Structure on the Training Dynamics and Generalization Error of Linear Networks

TL;DR

This work studies how data anisotropy, modeled by a spiked covariance structure, shapes learning dynamics and generalization in a two-layer linear network trained for linear regression. By introducing a data-adapted basis, the authors show the training proceeds in two phases: an initial rapid growth along the input-output correlation direction, followed by a later adjustment along an orthogonal principal direction that captures additional structure. They derive an analytical generalization error in the high-dimensional regime, revealing that larger spike magnitude and stronger spike-target alignment can substantially reduce error, with explicit dependence on the alignment $A$ and spike magnitude $\rho$. Overall, the results provide a principled understanding of how anisotropic data structures influence learning trajectories and final performance, offering guidance for leveraging data structure in simple networks and informing analyses of more complex models.

Abstract

The success of deep neural networks largely depends on the statistical structure of the training data. While learning dynamics and generalization on isotropic data are well-established, the impact of pronounced anisotropy on these crucial aspects is not yet fully understood. We examine the impact of data anisotropy, represented by a spiked covariance structure, a canonical yet tractable model, on the learning dynamics and generalization error of a two-layer linear network in a linear regression setting. Our analysis reveals that the learning dynamics proceed in two distinct phases, governed initially by the input-output correlation and subsequently by other principal directions of the data structure. Furthermore, we derive an analytical expression for the generalization error, quantifying how the alignment of the spike structure of the data with the learning task improves performance. Our findings offer deep theoretical insights into how data anisotropy shapes the learning trajectory and final performance, providing a foundation for understanding complex interactions in more advanced network architectures.

Figures (5)

• Figure 1: Learning trajectories in a phase plane with the spike direction ${\bm{\mu}}$ as the horizontal axis and the perpendicular direction ${\bm{\mu}}_\perp=({\bm{I}}-{\bm{\mu}}{\bm{\mu}}^\top){\bm{\beta}}/\|({\bm{I}}-{\bm{\mu}}{\bm{\mu}}^\top){\bm{\beta}}\|$ as the vertical axis. The vector field and nullclines (blue: $dw_1/dt=0$, orange: $dw_2/dt=0$) are from the reduced scalar dynamics (\ref{['eq:dynamics_w1_scalar', 'eq:dynamics_w2_scalar']}). The red trajectory is projected from a simulation, while the black trajectory is from the reduced scalar dynamics (\ref{['eq:dynamics_w1_scalar', 'eq:dynamics_w2_scalar']}). The basis directions, ${\bm{v}}_1$ (dashed) and ${\bm{v}}_2$ (dash-dotted), are also shown. The close match between trajectories validates the dimensionality reduction analysis. The parameters are set to $A=0.3$, $\rho=20$, and $\sigma^2=1$, with the number of training samples $n=10000$. The network parameters are $d=30$ and $m=50$ and the initialization is performed as $W_{ij}\sim\mathcal{N}(0,s^2/d)$ and $a_i\sim\mathcal{N}(0,s^2/m)$ with $s=10^{-5}$.
• Figure 2: Time evolution of the training loss and magnitudes of weight components. The plotted values are the numerically obtained training loss (black) and the projections of the weight vectors ${\bm{w}}_1$ (blue) and ${\bm{w}}_2$ (orange) onto the initial growth direction $\hat{{\bm{r}}}_1$, namely, $\hat{{\bm{r}}}_1^\top {\bm{W}} {\bm{v}}_1$ and $\hat{{\bm{r}}}_1^\top {\bm{W}} {\bm{v}}_2$, respectively. The black dashed line is the analytical approximation for the early phase, initialized by $u(0)=\hat{{\bm{r}}}_1^\top({\bm{a}}(0)+{\bm{w}}_1(0))/2$ (\ref{['eq:w1_analytical']}). Vertical dotted lines and vertical dash-dotted lines indicate characteristic timescales for the early and later phases, respectively, with the latter defined in \ref{['eq:later-timescale']} using $\delta=3$. The parameters are the same as those in \ref{['fig:phase_space_dynamics_w1_w2']}.
• Figure 3: Evolution of the normalized training loss for different settings of (a) spike magnitude $\rho$ and (b) spike-target alignment $A$. The loss $\mathcal{L}$ is normalized by the empirical second moment of the labels, $\frac{1}{n}\sum_{i=1}^n y_i^2$. Vertical dotted and dash-dotted lines indicate the characteristic timescales for the early and later learning phases, respectively, with the latter defined in \ref{['eq:later-timescale']} using $\delta=3$. Other parameters are the same as in \ref{['fig:phase_space_dynamics_w1_w2']}.
• Figure 4: Normalized generalization error $\hat{R}_0/{\bm{\beta}}^\top{\bm{\Sigma}}{\bm{\beta}}$ as a function of spike-target alignment $A$ and spike magnitude $\rho$ for a fixed ratio $\gamma=3$.
• Figure 5: Normalized generalization error $\hat{R}_0/{\bm{\beta}}^\top{\bm{\Sigma}}{\bm{\beta}}$ as a function of the ratio $\gamma=d/n$. Solid lines represent the theoretical prediction (\ref{['eq:generalization_error_final']}) and points are results of numerical simulations. The empirical results are averaged over 10 independent trials, with error bars indicating the standard deviation. All simulations were performed with a two-layer linear network of input dimension $d=3000$ and width $m=3000$, with $\sigma^2=1$. The initialization procedure is the same as in \ref{['fig:phase_space_dynamics_w1_w2']}.