Symmetric Single Index Learning
Authors: Aaron Zweig
Joan Bruna
Venue: ICLR 2024 poster (2024)
Keywords: single-index
symmetric
permutation invariant
gradient flow
learning
Primary area: learning theory
TL;DR: Gradient flow can provably learn a single index model in the feature space of symmetric polynomials

Abstract: Few neural architectures lend themselves to provable learning with gradient based methods. One popular model is the single-index model, in which labels are produced by composing an unknown linear projection with a possibly unknown scalar link function. Learning this model with SGD is relatively well-understood, whereby the so-called information exponent of the link function governs a polynomial sample complexity rate.  However, extending this analysis to deeper or more complicated architectures remains challenging.

In this work, we consider single index learning in the setting of symmetric neural networks.  Under analytic assumptions on the activation and maximum degree assumptions on the link function, we prove that gradient flow recovers the hidden planted direction, represented as a finitely supported vector in the feature space of power sum polynomials.  We characterize a notion of information exponent adapted to our setting that controls the efficiency of learning.

Published as a conference paper at ICLR 2024
SYMMETRIC SINGLE INDEX LEARNING
Aaron Zweig
Joan Bruna
Courant Institute of Mathematical Sciences New York University
Courant Institute of Mathematical Sciences
Center for Data Science
New York, NY 10012, USA a2 12831@nyu.ee edu
New York University
New York, NY 10012, USA
bruna @ cims. nyu. edu
ABSTRACT
Few neural architectures lend themselves to provable learning with gradient based methods. One popular model is the single-inder model, in which labels produced by composing an unknown linear projection with possibly unknown scalar link function. Learning this model with SGD is rel- atively well-understood, whereby the so-called information exponent of the link function governs polynomial sample complexity rate. However, extending this analysis to deeper or more complicated architectures remains challenging.
In this work, we consider single index learning in the setting of symmetric neural networks. Under analytic assumptions on the activation and maximum degree assumptions the link function, prove that gradient flow recovers the hidden planted direction, represented as a finitely supported vector in the feature space of power sum polynomials. We characterize a notion of information exponent adapted to our setting that controls the efficiency of learning.
1 INTRODUCTION
Quantifying the advantage of neural networks over simpler learning systems remains a primary question in deep learning theory. Specifically, understanding their ability to discover relevant low-dimensional features out of high- dimensional imputs is particularly important topic of study, One facet the challenge is explicitly characterizing the evolution of neural network weights through gradient-based methods, owing to the nonconvexity of the optimization landscape.
The single index setting. long studied in economics and biostatistics (Radchenko, 2015) offers the simplest setting where non-linear feature learning can be characterized explicitly. In this setting. functions of the form x + f((z, 8*)) where 0" E S4-1 represents hidden direction in high-dimensional space, and f certain non-linear link function, are learned via a student with an identical architecture x + [((x, 0)), under certain data distribution assumptions, such as Gaussian data. Gradient flow and gradient descent (Ychudai & Ohad [220)] Arous et al. 2021] Dudeja & Hsu 2018) in this setting can be analyzed by reducing the ligh-dimensional dynamics of & to dimension free dynamics of appropriate summary statistics, given in this case by the scalar correlation (0, (0,0").
The efficiency of gradient methods in this setting, measured either in continuous time or independent samples, is controlled by two main properties. First, the correlation initialization, which typically scales as va for standard
assumptions. Second, the information exponent s of f (Arous et al. 2021 Dudeja & Hsu 2018 Bietti et al. 2022) Damian et al. 2023 2022 Abbe et al. 2023), which measures the number of etfective vanishing moments of the link function - leading to a sample complexity of the form O(d=-1) for generic values of 8.
While this basic setup has been extended along certain directions, e.g. relaxing the structure the input data distribu- tion Yehudai & Ohad 22020); Bruna et al. (2023), considering the multi-index counterpart Damian et al. (2022) y Abbe et (2022) 2023); Arnaboldi et al. (2023), or learning the link function with cemi-parametric methods Bietti al. (2022): ; Mahankali et al. (2023), they are all fundamentally associated with fully-connected shallow neural networks. Such architecture, for all its rich mathematical structure, also comes with important shortcomings. In particular, it is unable to account for predefined symmetries in the target function that the learner wishes to exploit. This requires specialized neural architectures enforcing particular invariances, setting up novel technical challenges to carry out the program outlined above.
In this work, we consider arguably the easiest form of symmetry, given by permutation invariance. The primary architecture for this invariance is DeepSets (Zaheer et al. 2017), which is necessarily three layers by definition and

---

