Learning with Shallow Neural Networks on Cluster-Structured Features Elisabetta Cornacchia and Laurent Massoulié
arXiv:2605.14927v1 [cs.LG] 14 May 2026
INRIA, DI/ENS, PSL.
Abstract The success of deep learning in high-dimensional settings is often attributed to the presence of low-dimensional structure in real-world data. While standard theoretical models typically assume that this structure lies in the target function—projecting unstructured inputs onto a low-dimensional subspace—data such as images, text or genomic sequences exhibit strong spatial correlations within the input space itself. In this paper, we propose a tractable model to study how these correlations affect the sample complexity of learning with gradient descent on shallow neural networks. Specifically, we consider targets that depend on a small number of latent Boolean variables, and input features grouped into clusters and correlated with the latent variables. Under an identifiability assumption, we show that for a layerwise gradient-descent variant, the sample complexity scales with the number of hidden variables and, when the signal-to-noise ratio is sufficiently high, is independent of the input dimension, up to logarithmic terms. We empirically test our theoretical findings on both synthetic and real data.
1
Introduction
A common theme in deep learning is that high-dimensional problems encountered in practice often have an underlying low-dimensional structure, a property widely regarded as central to the success of modern algorithms. Theoretical analyses typically incorporate this assumption by modeling the labels as functions of a low-dimensional projection of the input (e.g., multi-index target functions), while treating the inputs themselves as unstructured—often drawn from an isotropic Gaussian or uniform Boolean distribution [BAGJ21, BBSS22, AAM23]. This simplification makes the analysis tractable, since the training dynamics can be described in terms of a small number of low-dimensional order parameters (or sufficient statistics) whose evolution can be tracked analytically. Within this framework, several works have shown that the sample complexity to learn with shallow neural networks scales as O(dI ), where d is the input dimension and I is an exponent depending on the target function—capturing, for example, the information [BAGJ21, BBSS22], generative [DPVLB24, DTA+ 24], or leap exponent [AAM23, ABAM22]. Real-world data, by contrast, rarely resemble such unstructured inputs. Images, text, and biological data exhibit strong correlations among input features, so the relevant low-dimensional structure is often encoded in the input distribution itself, not only in the target function. Empirical estimates illustrate this phenomenon: for instance, the intrinsic dimension of MNIST, despite its 784-dimensional pixel representation, is estimated to be on the order of 15, while that of CIFAR-10 is estimated around 50 [Cam03, CH04, LB04, FdRL17, ALMZ19]. Several works have studied how correlations among input coordinates can reduce the sample complexity of learning compared with unstructured product distributions [GMKZ20, MHWSE23]. However, these analyses typically remain in regimes where Email: [email protected], [email protected].
1
the latent dimension diverges, and the resulting complexity still depends critically on the ambient dimension d. In this work we focus on a different regime: the number of latent degrees of freedom remains bounded while the number of observed features grows. This is a natural regime for highly redundant data, where many observed coordinates provide noisy views of a small number of underlying factors. For example, in genomics, large collections of genes may be co-regulated by a much smaller number of latent biological programs, producing high-dimensional observations with strong feature redundancy (see experiments in Section 6). From a statistical perspective, this is precisely the regime in which one might hope for sample complexity governed by the intrinsic dimension rather than the ambient dimension. From an analytical perspective, however, it is not covered by existing high-dimensional asymptotic theories: since the latent dimension is fixed, there is no large-latent-dimensional or thermodynamic limit from which to derive closed dynamical equations or self-averaging order parameters. Motivated by this setting, we introduce a tractable model of clustered, correlated features. Labels are determined by a bounded set of hidden binary latent variables, while the high-dimensional observable variables are generated from them through a sub-Gaussian noisy process, subject to an essential identifiability condition. We prove that a simplified version of gradient descent—with layer-wise training–on a two-layer fully connected network can leverage these correlations to learn any target function of the latent variables. Importantly, in high signal-to-noise (SNR) regimes, the sample complexity becomes independent of the input dimension d, up to logarithmic terms. This perspective is complementary to works showing depth separations for hierarchical data models [Mos16, CPT+ 24, DPZK25, RDKL26, TDPK26]. Rather than asking whether deep networks outperform shallow ones on complex hierarchical structures, we ask a more basic question: when the relevant correlations are already accessible to a shallow architecture, can gradient-based training exploit them without being explicitly given the latent representation? Our results give a positive answer in a bounded-latent-dimensional regime, showing that the relevant complexity can be governed by the intrinsic rather than ambient dimension.
2
Related Work
Gradient Descent on Product Measures. While task learnability with neural networks has been widely studied, most analyses remain within the classical framework where inputs are sampled from an unstructured distribution—typically uniform Boolean or isotropic Gaussian. In this setting, precise characterizations of the sample and time complexities needed to learn single and multi-index models with gradient descent on shallow networks have been obtained in terms of the information [BAGJ21, BBSS22], generative [DPVLB24, DTA+ 24, TDD+ 24, JKMS25], and leap [AAM23, ABAM22, Gla23, KCGK24, BEG+ 22, JMS24] exponents of the target function. A few works have shown that shifting the first moment of an isotropic Gaussian (or Boolean) distribution can break symmetries and thereby reduce learning complexity [MKAS21, CM23, DM20, JMS24, CMM25]. [DPZK25, TDPK26] further study the computational benefits of deep versus shallow architectures for learning hierarchical targets, with Gaussian inputs of identity covariance. However, all these works consider product measures and do not address correlations among input features. Gradient Descent on Structured Inputs. Beyond product measures, [MHWSE23, BES+ 23] studied single-index models on Gaussian inputs with spike-covariance structure, showing that when the spike aligns with the target, sample complexity is drastically reduced compared to the unstructured case. [GMKZ20] proposed the Hidden Manifold Model (HMM), a generative model in which highdimensional inputs lie on a low-dimensional manifold and labels depend on their position within it. [BG24, SBGG24] study detection problems where the signal is encoded into high-order moments of the input, and show separation between neural networks and random feature models. These analyses are in the limit of large input and manifold dimensions, and the complexity still scales with the ambient dimension, as opposed to our analysis which focuses on small latent dimension (Od (1)). Moreover, 2
Figure 1: Our clustered feature model, with latent variables s ∈ {±1}N , noise variables ξ ∈ Rd , and observable variables x ∈ Rd . [ABRC24] empirically study how neural networks learn tasks with hidden factorial structure. [ABL+ 24] establish a lower bound for weak learning tasks with low globality-degree, where signal resides in global rather than local correlations. An active line of work studies models of tree-like correlation structures among features [Mos16, PMR+ 17]. In this direction, [CPT+ 24] introduce the Random Hierarchy Model (RHM), where a root label recursively generates hidden variables on a tree and the observed input is given by the leaves. They show that deep architectures aligned with the hierarchy can exploit this compositional structure, while shallow networks suffer from the curse of dimensionality. More recently, [RDKL26] proved that layerwise gradient descent on an L-layer convolutional network can efficiently learn RHMs, with sample complexity essentially mL , where m is the number of production rules per token and L is the hierarchy depth. Subsequent works have explored related hierarchical frameworks in next-token prediction [CW24], Transformer architectures [GBMMS24], and diffusion models [FSC+ 25, SFW25, SFLW24]. Our work is complementary: we do not study depth separation or architectures matched to a growing hierarchy. Instead, we ask whether shallow networks can exploit correlations among observed features in a bounded-latent-dimensional regime. A related line of work studies learning in structured Gaussian models in the thermodynamic regime, where the number of samples scales linearly with the input dimension. In particular, [MKL+ 20, MKUZ20] analyze the learning of mixtures of two Gaussians using a single-layer perceptron, while [RGKZ21] study the Gaussian XOR problem with a two-layer architecture. Furthermore, [LGC+ 21] extend the teacher–student framework to more realistic data distributions by introducing Gaussian covariate generalization. Our work departs from these approaches in two key respects: (i) we consider a broader class of structured models, beyond purely Gaussian inputs; and (ii) our results extend beyond the thermodynamic regime, with sample complexity scaling at most logarithmically with the input dimension. Identifiability of Latent Structure. Beyond gradient descent learning, some works studied parameters learning in latent variable models, considering settings similar to ours [AMR09, AGH+ 14, BBAP05, BGN11, AGM12, HKM24]. In particular, [AHJK15] study latent variable models with topic persistence, where a single topic is assumed to generate a sequence of n consecutive observations—an assumption that resembles our feature model. They provide sufficient conditions for identifiability based on observations of high-order moments of the observations. Instead, we focus on supervised learning tasks involving latent representations, using gradient descent on neural networks, rather than just identifiability of the hidden structure.
3
Summary of Results
We study a clustered feature model, in which the d observable coordinates are partitioned into N disjoint clusters C1 , . . . , CN . Each cluster Ci is associated with a hidden binary latent (or topic) variable 3
si ∈ {±1}, so that all features in the same cluster share a common signal component. Concretely, an input x ∈ Rd is generated by first sampling the latent variables s = (s1 , . . . , sN ) ∈ {±1}N , then setting each feature xj with j ∈ Ci to xj = si ξj , where ξj ∈ R is a sub-Gaussian noise variable, independent of s. Each ξj has mean mj := E[ξj ] (interpreted as the signal strength), and is independent across samples. The label is given by an arbitrary function f of the latent variables: y = f (s). See Figure 1 for an illustration of our data model and Def 1 for a formal definition. Throughout, N is treated as fixed while d grows. This model captures high-dimensional data in which many observed coordinates are noisy, redundant measurements of a much smaller number of latent factors. A natural example is genomics, where groups of genes may be co-regulated by a small number of latent biological programs, producing clustered covariance patterns and high feature redundancy. We refer to Section 6 for examples of real datasets with structure consistent with this model. Our main structural assumption ensures that the latent variables are statistically recoverable. For each cluster, define N X X vi := m2j , vsum := vi = ∥m∥22 . i=1
j∈Ci
Here vi is the total signal power carried by cluster Ci , while vsum is the total signal power across all observed coordinates. We assume that no cluster is vanishingly weak: there exists a constant c > 0 such that vi ≥ c, ∀i ∈ [N ]. vsum This condition ensures that each latent variable leaves a detectable footprint in the observed features. When the noise variables have bounded variance, vsum naturally plays the role of an aggregate signal-to-noise ratio. Our main theorem focuses on the high-SNR regime vsum = Θd (d), where the total signal grows linearly with the ambient dimension. Let us give few examples of simple models captured by this framework. Example 1 (Mixture of Gaussians). If the noise variables are Gaussian, our model recovers standard Gaussian mixture examples. For instance, let N = 1, with s1 ∼ Unif{±1}, y = s1 , and ξ ∼ N (m, σ 2 Id ) for some m ∈ Rd and σ > 0. Then x = s1 ξ is distributed as a symmetric mixture of two Gaussians, x∼
1 1 N (m, σ 2 Id ) + N (−m, σ 2 Id ), 2 2
as studied in [MKUZ20, MKL+ 20]. Similarly, when N = 2, with label y = s1 s2 , the model gives a Gaussian XOR problem [RGKZ21]. Previous works typically study the low-SNR regime ∥m∥2 = Θd (1), where the sample complexity remains linear in the ambient dimension. By contrast, our analysis focuses on the high-SNR regime vsum = ∥m∥22 = Θd (d), σ 2 = Θd (1). Example 2 (Binary Symmetric Cluster (BSC) Model). A particularly simple instance of our framework is obtained when the observed variables are Boolean noisy copies of the latent variables. For each cluster Ci and each coordinate j ∈ Ci , let ξj ∼ Rad(1 − δ), δ ∈ [0, 1],1 independently across coordinates, so that xj = si ξj , j ∈ Ci . Equivalently, each observable coordinate is obtained by passing the latent bit si through a binary symmetric channel (BSC) with flip probability δ. In this case, mj = E[ξj ] = 1 − 2δ, and therefore vi = |Ci |(1 − 2δ)2 , vsum = d(1 − 2δ)2 . When the clusters have comparable sizes, the identifiability condition is satisfied. When the clusters have the same size, we refer to this setting as the homogeneous BSC model. 1 Rademacher distribution: z ∼ Rad(p) means P(z = 1) = 1 − P(z = −1) = p.
4
PnOur goal is⊤to characterize the number of samples required for a two-layer neural network NN(x; θ) := h=1hah σ(wh x + bh ),i trained by stochastic gradient descent (SGD), to achieve small test error: 2 E(x,y) y − NN(x; θ̂) . We will consider the latent dimension N and the target function f to be fixed, and we study the scaling of the complexity as d grows. Our main theorem reads as follows (see Theorem 1 for a formal statement). Informal Theorem 1. For any ϵ, δ > 0, under the clustered feature model and assuming identifiability and high-SNR conditions, layerwise-SGD on a two-layer network with O(log(1/ϵ)) hidden neurons, 2 d 2 requires at most O( log(d) vsum δ 2 ) samples and O(1/δ ) steps to achieve error at most δ with probability at least 1 − poly(ϵ). In particular, since the overall signal satisfies vsum = θ(d), the sample complexity depends on the ambient dimension d only logarithmically. This result holds for Gaussian initialization and some polynomial activation functions (Assumption 2); the formal statement is given in Theorem 1. We then specialize to the homogeneous BSC model and obtain a result that holds in the wider SNR regime vsum = Ωd (log(d)4 ), under a mild non-degeneracy condition on the target function f (Theorem 2).
4
Setting and Formal Result
Data Model. Let us give a formal definition of our data model. As we mentioned earlier, we assume N and f to be fixed, and we study the scaling of the sample complexity as d grows. Definition 1 (Clustered-Structured Features). Let (x, y) ∈ Rd × R and f : {±1}N → R. Assume that: • s ∼ Unif{±1}N ; • There exists a partition C = (C1 , ..., CN ) of [d] into N disjoint clusters: [d] = ∪i∈[N ] Ci , such that for all i ∈ [N ] and j ∈ Ci , xj = si ξj , where (ξj )j∈[d] are mutually independent, independent of s, sub-Gaussian random variables with ∥ξj ∥ψ2 ≤ ψ = Od (1), finite means mj := E[ξj ] = Od (1), variances τj2 := Var(ξj ) such that Pd 1 2 3 j=1 τj = θd (1), and bounded third moments supj E|ξj − mj | = Od (1); d • y = f (s). We say that (x, y) ∼ DN,f,C,ξ . We call the s ∈ {±1}N the latent (or topic) variables, the x ∈ Rd the observable variables, and the ξ ∈ Rd the noise variables. Crucially, we assume that the size of the topic space N is constant, as the observable dimension d grows, i.e.: N = Od (1). Notice that, under this model, the covariance structure of x is block diagonal with respect to the partition C, with within-cluster correlations induced by the shared latent variable si and no cross-cluster correlations. Throughout the paper, we assume that the following identifiability assumption is satisfied. P 2 Assumption 1 ((A1), Identifiability). For each cluster Ci , i ∈ [N ], define: vi := j∈Ci mj , P vmin := mini∈[N ] vi and vsum := i∈[N ] vi = ∥m∥22 , where mj := E[ξj ], for j ∈ [d]. We assume that there exists a c > 0 such that a)
vmin ≥ c, vsum
b)
5
vsum ≥ c. d
Assumption 1a) states identifiability of individual clusters. It prevents degenerate cases in which one cluster contributes vanishingly small signal compared to the others, making the corresponding latent variable statistically unrecoverable. On the other hand, Assumption 1b) is a lower bound on the signal-to-noise (SNR) ratio, which is required for our main Theorem (Thm. 1) to hold. In fact, under Pd our data model assumption j=1 τj2 = θd (d), Assumption 1b) is equivalent to ∥m∥2 /∥τ ∥2 ≥ c′ > 0. In Theorem 2 we show that for the specific case of homogeneous BSC models, Assumption 1b) can be released to vsum = ωd (log(d)2 ). Architecture. We adopt standard choices for both architecture and optimizer, since our focus is on the impact of the data structure. Concretely, we use a 2-layer fully connected network: NN(x; θ) =
n X
ah σ(whT x + bh ),
(1)
h=1
where θ = (w, a, b) ∈ Rd×n × Rn × Rn are the network’s trainable parameters. In our main Theorem 1, we assume the activation to be a polynomial, satisfying the following technical assumption. PP ℓ Assumption 2 ((A2), Activation). Let σ(x) = ℓ=0 cℓ x be a polynomial activation of degree P d 1 N 2 P = Θd (1), with P ≥ 2 . Let µ := d j=1 τj , and vsum := ∥m∥22 , and for v > 0 define Sv (t) := EG∼N (0,v) [σ ′ (t + G)]. We assume that there exists c0 > 0 such that, for all 0 ≤ k ≤ N , (k)
Sµ+vsum /d (0) ≥ c0 . Remark 1. We note that in Assumption 2 the Hermite non-degeneracy condition is tied to the data model through the scalars µ (the reference smoothing scale) and vsum (from Assumption 1). This dependence is for convenience of the analysis: it allows us to match the Gaussian smoothing scale to the input signal-to-noise and is precisely what we need to certify linear separability after the first gradient step on the first layer (see the Algorithm description below). We remark that this condition is satisfied by most polynomial activations of degree at least 2N (see Lemma 14 in Appendix for details). Algorithm. As training algorithm, we take layerwise stochastic gradient descent (SGD): first we take one gradient step on the first layer (holding the second layer fixed), then we train the second layer until convergence (keeping the first layer frozen). We take the standard square loss. We take a Gaussian initialization for the first layer’s weights w, with normalized variance, while the second layer’s weights are initialized to a small enough value, so that the interaction term in the square loss is negligible with respect to the correlation term (similarly to e.g. [AAM23, MHWSE23]). The bias weights are set to zero for the first phase, and then sampled uniformly at random in a given interval, and kept fixed during training, for the second phase. This stylized layerwise simplification of gradient descent matches several recent theoretical studies in deep learning ([DLS22, DKL+ 23, BES+ 22, BEG+ 22, CM23]). We consider the online setting, where fresh batches of samples are generated at each training step. Further details on the algorithm can be found in Appendix A (Algorithm 1). In our experiments in Section 6 we use standard SGD, with both layers trained jointly. Main Result.
Let us state our main theorem.
Theorem 1. Let f : {±1}N → R be a target function such that Var(f ) = 1 and let (x, y)P ∈ Rd × R be drawn from DN,f,C,ξ , according to Def. 1. Assume A1 is satisfied with vmin := mini∈[N ] j∈Ci m2j and vsum := ∥m∥22 . Then, there exist c1 , c2 > 0 such that for any δ, ϵ ∈ (0, 1/4), a two-layer network with activation satisfying A2, n = θ(2N log(1/ϵ)) hidden neurons and Gaussian initialization with 6
variance 1/d for the first layer’s weights, trained by layerwise-SGD with the squared loss, batch size B = θ(N log(d log(1/ϵ))2 d/vsum ) for T1 = 1 steps on the first layer, with learning rate γ1 = √ −2P θ(1/ vsum ) and T2 = θ(22N /δ 2 ) steps, with γ2 = θ(vmin δ/2N ) on the second layer, with probability c2 1 − c1 ϵ , over the initialization and the training samples, will output a network NN(x; θT1 +T2 ) such that Ex,y y − NN(x; θT1 +T2 ) ≤ δ. Since Assumption 1 implies vsum ≳ d and N = Od (1), for fixed δ, ϵ > 0 Theorem 1 gives T1 + T2 = Od (1) and B = Od (log2 d). Thus the sample complexity B(T1 + T2 ) scales polylogarithmically in d. On the other hand, the time complexity (measured as dim(θ) · (T1 + T2 ) · B, where dim(θ) denotes the number of trainable parameters in the network) depends linearly on d, up to logarithmic factors. Remark 2. We assume a finite number of clusters N because a two-layer network imposes a representational bottleneck. In our construction, the first layer learns the topic variables s ∈ {±1}N , and the second layer fits f (s) (see Sec. 4.1). Since the space of all functions on {±1}N has dimension 2N , our construction uses width of order 2N in the second layer. Thus, we keep N fixed. Although N could scale as log d, we do not analyze that regime here. It is natural to ask whether adding depth, for instance using three layers, could reduce the exponential dependence on N to a polynomial one, allowing N to grow with the input dimension; we leave this direction to future work. Finally, the assumption s ∼ Unif{±1}N is made for simplicity of exposition, and we expect that the analysis can be extended to more general zero-mean distributions. Remark 3. Crucially, we exploit the random Gaussian initialization of the first layer to induce diversity among hidden units. The required non-degeneracy follows from the Carbery–Wright anti-concentration inequality for polynomials under log-concave measures; see Lemma 3 and [CW01]. This motivates our use of polynomial activations and Assumption 2. Gaussian initialization also allows us to control the second moments of the resulting random polynomial features via Hermite expansions. We expect that similar arguments extend to other log-concave initializations. Finally, analogous extensions to common non-polynomial activations, such as ReLU, may be possible by truncating their Hermite expansions, at the cost of additional approximation and tail-control arguments. In the next section we give a proof outline for Theorem 1. We refer to Appendix A for the complete proof.
4.1
Proof Outline for Theorem 1
The training proceeds in two phases. With a small second-layer scale at initialization, a single layerwiseSGD step on the first layer moves each weight toward a random linear combination of the cluster mean directions, so each hidden unit effectively computes a one-dimensional projection of the input onto a latent direction u(s) = s⊤ α̃(w0 ), where α̃(w0 ) are random variables that depend on the first layer’s initialization. In particular, using a Berry-Esseen type bound (Lemma 1), we show that: vi
i ∈ [N ], (2) vsum P P where vi := j∈Ci m2j , vsum := ∥m∥22 = i∈[N ] vi , Gw0 ∼ N (0, vsum /d), and Ps are polynomials in G of degree deg(σ) − 1 determined by the activation, that depend on s ∈ {±1}N . To conclude that the first phase has produced useful features, we need to show that the target function can be expressed using the hidden units obtained after this update. Since these hidden units depend on the latent variables only through the scalar projection u(s), this reduces to showing that the projection retains all information relevant to the label. In other words, we need f to be projection-consistent: for all s, t ∈ {±1}N , f (s) ̸= f (t) implies u(s) ̸= u(t). Crucially, we exploit the random Gaussian initialization, which induces diversity across hidden units. Using the Carbery–Wright anti-concentration inequality (Lemma 3), we show that the projected values {u(s) : s ∈ {±1}N } are well separated with high probability over the initialization. Assumption 1 ensures that all latent (α̃(w0 ))i =
Es f (s) si Ps (Gw0 ) ,
7
coordinates i ∈ [N ] contribute non-negligible signal, while the condition vsum /d = Θd (1) keeps the Gaussian argument Gw0 at a nondegenerate variance scale, which is needed for the anti-concentration argument. Then, by sampling sufficiently many hidden-unit biases uniformly over a wide interval, an intervalhitting argument ensures that at least one bias lies near each distinct projection value. The resulting hidden responses form a well-conditioned system. Solving this system yields a certificate, i.e., a choice of second-layer weights that exactly interpolates f (s) on the projection grid (Lemma 5). In the second phase, we freeze the first layer and fit the output weights by SGD on the square loss; standard convergence results for convex objectives then drive the error to the target level (Lemma 6).
5
Special Case: Homogeneous Clusters
Let us consider for simplicity the homogeneous BSC model, where clusters have equal size (see iid
Example 2). In particular, let ξj ∼ Rad(1 − δ), for some δ ∈ (0, 1), for all j ∈ [d], and |Ci | = k = d/N , for all i ∈ [N ], where for simplicity we assume that d is a multiple of N . Note that in this case vsum = d(1 − 2δ)2 and vmin /vsum = 1/N . If |δ − 12 | = θd (1), and therefore vsum = θd (d), then Theorem 1 applies, and for fixed ϵ > 0, the number of samples needed to achieve ϵ error scales only logarithmically in d. Specializing to the homogeneous BSC model, allows us to obtain a sharper guarantee, valid as soon as vsum = Ωd (log(d)4 ), rather than the high-SNR condition vsum =√θ(d) required by our general result. To this end, we consider a deterministic initialization of wj0 = 1/ d, for all j ∈ [d], and we focus on the simpler ReLU activation function. Because of the deterministic initialization, we cannot exploit the randomness of the initialization as in eq. (2) and Lemma 3 to show that the target function is projection consistent. Thus, we impose the following non-degeneracy assumption on the target f . Assumption 3 (Majority Margin). For T ⊆ [N ] and i ∈ [N ] let if T = {i}, 1 ˆ Maj (|T | − 1) if {i} ⊂ T, cT,i := ˆ N MajN (|T | + 1) if {i} ̸⊂ T,
(3)
ˆ N (k) are the Fourier-Walsh coefficients of the Majority function in dimension N [O’D14], where Maj see Appendix F. We say that f has ∆-margin, for ∆ > 0, if for all s, t ∈ {±1}N such that f (s) ̸= f (t), N X X (si − ti ) fˆ(T )cT,i ≥ ∆. i=1
(4)
T ⊆[N ]
Our Theorem reads as follows. Theorem 2. Assume data are sampled from the homogeneous √ BSC model with N = Od (1) clusters, noise parameter δ ∈ [0, 1/2] such that |δ − 1/2| = Ω(log(d)2 / d) and target function f with ∆-margin (according to Assumption 3). Then, there exists a constant C > 0 such that for all ϵ > 0, layerwise-SGD log(d)2 N with the squared loss, batch size B = O( (1−2δ) log(d)/∆) 2 ) on a two-layer network with n = Ω(2 2 N −2 −4 −C hidden neurons, after T1 = 1, T2 = O(log(d) 2 ϵ ∆ ) steps with probability 1 − d will output a network NN such that E(x,y) (y − NN(x))2 ≤ ϵ. √ Thus, at the threshold |δ − 1/2| = Ω(log d2 / d), equivalently vsum = Ωd (log4 d), the theorem still guarantees learning with polynomial-in-d sample size. In Appendix D we compare the sample complexity guarantee in Theorem 2 with the performance of classical clustering methods.
8
1e6
6 1e6
BSC Gauss. mixture
1.4
5
Number of samples
Number of samples
1.6
1.2 1.0 0.8 150
200
250
Input dimension
300
350
BSC Gauss. mixture
4 3 2 1 0.05
0.10
0.15
0.20
0.25
0.30
Figure 2: Learning parity functions of N = 3 latent bits, under the homogeneous BSC model (blue) and the Gaussian mixture model (orange), using a two-layer ReLU network with 1000 hidden neurons, trained by SGD (square loss, batch size 64, learning rate 0.001). We plot the number of fresh samples required to achieve test loss below 0.05: (Left) as a function of the input dimension d for δ = 0.2, and (Right) as a function of δ for d = 300.
6
Experiments
Experiments on Synthetic Data. We empirically verify whether our theoretical results hold in practice, for standard SGD training (both layers trained jointly). We consider learning parity targets: QN for latent dimension N = 3, we sample s ∼ Unif{±1}N and set y = i=1 si . Parities under uniform inputs are notoriously hard for neural networks: they are uncorrelated with every strict subset of coordinates, so successful learning requires identifying all relevant hidden variables. This makes them a natural choice for testing whether a method can recover the latent structure in our setting. In all experiments, we train a two-layer ReLU network with 1000 hidden neurons and standard uniform initialization, using SGD (batch size 64) and generating fresh samples at each step. Training stops when the in-distribution test error falls below 0.05, and each experiment is repeated 5 times. We report mean and 95% confidence intervals. We consider two data-generating models. (i) The BSC homogeneous cluster (Example 2), with noise parameter δ ∈ {0.05, 0.1, 0.15, 0.2, 0.25, 0.3}. (ii) A Gaussian mixture model, where the noise iid
variables satisfy ξj ∼ N (m, σ 2 ), for j ∈ [d], with mean m = 1 − 2δ and variance σ 2 = 4δ(1 − δ), using the same values of δ as the tree model. These choices ensure that the two models are directly comparable, as their noise variables have matching first and second moments. In the left panel, we plot the number of fresh samples required for learning as a function of the input dimension d, fixing δ = 0.2. We observe that, for sufficiently large d, the sample complexity remains essentially constant in d, in agreement with our theoretical predictions (up to logarithmic factors). In the right panel, we plot the number of samples required for learning as a function of δ, for a fixed input dimension d = 300. As δ approaches 1/2, learning becomes progressively less efficient, again consistently with the theoretical analysis. Experiments on Real Data. We consider the public single-cell RNA-seq dataset GSE96583, hosted on GEO, using batch-2 cells from control (unstimulated) samples [KST+ 18]. Each observation is a cell, and each feature is the measured expression level of one gene; cells are annotated with their type. A clustered feature structure is expected because the ambient dimension is large yet many genes are redundant: multiple genes act as correlated noisy readouts of a smaller number of latent biological programs, such as cell state. As a supervised learning task we consider multiclass prediction of cell type from gene expression. We exclude rare cell types (Dendritic cells, Megakaryocytes) and merge CD4 T, CD8 T, and NK 9
GSE96583: gene gene correlation
1.00 0.75
1.0
0.50
0.00 0.25 0.50 0.75 1.00
0.9
Test accuracy
0.25
0.8 0.7 0.6 d=50 d=100 d=200
0.5 0.4
0
1000
2000
3000
Training set size n
4000
d=300 d=400 d=500
5000
Figure 3: Cell-type classification on the single-cell RNA-seq dataset GSE96583. (left) t-SNE visualization of the three cell types (B cells, Monocytes, T/NK). (center) Gene-gene Pearson correlation matrix, with genes reordered by hierarchical clustering; diagonal blocks reveal a latent cluster structure. (right) Test accuracy as a function of training size n for varying input dimension d; curves for d ≥ 200 nearly coincide, consistent with dimension-independent sample complexity. cells into a single T/NK class, yielding a three-class problem: B cells, Monocytes, and T/NK. After standard quality-control filtering (minimum number of expressing cells per gene), approximately 11,990 cells and 10,500 genes remain. Figure 3(left) shows the t-SNE visualization of the three classes, which are well separated in gene-expression space. We select the top-d most variable genes by dispersion (variance/mean), using no class-label information, and compute the gene-gene Pearson correlation matrix across all cells. Reordering genes by hierarchical clustering reveals a block structure, shown in b ≈ 0.29, where Σ b denotes the empirical covariance Figure 3(center). We estimate SNR ≈ vbsum /Tr(Σ) matrix of the selected features. We then train a two-layer ReLU network with 128 hidden units to predict the three cell types; implementation details are given in Appendix G. Figure 3(right) shows test accuracy as a function of training size n for varying input dimension d: curves for d ≥ 200 nearly coincide, consistent with a sample complexity that is essentially independent of d.
7
Conclusion
In this paper, we introduce a tractable model of correlated, clustered inputs with a small number of latent binary variables and proved that, under identifiability and technical conditions, layerwise SGD on a two-layer network can learn any latent target with a sample complexity that is independent of the ambient dimension, up to logarithmic terms. Our experiments with standard SGD confirm these results. Looking forward, several natural extensions arise. First, deeper architectures may help overcome the exponential dependence on the number of latent variables and improve computational efficiency. Second, it would be interesting to extend the framework to larger alphabets and richer channels: we expect similar guarantees for finite alphabets, while handling continuous latent variables will likely require new technical tools. Third, one could allow cluster memberships to vary across samples.
Acknowledgement EC was supported by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute). LM was supported by the PR[AI]RIE-PSAI – Paris School of Artificial Intelligence, reference: ANR23-IACL-0008.
10
References [AAM23]
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pages 2552–2623. PMLR, 2023.
[ABAM22]
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782–4887. PMLR, 2022.
[ABL+ 24]
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can transformers reason? the globality barrier and inductive scratchpad. Advances in Neural Information Processing Systems, 37:27850–27895, 2024.
[ABRC24]
Charles Arnal, Clement Berenfeld, Simon Rosenberg, and Vivien Cabannes. Learning with hidden factorial structure. arXiv preprint arXiv:2411.01375, 2024.
[AGH+ 14]
Animashree Anandkumar, Rong Ge, Daniel J Hsu, Sham M Kakade, Matus Telgarsky, et al. Tensor decompositions for learning latent variable models. J. Mach. Learn. Res., 15(1):2773–2832, 2014.
[AGM12]
Sanjeev Arora, Rong Ge, and Ankur Moitra. Learning topic models–going beyond svd. In 2012 IEEE 53rd annual symposium on foundations of computer science, pages 1–10. IEEE, 2012.
[AHJK15]
Animashree Anandkumar, Daniel Hsu, Majid Janzamin, and Sham Kakade. When are overcomplete topic models identifiable? uniqueness of tensor tucker decompositions with structured sparsity. The Journal of Machine Learning Research, 16(1):2643–2694, 2015.
[ALMZ19]
Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[AMR09]
Elizabeth S Allman, Catherine Matias, and John A Rhodes. Identifiability of parameters in latent structure models with many observed variables. 2009.
[BAGJ21]
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1):4788–4838, 2021.
[BBAP05]
Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. 2005.
[BBSS22]
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768–9783, 2022.
[BEG+ 22]
Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08799, 2022.
[BES+ 22]
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932–37946, 2022. 11
[BES+ 23]
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dimensional structure: a spiked random matrix perspective. Advances in Neural Information Processing Systems, 36:17420–17449, 2023.
[BG24]
Lorenzo Bardone and Sebastian Goldt. Sliding down the stairs: how correlated latent variables accelerate learning with neural networks. arXiv preprint arXiv:2404.08602, 2024.
[BGN11]
Florent Benaych-Georges and Raj Rao Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances in Mathematics, 227(1):494–521, 2011.
[Cam03]
Francesco Camastra. Data dimensionality estimation methods: a survey. Pattern recognition, 36(12):2945–2954, 2003.
[CH04]
Jose A Costa and Alfred O Hero. Learning intrinsic dimension and intrinsic entropy of high-dimensional datasets. In 2004 12th European Signal Processing Conference, pages 369–372. IEEE, 2004.
[CM23]
Elisabetta Cornacchia and Elchanan Mossel. A mathematical model for curriculum learning for parities. In International Conference on Machine Learning, pages 6402–6423. PMLR, 2023.
[CMM25]
Elisabetta Cornacchia, Dan Mikulincer, and Elchanan Mossel. Low-dimensional functions are efficiently learnable under randomly biased distributions. arXiv preprint arXiv:2502.06443, 2025.
[CPT+ 24]
Francesco Cagnetta, Leonardo Petrini, Umberto M Tomasini, Alessandro Favero, and Matthieu Wyart. How deep neural networks learn compositional data: The random hierarchy model. Physical Review X, 14(3):031001, 2024.
[CW01]
Anthony Carbery and James Wright. Distributional and Lq norm inequalities for polynomials over convex bodies in Rn . Math. Res. Lett., 8(3):233–248, 2001.
[CW24]
Francesco Cagnetta and Matthieu Wyart. Towards a theory of how the structure of language is acquired by deep neural networks. arXiv preprint arXiv:2406.00048, 2024.
[DKL+ 23]
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305.18270, 2023.
[DLS22]
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413– 5452. PMLR, 2022.
[DM20]
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33:20356–20365, 2020.
[DPVLB24] Alex Damian, Loucas Pillaud-Vivien, Jason Lee, and Joan Bruna. Computationalstatistical gaps in gaussian single-index models. In The Thirty Seventh Annual Conference on Learning Theory, pages 1262–1262. PMLR, 2024. [DPZK25]
Yatin Dandi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The computational advantage of depth: Learning high-dimensional hierarchical functions with gradient descent. arXiv preprint arXiv:2502.13961, 2025.
12
[DTA+ 24]
Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220, 2024.
[FdRL17]
Elena Facco, Maria d’Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. Scientific reports, 7(1):12140, 2017.
[FSC+ 25]
Alessandro Favero, Antonio Sclocchi, Francesco Cagnetta, Pascal Frossard, and Matthieu Wyart. How compositional generalization and creativity improve as diffusion models are trained. arXiv preprint arXiv:2502.12089, 2025.
[GBMMS24] Jérôme Garnier-Brun, Marc Mézard, Emanuele Moscato, and Luca Saglietti. How transformers learn structured data: insights from hierarchical filtering. arXiv preprint arXiv:2408.15138, 2024. [Gla23]
Margalit Glasgow. Sgd finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the xor problem. arXiv preprint arXiv:2309.15111, 2023.
[GMKZ20]
Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure on learning in neural networks: The hidden manifold model. Physical Review X, 10(4):041044, 2020.
[HKM24]
Aapo Hyvärinen, Ilyes Khemakhem, and Ricardo Monti. Identifiability of latent-variable and structural-equation models: from linear to nonlinear. Annals of the Institute of Statistical Mathematics, 76(1):1–33, 2024.
[JKMS25]
Nirmit Joshi, Hugo Koubbi, Theodor Misiakiewicz, and Nathan Srebro. Learning singleindex models via harmonic decomposition. arXiv preprint arXiv:2506.09887, 2025.
[JMS24]
Nirmit Joshi, Theodor Misiakiewicz, and Nathan Srebro. On the complexity of learning sparse functions with statistical and gradient queries. arXiv preprint arXiv:2407.05622, 2024.
[KB14]
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[KCGK24]
Yiwen Kou, Zixiang Chen, Quanquan Gu, and Sham Kakade. Matching the statistical query lower bound for k-sparse parity problems with sign stochastic gradient descent. Advances in Neural Information Processing Systems, 37:113001–113037, 2024.
[KST+ 18]
Hyun Min Kang, Meena Subramaniam, Sasha Targ, Michelle Nguyen, Lenka Maliskova, Elizabeth McCarthy, Elaine Wan, Stephanie Wong, Lauren Byrnes, Cristina M. Lanata, Rachel E. Gate, Sara Mostafavi, Alexander Marson, Noah Zaitlen, Lindsey A. Criswell, and Chun Jimmie Ye. Multiplexed droplet single-cell rna-sequencing using natural genetic variation. Nature Biotechnology, 36:89–94, 2018.
[LB04]
Elizaveta Levina and Peter Bickel. Maximum likelihood estimation of intrinsic dimension. Advances in neural information processing systems, 17, 2004.
[LGC+ 21]
Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with a teacher-student model. Advances in Neural Information Processing Systems, 34:18137–18151, 2021. 13
[MHWSE23] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradientbased feature learning under structured data. Advances in Neural Information Processing Systems, 36:71449–71485, 2023. [MKAS21]
Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. In International Conference on Machine Learning, pages 7379–7389. PMLR, 2021.
[MKL+ 20]
Francesca Mignacco, Florent Krzakala, Yue Lu, Pierfrancesco Urbani, and Lenka Zdeborova. The role of regularization in classification of high-dimensional noisy gaussian mixture. In International conference on machine learning, pages 6874–6883. PMLR, 2020.
[MKUZ20]
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540–9550, 2020.
[Mos16]
Elchanan Mossel. Deep learning and hierarchal generative models. arXiv preprint arXiv:1612.09057, 2016.
[O’D14]
Ryan O’Donnell. Analysis of Boolean Functions. Cambridge University Press, 2014.
[PMR+ 17]
Tomaso Poggio, Hrushikesh Mhaskar, Lorenzo Rosasco, Brando Miranda, and Qianli Liao. Why and when can deep-but not shallow-networks avoid the curse of dimensionality: a review. International Journal of Automation and Computing, 14(5):503–519, 2017.
[RDKL26]
Yunwei Ren, Yatin Dandi, Florent Krzakala, and Jason D Lee. Provable learning of random hierarchy models and hierarchical shallow-to-deep chaining. arXiv preprint arXiv:2601.19756, 2026.
[RGKZ21]
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning, pages 8936–8947. PMLR, 2021.
[SBGG24]
Eszter Szekely, Lorenzo Bardone, Federica Gerace, and Sebastian Goldt. Learning from higher-order correlations, efficiently: hypothesis tests, random features, and neural networks. Advances in Neural Information Processing Systems, 37:78479–78522, 2024.
[SFLW24]
Antonio Sclocchi, Alessandro Favero, Noam Itzhak Levi, and Matthieu Wyart. Probing the latent hierarchical structure of data via diffusion models. arXiv preprint arXiv:2410.13770, 2024.
[SFW25]
Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. Proceedings of the National Academy of Sciences, 122(1):e2408799121, 2025.
[SSBD14]
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[TDD+ 24]
Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Fundamental limits of weak learnability in high-dimensional multi-index models. In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024.
14
[TDPK26]
Hugo Tabanelli, Yatin Dandi, Luca Pesce, and Florent Krzakala. Deep learning of compositional targets with hierarchical spectral methods. arXiv preprint arXiv:2602.10867, 2026.
[Tro12]
Joel A. Tropp. User-friendly tail bounds for sums of random matrices. Foundations of Computational Mathematics, 12(4):389–434, 2012.
15
A
Proof of Theorem 1
Algorithm 1 Layerwise SGD with Gaussian initialization Init. scale κ, bias range A, learning rates γ1 , γ2 , step counts T1 = 1, T2 , batch size B. Pn Require: Training data {(xt , yt )}t≥1 , model: NN(x; (w, a, b)) = i=1 ai σ(wi⊤ x + bi ). 1: Initialize: first-layer weights w (0) ∼ N (0, Id /d) (normalized variance), a(0) = τ , b(0) = 0. Phase 1: Train the first layer (second layer frozen) 2: for t = 0 to T1 − 1 : do
3: Sample a fresh mini-batch S (t) = {(xs , ys )}B s=1 4: Update w(t+1) ← w(t) − γ1 ∇w(t) L(S (t) ; w(t) , a(0) , b(0) ) 5: end for
Phase 2: Train the second layer (first layer frozen) 6: Draw random biases: b̂1 , . . . , b̂n ∼ Unif[−A, A]⊗n 7: for t = 0 to T2 − 1 : do 8: Sample a fresh mini-batch S (t) = {(xs , ys )}B s=1
9: Update a(t+1) ← a(t) − γ2 ∇a(t) L(S (t) ; w(T1 ) , a(t) , b̂) 10: end for 11: Output: Trained model NN(x; (w (T1 ) , a(T2 ) , b̂))
Layerwise-SGD with square loss. We consider learning with a 2-layer neural network: NN(x; θt ) =
n X
ah σ (wht )⊤ x + bh ,
h=1
with n hidden units, and polynomial activation σ, satisfying Assumption 2. We initialize whj ∼ N (0, 1/d) and ah = τ , for all h ∈ [n], j ∈ [d]. We perform layerwise training: i.e., during a first phase we train the first layer’s weights for one step, keeping the second layer’s weights fixed, and during a second phase we train the second layer’s weights, keeping the wh fixed. The bias neurons are initialized iid at 0 and kept fixed during the first step. Before the second phase, we draw bh ∼ Unif[−A, A], for a fixed A > 0, and keep them frozen during training. We adopt this choice by convenience of the analysis. We note, however, that with further technical work, one could consider random biases for the first part of training as well. We use the squared loss: 12 (y − ŷ)2 . We refer to Algorithm 1 for a detailed description of the layerwise SGD algorithm that we use. For each hidden unit h ∈ [n], wht+1 = wht −
B 2 γ 1 X ∇wht y b − NN(xb ; θt ) 2B
(5)
b=1
where B ∈ N denotes the batch size. Note that | NN(x; θt )| = τ polylog(d) with high probability over the initialization and the mini-batch, for B = poly(d). Indeed, x is centered (by symmetry of s) and sub-Gaussian (since |si | = 1 and ξj is sub-Gaussian), so (wh0 )⊤ x is sub-Gaussian with variance O(1), hence |σ((wh0 )⊤ x)| = O(polylog(d)) and |ah | = O(τ ); a union bound over h ∈ [n] and b ∈ [B] gives the polylog factor (using n, P = Od (1)). Thus, for τ small enough, we have 1 −γ∇wht (y − NN(x; θt ))2 = γah yσ ′ (wht x)x − γ NN(x; θt )ah σ ′ (wht x)x 2 = γτ yσ ′ (wht x)x − γO nτ 2 poly log(d) ,
(6) (7)
with high probability. Thus, for τ small, the interaction term NN(x; θt )ah σ ′ (wht x)x is negligible, and one can track only the correlation term. A similar approach is used in e.g. [AAM23, MHWSE23]. 16
A.1
First layer’s training
Initial population gradient. In the following, let us fix a hidden neuron h ∈ [n], and for simplicity let us denote w := wh0 . Recall, we initialize ah = τ and bh = 0 for all h ∈ [n]. For i ∈ [N ], let us denote: X (∆ξ (w))i := ξj wj . (8) j∈Ci
For j ∈ Ci , the correlation term of the initial population gradient reads: 1 Ex,y [y · ∂wj NN(x; θ0 )] τ = Es [f (s)si Eξ [ξj σ ′ (s⊤ ∆ξ (w))]].
(9)
Ḡ(wj ) : =
(10)
In the following, we show that G(wj ) concentrates around a quantity that depends only on the cluster’s latent i. Lemma 1 (Initial Population Gradient). Assume the conditions of Theorem 1; in particular, the noise variables {ξj }j∈[d] are mutually independent, sub-Gaussian with ∥ξj ∥ψ2 ≤ ψ, with means mj := E[ξj ], variances σj2 := Var(ξj ), and bounded third moments supj E|ξj − mj |3 = Od (1). Fix a neuron w ∈ Rd and define (∆ξ (w))i :=
X
ξj w j ,
(∆m (w))i :=
j∈Ci
X
V 2 :=
mj wj ,
d X
wj2 σj2 .
j=1
j∈Ci
Then for every i ∈ [N ] and every j ∈ Ci , G(wj ) = mj αi (w) + ηj , where
h i αi (w) := Es f (s) si EG σ ′ s⊤ ∆m (w) + G ,
G ∼ N (0, V 2 ),
(11)
and, on with high probability, r |ηj | ≤ C1 |wj | +
C2 wj2
+ C3
log d . d
The constants C1 , C2 , C3 > 0 depend only on the activation σ and on the sub-Gaussian parameters and third moments of ξ. The proof of lemma 1 is deferred to Appendix B.1. Initial estimated gradient.
We now consider the gradient estimated through B i.i.d. samples: B
GB (w) :=
1 X b ′ y σ (wxb ) xb , B
(12)
b=1
where we recall that σ is a polynomial activation of degree P such that 2N ≤ P = Od (1). If B is large enough, the estimated gradient is close to the population gradient Ḡ(w), as formalized by the following Lemma.
17
Lemma 2. Let Mf,N = maxs∈{±1}N |f (s)|. There exists a constant KN > 0 (depending only on the polynomial coefficients of σ and on the subgaussian parameter of x) such that, if B≥C
2 2 KN Mf,N log(dn) ζ2
(13) ∗
for some universal constant C > 0, then with probability at least 1 − d−C , for some C ∗ > 0, we have, for all h ∈ [n], (14)
∥GB (wh0 ) − Ḡ(wh0 )∥∞ ≤ ζ. Proof. Fix h ∈ [n] and a coordinate j ∈ [d]. Define
(15)
Zb := y b σ ′ (wh0 xb ) xbj .
Since |y b | ≤ Mf,N and x is subgaussian while wh0 ∼ N (0, Id /d) is independent of x, the inner product wh0 x is subgaussian. A polynomial function of a subgaussian variable multiplied by a subgaussian coordinate is subexponential. Hence, there exists a constant KN > 0 such that ∥Zb ∥ψ1 ≤ KN Mf,N . Applying Bernstein’s inequality for subexponential random variables yields, for some constants c, C > 0, ! ( )! B ζ2 1 X ζ Zb − E[Zb ] > ζ ≤ 2 exp −cB min Pr . (16) 2 M2 , K M B KN N f,N f,N b=1
2 2 KN Mf,N 2 ζ
∗
Choosing B ≥ C log(dn) makes the right-hand side at most (dn)−C for some C ∗ > 0. A union bound over all j ∈ [d] and h ∈ [n] then gives the desired result. By the above, it follows that: wj1 = w0 + γτ (mj αi + ηj + ωj + ζ), (17) p where αi is defined in (11), |ηj | ≤ C1 |wj | + C2 wj2 + C3 log(d)/d (Lemma 1), ωj = O(nτ poly log(d)) (control of the interaction component of the gradient), and ζ follows by Lemma 2, assuming the batch size B = Ω(log(dn)/ζ 2 ).
A.2
Second layer’s training
We now show that training the second layer, while keeping the first layer fixed to (17), allows to achieve small generalization error. the proof follows by the following steps: 1. We show that the (αi )i∈[N ] are such that any f : {±1}N → R is projection-consistent. Specifically, we show that for all s, t ∈ {±1}N such that f (s) ̸= f (t), and for all ϵ > 0, we have |αT (s − t)| > ϵ,
(18)
with probability at least 1 − poly(ϵ) over the initialization of the first layer’s weights (Lemma 3). This makes use of the well-known Carbery-Wright inequality [CW01]. 2. Next, we show that the property above, and randomly drawn bias neurons, guarantee that there exists an assignment of the second layer’s weights such that f (s) ≈
n X
a∗h σ(γvi sT α + bh )
h=1
for all s ∈ {±1}N (Lemma 5). 18
(19)
3. Then, we show that the generalization error of the certificate of the previous point is low (Lemma 6). This requires bounding the errors appearing in (17). 4. Lastly, we use standard results on the convergence of SGD on convex losses to conclude that the second layer’s training will convergence to small generalization error (Theorem 3). Let us start with the first step. P Lemma 3 (Non-degeneracy). Let Z ∼ N (0, diag(vi )i∈[N ] /d), and let vsum := i∈[N ] vi . Assume vi ≥ vmin > 0, for all i ∈ [N ]. Let σ be a polynomial of degree P that satisfies assumption 2 and let σ ′ be its first derivative. Let f : {±1}N → R be a target function and assume Var(f ) > ζ > 0. For c ∈ {0, ±2}N , define: N h X i PZ (c) := Es f (s) · C̃j vj cj sj · EG σ ′ (sT Z + G) ,
(20)
j=1
where G ∼ N (0, V 2 ), for some V > 0 and where C̃j > 0 for j ∈ [N ]. Let Lf := {(t, y) ∈ {±1}N × {±1}N : f (t) ̸= f (y)}. Then, there exist constants C1 , C2 > 0 and k0 ∈ [P ], such that for all ϵ > 0, PZ
min (t,y)∈Lf
PZ (t − y) > ϵ · vmin
v
min
k0 /2
d
> 1 − C1 P |Lf |
ϵ C2
1/P .
(21)
PP Proof. Write σ ′ (x) := n=1 an xn . Define ψ(t) := EG∼N (0,V 2 ) [σ ′ (t+G)]. This is a polynomial of degree PN ≤ P , and for fixed c, PZ (c) = Es [f (s)hc (s)ψ(sT Z)], hc (s) = j=1 C̃j vj cj sj . Note that also PZ (c) is a polynomial in Z of degree at most P . Let us state the following lemma, which is a restatement of the Carbery-Wright inequality [CW01]. Lemma 4 ([CW01]). Let ϕ : RN → R be a non-zero polynomial of degree at most P = O(1), and let µ be a log-concave probability measure over RN . Then, for all α > 0, µ {x : |ϕ(x)| ≤ α∥ϕ∥L2 (µ) } ≤ CP α1/P , (22) where C > 0 is an absolute constant. By applying Lemma 4 to our setting with µ being the Gaussian law of Z, and by union bound, we get that there exists a constant C such that with probability at least 1 − CP |Lf |ϵ1/P over Z, p min |PZ (t − y)| > ϵ · Mmin , (23) (t,y)∈Lf
where Mmin := min(t,y)∈Lf EZ [PZ (t − y)2 ]. Thus everything reduces to a lower bound on Mmin . Fix a pair (t, y) and let c := t − y. Recall that PN hc (s) := j=1 C̃j vj cj sj . Then, EZ [PZ (c)2 ] = EZ Es [f (s)hc (s)ψ(sT Z)]2 = Es,s′ f (s)f (s′ )hc (s)hc (s′ )EZ [ψ(sT Z)ψ((s′ )T Z)] .
(24) (25)
Let ζ 2 := Var(sT Z) = vsum /d be the variance of sT Z. Let x ∼ N (0, 1). As usual, we can write the Hermite expansion, ψ(ζx) =
P X
βk (ζ)Hk (x),
βk (ζ) =
k=1
19
1 Ex [ψ(ζx)Hk (x)]. k!
(26)
Note that: ζ k (k) S 2 2 (0), k! V +ζ
βk (ζ) : =
(27) (k)
where SV 2 +ζ 2 (t) := EG∼N (0,V 2 +ζ 2 ) [σ ′ (t + G)]. By Assumption 2, for each 1 ≤ k ≤ N , SV 2 +ζ 2 (0) ≥ c0 > 0. For two vectors s, s′ ∈ {±1}N , write the correlation: EZ [ψ(sT Z)ψ((s′ )T Z)] =
P X
(28)
βk (ζ)2 k!ρ(s, s′ )k ,
k=0 1 where ρ(s, s′ ) := vsum thus write:
PN
′ i=1 si si vi . Plugging it in the above, and using a tensorization trick, we can
EZ [PZ (c)2 ] =
P X k=0
≥
P X
1 X βk (ζ)2 k! k vi ...vik · Es [f (s)hc (s)si1 ...sik ]2 vsum i ,...i 1 1
βk (ζ)2 k!
k=0
vmin vsum
(29)
k
k max
T ⊆[N ]:|T |≤k
Es [f (s)hc (s)χT (s)]2
(30)
By Assumption 2, there exists some fixed k0 ∈ [P ] and c1 > 0 such that EZ [PZ (c)2 ] ≥ c1
v
sum
k0 v
d
min
k0 max
vsum
T ⊆[N ]:|T |≤k0
Es [f (s)hc (s)χT (s)]2 .
(31)
The following claim bounds the right-most term. Claim 1. There exists a subset T ⊆ [N ] and a constant C > 0 such that 2 2 Es f (s)hc (s)χT (s) ≥ 22−2N C vmin . Proof of Claim 1. Fix (t, y) with f (t) ̸= f (y). Let D := {i ∈ [N ] : ti ̸= yi }. Since t ̸= y, we have D ̸= ∅. For i ∈ D, we have ti − yi ∈ {±2} and for i ∈ / D the coefficient is zero. Compute hc at t: hc (t) =
N X
C̃i vi (ti − yi )ti =
i=1
X
C̃i vi 2t2i = 2
i∈D
X
C̃i vi .
i∈D
Using vi ≥ vmin and C̃i ≥ C̃min , and D ̸= ∅, we obtain |hc (t)| = 2
P
i∈D C̃i vi
≥ 2 C̃min vmin .
Similarly, hc (y) = −hc (t), so |hc (y)| = |hc (t)| ≥ 2C̃min vmin . Since f (t) ̸= f (y), at least one of f (t), f (y) is nonzero. Define such that |f (s0 )| = max{|f (t)|, |f (y)|} ≥ c0 > 0.
s0 ∈ {t, y}
Thus, denoting u(s) = f (s)hc (s), wePhave maxs∈{±1}N |u(s)| ≥ 2C̃min vmin c0 . Now consider the Fourier-Walsh expansion of u: u(s) = T ⊆[N ] û(T ) χT (s), with û(T ) = Es [u(s)χT (s)]. For any fixed s, P P |u(s)| = T û(T )χT (s) ≤ T |û(T )|. Taking the maximum over s gives X
|û(T )| ≥ max |u(s)| ≥ 2C̃min vmin c0 . s
T
On the other hand, X
|û(T )| ≤ 2N max |û(T )|. T
T
20
Combining these two inequalities yields max |û(T )| ≥
T ⊆[N ]
2C̃min vmin c0 = 21−N C̃min vmin c0 . 2N
Hence there exists some T ⊆ [N ] (depending on t, y) with |û(T )| = Es [f (s)hc (s)χT (s)] ≥ 21−N C̃min vmin c0 , and therefore
2 2 Es [f (s)hc (s)χT (s)]2 = û(T )2 ≥ 22−2N C̃min vmin c0 .
This proves the claim. The lemma follows by direct application of the claim. Now, we proceed to building a certificate for our target function f , as linear combination of the representation learned by the first layer. In order to do that, let (α̃i )i∈[N ] be such that α̃i := γτ αi vi , and define the projection values u(s) := s⊤ α̃, s ∈ {±1}N , and let U := {u1 < · · · < uM } be the set of distinct values taken by u(s), so that M ≤ 2N . Define also the minimum gap ∆ := min ′ |um − um′ |. m̸=m
On the event of Lemma 3, f is projection-consistent, i.e. u(s) = u(t) =⇒ f (s) = f (t). Lemma 5 (Certificate). Assume that Assumptions 1 and 2 are satisfied. Let α ∈ RN be defined by (11), iid
and define (α̃i )i∈[N ] by α̃i := γτ αi vi . Let bh ∼ Unif[−A, A] for h ∈ [n], with A ∈ [∆/CA , CA ∆], for some fixed constant CA ≥ 1, and A ≥ maxm |um |. Fix ϵ ∈ (0, 1). Then there exist constants C0 , C1 , C2 , Cσ,P,N,CA > 0 such that if n ≥ C0 2N log 2N + log(1/ϵ) , there exists a⋆ ∈ Rn with: 2 Mf,N . n ∆2P
(32)
a⋆h σ(s⊤ α̃ + bh ).
(33)
∥a⋆ ∥22 ≤ Cσ,P,N,CA such that on the event of Lemma 3, for all s ∈ {±1}N : f (s) =
n X h=1
Proof. Let Mf,N := sups∈{±1}N |f (s)|. Since |um | ≤ ∥α̃∥1 and |bh | ≤ A, and since σ is a degree-P polynomial, there exists a constant Cσ > 0 such that |σ(um + bh )| ≤ Cσ AP
∀ m ∈ [M ], ∀ h ∈ [n].
Define the feature matrix Φ ∈ RM ×n by Φm,h := σ(um +bh ), and its normalized version Ψ := A−P Φ, so that Φ = AP Ψ and |Ψm,h | ≤ Cσ . Now let K := Eb [ψ(b)ψ(b)⊤ ] ∈ RM ×M , ψ(b) := σ(u1 + b)/AP , . . . , σ(uM + b)/AP , 21
where b ∼ Unif[−A, A]. Since σ has degree P , nonzero leading coefficient, and M ≤ 2N ≤ P + 1, the shifted polynomials b 7→ σ(um + b), m = 1, . . . , M, are linearly independent whenever the shifts u1 , . . . , uM are distinct. Therefore, for every c ∈ RM \ {0}, !2 M X σ(u + b) m > 0, c⊤ Kc = Eb cm P A m=1 and hence K ≻ 0. We may therefore define κ := λmin (K) > 0. Claim 2. For every fixed constant CA ≥ 1, there exists a constant cσ,P,N,CA > 0 such that the following holds. If A ∈ [∆/CA , CA ∆], then κ = λmin (K) ≥ cσ,P,N,CA . Let us first show how Claim 2 implies the Lemma. We work on the event of Lemma 3, on which f is projection-consistent. Define the empirical Gram matrix n
X b := 1 ΨΨ⊤ = 1 ψ(bh )ψ(bh )⊤ . K n n h=1
Pn 1
b = Set Xh := ψ(bh )ψ(bh )⊤ ⪰ 0, so that K h=1 Xh and E[Xh ] = K. Moreover, since |ψm (b)| ≤ Cσ n for all m ∈ [M ], we have ∥Xh ∥2 = ∥ψ(bh )∥22 ≤ M Cσ2 . By a matrix Chernoff bound (e.g. [Tro12]), there exist a numerical constant c > 0 such that for any δ ∈ (0, 1), δ2 n κ b Pr λmin (K) ≤ (1 − δ)κ ≤ M exp −c , M Cσ2 where κ := λmin (K) > 0. Choosing δ = 12 and taking n ≥ C
M Cσ2 log M + log(1/ϵ) , κ
yields, with probability at least 1 − C1 ϵC2 , b ≥ λmin (K)
κ . 2
Equivalently, λmin (ΦΦ⊤ ) = A2P λmin (ΨΨ⊤ ) ≥
κn 2P A . 2
(34)
By projection-consistency, define F ∈ RM by Fm = f (sm ), where sm is such that um = u(sm ), and √ note that ∥F ∥2 ≤ M Mf,N . Let a⋆ be the interpolating solution: a⋆ := ΦT (ΦΦT )−1 F. Then Φa⋆ = F , thus sup
f (s) −
s∈{±1}N
n X
a⋆h σ(s⊤ α̃ + bh ) = 0.
h=1
Now, let us upper bound the norm of a . Using ∥Φ∥2 ≤ ∥Φ∥F and the bound |σ(um + bh )| ≤ Cσ AP , we have √ ∥Φ∥2 ≤ Cσ M n AP . ⋆
22
Moreover, ∥F ∥2 ≤
√
M Mf,N . Therefore, ∥a⋆ ∥2 ≤ ∥Φ∥2 ∥(ΦΦ⊤ )−1 ∥2 ∥F ∥2 ≤
∥Φ∥2 ∥F ∥2 . λmin (ΦΦ⊤ )
2P Combining with (34), namely λmin (ΦΦ⊤ ) ≥ κn , yields 2 A
∥a⋆ ∥22 ≤ Cσ,P,N,CA
2 M 2 Mf,N . n A2P
Since M ≤ 2N and N = Od (1), we may absorb the factor M 2 into the constant, and therefore ∥a⋆ ∥22 ≤ Cσ,P,N,CA
2 Mf,N . n A2P
Since A ∈ [∆/CA , CA ∆], it follows that ∥a⋆ ∥22 ≤ Cσ,P,N,CA
2 Mf,N . n ∆2P
All is left is the proof of Claim 2. Proof of Claim 2. Let us write σ(x) = aP xP + aP −1 xP −1 + · · · + a0 with aP ̸= 0. Let tm := um /A and z := b/A ∼ Unif[−1, 1]. Since A ∈ [∆/CA , CA ∆], we have |tm − tm′ | =
∆ = min ′ |um − um′ |, m̸=m
1 |um − um′ | ≥ , A CA
∀ m ̸= m′ .
Define ϕA (z) := Then, for each m,
σ(A(t1 + z)) σ(A(tM + z)) , . . . , , AP AP
K(A, t) := E[ϕA (z)ϕA (z)⊤ ].
P −1 X σ A(tm + z) P = a (t + z) + ak Ak−P (tm + z)k . P m AP k=0
Since |tm | ≤ 1 and |z| ≤ 1, we have |tm + z| ≤ 2, and therefore the remainder term is uniformly O(1/A) over z ∈ [−1, 1] and over all choices of (t1 , . . . , tM ) ∈ [−1, 1]M satisfying t1 < · · · < tM , Consequently,
tm+1 − tm ≥
1 . CA
ϕA (z) −→ ϕ∞ (z) := aP (t1 + z)P , . . . , aP (tM + z)P
uniformly on [−1, 1] as A → ∞, uniformly over all such separated tuples (t1 , . . . , tM ). It follows that K(A, t) converges in operator norm to K∞ := E ϕ∞ (z)ϕ∞ (z)⊤ uniformly over the same class of tuples. 23
Let c ∈ RM \ {0} and define M X
g∞,c (z) :=
cm aP (tm + z)P .
m=1
Since M ≤ 2 ≤ P + 1 and the shifts t1 , . . . , tM are distinct, the functions z 7→ (tm + z)P are linearly independent. Hence g∞,c is not identically zero, and therefore c⊤ K∞ c = E g∞,c (z)2 > 0. N
Thus K∞ ≻ 0, and κ∞ := λmin (K∞ ) > 0. Since the set of separated tuples n o TCA := (t1 , . . . , tM ) ∈ [−1, 1]M : t1 < · · · < tM , tm+1 − tm ≥ 1/CA is compact, and t 7→ λmin (K∞ (t)) is continuous and strictly positive on TCA , there exists κ∞,CA > 0 such that λmin (K∞ ) ≥ κ∞,CA uniformly over TCA . Hence, by uniform convergence and Weyl’s inequality, there exists A1 ≥ 1 such that for all A ≥ A1 , κ∞,CA λmin (K(A, t)) ≥ . 2 Finally, for A ∈ [1, A1 ], the map (A, t) 7→ λmin (K(A, t)) is continuous on the compact set [1, A1 ] × TCA and is strictly positive there. Therefore it attains a positive minimum on that set. Combining the two regimes yields a constant cσ,P,N,CA > 0 such that κ = λmin (K) ≥ cσ,P,N,CA , which proves the claim. Lemma 6 (Generalization error). Let θ̂ := (w1 , a∗ , b). Then, for any δ > 0, under the assumptions of Lemma 5, (35)
Ex,y [|y − NN(x; θ̂)|] < δ, for sufficiently large d. Proof. Recall that, for each j ∈ Ci , and for each h ∈ [n]:
(36)
1 0 ′ wj,h = wj,h + γτ mj αi + γτ ηj,h ,
with ′ ηj,h = ζ + ωj,h + ηj,h , (37) p where αi is defined in (11), |ηj | ≤ C1 |wj | + C2 wj2 + C3 log(d)/d (Lemma 1), ωj = O(nτ poly log(d)), and ζ follows by Lemma 2, assuming the batch size B = Ω(log(dn)/ζ 2 ). Note that by Lemma 7, √ 0 |wj,h | = O(log(dn)/ d) for all j, h. Thus, for a fixed h,
(wh1 )T x = γτ
N X i=1
= γτ
N X i=1
αi si
X
mj ξj +
(38)
j=1
j∈Ci
αi si vi + γτ
d X (γτ ηj′ + wj0 )xj
N X i=1
αi si
X j∈Ci
24
mj (ξj − mj ) +
d X (γτ ηj′ + wj0 )xj j=1
(39)
Let NN(s; θ∗ ) :=
Pn
∗ T h=1 ah σ(α̃ s + bh ), with α̃i = γτ αi vi ,
(40)
Ex,y [|y − NN(x; θ̂)|] ≤ Ex,y [|y − NN(s; θ∗ )|] + Ex,y [| NN(s; θ∗ ) − NN(x; θ̂)|] n (a) X ≤ |a∗i | · Ex,y |σ(α̃T s + b) − σ((w1 )T x + b)|
(41)
i=1
where (a) holds because, by Lemma 5, the first term is zero. By mean value theorem for polynomials, (42)
|σ(α̃T s + b) − σ((w1 )T x + b)| ≤ |σ ′ (t)| · |α̃T s − (w1 )T x| ≤ Cσ (1 + |t|
P −1
T
(43)
1 T
) · |α̃ s − (w ) x|
with t between α̃T s and (w1 )T x and Cσ := maxi∈[P ] P |βk |. By Cauchy-Schwartz and by sub-Gaussianity of the x, we get q q Ex [|σ(α̃T s + b) − σ((w1 )T x + b)|] ≤ Cσ E(1 + |t|P −1 )2 · Ex,y [((w1 )T x − α̃T s)2 ] (44) Thus, N d X 2 X X Ex [((w1 )T x − α̃T s)2 ] = Es,ξ [ γτ αi si mj (ξj − mj ) + (γτ ηj′ + wj0 )xj ] i=1
(45)
j=1
j∈Ci
N d 2 i h X 2 i h X X ≤ 2 γ 2 τ 2 Es,ξ αi si mj (ξj − mj ) + Es,ξ (γτ ηj′ + wj0 )xj . i=1
j=1
j∈Ci
|
{z
:=A
}
|
{z
:=B
} (46)
For the first term, N X 2 X Eξ [A2 ] = E αi si mj (ξj − mj ) i=1
=
N X
(47)
j∈Ci
αi αk m⊤ Ci ΣCi ,Ck mCk
(48)
X
(49)
i,k=1
≤ ∥Σ∥op
αi2 ∥mCi ∥22
i
≤ ∥Σ∥op ∥α∥2∞ vsum ,
(50)
where mCi = (mj )j∈Ci and ΣCi ,Ck denotes the corresponding block of Σ. Pd For the second term, write B = j=1 η ′′ xj , where ηj′′ := γτ ηj′ + wj0 and note that cov(x) = Σ. Then Es,ξ [B 2 ] = η ′′⊤ Σ η ′′ + (η ′′⊤ E[x])2
(51)
≤ ∥Σ∥op ∥η ′′ ∥22 ,
(52)
√ since s is symmetric, thus E[x] = 0. Note that ∥w0 ∥22 = θd (1) (see Lemma 7). Set √ γ = 1/ vsum . √ √ vsum 2 If B ≥ C log(dn)Mf,N d/vsum , then ζ = O( vsum / d). Furthermore, if τ = O( √d log(d) ), then vsum ω = O( log(d)d ). Thus, ∥η ′′ ∥22 = θd (1). Then,
Ex [((w1 )T x − α̃T s)2 ] ≤ C0 25
(53)
Then, since by Lemma 5 ∥a∗ ∥1 ≤
√
′
n∥a∗ ∥2 ≤ vCP , we have min q Ex,y [|y − NN(x; θ̂)|] ≤ C∥a∗ ∥1 · E(1 + |t|P −1 )2 ≤C
1 P −1 v = C/vmin ≤ δ, P vmin min
(54) (55)
for large enough d, since vmin = θ(d) by assumption. To conclude, we use the following well known result on convergence of SGD on convex losses. Theorem 3 ([SSBD14]). Let L be a convex function and let a∗ ∈ argmin∥a∥2 ≤B L(a), for some B > 0. For all t, let αt be such that E [αt | at ] = −∇at L(at ) and assume ∥αt ∥2 ≤ ξ for some ξ > 0. If a(0) = 0 B , then and for all t ∈ [T ] at+1 = at + γαt , with γ = ξ√ T T
1X Bξ L(at ) ≤ L(a∗ ) + √ . T t=1 T −2P −P P ϵ/n). ), ξ = θ(nvmin ), T2 = θ(n2 /ϵ2 ), and γ2 = θ(vmin We take B = θ(vmin
B
Technical Lemmas iid
Lemma 7. Let wjh ∼ N (0, 1/d), for j ∈ [d], h ∈ [n]. Let ∆m (wh ) ∈ RN be such that (∆m (wh ))i = P j∈Ci wj mj . Then, for any δ > 0, with probability 1 − 2nδ, where n is the number of hidden units, and for d large enough, we have: (i) maxh∈[n] ∥wh ∥2 ≤ 3/2. p (ii) maxh∈[n] ∥∆m (wh )∥L1 ≤ C log(2/δ). p √ (iii) maxh,j |wjh | ≤ O( log(dn)/ d). Pd Proof. (i) For each h, we have ∥wh ∥22 = (1/d) j=1 Zj2 = (1/d)χ2d . By standard χ2 concentration, Pr(∥wh ∥2 ≥ 3/2) ≤ e−cd . Taking a union bound over h ∈ [n] gives Pr(maxh ∥wh ∥2 ≤ 3/2) ≥ 1 − ne−cd , −cd which is at least 1 − nδ for d large ≤ δ.P P enough soh that e h (ii) For each i, (∆m (w ))i = j∈Ci mj wj ∼ N (0, (1/d) j∈Ci m2j ), hence it is sub-Gaussian with variance proxy at most C 2 = maxi ∥(mj )j∈Ci ∥22 /d. Note that 0 < C 2 < ∞ by assumption. Therefore p 2 2 Pr(|(∆m (wh ))i | > t) ≤ 2e−t /(2C ) . Choosing t = C p 2 log(2/δ) yields probability at most δ. By a h union bound over h, we obtain maxh |(∆m (w ))i | ≤ C 2 log(2/δ) with probability at least 1 − nδ. 2 (iii) Each coordinate satisfies Pr(|wjh | > t) ≤ 2e−dt /2 . A union bound over all dn entries gives q p 2 h Pr(maxh,j |wjh | > t) ≤ 2dn e−dt /2 . Taking t = d2 log 2dn (2/d) log(2dn/δ) δ shows that maxh,j |wj | ≤ p √ h with probability at least 1 − δ, i.e. maxh,j |wj | = O( log(dn)/ d). Combining the three parts and adjusting constants so that e−cd ≤ δ shows that all items hold simultaneously with probability at least 1 − 2nδ. Lemma 8 (Concentration of V ). Let w = (w1 , . . . , wd ) ∼ N (0, Id /d) be independent of the noise vector ξ, and suppose the coordinates of ξ are jointly sub-Gaussian with covariance matrix cov(ξ) = Σ ⪰ 0. Define 1 V = w⊤ Σw = Z ⊤ ΣZ, Z ∼ N (0, Id ). d Then Ew [V ] = µ := d1 Tr(Σ). Moreover, for some universal constant c > 0 and every ϵ ∈ (0, 1], d µ 2 ϵ . (56) Pr( |V − µ| ≥ ϵµ ) ≤ 2 exp − c ∥Σ∥op 26
Proof. By the Hanson–Wright inequality for Gaussian quadratic forms, 2 t t ⊤ Pr Z ΣZ − Tr(Σ) ≥ t ≤ 2 exp −c min , . ∥Σ∥2F ∥Σ∥op Substituting t = d εµ = εTr(Σ) and dividing by d gives 2 2 2 d ε µ d εµ Pr( |V − µ| ≥ εµ ) ≤ 2 exp −c min , . ∥Σ∥2F ∥Σ∥op Since ∥Σ∥2F ≤ ∥Σ∥op Tr(Σ) = d µ ∥Σ∥op , we obtain (56). Lemma 9 (Stability under random variance). Let Assumption 2 hold, and let V = w⊤ Σw, N (0, Id /d), independent of ξ, with µ = d1 Tr(Σ). Define, for s ≥ 0,
w∼
Ss (t) := EG∼N (0,s) [σ ′ (t + G)]. Then there exists a constant δ > 0 (depending only on σ, N , and µ + vsum /d) such that, on the event |V − µ| ≤ δ, the following holds simultaneously for all k ≤ N : (k)
|SV +vsum /d (0)| ≥
c0 , 2
d µ δ2 with probability at least 1 − 2 exp − c ∥Σ∥ . 2 op µ Proof. For s ≥ 0 and k ≥ 0, differentiation under the expectation is justified since σ is a polynomial, which gives Ss(k) (0) = EG∼N (0,s) [σ (k+1) (G)]. (k)
Hence Ss (0) depends on s only through the Gaussian variance. (k) The map s 7→ Ss (0) is continuously differentiable, and for Gs ∼ N (0, s) we have the identity h i 1 d (k) Ss (0) = EGs σ (k+3) (Gs ) . ds 2 Since σ (k+3) is a polynomial, the right-hand side is finite and continuous in s. Therefore, for each k ≤ N , there exists a constant Lk < ∞ such that Ss(k) (0) − Ss(k) (0) ≤ Lk |s − s0 |, 0
s0 := µ + vsum /d,
for all s in a neighborhood of s0 . Let L := maxk≤N Lk . (k) c0 By Assumption 2, Ss0 (0) ≥ c0 for all k ≤ N . Choosing δ := 2L , we obtain that whenever |s − s0 | ≤ δ, c0 Ss(k) (0) ≥ Ss(k) (0) − L|s − s0 | ≥ , ∀ k ≤ N. 0 2 Finally, since s = V + vsum /d and s0 = µ + vsum /d, the condition |s − s0 | ≤ δ is equivalent to |V − µ| ≤ δ. The probability bound follows directly from Lemma 8. p Remark 4. Lemma 8 above shows that V concentrates at scale (C/(dµ)) around µ. In particular, for fixed C and µ bounded away from 0, the dependence on the specific initialization vanishes exponentially fast in d.
27
B.1
Proof of Lemma 1
Fix j ∈ Ci . Write ak := sc(k) wk and Xk := ξk − mk . Define µ :=
d X
ak mk = s⊤ ∆m (w),
U :=
k=1
d X
ak Xk ,
U−j :=
k=1
X
ak Xk .
k̸=j
P By independence of the ξk , conditional on s, Var(U | s) = k wk2 σk2 = V 2 , which does not depend on s. Let Ew denote the high-probability event of Lemma 7 together with the chi-squared concentration estimate ∥∆m (w)∥22 ≤ 2 vsum /d ≤ C (which holds with probability 1 − e−cd since E∥∆m (w)∥22 = vsum /d = Od (1)). On Ew : p √ max |wk | ≤ C log d/ d, ∥w∥2 ≤ 3/2, ∥∆m (w)∥2 ≤ C ′ , V 2 ≍ 1. k
In particular, for every s ∈ {±1}N , |µ| = |s⊤ ∆m (w)| ≤
√
√ N ∥∆m (w)∥2 ≤ C N = Od (1),
since N = Od (1). Throughout the proof we work on Ew and the constants C1 , C2 , C3 depend only on σ, on the sub-Gaussian parameters and third moments of ξ, and on N — but not on d. Writing ξj = mj + Xj , Eξ [ξj σ ′ (µ + U )] = mj Eξ [σ ′ (µ + U )] + Eξ [Xj σ ′ (µ + U )]. (57) P P −1 r Since σ ′ is a polynomial of degree P − 1, expand σ ′ (µ + x) = r=0 cr (µ) x with |cr (µ)| ≤ P −1−r C(1 + |µ| ) = Od (1) on Ew . Then Eξ [σ ′ (µ + U )] − EG [σ ′ (µ + G)] =
P −1 X
cr (µ) E[U r ] − E[Gr ] ,
r=3
since the r = 0, 1, 2 moments match. By the standard moment-comparison bound for sums of independent centered sub-Gaussian variables, √ X log d E[U r ] − E[Gr ] ≤ Cr |ak |3 E|Xk |3 ≤ C √ , d k √ √ P using k |ak |3 ≤ maxk |ak | · ∥w∥22 ≤ C log d/ d on Ew and bounded third moments. Summing over r, √ log d Eξ [σ ′ (µ + U )] − EG [σ ′ (µ + G)] ≤ Cσ εd , εd := √ . (58) d Since ξj ⊥ ξ−j , we have Xj ⊥ U−j . Taylor-expand σ ′ around µ + U−j : σ ′ (µ + U−j + aj Xj ) = σ ′ (µ + U−j ) + aj Xj σ ′′ (µ + U−j ) +
(aj Xj )2 (3) σ (ζj ), 2
for some intermediate point ζj between µ + U−j and µ + U−j + aj Xj . Multiplying by Xj and taking expectation over ξ, then using Xj ⊥ U−j and E[Xj ] = 0, Eξ [Xj σ ′ (µ + U )] = aj σj2 E[σ ′′ (µ + U−j )] +
a2j 3 (3) (ζj ) . 2 E Xj σ
Each σ (k) (µ + U−j ) is a polynomial of degree P − k in U−j with coefficients that are Od (1) on Ew 2 (since |µ| = Od (1)). As U−j is sub-Gaussian with variance V−j ≍ 1, all its moments are Od (1). Hence E[σ ′′ (µ + U−j )] ≤ Cσ ,
E[Xj3 σ (3) (ζj )] ≤ (EXj6 )1/2 (Eσ (3) (ζj )2 )1/2 ≤ Cσ′ , 28
where the second bound uses Cauchy–Schwarz, bounded sub-Gaussian moments of Xj , and the fact that ζj is sub-Gaussian (being bounded by two sub-Gaussian random variables). Therefore Eξ [Xj σ ′ (µ + U )] ≤ C1 |wj | + C2 wj2 .
(59)
Multiplying (57) by f (s) si and averaging over s, G(wj ) = mj Es [f (s) si Eξ [σ ′ (µ + U )]] + Es [f (s) si Eξ [Xj σ ′ (µ + U )]]. Applying (58) inside the first s-expectation and (59) for the second, both bounds being uniform in s ∈ {±1}N on Ew , we conclude G(wj ) = mj αi (w) + ηj ,
C
|ηj | ≤ C1 |wj | + C2 wj2 + C3 εd .
Proof of Theorem 2
The proof of Theorem 2 follows the same path as the proof of Theorem 1. The main differences are in the computation of the initial population gradient (derived in Lemma 1, which is replaced by Lemma 10 below), and in the construction of the certificate (Lemma 5, replaced by Lemma 13). Moreover, the non-degeneracy (derived in the previous proof in Lemma 3) is no longer derived from the randomness of the initialization, and it is instead assumed (Assumption 3). Initial population gradient. In this Section, we consider learning with a 2-layer network with ReLU Pn activation: NN(x; θt ) = h=1 ah ReLU((wht )⊤ x + bh ). We assume that the target function satisfies √ (0) (0) Assumption 3. We initialize whj = 1/ d (deterministically) and ah = τ , with τ > 0, specified later. Similarly to the previous proof, we initialize the bias neurons to 0 and keep them frozen during the iid first phase of training. Before the second√phase of training, we draw bh ∼ Unif[−A, A], for a fixed 0 A > 0, specified later. Because wh = 1/ d for all h ∈ [n], all hidden neurons behave analogously. Thus, in the following we consider a single vector w ∈ Rd , dropping the subscript h and P upscript 0 for simplicity. For simplicity, we assume N to be odd. Let ξ¯ = (ξ¯i )i∈[N ] , where ξ¯i = k1 j∈Ci ξj . For j ∈ Ci , let us compute the initial population gradient. Ḡ(wj ) : = Ex,y [yxj σ ′ (wx)] h √ i ¯ = Es,ξ f (s)si ξj σ ′ sT ξk/ d h h √ ii ¯ d . = Es f (s)si Eξ|s ξj σ ′ sT ξk/
(60) (61) (62)
Let us consider the inner expectation. Recall that we assumed ReLU activation, thus σ ′ (x) = 1(x ≥ 0). Thus, Eξ|s [ξj 1 sT ξ¯ ≥ 0 ] = (1 − δ)p1 − δp−1 , (63) where, for b ∈ {±1}, pb := P(sT ξ¯ ≥ 0 | ξj = b). We make use of the following Lemma. Lemma 10. For b ∈ {±1}, pb − 1((1 − 2δ)sT 1 > 0) ≤ exp(−CkN (1 − 2δ)2 s̄2 ), where C > 0 is a constant and s̄ := N1
PN
i=1 si .
29
(64)
P Proof of Lemma 10. First, let us write sT ξ¯ = si ξj + S−j , where S−ij := (l,h)̸=(i,j) sl ξh . Let µ := k(1 − 2δ)sT 1 − si (1 − 2δ) denote the mean of S−ij . Note that pb = P(S−ij ≥ −b · si ). Moreover, because we assumed N to be odd, sT 1 ̸= 0. Assume first that (1 − 2δ)sT 1 < 0. Then, by Hoeffding’s inequality, (65)
pb = P(S−ij ≥ −b · si ) ≤ P(S−ij − µ ≥ −µ/2) kN (1 − 2δ)2 s̄2 ≤ exp − . 8
(66) (67)
On the other hand, if (1 − 2δ)sT 1 > 0, (68)
pb = 1 − P(S−ij ≤ −b · si ) ≤ P(µ − S−ij ≥ µ/2) kN (1 − 2δ)2 s̄2 ≤ exp − . 8
(69) (70)
Thus, applying Lemma 10 to (63), Eξ|s [ξj 1 sT ξ¯ ≥ 0 ] = (1 − 2δ)1((1 − 2δ)sT 1 > 0) + ηs ,
(71)
where ηs is an error term such that |ηs | ≤ exp(−CkN (1 − 2δ)2 s̄2 ) ≤ exp(−Ck(1 − 2δ)2 /N ), since |s̄| ≥ 1/N , because N is odd. Note that since by assumption d(1 − 2δ)2 = Ω( log(d) 1−2δ ), we have |ηs | = Ω(d−c ), for some c > 0. Plugging this into the outer expectation over s, we obtain, using sgn(1 − 2δ) = 1, since δ ∈ (0, 1/2): Ḡ(wj ) = Es f (s)si Eξ|s [ξj 1 sT ξ¯ ≥ 0 ] = (1 − 2δ)E[f (s)si 1((1 − 2δ)sT 1 > 0)] + η ′ 1 = (1 − 2δ) (E[f (s)si ] + sgn(1 − 2δ)E[f (s)si Maj(s)]) + η ′ 2 X 1 − 2δ ˆ 1 X ˆ ˆ ˆ = f ({i}) + f (T )Maj(|T | − 1) + fˆ(T )Maj(|T | + 1) + η ′ 2 2 T :i̸∈T
T :i∈T
X 1 = (1 − 2δ) fˆ(T )cT,i + η ′ 2 T ⊆[N ]
= (1 − 2δ)αi + η ′ , where, |η ′ | = |Es [f (s)si η] ≤ Es [|f (s)|] · sups |ηs |, cT,i are defined in (3) and where we defined αi := P 1 ˆ T ⊆[N ] f (T )cT,i . 2 Initial estimated gradient. Let us now consider the gradient estimated through B i.i.d. samples. B
GB (w) : =
1 X b ′ y σ (wxb )xb . B
(72)
b=1
If B is large enough, the estimated gradient is close to the population gradient Ḡ(wh0 ), as formalized by the following Lemma.
30
2 Lemma 11. Let us denote by Mf,N = maxs∈{±1}N |f (s)|. If the batch size B ≥ C log(dn)Mf,N /2ζ 2 , −C ∗ ∗ for ζ, C > 0, then with probability 1 − d , for C > 0, for all h ∈ [n],
(73)
∥GB (wh0 ) − Ḡ(wh0 )∥∞ ≤ ζ. Proof. Let us notice that for a j ∈ [d], and for σ = ReLU activation,
(74)
|GB (wh0 )j | ≤ Mf,N . ∗
By Hoeffding’s inequality, there exists C ∗ > 0 such that with probability 1 − d−C /dn, (75)
GB (wh0 )j − Ḡ(wh0 )j ≤ ζ. The result follows by union bound.
√ 2 Lemma 12. Let z := w1 x. Assume |1 − 2δ| = Ω(log(d)2 / d). Assume B ≥ C log(dn)Mf,N /2ζ 2 , with 2
log(d) 1−2δ ζ = θ( log(d) ). Assume the initialization scale τ = θ( (1−2δ) log(d)d ), the learning rate γ = θ( (1−2δ)4 ) and N = θ(1). Then, with probability 1 − od (1), N
Ex |z −
1 X αi si | = O N i=1
1 log(d)
(76)
.
Proof. For j ∈ Ci , we have: wj1 = wj0 + γτ Ḡ(w0 ) + γτ (G(w0 ) − Ḡ(w0 )) + O γτ 2 dMf,N √ (a) = τ γ(1 − 2δ)αi + γτ (η ′ + ζ + 1/ d + O(τ dMf,N )) (b) = τ γ (1 − 2δ)αi + η ′′ ,
(77) (78) (79)
where in (a) we √ used the assumption on the initialization and Lemma 11, and in (b) we denote η ′′ := η ′ + ζ + 1/ d + O(τ dMf,N ), and note that by the assumptions on the batch size, learning rate and initialization scale: n |1 − 2δ| o |1 − 2δ| η ′′ = O max , d−1/2 =O . (80) log(d) log(d) Let Eξ be the event on which, for all i ∈ [N ], k X
ξil − k(1 − 2δ) ≤
√ k log d.
(81)
l=1
By Hoeffding’s inequality and a union bound over i, P(Eξ ) ≥ 1 − d−C for some constant C > 0. On Eξ , using the expression for wj1 derived above, we have N X X z = (w1 )⊤ x = γτ (1 − 2δ) αi si ξj + d η ′′ . i=1
Write
X
ξj = k(1 − 2δ) + ∆i ,
j∈Ci
31
j∈Ci
|∆i | ≤
√ k log d,
(82)
which holds on Eξ . Substituting gives N N X X z = γτ (1 − 2δ)2 k αi si + (1 − 2δ) αi si ∆i + d η ′′ i=1 (a)
= c0 ·
(83)
i=1
N N X 1 X αi si + γτ (1 − 2δ) αi si ∆i + γτ d η ′′ , N i=1 i=1
(84)
for some c0 > 0, where in (a) we used the assumptions on γ, τ and k = d/N . We now bound the remainder terms. γτ (1 − 2δ)
N X
√ αi si ∆i ≤ O γτ |1 − 2δ| d log d
(85)
i=1
log(d) ≤O √ d|1 − 2δ|
,
(86)
√ where the second inequality uses γτ = Θ((1−2δ)−2 d−1 ). Using the assumption |1−2δ| = Ω(log(d)2 / d), this is O(1/ log d). Next, recall the bound η ′′ = O(|1 − 2δ|/ log(d), which implies |1 − 2δ| |1 − 2δ| ′′ |γτ d η | = O γτ d =O . (87) log d log(d) Since |1 − 2δ| ≤ 1, this yields |γτ d η ′′ | = O(1/ log d). Combining the above bounds, on the event Eξ we obtain z−
N 1 X 1 αi si = O . N i=1 log d
Since P(E) = 1 − od (1), taking expectation over a fresh draw of x gives N 1 1 X αi si = O . Ex z − N i=1 log d
iid
Lemma 13. Let f satisfy the non-degeneracy Assumption 3 and let the biases be sampled as bh ∼ p Unif[−A, A], with A := 2N 2N/2 Var(f ), for h ∈ [n]. Then, there exist a constant ∆ > 0 such that if 2 n ≥ C log(d)A , there exists a vector a∗ ∈ Rn such that ∥a∗ ∥∞ ≤ ∆ Mf,N and ∆N f (s) =
n X
a∗h σ(sT α + bh ),
(88)
h=1
for all s ∈ {±1}N . Proof. Let us first enumerate the vectors sl ∈ {±1}N , l ∈ [2N ], so that the real numbers vl =
N X
sl,i αi
(89)
i=1
are in strictly P increasing order, i.e. v1 < v2 < ... < v2N . This is possible since by Assumption 3 the map s 7→ i si αi is injective on {±1}N . 32
Let ∆ = min1≤l<2N (vl+1 − vl ) > 0. Fix any δ ∈ (0, ∆/2) and consider the interval I = [−v2N − δ, −v1 + δ]. Note that I ⊆ [−A, A]. Partition I into 2N disjoint sub-intervals of width ∆/2: Ij = [−vj + δ/2, −vj + δ]. Now, if n biases are sampled uniformly from [−A, A], for each j ∈ [2N ], the probability that none of the n samples falls into Ij is n ∆/2 ∆ 1− ≤ exp − n . (90) 2A 2A By union bound, the probability that every interval Ij is hit at least once is ′ ∆ N 1 − 2 exp − n ≥ 1 − d−C , 2A
(91)
if n ≥ C log(d)A , for suitable constants C, C ′ > 0. For each j ∈ [2N ], choose one bhj such that bhj ∈ Ij , ∆N N N and denote by J = {hj : j ∈ [2N ]} ⊂ [n] such set. Let M ∈ R2 ×2 be such that Ml,j : = ReLU(vl + bj )
(92)
= vl − vj + ϵ,
(93)
where ϵ ∈ (0, ∆/2). Note that, if l < j, then vl − vj < ∆, thus Ml,j = 0, while if l ≥ j, vl − vj ≥ 0, thus Ml,j > 0. Thus, M is lower triangular, with strictly positive entries on the diagonal, and thus N invertible. Let c∗ ∈ R2 be such that (94)
c∗ = F M −1 , N
where F ∈ R2 is a vector such that Fl = f (sl ). Then, a∗ ∈ Rn defined by a∗hj = c∗j 1(hj ∈ J) satisfies (88). Moreover, ∥a∗ ∥∞ ≤ ∥M −1 ∥∞ ∥F ∥∞ ≤
2 · Mf,N . ∆
(95)
From the previous Lemmas, it follows that with probability 1 − d−C , for some C > 0, Ex,y y −
n X
(96)
a∗h σ(w1 x + bh )
h=1
≤ Ex,y y −
n X
a∗h σ(αs + bh ) + Ex,y
h=1
n X
a∗h (σ(αs + bh ) − σ(w1 x + bh ))
≤ n∥a∗ ∥∞ Ex σ(αs + bh ) − σ(w1 x + bh ) ∗
(97)
h=1
1
≤ n∥a ∥∞ Ex αs − w x 1 2N log d2N 1 = O . ≤O ∆N ∆ log(d)2 N ∆2 log(d)
(98) (99) (100)
We then use the well known result of Theorem 3. In particular, we note that ∇at NN(x; (w1 , at , b)) ≤ √ √ N/2 √ √ log(d)2 log(d) C n, for some C > 0. Thus, choosing B = n · ∥a∗ ∥∞ = O ∆3/2 √N , ξ = O (∆N )1/2 , 2 N 2 ϵ∆N T2 = O log(d) and γ2 = θ( log(d) ). ϵ 2 ∆4 N 2
33
D
Clustering Algorithms
Comparison with classical clustering methods. Let us briefly compare Theorem 2 with the guarantees of two classical unsupervised clustering methods. In the homogeneous BSC model with equal cluster sizes |Ci | = k = d/N (Example 2), the population covariance matrix takes the form Σ = (1 − m2 )Id + m2
N X
1Ci 1⊤ Ci ,
i=1
where m = 1 − 2δ is the signal mean. Thus, Σ consists of the identity plus N spikes of size m2 (k − 1) ≍ (1 − 2δ)2 (d/N ). Spectral clustering succeeds in weakly recovering the clusters once these spikes cross the Baik-Ben Arous-Péché (BBP) threshold, namely when d (1 − 2δ)2 ≳ c N for a suitable constant c [BBAP05, BGN11]. If Nd (1 − 2δ)2 = ωd (1) (i.e. if vsum = ωd (1)), recovery becomes asymptotically exact, see details below. Our layerwise-SGD method attains a vanishing error whenever vsum = d(1 − 2δ)2 = Ωd (log(d)2 ), which (for fixed N ) matches the spectral detectability up to logarithmic terms. We remark that spectral recovery uses B = θd (d) samples, as opposed to the B = Ω(log2 (d)) required by layerwise SGD. PB (b) (b) On the other hand, a simple covariance–thresholding baseline, where one estimates Cij = B1 b=1 xi xj and threshold at 12 (1 − 2δ)2 , succeeds with high probability using B ≳ (log d)/(1 − 2δ)4 samples (this can be shown by simple Hoeffding and union bounds, see below), which is comparable to layerwise SGD (up to (1 − 2δ)2 factors). Spectral Clustering. Let x1 , . . . , xm ∈ Rd be m ∼ Θ(d) independent observations of a d-dimensional random vector. Assume the covariance matrix Σ ∈ Rd×d of the distribution generating these observations takes the form Σ = Id + ∆, where ∆ is a rank-N symmetric perturbation with nonzero eigenvalues λ1 , . . . , λN . This setting models r clusters in the data, each inducing a structure in the covariance. Under this model, it is possible to recover the eigenvalues λ1 , . . . , λN , as well as their associated eigendirections, provided that |λi | ≥ θ, for a threshold θ = Θ(1). In the regime where |λi | ≫ 1, recovery becomes asymptotically exact. This justifies the use of spectral methods such as PCA for recovering low-rank structure in high-dimensional data, under appropriate signal-to-noise conditions. Now consider our model where the data consist of N clusters, each of size d/N , and the covariance structure satisfies: • For coordinates i, j within the same cluster: Cov(xi , xj ) = (1 − 2δ)2 , • For coordinates in different clusters: Cov(xi , xj ) = 0. In this case, the rank-N perturbation ∆ has eigenvalues of order λi ∼
d (1 − 2δ)2 . N 34
To ensure that the perturbation is detectable, it suffices that λi ≫ 1
d (1 − 2δ)2 ≫ 1 N
⇐⇒
This translates to the condition δ−
⇐⇒
(1 − 2δ)2 ≫
N . d
1 1 ≫√ . 2 d
‘Naive’ Clustering. Let x1 , . . . , xm ∈ Rd be m independent observations of a d-dimensional random vectors sampled from a clustered distribution such that coordinates within each cluster have covariance (1 − 2δ)2 and coordinates in different clusters are independent. Then, one could think of cluster the coordinates by estimating the pairwise covariance, i.e.: T
Ci,j :=
1X t t xx . T t=1 i j
(101)
By Hoeffding’s inequality: P |Ci,j − (1 − 2δ)2 1(c(i) = c(j))| ≥ (1 − 2δ)2 /2 ≤ 2 exp(−(1 − 2δ)4 T /8) ≤ 1/d2 ,
(102)
if T ≥ 16 log(d)/(1 − 2δ)4 .
E
Deferred Proofs
Lemma 14 (Quantitative nondegeneracy under random activation coefficients). Fix a > 0 (e.g. a = vsum ) and µ > 0. Expand σ ′ in the Hermite basis aligned with variance a, σ ′ (x) =
M X
√ n/2 He(a) Hn x/ a , n (x) := a
αn He(a) n (x),
n=0
where Hn are the probabilists’ Hermite polynomials. Assume the coefficients are independent Gaussians i.i.d. α0 , . . . , αM ∼ N (0, 1) and M ≥ N . Let ψµ (x) := EG∼N (0,µ) σ ′ (x + G) ,
βk (µ) :=
√ 1 EX∼N (0,1) ψµ ( a X) Hk (X) , k!
0 ≤ k ≤ N.
Then each βk (µ) is (marginally) Gaussian, βk (µ) ∼ N (0, s2k ), with variance ⌊(M −k)/2⌋
s2k
k
= a
X r=0
"
(k + 2r)! (µ/2)r k! r!
#2 (103)
> 0.
In particular, for any τ > 0, r 2 τ Pr min |βk (µ)| ≤ τ ≤ (N + 1) , 0≤k≤N π smin
smin := min sk .
Equivalently, choosing τ = θ smin /(N + 1) with θ ∈ (0, 1) yields q Pr min |βk (µ)| > Nθ+1 smin ≥ 1 − π2 θ. 0≤k≤N
35
0≤k≤N
(104)
2
Proof. The Gaussian smoothing operator is the heat semigroup Wµ := e(µ/2)∂x , which acts lowertriangularly on the Hermite basis: (a) Wµ Hek+2r
=
r X
(k + 2r)! (µ/2)u (a) Hek+2r−2u . (k + 2r − 2u)! u! u=0
(a) √ 1 Projecting onto order k and noting that k! E[Hek ( aX)Hk (X)] = ak/2 and is zero for other orders, we obtain the explicit linear representation ⌊(M −k)/2⌋
βk (µ) = a
k/2
X r=0
αk+2r
(k + 2r)! (µ/2)r , k! r!
0 ≤ k ≤ N.
Since the αn are independent N (0, 1), each βk (µ) is Gaussian with mean 0 and variance given by the sum of squares of the coefficients, which is precisely (103). In particular s2k > 0 because the r = 0 term contributes ak . ´τ 2 2 1 For the small-ball bound, recall that if Z ∼ N (0, s2 ), then Pr(|Z| ≤ τ ) = −τ √2πs e−x /(2s ) dx ≤ q 2 τ 2 √2τ = π s . Applying this to each βk (µ) ∼ N (0, sk ) and taking a union bound over k = 0, . . . , N 2πs yields (104). Remark 5. The mapping from the Hermite coefficients (α0 , . . . , αM ) of σ ′ to (β0 (µ), . . . , βN (µ)) is linear and upper-triangular in parity blocks with strictly positive diagonal; thus s2k ≥ ak and smin ≥ 1. The formula (103) makes the dependence on (µ, a, M ) explicit and can be used to pick a quantitative margin τ with the desired probability level.
F
Fourier–Walsh expansion of Majority.
Let MajN : {±1}N → {±1} denote the majority function, MajN (s) := sgn
N X
! si
,
i=1
with can be broken arbitrarily (e.g., MajN (0) = 1 when N is even). The Fourier–Walsh expansion of MajN is given by X Y ˆ N (S) χS (s), MajN (s) = Maj χk (s) := si , i∈S
S⊆[N ]
where the Fourier–Walsh coefficients are defined as ˆ N (S) := Es∼Unif({±1}N ) [MajN (s) χS (s)] . Maj ˆ N (k) depends only on |S|, and Maj ˆ N (S) = 0 whenever |S| is even. For odd |S| = k, By symmetry, Maj the coefficients admit the explicit formula k−1 −1 2 N k−1 N −1 k−1 2 2 d N (S) = (−1) 2 Maj · , −1 2k−1 2N N k−1 Precise asymptotics and proofs can be found in [O’D14], Chapter 5.
G
Experiments Details
In this section, we provide more details about the experiments on real data shown in Section 6. 36
G.1
Architecture and Loss Function
Architecture. We use a two-layer fully-connected neural network with ReLU activations, f (x) = W2 ReLU(W1 x + b1 ) + b2 , where W1 ∈ RH×d , W2 ∈ RC×H , H = 128 is the hidden width, d is the input dimension, and C is the number of classes. Weights are initialized as W1 ∼ N (0, 1/d)d×H and W2 ∼ N (0, 1/H)H×C ; biases are initialized to zero. Loss function. We minimise the cross-entropy loss with a softmax output layer. Gradients are clipped to [−5, 5] component-wise at every step.
G.2
Datasets and Preprocessing
GSE96583 (single-cell RNA-seq). We use batch 2 of the GSE96583 dataset [KST+ 18], restricting to control (unstimulated) singlet cells. Cells are assigned to three merged classes: B cells, Monocytes (CD14+ and FCGR3A+ merged), and T/NK (CD4 T, CD8 T, and NK cells merged); Dendritic cells and Megakaryocytes are excluded. After filtering cells with fewer than 500 UMI counts and genes expressed in fewer than 25 cells, we obtain n ≈ 11,990 cells over ∼ 10,500 genes. Raw counts are library-size normalised to 104 counts per cell and log1+x transformed. For each value of d, the top-d most variable genes are selected by dispersion (variance/mean), using no class label information.
G.3
Training Procedure
We use the Adam optimiser [KB14] with learning rate 5 × 10−4 , β1 = 0.9, β2 = 0.999, ε = 10−8 , batch size 256, and 100 epochs. Train/test split and seeds. For each experiment, we perform a stratified 80/20 train/test split. The training set is then subsampled to n examples (stratified by class) for each value of n considered. All results are averaged over 5 independent random seeds; shaded bands report ±1 standard deviation.
37