Robust Feature Learning for Multi-Index Models in High Dimensions
Authors: Alireza Mousavi-Hosseini
Adel Javanmard
Murat A Erdogdu
Venue: ICLR 2025 Poster (2025)
Keywords: feature learning
adversarial robustness
neural networks
multi-index models
gradient descent
Primary area: learning theory
TL;DR: We prove that under certain standard conditions, adversarially robust learning can be as easy as standard learning in high dimensions, namely neural networks can share the same representation for both tasks.

Abstract: Recently, there have been numerous studies on feature learning with neural networks, specifically on learning single- and multi-index models where the target is a function of a low-dimensional projection of the input. Prior works have shown that in high dimensions, the majority of the compute and data resources are spent on recovering the low-dimensional projection; once this subspace is recovered, the remainder of the target can be learned independently of the ambient dimension. However, implications of feature learning in adversarial settings remain unexplored. In this work, we take the first steps towards understanding adversarially robust feature learning with neural networks. Specifically, we prove that the hidden directions of a multi-index model offer a Bayes optimal low-dimensional projection for robustness against $\ell_2$-bounded adversarial perturbations under the squared loss, assuming that the multi-index coordinates are statistically independent from the rest of the coordinates. Therefore, robust learning can be achieved by first performing standard feature learning, then robustly tuning a linear readout layer on top of the standard representations. In particular, we show that adversarially robust learning is just as easy as standard learning. Specifically, the additional number of samples needed to robustly learn multi-index models when compared to standard learning, does not depend on dimensionality.

Published as a conference paper at ICLR 2025
ROBUST FEATURE LEARNING FOR MULTI-INDEX MODELS IN HIGH DIMENSIONS
Alireza Mousavi-Hosseini=> Adel Javanmard', Murat A. Erdogdu University of Toronto, Vecctor Insitute, 'University of Southern California {mousavi, erdogdu 6 toronto, edu, ajavanma&use usc edu
ABSTRACT
Recently, there have been numerous studies on feature learning with neural net- works, specifically on learning single- and multi-inde models where the target is a function of cow-dimensional projection of the input. Prior works have shown that in high dimensions, the majority of the compute and data resources are spent on recovering the low-dimensional projection; once this subspace is recovered, the remainder of the target can be learned independently of the ambient dimension. However, implications feature learning in adversarial settings remain unexplored. In this work, we take the first steps towards understanding adversarially robust feature learning with neural networks. Specifically, we prove that the hidden di- rections of multi-inde model offer Bayes optimal -iw-demeniioal projection for robustness against fz-bounded adversarial perturbations under the squared loss, assuming that the multi-index coordinates are statistically independent from the rest of the coordinates. Therefore, robust learning can be achieved by first performing standard feature learning, then robustly tuning a linear readout layer on top of the standard representations. In particular, we show that adversarially robust learning is just as easy as standard learning. Specifically, the additional number of samples needed to robustly learn multi- index models when compared to standard learning, does not depend on dimensionality.
INTRODUCTION
A crucial capability of neural networks is their ability to hierarchically learn useful features, and to avoid the curse of dimensionality by adapting to potential oww-iimesiional structures data through empirical risk minimization (ERM) (Bach, 2017; Schmidt-Hieber 2020). Recently, a theoretical line of work has demonstrated that gradient-based training, which is not priori guaranteed to implement ERM due to non-convexity, also demonstrates similar behavior and efficiently learns functions of low-dimensional projections (Wei et al., 2019; Damian al., 2022; Bietti et al., 2022; Barak et al., 2022; Ba et al., 2022; Mousavi-Hosseini et al., 2023a) functions with certain hierarchical properties (Abbe et al., 2022; 2023; Dandi et al., 2023). These theoretical insights provided a useful avenue for explaining standard feature learning mechanisms in neural networks.
On the other hand, it has been empirically observed that deep neural networks trained with respect to standard losses are susceptible to adversarial attacks; small perturbations in the input may not be detectable by humans, yet they can significantly alter the prediction performed by the model (Szegedy al., 2014). To overcome this issue, popular approach is to instead minimize the adversarially robust empirical risk (Madry et al., 2018). However, unlike its standard counterpart, achieving successful generalization of deep neural networks on robust test risk has been particularly challenging, and even the standard performance of the model can degrade once adversarial training is performed (Tsipras al., 2018). Therefore, one may wonder if robust neural networks are still adaptive to certain problem structures that improve generalization. By focusing on hidden now-dimensionality as well-known example of such structure, we aim answering following fundamental question:
Can neural networks maintain their statistical adaptivity to low dimensions when trained for robustness against adversarial perturbations?
We answer this question positively by providing the following contributions.

---

Published as a conference paper at ICLR 2025
When considering fg-constrained perturbations, Bayes optimal predictors can be constructed by projecting the input data onto the -iw-dmmnsiioal subspace defined by the target function. In this sense, the optimal low-dimensional projection remains unchanged compared to standard learning.
Consequently, provided that they have access to an oracle that is able to recover the low- dimensional target subspace, neural networks can achieve a sample complexity that is independent of the ambient dimension when robustly learning multi-index models. This is achieved by mini- mizing the empirical adversarial risk with respect to the second layer. While the basic definition of empirical adversarial risk implies computational complexity dependent input dimension, by simply projecting the inputs onto the oow-dimensional target subspace, the computational complexity can also be made independent of the input dimension.
An oracle for recovering the low-dimensional target subspace can be constructed by training the first layer of a two-layer neural network with a standard loss function, as demonstrated by many prior works. By combining our results with two particular choices of oracle implemen- tation (Damian et al., 2022; Lee et al., 2024), we provide end-to-end guarantees for robustly models with gradient- based algorithms.
learning
1.1 RELATED WORKS
Feature Learning for Single/Mullti-Inde Models. Many recent works have focused on proving benefits of feature learning, allowing the neural network weights to travel far from initialization, as opposed to freezing weights around initialization in lazy training (Chizat et al., 2019) which is equivalent to using the Neural Tangent Kernel (Jacot et al., 2018). When using online SGD on the squared loss, Ben Arous et al. (2021) showed that the complexity of learning single-index models with known link function depends on a quantity called information exponent. Gradient- learning of single-index models with information exponent I was studied in Ba et al. (2022); et al. (2023a), and Damian et al. (2022) considered multi-indes polynomials where the equivalent of information exponent is at most 2. For general information exponent, Bietti et al. (2022) provided an algorithm for gradient-basee learning using two-layer neural networks. A feature learning analysis faithful SGD without modifications was presented in Glasgow (2024) for learning the XOR. The counterpart of information exponent for multi-index models, the leap exponent, was introduced in Abbe et al. (2023). Considering SGD the squared loss an example of Correlational Statistical Query (CSQ) algorithm, Damian et al. (2023) provided CSQ-optimal algorithms for learning single- index models. Further improvements to the isotropic sample complexity were achieved by either considering structured anisotropic Gaussian data (Ba et al., 2023; Mousavi-Hossein et al., 2023b), or the sparsity of the hidden direction (Vural & Erdogdu, 2024). The benefits of feature learning have also been considered for multitask learning (Collins et al., 2024) and in networks with depth larger than 2 (Nichani et al., 2023; Wang et al., 2024b). More recently, it was observed that gradient-based learning can go beyond CSQ algorithms by reusing batches (Dandi et al., 2024; Lee et al., 2024; Arnaooli et al., 2024), or by changing the loss function (Joshi et al., 2024). In such cases, the algorithm becomes an instance of Statistical Query (SQ) learner, and the sample complexity is characterized by the generative exponent of the link function (Damian et al., 2024).
Hosseeiii
While the above works mostly exist in narthw-width setting where the interaction between neurons is largely ignored, another line of research focused on the mean-field or wide limits of two-layer neural networks (Chizat& & Bach, 2018: Rotskoff & Handen-Eijnden 2018; Mei et al., 2018) for providing learnability guarantees (Wei et al., 2019; Chizat & Bach, 2020; Abbe et al., 2022: Telgarsky, 2023; Mahankali et al., 2023; Chen & Ge, 2024). In particular, the mean-field Langevin algorithm provides global convergence guarantees for two-layer neural networks (Chizat, 2022; Nitanda et al., 2022), leading to sample complexity linear an effective dimension for learning sparse parities (Suzuki et al., 2023; Nitanda et al., 2024) and multi-index models (Mousavi-Hosseini et al., 2024).
Adversarially Robust Learning. The existence of small worst-case or adversarial perturbations that can significantly change the prediction of deep neural networks was first demonstrated in Szegedy et al. (2014). Among many defences proposed, one effective approach is adversarial training introduced by Madry et al. (2018), which is based solving min-max problem to perform robust optimization. One observation regarding this algorithm is that it tends to decrease the standard performance of the model (Tsipras et al., 2018). Therefore, the following works studied the hardness of robust learning and established a statistical separation in a simple mixture of Gaussians setting (Schmidt et al., 2018), or computational separation by proving statistical query lower bounds (Bubeck et al.,
2

---

Published as a conference paper at ICLR 2025
2019). Further studies focused on exact characterizations of the robust and standard error, well as the fundamental and the algorithmic tradeoffs between robustness and accuracy in the context of linear regression (Javanmard et al., 2020), mixture of Gaussians classification (Javanmard & Soltanolkotabi, 2022), and in the random features model (Hassani & Javanmard, 2024). Closer to our work, Javanmard & Mehrabi (2024) show that this tradeoff is mitigated when the data enjoy low- dimensional structure. However, they focus on the population adversarial risk in binary classification and generalized linear models, where the features live on a low-dimensional manifold with known structure. Here, we consider multi-inder model wherein the response depends on a low-dimensional projection of inputs, and derive fintte-saple bounds for adversarial risk.
In this work, we provide an alternative narrative compared to the line of work above by showing that in a high-dimensional regression setting, learning multi-index models that are robust against & perturbations can be as easy as standard learning. We achieve this result by focusing on the feature learning capability of neural networks, i.e. their ability to capture low-dimensional projections.
Notation. For Euclidean vectors, (..-) and ||-|| denote the Euclidean inner product and norm respec- tively. For tensors, II-IIp and ||- denote the Frobenius and operator norms respectively. We use S for the unit sphere in RR and Tl denotes the uniform probability measure on S For quantities a and 6, (b) means there is an absolute constant C. such that Cb, and ! is similarly defined. O and Ω allow C to grow polylogarithmially with problem parameters.
2 PROBLEM SETUP: FEATURE LEARNING AND ADVERSARIAL ROBUSTNESS
Statistical Model. Consider a regression setting where the input æ E Rd and the target E R are generated from distribution (z, ~ P. For prediction function Rd R, its population adversarial risk, where we assume the adversary can perform a worst-case perturbation on the input with budget o measured in f--nomm, before passing the model, is defined
AR(f)= = E max (f(x + 5) - y)²
(2.1)
|
where the expectation is over all random variables inside the brackets. Given (non-parametric) family of prediction functions F, our goal is to learn predictor that achieves the optimal adversarial risk given by
AR = min AR(S): FEF
(2.2)
We focus on learners of the form of two-layer neural networks with width N, given as
f(x; a, W, b) = a Toow + b),
(2.3)
where a E N is the second layer weights and W E xxd and b E R are the first layer weights and biases. To avoid overloading the notation we use AR( a, W, bb))) access to n i.i.d. samples ()) y y0 from P, the goal is to learn the network parameters a, W. and b in such a way that the quantity AR(a, W, b) is close to the optimal adversarial risk AR".
AR(a, W, b). Given
long line of recent works has shown that neural networks are particularly efficient in regression tasks when the target is function of dimeniienioonal projection of the input, see c.g. Bach (2017). Throughout the paper, we also make the same assumption that the data follows a multi-index model,
Ely æ] g((u), .)............... ....... (uR, x)),
(2.4)
for all æ E Rd, where g RR - R. is the link function, and we assume U1...., Us are orthonormal without loss of generality. Let U E Rxd be an orthonormal matrix whose rows are given by (u.); we use the shorthand notation g( u((u1,,)) ...... (ux, x)) := (Uw). ). In the special case where = 1,
this model reduces to single-index model. In this paper, we consider the setting where A: < d, and in particular l: = (1).
Feature Learning. In the context of training two-layer neural networks when learning multi-index models, feature learning refers to recovering the target directions U via the first layer weights W. Successful feature learning reduces the effective dimension of the problem from the input dimension to the number of target directions kt and circumvents the curse of dimensionality when < d.
The complexity of recovering U depends multiple factors such as the choice of algorithm as well as the properties of the link function. We will provide an overview of some existing results for recovering U with neural networks in Section 4.2, along with several concrete examples.
3

---