Published as a conference paper at ICLR 2024
therefore not a simple extension of the two layer setting. In order to quantify the notion of 'symmetric' feature learning this setting, we introduce a symmetric single index target, and analyze the ability of gradient descent over DeepSets architecture to recover it. Under appropriate assumptions on the model, initialization and data distribution, we combine the previous analyses with tools from symmetric polynomial theory to characterize the dynamies of this learning problem. Our primary theorem is proof of efficient learning under gradient flow, with explicit polynomial convergence rates controlled by an analogue of information exponent adapted to the symmetric setting. Combined with other contemporary works, this result solidifies the remarkable ability of gradient descent to perform feature learning under a variety of high-dimensional learning problems.
2 SETUP
2.1 NOTATION
For E C, we will use Z to denote the complex conjugate, with the notation =" aaways being reserved to denote special value rather than an operation. For complex matrices we will use A to denote the conjugate transpose. The standard inner product on CV is written as (,,)), whereas inner products (?(r) spaces for some probability measure y will be written as (,.) Furthermore, for h vector and p(z) a vector valued function, we will use (h, p) shorthand for the notation (h))
2.2 REGRESSION SETTING AND TEACHER FUNCTION
We consider a typical regression setting, where given samples (2r, y) E X' x C with y = F(x), we seek to learn function F. with parameter E C by minimizing some expected loss Ezzy (L(F() Note that we consider complex-values inputs and parameters because they greatly simplify the symmetric setting (see Proposi- tion 2.3 hence we will also assume X' CN Both F and Fw will be permutation invariant functions, meaning that F( EN) for any permutation F( Fx
N1, + N1.
Typically the single index setting assumes that the trained architecture will exactly match the true architecture (e.g. as in Arous et al. ()02T)), but below we will see why it's necessary to consider separate architectures. For that reason, we'll consider separately defining the teacher F and the student P.p.
The first ingredient are the power sum polynomials:
Definition 2.1. For l E N and x E CN, the normalized powersum polynomial is defined as
N
(a(x)= =
x
==1
Let p(z) [pa (xx), p2(x)....] be infinite dimensional vector of powersums, and consider fixed vector h E Co of unit norm. Then our teacher function F will be of the form
F:
X - C
(1)
F(x) = f((h",((z))
(2)
for some scalar link function C C. F may thus be understood as single-index function in the feature space of powersum polynomials.
2.3 DEEPSETS STUDENT FUNCTION
Let remind the typical structure of DeepSets network (Zaheer al. 2017). where for some maps $ X + CM and p: CM + C, the standard DeepSets architecture is of the form:
z p (I),...... P (z))
(3)
The essential restriction is that & is permutation invariant mapping, typically of the form (m = Σn p. (an) for some map Φm C - C. In order to parameterize our student network DeepSets model, we will make the simplest possible choices, while preserving its non-linear essence. To define our student network, we consider the symmetric embedding as a one-layer neural network with no bias terms:
N
T (z) Σ o (am
(4)
nml
2

---

Published as a conference paper at ICLR 2024
for i..dd complex weights sampled uniformly from the complex circle am ~ S and some activation a C - C. And
given some link function g C - C, we'll consider the mapping p as:
Pw (.)
(5)
where w E CM are our trainable weights, Putting all together, our student network thus becomes
Fw X C + Fm(x) :- (w, (())))
(6)
Explicitly, this corresponds to DeepSets network with only trainable vector w, while the first layer weights {am LM and the third layer weights that parameterize the function g are frozen.
The first fact need is that, through simple algebra, the student may be rewritten in the form of a single-index model. Proposition 2.2. There matrix A Coex depending only on the activation a and the frozen weights {am m== such that
s(w, ((x))) (((,,,(())))
2.4 HERMITE-LIKE IDENTITY
In the vanilla single index setting, the key to giving explicit expression for the expected loss (for Gaussian inputs) is a well-known identity of Hermite polynomials [O'Doneel] 2021 Jacobsen 1996). fıg denotes the Hermite polynomial of degree k, this identity takes the form
(hx (6, n)), hill, ,))) Saakl(u, u)*
(8)
where u, R" and Tn is the standard Gaussian distribution on dimensions.
In our setting, as it turns out, one can establish an analogous identity, by considering a different input probability measure, and a bound on the degree of the link function. We will choose our input domain X' (S*) and the imput distribution we will consider is the set of cigenvalues of Haar-distributed unitary matrix in dimension N (Di-) aconis & Shahshahani 1994), or equivalently the squared Vandermonde density over N copies of the complex unit circle Macdonald] 1998). We'll interchangeably use the notation (f (f.g)v.
Proposition 2.3. Consider h, h C with bounded L.2 norm. For exponents k. with k S vN, if is only supported the first VN elements, then:
p),, (h, Susk!(h, h)**
(9)
The crucial feature of this identity is that the assumptions support and bounded degree only apply to (h, p)* with no restrictions on the other term. our learning problem, we can use this property to make these assumptions on the teacher function, while requiring bounds the terms of the student DeepSets architecture.
In order to take advantage of the assumptions on the support of h and the degree in the above proposition, we need to teacher link function f and our true direction ht::
make the following assumptions
Assumption 2.4. The link function f is analytic and only supported on the first VN degree monomials, i.e.
f(z)
(10)
Furthermore, the vector h² is only supported on the first N elements.
Although this assumption is required to apply the orthogonality property for our loss function in the following sections, we note that in principle, including exponentially small terms of higher degree in f or higher index in h should have negligible effect. Moreover, one should interpret this assumption as silently disappearing in the high-dimensional regime N > co. For simplicity, we keep this assumption to make cleaner calculations and leave the issue of these small perturbations to future work.

---

Published as a conference paper at ICLR 2024
2.5 INFORMATION EXPONENT
Because Proposition 2.3 takes inner products monomials, alludes to a very simple characterization of information exponent. Namely:
Definition 2.5. Consider an analytic function C C that can be written in the form
0c
[(z)
(11)
j=0
Then the information exponent is defined as inf 2 1 0j 00}.
Similar to the Gaussian case (Arous et al. 2021 Bietti et al. 2022). the information exponent will control the efficiency of learning. Assuming [Og] some non-negligibn constant, the value of s will be far more important in governing the convergence rate.
2.6 CHOOSING A LEARNABLE LOSS
There are two subtleties to choosing an appropriate loss function. Namely, the necessity of correlational loss (with regularization), and the necessity of choosing the student and teacher link functions to be distinct.
At first glance, is tempting to simply define a loss the form
L(w)= = E F(x) Fu (z) Ez~V [I5(th .(()))) - f((Aw, [(x)))
(12)
However, the Deepsets student model is degree limited, that is the support of Aw is not restricted the first VN terms of the powersum expansion. In other words, expanding this loss will require calculating the term [I,[(Aw, I which will contain high degree terms that cannot be controlled with Proposition 2.3 One could avoid this issue by choosing the activation such that Aw only contains low-index terms, but we want to consider larger classes of activations and enforce fewer restrictions.
One can instead consider a correlational loss. In this case, in order to make the objective have a bounded global minimum, it's necessary to either regularize w, or project at every step SGD, which the strategy taken in Damian al. (2023). In our setting, this projection would correspond to projecting w to the ellipsoid surface ||Aw||| = 1. This projection would require solving an optimization problem at every timestep (Pope) 2008). To avoid this impracticality, we instead consider regularization.
Then with complete knowledge of the link function f. specifically its monomial coefficients, we can now define the correlational loss
i(w) EENV [- Re f((h' p(=)))((A,, p(+)))
layl
||AW| 2
(13)
This loss enjoys benign optimization properties, as shown by the following proposition:
Proposition 2.6. If there exist coprimes k, I with on, on + O, and h² is in the range of A, then L exclusively has global minima all w such that Aw hi.
However, unlike the real case. complex weights causes issues for learning this objective. Namely, this objective can be written as a non-convex polynomial in cos where is the angle of (Aw, h²) in polar coordinates.
Therefore, we consider different choice of student link function that will enable simpler analysis of the dynamics. For the choice of g(z)
Jas Val z", we instead consider the loss:
ar
|a,l
L(w)= = Ex-V [- Re f((h", ((((h,,(())))((())))) p(z)) (Aw,)(()) + 2
[Awl|2
(14)
|ota |wwl|a
= -lx Re (Aw, 1 1="))= }
(15)
We note that Dudeja & Hsu 2018) used similar trick of a correlational loss containing a single orthogonal polynomial in order to simplify the learning landscape. The global minima of this loss, and in fact the dynamics of gradient flow on it, will be explored in the sequel.

---

Published as a conference paper at ICLR 2024
3
RELATED WORK
3.1 SINGLE INDEX LEARNING
The conditions under which single-index model learning is possible have been well-explored in previous literature. The main assumptions that enable provably learning under gradient flow gradient descent are monotonicity of the link function (Kakade et al. 2011; Kalai & Sastry 2009, Shalev-Shwartz et al., 20TO Yehudai & Ohad 2020) or Gaussian imput distribution (Arous et al. 2021). The former assumption essentially corresponds to the setting where the information exponent the optimal sample complexity was achieved Damian et al. (202)). with study of learning when the link function is not known in Bietti al. (002)).
I, as it will have positive correlation with a linear term. Under the latter assumption,
When both assumptions are broken, the conditions on the input distribution of rotation invariance or approximate Gaussianity are nevertheless sufficient for learning guarantees (Bruna et al. (023). But more unusual distributions, especially in the complex domain that is most convenient for symmetric networks, are not well studied.
3.2 S SYMMETRIC NEURAL NETWORKS
The primary model for symmetric neural networks was introduced in Zaheer et al. (2017) as the DeepSets model. There are many similar models that enforce permutation invariance (Qi et al. [2017] (Santoro al. 2017 Lee et al. 22119),
though we focus on DeepSets because of its relationship with the power sum polynomials and orthogonality (Zweig & Bruna 2022). We are not aware of any other works that demonstrate provable learning of symmetric functions under gradient- methods.
4 PROVABLY EFFICIENT RECOVERY WITH GRADIENT FLOW
4.1 DEFINING THE DYNAMICS
The gradient methods considered in Arous et al. (022);; Ben Arous et al. (2022) are analyzed by reducing to a dimension-free dynamical system the so-called summary statistics. For instance, in the vanilla single-index model, the summary statistics reduce to the scalar correlation between the learned weight and the true weight. In our case, we have three variables, owing to the fact that the correlation complex and represented by two scalars, and third variable controlling the norm of the weight since we aren't 'ttt using projection.
Note that although our weight vector w is complex, we still apply regular gradient flow to the pair of weight vectors WR; we where w + iwe. Furthermore, we use the notation V == V.. VaR + iV we. With that in mind, we can summarize the dynamics of our gradient flow the following Theorem
Theorem 4.1. Given a parameter w. consider the summary statistics m = (Aw, E C and v = P Aul where PiL- is projection onto the orthogonal complement of h.. Let the polar decomposition of m be r
Then given the preconditioned gradient flow given by
1
AA'A)-
(A
(16)
s s|a,
the summary statistics obey the following system of ordinary differential equations:
f 8)r s0 (u + 12)0-1
(17)
d
cos s0 (1 - d)ST** 21 cos² s0)
(18)
di 0 28r cos s0 2( + , 12
(19)
where 8 1 1 | and PA is the projection onto the range of A.
The proof is in Appendix D The main technical details come from using Wirtinger calculus to determine how the real and imaginary parts of evolve under the flow. Additionally, the correct preconditioner (intuitive from the linear transform of w)) is crucial for reducing the dynamics to only three summary statistics, and converting to dynamies on cos sO rather than 0 itself simplifies the description of the learning in the next section dramatically.
5

---

Published as a conference paper at ICLR 2024
4.2 PROVABLE LEARNING
These dynamics naturally motivate the question of learning efficiency, measured in convergence rates time in the of gradient flow. Our main result is that, under {am mm= and the initialized weight vector Wo, the efficiency is controlled by the initial correlation with the true direction and the information exponent, just in the Gaussian case.
assumptions the initialization of the frozen weights
Theorem 4.2. Consider a fixed e > O. Suppose the initialization of we and (am )M are such that:
(i) Small correlation and anti-concentration at initialization: 0 < ro 11,
(i) Initial phase condition: cos s00 2 1/2
(ii)) Iniiill magnitude condition for Aw: ||wwll = 1, or equivalently Uo = 1 ro-
(iv) Small Approximation of optimal error: & minn(e/,
Then if we run the gradient flow given in Theorem 4.1 we have e accuracy in the sense that:
TT 1 c, cos s0, 2 c, VT S
(20)
after time T, where depending on the information exponent 8:
(O (log !) T 10 2 To + log
8 = 1,
(21)
8 s> 1.
Remark 4.3. We note that we only recover cos sA rather than guarantee that 0 zs 0, and so the hidden direction is only determined up to scaling by sth root of unity. This limitation is may appear to be issue with the choice of the student link function g, but unavoidable: if the teacher link function [(z) Vsl" z", one can calculate that for any choice of g. L(w) is invariant scaling w by sth root of unity.
4.3 INITIALIZATION GUARANTEES
In order to apply the gradient flow bound proved in Theorem 4.2 only remains to understand when the assumptions initialization are met. Unlike the single-index setting with Gaussian inputs, the initial correlation is not guaranteed to be on the scale of VN' but will depend on the activation function and the random weights in the first layer. Let us introduce the assumptions we'll need:
Assumption 4.4. We assume an analytic activation a(2) = EEs ck2", with the notation a+ == maxisksN lcl|vk
and a_ == min, SKSVT (culV/. k. We further assume:
(i) C = iff ==
(ii) a analytic on the unit disk,
(ili) 1/a_ = Opooyy(N)),
(iv) Σ Mcx12 S <D(NN)
The first two conditions are simply required for the application of Proposition 2.3 as the powersum vector is built out of polynomials induced by the actívation and does include constant term. The second conditions concern the decay of the coefficients of a, in the sense that the decay must start slow but eventually become very rapid. These conditions are necessary mainly for ensuring the Small Approximation of optimal error condition:
Lemma 4.5. Let a satisfy Assumption 4.4 and assume M = O(N3). Consider any unit norm h" E C that is only supported on the first /N elements. If we sample the weights i...d. uniformly on the complex unit circle, am ~ St with probubility 1 2 2exp(-
1- 1P CAMVN))
Lastly, we can choose an initialization scheme for w which handily ensures the remaining assumptions we need to apply Theorem 4.2] The crucial features of are similar to the previous result. Namely, we want the initial correlation To to be non-negligible because this directly controls the runtime of gradient flow. Slow initial decay with fast late decay of the coefficients directly implies that Awo has lot of mass in the first v indices and very little mass past the first N indices. These requirements rule out, say, exp as an analytic activation because the coefficients decay too rapidly.
6

---

Published as a conference paper at ICLR 2024
Lemma 4.6. Suppose w is sampled from standard complex Gaussian on M variables. It follows that if we set
wi w T and use the summary statistics from Theorem 4. then with probability 1/3 - 2 2exp(= and any
h as in Lemma 4.5
(1) 12 To avVM for some universal constant 0,
(i) Cos =0 2 1/2
(iti) 00 rō-
Finally, we consider straightforward choice of that meets Assumption 4.4 so that we can arrive at explicit complexity bound on learning: Corollary 4.7 (Non-asymptotic Rates for Gradient Flow). Consider & = /N
and the specific choice of activation
((z) arctan Ez + Ez arctan
Suppose we initialize w from slandard complex Gaussian in dimension M with M = O(N3), and {am Jm=1 S id. Furthermore, treat and e as constants relative N. Then with probubility 1/3 - 2 accuracy in time
we will recover
[ (log 2) T
8=1
10 10 N7 + +log log
(22)
8 1.
Proof. By Proposition H.5 the activation given in the corollary statement satisfies Assumption 4.4 so we can apply Lemma [4.6 and Lemma 4.5 to satisfy the requirements Theorem 4.2 In particular, the fourth condition is given
by assuming -QVNN < min(e/2,0( "r()) which is true when is constant, and and To are at most polynomial compared to N.
Note that + O(1 and a_ Ω N ) it follows that Ω (777)) ) with probability 1/3 2 2exp(- (SN))))(()) Conditioning on this bound gives the desired bound on the time for accuracy.
Hence, have rate that, for O(1) is not cursed by dimensionality to recover the true hidden direction h² As mentioned above, there are two caveats to this recovery: w is only recovered up to an sth root of unity, and to directly make predictions of the teacher model would require using the teacher link function rather than using the student model directly.
Since this result concerns gradient flow over the population loss, a natural question what barriers exist that stymie the SGD analysis of recent single index papers (Arous et al. /202))/amian Damian et al.) [2023 Bruna et al. (2023). These works treat the convergence gradient flow, and the martingales are shown to be controlled via standard concentration inequalities and careful arguments around stopping times. Applying these tactics to a discretized version of the dynamics given in Theorem mainly runs into an issue during the first phase of training. Unlike in Arous et al. (2021) where the drift dynamics have the correlation monotonically increasing towards 1, at the start our dynamics the correlation magnitude r and the "orthogonal" part of the learned parameter are both decreasing (with high probability over the initialization). Showing that this behavior doesn't 't draw the model towards the saddle point where r = 0 requires showing that v di log r² is positive. It's not clear what quality of bounds the martingale concentration inequalities would provide for this quantity, and we leave for future work if the six stage proof of the dynamics behavior could be successfully discretized.
SGD by a standard drift and martingale argument, where the drift follows the population
decreases meaningfully faster than i.e. showing that d
5 EXPERIMENTS
To study an experimental setup for our setting, we consider the student-teacher setup outlined above with gradient descent. We consider N 25, M 100, and approximate the matrix A by capping the infinite number of rows at
150, which was sufficient for 1 - hA < 0.001 in numerical experiments. For the link function ʃ. we choose its only non-zero monomial coefficients be ag other coefficients at zero.
1 And correspondingly, g simply has as and all
a4 as
v3'
We choose for convenience an activation function such that Akm = ( d We make this choice because, while obeying all the assumptions required in Assumption this choice implies that the action of A on the elementary
7

---

Published as a conference paper at ICLR 2024
basis vectors ej for S S VN is approximately distributed the same. This choice means that P| is less dependent on the choice of h" and therefore reduces the variance our experiments when we choose ht uniformly among unit norm vectors with support on the first vN elements, i.e. uniformly from the complex sphere in degree VN.
Under this setup, we train full gradient descent on 50000 samples from the Vandermonde V distribution under 20000 iterations. The only parameter to be tuned is the learning rate, and we observe over the small grid of [0.001, 0.0025, 0.005] that a learning rate of 0.0025 performs best for the both models in terms of probability of r reaching approximately 1, i.c. strong recovery.
As described in Theorem 4.1 we use preconditioned gradient descent using (AI as the preconditioner, which can be calculated once at the beginning of the algorithm and is an easy alteration to vanilla gradient descent to implement. We use the pscudoinverse for improved stability in calculating this matrix, although doesn't 't introduce stability issues into the updates of our summary statistics, even in the case of gradient descent. Indeed, even if one considers the loss I(w)) under empirical expectation rather than full expectation, the gradient L(w) can still be seen to be written in the form Atu for some vector e. If one preconditions this gradient by (AtA)-!, and observes that the summary statistics and v both depend on Aw rather than w directly, it follows that the gradient update on these statistics is always of the form A(A' A) -A)-"" preconditioner doesn't It introduce exploding gradients.
note that this preconditioner
A PA. so even in the empirical case this
Training of
Training cos(s*theta)
0.75
0.75
0:25
0.25
Vile
Vile
-0.75
-0.35
-0.75
-0.75
2500
5000
7500
10000 Itaratior
12500
15000
27500
20000
2500
5000
7500
20000 teration
12500
25000
17500
20000
Training v
Training of (with L
1.0
1.0
0.75
0.75
0.30
0.30
0.25
0.25
Vale
0.00
0.50
-0.55
-0.25
-0.50
-0.50
-0.75
-0.75
-1.00
-1.00
2500
5000
7500
10000
12500 15000 27500 20000
2505
5000
7500
20000
12500
25000 13500 20000
Iteratter
ttraaiio
Figure 1: The learning trajectory, over ten independent runs, of the three summary statistics in the case of our chosen loss function L, and the trajectory of the r statistic for the more complicated loss function L
6 DISCUSSION
6.1 EXPERIMENTAL RESULTS
The outcomes of our experiments are given in Figure I We observe very high rates of strong recovery using the loss L. For the loss L, we note that often becomes stuck, indicating the model has reached a local minima.
8

---

Published as a conference paper at ICLR 2024
We note that our analysis is somewhat pessimistic, as the experimental gradient descent on L(w) will often achieve near perfect accuracy even if cos =On < 0. This mainly an issue of proof technique: although cos s0 is always increasing under the dynamics, r is necessarily decreasing for as long as cos s0 is negative. It is quite subtle to control whether cos s0 will become positive before becomes extremely small, and the initialization of is the main feature that controls the runtime of the model. However the empirical results suggest that chance of success > 1/2 is possible under a more delicate analysis.
However, the analysis given in the proof of Theorem 4.2 does accurately capture the brief dip in the value of in the initial part of training, when the regularization contributes more to the gradient than the correlation until cos s becomes positive.
Because we can only run experiments on gradient descent rather than gradient flow, we observe the phenomenon of search descent studied in Arous al. [2021) where the increase in the correlation term r is very slow and then abruptly increases. For the model trained with i, we observe that there is much greater likelihood of failure in the recovery, as r appears to become stuck below the optimal value of 1.
6.2 EXTENSIONS
The success of this method of analysis depends heavily on the Hermite-like identity Proposition 2.3 general, many of the existing results analyzing single index models need to assume either Gaussian inputs, or uniformly dis- tributed inputs on the Boolean hypercube (see for example Abbe et al. (2023). In some sense, this works cements the inclusion of the Vandermonde distribution in this set of measures that enable clean analysis. The proof techniques for these three measures are quite disparate, so remains open to determine if there is a wider class of "nice" distributions where gradient dynamics can be successfully analyzed.
Additionally, the success of the multi-layer training in Bietti et al. (002)) Mahankali et al. (2023) suggests that simul- taneously training the frozen first layer weights may not prohibit the convergence analysis. The matrix A depends the first layer weights through a Vandermonde matrix (see X in the proof of Lemma 4.5), and the simple characteri-
zation of the derivative of Vandermonde matrix alludes to further possibilities for clean analysis.
6.3 LIMITATIONS
A first limitation is the focus of this work on complex inputs, analytic activations, and fixed input distribution (namely the squared Vandermonde density). Although complex analytic functions are less commonly studied in the literature. they do still appear in settings like quantum chemistry (Beau & del Campo 2021 Langmann] 2005). Regarding the focus on the Vandermonde distribution, we note this is similar to the vanilla single-index setting in the restriction to Gaussian inputs, under which the theory is particularly powerful, simplest and understanding of non- non-Gaussian data is still nascent.
A second limitation is that this work focuses on input distributions over sets of scalars, whereas typically symmetric neural networks are applied to sets of high-dimensional vectors. Proposition 2.3 does not work out of the box for these settings without high-dimensional analogue of the inner product by with similar orthogonality properties. It is possible to define such an inner products on the so-called multisymmetric powersums with similar orthogonal- ity (Zweig & Bruna, 2022). and we leave future work the question of whether such inner products could grant similar guarantees about the learning dynamics in this more realistic setting.
7 CONCLUSION
In this work we've shown first positive result that quantifies the ability of gradient descent to perform symmetric feature learning, by adapting and extending the tools of two-layer single index models. In essence, this is made possible by 'miracle', namely the fact that certain powersum expansions under the Vandermonde measure enjoy the same semigroup structure as Hermite polynomials under the Gaussian measure (Proposition 2.3 - -leading to a dimension-free summary statistic representation of the loss. Although the resulting dynamics are more intricate than in the Euclidean setting, we are nonetheless able to establish quantitative convergence rates to 'escape the mediocrity" of initialization, recovering the same main ingredients as in previous works Arous al. (202T); [bbe et al.] (2022). driven by the information exponent. To our knowledge, this is the first work to show how learning with gradient based methods necessarily succeeds in this fully non-linear (i.e. not in the NTK regime) setting. Nevertheless, there are many lingering questions.
9

---

Published as a conference paper at ICLR 2024
As discussed, one limitation of the analysis is the reliance on gradient flow rather than gradient descent. We hope that in future work we'll able to effectively discretize the dynamics, made more challenging by the fact that one must track three parameters rather than simply the correlation. Still, we observe theoretically and empirically that the symmetric single index setting demands a number of unusual choices, such as a correlation loss and distinct student and teacher link function, in order to enable efficient learning. And in broader scheme, if one remembers the perspective of DeepSets as a very limited form of a three-layer architecture, the issue of provable learning for deeper. more realistic architectures stands as a very important and unexplored research direction and where Transformers with planted low-dimensional structures appear as the next natural question.
REFERENCES
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewiez. The merged-stairean, property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. arXiv preprint arXiv:2202.08658, 2022.
Emmanuel Abbe, Enric Boix-Adsera and Theodor Misiakiewiez. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprini arXiv:2302. 2023.
Luca Amaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean- field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. arXiv preprint arXiv:2302.05882, 2023.
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):478844838, 2021.
Mathieu Beau and Adolfo del Campo. Parent hamiltonians of jastrow wavefinntions. SciPost Physics Core, 4(4):030, 2021.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. arXiv preprint (tt1:22006040330 2022.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, 2022.
Imre Bihari. A generalization of lemma of bellman and its application to uniqueness problems of differential equa- tions. Acta Mathematica Hungarica, 7(1):81 1956.
Joan Bruna, Loucas Pillaud- Vivien, and Aaron Zweig. On single index models beyond gaussian data. arXiv preprint arXiv:2307. 15804, 2023.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. arXiv preprint arXiv:2305.106333 2023.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, 2022.
Persi Diaconis and Mehrdad Shahshahani. On the eigenvalues of random matrices. Journal of Applied Probability, 31 (A):49-62, 1994.
Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pp. 1887-1930. PMLR, 2018
Robert FH Fischer. Precoding and signal shaping for digital transmission. John Wiley & Sons, 2005.
Martin Jacobsen. Laplace and the origin of the orstten--ullnneeck process. Bernoulli, 2(3):271-286, 1996.
Sham M Kakade, Varun Kanade, Ohad Shamir, and Adam Kalai. Efficient learning of generalized linear and single index models with isotonic regression. Advances in Neural Information Processing Systems, 24, 2011.
Adam Tauman Kalai and Ravi Sastry. The isotron algorithm: High-dimensional isotonic regression. In COLT, 2009.
Edwin Langmann. A method to derive explicit formulas for an elliptic generalization of the jack polynomials. arXiv preprint 1015 2005.
10

---

Published as conference paper at ICLR 2024
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invarinn neural networks. In International Conference on Machine Learning, pp. 3744-3753. PMLR, 2019.
lan Grant Macdonald. funnttrricfunctions and Hall polynomials. Oxford university press, 1998.
Arvind Mahankali, Jeff Z. Haochen, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time, arXiv preprint UrXiv:2306.16333 2023.
Ryan O'Donnell. Analysis of boolean functions. arXiv preprint arXiv:2105. 10386, 2021.
Stephen B Pope. Algorithms for ellipsoids. Cornell University Report No. FDA, pp- 08-01, 2008.
Charles Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings the IEEE conference on computer vision and pattern recognition, pp. 652 2606 2017.
Peter Radchenko. High dimensional single index models. Journal of Multivariate Analysis, 139:266-282, 2015.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. Advances in neural information processing systems, 30, 2017.
Shai
Ohad Shamir, and Karthik Sridharan. Learning kernel-based halfspacce with the zero-one loss. arXiv preprint 0OK5v: 2010.
Roman Vershynin. Hlihh-immensionll probability: An introduction with applications in data science, volume 47,
Cambridge university press, 2018.
Gilad Yehudai and Shamir Ohad. Learning a single neuron with gradient methods. In Conference on Learning Theory, pp. 3756-3786. PMLR, 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Aaron Zweig and Joan Bruna. Exponential separations in symmetric neural networks. Advances in Neural Information Processing Systems, 35:33134-33144 2022.


Reviews:
[{"review_id": "dOJ5CMBnb9", "reviewer": "Reviewer_xBy6", "summary": "The paper analyzes single-index learning, especially when the underlying teacher model has permutation invariance and the student model is a variant of the DeepSets architecture. \nThe paper shows that under the squared Vandermonde density over a complex domain, the hidden direction (of a finite support in the power-sum polynomials feature space) can be recovered up to an $s$-th root of unity via the gradient flow applied for a variant of correlation loss with regularization loss, where $s$ is the information exponent. A provably good initialization scheme is also provided as well as a concrete form of a good activation function that meets the regularity condition.", "questions": "Minor suggestions\n- After eq. (6), \"In other words\" sounds confusing. It might be better to explicitly say that the first-layer weights ($\\{a_1,\\ldots,a_M\\}$) and the third-layer weight (the function $g$) will be fixed (frozen), and the second-layer weight ($w$) will be the only trainable parameter.\n- In Corollary 4.7, delete \"And\" before \"Consider\".\n- Please consider different notation for the summary statistics in Theorem 4.1 other than $m$ as it clashes with its other usage as an index.\n- The usage of $\\inf$ in the definition of the information exponent seems unnecessary.  \n- Figure 1 can be much improved. For examples, y-axes can be shifted to hide unused range.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Understanding convergence of a neural network training with gradient-based methods is an important problem, and this paper provides a new theoretical framework that admits a clean analysis. The problem setting and techniques may be useful for further work in this line of study.\n- The paper is well written overall. I am new to this area, but I was able to understand a high-level landscape of the previous works and the contribution of this work compared to the existing works. Though there are several simplifying assumptions on the teacher network, the input distribution, the loss function, and etc., the rationale behind them are clearly explained and discussed appropriately. The reasonings behind the choice of loss functions and the assumptions are well thought out and helpful.\n- A technical difficulty in a SGD analysis is also insightful.", "weaknesses": "- There are several assumptions that limit the applicability of the framework, though they are properly discussed. Proposition 2.3 is the key technique that governs the whole technical assumptions including the assumption on the link function (Assumption 2.4). It is unclear for me as a newbie reader to this field how much it is restrictive and whether this can be relaxed. Any further discussion would be appreciated.\n- Though the paper starts with the single-index learning, it seems that the framework deviates from the framework by considering a different link function in the final form of the modified correlation loss in eq. (14). I found it quite confusing at first sight. After this modification, is this technically qualified as a single-index learning?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mo8KW8Q9j0", "reviewer": "Reviewer_oXjh", "summary": "This paper studies the problem of studying single-index functions on the feature space given by power-sum polynomials.\nThus, all of the teacher functions fall into the class of symmetric functions, which are those functions such that f(x1,...,xn).\n\nUnder a specific data distribution that allows for an inner product with nice orthogonality properties, this paper analyzes the gradient flow of learning with a student model under the correlation loss.\n\nAn information exponent (similar to Ben arous et al.), plus the initial correlation with the ground truth, is shown to upper-bound the time needed to learn.", "questions": "See questions in the weaknesses section.\nAlso:\n4) It is unclear to me what I am supposed to be learning from the experiments in Figure 1. How does the initial alignment of the model correspond to the gradient flow time in these cases?\n\nTypos:\n* \"The former assumptions essentially corresponds\"\n* \"any other works that demonstrates\"\n* \"dynamcis\"\n* Assumption 4.4(iv) seems like it should be Omega(sqrt(N)) instead of O(sqrt(N)); Lemma 4.5 should be exp(-Omega(N)) instead of exp(-O(N)), and similarly for Lemma 4.6; other places including in the appendix there are some Omega(.) vs O(.) issues", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The presentation is clear and easy to follow.\n\nThe problem of provable guarantees for the dynamics of learning symmetric functions has not been studied as far as I know, so this can be of interest to the community.\n\nRelated literature is covered well.\n\nThe analysis seems correct.", "weaknesses": "1) \"These dynamics naturally motivate the question of learning efficiency, measured in convergence rates in time in the\ncase of gradient flow\". Does this say anything about sample complexity when discretized? Because gradient flow time can be rescaled, so it doesn't appear to be a well-defined complexity measure?\n\n2) Can you prove a converse to Theorem 4.2 with corresponding lower bounds on the time? This seems doable and like it would strengthen the result to make it more of a characterization.\n\n3) It is unclear to me what is conceptually new in this work that does not appear in previous analyses of single-index learning?\n* (a) the teacher model / student model are single-index models, but on a different data distribution than usual. But this has a similar inner product to what allows analyzing the Gaussian case.\n* (b) the analysis of the single-index model appears to follow a now-standardized template. Would be interesting for the authors to highlight what are the new elements, or how this generalizes the currently-known techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ecRLNqaKvD", "reviewer": "Reviewer_qFtS", "summary": "This paper studies the convergence guarantees of gradient flow for the problem of retrieving the unknown parameters of a symmetric single index model with a known analytic link function, under appropriate assumptions on the distribution of the features, which are complex vectors. It is shown that gradient flow on some appropriately defined loss function based on a DeepSets student model structure converges to a student model whose parameters are highly correlated with the ground truth, in time that scales logarithmically with the approximation error parameter and exponentially to the (square of the) information exponent associated with the link function, whenever the chosen student model satisfies certain properties. Moreover, an explicit choice of the student model is proposed and precise bounds for the associated convergence rate are provided. The authors also provide experimental demonstration of their results on synthetic data.\n\nThe results in this paper work under a set of modeling assumptions. First, the teacher model is assumed to have a symmetric structure, in the sense that variable relabelling does not change the value of the function. Each of the unknown parameters influences each variable in the exact same way and, in particular, the $k$-th parameter multiplies the sum of all the monomials of degree $k$. Moreover, the link function is assumed to be analytic and supported on monomials of degree at most equal to the square root of the input dimension. Finally, the input distribution is assumed to be the squared Vandermode density, which is shown to enjoy properties in the symmetric setting analogous to those of the Gaussian for a non-symmetric setting.\n\nThe proof techniques used include complex calculus as well as a crucial identity of Vandermode distributions on the symmetric polynomials setting, which motivates establishing Vandermode distributions as the Gaussian analogue in the symmetric polynomials setting.", "questions": "My main question concerns my second comment in the weaknesses section: Can the results provided in this paper be translated to standard PAC learning guarantees? If not, then what are the main obstacles in doing so?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper provides the first positive result in the proposed setting and acquires bounds in terms of the information exponent, which is well-studied quantity. The presentation is clear, sufficiently detailed and accurate. The techniques proposed in this paper might be of independent interest, especially as they motivate further study of learning problems under Vandermode marginals.", "weaknesses": "As the authors mention in their limitations section, the (distributional and modeling) assumptions required for the proposed analysis to work are not as common in the literature. Therefore, it is not clear to what extent such assumptions are realistic or significantly valuable from a theoretical perspective. That said, this work provides results that apply to a very large class of link functions, which partly justifies a certain number of strong distributional assumptions.\n\nAnother weakness of this work is that the results do not seem to be immediately interpretable in the Probably Approximately Correct (PAC) learning setting, since it is implicitly assumed that one has access to exact gradient oracles for optimizing the chosen loss function.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k6EKq3qUkR", "reviewer": "Reviewer_7zCe", "summary": "This paper studies the population gradient flow dynamics for learning a single index model using a symmetric neural network. It is shown that under powersum polynomial features and inputs following the Vandermonde distribution, identities appear that are similar to those of Hermite polynomials in Gaussian space. Using this property, the authors track the high-dimensional gradient flow dynamics using a small number of summary statistics and show the recovery of the hidden direction in the single index model.", "questions": "* The authors have mentioned why analyzing the SGD counterpart of the population gradient flow using the drift and martingale decomposition technique can be challenging. Another avenue for obtaining sample complexity is to use full-batch (empirical) gradient flow similar to Bietti et al., 2022 instead of the population version. I'm wondering if the six-stage proof can be adapted to handle such concentration errors and lead to a sample complexity.\n\n* A particularly relevant paper could be [1] where the authors consider learning a single index model with anisotropic Gaussian data, and show that a structure in the covariance can reduce sample/runtime complexity and even remove dependency on the information exponent. I am wondering if similar observations can be made in this paper by introducing additional structure in the inputs. Furthermore, some aspects of the current analysis also appear in [1] due to dealing with anisotropy, such as having to control a quantity of the type $\\Vert f(\\langle Aw, x\\rangle)\\Vert_{L^2}^2$ or preconditioning, which might be interesting to point out.\n\n* The squared Vandermonde density is never explicitly defined in the manuscript. Perhaps providing an explicit definition and recalling some of its relevant properties can help the readers better understand the problem setting.\n\n## Minor Comments:\n* Is the probability in Lemma 4.5 is over the randomness of $(a_m)$? In that case, it might be helpful to point this out in the statement of the lemma.\n* Theorem 4.2 point (iii) asks for $v_0 = 1 - r_0^2$. Could it be more intuitive to introduce this condition as $\\Vert Aw_0 \\Vert = 1$?\n---\n[1] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, Murat A. Erdogdu. \"Gradient-Based Feature Learning under Structured Data.\" NeurIPS 2023.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper presents an interesting analysis that goes beyond the standard two-layer fully connected network and Gaussian input assumptions of the recent literature on single index learning, specifically by considering symmetric neural networks under a Vandermonde distribution. Furthermore, the authors are able to generalize the intuitions of Gaussian single index learning by showing that a proper notion of *information exponent* is still controlling the complexity of learning.", "weaknesses": "* There is no finite-sample analysis which is in contrast to the majority of recent works on learning single index models which had a particular focus on the sample complexity of the learning problem.\n\n* Certain assumptions of the work are a bit restrictive. For example, the assumption on the activation is only verified for a particular choice that is not commonly used in practice, and the student link function $g$ needs to be carefully constructed using full information on the teacher link function $f$. However, these restrictions are understandable given that this work is a first attempt of studying the training dynamics of symmetric neural networks under a single index model.\n\n* The writing of the paper can be improved with examples provided below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)