SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alireza Mousavi-Hosseini<sup>1</sup>, Denny Wu<sup>2</sup>, Taiji Suzuki<sup>3</sup>, Murat A. Erdogdu<sup>1</sup>
<sup>1</sup>University of Toronto and Vector Institute,
$^{2}$ New York University and Flatiron Institute,
<sup>3</sup>University of Tokyo and RIKEN AIP
{mousavi,erdogdu}@cs.toronto.edu,dennywu@nyu.edu,
taiji@mist.i.u-tokyo.ac.jp
Recent works have demonstrated that the sample complexity of gradient-based learning of single index models, i.e. functions that depend on a 1-dimensional projection of the input data, is governed by their information exponent. However, these results are only concerned with isotropic data, while in practice the input often contains additional structure which can implicitly guide the algorithm. In this work, we investigate the effect of a spiked covariance structure and reveal several interesting phenomena. First, we show that in the anisotropic setting, the commonly used spherical gradient dynamics may fail to recover the true direction, even when the spike is perfectly aligned with the target direction. Next, we show that appropriate weight normalization that is reminiscent of batch normalization can alleviate this issue. Further, by exploiting the alignment between the (spiked) input covariance and the target, we obtain improved sample complexity compared to the isotropic case. In particular, under the spiked model with a suitably large spike, the sample complexity of gradient-based training can be made independent of the information exponent while also outperforming lower bounds for rotationally invariant kernel methods.
A fundamental feature of neural networks is their adaptivity to learn unknown statistical models. For instance, when the learning problem exhibits certain low-dimensional structure or sparsity, it is expected that neural networks optimized by gradient-based algorithms can efficiently adapt to such structure via feature/representation learning. A considerable amount of research has been dedicated to understanding this phenomenon under various assumptions and to demonstrate the superiority of neural networks over non-adaptive methods such as kernel models [GMMM19, WLLM19,  $\mathrm{BES}^{+}19$ , LMZ20, AAM22,  $\mathrm{BES}^{+}22$ , DLS22, Tel23,  $\mathrm{MHPG}^{+}23]$ .
A particular relevant problem setting for feature learning is the estimation of single index models, where the response  $y \in \mathbb{R}$  depends on the input  $\boldsymbol{x} \in \mathbb{R}^d$  via  $y = g(\langle \boldsymbol{u}, \boldsymbol{x} \rangle) + \epsilon$ , where  $g: \mathbb{R} \to \mathbb{R}$  is the nonlinear link function and  $\boldsymbol{u}$  is the unit target direction. Here, learning corresponds to recovering the unknowns  $\boldsymbol{u}$  and  $g$ , which requires the model to extract and adapt to the low-dimensional target direction. Recent works have shown that the sample complexity is determined by certain properties of the link function  $g$ . In particular, the complexity of gradient-based optimization is captured by the information exponent of  $g$  introduced by [BAGJ21]. Intuitively, a larger information exponent  $s$  corresponds to a more complex  $g$  (for gradient-based learning), and it has been proven that when the input is isotropic  $x \sim \mathcal{N}(0, \mathbf{I}_d)$ , gradient flow can learn the single index model with  $\tilde{\mathcal{O}}(d^s)$  sample complexity [BBSS22].
In practice, however, real data always exhibits certain structures such as low intrinsic dimensionality, and isotropic data assumptions fail to capture this fact. In statistics methodology, it is known that the directions along which the input  $\pmb{x}$  has high variance are often good predictors of the target  $y$  [HTFF09]; indeed, this is the main reason principal component analysis is used in pretraining [JWHT13]. A fundamental model that captures such a structure is the spiked matrix model in which  $\pmb{x} \sim \mathcal{N}(0,\mathbf{I}_d + \kappa \pmb{\theta}\pmb{\theta}^\top)$  for some unit direction  $\pmb{\theta} \in \mathbb{R}^d$  and  $\kappa > 0$  [Joh01]. Along the direction  $\pmb{\theta}$ , data has higher variability and predictive power. In single index models, such predictive power translates to a non-trivial alignment between the vectors  $\pmb{u}$  and  $\pmb{\theta}$  — our focus is to investigate the effect of such alignment on the sample complexity of gradient-based training.
1.1 Contributions: learning single index models under spiked covariance
In this paper, we study the sample complexity of learning a single index model using a two-layer neural network and show that it is determined by an interplay between
- spike-target alignment:  $\langle \pmb{u},\pmb{\theta}\rangle \asymp d^{-r_1}$ ,  $r_1\in [0,1 / 2]$
- spike magnitude:  $\kappa \asymp d^{r_2}$ , for  $r_2 \in [0,1]$ .
Our contributions can be summarized as follows.
1. We show that even in the case of perfect spike-target alignment  $(r_1 = 0)$ , the spherical gradient flow commonly employed in recent literature (see e.g. [BAGJ21, BBSS22]) cannot recover the target direction for moderate spike magnitudes in the population limit. The failure of this covariance-agnostic procedure under anisotropic structure insinuates the necessity of an appropriate covariance-aware normalization to effectively learn the single index model.
Figure 1: Sample complexity to learn  $\mathbf{u}$  and  $g$  under the spiked model. Smaller  $r_1$  denotes a better spike-target alignment, while larger  $r_2$  denotes a larger spike magnitude. The sample complexities are based on Corollary 8.
2. We show that a covariance-aware normalization that resembles batch normalization resolves this issue. Indeed, the resulting gradient flow can successfully recover the target direction  $\pmb{u}$  in this case, and depending on the amount of spike-target alignment, the sample complexity can significantly improve compared to the isotropic case.
3. Under the spiked covariance model, we prove a three-stage phase transition for the sample complexity depending on the quantities  $r_1$  and  $r_2$ . For a suitable direction and magnitude of the spike, the sample complexity can be made  $\tilde{\mathcal{O}}(d^{3 + \nu})$  for any  $\nu > 0$  which is independent of the information exponent  $s$ . This should be compared against the known complexity of  $\tilde{\mathcal{O}}(d^s)$  under isotropic data.
4. We finally show that preconditioning the training dynamics with the inverse covariance improves the sample complexity. This is particularly significant for the spiked covariance model where  $\tilde{\mathcal{O}}(d^{3+\nu})$  samples can be reduced to  $\tilde{\mathcal{O}}(d^{1+\nu})$  for any  $\nu > 0$ , i.e. almost linear in  $d$ . The three-stage phase transition also emerges, as illustrated in Figure 1: in the "hard" regime, the complexity remains  $\tilde{\mathcal{O}}(d^s)$  regardless of the magnitude and direction of the spike, while in the "easy" regime the complexity only depends on the spike magnitude and not its direction. The "intermediate" regime interpolates between these two; smaller  $r_1$  and larger  $r_2$  improve the sample complexity.
The rest of the paper is organized as follows. We discuss the notation and the related work in the remainder of this section. We provide preliminaries on the statistical model and the training procedure in Section 2, and provide a negative result on the covariance-agnostic gradient flow in Section 2.1. Our main sample complexity result on a single neuron is presented in Section 3.2. We provide our results on multi-neuron neural networks in Section 4 and also discuss extensions such as preconditioning and its implications. We provide a technical summary in Section 5 and conclude in Section 6.
Notation. We use  $\langle \cdot, \cdot \rangle$  and  $\|\cdot\|$  to denote Euclidean inner product and norm. For matrices,  $\|\cdot\|$  denotes the usual operator norm, and  $\lambda_{\max}(\cdot)$  and  $\lambda_{\min}(\cdot)$  denote the largest and smallest eigenvalues respectively. We reserve  $\gamma$  for the standard Gaussian distribution on  $\mathbb{R}$ , and let  $\|\cdot\|_{\gamma}$  denote the  $L^2(\gamma)$  norm.  $\mathbb{S}^{d-1}$  is the unit  $d$ -dimensional sphere. For quantities  $a$  and  $b$ , we will use  $a \lesssim b$  to convey there exists a constant  $C$  (a universal constant unless stated otherwise, in which case may depend on polylogarithmic factors of  $d$ ) such that  $a \leq C b$ , and  $a \asymp b$  signifies that  $a \lesssim b$  and  $b \lesssim a$ .
1.2 Further related work
Non-Linear Feature Learning with Neural Networks. Recently, two popular scaling regimes of neural networks have emerged for theoretical studies. A large initialization variance leads to the lazy training regime, where the weights do not move significantly, and the training dynamics is captured by the neural tangent kernel (NTK) [JGH18, COB19]. However, there are many instances of function classes that are efficiently learnable by neural networks and not efficiently learnable by the NTK [YS19, GMMM19]. Under a smaller initialization scaling, gradient descent on infinite-width neural networks becomes equivalent to Wasserstein gradient flow on the space of measures, known as the mean-field limit [CB18, RVE18, MMN18, MMM19, NWS22, Chi22], which can learn certain low-dimensional target functions efficiently [WLLM19, AAM22, HC22, ASKL23].
As for neural networks with smaller width, recent works showed that a two-stage feature learning procedure can outperform the NTK when the data is sampled uniformly from the hypercube  $\left[\mathrm{BEG}^{+}22\right]$  or isotropic Gaussian [DLS22,  $\mathrm{BES}^{+}22$ , BBSS22,  $\mathrm{MHPG}^{+}23$ , ABAM23]. However, these results do not take into account the additional structure that might be present in the covariance matrix of the input data. Two notable exceptions are [GMMM20, RGKZ21], where the authors analyzed a spiked covariance and Gaussian mixture data, respectively. Our setting is closer to [GMMM20], however, they do not provide optimization guarantees through gradient-based training. Furthermore, in a companion work  $\left[\mathrm{BES}^{+}23\right]$ , we zoom into the setting where the spike and target are perfectly aligned  $(r_{1} = 0)$ , and prove learnability in the  $n \asymp d$  regime for both kernel regression and two-layer neural network. Finally, we go over some results concurrent to our work in Appendix A.
Learning Single Index Models. The problem of estimating the relevant direction in a single index model is classical in statistics [LD89], with efficient dedicated algorithms ([KKSK11, CM20] among others). However, these algorithms are non-standard and instead, we are concerned with standard iterative algorithms like training neural networks with gradient descent. Recently, [DH18] considered an iterative optimization procedure for learning such models with a polynomial sample complexity that is controlled by the smoothness of the link function.  $\mathrm{[BES^{+}22]}$  considered the effect of taking a single gradient step on the ability of a two-layer neural network to learn a single index model, and [BBSS22, MHPG  $^+$  23] considered training a special two-layer neural network architecture where all neurons share the same weight with gradient flow or online SGD. However, these works only consider the isotropic Gaussian input, and the effect of anisotropy in the covariance matrix when training a neural network to learn a single index model has remained unclear.
Training a Single Neuron with Gradient Descent. When training the first layer, we consider a setting where there is only one effective neuron. A large body of works exists on training a single neuron using variants of gradient descent. In the realizable setting (i.e. identical link and activation), the typical assumptions on the activation correspond to information exponent 1 as the activations are required to be monotone or have similar properties, see e.g. [Sol17, YO20, DKTZ22]. In the agnostic setting, [FCG20] considered initializing from the origin which is a saddle point for information exponent larger than 1. [ATV22] also considered the agnostic learning of a ReLU activation, albeit their sample complexity is not explicit other than being polynomial in dimension.
2 Preliminaries: Statistical Model and Training Procedure
For a  $d$ -dimensional input  $\pmb{x}$  and a link function  $g \in L^2(\gamma)$ , consider the single index model
$$
y = g \left(\frac {\langle \boldsymbol {u} , \boldsymbol {x} \rangle}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}\right) + \epsilon \text {w i t h} \boldsymbol {x} \sim \mathcal {N} (0, \boldsymbol {\Sigma}), \tag {2.1}
$$
where  $\epsilon$  is a zero-mean noise with  $\mathcal{O}(1)$  sub-Gaussian norm and  $\pmb{u} \in \mathbb{S}^{d-1}$ . Learning the model (2.1) corresponds to approximately recovering the unknown link  $g$  and the unknown direction  $\pmb{u}$ . Note that a normalization is needed to make this problem well-defined; without loss of generality, we write  $\langle \pmb{u}, \pmb{x} \rangle / \| \pmb{\Sigma}^{1/2} \pmb{u} \|$  to ensure that the input variance and the scaling of  $g$  both remain independent of the conditioning of  $\pmb{\Sigma}$ . For this learning task, we will use a two-layer neural network of the form
$$
\hat {y} (\boldsymbol {x}; \boldsymbol {W}, \boldsymbol {a}, \boldsymbol {b}) := \sum_ {i = 1} ^ {m} a _ {i} \phi \left(\left\langle \boldsymbol {w} _ {i}, \boldsymbol {x} \right\rangle + b _ {i}\right), \tag {2.2}
$$
where  $\pmb{W} = \{\pmb{w}_i\}_{i=1}^m$  is the  $m \times d$  matrix whose rows correspond to first-layer weights  $\pmb{w}_i$ ,  $\pmb{a} = \{a_i\}_{i=1}^m$  denote the second-layer weights,  $\pmb{b} = \{b_i\}_{i=1}^m$  denote the biases, and  $\phi$  is the non-linear
activation function. We assume  $g$  and  $\phi$  are weakly differentiable with weak derivatives  $g'$  and  $\phi'$  respectively, and  $g, g', \phi, \phi' \in L^2(\gamma)$ . We are interested in the high-dimensional regime; thus,  $d$  is assumed to be sufficiently large throughout the paper. Our ultimate goal is to learn both unknowns  $g$  and  $u$  by minimizing the population risk
$$
R (\boldsymbol {W}, \boldsymbol {a}, \boldsymbol {b}) := \frac {1}{2} \mathbb {E} \left[ \left(\hat {y} (\boldsymbol {x}; \boldsymbol {W}, \boldsymbol {a}, \boldsymbol {b}) - y\right) ^ {2} \right], \tag {2.3}
$$
using a gradient-based training method such as gradient flow.
We follow the two-step training procedure employed in recent works [BES+22, MHPG+23, BBSS22, DLS22]: First, we train the first-layer weights  $\pmb{W}$  to learn the unknown direction  $\pmb{u}$ ; at the end of this stage, the neurons  $\pmb{w}_i$  align with  $\pmb{u}$ . Here, the goal is to recover only the direction. Next, using random biases and training the second-layer weights, we obtain a good approximation for the unknown link function  $g$ . In the majority of this work, we focus on the first part of this two-stage procedure as the alignment between  $\pmb{w}_i$ 's and  $\pmb{u}$  essentially determines the sample complexity of the overall procedure. This problem is somewhat equivalent to the simplified problem of minimizing (2.3) with  $m = 1$ ,  $a_1 = 1$ ,  $b_1 = 0$ , i.e.,  $\hat{y}(\pmb{x}; \pmb{W}, \pmb{a}, \pmb{b})$  is replaced with  $\hat{y}(\pmb{x}; \pmb{w}) \coloneqq \phi(\langle \pmb{w}, \pmb{x} \rangle)$  and we write  $R(\pmb{w}) \coloneqq R(\pmb{W}, \pmb{a}, \pmb{b})$  for simplicity. We emphasize that unless  $\phi = g$  (i.e. the link function is known), the first stage of training only recovers the relevant direction  $\pmb{u}$  and is not able to approximate  $g$ . Indeed,  $m > 1$  is often needed to learn the non-linear link function; this is the focus of Section 4.2 where we derive a complete learnability result for a two-layer neural network with  $m > 1$ .
Characteristics of the link function play an important role in the complexity of learning the model. As such, a central part of our analysis will rely on a particular property based on the Hermite expansion of functions in a basis defined by the normalized Hermite polynomials  $\{h_j\}_{j\geq 0}$  given as
$$
h _ {j} (z) = \frac {(- 1) ^ {j} e ^ {z ^ {2} / 2}}{\sqrt {j !}} \frac {\mathrm {d} ^ {j}}{\mathrm {d} z ^ {j}} e ^ {- z ^ {2} / 2}. \tag {2.4}
$$
These polynomials form an orthonormal basis in the space  $L^2(\gamma)$ , and the resulting expansion yields the following measure of complexity for  $g$ , which is termed as the information exponent.
Definition 1 (Information exponent). Let  $g = \sum_{j\geq 0}\alpha_jh_j$  be the Hermite expansion of  $g$ . The information exponent of  $g$  is defined to be  $s := \inf \{j > 0 : \alpha_j \neq 0\}$ .
This concept was introduced in [BAGJ21] in a more general framework, and our definition is more in line with the setting in [BBSS22]. We remark that the definition of [BAGJ21] can be modified to handle anisotropy in which case one arrives at Definition 1. We provide a detailed discussion on this concept together with some properties of the Hermite expansion in Appendix B. Throughout the paper, we assume that the information exponent does not grow with dimension.
In the case where the  $d$ -dimensional input data is isotropic, [BBSS22] showed that learning a single index target with full-batch gradient flow requires a sample complexity of  $\tilde{\mathcal{O}}(d^s)$  for  $s \geq 3$  where  $s$  is the information exponent of  $g$ . We will show that this sample complexity can be improved under anisotropy. More specifically, if the input covariance  $\pmb{\Sigma}$  has non-trivial alignment with the unknown direction  $\pmb{u}$ , we prove in Section 3 that the resulting sample complexity can be even made independent of the information exponent if we use a certain normalization in the training. In what follows, we prove that such a normalization in training procedure is indeed necessary.
2.1 The spiked model and limitations of covariance-agnostic training
In practice, data often exhibit a certain structure which may have a profound impact on the statistical procedure. A well-known model that captures such a structure is the spiked model [Joh01] for which one or several large eigenvalues of the input covariance matrix  $\pmb{\Sigma}$  are separated from the bulk of the spectrum (see also [BBAP05, BS06]). Although our results hold for generic covariance matrices, they reveal interesting phenomena under the following spiked model assumption.
Assumption 1. The covariance  $\boldsymbol{\Sigma}$  follows the  $(\kappa, \theta)$ -spiked model if  $\boldsymbol{\Sigma} = \frac{\mathbf{I}_d + \kappa \boldsymbol{\theta} \boldsymbol{\theta}^\top}{1 + \kappa}$  where  $\|\boldsymbol{\theta}\| = 1$ .
In pursuit of the target (unit) direction  $\mathbf{u}$ , the magnitude of the neuron  $\mathbf{w}$  is immaterial; thus, recent works take advantage of this and simplify the optimization trajectory by projecting  $\mathbf{w}$  onto unit sphere  $\mathbb{S}^{d-1}$  throughout the training process [BAGJ21, BBSS22]. In the sequel, we study the same dynamics
which is agnostic to the input covariance in order to motivate our investigation of normalized gradient flow in Section 3. More specifically, we consider the spherical population gradient flow
$$
\frac {\mathrm {d} \boldsymbol {w} ^ {t}}{\mathrm {d} t} = - \nabla^ {S} R (\boldsymbol {w} ^ {t}) \text {w h e r e} \nabla^ {S} R (\boldsymbol {w}) = \nabla R (\boldsymbol {w}) - \langle \nabla R (\boldsymbol {w}), \boldsymbol {w} \rangle \boldsymbol {w}. \tag {2.5}
$$
where  $\nabla^S$  is the spherical gradient at the current iterate. It is straightforward to see that when the initialization  $\boldsymbol{w}^0$  is on the unit sphere, the entire flow will remain on the unit sphere, i.e.  $\boldsymbol{w}^t \in \mathbb{S}^{d-1}$  for all  $t \geq 0$ . The flow (2.5) has been proven useful for learning the direction  $\boldsymbol{u}$  [BBSS22] in the isotropic case  $\boldsymbol{\Sigma} = \mathbf{I}_d$  when the activation  $\phi$  is ReLU. In contrast, when  $\boldsymbol{\Sigma}$  follows a spiked model, we show that it can get stuck at stationary points that are almost orthogonal to  $\boldsymbol{u}$ . Indeed, when the input covariance  $\boldsymbol{\Sigma}$  has a spike in the target direction  $\boldsymbol{u}$ , i.e.  $\theta = \boldsymbol{u}$ , one expects that the training procedure benefits from this as the input  $\boldsymbol{x}$  contains information about the sought unknown  $\boldsymbol{u}$  without even querying the response  $y$ . The following result proves the contrary; for moderate spike magnitudes, the alignment between the first-layer weights and target  $\langle \boldsymbol{w}^t, \boldsymbol{u} \rangle$  will be insignificant for all  $t$ .
Theorem 2. Let  $s > 2$  be the information exponent of  $g$  with  $\mathbb{E}[g] = 0$ , and assume  $\Sigma$  follows the  $(\kappa, \mathbf{u})$ -spiked model with  $\Omega(1) \leq \kappa \leq \mathcal{O}(d^{\frac{s - 2}{s - 1}})$ . For ReLU activation, let  $\mathbf{w}^t$  denote the solution to (2.5) initialized uniformly at random over  $\mathbb{S}^{d - 1}$ , then with probability at least 0.99,
$$
\sup  _ {t \geq 0} | \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle | \lesssim 1 / \sqrt {d}, \tag {2.6}
$$
A non-trivial alignment between the first-layer weights  $\boldsymbol{w}^t$  and the target direction  $\boldsymbol{u}$  is required to learn the single index model (2.1). However, the above result implies that in high dimensions when  $d \gg 1$ , the alignment is negligible in the population limit (when the number of samples goes to infinity). We remark that when the spike magnitude is large, i.e.  $\kappa \geq \Omega(d)$ , the flow (2.5) can achieve alignment as the problem essentially becomes one-dimensional, as we demonstrate in Appendix C.
To see why the flow (2.5) gets stuck at saddle points and fails to recover the true direction, notice that
$$
R (\boldsymbol {w}) = \frac {1}{2} \mathbb {E} \left[ \left(\phi \left(\langle \boldsymbol {w}, \boldsymbol {x} \rangle\right) - y\right) ^ {2} \right] = \frac {1}{2} \mathbb {E} \left[ \phi \left(\langle \boldsymbol {w}, \boldsymbol {x} \rangle\right) ^ {2} \right] - \mathbb {E} \left[ \phi \left(\langle \boldsymbol {w}, \boldsymbol {x} \rangle\right) y \right] + \frac {1}{2} \mathbb {E} \left[ y ^ {2} \right]. \tag {2.7}
$$
If the input was isotropic, i.e.  $\pmb{x} \sim \mathcal{N}(0, \mathbf{I}_d)$ , the first term in (2.7) would be equal to  $\| \phi \|_{\gamma}^2$ , which is independent of  $\pmb{w}$ . Thus, minimizing  $R(\pmb{w})$  in this case is equivalent to maximizing the "correlation" term  $\mathbb{E}[\phi (\langle \pmb{w}, \pmb{x} \rangle)y]$ . However, under the spiked model, the alignment between  $\pmb{w}$  and  $\pmb{u}$  breaks the symmetry; consequently, the first term in the decomposition grows with  $\langle \pmb{w}, \pmb{u} \rangle$ , creating a repulsive force that traps the dynamics around the equator where  $\pmb{w}$  is almost orthogonal to  $\pmb{u}$ .
3 Main Results: Alignment via Normalized Dynamics
Having established that the covariance-agnostic training dynamics (2.5) is likely to fail, we consider a covariance-aware normalized flow in this section and show that it can achieve alignment with the unknown target and enjoy better sample complexity compared to the existing results [BAGJ21, BBSS22] in the isotropic case. We start with the population dynamics.
3.1 Warm-up: Population dynamics
To simplify the exposition, we define  $\pmb{z} \coloneqq \pmb{\Sigma}^{-1/2}\pmb{x}$ ,  $\overline{\pmb{w}} \coloneqq \pmb{\Sigma}^{1/2}\pmb{w} / \|\pmb{\Sigma}^{1/2}\pmb{w}\|$  and similarly define  $\overline{\pmb{u}}$ , and consider the prediction function  $\hat{y}(\pmb{x};\overline{\pmb{w}}) \coloneqq \phi(\langle\overline{\pmb{w}},\pmb{z}\rangle)$ . Due to symmetry, the second moment of the prediction is  $\mathbb{E}\left[\hat{y}(\pmb{x};\overline{\pmb{w}})^2\right] = \|\phi\|_{\gamma}^2$  which is independent of  $\pmb{w}$ ; thus, the population risk reads