Published as a conference paper at ICLR 2025
3 OPTIMAL REPRESENTATIONS FOR ROBUST LEARNING
In this section, we demonstrate that under 42 fg-constrd perturbations, the optimal low-dimensional representations for robust learning coincides with those in standard setting, both of which are given by the target directions U. Crucially, our result relies on the following assumption on input distribution.
Assumption Suppose U R k)x is any orthonormal matrix whose rows complete the rows of U into a basis of R. Then, y and U are jointly independent from Ux.
Üæ The above assumption is quite general. For example, with the notation æ1| == U and T holds when y = g( + where is independent zero-mean noise, and U + Ú for independent vectors 21, Z2 E Rd, We now present central result below along with its proof. We discuss the necessity of Assumption and fz constrained attacks to obtain this result in Appendix B.
Theorem 1. Suppose Assumption holds and (2.2) admits a minimizer. Then, there exists function I : IRd + R of the form (x) hu x)) with h IR - R given by l((2) = B(f((r) Ux 2) 2 some f E F, such that
AR( S AR",
(3.1)
with equality when f F.
Consequently, to achieve AR", one only needs to (i) learn the target directions U, and (ii) approxi- mate functions in k-dimensional subspace rather than d. For two-layer neural networks, the first layer W recovers U, and the remaining parameters a and b are used to approximate the optimal h. While this recipe is general, we provide specific implications in the next section.
Proof. We will show that for every f E F. h(z)= = [x)(() Ux = =] gives AR(h(U-)) S AR()).
Then, choosing f to be some minimizer of AR yields the desired result.
Define the residuals Ty [2]] d1j) == γ (3 d1)). and rx 8) == f(x + 8) h h( + +11)) Then, by a decomposition of the squared loss and the tower property of conditional expectation,
AR(J)- = E(awy)
Ex1
max Ty (3)) (1111110 0113 + rr, 8)2 - 2r 0rrrr (x: 8) [181|5e
æt y/
max (z) 61 + Eze (x, 8)² x1|| y] - (, . 8, Exs [ry (xx 5) x 1.
2 ((r
2 ())
max [118115e,5 =0} To (2 81 + Eu, [r, 8)3² x - 2r,( 81) Ex |r 5)
20) 3/
Since y| is independent from 1 for any fixed 8, we have E (1z, 8) 2 w] [x 8) x
Thus, using the notation f(x) = [ || x1), provided that 8_ 0, Assumption 1 yields
h(z + 51) = EE [f()) क|| = x + 1p1 [f + 6/. x1 + 81)] B[f( + 8) 인 = x),
for all = E R Plugging in z Bj gives (ar, 8) [m]] = 0. Therefore,
AR(F) 2 E (()jy)
max
(z 8 ) -E [ry(2,, 8)3 x v]
[|]o]| =0}
E (
max [1|8||==.0
(y- h(U( + 5)))² AR(M(U-)),
where we dropped the constraint 81 0 as it does not contribute, which concludes the proof.
discussion on robust/non-robut feature decomposition. Many prior works in classification assume that features can be divided into robust and non-robust groups (see e.g. Tsipras et al. (2018); Ilyas et al. (2019); Kim et al. (2021); Li & Li (2024).) We do not rely on this robust/non-robus decomposition. Instead, the ls relevant features for predicting y can be either robust or non-robust. The robust training of the second layer ensures the model utilizes the robust subset of these features, if such subset exists, while the first layer performs dimensionality reduction.
4

---

Published as a conference paper at ICLR 2025
Before moving to the next section, provide the following remark on proper scaling of Since ||||=|| grows with vd, it may seem natural to scale the adversary budget E with dimension as well. However, we provide a simple argument on the contrary. Consider the single-index case y g (u, 2)), and let h be the optimal function constructed in Theorem 1, providing the prediction function æ + h( (u, x)). One can then observe that even a constant order e is sufficient to incur a large change in the input of h, e.g.. choosing 6 su perturbs the input of the predictor by e. Thus, this justifies the regime where & is of constant order compared to the input dimension, which is the focus in the rest of the paper.
4 LEARNING PROCEDURE AND GUARANTEES
As outlined in the previous section, to robustly learn the target model, standard representations U suffice. In this section, we consider concrete examples of how a standard feature learning oracle combined with an adversarially robust second layer training leads to robust learning. We assume access to the following feature learning oracle to recover U. We will provide instances of practical implementations of this oracle using standard gradient- based algorithms in Section 4.2.
Definition 2 (DFL). An a- Deterministic Feature Learner (DFL) is an oracle that for every Ç given (() samples from P. returns a weight matrix W (W1..... WN,) E R with unit-norm rows, such that for all u E span(u1,... ...) with |||||| for some a >> 0 we have
Ki: (w;, u) 2 1 11--5
5H
aç 12-1)
N
An a-DFL oracle returns weights such that roughly an a-proportion of them align with (and suffi- ciently cover) the target subspace. By packing argument, we can show that the best achievable ratio is a S c(k) for some constant c(k) >> 0 depending only on k. which is why we use the normalizing factor e 1)/2 above. We show in Section 4.2 that the definition above with a constant order a is attainable by standard gradient-based algorithms. That said, the multi-index setting, it is possible to improve our learning guarantees by considering the following stochastic oracle.
Definition 3 (SFI.). An (ax,3)- SStochatiic Feature Learner (SFL) is an oracle that for every {> given (S) samples from P, returns a random weight matrix W = =(w1,... ,,,N) E RRN xd with
unit-norm rows, such that there exists S C [N] with |S|/N satisfying l|w. U U w.ll S
for E S. Further, ([Uw,ll)ee ies" ~ p, and du 2 B for some where ş is some measure and Tk drk
Uw,, i.d
is uniform, both supported on gk-1
The above oracle essentially defines a random features model the smaller target subspace, where a subset of the weights are sampled independently from a distribution that supports all target directions du U. The lower bound on dry ensures sufficient coverage of the low-dimensional space of target directions. We note that an (a., 3)LL oracle can be used to directly implement a-DFL oracle; by standard union bound argument, one can show N Õ(1 ((1/(.8((-1))/) guarantees the output of (a. (3-FFL satisfies Definition with high probability. Therefore, while its definition is slightly more involved, (a, B)L a more specialized oracle compared a-DFL.
Once the first layer representation is provided by above oracles, we can fix the biases at some random initialization, and train the second layer weights a by minimizing the empirical adversarial risk
REA
AR(a, W
max 180
(f) (4) + 8(0); a, W. V) - (
(4.1)
T&FA
We formalize the training procedure with two-layer neural networks in Algorithm 1. We highlight that keeping biases at random initialization while only training the second layer performs non-linear function approximation, and has been used in many prior works on feature learning (Damian et al., 2022: Mousavi-Hosseine et al., 2023b; Oko et al., 2024). Further, while a + AR(a, (a, W, b) is convex for fixed W and b since it is maximum over convex functions, exact training of in practice may not be straightforward since the inner maximization is not concave and does not admit a closed-fome form solution. In practice, some form of gradient descent ascent algorithm is typically used when training a (Madry et al., 2018). We leave the study of the computational aspect of that part to future work.
We will make the following standard tail assumptions on the data distribution.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Adversarially robust learning with two-layer NNs.
Input: Ç, Tas rb. (ar(0, ())) J=1
+FL(()+BAA FL E {a-DDLL (a.,3)-FI..)
1: Phase I: Feature Learning
2: W = =FL (..
( y (0)0
me))))(()))
3: Phase 2: Robust Function Approximation
4: by Fdd
for 1 J S. AR(a, b).
5: à arg min min ||lls To VN
6: return (a, W.
Assumption 2. Suppose x: has zero mean and (11) subGaussian norm. Furthermore, for all 1,
it holds that E[vl'"] I ((PP/)) ) for some constant p 2 1.
Note that the condition on y above is mild; for example, it holds for a noisy multi-index model (U)) + where has (1)) subGaussian norm and grows at most polynomially, i.e., (g()) 1+| ... Similarly, we also keep the function class FF quite general and provide our first set of results for class of pseudo-Lipschitz functions which is introduced below.
Assumption 3. We assume F is a class of functions that are Eseudo-Lipschiin along the target coordinates. Specifically, using the notation f(x) = f(z]], xL) and defining E1 == 1V E, we have
x1) - f(z2, x1)) & (< p-1 + "|22111 + 1)]1 - 222
for all f E F. all 21, 22 R and some constants L and p 1 such that E[L(x = L.
Remark. The prefactor - is justified intuitively since the optimal function of the form h(2) ([f(x) Ux= x] should satisfy E max (se - hxU(x + 8) AR", which is bounded, and does not grow with E beyond a certain point. This implies that h must be sufficiently smooth while its input is perturbed, and in particular, its (local) Lipschitz constant should remain bounded while E grows, hence the introduction of the prefactor.
The first result of this section assumes access to -DFL oracle.
Theorem Suppose Assumptions 1,2,3 hold and the ReLU activation used. For a tolerance e define ē = e A (e2 AR"), and for the adversary budget & recall e == 1 V.e Consider Algorithm I with FL = a- aDFF oracle, Ta ((ee//V)/tt244// /a) and rb (EI((E/VV))
Then, if the number of second phase samples TiFA, the number of neurons N, and -DLL error satisfy
O(k)
O(K)
O(K)
TIFA 2 á
E1
N N
E1 Ve
aÇ
11/2
VEI
we have AR(â, W. b) S AR" + c with probability at least 1 nPA where e >> O is an absolute constant.
The total sample complexity of Algorithm 1 is given by the sum of complexities of the two phases, i.e., Teootll T&AA + (DDFL(C). We will provide bounds on (DF in Propositions and 10 to ultimately characterize Neotal in Corollaries 9 and 11.
Remark.
The above theorem states that once the feature learning oracle has recovered the target subspace, the number of samples and neurons needed for robust learning is independent of the ambient dimension d. Thus, high-dimensional setting, statistical complexity is dominated by the feature learning oracle, implying that adversarially robust learning is statistically as easy as standard learning.
Arguing about computational complexity is more involved. While the number of neurons required is independent of d, in its naive implementation, Phase of Algorithm 1 needs to solve inner maximization problems over Rdd which may be costly. However, once U E ggxd is estimated in Phase 1, we can reduce the dimension from d to & by projecting onto span ppan(u..... (uI...., 2,,,,, i.e.,
Σ a,0((1,, + b,)
Σ aoo(W Ux) + b,).

---

Published as a conference paper at ICLR 2025
With this modification, only need to consider worst-case perturbations over IR*,, thus the computa- tional complexity of Phase will also be independent the ambient dimension d.
It is possible remove the dependence on Ç in the number of neurons by instead assuming access to an (a,3)-SFF oracle, as outlined below.
Theorem 5. Consider the same setting as Theorem 4, except that we use the (a,3) oracle in Then, the sufficient number of second phase samples nEA, neurons N, and oracle error Ç are given
Algorithm with Ta = Ō ((#1/ /tn//t/////////aa/////
O(K)
O(K)
O
TAAA 2
E1
N2
E1
5S
3
a
023²
Under a Gaussian input assumption, there exist a DFL and (a.,3-SFL oracles that rely only on standard gradient-based training such that for small constant Ç, (DF and (C) both scale with some polynomial of d, where the exponent depends on certain properties of the link function, termed the information or generative exponent (Ben Arous al., 2021; Damian et al., 2024). We will provide explicit examples of such algorithms in Section 4.2 to characterize the total sample complexity Nhttal = nFA + "IDFL/SFL- For the interested reader, we restate Theorems 4 and 5 in Appendix B.2 with explicit exponents.
4.1 COMPETING AGAINST THE OPTIMAL POLYNOMIAL PREDICTOR
In this section, we restrict F to only polynomials, which allows us to derive more refined bounds on the number of samples and neurons. Specifically, we make the following assumption.
Assumption Suppose F the class of d-variate polynomials of degree p for some constant
p > 0. Further, a is either polynomial of degree 2 or the ReLU activation for which we define q = (p V.
While the ReLU activation is sufficient for function approximation, also consider polynomial activations in Assumption 4 since using those, recent works have been able to achieve sharper theoretical guarantees of recovering the target directions (Lee al., 2024); we provide a more detailed discussion in Section 4.2. Note that priori we do not require a growth constraint on the Theorem 1 automatically chooses a
coefficients the polynomials in F. The optimal function polynomial with suitably bounded coefficients in order to avoid incurring large robust risk.
The following result establishes the sample and computational complexity for competing against polynomial predictors when having access to a-DFL oracle.
Theorem 6. Suppose Assumptions 1,2,4 hold. For tolerance c > 0 define cA (c2/AR),, and for the udversary budget & recall & == 1 V s. Consider Algorithm I with a -DLL oracle, To Õ(1), Tb ((1)) the number of second phase samples REA, neurons N, and -DFL error satisfy
4(g+1)
a+1 <5ō
TIPA > 92
N
á
442
ac
+
2[g+1] El
we have AR(a, W, < AR" + e with probability at least 1 - 1 where e > O is an absolute constans.
Consequently, when restricting F to the class of fixed degree polynomials, there is no curse of dimensionality for sample complexity, even in the latent dimension k. This is consistent with the standard learning setting, see e.g. Chen & Meka (2020). Further, similar the general case above, it is possible to remove the Ç dependence from N when having access to an SFL oracle, thus also achieving computational complexity as fixed polynomial independent of the latent dimension.
Theorem 7. In the setting of Theorem 6, consider using Algorithm I with (o.F))SFL. oracle. Then, the sufficient number of second phase samples, neurons, and oracle error are given
4(++1)
2(tt1)
P2
REA 2 52
N á
5
a.22
2[g+1)
E1
7

---

Published as a conference paper at ICLR 2025
We remark that the guarantees provided in Theorem 7 are generally better than those in Theorem 6 for large k; yet, they are strictly worse for & = 1. That said, both Theorems 7 and 6 respectively achieve better sample complexity guarantees compared their counterparts in the previous section, namely Theorems 4 and 5, simply by restricting the function class FF to polynomials.
4.2 ORACLE IMPLEMENTATIONS FOR FEATURE LEARNING
The task of recovering the target directions U is classical in statistics, and is known as sufficient dimension reduction (Li & Duan, 1989: Li, 1991). with many dedicated algorithms, see e.g. Kakade al. (201 1); Dudeja Hsu (2018); Chen & Meka (2020); Yuan al. (2023) to name a few. Here, we focus on algorithms based on neural networks and iterative gradient-based optimization.
While we only consider the case where æ is an isotropic Gaussian random vector, recovering the hidden direction has also been considered for non-isotropic Gaussians (Ba et al., 2023; Mousavi- Hosseini et al., 2023b) where the additional anisotropic structure the inputs can provide further statistical benefits, or non-Gaussial spherically symmetric distributions (Zweig et al., 2023). Our results readily extend to these settings as well. First, we present the case of single-index polynomials.
Proposition 8 (Lee al. (2024)). Suppose x ~ NN Ia), k and g is a polynomial of degree p where is constant. Then, there exists an iterative first-order algorithm on two-layer neural networks (see Algorithm 2) that implements an (a,,3)-FLL oracle and an a a-DFL oracle, where = +() and B = 1. Furthermore, we have (C) = (CF = (d)
Combined with Theorems 4-7, we obtain the following total sample complexity guarantee for robustly learning Gaussian single-index models.
Corollary 9. Consider the data model of Proposition 8, and assume that the adversary budget is E = (1)) Then, the total sample complexity of Algorithm I achieve optimal adversarial risk AR" with tolerance e using either a aDDFL or (a,3))SFL oracle in Proposition 8 is given as
Rhotal Ódd/²) when choosing F to be polynomials of fixed degree in Assumption 4, and the polynomial activation according to Lee et al. (2024),
Riotal= = (dd/2(1)) when choosing F to be Creudo-Lipschittfuunciii as in Assumption 3,
where we recall = (e²/AR").
When considering Gaussian single-in single-index models beyond polynomials, we must introduce the concepts of information and generative exponent to characterize the sample complexity of recovering the target direction. Let y (0, 1), and for any g R R in T.2 let g Ejz0 200 O, He, denote its Hermite expansion, where He, is the normalized Hermite polynomial of degree j. The information exponent of g is defined as s(g) == min > 0 ag # 0). The generative exponent on the other hand, is defined as the minimum information exponent attainable by any transformation of g, i.e. 8* (g)) == min7 (()))), where the minimum is over all T E 1.2 (g Thus, 8* (g) I s(g), and in for all polynomials. See Ben Arous al. (2021); Damian et al. (2024) for details.
particular,
There exists an algorithm based on estimating partial traces that implements a 1 1-DFL (or (1, 1)-SFL) oracle with (C) = Odd" /2 + d/(²) (Damian et al., 2024). While it may be possible to achieve similar sample complexity when training neural networks with a ReLU activation, the state of the art results for ReLU neural networks so far are only able to control the sample complexity with the information exponent 8, e.g. Bietti et (2022) provides a gradient-based algorithm for optimizing a variant of two-layer ReI.U network that implements 1 1-DFL. with ADFL = (d' poly({-1))
Recovering U with k > more challenging, and the general picture is that the directions in U are recovered hierarhically based on each direction's 's corresponding complexity, such as in Abbe et al. (2023). For simplicity, we look at case that sufficiently simple for all directions to be learned simultaneously, while emphasizing that in principle any guarantee for learning the subspace U can be turned into an implementation of the oracles introduced in the previous section.
NO, L2), g is polynomial of degree p. and p.k are constant. Further assume Twax (V's) (min("ag) 2 K. for some > 0, where o and mmax denote the minimum and maximum singular values, respectively. Then, there exists a first-order algorithm on
Proposition 10 (Damian et al. (2022)). Suppose
8

---

Published as a conference paper at ICLR 2025
ReLU (Known Direction)
Teacher: (Known Direction)
Teacher: (Known Direction)
training W fixed
1.05
0.48
init from
01.4
Robat Tex. Risk
0.40
0.40
0.85
0.36
0.90
0.32
1500
1500
4500
Iterations
Iterations
herotions
(Unknown Direction)
(Unknow Direction)
Teacher: (Unknown Direction)
training
IV' fund SD training
0.45
0.5
IV' init from SD training
1.04
Robust Text Rok
0.3
0.92
1500
6000
Iterations
Iterations
Iterations
Figure 1: The adversarial error of two-layer ReLU network function of the number of adversarial training iterations, where each iteration is performed on a batch of independent 300 samples, except 500 samples for He2 with unknown direction reduce variance. Full AD training performs adversarial training on all layers from random initialization. SD training is standard training. which provides better initialization for W before performing adversarial training. We use the adversary budget e for all experiments, each of which are averaged over three runs.
two-layer ReLU neural networks (see Algorithm 3) that implements an (o,3)-SFL. and an aDDL oracle, where = 1 and 3 2 Cj for some constant > 0 depending only K. Further, we have (SFL(() = (CF = (d + d
Combining the above proposition with Theorems 4-7, we obtain the following total sample complexity for robustly learning Gaussian multi-index models.
Corollary 11. Under the data model of Proposition 10. assume that the adversary budget is E = O(1). Then, the total sample complexity of Algorithm 1 to achieve optimal adversarial risk AR" with tolerance e using either (x DD-F or (o,,3)-FFI, oracle in Proposition 10 is given as
Thotal Õ + d/²²) when choosing F to be polynomials of fixed degree as in Assumption 4, and the polynomial activation according Lee et al. (2024).
Ttttal = ( + d(k0(()) when choosing F to be
as in Assumption 3,
where we recall
(22/AR*).
Remark. We highlight that the gap in the total sample complexity of Corollary 9 and Corollary 11 is due to more efficient guarantees for recovering the hidden direction for single-index polynomials. is an open question whether such efficient recovery is also possible for multi-index polynomials.
5 NUMERICAL EXPERIMENTS
As proof concept, we also provide small-scale numerical studies on Gaussian data to support intuitions derived from theory. Additional experiments real datasets are provided in Appendix E.

---

Published as a conference paper at ICLR 2025
We consider single-index setting, where the teacher non-linearity is given by either ReLU, tanh, or He2(z) (z² 1)/v2 which is the normalized second Hermite polynomial. The student network has _N = 100 neurons, and the input is sampled from æ ~ (0, Ia) with d = 100. We implement adversarial training in the following manner. At each iteration, we sample a new batch of i..d. training examples. We estimate the adversarial perturbations on this batch by performing 5 steps of signed projected gradient ascent, with stepsize of 0.1. We then perform gradient descent step on the perturbed batch. To estimate the robust test risk, we fix a test set of 10,000 i.i.d. samples, and use 20 iterations estimate the adversarial perturbation. Because of the online nature of the algorithm, the total number of samples used is the batch size times the number of iterations taken.
The first row of Figure compares the performance of three different approaches. Full AD training refers to adversarially training all layers from random initialization, where first layer weights are initialized uniformly on the sphere 5d-1 second layer weights are initialized i.i.d. from (0, I/N²), and biases are initialized i.i.d. from (0, 1). In the two other approaches, we initialize all first layer weights to the target direction u. In one approach we fix this direction and do not train it, while in the other, we allow the training of first layer weights from this initialization. As can be seen from Figure I, there is considerable improvement in initializing from u, which is consistent with our theory that this direction provides Bayes optimal projection for robust learning.
In the practical setting where we do not have the knowledge of u, we consider the following alternative. We first perform standard training on the network, i.e. assume 0 (denoted in Figure I by SD training). We can then either fix the first layer weights to these directions, or further train them adversarially from this initialization. Note that for fair comparison with the full AD method, we provide the same random bias and second layer weight initializations across all methods at the beginning of the adversarial training stage. Even though this approach is not perfect estimating the unknown direction, it still provides a considerable benefit over adversarially training all layers from random initialization, as demonstrated in the second row of Figure 1.
CONCLUSION
In this paper, we initiated theoretical study of the role of feature learning in adversarial robustness of neural networks. Under fz-constrained perturbations, we proved that projecting onto the latent subspace of multi-index model is sufficient for achieving Bayes optimal adversarial risk with respect to the squared loss, provided that the index directions are statistically independent from the rest of the directions in the input space. Remarkably, this subspace can be estimated through standard feature learning with neural networks, thus turning a low-dimensional one. As result, under the assumption of having access feature learning oracle which returns an estimate of this subspace, and can be implemented e.g. by training the first-layer of a two-layer neural network, we proved that robust learning of multi-index models is possible with a number of (additional) samples and neurons independent from the ambient dimension.
robust learning problem into a
We conclude by mentioning several open questions that arise from this work.
Stronger notions of adversarial attacks such loo -norm constraints have been widely considered in empirical works. It remains open to understand optimal low oow-dinesiionl representations under such perturbations as well as their implications on sample complexity.
While our work demonstrates that standard training is sufficient for the first layer, it is unclear what kind of representation is learned when all layers are trained adversarially. In particular, Figure I suggests that adversarial training of the first layer may be suboptimal in this setting, even when infinitely many samples are available during training.
Since our main motivation was to show independence from input dimension, the dependence of our bounds on the final robust test risk suboptimality e potentially improvable by more careful analysis. It is an interesting direction to obtain sharper dependency and investigate the optimality of such dependence on the tolerance e.
Finally, it is worth emphasizing that our theorems be easily adapted other standard feature learning oracles. As such, based on the training procedure used and its complexity in feature learning, results are amenable to further improvements in their total sample complexity.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
AJ was partially supported by the Sloan fellowship in mathematics, the NSF CAREER Award DMS-1844481, the NSF Award DMS-231 1024, an Amazon Faculty Research Award, an Adobe Faculty Research Award and an iORB grant form USC Marshall School of Business. MAE was partially supported by the NSERC Grant [2019-06167]. the CIFAR AI Chairs program, and the CIFAR Catalyst grant.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, 2022.
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewiez. Sgd learning on neural networks: leap complexity and saddle-to to--to---dddd saddle dynamics. arXiv preprint arXiv: 2023.
Luca Amaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iu- vant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint (xXiv:2405.15459) 2024.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional Asymptotics Feature Learning: How One Gradient Step Improves the Representa- tion. arXiv preprint arXiv:2205.0144 01445, 2022.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dimensional structure: spiked random matrix perspective. Advances in Neural Information Processing Systems, 36, 2023.
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, (((1):6296681 2017.
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 .08999, 2022.
Peter L Bartlett, Dylan Foster, and Matus J Telgarsky. neural networks. Advances in neural information processing systems, 30, 2017.
normalized margin bounds for
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent non-convex losses from high- -ighedimensional inference. J. Mach. Learn. Res., 22:106-1, 2021.
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.
Sébastien Bubeck, Yin Tat Lee, Eric Price, and Ilya Razenshteyn. Adversarial examples from computational constraints. International Conference on Machine Learning, pp. 831-840. PMLR, 2019.
Sitan Chen and Raghu Mcka. Learning polynomials in few relevant dimensions. In Conference Learning Theory, 2020.
Ziang Chen and Rong Ge. Mean-field analysis for learning subspace- sparse polynomials with gaussian input. arXiv preprint arXiv:2402.0888 08944, 2024.
Lénaic Chizat. Convergence rates of gradient methods for convex optimization in the space of measures. Open Journal of Mathematical Optimization, 3:1-19, 2022.
Lenaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems, 2018.
Lénaïc Chizat and Francis Bach. Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss. In Conference on Learning Theory, 2020.
11

---

Published as a conference paper at ICLR 2025
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. In Advances in Neural Information Processing Systems, 2019.
Liam Collins, Hamed Hassani, Mahdi Soltanolkotabi, Aryan Mokhtari, and Sanjay Shakkottai. Provable multi-task representation learning by two-layer relu neural networks. Proceedings of machine learning research, 235:9292, 2024.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances Neural Information Processing Systems, 36, 2023.
Alex Damian, Loucas -ivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529, 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural Networks can Learn Representa- tions with Gradient Descent. In Conference on Learning Theory, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at time. arXiv preprint arXiv:2305.1822 2023.
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.
Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pp. 1887-1930. PMLR, 2018.
Margalit Glasgow. SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the XOR problem. The Twelfth International Conference on Learning Representations, 2024.
Hamed Hassani and Adel Javanmard. The curse of overparametrizztion in adversarial training: Precise analysis of robust generalization for random features regression. The Annals of Statistics, 52(2):441 465, 2024.
Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples not bugs, they are features. In Advances in Neural Information Processing Systems, volume 32, 2019.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. Advances in Neural Information Processing Systems, 2018.
Adel Javanmard and Mohammad Mehrabi. Adversarial robustness for latent models: Revisiting the robust-standard accuracies tradeoff. Operations Research, 72(3):016-1030, 2024.
Adel Javanmard and Mahdi Soltanolkotabi. Precise statistical analysis of classification accuracies for adversarial training. The Annals of Statistics, 50(4))2127-2156, 2022.
Adel Javanmard, Mahdi Soltanolkotabi, and Hamed Hassani. Precise tradeoffs in adversarial training for linear regression. In Conference on Learning Theory, pp. 2034-2078, PMLR, 2020.
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.
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.
Junho Kim, Byung-K wan Lee, and Yong Man Ro. Distilling robust and non-robust features in adversarial examples by information bottleneck. Advances in Neural Information Processing Systems, 34, 2021.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, 86(11), 1998.
12

---

Published as a conference paper at ICLR 2025
Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information thherrettct limit. arXiv preprint arXiv:2406.0158/ 2024.
Binghui Li and Yuanzhi Li. Adversarial Training Can Provably Improve Robustness: Theoretical
Analysis of Feature Learning Process Under Structured Data. arXiv preprint arXiv:2410.08503, 2024.
Ker-Chau Li. Sliced inverse regression for dimension reduction. Journal of the American Statistical Association, 1991.
Ker-Chau Li and Naihua Duan. Regression Analysis Under Link Violation. The Annals of Statistics, 1989.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018.
Arvind Mahankali. Haochen Zhang, 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. Advances in Neural Information Processing Systems, 36, 2023.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 11665)) 2018.
Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In The Eleventh International Conference on Learning Representations, 2023a.
Alireza Mousavi- -osssiii,,, Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient-based feature learning under structured data. Advances in Neural Information Processing Systems, 36, 2023b.
Alireza Mousavi-Hossele Denny Wu, and Murat A Erdogdu. Learning multi-index models with neural networks via mean-field langevin dynamics. arXiv preprint arXiv:2408.07254, 2024.
Eshaan Nichani, Alex Damian, and Jason D Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. Advances in Neural Information Processing Systems, 36, 2023.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 9741-9757. PMLR, 2022.
Atsushi Nitanda, Kazusato Oko, Taiji Suzuki, and Denny Wu. Improved statistical and computational complexity of the mean-field langevin dynamics under structured data. In The Twelfth International Conference on Learning Representations, 2024.
Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: compu- tational hardness and efficient gradient-baat based training for ridge combinations. In Conference on Learning Theory. PMLR, 2024.
Gilles Pisier. Remarques sur un résultat non publié de b. maurey. Séminaire d' d'Analyse fonctionnelle (dit" Maurey-Schwartz") pp. 1-12, 1981.
Grant M Rotskoff and Eric Vanden-Eijnden Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the Approximation Error. 2018
arXiv preprint
Ludwig Schmidt, Shibani Santurkar, Dimitris Tsipras, Kunal Talwar, and Aleksander Madry. Ad- versarially robust generalization requires more data. Advances in neural information processing systems, 31, 2018
Johannes Schmidt-Hieber. Nomparametric regression using deep neural networks with ReLU activa- tion function. The Annals of Statistics, 48(4):1875 1897, 2020.
13

---

Published as a conference paper at ICLR 2025
Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond. In Thirty-sevente Conference on Neural Information Processing Systems, 2023.
Christian Szegedy, Wojeiech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, lan Goodfellow, and Rob Fergus. Intriguing properties of neural networks. In The International Conference on Learning Representations, 2014.
Matus Telgarsky. Feature selection and low test enror in shallow low-rotation relu networks. In The Eleventh International Conference Learning Representations, 2023.
Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Tumner and Aleksander Madry. Robustness may be at odds with accuracy. arXiv preprint arXiv: 12152, 2018.
Nuri Mert Vural and Murat A. Erdogdu. Pruning is optimal for learning sparse features in high- dimensions. arXiv preprint arXiv:2406.0888 08658, 2024.
Guillaume Wang, Alireza Mousavi-Hosseini, and Lénaïc Chizat. Mean-field langevin dynamics for signed measures via bilevel approach. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a.
Zihao Wang, Eshaan Nichani, and Jason D. Lee. Learning hierarchical polynomials with three-layer neural networks. In The Twelfth International Conference Learning Representations, 2024b.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
Jiancong Xiao, Qi Long, Weijie Su, et al. Bridging the gap: Rademacher complexity in robust and standard generalization. In The Thirty Seventh Annual Conference Learning Theory, pp. 5074-5075. PMLR, 2024.
Gan Yuan, Mingyue Xu, Samory Kpotufe, and Daniel Hsu. Efficient estimation of the central mean subspace via smoothed gradient outer products. arXiv preprint arXiv:2312.15669, 2023.
Tong Zhang. Covering number bounds of certain regularized linear function classes. Journal of Machine Learning Research, 2(Mar):527-550, 2002.
Aaron Zweig, Loucas Pillaud-Vivis Vivien, and Joan Bruna. On single-index models beyond gaussian data. Advances in Neural Information Processing Systems, 36, 2023.
14

---

Published as a conference paper at ICLR 2025
A GRADIENT-BASED NEURAL FEATURE LEARNING ALGORITHMS
In this section, we will provide examples of implementations the feature learner oracles introduced in Section 4 using gradient-based training of two-layer neural networks. First, look at the algorithm provided by Oko et al. (2024), which we restate here as Algorithm 2, for the case where a polynomial of degree p. Consider the following two-layer neural network with zero bias
[(x; a, W) Σ (jjo)((w3;
Note that we allow the activation vary based on neuron. Specifically, we let o, o, Ei= BatHe where He, is the jth normalized Hermite polynomial, 35.1 Lid. (ni]((+ for appropriately chosen rt, and 2 Cp, see Oko et al. (2024, Lemma 3) for details. Now, we consider the following algorithm.
Algorithm 2 Gradient-Bass Feature Learner for Single-Index Polynomials (Oko et al., 2024, Algorithm 1, Phase I).
Input: T, step size (1/)
momentum parameters (C),, ra-
1: w i.dd Unif(Sd-1),
aj Uni(((*ra/NH), Vj E [N].
i..d.
2: (o) ,,()))
3: for -....., T- do
4: ift > 0 and is even then
Draw (zxt//2, (00/2)) ~ P
5:
6:
w t- u (ums w Aj [N]
w/ + wf
7:
VjE [N]
8:
end if
w +1 f w/ v (f(x(11/21) a, Wt) - 211/221122
9:
10: end for
return (w.. ....,...
Note that f(w) (I ww (w) denotes the spherical gradient. Essentially, Algorithm takes two gradient steps each new sample, and in the even iterations performs certain interpolation. Proper choice of hyperparameters in the above algorithm leads to Proposition 8.
Next, we consider the algorithm of Damian et al. (2022), which we restate here as Algorithm 3, for the case where g is multi-index polynomial.
Algorithm 3 Gradient Based Feature Learner for Multi- Polynomials (Damian et al., 2022,
Algorithm 1, Adaptedd
Input: {o (e), , (0)
Me Ta 1: aj l.d. Unif( ({ (#ra)), ) w
l.d.
Unif(gd-1))cm a, -a, -N- =wf. Aj E [N/2].
2: a + 1 Euer (())
B + 1 E (1)
MFL
DFL
3: (+) (1))
- (3, (1)
Vie [NFI.].
4: W + -V 1 > RPL
(1)
W y)²
5: w; f- |
ViE
6: return (Wo..... WN)
After performing preprocessing on data, Algorithm 3 essentially performs one gradient descent step with weight decay, when the regularizer of the weight decay is the inverse of step size, thus cancelling out initialization and leaving only gradient as the estimate, Damian et al. (2022) prove that, with a sample complexity of NFL = O + d/(2), the output of Algorithm 3 satisfies
UT UT HUW!
Wi; I0 HUW!
2 1 C. Vi [N),
15

---

Published as a conference paper at ICLR 2025
witi high probability, where H = E E Thus, for full-rank H, the output of Algorithm
satisfies the definition of a (1, 3)SFF. oracle for a constant B >> 0 depending only the conditioning of and the number of indices k.
B ADDITIONAL NOTATIONS AND DETAILS OF SECTION
Throughout the appendix, we will assume the activation satisfies (0) for simplicity of presenta- tion, without loss of generality. We will also assume that
(z1) o & L. (1211- + 12 4-1 + 1) -2/
(B.1)
for all 21, 22 E IR and some absolute constant La. In the case of ReLU, we have and L, 1. For polynomial activations, q is the same as the degree of the polynomial. For of parameters vir (e.g. / = k), we will use Cp. to denote a generic constant whose value depends only on ab and may change from line line.
B. NECESSITY OF Eg NORM AND ASSUMPTION FOR THEOREM
In this section, we demonstrate that both restricting the attack norm and Assumption 1 are necessary for the statement of Theorem to hold.
First, we focus violating Assumption 1. Suppose (0, E). Suppose k 1, y (u1 x). and let FF be the class of linear predictors. Then, the adversarial risk associated to the predictor + (w, æ) is given by
AR((w, :)) = =1/0//1 -i)l 12000 + +2
2 11/2 (w un) [1)|11
From here, one verify that the optimal weight wn satisfies w = ( + La) -uu for some a > O. Note that w" is only in the direction of un if un is an eigenvector of E, which would imply (u1,z) is statistically independent from (v,x) x) for every v E I2 orthogonal u1
Next, we replace the &a constraint for the adversary with an Eo constraint, i.e. we define
AR (f) = E
max [l18|l Sea
(f(x + 8) - -))
Suppose x ~ (O, Ia), and let y = (u1, x) and F be linear as above. Then, we have
AR((w, +) | un || + axc [
2
2 - u |||l
Then, assuming all the coordinates of u are bounded away from zero and for sufficiently small eoc. one can show w" úi co sign(u) ( for some e > 0, which will no longer necessary be in the direction of u1.
B.2 COMPLETE VERSIONS OF THEOREMS IN SECTION
We first restate Theorem with explicit exponents.
Theorem 12. Suppose Assumptions 1,2, and 3 hold. For any e > 0, define ē = recall E = v. Consider Algorithm I with the aDFF oracle, ra = (((e//Va)+1++1/ /a), and TO ((E((E//VE)+t)//) Then, if the number of second phase samples NEA, number of neurons N. and a -DLL error satisfy
and
TFA 2 á
S2t444//) N 2
E1
N+342/H),
<5 E 1+2+1/k) )
1)1k
we have AR(â, W.b) S AR" + with probability at least 1 - npλ where c. >> O is absolute constant. The total sample complexity of Algorithm is given by nhotal = nFA + (DDFL())
Similarly, restate Theorem 5 with explicit exponents.
16

---

Published as a conference paper at ICLR 2025
Theorem 13. Consider the same setting as Theorem 12, except that use the (a,3)-SFL oracle in Algorithm with ra O ((e1) ((lee////ttt V2)++++1/k/ (a.))- Then, if the number of second phase samples RFA, number of neurons N, and a-DFL error satisfy
S
o Bie
N á
<50
? (2241/00.
The total sample complexity this case given by Thttall APA + (S).
The proof of both theorems follows from combining the results of the following sections. Since both proofs are similar, we only present the proof of Theorem 12. The proof of Theorems and can be obtained in similar manner.
Proof. [Proof of Theorem 12] The proof is based on decomposing the suboptimality into generaliza- tion and approximation terms, namely
AR(â, W, b) AR AR(â, W,b) b) AR(a", 6) + AR(a", W, b) AR",
where a' min Srl /VV AR(a, W, b), thus see the first term above generalization error, and the second term as approximation error.
From Proposition 22, have AR(â, W. AR(a", W. c/2 as soon Q(r1(et rt) (recall that y = 1 here, since we are considering the ReLU activation). For the approximation error, we can use Proposition 36, which guarantees there exists a" with ||a"|| ra/vN such that AR(a", W. b) AR S e/2 with S (((ei/VE)tttt// /a)tt///////...///)))/e)) as soon
s 5 k+242/46
k+342/K)
and
N (k-11) ((k-1)/2a
provided that we choose = ((ei(ts/V3) 11/K). Plugging the value of ra and rs in the bound for n completes the proof.
C GENERALIZATION ANALYSIS
We will first focus on proving a generalization bound for bounded and Lipschitz losses, and then extend the results to cover the squared loss. In this section, we will typically use n to refer nEA. the number of Phase samples.
C.I GENERALIZATION BOUNDS FOR BOUNDED LIPSCHITZ LOSSES
Let us focus general Ce Lipschitz loss {(f(;;, a, W b) y) for now. Later, will argue how to extend the results of this section to the squared error loss. Our uniform convergence argument depends on the covering number of the family of adversarial loss functions. Let θ R be the set of second layer weights, to be determined later. This family is given by
C(W,b b) = {(. v) T max ((f(x + 5: a. a. y) a [|ôl|ş=
For brevity, we will also use denote L(W, b). but we highlight that W and b are fixed this stage. We define the following metric over this family
V E C(W,b). de((,T)2 = Σ (1
(0))
We say S C L is an 6-cover of for every i E L, there exists /' E S such that de((,T') i') < c. The e-covering number of L is the least cardinality among all e-covers of L, which we denote by E(C,de, e). Note that since C is paramterized by a, constructing such covering reduces to constructing a finite set over O
Therefore, we define the following metric over θ,
Va E,, dola,a')² Σ max (f(z +000 a, W. () b) ((x(() +800aa W, 3.
1809
Sc
17

---

Published as a conference paper at ICLR 2025
We can similarly define the e-covering number of Q with respect the metric de as C(o, de. c).
The following lemma relates the covering numbers of and e.
Lemma 14. We have (c,,d.,) & C(e, de. e/Ce) for all e >> O.
Proof. We will use the following fact the proof. For any F1, F2 S IR, we have
max F1(61)
max F2(62) 62
I Fia [
(C.1)
BIES
BES
This is true because
max Fi (ō1) E
max (6)) S max (6)) - F2 61ES
and the other direction holds by symmetry. This trick is used relate the adversarial loss to its non-adversarial counterpart, e.g. in Xiao et al. (2024, Lemma 5).
Now, we will show that an c/Ce cover for θ implies an e cover for C. We will supress dependence on the fixed W and b in the notation. Let Sg be an c/C₂ cover of θ with respect to the de metric. Then, we define S via
S (<x,)) + max {(f(x + a; - y) : a Se).
[
To show S an c cover of L, consider an arbitrary E(x, v) = max {ff 8; a) - y). Suppose a' is the closest element to a in Se. and let 8l (xr = max 18 {[ + 6; a') y). Then,
de dc(E,8))
max {(f(x + 115e
o(: + (0)) max {(f(x 58²0 : a') - y(0)
18
max (e(f (( - {(f(z + 810) 1y 10) |000E
C7 x
Σ
(f(x + 5 f(x + a())
max
< de a')
where we used (C.1) for the first inequality.
To construct an c-cover of O, we depend on the Maurey sparsification lemma (Pisier, 1981), which has been used in the literature for providing covering numbers for linear classes (Zhang, 2002) and neural networks via matrix covering, see e.g. Bartlett et al. (2017).
Lemma 15 (Maurey Sparsification Lemma, (Zhang, 2002, Lemma 1)), Let H. be a Hilbert space with norm II- let E H be represented by Σド= Uy where oj 0 and ||by|| S b for all j E [m], and a = ΣTT Oj S 1. Then, every k 2 there exist non-negative integers ky...., km such that
m ΣØ=1 k, < k and
1 m
ab²
||21| k
ky
k
1=1
Then, we have the following upper bound on the the covering number of θ
Lemma 16. Suppose a satisfies (B. θ = {Ilall, S <rah and additionally ||w:|| S rw and [ba| S rb for all 1 <i S N. Then we have
C.T log N ( + 222g 22222 + rb + T(2) x + n220 + r
log C(),de,))
c2
where (())
maxigisN EE m.)
W.x
18

---

Published as a conference paper at ICLR 2025
Proof. Given some positive integer l >> O, let Se be given by the following
N Σ
So = { [a (k1 - ki,ko h2 12 AN - KN)T= Vi, Ks k 2 0
(
k4 k..
1=1
Let X, A Raxd be the matrices with (x,) and (6,) as rows respectively. Let A o((X W 1,,b²) E REXN Then,
(f(x0))00) a, W b) ( -0 a', W, b))² 1A(a d'll²² 1
Σ A( -b3
n
where A, = o + A)w, + 1,b;)) is the ith column of A. We are going to choose a' from Sg. To that end, define
Á, = sign(a:) A,,
By Maurey's 's sparsification lemma (Xiao et al., 2024, Lemma 13), there exist k; 2 0 with Σx= k such that
N
N
2022
Σ
ag
A
Ta
k
k
k
where ||A.]| S 6 for all i. We will then choose
sign(a;) 2 0,
0,
sign( (a) 2 0, sign(a:) 0
k; [0, 0,
sign(a;) < 0
Therefore, we have >N 1==1 = k. Finally, with the constructed (kt) and (k{), let
a' (k: Ky KN KN)T
and also note that Limla, such that
N a.A, Consequently, given a, we have constructed a' Se
N A, - -4) S
12
n
nk
Next, we provide a bound on b. By the assumptions on a, we have
| LC (1
X w: 24
A
28 nb
|| w 1| #++ 1 nb3))
1 SnCgLe TW. x
+ -2
+ rb
24
+ (2)
W
+ 2255 + r3)
Consequently, we can choose
C2L3r2 (g)) (T
W.. + 2-22
+ Tb
+ T (2)
W X
+ +r255 + r
22
Finally, we need to count |Se]. Note that
(2N +
e(2N + 11-11)))
11)2 < (3eN)*
[ol =
S
k
which concludes the proof.
We can now turn the above covering number into Rademacher complexity via a chaining argument, follows.
19

---

Published as a conference paper at ICLR 2025
Lemma 17. Let R(C(W, b)) denote the Rademacher complexity of the class of adversarial loss functions L(W, b), defined via
R(C(W, b)) = E
sup aco
&
{(f(x)()) + 0 a, W. b), y 1/0) |
max
where E, are ii.d Rademacher random variables and θ = fa: ||a||; S raa For simplicity, assume
Cr, 2 1. Then we have
C.CyLor log log N
E
Tw.x
1
+ +r + +rg +
1:22) W x
+ Tww + Tb
R(C(W, .)) S
Proof. Let R, (C(W, (WVb) b)) denote the empirical Rademacher complexity by
(C b)) = EE
sup acB
&
max {(((x( + 50) a, W b), y 0
|
where the expectation is only taken w.r.t. the randomness of and is conditional the training set. For simplicity, define
B B=CoLo (VIW.x r 19 + rg + TW.x + rwwE - Tb)
Then, by a standard chaining argument, we have for all a > O.
log C(C, de
Rr (CW,, b)) S a
r=a
n
Cer B Blg.Boog.. N
S a +
log
By choosing 1/ we obtain
CeraB
R,(C(W b))
vn
Taking expectations with respect to the input distribution completes the proof.
Note that it remains to provide an upper bound for ( W. introduced in Lemma 16. This is achieved by the following lemma. Lemma 18. Suppose
|
I
Tw- Then, for all q > 0 and N >> have
W 20))) 24
E max
Cpr 24 (log N)
ISJEN
where Cg is constant depending only q-
Proof. For conciseness, let Z, = 112 Em= (w,, 20)) 29
By non non-tegatyvyty of Zy and Jensen's
inequality, for all t we have
B[zz]) 1/1
max ISIS
Zy
max ISINN
Z
Σ B
N
max E
[[Z)]) 1/2.
ISSN
1=1
Further, by Jensens's inequality
E
1
IE
1
Wi x0(1)
27
200
E
W (+))
i=
(Crw.)
2at (2Tt) 02
27T
20

---

Published as a conference paper at ICLR 2025
where C' 0 is absolute constant, and used the moment bound of subGaussian random variables along with the fact that (wy, x) is a centered subGaussian random variable with subGaussian norm Orrw). As a result,
E
max ISJSN Zj
S C MNNNNN x Car (log NY
where the last inequality follows by choosing t log
As consequence, if the loss is also bounded, get the following high-probability concentration bound. Corollary 19. Suppose | S for all C(W, b). Then, with probability at least 1 we have
Σ (10 y00) 2 CraRlogn log N + Br log(1/6)
E f v)
sup (.C(Wb))
Vn
where
R = Cobo (logg)
9/2 N + 5²) + + + rr 1/2 +x) +
C.2 APPLYING THE GENERALIZATION BOUND TO SQUARED LOSS
To apply the generalization argument above the squared loss, we bound it with a threshold т, and define the loss family
C. (W. = (( + { max (f + S:a. b) - y)² Ar: :EE0)
[
We similarly define AR, and AR,. Recall that our goal is to show
AR(â, W b) & AR(a, W. b) + (n, N.d).
We readily have (R, W. b) AR(a, W. bb. Further, Corollary 19 yields
AR, (á. W. b) AR, (ā, W. b) x VTraRlogm log N
log(1/8)
+ TV
Vn
n
with probability least 1-6 8. Thus, the remaining step is to bound AR(â, W, b) and AR(â, W b) with their clipped versions. To do so, we first provide the following tail probability estimate.
(2 are non-negative random variables with subGaussian norm r. Then, for any q > 0 and T 2 r where Cg is constant depending only on q. we have
Lemma 20. Suppose
N
2/8
P
< exp
p2
j=1
where e 0 is an absolute constant.
Proof. For any /. 2 1, we have the following Markov bound,
( E 23)
E
Σ
P
P
7
where the last inequality follows from Jensen's 's inequality. Further, by subGaussianity of 2j, we have E 1 02//2 where C > is absolute constant. As result,
N
(Cr9qt)a1/2
T
N
==1
21

---

Published as a conference paper at ICLR 2025
The above bound is minimized choice of in the above bound yields
2/4 Crage
Note that 2 1 requires T 2 Car9. Plugging this
72/9
P
27
S exp
2Crree
j=1
which completes the proof.
Lemma 21. Suppose Assumption holds. Lei θ = {a ||a]| I /rr/v NN, || S Tw. and [bx] S Tb
for all i [N]. Assume satisfes (B.1). Define 1 E, and let
x C + r + o + r Cp.
where Cg and Cp are constants depending only on q and p respectively. Then, for all
T c x 12 log n
v log" n 5 }
we have
AR(a, - AR((a, W.
W.b) b) S AR (a, b) - AR, (a, W. b)
Cx exp [2
)) e 12/0r2 r
72/
with probability at least 1 å uniformly over all E O.
Proof. Since W and b are fixed, we use the shorthand notation f(z; a) = f(x; a, W,
In the first section of the proof, we will upper and lower bound AR(a, W, b) with AR,(a, W, b).
Note that the lower bound is trivial as (a, W, b) I AR(a, RR(a,W,bi b), thus we move on to the upper bound. Let
((x.v)
max (5 + 0 y)?.
|8 Ie
Then,
AR(a, W, b) = E =E((lz 10(((( 3) s <T] + E
(( ((z, will ((x, v) > >T]
T]]
[(((,.))p] w)] 1/2 ((((, 27) 1/2
S AR,(a, W. b) + E [(0,
Further, we have the following upper bound for the adversarial loss
i(x, v) max (ʃ + d) y) =)
max [(x 2) +2²
[|
< max ||a|| (W + 5) b)1l y2
y2
||
S L
10 m) 24 r
2922
+ v
24 +
1
(w. xx)
22,62 + r8
+ 12
N
N
Moreover, by Jensen's inequality,
1 N
E
W x) 27
< E
Σ
W. xx)11 19
N
N
< (Crw) (R < Car
22

---

Published as a conference paper at ICLR 2025
for all 9 O, where C is absolute constant and we used the subGaussianity of (wy, x)) to bound its moment. As a result,
B ((1,)))) w/
1/2 S 2 21 = 22 + + + + + r3) + +E 1/2
By assumption 2, we have E [/^]] 1/2 C..
To estimate the tail probability of f(x, y). Using the assumption on 7 and the upper bound on E(xx y) developed above, via a union bound have
N
L
P ((((..))27)
SP
Σ
x)
+
W
xx
+
T 1) 2
N
2
==1
/=1
I N
P
Σ
W x))
24
N
+ P
w x) 2
+ +P('²Z 6)
N
N
j=1
5=1
-CT
1/9
2 exp
+ P 2 T)
where we used Lemma 20, the fact that (wy, x>) is subGaussian with norm O(rw), and that q 2 1. Furthermore, using the moment estimate on y in Assumption 2 along with the technique developed in Lemma 20, we have
P( 2
S Sexp(-et lt)
for r 2 Cp. where e is an absolute constant.
As result, we obtain
AR(a, b) - AR-(a, W, b) $ x exp (-
c72/9
+ exp( +xx((-c7/1))
T 2.29
m
for all a E e.
In the next part of the proof, we will show that with probability at least 1- we have AR(a, W, b) =
AR,((. (a. W, b) uniformly over all a. Note that this is equivalent to asking (()( ()) I for all
1 s i n. For any fixed i, using the upper bound on ((x(1), y we have
1 N
((((x(0))0) 2 S. P
W
x))
+
Σ
(n xx)² + + 2 2
N
Ø=1
j=1
1/9
S exp
+ exp -cT1/p).
L3
Consequently, by a union bound we have
- 1/0 1/4
P max ((x(0, ( (0) 2
+ exp (corrll1g -orl/)))
S n
exp
Isi5n
2/4 L 12
Choosing
t 2 L3 log"
Vog
v
8
with sufficiently large constant C' ensures the above probability is at most 6, finishing the proof.
We now ready present the main result of this section.
23

---

Published as a conference paper at ICLR 2025
Proposition 22. Suppose Assumption 2 holds and satisfies (B.1), @ {a |||| I ra/ (N/VN). wwill $ 1, and bil S rs for all 1 N. Let
x = 01
+ +rb) + Cp.
where Cg and Cp are constants depending only on q and p respectively, Then we have
AR(ā, W b) min AR(a, W. b) Ô age
with probability at least On for some constant c > O.
Proof. We can summarize the generalization bound of Corollary 19 as
log (1/)
(R, W. b) AR, (á. W. b) - n
+ +TV
n
where
x Car
24 + ) + Cp:
obtained from Lemma by letting Tw 1. Thanks to Lemma 21, we arrive at
log(1/o)
1/4
xe -0(1110)
AR(â W, b) AR(a, W) S Ô
n
Note that x Z L. Choosing т Cx log (sn//) with sufficiently large absolute constant >>> 0 satisfies the assumption of Lemma 21. By letting d = n n- for some constant e >> O, we obtain
AR(â, W, b) - AR(a, W, b) s
which holds with probability at least 1 - n- over the randomness of the training set.
Recall a² arg min AR(a, W, b). Similarly, Lemma 21 guarantees
AR(a" W. b) AR(a", W. b) S ō /n
on the same event above. Finally, we have AR(a W,b) < ARi(a, W. b) by definition of à, which concludes the proof of the proposition.
D APPROXIMATION ANALYSIS
Let my T denote the projection of E Sd-1 onto span(u1,.... ||Wl||
us n (if [[U
0 we can simply let HI w = uz). Suppose (w, u) 2 1 - Ç for some E (0, 1) and u E span(ui..... Ux) with |ull 1. Then, have the following properties for this projection:
w w) 2 - Ç. ww W < V2C.
Let h: I + IR be the function constructed in the proof of Theorem 1. Then,
AR E max (h(U(x + 6)) y)² [iol|se
Let us denote f(x) f(x; a", W. b) for conciseness. Then,
AR(a", W.B. AR"
E
max (f(x + 6) -)) max (h(U(x + 8)) y)² [8
I
max {f + +) - (NL(U(x + + - - [|
E
max M Se
(f(x + 8) n(U(2 + of) + 5) h ( 0)) - 2y)
24

---

Published as a conference paper at ICLR 2025
Let Hl W (Ilywi...... WN Then, have the decompositions
f(x + a; W, b) f(x 0: W. b) f(x 8;a", N, b) + f( Il Wbb).
and
2 =f 8: a"" W. - f(x + d; IT W. b) f(x 8;a", TT W, b) ((U 8))
+ +2h(U(( + 6)) 2y.
Plugging this decomposition into the above and using the Cauchy-Schwand inequality yields
AR(a", W.b) AR & (VE + VE2²² + V +
(D.1)
where
& = max (ʃ + 0; Tly W. b) 1(U(xx + 61))2 Sc
(D.2)
max (+ + 0; 1l W, b) f(æ + d: a" W, b))
8 = E
(D.3)
& = 4E max (h(U y)² = 4AR".
(D.4)
[|
Under Definition 3, have of good neurons S to work with. To continue, introduce similar subset of good neurons under Definition 2.
Definition 23. Suppose the weights W = (wi..... w are obtained from the a Q--De DFL oracle of Definition Fix maximal EV/2(-packing of with respect to the Euclidean norm, denoted by (0.) Define v, = Tw, for all j E [N], and
Uw,
= -i 11,, - 0.ll S v2C),
for all i E [M]. Note that (Si) are mutually exclusive. Define = U. Sy. By upper and lower bounds the surface area of the spherical cap (see e-g. Wang al. (2024a, Lemma F.11)), there are constants Ck. Ck > 0 such that (e(1/C) (k-1))))1))))))))) 1)/2 & M S Cc((1/(((-1)) Therefore, using Definition 2, we have [S|/ 2 ((a))
Note that when considering the (a,3)-SFL. oracle, leave S unchanged from Definition In either case, for every j & S, we will choose aj = 0. Then, we then have the following upper bound on &g.
Lemma 24. Suppose az for j & S and ||a"|| s Fa/ Pa/V/b VIS.. Then,
E max (f(z + 0; H. b) f(x + 8: W. b))² s 1.C 2ta-1) -10 [olise
where Cg a constant only depending on q.
Proof. To be concise, we define ãs = x + ö and hide dependence on a" and the following notation. By ggeudo-Lipschiznne of o,
f(ă;; Tl W) f(ăs: W) Σ
( ms) + b,) - o((nn) #s) + b,))
JES
Ln
Σ
at (K 38) + b,,8 + ((w0 2s) b 8- + 1)I w8)
jeS
Let
A, = W,, xs) + bbb + +(w,, 2s) by) 8-1
and
B, (Tuw; w, Ka)l
25

---

Published as a conference paper at ICLR 2025
Then by the Cauchy-Schwann inequality,
13 & |S|
82 S 12 E
max
Σ
a5 1,8,
-
E
max 18
Σ
A
[8 Sc
JES
JES
L 0
Σ
IE
max A3 18815
|S|
JES
1/2
1/2
E
max [18|scc
A
E
max Bj
SS|
JES
Additionally, we have
A₁ S Co (1 2))1 + (20, x))1 x2)1811 5 +rg 10-1 + +1).
max |
and
max Bj S T wyl + (TT wyj; x) Se
Further, by Assumption 2, for all v E R2 (v, x) a centered subGaussian random variable with subGaussian norm (1( therefore E (V, x)11 < C/|le|| for all q > In summary.
1/2
1/2
E
A3 S8
max
- Cg(1+
2(1-1) = 2(4 10--1)), 1) and
E
max B [|ö]|
(1+
where we used the fact that ||Iuw wjll? s 2 for all j E S. This completes the proof.
While the term & defined in (D.2) is an expectation over the entire distribution of æ, most approxi- mation bounds support only a compact subset of TR.d. The following lemma shows that approximation on compact sets is sufficient to bound E
Lemma 25. Suppose a for j & S and ||a" 1| S Fa/ Fa// Further, suppose T= 1 V 2z. Let
Capprox sup ||Un||Sr
[s(x; TT W, bbb h(Ux))].
Assume h satisfies [h(2)] S (1 + 11 all z E R and some constant p 2 0. Then,
E1 < approx + (L.C + 20 - 2,0) + 13, + = -006))
Proof. For brevity, define
As = (f(ēs; 1, W, b) - h(U(( -(U(x++0)) +81))²
where T8 x + 5. Then.
E
max Ao
< E
max ||
As
| S Sra]]
+
max As[[[
> re]]
1/2
1/2
approx
+
E
max [
A
E
max T[|[Uä||| > >r;]]
[|
1/2
apppoxx
+
E
max Se
A
P([ || > rr -==-))//2
1/2
P(|UaO) > 3// ==) 1/2
approx
+ E
max [8]
A
Furthermore, we have
E max ||
A
$
max f(ăs; a", [8
anaaa""TTo Th W, b)d
W.b)
+R
max 1(U(x (8)+
[
26

---

Published as a conference paper at ICLR 2025
Recall the notation vy
Uwy [|/w/l
and Ux. Then, by Cauchy-Schwartz and Jensen inequalities,
E max f(0s: a" HI W. 6)²
max l|a
() ()
1
E
max [e
Σ
((0), U8) + bb,)*)
[8|
JES
1 Cg
Σ
(Uj 2) 14 + c + 44
|||
3ES
< Lo + Eth
+ ).
Similarly we can prove
E max huU h(u 8))²
(p +
||8||
In summary,
1/2
max A [8
X (1 + +r²²)) + (1
+
Finally, the probability bound
P(U
2 < -
follows from subGaussianity of x and the fact that k = O1).
D.1 APPROXIMATING UNIVARIATE FUNCTIONS
In this section, we recall prior results approximating univariate functions with random biases in the afinite-width regime under ReLU and polynomial activations.
Lemma 26 (Damian et al. (2022, Lemma 9, Adapted)). Let a be the ReLU activation, a Unif({ - + +)), and b ~ Unif(-ro rb). Then, there exists f {-1, x [-rb, rol + R, such
that for all | S Tb we have
Ba, (an,)rras())0( b)] h(z).
Additionally, if h is a polynomial of degree 8, we have sup,, f(a b) S (e
Proof. From integration by parts, namely
[( / h*(b)(-22+ Bjdb
(2r a)h" (b) + +b]] +b]=
= =1((ro)(=2 -f- 1((b)db
= =N'(ro) + rb) h(2)) h(rr).
Therefore, it remains to approximate the constant and linear parts. It is straightforward to verify that
6b
Ba,b
aaaz + b) =, Ea,b (az + b)] z.
rb
Thus, we let
ah'(ro)
36(h'(ro)r -((ro)
f(a, = - ab)"(b) +
Tb
which completes the proof.
Furthermore, we have the following result for infinite-width approximation with polynomial activa- tions.
27

---

Published as a conference paper at ICLR 2025
Lemma 27 (Oko et al. (2024, Lemma 30, Adapted)). Let a be polynomial of degree q and suppose b Unif(-ro::0) Tb) and is a polynomial of degree such that q 2 p. and in particular satisfies [h(2)] (1 + (2|?). Suppose ть q. Then, there exists a function f 1-Tb, Tbl R such that
Es (2) + b)] = h(2), 42 E R.
Furthermore, we have If(=) C. for all where Ca, only depends the activation and Li-
Proof. In order for a to approximate arbitrary polynomials of degree at most g, it is sufficient to show that a can approximate at least one polynomial per degree, ranging from degree 0 to q. Defining the corresponding polynomial with degree as gi(2), then h will be in the span of {9: hi More specifically, suppose h(2) Ej= agz', and gi(2) Ej=o n.,z.. Then there exist ( i=u such that
Σ (.gi()) = 200
Σ aayo 3=0
j=
Indeed, we can let 3, 0 for all >> Additionally, note that 7/. * 0 for all i S q by definition. Therefore, the solution to the above equation is given iteratively by 8p = Op and
Pp-1 Op-j EiEo 7p ip ipp-ii---pp i8p-
for 1 S S Importantly, [3:] for all can be bounded polynomially by {ayby (mu))s and {ng Ji. Further, |au| can be bounded polynomially by La for all i. Thus, it remains to construct {9.}.
Following Oko et al. (2024), we define
(x(2) f a(2 + b)db. -4
is straightforward to verify that Og has degree (exactly) q. We then iteratively define
9q- g2-1(2) = (i- 1> + - 9 ()-(1-))22), vi
Using the definition above and by induction, one can verify 9i has degree exactly Furthermore, expanding the definition above yields
9n gg-1(z) = Σ (2+0((++))+ as I" o(z + b j)db,
= ai 3=0
2=0
where Cij = (-1)i-3 ()) i.e. the coefficients that satisfy (a - 1)' "}===zzo0 Ci,z2. In particular, we can write
0g-i(z) = Σ Cr L' a(2 + b)db = Eg 2rb Σ
T[- q + j s 510(2
-gts
1=0
1=0
Therefore, define
f(b) = Σ Bp- Σ -c,1l--++s0 5.
1=0 1=0
which completes the proof.
D.2 APPROXIMATING MULTIVARIATE POLYNOMIALS
We adapt the approximation result of this section from Damian et al. (2022), modifying the proof to be consistent with assumption the first layer weights.
First, we remark that for any fixed u E g and any degree 0 < s S p. we can approximate the function z (v, 2) with random biases established by Lemma 26 for the ReLU activation and
28

---

Published as a conference paper at ICLR 2025
Lemma 27 for the polynomial activation. Therefore, main effort will be spent in approximating polynomial h(() ()) using monomials (v, x)" Note that we can represent h by
h(2) [T
where (()) is symmetric tensor order 8, and we use the notation
T [z _ = vec(TT(>) vec(2 00)= f1.0.1,=1
T Rig ......0 Ri,
The approximation result relies on the following fact. Lemma 28. Let
Then, the matrix Bowr, [vec (ee (ve((o)))ec( )ve((=*)) is invertible.
Proof. Let T an arbitrary symmetric tensor order with |T 1. We need to find a constant Ca,k such that
(e)(T) Evene
vec(T) 2 Ca,k-
Note that
T 02
vec(T) Evwra [vec( 3 (vec(e (eec(ve))eec(""
vec(T) Ev To0000 E (V
| 28
Furthermore, (Damian et al., 2022, Lemma 23) implies that Ew-,A(0,1)) T[w0113 2 c for some constant d,2, > Therefore, for any r >> 0, we have r]] + (w [ Sr]]
dg,t- Note that the first term on the LHS above can become arbitrarily small by choosing r sufficiently large (depending and k). Thus for sufficiently large have
E Ew TTw 8 [
w|
r]]
Finally, we have
T Tw 0u7
1 22s
c 2p
(w-A( Therefore, taking Ca.k
2
Be- [Tw [1l||| w|| S r] >
u 26
completes the proof.
2
The following lemma establishes how we can use monomials of the form (v T. 2)" to approximate each term appearing in h(z). Lemma 29 (Damian al. (2022, Corollary 4, Adapted)). There exists S R such that for all z E IR* and non-negative integers & 2 0,
-t- 2)) 1
Further, If(v) < Ck,.
(T for all v gk-1
Proof. Note that by definition, (0, 2)² ve0( (oloBe)Tvocce veo(2 Therefore,
f [())(o, 2)" drx (1 ((o)vec(
(d)x (d))((())))))
veo(2
We need to match the first vector on the RHS above with v8 thus our choice of is
f(v) @* Ev-ry [vec(v @x "eec Sx) ) vec(T((s).
The proof then completed via the lower bound of Lemma 28 which gaurantees the existence of some constant C, > 0 such that Amiu (E.-. |vec( y vvec @x))
C,
The above result along with the univariate approximations proved earlier immediately yields the following corollary.
29

---

Published as a conference paper at ICLR 2025
Corollary 30. Suppose h is a polynomial of degree p denoted by h(z) = =E [20). Further assume the activation a is either ReLU or polynomial of degree 9 2 p. Then, there exists h: -55 x [-Tb, r6] - IK such that for every || S rb. we have
Js (x0-1x(-romm))
h(w, (v, z) ())rx = h(z).
Furthermore, [u(u,b) b) CK. max,g
[T
for the polynomial activation and [(v,bbi s
(p-2)v0 | Carb--21v0 Proof. First, we consider the case where we use polynomial activations. Let
for the ReLU activation.
h(v,b) Σ (2..2))....0)))
for (fi,,) and (f2.s) which now determine. We choose f2.. according and Lemma 27, then
f f2.. (b),(b)(((),)) b) (v, z)"
for all || I r2/2. and (b)] & C.gg for all Then, we choose fi.. according to Lemma 29, which yields
-g--ix(-smmm (1. 2) + b)dr% (w)db = / fi,,(0)( 2) (T) = [20])
for all ||||| < rb/2. Additionally (v)) s C=.e| which completes the proof of the
polynomial activation case.
Now, consider the case where we use the ReLU activation. Let
p h(v,b) Σ (v.
where
(v,b) (1. + f (-)((-)))aa((1,))))
with fi,e given above and f2.0 introduced below. Since and v have the same distribution, we have
f 3o (v, b),,))(((, 2))dbdrx( / 2 (11a(1) + ( (.((-))22.(-11 ( a((1 2)dbd7.(w)
f (1) ( b--ro
+ bjab (om(()))
+ b) +
/ fl.( dr((v) T
As result, it suffices to choose fe.e according Lemma 26, which completes the proof of the corollary.
As last step in this section, we verify that one can indeed control max.<p [T>) || with an absolute constant when the minimizer of the adversarial risk.
Lemma 31. Suppose F is the class of degree polynomials on Rd. Let H = {x +
Elf(x) Ux 2] f F}, and define
h arg min E.
max (N' (x: + y)2 |6
MEH
Denote the decomposition of by (xx))==' T' Then. || S Ck.y where Ok is a constant depending only on and the target second moment EE [3 (thus an absolute constant in our setting). As consequence, we have [h(x) < (1 +1+ ||x]]") for all z E R* where L₁, > 0 an absolute constant.
30

---

Published as a conference paper at ICLR 2025
Proof. By comparing with the zero function, have
max (N(U(x+ - -1)² SE('"].
[[(h(V)) - y)² = E
|
Furthermore, by the Cauchy-Schwartz inequality,
B [ y) E - 1/2 12 [E
Combining the two incqualities above, we obtain E [[h(Uxx) - 4 Let = Ux, and let μa be the marginal distribution of z. Then
E [[((2) = [M())
dez C.
Further, by subGaussianity of a and subsequent subGaussianity of z, we have aV (2) S C oo for all when Ck2 C sufficiently large constants depending only k. Therefore, E.-,)o...L. [b()²]] - B[y
dus
The proof is completed by using the Hermite decomposition of h.
D.3 APPROXIMATING MULTIVARIATE PSEUDO-LIPSCIITTZ FUNCTIONS
We now turn to the more general problem of approximating pseudo-Lipschitz functions. Specifically, when satisfies Assumption 3, functions of the form h(z) = Eff()) Ux = z] will be L-pseudo- Lipschitz. The following lemma investigates approximating such functions with afinite-width two-layer neural networks.
Lemma 32. Suppose
R* + is I.-.hihchitz on || S r. and the ReL.U activation. Then,
for every A 2 Ck. there exists h: s Tb. ro] - IR such that
A
A
h(2)
-St- (o, (o)r S Cr
log
Lr=s
Lr=
Tb
for all || Srz. Furthermore, we have [(( s LC&I(A/Lr,)a/ /rr all v and b, and
| (((rrd) (o) <
C2A2
Proof. Lett 1 = (2 TT E IR++1, By (Bach, 2017, Proposition 6), we know that for all A 2 Ck. there exists Sh + IR, such that [| I and
h((z) / p(o)or
(0,2) d < CCLLT
A
log Lr
ggt for all ||*|| < Te. Furthermore, the proof of (Moussvi-Hosseeni et al., 2024, Proposition 19) demonstrated that
Fs
A
[p()) S CeLr= Tr=
Vōesk.
Let 0 (WT ( be the decomposition of é into its first and last coordinate. Then, we will use the fact that for 0 ~ Unif(S*) when conditioned on 0k41. by symmetry [[eaa] on 1 1
Pua is uniformly distributed
other words, e ~ Unif(S*-1) 1) and ~ pk+1 independently, where we choose Pk+1
such that & has the marginal distribution as ētt1- Since the marginal distribution of VI+b2 i is given by dPP(k+1) (1 ) 2k) have pu+t (6) Zx -
31

---

Published as a conference paper at ICLR 2025
where Z& is the normalizing constant. Then, ū T(v. b) is distributed uniformly Sh, where T: sk-1 x R g is given by T(u,b)= VI+5" (vT, .). As result,
I p(o)o (0,,8) Tz
drx+1(0) / p(T((),))) (u, 2) + br: r3 r.v1+8
dr (T((d)dpu+1(0))
p(T(u,)) b))
Zk
- T2 / 1V+++ 12 (1 + b2)(K+1)/2 ((v. x) + brr (v)
rT b/r=)/
Zz
of (v, 2) ()ddx
(r2 + 62)(8+2)/2
JK 1xB
Therefore, our choice of will be
h(v.b) Zn [TT [r? + 12)(++2)/2
Next, we bound the following error term due to cutoff of bias,
= Zz (
(T((((,,/rz)) (r2 62) ((0, 2) ()dx(()dd
We have
A 2k
(= + bl)
8 E Cr
(T - 62)(k+2)222 db
Lrzs
$ CaLrz
(r2 + b2)( db
Tr=
b/r
bb+1
b/2rg
Tz
S C
Tb
Finally, we prove the guarantees provided for h. The uniform bound on [ b) follows directly by plugging in the uniform bound on p. For the 2,2 bound on h, we have
Ja- [ b) <
h(v, (2rb)) (v)db
23 P(T(p (T((bb//r))))) dra(d)db
(r2 02)442
Zip(T(U 6))² drx r2(1 i2)
Zu
/
PT 8))
dr. (1) (b)
r
(1 (1+ ++b2)(++3)/2
/
(kt3) d7
Za 11 (
<
p
r3
completing the proof.
32

---

Published as a conference paper at ICLR 2025
D.4 DISCRETIZING INFINITE- WIDTH APPROXIMATIONS
In this section, we provide finite-width guarantees corresponding to the infinite-width approximations proved earlier. Define the following integral operator
f h(v, (), z) ())rx(o)db.
Th(z)
(D.5)
The type of discretization error depends on whether we are using the a-DFL or the (a.,)-SFL oracle. We first cover the case of a-DFL oracles.
Proposition 33 (Approximation by Riemann Sum), Suppose a satisfies (B.1). Let (wi..... wN) be Uw the first layer weights obtained from the a RDDF oracle (Definition 2), and define v; = [| for
i [N] Suppose (by))E(N] i..dd Unif( Tbe rb), and let llAll = / (((,,)) 6)]. Then, there exists a' such that a; = 0 for j & S and |až] < C₂ Ci (oter///ooo////////())) for E S (where S is given by Definition 23), and
[ojo((o, 2) bb)) Th((z) S &cglbblockoorg (raV
Tb Logg(N/))
(D.6)
((k-1)/2.N)
jES
for all z E RR* where ||||| rzz S rb, with probability at least 8 over the randomness of biases.
Proof. The proof is multivariate version of the argument given in (Oko al., 2024, Lemma 29). Let (ø,
) M be the maximal 2/22iin of @ from Definition 23, which also a 2 of s Recall from Definition 23 that M < (+())((-1))2
For every i [M], define
= [M] loy B,1l S 2()-
Note that by definition of packing and Definition 2, each vj can only belong to exactly one of S; when j E S. meaning that (S:) are disjoint and Uketan S = S. In particular, |St|/N 2 clk-1)/2a, and // < 1/M I <(( /ck-
We want each group of biases (by)jes, cover the interval [-ra, mx] We divide this interval into 2A subintervals of the form [-ro(1+ by a union bound, the probability that there exists some subinterval and some St such that the subinterval contains no element of {by : E S is at most 2A EM 24 15A., Thus, taking
4)
for 0 < 1 S 2A-1. When by Lid. Unif(-ro, 7o),
A < TS. M/35_ for all i E [M] guarantces that all subintervals have at least one bias from every S; inside them with probability at least 1 8.
18.1
SK-1 the projection onto the packing, i.e. [T(v) arg min 111 0 0.ll. Further, we define II, [M] x [-rb, ro] + [-ro, rb] by Il(i, = arg min (by)) b byl. Tie braking can be performed by choosing any of the answers. By defini- tion, we have |/v 11, (v)] s 2/2(, and additionally 66 1l2 6) < ro/A for all i [M] on the event described above.
Next, define IT Sk-1
We are now ready construct a". Specifically, let
I h(v b)[[( = (I),),) b, (b) if S; for some i, j $ S.
0j
0
Note that by definition,
M ΣΣ Ili 11)11,00),, b, = 112 = 1,
1=1 jelj
For conciseness, we define E(v, i.j) = [[i (w). by [I2(i, b)). When j E S,, on the event E(v, i. we have
|/ wyll
|| .: || uy s 332
33

---

Published as a conference paper at ICLR 2025
Moreover, since P Pown ||| -ū,ll S 2v2[ S (/5(k-1)/2, for S we have
lašl S Callhilac ( (k- 11/2 Gellilloorooo (N/(1
A
a.N
As result,
Σ +b)) Th(z))
x) ((e.,))((0,2))++b ())
JES
jES
hoo.)))eorr.b00((1))))+)) o(w, + ())dd JES.
S Callilookoo&(---3 ))
for all |||||| S rz, where we used the fact that (2) is O( Lipschitz when restricted to |a| S rb. This concludes the proof.
Next, we provide discretization guarantee when using (a,3)-SFL oracles.
Proposition 34. Consider the same setting Proposition 33. except the first-layer weights (Wi...., WN) are obtained from the (a,3))SFL oracle (Definition 3). Then, there exists a" such that af 0 for i & S and la? S |/ (BBaN) (BoN)foor f E S, and
< CoLollills 9+1
log(aN/5)
Σ aajo((U3. 2) + b,) Th(z) jES
B
aN
for all Rh with ||||| S rb. with probability at least 8 over the randomness of (v,, be)iE(M)- dp dr, S B', we have
Moreover, suppose Bon. [hee,b b)2| s Mz(h)² Then, assuming
M22)2
| log(1/6) N
||a"|
provided that,
N,32N
a M2
which also holds with probability at least 1 - 8.
Proof. By definition,
[ Jsn-ix(-moral
Th(x)=
h(v b)o((() + ())r
- ix(-ruiro
M(, b) drk (v)o(((, + b)du(o)db du
2rgh(v, b) drk du
Ey br m.))))
(v) 2) + b)
Consider (v, b.)ss i.d.d.
Unif( rb) from Definition 3. Let
2r h(r b.)
dra (vx)
if E if o S.
a;
[s]
du
0
Consequently
2r 2rlli
|at1 <
for all E S. Given z, define the random variable
Th(z) Σ
at (1,, 2) + b,).
iES
34

---

Published as a conference paper at ICLR 2025
Our next step is to bound the difference between Th(x) and Th(2) uniformly over all || s mb. Let (z be a A-covering of {=1 : || S ro). therefore M I (3r//A). Note that for any fixed x with || S rb. we have Th(z) E Lo 1/. Thus, by Hoeffding's 's lemma,
(h(=)) - -Th(z) x VVllloloor 91
log(1/6)
B
|S|
with probability at least - 8 for fixed z. By union bound.
|1Ail1
Lorg
4+1
log (M/6)
Th(2,) - Th((2)) x
max JEIM
3
[S]
with probability at least 1 - 6. For any = with ||=|| S nb. let & denote the projectionn of z onto the covering (#,) Then,
Th(z) - Th(=))
Th(2,) - -hh(;)) + -h(()) + Th(z) Th((z)
sup 1/
max JE(M]]
| lo Lorg 9+1
log(M/0)
B
S
| lo Iorg 9+1
[lg((r/((D)) |S]
11 T
+
B
8
Choosing A Tb/ (IS implies
Th(x)- Th(x)) i |1 collooLoo borg q+1
log((S)/5)
sup Sr/
B
S
with probability at least
8 over the randomness of (vi, Ds)SENN
The last step is to bound ||a" 1| Note that,
h(v, 6,)2
|| 122
IES [S]
Further, by the Hoeffding inequality,
bi)²
log(1/6)
E. Pe.b-suUnil(-na.r) [((,,,)3 b)3] x |1
IS
|S|
YSS
with probability at least 8. Moreover,
Bepbbrvootnii( [(v,b)a] = Ev Bor h(w b)2
du
(w)
E
drk
Thus, when |S| >
| log(1/8) (b))+
have ||a*|2
(B?)()) with probability least
1 - 8, which completes the proof.
D.5 COMBINING ALL STEPS
We can finally bound our original objective of this section, i.e. AR(a", W,b) - ARR., Let us begin with the case where F is the class of polynomials of degree p.
Proposition 35. Suppose IF and satisfy Assumption 4 and (b:)se(N) Lid. Unif(-ro, rb). Recall
that E = 1 V and e ART for any E (0, 1). Using the simplification k, q, p, La S 1 and recalling ε = V e, there exists choice of rs ((1)) such that:
35

---

Published as a conference paper at ICLR 2025
W (wi..... ................................ wN) are given by the a-DFL oracle, there exists a' such that |a 1 ((e1/(aN)) for all E [N], and AR(a", W, AR" as
E 9+1
2 0
and
N N Ω
2(g+1)
11-1///
Wf (wi..... ........ wN) T are given by the (a,,3)-FF) oracle, there exists a" such that |a s (BEL/(BaN)) for all E [N. and AR(a", W. b) AR" as soon as
B22
N á 2(++1)
and
2[g+1)
a.32
Both cases above hold with probability at least 1 - n for some absolute constant > 0 over the choice of random biases (bi) (and random weights (w,) in the case of SFL).
Proof. Recall from (D.1) that
AR(a', W. -R' & + +
By definition, & S AR". By Lemma 24, we have
&
+ rb
210-1) + 200 E²)ç.
Further, thanks to Lemma 31 we have 111 [()) S + | Therefore, by Lemma 25 with r2 = Tb. have
& 2 approx + (L) : )))) + + -663)
Let us now consider the case of aDF.. Define p p= (p - 2) v if the ReLU activation is used and
p 0 if the polynomial activation is used. Notice that by the definition in Assumption
have
7 p - q. By Proposition 33, know there exists a" with |a S Ö(r)
((rr+//((N) /(aN)) (we used the
fact that max.5p| ((+)) $
from Lemma 31) such that
fapprox I x0(1gt
(VS+ (1( (kk-1)/2aN))
(
provided that rb Z S1 where we recall é = e, and the above statement holds with probability least for any polynomially decaying 6, e.g. d n for absolute constant Therefore, we have Fa ((r)+P). Further, it suffices to choose T large enough such that rs 2
e V 1og 24 + + é(s1) to have
& ô (20()2(000 (( x2
Plugging in the values of Fa and rb, we obtain,
E 2
& < Ō 2(a
and
ε, s 2(((tt))
ck-la2N2
Hence, choosing
< 2(
and
N á
1)/2
which concludes the proof the a -DLL case.
In the case (F,, we instead invoke Proposition 34, thus obtain |ay) S 1+D /(BaN) / and
Eapprox I 5 Lo g
3vaN
36

---

Published as a conference paper at ICLR 2025
which holds with probability least - á for any polynomially decaying 8 such as å = n- for some absolute constant > O. Consequently, with the same choice of = (E1) before, we have
& S 1506-0 (
E S Ō 2(++1)
and
B2
PaN
which completes the proof.
We can also combine approximation bounds for the more general class of pseudo-L.psshitz. F. Proposition 36. Suppose and satisfy Assumption 3 and (bs))e(N) Li.d. Unif(-rb. rt). Recall that ε V e, and ē == A AR for any (0,1). Using the simplification k, L 2 there exists a choice of rb = ((<L(E//V2)+11)) +1/)))/)) such that:
è²
1/W (u),..... UNN given by the a-DFL oracle, there exists a such that |a S (((e//Va)tt++1/
for all E [N], and AR(a" aa W. b) AR S c as soon as
s 550(1 ?
*+2+1/k)
and N
2 C(k- 11/211/2
k+3+2/k)
If W (w),... ......., wN) is given by the (a.3)-SFL oracle, there exists a' such that afl s ((ei/VVa) (((+tt1///((aN)))) and AR(a", W.b) -AR" S as soon as
< (2 { (4221/1). (
and
N Ω
k+342/K)
0
Both cases above hold with probability at least 1 1 some absolute constant e 0 over the choice of random biases (b:)e([) (and random weights (ws) in the case of SFL).
Proof. Our starting point
again the decomposition
AR(a", W. b) AR" S E2 + V + E2)
Given Assumption 3, it is straightforward to verify that (M(z)) M(22)1 S (1 1|--1 cf 11=1|1211011 + 1)]lz1 - zell for 21, Z2 E R As consequence, we have [h(2)] 2 1 1 |
for all # R Therefore, by Lemma 25 with choice of r. @(E1), we have &, S 2 rest of the proof we will fix T= ((1)).
approx In the
We begin by considering the case of a DFL.. Unlike the proof of Proposition 35 where Th h, in this case we have an additional error due to Th only approximating h. From Lemma 32, we have
|1 s 0 61
2l/(k+1)).
Thus,
Capprox < sup
2) bb)) -)) Th((z)
+ Th(z) - h(2))
|
JES
(e1 (eivC
rb
θ 0 ( की
((k-1)/2aN))
61
61
where we bounded the first term via Proposition 33 with 1, and the second term via Lemma 32. Additionally, we have
|aj < E
Tb
for all j E [N]. To obtain AR(a", W, b) S AR" + we must choose A = =((e((es/v7))utt)) (411/2)
Next, we choose rs
(x+//)))) This combination ensures Th(2) l1(z) is Ve. To
make sure (Th(2) Th(z) S Vł, we should let
<< 5 m )
N Ō ((k-1)/2 ( Ve
R+3+2/K)
and
37

---

Published as a conference paper at ICLR 2025
The above guarantees that Capprox X Ve and consequently & VEgg S Note that the above choices imply |a3] $ Fa/|S] for all E S with Fa ((s//VV))tttt)/) From Lemma 24 with q = we have &2 I fasf.. Therefore, if we let
Ç 6=6(( h
we have Eg S é and consequently &g VEze1 S This concludes the proof of the a-DFL case. Next, consider the case of (a,L)-SFL. Note that the error Th(z) h(z) However, this time invoke Proposition 34 for controlling Th(=) Th(=))
remains unchanged.
Therefore,
Σ sup Σ
(())(((;)))+)) +h(z) -b) Th(z)) n(2)
Capprox
||
jES
A + + +
क + E1
E1 ( ro
B21 E1
E1
Since the second term is unchanged, we have the same choices of A ((e/(Es/V2)(tt)))) and тb (=) +1/k) as in the a-DFL case. However, for the finite-width discretization, we should choose
1 N 6 a2 k+3+//k)
(D.7)
|aī] S i Fa/ with Pa (((e//V2))t1111)) /3). As a result, to get & S Pa-iC < from Lemma 24 with 1, we let
Moreover, Proposition 34 implies
Ç = θ 8 kt2+1/1)).
completing the proof.
E ADDITIONAL EXPERIMENTS
In this section, we perform simple experiment on the MNIST dataset (LeCCun et al., 1998) to demonstrate that the intuitions from our theoretical results go beyond the setting of multi-index models, squared loss, and & norm attacks. We choose a convolutational neural network our predictor, given by two convolution layers with 32 and 64 channels respectively, max pooling layer, and two fully connected layers. For each choice of we do two experiments:
1. ADV training: We train the network adversarially for 1 epochs, where we initialize the model using the default Py Torch initialization, then use the adversarial training algorithm of Madry et al. (2018).
2. STD + ADV training: We first train the model with standard SGD for 10 epochs from PyTorch initialization, then perform 10 epochs of adversarial training on top of this standard presentation.
For both approaches, we use the corss entropy loss, batch size of 64, learning rate of 0.01 for both PGD and SGD updates, and use loo norm to constrain perturbations, where pixels normalized between 0 and 1. Note that due to the additional cost of generating adversarial samples, taking one additional epoch ADV ensures fare comparison where ADV and STD ADV will roughly have the same computational complexity.
As be seen from Table 1, the STD ADV training approach achieves higher test accuracy compared to the ADV approach across all model architectures considered here. This is consistent with our intuition from the guarantees of Algorithm 1. The standard training phase recover
38

---

Published as a conference paper at ICLR 2025
Algorithm
0.2
0.3
0.4
ADV
91 +
82 +
69 + 2
STD + ADV
93.80
0.09
87.2 + 0.4
75.3
0.6
Table Adversarial test accuracy comparison between ADV and STD + ADV across three different values for on the MNIST dataset. The error shown for each accuracy is the standard deviation over three runs.
optimal low-dimensional representation top of which adversarial training becomes easier. Note that since we work with multi-layer convolutional neural networks, there is no longer a single layer that captures the entirety of the low-dimensional projection, which is why we choose retrain all parameters of the network adversarially. One interesting direction for future research is to understand in settings beyond two layers, which parameters need be adversarially trained after the standard training phase. The code to reproduce the results of Figure and Table is provided at: thub. www.:uu..coo/oooo....c...a. wobustt feature sereeellarninng
https: / http///ttttttttttt
39


Reviews:
[{"review_id": "PRvlujtdRB", "reviewer": "Reviewer_weXL", "summary": "The paper analyzses adversarially robust learning for two-layer neural networks trained on multi-index targets dependent on a low-dimensional subspace. The paper begins by proving that an optimal adversarially-robust predictor depends solely on the target subspace. Subsequently, the paper shows that existing guarantees for recovering the target subspace through SGD can be combined with covering number, rademacher complexity-based generalization and approximation arguments to obtain bounds on the adversarial risk. Crucially, the paper shows that the additional sample-complexity requirements for vanishing adversarial risk are independent of the ambient dimension.", "questions": "- What are the assumptions on the class $\\mathcal{F}$ in Theorem 1?\n- In the proof of Theorem 1, is $h(x)$ assumed to be in $\\mathcal{F}$?\n- Does definition 2 (DFL) implicitly assume constraints on the magnitudes of $w_i$? Otherwise, DFL is trivially satisfied by scaling up the weights. How is this scaling issue consistent with Theorem 4?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem of obtaining theoretical guarantees on adversarial risk and the dependence on the structure of the data is well-motivated of interest to the community.\n- The direction of studying the interplay between feature learning and adversarial robustness appears to be novel.\n- The paper is well written and describes the setup, literature well.\n- The proofs are easy to follow.", "weaknesses": "My main criticism of the paper is that the extent of the theoretical novelty in the analysis and unexpectedness of the result is unclear in the paper's presentation. The paper would improve by emphasizing the central difficulties in combining the finite-dimensional guarantees with feature learning results. Specifically, it would be useful to know what parts of the analysis would change if one were directly studying the setup of random features on finite dimensions. It appears that the role of the low-dimensionality of the target subspace and feature learning is primarily in obtaining a second layer with dimension-independent norm ($r_a$ in Algorithm 1) that fits the target allowing for the Rademacher-complexity argument to go through. This aspect of the analysis is however, the same as generalization bounds in existing works on feature learning such as Damian et al. 2022.", "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": "zu9dtxV9m3", "reviewer": "Reviewer_1ZeF", "summary": "This paper addresses the challenge of learning adversarially robust models using 2-layer neural networks where the targets are a function of a latent low-dimensional projection of the inputs. Throughout the paper, they consider L2-constrained adversarial perturbations. They first prove that the optimal features in terms of adversarial robustness are in fact the true latent features of the data model. Then building off of prior work in 2-layer network feature learning, they assume access to a feature-learning oracle and analyze algorithms which query the oracle for the first-layer features and then adversarially train over the second layer’s parameters. For different function classes, they provide bounds for the number of required samples for adversarial training, neurons, and feature learning oracle error in order to achieve optimal robustness up to epsilon error with high probability. Through brief experiments, they validate their theoretical results and show that standard training followed by training just the last layer using the adversarial objective performs much better than full adversarial training.", "questions": "Can you discuss how this result and analysis techniques connect to other recent results in how shallow networks learn representations, such as (Collins et. al., Provable Multi-task Representation Learning by 2-Layer ReLU Neural Networks, ICML 2024) and (Wang et al, Learning Hierarchical Polynomials with Three-Layer Neural Networks, ICLR 2024) ? \n\nIs the first inequality of the proof of Theorem 1 actually an equality?\n\nIn the experiment section, how did you select the “direction u”? Is this a vector in R^d? I do not see where u is defined.\n\nCan you discuss what the actionable insights are in terms of training robust models? It would be interesting to see more practical experiments of how standard feature learning followed by running adversarial training on the last layers of a network can result in a full model which is adversarially robust.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper’s first contribution is that the ground truth latent features are in fact the optimal features in terms of adversarial robustness. This is not largely surprising, but it plays an important role in informing the kinds of algorithms one should use to achieve an adversarially robust model. They then show how to characterize the sample complexity and model complexity for just adversarially training the second layer for a broad class of target functions when given access to a feature learning oracle. Importantly, these bounds are independent of the ambient dimension as the hardness only depends on the effective dimension of the problem. These results are an interesting extension of previous works in feature learning for 2-layer neural networks which do not consider adversarial robustness.", "weaknesses": "The major weakness is the lack of extensive experiments. I think that the first row of figure one is not a fair comparison as the paper graphs a model trained completely using adversarial training techniques to just training the last layer using adversarial techniques while initializing the first layer to the ground truth representation. In the second row,  all models are forced to learn representations which is a more realistic scenario. However, based on the high-level result that adversarial training is no harder than standard training and that the optimal adversarially robust features are just the ground truth features, it would be interesting to see how different combinations of standard vs. adversarial training on different network layers for real-world data affects performance.  \n\nAlso, the statement of Theorem 1 is confusing. First of all it is not clear that the multi-index data model is still assumed. Also, the phrase “with equality when f^* in F” is confusing since f* in F doesn’t imply equality and that part is obvious, so stating it in this way is confusing for no reason. Also, the statement “Further, h is represented as E[]” is also confusing because it seems to mean that any h that satisfies the theorem must be represented in that way which is not true. The theorem should be restated as something like “Assume multi-index data model as explained above, Assumption 1, and (2.2) admits a min. Then there exists a function f* = h(Ux) where h = E[f] for some f in F where AR(f*) leq AR*.”", "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": "APBuWJ0rlX", "reviewer": "Reviewer_vLR3", "summary": "This paper demonstrates that it is possible to robustly train a multi-index model without impacting the statistical performance of the algorithm compared to standard learning methods. Assuming a low-dimensional structure in the probability distribution of the observations, the authors propose a procedure for robust learning in two-layer neural networks. This approach involves a standard learning step for the input layer aiming at recovering the relevant low-dimensional subspace, followed by a robust learning phase for the output layer in lower dimension.\n\nIn Theorems 4, 5, 6, and 7, the authors provide rigorous asymptotic bounds on the sample complexity (in the second step), number of neurons, and achievable error for the feature learning algorithm, ensuring the algorithm attains a small adversarial generalization error. These bounds are shown to be independent of the ambient dimension. Finally, using recent results on the number of observations necessary to learn the relevant low-dimensional subspace, the authors establish a scaling for the sample complexity for both single and multi-index models in Corollaries 9 and 11.", "questions": "-\tThe proposed method only involves an adversarial training of the second layer. This has the benefit of achieving a robust training without requiring more samples or neurons than standard training. In addition, Figure 1 suggests that the convergence is faster compared to a full adversarial training. However, in the ReLu case (left panel), it seems that the robust test risk in the case of the full AD training is still decreasing. Given this observation and in the same setting as the paper, is there a possibility that robustly training both layers could achieve the same guarantees (sample complexity, number of neurons required) as standard training? \n\n-\tIn line 187, it is mentioned that the parameters $a, b$ are used to approximate $h$. This corresponds to roughly $N$ parameters to approximate the $k$-dimensional function $h$: can we think of the power $k$ in the bound for $N$ in Theorems 4 and 5 as a form of curse of dimensionality associated with this approximation task? If so, is there an intuitive reason why those exponents do not appear in Theorems 6 and 7 ? \n\n-\tI am not sure about the assumptions of Proposition 8: In [Lee et. Al., 2024] Theorem 2 requires that the activation is also a polynomial with a higher degree than the link function. Is this additional assumption required for Proposition 8? Likewise, in Corollaries 9 and 11, the mention “as in Assumption 4” may be confusing: do you also make the assumptions regarding the activation? \n\nMore precise questions and remarks:\n\n-\tIn line 215, it is not clear why the expectation of the norm of $x$ is proportional to the square root of the dimension. \n-\tLikewise, in line 218, the argument for choosing $\\varepsilon$ of order 1 can be made clearer (is this because the scalar dot between $u$ and $x$ is also of order 1?).\n-\tIn and before Definition 2, it is not mentioned that the columns of the weight matrix should be of unit norm. Without this detail, it seems that the definition of the DFL do not make much sense. \n-\tIn Definition 3, the last sentence is not very understandable and is not explained afterwards.\n-\tIn Theorems 4 and 6, it would be nice for completeness to precise that the absolute constant c is positive. \n-\tSince the results of section 4.2 do not cover the case of general activation, the discussion spanning from line 405 to 423 could be shortened. \n\n\n[Lee et. Al., 2024]: Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoretic limit.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-\tAlthough I do not have any particular knowledge in robust training, I found the paper accessible and well explained. The stakes are clearly presented in the introduction. Additionally, the setup and the assumptions are generally well motivated. \n-\tDespite not having read the proofs in detail, I appreciated that the proof techniques originate from diverse ideas and prior works including statistical learning, with generalization bounds provided in Appendix C1, geometrical analysis from the low-dimensional subspace assumption, approximation of functions using infinite-width neural networks, etc.\n-\tFrom the references provided in the introduction, it is clear that this paper offers new results and insights, distinguishing this work from previous research in the field.\n-\tThe numerical experiment displayed in the main text gives a nice support to the theoretical claims of the paper.", "weaknesses": "-\tTheorems 4, 5, 6 and 7 rely on specific activations (ReLU / polynomial). What is the technical challenge arising when considering more general activations? \n-\tAlthough it is interesting to include some proofs in the main text, the proof of Theorem 1 could be easily replaced with an outline of the proof ideas. \n-\tOn the other side, it would be helpful for the reader to have a paragraph or two summarizing the main steps or ideas of the proofs of the main theorems. \n-\tOne of the main claim of this work is that the bounds derived do not depend on the dimension. These bounds include the optimal adversarial risk (defined in equation 2.2), therefore, this quantity should not depend on the dimension. Whether this is a direct consequence of the definition or of Theorem 1, I don’t see it mentioned in the main text. \n-\tSeveral concepts and notations are not clearly defined, and it sometimes affect the understanding of the results. My main concern is the lack of a definition for the Landau notations in the main results. As a consequence, it is not clear whether or not the obtained bounds are non-asymptotic or asymptotic and if they are, with respect to which variable.", "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": "MD3oJkDQqr", "reviewer": "Reviewer_UkaQ", "summary": "This work considers robust learning in the multi-index model for two-layer ReLU networks. In particular, authors study the setting of worst-case $\\ell_2$-perturbation of the input and the problem of regression with squared loss. \n\nThey propose the following two-stage method to find a set of parameters (weight of the first layer and readout vector):\nFirst, algorithm finds a low-dimensional subspace, on which the target function depends. This is done using non-adversarial training, using prior feature learner methods. \nUsing this subspace, adversarial training is done on the readout vector.", "questions": "Second row of Figure 1: why having $W$ fixed at SD training plateaus, while continuing to train $W$ improves the performance? Doesn't it contradict the theoretical results, since according to Algorithm 1, one should not update $W$ after the first stage? \n\nWhy for the numerical experiments Algorithms 2,3 are not used?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The question of robust learning in the multi-index model is well-motivated by the authors. Authors clearly state the assumptions and limitations of their results.\n\nAs the main result, authors prove that in the multi-index model, sample complexity of adversarial training is the same as in standard training. They combine prior work on feature learning with a second stage adversarial training which requires careful technical analysis. Authors also provide experiments to support their claim.", "weaknesses": "Authors highlight themselves that the result is limited to $\\ell_2$ norm of perturbations. Some discussion on the extension to $\\ell_\\infty$ norm is lacking, i.e., why current proof techniques cannot tackle this setting.\n\nFurthermore Theorem 1, which is central for the results of the paper, relies on the squared loss assumption. For classification tasks, which is a standard setting for adversarial robustness, logistic loss would be a more natural choice.  \n\nOrganization of the main text is hard to follow, theorems for SFL and DFL are more or less copy each other, but stated in different places. Since only one proof (out of four Theorems 4-7) is given in the Appendix, possibly a cleaner way to present the results is to combine Theorems 4&5 together, and show Theorems 6&7 as a corollary under additional assumption (with a separate proof).\n\nProof of Lemma 16 in the Appendix is hard to follow because of typos (line 894, should be $A_i (a_i - a'_i)$), and also $\\tilde k_i$ is defined with $k_i$ and $k'_i$ in line 899, but both $k_i, k'_i$ are not defined at this point. I think what authors wanted to say is that $k_i, k'_i$ is defined with $\\tilde k_i$, which is obtained through Mauery's lemma. It is confusing as it is written now.\n\nNumerical experiments can also be improved, see Questions.", "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": "MA9Gk6D9u8", "reviewer": "Reviewer_E3tK", "summary": "The authors study the problem of $\\ell_2$-constrained adversarial regression under a multi-index model setting. Specifically, in this setting, the conditional form of $y\\vert x = g(\\{\\langle u_i, x\\rangle\\}_{i \\in [k]})$ holds, where $u_i$ are the weights (or features) to be learned, $g:\\mathbb R^k \\to \\mathbb R$ is the link function and $x,y$ are the covariates and target, respectively. \n\nTheir main result shows that under the assumption that the model relevant co ordinates $Ux$ (for $U \\in \\mathbb R^{k \\times d}$) and model irrelevant co ordinates $\\tilde Ux$ ( $\\tilde U \\in \\mathbb R^{(k-d)\\times d}$) are statistically independent, there is a learning method that attains a Bayes optimal adversarial risk. \n\nThis is achieved through access of oracles which can be realized through gradient based methods, with additional assumptions on the data distribution and for a certain class of prediction functions, such as variants of Lipschitz and polynomials. \n\nUnder the assumption that the predictor is a \"pseudo\"-Lipschitz, the sample complexity is not dependent on the data dimension, though it has an exponential dependence on $k$; for polynomial predictors, their sample complexity results do not depend on k. \n\n\nThey do not analyze the computational complexity of this problem, and provide empirical evidence by comparing training approaches including initializations based on a known $u$ and fully adversarial training.", "questions": "Asked above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The contribution is good because it is a new theoretical study for adversarial feature learning for learning multi index models, which has been gaining traction recently. Their main results highlight that robust learning is achievable without a significant increase (at least for some functions) in sample complexity. Moreover, the oracles assumed for learning the target subspace can be realized by training the first layer in a 2 layered neural network. Overall, this is a well structured, theoretically sound work with a good upshot and I recommend an accept.", "weaknesses": "I would like to reiterate a few points they already raised. Firstly, for the case of pseudo-Lipschitz functions, while the sample complexity is independent of $d$, there is an exponential dependence on $k$, which seems quite high. \n\nI also think that many mathematical terms could benefit with more intuitions for the first time reader. For example, a minor point of clarification would be in Theorem 1, where it would be helpful to specify the random variables with respect to which the expectation is taken in the towering property, $\\mathbb E_X[\\cdot] = \\mathbb E_Y[\\mathbb E_X[\\cdot \\vert x,y]$. \n\nIt is also somewhat unclear to me what $\\beta$ represents in the SFL oracle and also if a lower bound on $d\\mu/d \\tau_k$ connects to the Assumption 1. Moreover, I also think that Assumption 1 does not hold for natural scenarios where $x$ has a low rank structure where $\\tilde Ux$ may not vary independently from $Ux$.", "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": "9GilLfVK3j", "reviewer": "Reviewer_1AAV", "summary": "This paper explores the robust feature learning properties of two-layer neural networks for learning multi-index models in high-dimensional settings, under $\\ell_2$-bounded perturbations. In particular, the authors show that Bayes-optimal adversarial risk can be achieved by neural networks once the input data are projected onto the low-dimensional manifold defined by the target multi-index model (first-layer feature learning) and the second layer is trained by minimizing the empirical adversarial risk. They demonstrate that the resulting sample complexity is independent of the ambient dimension. Leveraging prior results, the paper provides guarantees for robustly learning multi-index models with gradient-based algorithms and includes numerical simulations to support its findings.", "questions": "1. Can you provide insights or intuitions about the validity of your results beyond Assumption 1? Have you numerically tested your algorithmic framework in settings that do not satisfy this assumption?\n\n2. In lines 384-388, you claim that your results in this section can be readily extended to the setting of anisotropic Gaussian $x$, which seems to contrast with Assumption 1. Is it still possible to leverage Theorems 4-7 in this case?\n\n\n3. I believe there is a typo in equation (4.1), where $y$ should instead be $y^{(i)}$.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper\n* presents results with clarity and precision, offering complete proofs, insightful comments, and remarks that make the content accessible and well-structured\n* leverages previous results and introduces interesting, non-trivial findings in a rigorous framework (e.g. the possibility of performing standard feature learning to obtain optimal adversarial risk)\n* validates its theoretical results with numerical simulations\n* acknowledges its limitations and open problems.", "weaknesses": "I did not find any major weaknesses in this work. A minor issue would be the reliance of the results on the assumption of independence between the coordinates $x_\\parallel$ and $x_\\perp$ (Assumption 1), see also my question 1.", "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": "7RROgoqHmP", "reviewer": "Reviewer_6kxP", "summary": "This paper studies the adversarial training of two-layer neural networks on learning multi-index target function (a function of a low-dimensional projection of the input) under squared loss. Under some assumptions, the authors first show that there exists an optimal low-dimensional projection (of inputs) for robustness against $l_2$ adversarial perturbations (Theorem 1). They then provide conditions for the first layer weights to satisfy such that the optimal low-dimensional projection matrix and the first layer matrix are partially aligned (Definition 2-3). Since there exist practical feature learning algorithms (examples given in Section 4.2) satisfying these conditions (definitions), they focus on the case where only the second layer is trained with adversarial loss while the first layer is pre-trained with a standard (not adversarial) feature learning algorithm. This simplification allows them to study the adversarial training of the second layer mainly. In this setting, the paper shows that adversarial training does not significantly affect the sample complexity of learning the target function. Specifically, the number of additional samples required for adversarial training (in comparison to standard training) is shown to be independent of the input dimension (Theorem 4-7). Thus, the authors conclude that adversarial learning in this setting is as easy as standard feature learning.", "questions": "1.  **The assumption on the data distribution:** \n     + a) The authors mention that inputs can be Gaussian vectors with isotropic covariance, but what are other examples satisfying this assumption? For example, can we have anisotropic covariance (e.g., spiked covariance)? What if the covariance matrix and the projection matrix $U$ are aligned?\n     + b) The authors mention (in Line 388) that \"Our results readily extend to these settings as well,\" referring to non-isotropic Gaussian inputs and non-Gaussian symmetric input distributions. How so?\n\n2. **Numerical results:** \n    + a) What is the step size for the gradient descent step on the perturbed batch for Figure 1?\n    + b) What do shaded error bands denote in Figure 1? If variance, how much variance is illustrated?\n    + c) Why do we see oscillating robust test risk for the two cases other than the full adversarial (full AD) training, while the robust test risk for the full adversarial training seems to be smooth in Figure 1? Furhermore, the robust test risk for full AD seems to converge fast while those for the other two cases do not seem converge even at the end of the iterations. Why so? \n    + d) The authors provide theoretical results on polynomial predictors in Section 4.1. A natural question is, \"How does polynomial activation compare to ReLU activation in numerical results?\". \n    + e) Prior work (Hassani & Javanmard, 2024) showed that adversary budget $\\epsilon$ can significantly impact the robust test risk in their setting. How does adversary budget $\\epsilon$ affect the robust test risk here? Similarly, what is the impact of dimensions and the number of samples?\n\n3. The authors mention $\\zeta^{(k-1)/2}$ as \"the normalizing factor\" in Line 239. But, $\\zeta$ on its own deserves an explanation since it appears in most of the theorems. What does $\\zeta$ intuitively correspond to?\n\n4. **Minor issues:**\n    + Some non-standard notations are used without definition. For example, soft O notation $\\tilde{\\mathcal{O}}$ and its counterparts for $\\Omega$ and $\\Theta$ should be mentioned in the \"Notation\" paragraph starting in Line 122.\n    + There seems like a notational abuse when using the aforementioned asymptotic notations. There is asymptotic notation inside another asymptotic notation (e.g., in Theorem 4): $N \\geq \\tilde{\\Omega}(\\cdot^{\\mathcal{O}(k)})$. Indeed, the authors assume $k = \\mathcal{O}(1)$ in Line 154, which means $k < C$ for some $C > 0$. So, these equations can be simplified, and the notational abuse can be fixed at the expense of introducing one additional constant. Furthermore, note that the asymptotic notations are conventionally used with equality sign $(=)$ instead of $\\leq$ or $\\geq$.\n    + In Line 135, the authors define the family of prediction functions to be \"non-parametric\". However, they then use two-layer neural networks as a prediction function in Line 140 or d-variate polynomials as the family of prediction functions in Section 4.1, which are \"parametric\" families of functions. There seems to be a contradiction here. Or am I missing something?\n    + Ranges of $\\alpha$ and $\\beta$ are missing from Definitions 2-3.\n    + Complexity results provided in Theorem 4-7 reduce the readability of main text since they include many parameters related to the setting. I think that the author can provide more insightful presentation of these results showing the essence.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- **Interesting problem:** Understanding the effects of adversarial training of neural networks on the sample complexity may be significant for the deep learning literature. Furthermore, this paper combines two lines of work: 1) feature learning in two-layer neural networks and 2) adversarially robust learning in linear and random feature models.\n\n- **Theoretical study under reasonable assumptions:** The authors study the problem under reasonable assumptions (in line with the prior work about such as Hassani & Javanmard, 2024).\n\n- **Provides a valuable insight into the problem:** The paper shows that the adversarial training does not significantly (in the order of dimensions) affect the sample complexity in comparison to standard training.\n\n- **Good enough presentation:** The paper answers most of the questions a reader may have, which is the main criterion when assessing the quality of the presentation. However, I think it is not an easy-to-follow paper. Instead, some parts of the paper require back-and-forth reading, which may be acceptable for such a theoretical paper.", "weaknesses": "- **No characterization of robust test risk:** Prior works such as (Hassani & Javanmard, 2024) provide an asymptotic characterization of robust test risk for random feature models, which allows them to show how factors such as dimensions or perturbation budget $\\epsilon$ impact the robust test risk. However, this work just focuses on the sample complexities and ignores the characterization of robust test risk. \n\n- **Significant gap to practice:** Although the assumptions are reasonable from a theoretical perspective, they differ significantly from practical adversarially robust learning.\n   + **Data assumption:** Input is assumed to be subGaussian with constant subGaussian norm in Assumption 2. Furthermore, Assumption 1 imposes an implicit assumption on inputs and labels (target function), which needs further discussion (please see Question 1 about this below).\n   + **The first layer is also trained with adversarial loss in practice:** While it seems like the adversarial training of the first layer hurts the robust test risk in Figure 1, this requires explanation (mentioned as future work in the conclusion). Thus, adversarial training of the first layer  is missing in this work.\n   + **Limitation of $l_2$ perturbations:** The paper only studies adversarial training with $l_2$-bounded perturbations. Although $l_2$ perturbations for the adversarial training may be a good starting point considering the prior work on the theory of robust learning (Hassani & Javanmard, 2024), $l_{\\infty}$ is primarily used in empirical studies as mentioned by the authors.\n   + **Assumption on perturbation budget $\\epsilon$:** The paper assumes constant adversary budget $\\epsilon = \\mathcal{O}(1)$ and justifies it by saying even this much perturbation can significantly perturb the output of the prediction function. However, as the authors noted, $\\mathbb{E}[||\\mathbf{x}||] = \\mathcal{O}(\\sqrt{d})$ and thus perturbation budget of $\\epsilon = \\mathcal{O}(\\sqrt{d})$ is used in practice, which causes another discrepancy between the assumptions and practice.\n   +  **Only polynomial and ReLU activation functions are considered.**\n\n- **Limited numerical results and limited discussion of those results:**  \n   + Most of Section 5 (Numerical Experiments) describes the experimental setting, while only one or two sentences discuss the results. There should be a more insightful discussion of the numerical results and their relation to the theoretical results.\n   + There are some missing details about Figure 1 (please see Question 2.a-b below).\n   + Effects of some parameters of the setting should be illustrated numerically or at least discussed: 1) activation function (polynomial vs ReLU), 2) different values for adversary budget $\\epsilon$, 3) dimensions and number of samples.", "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)