A Derandomization Framework for Structure Discovery: Applications in Neural Networks and Beyond
Authors: Nikos Tsikouras
Yorgos Pantis
Ioannis Mitliagkas
Christos Tzamos
Venue: ICLR 2026 Poster (2026)
Keywords: Neural Networks
Optimization
Structure Discovery
Compressibility
Derandomization
Multiple Index Model
Johnson Lindenstrauss
MAXCUT
Primary area: learning theory
TL;DR: We extend theoretical insights into Neural Networks, proving a key derandomization lemma that explains structure discovery and applies to other problems such as MAXCUT approximation and Johnson-Lindenstrauss embeddings.

Abstract: Understanding the dynamics of feature learning in neural networks (NNs) remains a significant challenge.
  The work of (Mousavi-Hosseini et al., 2023) analyzes a multiple index teacher-student setting and shows that a two-layer student attains a low-rank structure in its first-layer weights when trained with stochastic gradient descent (SGD) and a strong regularizer.
  This structural property is known to reduce sample complexity of generalization.
  Indeed, in a second step, the same authors establish algorithm-specific learning guarantees under additional assumptions.
  In this paper, we focus exclusively on the structure discovery aspect and study it under weaker assumptions, more specifically: we allow (a) NNs of arbitrary size and depth, (b) with all parameters trainable, (c) under any smooth loss function, (d) tiny regularization, and (e) trained by any method that attains a second-order stationary point (SOSP), e.g. perturbed gradient descent (PGD). At the core of our approach is a key $\textit{derandomization}$ lemma, which states that optimizing the function $E_{x} \left[g_{\theta}(Wx + b)\right]$ converges to a point where $W = 0$, under mild conditions. The fundamental nature of this lemma directly explains structure discovery and has immediate applications in other domains including an end-to-end approximation for MAXCUT, and computing Johnson-Lindenstrauss embeddings.

Published as a conference paper at ICLR 2026
A DERANDOMIZATION FRAMEWORK FOR STRUCTURE DISCOVERY: APPLICATIONS IN NEURAL NETWORKS AND BEYOND
Nikos Tsikouras Yorgos Pantis Ioannis Mitliagkas2 Christos Tzamos 1.2
National and Kapodistrian University of Athens, Greece
2 Archimedes, Athena Research Center, Greece
3 Mila & Université de Montréal, Canada
ABSTRACT
Understanding the dynamics of feature learning in neural networks (NNs) remains a significant challenge. The work of (Mousavi-Hossetle et al. 2023) analyzes a multiple index teacher-student setting and shows that two-layer student attains a low-rank structure in its first-layer weights when trained with stochastic gradient descent (SGD) and strong regularizer. This structural property is known to reduce sample complexity of generalization. Indeed, in a second step, the same authors establish algorithm-specifi learning guarantees under additional assumptions. In this paper, we focus exclusively on the structure discovery aspect and study it specifically:
allow (a) NNs of arbitrary size and depth, (b) with all parameters trainable, (c) under any smooth loss function, (d) tiny regularization, and (e) trained by any method that attains second-order stationary point (SOSP), e.g. perturbed gradient descent (PGD). At the core of our approach is key derandomization lemma, which states that optimizing the func- tion Eg [ge ( [W œ + b)] converges to point where W = O, under mild conditions. The fundamental nature of this lemma directly explains structure discovery and has immediate applications in other domains including an end-to-end approximation for MAXCUT, and computing Johnson-Lindenn Lindenstrauss embeddings.
under weaker assumptions,
INTRODUCTION
Neural networks (NNs) have become successful tools across different domains, demonstrating exceptional performance complex tasks, such as image recognition, natural language processing, or speech synthesis (LeCun et al. 20155 Goodfellow (016). This broad applicability is primarily due to their ability learn and generalize from large datasets, enabling them to identify difficult patterns and relationships that are difficult to capture with traditional techniques. Theoretical work in this area focuses on various aspects, including the structure of optimization landscapes, and generalization behavior, aiming to answer fundamental questions about why these models work as well as they do and how they can be made more efficient and trustworthy (Arora et al. 2017 Montavon et al., 2018 Neyshabur al. 2018).
Since data is crucial in this line of research, teacher models have emerged in learning theory as a formalism for structured data. Extensive research has been conducted this topic, particularly when the trained (student) model is a NN, offering precise and non-asymptotic guarantees in various
contexts (Zhong et al. 2017 Goldt al. 2019 Ba et al. 2020 Sarao Mannelli al. 2020 Zhou
et al. 2021 1021yaaiyaaa & Suzuki 202 Abbe et al. 2022 Ba et al. 2022 Damian et al. 2022 Veiga
al. 2022 Mousavi-Hosseini al. (202)). Experimental evidence suggests that traditional learning
theory fails to fully explain the generalization properties of large NNs, highlighting the need for more
modern approaches (Zhang et al. 202T).
Code available http:///hhhicccmm//////2S/uuuund

---

Published as a conference paper at ICLR 2026
An important concept that frequently appears in modern learning theory is the implicit regularization effect introduced by training dynamics (Neyshabur et al. (015a). The work of (Soudry et al. 2018) sparked wave of recent studies investigating how gradient descent (GD) naturally tends to favor lower-complexity models, often leading to minimum-norm and/or maximum-margin solutions even without explicit regularization (Gunasekar al. 2018 Li al. 2018h) Ji & Telgarsky 2019 Gidel iet al. 2019 Chizat & Bach, | Pesme et al. 2021) However, much of this research focuses on
linear models or excessively wide NNs, with varying interpretations of reduced complexity and its impact on generalization. A notable example in this context is compressibility and its relationship to generalization (Arora et al., 2018) (Suzuki al. 2020). When trained NN can be compressed into a smaller model with similar predictive behavior, both models show comparable generalization performance. This suggests that the original NN's complexity may be understood via its simpler compressed form, which is traditionally associated with improved generalization.
key contribution in the area, and influence for our work, is the work of (Mousavi-Hossene et al. 2023), which studies the training dynamics of a two-layer NN using stochastic gradient descent (SGD) on data drawn from a multiple-index teacher model. First, they show that low-complexity structures emerge during training when strong regularizer is used: on (first-order) stationary points, the first layer weights align with key directions in the input space, the principal subspace. Now-dimensional structure of this type is known help with generalization (Neyshabur et al., 2015b] Bartlett et al. 2017). As a second step, they establish GD-specific generalization guarantees under additional assumptions. Our focus is on the first step, and the goal of this work is to answer the question:
Can we discover low-rank structure in neural networks under more natural assumptions?
To this end, we consider a more precise and well mottvatee solution concept, a p-approximate second-order stationary point (p-SOSP). (Jin et al. 2017), see Definition f.2 for a formal definition. Using the properties of p-SOSPs, we provide a general derandomization lemma. When applied specifically to NNs, lemma implies that, for any arbitrarily small regularizer value, there exists >> 0 such that all p-SOSPs correspond to low-rank first-layer weights. Thus, we significantly relax the sufficient conditions for uncovering this kind of structure as we allow (a) NNs of arbitrary size and depth, (b) with all parameters trainable (including biases), (c) under any smooth loss function, (d) arbitrarily small regularization, and (e) trained by any method that attains a p-SOSP, for example, SGD with random initialization, PGD and Hessian descent.
Importance of training the biases. For the sake of analytical simplicity, some past work has frozen the biases, for example in (Mousavi-Hossein et al.) (023). Here, we show that training the biases is necessary for derandomization to take effect. To illustrate this point, we consider the following toy example. Let x: be a one-dimensional standard Gaussian random variable, and suppose the target label is fixed at 1, i.e., the output lies in a zero-dimensional space. We model the prediction using single-layer NN. The objective is to minimize the loss function:
f(w, b) E (ReLU (wr + +b) -)) +
1)²]
Aw=,
where w and b denote the weight and bias parameters. respectively, and À >> 0 is the regularizer. Directly applying the result from Hosseevi solution lies on a zero-dimensional space. However, this solution is evidently suboptimal when b a 1. Consequently, enforcing this behavior requires an artificially large regularization X. We illustrate this phenomenon in Figure where the minimizer w* approaches zero only under large values of X.
Hosseini al., [023). implies that w = O, i.e. the
In contrast, our analytical approach, which allows the training of biases, avoids this drawback by considering p-SOSPs. Instead of requiring
artificially large regularization parameter to explain this structural behavior (w = 0), we only need a tiny amount of regularization, as shown in Figure 2 because the bias term can adjust to b = 1. This shows that freezing the biases places unnecessary restriction, whereas allowing them to be trained explains the phenomenon more directly and under milder conditions.
Training the biases allows for smaller regularizers, but proving low-rank solutions remains challenging. In (Mousavi-Hosseni et al. 22033), a strong regularizer is used to show that first-order stationary point (FOSP) can be high-rank; without it, higher-rank FOSPs may exist.
Discussion on SOSOSS
2

---

Published as a conference paper at ICLR 2026
Optimal w Reguuarizzaln izalion Patdetie A
Optimaa w' and Curves for Different
app
Optimal W
p
0.0002
0.2
0.005
0.0003
Optimal Iy
0.000
0.0
0.003
-0.2
0.001
-0.4
0 000
90
0.
18
15
2.0
2.5
0.00
0.23
0.50
0.73
2.00
Regularization Paramater
Regiataization Parametor
Figure 1: Plot of the global minimizer of
À for f(w, b) [[ + b) 1)) + Aw when biases are trained. The solid line shows convergence to O at an exact SOSP.
Figure 2: Optimal
fws 0) = E[(ReLU(wwz) - 1) + Aw as a function of the regularization parameter A.
We address this analytical challenge by exploiting the extra properties of p-SOSPs, which bound the negative curvature. The p-SOSP solution concept excludes all but the flattest of saddle points, in other words, p-SOSP more likely to local minimum than the corresponding approximate FOSP. Using this, we show that for sufficiently small the only valid solutions are low-rank.
We remark that our focus on the landscape of the objective (Equation D. not specific optimization methods. Importantly, P-SOSPs capture the solutions reached by standard algorithms: GD/SGD almost surely avoid strict saddles (Panageas et al. 2019). and PGD (Jin et al., 2017) guarantees efficient convergence to p-SOSPs. Crucially, all local minima p-SOSPs. Furthermore, empirical observations show that gradient-basee based methods reliably reach good minima while avoiding saddles (Li al. 2018 8a Zhou al. 2020) and that these solutions have good learning properties. Thus, the focus on p-SOSPs is well motivated as it captures the types of solutions seen in practice and theory, while also providing framework to study the mechanisms of implicit regularization.
Summary of our contributions. Our key contributions can be summarized as follows:
We consider general family of functions of the form:
f(W, b: Ez [ge W x + +b] Al W 12:.
where W E Rxd b E R2 x ~ NO Ia), (o( & + R denotes a parameterized nonlinear function and À > is regularization parameter. This formulation is highly general and encompasses a wide range of applications, including the population risk of NNs of arbitrary depth and architecture, under any smooth loss functions.
Below we present our key derandomization lemma, which forms the foundation of our results, stated informally as follows:
Let f be a twice differentiable function in the form of Equation where z ~ A000, Ia). Then, all SOSPs of f satisfy W = 0.
Informal version of Lemma 3.1
Structure discovery in NNs. We first show that the regularized risk of any NN can be expressed as function of the form in Equation[]] Then, by applying our key derandomization lemma, we establish that the teacher-student setting any p-SOSP solution of the risk yields first-layer weight matrix W which is near low-rank. This type of structure is closely associated with improved generalization et al. 2023).
performance as shown in (Mousavi-Hosseee providing guarantees for structure discovery under much broader and more minimal assumptions, thereby capturing a wider class of models and cases of arbitrarily weak regularization.
this paper, our express focus is on
In other words, main result applied to NNs suggests with weak regularizers. Recall that standard methods like SGD, GD, or PGD are known to escape saddle points; when that happens, results guarantee solution with good structure (i.e. low rank).
explanation for parsimony
3

---

Published as a conference paper at ICLR 2026
Our results in other domains. Due to the generality of our derandomization lemma obtain strong theoretical results across domains. For example, we get (i) a deterministic MAXCUT ap- proximation matching the randomized guarantee of (Goemans & Williamson (995), and (ii) deterministic construction for learning contribution is show that derandomization be achieved via In the case of MAXCUT, simple gradient- optimization, without relying on explicit combinatorial constructions or pseu- dorandom generators. To the best our knowledge, this the first optimization-bant derandomizing the Goemans- Wiillimmsn algorithm, and we view it as a conceptual contribution that enriches the literature on derandomization. In the case of JL, we match the state-of- oo--th--tthtrrrereee result of (Tsikouras et al., 2024) further demonstrating the generality of our lemma. We expect the same approach to extend to other domains.
(JI.) embeddings (Johnson et al. 1984).
approach for
NOTATION AND PRELIMINARIES
Notation. For vectors u.
(u, v) v to denote their inner product and ||u|| to denote the L2 norm. For matrix M E Rexdd we denote the element of the 7th row and jth column by || to denote the Frobenius norm. We use V f and V denote the gradient and Hessian operators, respectively. Additionally, we use A ~ N (M, E), where M = (f4,,) and E = (o).)) indicate that A (as.s) is matrix with independent random entries, and each entry follows 0jjj N (J4jj Ri.)). Using this notation we can write A = M + Z, where Z ~ NN(,, E). Definition 2.1. A twice differentiable function f Rd + RR is defined to be L-smooth, if for all x, y E Rd satisfies:
Mi.j and
[V V S L|| g/12- The function K-Hessian Lipschitz for all æ, y E Rd:
f(=)
& K| x y/12-
Below, we give the definition for approximate stationarity.
Definition 2.2 (Approximate second-order stationarity). For K-Hessian fincciit J()), we say that point x" is a p-second-order stationary point (p-SOSP) if:
[Vv 1l S p and Xmiu (V2 2 VKp.
Assumption 2.3. The function is both L-smooth and K K-Hessian Lipschitz.
Provided Assumption 2.3 holds, Algorithm converges to SOSPP in O(1/p²) iterations with high probability (Jin et al. 2017). Alternatively. Algorithm 2 achieves a deterministic p-SOSP in O(1/p) (1./p iterations (Tsikouras al. 2024), but requires Hessian access.
3 MAIN CONTRIBUTION: KEY DERANDOMIZATION LEMMA
In this section, we prove that convergence to SOSPs is a sufficient condition for derandomization. Let W gexd b E A (-) : IR* + E be function satisfying Assumption 2.3 and let ~N Ia). We analyze the behavior of the following objective function at its SOSPs:
f(W, b; 0) = Ex W 6)] + A|| W 12:,
(2)
where À > is an arbitrarily small regularization parameter. Our main result is presented below.
Lemma 3.1 (Key Derandomization Lemma). Let x (O, Iz) be a standard multivariate Gaussian
random variable. For the objective function defined in Equation 2 with
<Ke where go (-)
satisfies Assumption 2.3 any p-SOSP satisfies [| 2A-VKp
Proof sketch: The key observation is that applying Stein's Lemma Stein 1973
1981 relates
with the first derivative of W. This allows us to express the first-order conditions for W in terms of expectations involving the second derivatives of ga- Using these relations and an approximate first-order optimality condition, we derive the required bound on the Frobenius norm of W. Full proof can be found in Appendix B.T
the second derivative
4

---

Published as a conference paper at ICLR 2026
Remark 3.2. Note that achieving perfect SOSP (i.e. p 0), would result in W O; the proof of this is in Appendix B.2
This result shows that when the objective function takes the form given in Equation 2 common structure in practice, it is possible to improve it by minimizing the inherent randomness. Since the input a is random, having a small [ ensures that Wa a remains small on average, so the values of go (Wa x: + b) vary less. In other words, second-order stationarity implies that the randomness in the objective function effectively decreased. We illustrate the implications of this result by applying it to three distinct examples from different fields, as demonstrated in the following sections. We emphasize that the p-SOSP is taken with respect to all model parameters, see Appendix ?? for further discussion.
Additionally, we clarify that our amalysis operates under the assumption that the optimization method employed converges p-SOSP with respect to all model parameters, including those encapsulated in A. Specifically, we assume that an approximate SOSP is identified jointly over the entire parameter space. Once such a point is found, observe that fixing the auxiliary parameters preserves the approximate second-order stationarity with respect to the first-layer weight matrix W.
It important to note that tiny amount of regularization is necessary ensure well-defined solution in this lemma. Without it, choosing the function ga (wr (a(wr + b) = O, would result in all w E IR being local minima, as the objective provides no preference among these values. Introducing a tiny regularization term À climinates this anbiguity by penalizing non-zero weights, thereby enforcing 0 as the unique optimal solution.
Additionally, À is fully controllable by which is set prior to the optimization process, As a result, it is possible to use arbitrarily small regularization, provided one willing incur the additional cost of executing the optimization algorithm for a greater number of steps. We emphasize this point to remind the reader that only minimal amount regularization is necessary. Therefore, the regularization term itself is not the primary factor influencing the outcome; rather, it is the interaction between the second derivative of b and the first derivative of W that plays a more significant role in the optimization process.
is important to clarify that Lemma 3.1 serves as a structure discovery result rather than an optimization result. In particular, focus lies not the specific optimization algorithm employed, but rather on establishing that any solution satisfying the p- p-SOSP condition necessarily reveals structure. There might be different methods under many different sets of assumptions that yield a P~SOSP solution, even in more practical and realistic finite-sample regimes. Nonetheless, for completeness, note that PGD) (Tin et al. 2017). originally developed for deterministic objectives such as the population risk, has been shown to efficiently yield p-SOSP solutions in polynomial time in stochastic and finite-sample regimes (see Theorem 15 in (Jin et al. 2018
STRUCTURE DISCOVERY IN NEURAL NETWORKS
In this section, we present our primary application, which builds on the central derandomization Lemma Our approach extends the work of (Mousavi- et al., (023). which demonstrated a key insight: the convergence of the first-layer weights to a dow-dimensi subspace. We refine
and generalize these results to broader setting. Let æ E Rd be a standard Gaussian distribution A(0,, Ia). The target labels are generated by multiple-index teacher model of the form:
y h( (u,, .)...., (ux, e) = h(Ux;e),
where h : R+++ - IR is weakly differentiable link function and represents additive noise. For any vector V1 := v Un- For a matrix W e R we define Wi and W1 by projecting each row of W similarly. Using this notation, we rewrite the labeling function to depend only on the æn component:
IRdd let in denote its orthogonal projection onto span(u1..... uz), and define
y hUx)) h 1'(2z1).
Our goal is to show that the perpendicular component W, converges to zero, implying that the first-layer weight matrix W lies entirely in the teacher subspace.
5

---

Published as a conference paper at ICLR 2026
MMossav--osseini al. 2023) showed that, under certain conditions, training only the first-layer weights of a two-layer NN suffices to ensure convergence to the low-dimensional principal subspace defined by the teacher model. We extend this result by showing that training the first-layer bias is also essential, Including the bias enables a simpler and more direct analysis the training dynamics.
This expanded approach allows general conditions, including: (a) arbitrary regularization parameters > (resolving an open question in earlier work), (b) arbitrary NN size and depth, (c) all parameters being trainable (including biases), and (d) any choice of smooth loss function.
establish convergence guarantees under significantly more
4.1 DISCOVERING STRUCTURE IN NEURAL NETWORKS VIA SOSPS
Let W RRx and b T& and consider the first layer of NN given by Wæ b. We decompose this expression into components that are parallel and perpendicular to subspace U, as follows:
Wæ + b = Wlll W a b,
since WI 0 and W æ1 = 0 due orthogonality.
Now define the NN's prediction as j(æ; W. b) go Wa + b). where go(-) is a NN of arbitrary
size and depth, parameterized by 0, which satisfies Assumption 2.3 Given a loss function f(y, ú) that
also satisfies Assumption 2.3 and regularization parameter À U, we define the regularized risk as:
R(W (W, b;0) F. (v. ((:; W, b, 0)) W||
As shown in Appendix D Equation Sattiionnnaaa ean be reformulated way that enables direct application of Lemma 3.1 In particular, this reformulation expresses the regularized risk as function of the perpendicular components:
R(W b² = Ezz 1 (W x - b)] + W| 1
In this form, the parallel and perpendicular components are fully decoupled. The modified loss Eg. implicitly depends on the parallel components, the corresponding regularization, and all other parameters of the NN. By applying Lemma 3.1 we conclude that [| |r can be made arbitrarily small, implying that the perpendicular components are effectively suppressed during training.
Theorem 4.1. Consider an arbitrary NN of any size and depth that satisfies Assumption 2.3 and loss function that also satisfies this assumption. Let the input data ~ ~ (0, La) be standard multivariate Gaussian distribution, and the labels generated according to Equation 3 If we minimize Equation 6 with respect (W, b), then for any precision parameter p and regularization parame- tke where K denotes the Hessian Lipschitz constant of the objective, all p-SOSPs satisfy ter X > 2
the inequality:
P Wille S 22 - /Kp
Proof. Since both the NN and the loss function are smooth and Hessian Lipschitz, their composition inherits these properties. The result then follows directly from Lemmal 3.1
The result establishes the theoretical validity of our approach but qualitative nature, as it does not specify the number of steps required for optimization. To complement this, provide quantitative result demonstrating that the objective function can be minimized efficiently.
Theorem 4.2, Let the objective function in Equation 5 be twice differentiable, bounded below, and satisfies Assumption 2.3 Let the data
(N Ia). and suppose labels are generated according to Equation 3 Let > O be a prespecified accuracy, and define the regularization parameter X A where K is the Hessian Lipschitz constant of the objective and A > 0 an arbitrarily small constant. Then, with probability 8, running Algorithm I for T > O (poly L, Log(d), log(0),, -1)) iterations with a step size O() yields weight matrix W = W. + Wi that satisfies:
Willr <..

---

Published as a conference paper at ICLR 2026
Proof. Full proof can be found in Appendix E.1
This result demonstrates that, with a number of samples that scales with the Hessian Lipschitz constant, PGD can effectively generate iterates that are close to the principal subspace required. This allows the model to learn low-dimensional representations, and introduce an implicit bias toward simpler, lower-complexit solutions. The discovery of structure appears be an inherent characteristic of this optimization process for problems this nature. This convergence to a low- dimensional solution often linked with the generalization behavior of NNs (Neyshabur al. 2015b) Bartlett et al. 2017 Arora et al., 2018 Suzuki et al., 2020 Mousavi-Hossee do not attempt to establish generalization guarantees here, as such results would require additional assumptions on the data distribution or the hypothesis class. Instead. our contribution is to provide guarantees for structure discovery under broader and more minimal conditions, thereby extending the potential applicability of these ideas to a wider range of models.
et al.] 2023). We
Our results pertain to NNs with smooth activation functions; here we discuss the wide practical appli- cability of this family of networks. Smooth nonlinearities are widely used in modern architectures, and there no strong evidence that non-smooth activations outperform their smooth counterparts. For example, BERT adopts the gaussian error linear unit (GELU) (Hendrycks & Gimpel] (2016] Devlin et al. 2019), a smooth activation that has been shown to benefit from this choice compared with non- smooth alternatives. Thus, assumption is aligned with standard practice. We provide additional insights for the non-smooth ReI.U case by employing a smooth approximation in the next section.
4.2 THE CASE FOR RELU
The non-smoothness of ReLU poses challenges for our framework. To ensure the smoothness and Hessian Lipschitz continuity needed for defining p-SOSPs, we use a smooth approximation of ReLU:
1 1
ReLU,(x) log (1 + tct).
As > 00, the function converges to the standard ReLU. Moreover, it is 4 i-gradient Lipschitz and v3.² -Hessian Lipschitz, ensuring that framework remains valid for any smooth ReLU approximation. In this sense, ReLU, captures the essential behavior of ReLU while enabling theoretical guarantees. To show the dependence on i we give the following theorem, regarding one layer NN using activation function R-LU,(-).
Theorem 4.3. Assume that the data x ~ (O, Ia), and labels are generated according to Equa- tion Additionally, consider the NN a TRRLU, is twice differentiable. bounded below, and satisfies Assumption 2.3 with gradient and Hessian Lipschitz constants Le and Kr, respectively. Let p > be prespecified accuracy and define the regularization parameter X = vKo+A, where = O is the overall Hessian Lipschitz constant and A 0 is arbitrarily small. Then, with probability 1 5, running Algorithm for T > O (poly (1, Le, log(d), 1 A11) iterations, with a step size ((1((L))),, yields a weight matrix W W_ Wi that satisfies:
b) and the objective in Equation 5 which
Wille <E.
Proof. This is a direct application of Theorem 4.2 The composition of the objective with ReLU, satisfies Assumption 2.3 and is lower bounded.
4.3 NEURAL NETWORKS EXPERIMENTS
In this section, we empirically validate theoretical framework by showing that the first layer weight matrix of the student network converges to the principal subspace. The teacher network is single-index model that generates outputs according y tanh(0 æ) + noise, where R² and 0 = 7 (1, 1) is a fixed direction. The student network attempts to learn this mapping using a two-layer NN of the form
y = a. ".RRel.2(WW + b).
7

---

Published as a conference paper at ICLR 2026
Weight Progression With Perturbed Gradient Descent
4
span(u)
Initialized
Trained
0
-2
Figure 3: Two-layer ReLU2 network width = 1, 00 and = 2 for the task of recovering tanh single-index teacher model. We observe convergence of weights W to the principal subspace.
where W Rixdd Rha and TRA All parameters randomly initialized, except that the second layer weights a are kept fixed throughout training to isolate the representation learning dynamics of the first layer. Full details can be found in Appendix F
Empirically, we observe that the initially random rows of W align with the teacher direction and converge the principal subspace spanned by gradient-based optimization recovers the underlying signal structure despite random initialization. Moreover, we state for completeness, as rigorously established the emergence of such low-rank structure reduces the effective dependence of generalization error on the ambient input dimension d.
illustrated in Figure 3 This confirms that
Moousavi-Hosseing al. (023).
Training is conducted using PGD. At each iteration, if the gradient norm falls below the threshold 10-6, isotropic Gaussian noise is injected into the parameters. Concretely, for the first-layer weights,
W < W + 5.00
with 8 = 0.005 analogous perturbations are applied the bias terms b. Although this threshold is small, the noise is applied many times during training. ensuring exploration of flat regions of the loss landscape. This modification of standard SGD helps the optimization escape saddle points and flat regions, which is particularly useful in our setting and aligns with our theoretical guarantees.
Training proceeds for 10,000 steps, minimizing the Mean Squared Error (MSE) loss function with L2 regularization applied only to the first layer weights W with coefficient À = 10- Learning rates are set to = for W and 7p- = 1 for the b.
OTHER APPLICATIONS
5.1 MAXCUT
The MAXCUT problem is a classical combinatorial optimization problem that seeks to partition the vertices of a graph (V. E) into disjoint sets, S and such that the sum of the weights the edges crossing between S and T is maximized. The objective function for the MAXCUT problem is:
Wiy x4H2))
Maximize:
subject to
24 E (-1,1), Wi W.E.
2
(EE)
8

---

Published as a conference paper at ICLR 2026
where Way denotes the weight of the edge (i.j) and the term 1 2,2, equals if edge (i.j) crosses the cut and otherwise. We follow the common assumption that Wi, = 1. This is a combinatorial optimization problem that is NP-complete (Karp 1972).
To make the problem more tractable, (Goemans Williamson, 1995) proposed using a semidefinite program (SDP) relaxation. In this relaxation, the discrete MAXCUT problem is lifted to continuous one by representing each vertex as a unit vector v; E R on the unit sphere, where m is the number of nodes in the graph. The algorithm first solves the SDP obtain these vectors. then applies randomized rounding procedure to map the continuous solution back to a discrete cut: a standard Gaussian vector
(O, Im) sampled, and each vertex is assigned to one of the two S or T based on the sign of the inner product (v, z)-
We aim to approximate the MAXCUT problem by derandomizing this rounding algorithm. Let V mm be the matrix of SDP vectors. Define Vz + μ., where ~ (0,, Im) and R is mean vector. Our goal is to minimize the negative expected cut value, leading to the regularized objective function:
S(V. 14) -->=--m E Wig Pr [sgn((v;; = ps) & sgn(w; z ++)1 VI i<}
The probability term in Equation 7 can be interpreted as the expectation of indicator function for the event inside the probability. To allow the application of Lemma 3.1 we replace this
Remark 5.1.
indicator function with smooth e-approximation, described in Appendix G.
We now present our main result on the derandomization of the randomized MAXCUT algorithm.
Theorem 5.2 (Derandomized Approximation for MAXCUT), Let G = (V, E) be a graph with m edges, where the edge weights are given by Wi. Wi = 1 for all (i.j) E E. Let V E mm be the matrix of vectors obtained from the SDP relaxation of the MAXCUT problem, as described in (Goemans & Williamson 1995). Denote by z _N(0, Im) standard multivariate Gaussian vector, and let μ represent a vector of means. Initialize Algorithm 2 with μ = O, and optimize the e-smoothed version of the objective function in Equation (see Equation 27,. using the regularization
parameter X v/le+A and run for =
A-))) iterations. After this optimization process, the resulting vector μa defines cut whose value is guaranteed to be at least:
O (poly (log(m),,og))
OPT -
with probability 1 - Here, a = 0.878 is the approximation factor (Goemans Williamson 1995).
Proof.
Full proof can be found Appendix G.I
To the best of knowledge, this work is the first optimization- based derandomization of the MAXCUT problem: rather than relying on the method of conditional expectations, small-bias spaces, explicit pseudorandom constructions (Naor & Naor 1990 Motwani & Raghavan 1995 Mahajan
& Ramesh 1995). In addition, we have conducted experiments demonstrating our optimization-ban construction for solving MAXCUT, Full details are provided in Appendix G.2
5.2 JOHNSON-LIDEENSTRAGSS EMBEDDINGS
The JI. Lemma is well-know result in the field of dimensionality reduction (Jonsson et al. 1984). Specifically, consider unit norm data points T1. T, E Rd, which we aim to project into dimensions while preserving their norms with at most e-distortion. Here, the distortion is given by
E == (logn/k) A detailed description of the JL Lemma is given in Appendix H.1 In this context, we are interested in finding matrices that satisfy the JL guarantee:
Definition 5.3 (JL guarantee). The JL guarantee states that for given dataset æ1..... Cn E Rd and target dimension k, the distortion for all points does not exceed O(Vlog n/k).
Significant research has been dedicated to improving the construction of random projections (Idyy & Motwani 1998 Achlioptas 2001 Matoušek 2008). In contrast to these traditional methods,

---

Published as a conference paper at ICLR 2026
approach recovers the result from Tsikouras al. 2024), which proposes learning the linear mapping directly from the data, deterministically. Other derandomization methods for JL include (Engebretsen
et al. 2002. Mcka & Zuckerman 2010).
Let A be a random matrix whose entries at are independently drawn from Gaussian distribution with means j4.j and variances o².j" Let Σ denote the matrix collecting these variances. Our goal is to minimize the following quantity:
Σ 11/2
Pr
max
| 12
1/ >E>>E
(8)
i=
2kd
which represents the probability that the maximum distortion across all input vectors exceeds prescribed threshold E, augmented by a regularization term that penalizes large variances. Notably, the regularizer vanishes as Σ O, recovering deterministic transformation in the limit.
As shown in Appendix H.T we use a union bound to relax the original objective in Equation g
reducing equivalent surrogate objective:
E1/2
{x
14)
Epr
[1(2))/2
+
(9)
2kd
where 2 ~ (N, The optimization performed over the parameters (Σ p.).
Remark 5.4. The probability term in Equation 9 can interpreted as the expectation of an indicator
function for the event inside the probability. To allow the application of Lemma 3.7 we replace this
indicator function with smooth E -pppppoxximation, as described in Appendix H.3
We now present our main result the derandomization of the JL Lemma.
Theorem 5.5. Let be unit vectors in IRd, k be the target dimension, be smoothening parameter and A > 0 be an accuracy purameier. For any ε 2 C./log n/k, where C. is sufficiently large constant, initialize M = and Σ = Ind and run Algorithm 7 to optimize the E sm-oohhhed version of the objective function in Equation 9 (see Equation 27 using the regularization parameter
A Vole+> After satisfies the JL guarantee with distortion at most Os)), with probability 5.
(poly (n, k d, 1, A-1)) iterations, this returns matrix M that
2
Proof. Full proof can be found in Appendix H.4
In addition, we have constructed an experiment and show that indeed optimizing a distribution over projection matrices can reduce JL distortion beyond standard Gaussian constructions by directly minimizing worst-case distortion. Full details can be found Appendix H.5
CONCLUSION
We study the theoretical properties of NNs under specific conditions, showing they can discover low-rank structures. Building on (Mousavi-Hosseini et al. 2023), we extend their framework to allow (a) NNs of arbitrary size and depth, (b) all parameters trainable, (c) any smooth loss function, and (d) minimal regularization. The core of our analysis is the derandomization Lemma 3.1 which ensures effectiveness even with small regularization. Training biases is a common practice, and our theory guarantees that it can improve model performance. The strength of our lemma is demonstrated in three applications, mainly in NNs and secondarily in MAXCUT and JL embeddings.
Finally, we outline results rely on the assumption that the input distribution is Gaussian. Extending these findings to other distributions is interesting avenue for future research. Additionally. it would be valuable explore connections between our theoretical results and the learning and generalization guarantees that are observed in practice.
limitations of our current work and suggest future research directions. Our
10

---

Published as a conference paper at ICLR 2026
7 ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
The authors would like to thank Alireza Mousavi-Hosseing for useful discussions and feedback. This work has been partially supported by project MIS 5154714 of the National Recovery and Resilience Plan Greece 2.0 funded by the European Union under the NextGenerationEU Program. Ioannis Mitliagkas acknowledges support by Archimedes, Athena Research Center, Greece and a Canada CIFAR AI chair. The majority of work performed at Archimedes in Athens.
REFERENCES
Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems, 34:26989-27002, 2021.
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. Conference on Learning Theory, pp. 4782-4887. PMLR, 2022.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pp. 2552-2623. PMLR, 2023.
Dimitris Achlioptas. Drienlse- random projections. In Proceedings of the twentieth ACM symposium on Principles of database systems, pp. 274-281, 2001.
Shunta Akiyama and Taiji Suzuki. On learnability via gradient method for two-layer relu neural networks in teacher-student setting. In International Conference on Machine Learning, pp. 152- 162. PMLR, 2021.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International conference on machine learning, pp. 242-252. PMLR, 2019.
Sanjeev Arora, Yingyu Liang, and Tengyu Ma. simple but tough-to-beat baseline for sentence embeddings. In International conference on learning representations, 2017.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In International conference machine learning, pp. 254-263. PMLR, 2018.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high- sigh--imensional inference. Journal of Machine Learning Research, 22 (106):1-51, 2021.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. In International conference on learning representations, 2020.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37944 2022.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21 55:21750-21764, 2022.
Peter L Bartlett, Dylan Foster, and Matus J Telgarsky. Spectrally-normalizee margin bounds for neural networks. Advances in neural information processing systems, 30, 2017.
Raphaöl Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations Computational Mathematics, pp. 1-84, 2024.
11

---

Published as a conference paper at ICLR 2026
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-inder models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768 -9793, 2022.
Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivies On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310. 2a:::22300.199999 19793, 2023.
Guillaume Braun, Minh Ha Quang, and Masaaki Imaizumi. Learning a single index model from anisotropic data with vanilla stochastic gradient descent. arXiv preprint arXiv:2503.23642, 2025.
Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory, pp. 1161-1227. PMLR, 2020.
Siyu Chen, Beining Wu, Miao Lu, Zhuoran Yang, and Tianhao Wang. Can neural networks achieve optimal computationallstaisnn stattssiia tradeoff? an analysis on single-index model. In The Thirteenth International Conference Learning Representations, 2025.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over aver-parmmeterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
Lenaic Chizat and Francis Bach. Implicit bias gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference learning theory, pp. 1305-1338. PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training differentiable programming. Advances in neural information processing systems, 32, 2019.
Alex Damian, Loucas Pillaud-Vivien, Jason D Lee, and Joan Bruna. Computationnal-staisiicaal gaps in gaussian single-index models. arXiv preprint arXiv:2403.05529, 2024.
Alex Damian, Jason D Lee, and Joan Bruna. The generative leap: Sharp sample complexity for efficiently learning gaussian multi-index models. arXiv preprint NXKv:25066005000 2025.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pp- 5413-5452, PMLR, 2022.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33:20356-20365, 2020.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Ilias Diakonikolas, Daniel M Kane, Vasilis Kontonis, Christos Tzamos, and Nikos Zarifis. Ag- nostically learning multi-index models with queries. In 2024 IEEE 65th Annual Symposium on Foundations of Computer Science (FOCS), pp. 1931-1952. IEEE, 2024.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-param peeeparraeteri neural networks. In International Conference on Learning Representations, 2018.
Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference Learning Theory, pp. 1887-1930. PMLR, 2018.
Lars Engebretsen, Piotr Indyk, and Ryan O'Donell. Derandomized dimensionality reduction with applications, 2002. Manuscript, Cameeiic Mellon University.
Spencer Frei, Yuan Cao, and Quanquan Gu. Agnostic learning of a single neuron with gradient descent. Advances in Neural Information Processing Systems, 33:5417-54288 2020.
David Gamarnik, Eren C Kızilda, and Ilias Zadik. Stationary points of shallow neural networks with quadratic activation function. arXiv preprint arXiv: /9215999/2.01599, 2019.
12

---

Published as a conference paper at ICLR 2026
Ankit Garg, Neeraj Kayal, and Chandan Saha. Learning sums of powers of low-degree polynomials in the dendeeenerete case. 2020 IEEE 61 Anmual Symposium on Foundations of Computer Science (FOCS), pp. 889-899. IEEE, 2020.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Margalit Glasgow. Sgd finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the xor problem. arXiv preprint arXiv:2309.1511 15111, 2023,
Michel X Goemans and David P Williamson. Improved approximation algorithms for maximum cut and satisfiability problems using semidefinite programming, Journal of the ACM (JACM), 42(6); 1115-114 1145, 1995.
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
Ian Goodfellow. Deep learning, 2016.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
arXiv:1606.00 wet.16606.8415,, 2016.
Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing. pp. 604-613, 1998.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference learning theory, pp. 1772-1798. PMLR, 2019.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How escape saddle points efficiently. In International conference on machine learning, pp. 1724-1732. PMLR, 2017.
Chi Jin, Lydia T Liu, Rong Ge, and Michael Jordan. On the local minima of the empirical risk. Advances in neural information processing systems, 31, 2018.
William Johnson, Joram Lindenstrauss, Extensions lipschitz mappings into hilbert space. Contemporary mathematics, (6(189-206))) 1984.
Richard M Karp. Reducibility among combinatorial problems. In Raymond E. Miller and James W. Thatcher (eds.), Complexity of Computer Computations, pp. 85-103. Springer, 1972.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6941.6 2014.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning, nature, 521(7553):436 444, 2015.
Jason Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theorce limit. Advances in Neural Information Processing Systems, 37:58716-55877 2024.
Hao Li. Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018a.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang, Algorithmic regularization in wer-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 80.
13

---

Published as a conference paper at ICLR 2026
Cosme Louart, Zhenyu Liao, and Romain Couillet. A random matrix approach to neural networks. The Annals of Applied Probability, 28(2):1190-122 1248, 2018.
Sanjeev Mahajan and Hariharan Ramesh. Derandomizing semidefinite programming based approxi- mation algorithms. In Proceedings of IEEE 36th Annual Foundations of Computer Science, pp. 162-169, IEEE, 1995.
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:57367-57480, 2023,
Simon Martin, Francis Bach, and Giulio Biroli. On the impact of overparamettrizzatin on the training of shallow neural network in high dimensions. In International Conference on Artificial Intelligence and Statistics, pp. 3655-3663. PMLR, 2024.
Jifí Matoušek. On variants of the Sechnson-lindenstinn lemma. Random Structures Algorithms, 33(2):142-156, 2008.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75 (4):667-766, 2022.
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, 115(33):E7665-77771 2018.
Raghu Meka and David Zuckerman, Pseudorandom generators for polynomial threshold functions. In Proceedings of the Forty-second ACM Symposium on Theory of Computing, pp. 427-436, 2010.
Grégoire Montavon, Wojeiech Samek, and Klaus-Robert Müller. Methods for interpreting and understanding deep neural networks. Digital signal processing, 73:1-15, 2018.
Rajecv Motwani and Prabhakar Raghavan. Randomized Algorithms. Cambridge University Press, 1995.
Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. The Eleventh International Conference Learning Representutions, 2023.
Joseph Naor and Moni Naor. Small-bias probability spaces: Efficient constructions and applications. In Proceedings of the twenty-secone annual ACM symposium on Theory of computing, pp. 213-223, 1990.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. International Conference on Machine Learning. 2015a.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Norm-based capacity control in neural networks. In Conference on learning theory, pp. 1376-1401. PMLR, 2015b.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of over-parametrrizatoo in generalization of neural networks. arXiv preprint arXiv:1805.1200 120.2200....2..... 2018
Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning of diverse features: putational hardness and efficient gradient-based training for ridge combinations, arXiv preprint arXiv:2406.1 10.6.11880 11288, 2024.
Joannis Panageas, Georgios Piliouras, and Xiao Wang. First-order methods almost always avoid saddle points: The case of vanishing step-sizes. Advances in Neural Information Processing Systems, 32, 2019.
Scott Pesme, Loucas Pillaud- Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218-22 2923000, 2021.
14

---

Published as a conference paper at ICLR 2026
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Stefano Sarao Mannelli, Eric Vanden- Eijnden, and Lenka Zdeborová. Optimization and generalization of shallow neural networks with quadratie activation functions. Advances in Neural Information Processing Systems, 14113445-13455 2020.
Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information processing systems, 30, 2017.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018.
C. Stein. Estimation of the mean of multivariate normal distribution. In Proceedings of the Prague Symposium on Asymptotic Statistics, 1973.
Charles M Stein. Estimation of the mean a multivariate normal distribution. The annals of Statistics, pp. 1135-1151, 1981.
Taiji Suzuki, Hiroshi Abe, and Tomoaki Nishimura. Compression based bound for com-compeessed network: unified generalization error analysis of large compressible deep neural network. In International conference on machine learning, 2020.
Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond. Advances in Neural Information Processing Systems, 36, 2024.
Nikos Tsikouras, Constantine Caramanis, and Christos Tzamos. Optimization learn johnson lindenstrauss embeddings. Advances in neural information processing systems, 2024.
Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. Advances in Neural Information Processing Systems, 35:23244 232255, 2022.
Puqian Wang, Nikos Zarifis, Ilias Diakonikolas, and Jelena Diakonikolas. Sample and computation- ally efficient robust learning of gaussian single-index models. Advances in Neural Information Processing Systems, 37:58376-58422,, 2024.
Lei Wu. Learning a single neuron for non- momo-oooooonccccc activation functions. In International Conference Artificial Intelligence and Statistics, pp. 4178-4197 PMLR, 2022.
Weihang Xu and Simon Du. Over- peee-parameterzzation exponentially slows down gradient descent for learning a single neuron. In The Thirty Sixth Annual Conference on Learning Theory, pp- 1155-1198. PMLR, 2023.
Gilad Yehudai and Shamir Ohad. Leariing single neuron with gradient methods. In Conference Learning Theory, pp. 3786-3786. PMLR, 2020.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
Kai Zhong, Zhao Song. Prateek Jain, Peter L Bartlett, and Inderjit S Dhillon. Recovery guarantees for one-hidden-layer neural networks. In International conference machine learning, pp. 4140-4149. PMLR, R,, 2017.
Mo Zhou and Rong Ge. How does gradient descent learn features-a a local analysis for regularized two-layer neural networks. arXiv preprini urXiv:2406.01766 2024.
Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly wer-parameterized two-layer neural network. In Conference on Learning Theory, pp. 4577-4632, PMLR, 2021.
Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly over- parr-parameterized two-layer neural network. In Conference Learning Theory. PMLR, 2022."
15

---

Published as a conference paper at ICLR 2026
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 33:21285-212992 2020.
A RELATED WORK
Feature learning in NNs. Despite the established importance of feature learning in NNs, the specifics of how gradient-based algorithms develop useful features remain somewhat unclear. The neural tangent kernel (NTK) framework, primarily used for examining overparameterized NNs, suggests that neuron movement from their initial positions minimal, highlighting the role of NN architecture and initial settings (Jacot et al. 2018 Du et al. 2018 Allen-Zhhu et al. [011 Chizat et al., 2019). Limitations of the NTK framework have led researchers to explore other analytical approaches, such mean-field analysis, initially requiring vast neuron counts (Chizat & Bach 2018 Mei et al. 2018). Later studies have shown that early stages of training, such as initial steps in GD, are crucial for effective feature learning, with the first layer in 2-layer NNs capturing valuable features (Daniely & Malach 2020; Abbe et al. 2021: 2022 Zhou & Ge 2024). This early capture of features by the first layer offers better performance than models relying solely kernel random features. In the exploration of NN and kernel method interconnections, it has become evident that gradient-based training facilitates representation learning, setting NNs apart from kernel methods (Mousavi-Hosseinl et al. 2023 Abbe et al., 2022 Ba et al. 2022 Barak et al. 2022 Damian et al., 2022). A 2-layer NN with untrained, randomly initialized weights epitomizes a random features model (Rahimi Recht 2007), capturing complex phenomena seen in NN practice (Louart et al. 2018 Mei & Montanari
2022). Despite inheriting positive traits from optimization procedures, these cannot be fully expressed
as random feature regression. The implicit regularization aims of the training dynamics, favoring low-complexity models, are widely discussed (Neyshabur et al. 2015a).
Single/Multi index models. NNs are widely studied for learning single-index and multi-index
models, which depend few directions in dighndimensional inputs. Recent works demonstrate
the effectiveness of two-layer NN in learning single-index (Soltanolkotabl 2017] Yehudai & Ohad
2020 Frei al. 2020 Wu 2022 Bietti al.) 2022 Xu & Du 2023 Mahankali al. [023 Berthier
et al. 2024) and multi-index models (Damian et al. 2022 Bietti er al., [023] Glasgow 2023 Suzuki
ctt al. 2024). These studies emphasize the benefits of feature learning over fixed random features. For
multi-index functions representable by compact two-layer NN, a GD variant with weight decay can recover ground-truth directions. Gradient-
learning shows that NNs trained via GD can learn useful representations for single-index (Ba et al.) [2022] Bietti et al. [022)] Mousavi- et al.
2023
Berthier al. 2024 Oko et al., 2024) and multi-inded models (Damian et al.] [2022] et
2022
Bietti et al. (023). Learning complexity is influenced by the information exponent (Arous et al.
2021
or leap complexity (Abbe et al. 2023). While guarantees for low-dimensional models often lead
to superpolynomial dependence, other research examines cases where student NN match the target function's s architecture (Gamarnik et al. 2019; Akiyama & Suzuki, 2021 Zhou et al. 2022 Veiga
et al. 2022 Martin et | (024). This study considers an intermediate case where width scales with dimensionality without assuming a known nonlinear activation, showing GD achieves polynomial sample complexity when target weights are diverse. Additionally, statistical query algorithms address
related polynomial regression tasks (Dudeja & Hsu 2018 Chen & Mcka] 2020 Garg et al. 2020
Diakonikolas al. 2024).
A significant line of recent work investigates the learnability of single-index models via Hermite decompositions under Gaussian inputs. These works show that for single-index targets, the first nonzero Hermite term, captured by the information exponent or, in newer formulations, the generative exponent, governs the difficulty of recovering the index direction using first-order or statistical query-style methods (Arous et al. [021]] Wang et al., 2024 Braun et al. (025). Recent lower bounds based on the generative exponent reveal computationallstatsste
gaps, establishing sharp statistical query and low-degree polynomial hardness results (Damian et al. 2024 2025). Complementary algorithmic results show that gradient- learners can match these limits in certain regimes: two- layer networks with data reuse effectively reduce the relevant Hermite order (Lee et al. 2024), and new SGD-based methods achieve sample complexities near the generative-exponee boundary (Chen
al. 2025
16

---

Published as a conference paper at ICLR 2026
B PROOF OF MAIN LEMMA
B.1 PROOF OF LEMMA 3.1
Lemma. Let (0, La) be standard Gaussian random variable. For the objective function
defined in Equation 2 with À > <K2 where go(-) satisfies Assumption 2.3 any p-SOSP satisfies
[W]| < 2A- 2X-VKp'
Proof.
The first and second derivatives of f(W, ,,, 0) with respect to b are given by:
0f(W, b. 0)
db Eg V go Wx x + b]].
8² f (W, b.
E. V2 go (Wx b)].
b
The first derivative with respect to W is:
@f(W. b,
']+ 2XW. OW Ex Vo W x + b) [1
Using Stein's Lemma in the multivariate case, this can be rewritten as:
@f b, 0)
E Wa x + +b + 2 W.
(10)
dw
At a p-second-order stationary point, the Hessian with respect to b satisfies:
8 ffW,(WWWW b,0)) E. V b]] = Vpl.
(11)
b2b
From Equation
11 it follows that:
Ex 2 go (W b)] + +2X 2AI VKI..
Then dividing both sides with 2A- VKp, we get:
2 go W b) + 2X1
Ex
I.
(12)
/Kp
Using the approximate first-order optimality condition
(f(w.) aw
< p along with Equations
10
F
and 12
have:
P
V
go
Wx x + b) + 2XI
Eg
W
21
K
2X
VKp
Vo Wa + b) + 221]
min
E
||
2) VKp
W
|[F
where amin in the penultimate inequality is the minimum singular value which lower bounded by one.
B.2 PROOF OF LEMMA 3.1 FOR =0
Lemma. Let (0, Ia) be standard Gaussian random variable. For the objective function
defined Equation 2 with vRa where go(-) satisfies Assumption 2.3 any second-order
stationary point satishes W 0.
17

---

Published as a conference paper at ICLR 2026
Proof.
The first and second derivatives of f(W, .b, with respect to b given by:
@f( b.0)
E. Ve W + b)]. 0 f( [W, b.0) 0)
ab
Er V (W + 6)]
8bb
The first derivative with respect to W is:
OS( b,0)
Ow = E. Voo Wx + b)x b] 2XW.
(13)
Using Stein's Lemma in the multivariate case, this can be rewritten as:
@f(W. b,0)
E E
go W + 6) + W.
(14)
ow
At a second-order stationary point, the Hessian with respect to satisfies:
8 P(W, b.
E. V go (Wx + b)] >
(15)
b
From Equation
15
follows that:
En V go Wx x + b)] + 2XI > 0,
(16)
where the addition of 2XI ensures strict positive definiteness.
Using the first-order optimality condition 0/(W b.0) along with Equations 14 and 16 we have:
9W
Em
2
go W x + b) + 2 0.
(17)
Since En VV* b)] 2X1 (from Equation 2. the only solution W 0.
C
OPTIMIZATION ALGORITHMS
Below we give the two main algorithms for finding SOSPs; PGD and Hessian Descent.
Algorithm Perturbed Gradient Descent
Require: Objective function f(x), initial point æo, gradient Lipschitz constant L, learning rate 1) maximum iterations T
1: Initialize œ1 <- -pp
2: for 1 to T do
3: if' perturbation condition holds then
4: Draw random perturbation E;
5: Tt + 2; + &
6: end
7:
Et++1 <- x. V (x,)
8:
end for
9:
return #T-1
18

---

Published as a conference paper at ICLR 2026
Algorithm 2 Hessian Descent
Require: Gradient Vg. Hessian V g, initial point wg, step size t. perturbation step size
h 3vP Lipschitz constants L, K,
K
1: Initialize 1 4-0
2: while do
3: if' V > p then
4: L++1 2t V Vg((z))
5:
else if VV II I p and Amin (V <-y Ky then
6:
U1 c eigenvector corresponding Amin g
7:
1 + hur
8:
else
9:
return 3t
10:
end
11:
t +
12:
end while
D REFORMULATION
R( W) Exi.. [h' (2r1) .00 (Wix + Wi i + b))] + + + WIl?
-EE [Rey [e(1 (h' (mi) 90 ( Waw + Wii b))] + Al| Wi [|]] + W ||
(18)
Emy 1e/ (W b)] X W ||:
where Ey (Wix b) Exj [h' (æn) 9e (7 Tazw Wia + +b)] W I| Equation 18
holds because W1 orthogonal to WI- For notational convenience, we suppress Wi, W1 and T in the expectation and just write
R(W R[W_ E 1C (T + b)] W|
Additionally, clarify that our analysis operates under the assumption that the optimization method employed converges to p-SOSP with respect to all model parameters, including those encapsulated A. Specifically, we assume that an approximate SOSP identified jointly over the entire parameter space. Once such a point is found, we observe that fixing the auxiliary parameters preserves the approximate second-order stationarity with respect to the first-layer weight matrix W.
E PROOFS OF SECTION 4
E. PROOF THEOREM
Let the objective function in Equation 5 be twice differentiable, bounded below, and satisfies Assumption 2.3 Let the data x ~ (O, Ia). and suppose labels are generated according Equation 3 Let
Theorem.
0 be a prespecified accuracy, and define the regularization parameter A where K is the Hessian Lipschitz constant and A is an arbitrarily small constant. for T > O (poly (L, log(d), e A1))
2 Then, with probability 1 - 8, running Algorithm [ iterations with step size (1)L), yields a weight matrix W1 + Wi that satisfies:
Willr <..
Proof.
Since the objective function has Lipschitz continuous gradient and Hessian and bounded from below, this implies that it has at least one S~SOSP. Choose some A > 0 and X <Kot4 and using Lemma 3.Il we get that any p-SOSP is weight matrix W WW W that satisfies W W < S.
19

---

Published as a conference paper at ICLR 2026
L 2. solving for p, this gives p =A. Then, running Algorithm | foo:
T T= o 10
log log' (d)
L 1og' log (5)
=0
0 L
ldg
(d)
L
log' (6)
pa
E²
iterations gives the required result.
F EXPERIMENTS IN NEURAL NETWORKS OF SECTION 4.3
The student NN is a two-layer feedforward network with a single hidden layer. The input data X E R2*d is sampled from a standard multivariate Gaussian distribution with d = 2 and = 50. The network architecture consists of an input layer with features, a hidden layer of width h 1, 000, and an output layer that produces scalar predictions.
The first-layer weight matrix W R is initialized with entries drawn from 10 1/d), while the bias vector b E R and the second-layer weight vector are initialized from (O0, 1/h²). The trainable parameters include both W and b and we freeze the second layer for stability reasons. The student NN computes predictions according to the mapping:
y a RWL b),
where, ReLU2 denotes a smoothed version of the ReLU, defined in Section 4.2 Although ReLU, is used in our experiments, other nonlincarities such as tanh can be employed in the same framework.
To generate the labels, we use teacher network based on single-index model. A fixed direction is chosen as
1 (1,1) 0 == V2
and labels are generated according to the rule:
y/ ta xx) + where N(0,0.1).
Figure shows that the first-layer weights of the student network have effectively converged the principal subspace, indicating that the network has focused on the relevant direction.
G SUPPORTING MATERIAL OF SECTION 5.1
G.1 PROOF OF THEOREM 5.2
Theorem. Let G (V, E) be graph with edges, where the edge weights given by Wi. Wj.i for all (i,j) E E. Let V E R= be the matrix of vectors obtained from the SDP relaxation of the MAXCUT problem, as described in (Goemans Williamson 1995). Denote by ~ (0, ) a standard multivariate Gaussian vector, and let fe represent vector of means. Initialize Algorithm with μ O. and optimize the e-smoothed version of the objective function
Veje++ and run for
in Equation 7 (see Equation 27 using the regularization parameter À
2
(og()),c 1 1,A,))1)) .)) iterations. After this optimization process, the resulting vector 1a defines a cut whose value is guaranteed to at least:
O (poly (log(m). (og(m)))o)(())))
OPT(a O(e)),
with probability Here, a = 0.878 is the approximation factor from (Goemans & Williamson 1995).
20

---

Published as a conference paper at ICLR 2026
Proof. Given graph with m edges and weights that are equal to one, W, W; provide an approximation to the MAXCUT problem by derandomizing the randomized rounding algorithm. Consider the matrix V E Rmxm which gathers all the vectors from the Semidefinite Program. Let us also define the function Vz + μs, where 2 ~ (O,, m)) and is a vector of means. Our objective is to minimize the negative expected cut:
aim to
J(V 1)) - E
wij Pr (sg((e; z μ4) + ug p;)]
icj
for which initially we have 1 = 0. To minimize this function using our Lemma, we define the indicator function:
1, if ay < 0
(xx v)
0,
if ayy
Using this indicator, can reformulate the objective as:
((V= + (1) = = - Σ Wi [(u; z +4,4 2 (k))
+ + 12,
(19)
We are concerned with the value of [(v; 2, Uj z). If (v; 2, Uy 20) =) = 1, then this edge contributes the because the signs different, otherwise it does not. However, since this function is not smooth, we introduce a smoothed version:
L, 0, if |a| < or lyl < or > 0
xy < 0 and | 1yl >
M(, v)
[0.
otherwise.
This function smoothened to twice differentiable in the interval [0, 1]. and we can approximate it polynomial ensure that the Hessian Lipschitz constant is K O (+=) and the gradient
Lipschitz constant is L (1)).
The function [x, y) is formally defined as follows. First, we define,
1
x 2 x <
0,
S(z)
r 1
< x S
(r 2)
+1,
+1,
38 < x < 1
then, the smoothened step function i is defined as:
i(x, y) = S(x)S(-)) + SS()))
By using this smoothened function to calculate the cut, we worsen the result from original cut by at most Ommc), where m is the number of edges. This corresponds to the area of disagreement between using the actual indicator function and the smoothed indicator function. Our goal is now to minimize the following function:
f(Vz μ) = Ex
Σ K33
+ p.i v, 2 m.) + All V
(20)
where we absorb the minus sign into the original definition of w.,j . for convenience. Choose p such that at the end of the optimization have V < <
21

---

Published as a conference paper at ICLR 2026
Next, we investigate the effect of ignoring V and only using the means μ.. At the end of the optimization, we have vectors Us I,..., m. with | 11 - e², for all We examine two cases for each edge of the graph:
1f If Mifly < 0, these two values contribute to the cut directly, so we do not need anything more.
2) If pall > need further analysis. Without loss of generality, assume that μ4; 0 and μi 0.
To analyze the difference when using full randomness, we need to consider the following. For the edge to contribute the randomized version, require that either v + 14; since we want one of the have to generate z such that the magnitude l|u z|| > which has exponentially small probability, that
v, P4 <
change signs and thus contribute I. Consequently, would
(-2)-
1
Pr(l =1 >
c) S exp
Because of this we can conclude that the total expected cut remains nearly unchanged.
So, initially before smoothening, we have Σicy Egtss.j (T(D; z, fy z)] aOPT, then after using the smoothened function we obtain:
EE.us, [[(o: u. x) aOPT O(em).
Finally, after optimization, we have:
Σ Br Be [i(0= p4, uy mi) 2 aOPT O(m)),
As final step, we compare our (deterministic) cut, with what would have happened if we took the original randomized version. To complete this step, we use a union bound:
Σ Wis ( = aOPT - O(em) - Pr (|v,2) > e for any edge) i4
aOPT - -00PT-O(em) - m Pr (1viz > E)
a OPT - (em) O . exp
emxppp (-2)
2 a OPT (em) -(em)
OPT (PTT
OPT (()),
where the penultimate equality follows since OPT is function of the edges For the iteration complexity, since we want to reach point with |V < e/2 for the specific
choice of À +p for p 3/2 and running Algorithm I with step size (
L log" (m)
L
1 2 log (8)
T= O
log* (6) o 2
2 log (m)
iterations gives a p- SOSP with probability and thus VV S the above and the fact that we have reached p- SOSP giving the required result.
3/2 (via Theorem Combining
G.2 EXPERIMENTS IN MAXCUT
In this experiment, we evaluated a stochastic optimization algorithm for solving the MAXCUT problem a randomly generated undirected graph with 15 vertices and edge probability
22

---

Published as a conference paper at ICLR 2026
of 0.6. The exact MAXCUT value (computed as 41) was obtained using exhaustive search and used as a ground-truth reference. The optimization procedure based on the Goemans- Williamson relaxation, where node embeddings are derived from the top cigenvectors of the adjacency matrix. A stochastic gradient-hasel method is then applied, which samples noisy directions from a Gaussian distribution parameterized by mean vector m and log-standard deviation log and the number of cut edges is evaluated. Gradients with respect to both m and a are estimated using a Monte Carlo approximation with 100 samples per step. The parameters are updated via SGD, which sufficient for this task, with adaptive learning rates: 0.01 for m and 0.001 for σ. To ensure stability and exploration, regularization term applied to a, and its values are clipped to the range [10-3 1.5]. Learning rates are further annealed by decay factors every 100 iterations. The algorithm was run for 5000 iterations. Throughout optimization, we tracked the evolution of cut values, the maximum standard deviation across dimensions, and sampled cut edges to monitor progress relative to the exact MAXCUT benchmark.
Figure iillustrate the progression of the cut value over the course of training, showing consistent improvement and eventual convergence to the optimal cut obtained via brute force. In parallel, Figure 5 shows the evolution of the maximum value of o², which steadily decreased over time. This trend indicates that the algorithm gradually reduced its randommess as it converged toward a confident, high-quality solution. Notably, our optimization method substantially outperformed the baseline cut value of approximately 36 achieved by the classical randomized algorithm. Overall, the results demonstrate that the method not only successfully identified an optimal cut but also naturally annealed its uncertainty, confirming both its effectiveness and stability.
Cut Value Over Iterations
Maximum a Over Iterations
3.0
Out Reratrons
mal OUE V2lue
0.8
30
5000
2000
Reretions
Figure 4: Progress of the value itera- tions.
Figure 5: Evolution of the maximum 0² value over iterations.
H SUPPORTING MATERIAL OF SECTION 5.2
H. 1 REFORMULATION OF JOINSON-LINDENSTRAUSS OBJECTIVE FUNCTION
To achieve the JL guarantee from Definition 5.3 A E k The JL. Lemma guarantees the existence of random linear mapping that achieves this projection with high probability:
define linear mapping f(x)) Aar, where
Lemma H.1 (Distributional JL Lemma). For E, E (0, 1) and k = (1oo((1/)){{)) there exists probability distribution over linear functions f R2d - TR such that for every x E RRd
Pr (1f(=)) E [(1 [(1=<)]|==) JND
| 12 (1 )) 2 5.
23

---

(24) Published as a conference paper at ICLR 2026
Let A be random matrix whose elements d independently drawn from a Gaussian distribution with mean Ma, and variance o. Define the distortion function as:
AAA; æ4) || 1/
(21)
where (M, E). Our objective is to minimize the following function:
[(A Σ Pr(h(A;;,) > |E
(22)
2kd
where 11/2 denotes matrix whose entries the square roots of the corresponding variances in E.
The first term measures the number of distortion violations (i.e., how often the projected norm deviates from more than E), while the second term is regularization penalty on the variance of the matrix entries.
Applying a union bound, the objective in Equation 22 serves an upper bound for:
|
Pr max a(A;z,)
(23)
2kd
which represents the probability that the maximum distortion across all data points exceeds E, plus regularization term. As the variances Σ approach zero, this regularizer vanishes. the context of the JL Lemma, our goal is to minimize the probability expressed in Equation 23
To use Lemma
we can think of the matrix A as kd-diiensional vector and write:
vec 1/ = u4,
where μ = (M1,1, 12.1.... a lkdd is a vectorized version of each mean of matrix A, Σ diag o i.2.... k.d Bkdxkd is the diagonal covariance matrix and z is kdd- dimensional multivariate Gaussian vector with independent entries with zero mean and unit variance.
Then, define,
(o(A vec) g0 (2 (1
(25)
Then, we can define the objective function:
1//
f {x>) .1)
E
go
Σ 1/2
2 [1]
2kd
||>
E
+
2ke
|
Σ Pr (h +
(26)
2kd
Remark H.2. It is important to observe that Equations 22 and 26 represent the same quantity, but expressed using different parameterizations.
Thus, minimizing Equation 22 equivalent to minimizing Equation 26 where optimization is carried out over the parameters
..).
H.2 PROOF OF JOHNSON LINDENSTRAUSS GUARANTEE PRESERVATION LEMMA
Here we give extension of Lemma 4 from Tsikouras et al. 2024 which is required due to the practical limitation that achieving an exact SOSP is not feasible. Since Algorithm identifies an approximate SOSOSP an additional result is required to provide stopping criterion once the variance becomes sufficiently small. This ensures that the mean can used with a controlled deterioration of the JI. guarantee.
24

---

Published as a conference paper at ICLR 2026
Lemma H.3. Given unit vectors in IR and target dimension k. choose such that the random matrix A (M,, Σ) satisfies the JL guarantee with distortion with probability at least 1/6. Then using matrix M instead of sampling from A retains the JL guarantee with threshold increased by at most poly(Tmax. /k).
Proof.
We start with the assumption that | (1 - 8) with probability at least -
Expressing A as A M + 2 where ~ N0 E). For this, we have
Z
Z
Za 12
where and are the same as but scaled with the maximum and minimum variance from E
respectively. This way, all the entries of Z have the same common variance, of and all the entries of Z have the same common variance, mina
From the Lemma, we can select Eu such that
Zee [12 E [o o max (1 Eo)]
Zæ l12
io (1 Eo), oin (1 =o]]
with probability at least 7. This ensures there exists an overlap where both inequalities for A, Z and hold simultaneously. Our goal is determine how much excess distortion we get when using M
instead of sampling from the random matrix A.
Using the triangle inequality we have:
Mæ - M + Zæ Zzel22 S Ma + Z + Ze 2|
1
S &
A
+ k.
which by squaring both sides and using the JL guarantee for A and Z, we obtain:
2
1 k
Mæ
Ax
Ax la
Zæ
12
Zæ 1
A2
20max
< 1 +
V1 + 1 + o o (1 + en)
2020
20 max
<1+et
VI+ 20 -
For the lower bound, using the Cauchy-Schwarz inequality and the JL guarantee for A and Z, we have:
Mx 2
|Mx
Zx
|Z
2k
k
1
||
|
27:
1/2(1 - 8) -ii (1 E0)
> 1/2(1 E) Imin
> 1/2(1 - -8) max
Finally, combining these results, we observe that replacing with M maintains the JL guarantee with an increased distortion threshold, bounded by at most poly(Tmaxx, 1/k), with high probability. n
25

---

Published as a conference paper at ICLR 2026
H.3 SMOOTHENING OF THE JOHNSON- LINDESNN-I INDICATOR FUNCTION
For the JL objective function we have the indicator function:
if Ar; - 2
1 [ A)
0 if Ax; 2
and we define a smoothed version of
0,
if A 12 1| S
Az 1| -1 -e)3 (E +5 - -|1AA,112 -1)9.
if c < ||| 12 - S 5
M(m;; A)
if + a < A1, - <
if Ar. 2 2 & +
for a small value 61. This smoothed indicator has gradient Lipschitz constant L (11=E)) and Hessian Lipschitz constant K (11/E)) We define the Em-smoothed version of the objective function in Equation 26 that is:
= [(A) B[7 A)].
and the regularized version of it, that is:
f {x fA) B[T(,;; A)] + |
(27)
2kd
By assumption we have that 1/(3n) > E[T( A)] > ( A)].
We also have that
E[[(;AA) = A) + Pr(e | - -1| < <
Thus,
2 Σ [[]();A)) S Σ R A)] + Σ
Pr(s < ||| - < +11)).
We will show that when A has small variance (for appropriately chosen small p). the B[i(z;; A) becomes smaller than 8/nn.
We have that [[] A)] < Prr]||rr11 We assume that we have reached point for which we have A ~ M(M, Σ). This means that |||,,;1|2 - 1| follows a non-central chi-squared distribution. From subgaussian properties we have: For t > 0:
- e1) Pr(s < || 1 - < < 21).
12
Pr(X - 2 exp
2
For /. < 0:
f2
Pr(X - 4 5t) < exp
20
We have Var A - 1)) < 2 o4 max + 2 Σ
13 = Vi, where = Ed Mij
Choose + E and if > E[] Ax; 2 1|]
we have that:
26

---

Published as a conference paper at ICLR 2026
(1
A A
- 1 1))
Pr[[|Az,||2 12 1| BBIIAAz 1|1 2 t - ||, - 11) - exp
2Vr (1ll A |= 1
Ei|l| 112 -1/12
I exp
2V,
Otherwise if < Bi| |A:;l 1[] we have that:
Pr[[||A 1| - E Ar. 1|1
S1
51-11|Ax:1* ||| -1()
(t E3 |||| Ax; 111
exxp
(28)
2V;
This implies that
(t - Bi]
|Ax 1|2 111)
Pr[I|Aax:|2 11 B[ -1l] 1 || 11) -exp
2V;
2V,
2 1 (t B[| Ar.
11])
Since E[[(;;A)) < 1/(3n) we have that the 2nd case where We also have that for V3 I (e+e) |la the probability Equation 28 is bounded by 61/ /2n()..
< E| A | || is rejected.
log(6./(2m))
Similarly we have,
Pr(s < || -1 < <2+ Ar: < <1+ Pr[[||A;;1| 11 <))
=Pr( 2 1| 8) Pr[[]|Ax;| +1).
This implies that
Pr(s < ||x 2 1/ < <E+ s Pr[[|AZ,11 2- 11 > E).
Similarly before can get that:
|Ar, 2 1- 11) E
Arr - > E) < exp
(29)
2V:
(8 (8-B|]) Ax,11=-111 l2g())/(2n)
Therefore we have that for Vi2 < 61/(2n)
the probability in Equation 29 is bounded by
(e [e=, (eter E 1
This means that choosing,
min
l2g()
log()1((2m)
(s- (s- 1 and overall, if we choose V max (Vi. Vn to satisfy all inequalities
(2n)) we get:
Σ E[[ A) < Σ []i(2;A)] + Σ Pr(e < || - < < 61.
i=1
27

---

Published as a conference paper at ICLR 2026
Denote M; = Ei= I Pia and C1 min I...
[e-B]]
|
Then to get V C1.
(2n(
need o - min 2M.+VMM2+RCC C2-
Choose 81 < 5/6 and p < AvC2. Thus reaching p-SOSP for (x,; A)] has returned a random
matrix A that satisfies the JL guarantee with probability at least 1/6. We will use this in the proof of Theorem in the next section.
H.4 PROOF OF THEOREM 5.5
Theorem. Let be unit vectors in Rdd & be the target dimension, c be a smoothening parameter and A > 0 be an accuracy parameter. For any 2 C Co/log n/k, where C is sufficiently large constant, initialize M and
Ix and run Algorithm to optimize the e e--smothhed version of the objective function in Equation 9 (see Equation 27) using the regularization parameter À = + After 7 O (poly (n, ki d, log( A-1)) iterations, this returns matrix M that satisfies the JL guarantee with distortion at most (E), with probability 1 6.
2
To ensure good performance, we choose the dimension such that the probability of any individual distortion constraint being violated is more than 1/(3n). This choice is critical, particularly at the initialization point (Jadx 0). where the regularization term contributes exactly 1/2. Under this setting, the objective function in Equation 26 satisfies:
Proof.
71
1
f(Indx kd. 0)
2 =
3n
3
2
6
This observation implies that, if we follow a monotonically decreasing path of the objective and converge to a deterministic solution, specifically one where 1 = O, the only feasible outcome is that each term in the summation becomes zero. Consequently, the distortion probability in Equation will converge to zero.
However, to our key Lemma need use smoothed version of the indicator function. The smoothed objective function in Equation 27 both gradient and Hessian Lipschitz continuous. Let L L= , O (1 be the gradient and Hessian Lipschitz constants, respectively.
Choose 81
3/6, A O, and X vKp+A
Vpte3+A and using Lemma 3.1
get that any
p-SOSP gives a matrix 1/2 that satisfies
/2
S &.
(c-E1 | Ar, 1-1
Denote M; = Ei= Mist and C1 = min
log(d./(2n) S
Then get S C1, we
2M,+ / /M?
need 2 max S min;
Cz.
Choose p < D2 then running Algorithm I Twith step size Oss2) for:
T T=0(1 log p²
I, log (8) = O (poly (n, k d, log(o), A-)).
returns a random matrix with 2 < o² that satisfies the JL. guarantee with distortion & with probability at least 1/6, with high probability. Then, from Lemma H.3 get that use the mean matrix M which will increase the distortion threshold by most poly (Amax: 1) poly (g. 1), meaning that it satisfies the JL guarantee with distortion at most O(e).
H.5 EXPERIMENTS IN JOINSON-LINDENSTRAANE
In this experiment, aim to minimize the distortion introduced by random linear projections in the JL framework. A batch-based variational model is trained using the gradient- based method; SGD, which is sufficient for this task, to produce random matrices A AM, Σ) R (with
28

---

Published as a conference paper at ICLR 2026
k 30 and d = 500) that minimize the maximum distortion when applied to normalized dataset of n = 100 samples, each with = 500 dimensions. Unlike traditional JL embeddings that rely on random Gaussian matrices, approach optimizes the parameters (M, 2) of distribution over projection matrices using the Adam optimizer (Kingma & Ba 2014) in order to minimize the worst-case distortion. We used batch size of 20, learning rate of 0.01,
maximum of 5000 iterations, and early stopping is triggered if the distortion falls below 0.01. To track how the distortions evolve with our method, we sample from the current mean matrix and variance each iteration and then calculate the resulting distortion.
Figure 6 shows the evolution of the maximum distortion throughout training, demonstrating steady decrease. Over time, our method significantly outperforms both the average and minimum distortions obtained from standard Gaussian matrices over 1000 trials. Specifically, our learned projection achieves near-7ero distortion, compared to typical random projections that yield average and minimum distortions around and 0.6, respectively. Figure 7 illlustrates the evolution of the maximum variance o², which converges toward zero during training. This indicates that the model refining its uncertainty and collapsing toward a deterministic, low-distortion projection. These findings suggest that structured embeddings with far lower distortion than those from conventional random constructions do exist, and that such embeddings can be effectively discovered via gradient- based optimization.
Distortion Over Iterations
Variance d² Over Iterations
Ogbnized Distortion
Pver Detortoo
2.2
0.8
1.0
vallucce
0.8
0.6
0.1
0.0
1000
2000
3000
4000
5000
1000
2000
2000
4000
5000
Esrations
Enrations
Figure 6: Evolution of the optimized distortion over iterations.
Figure 7: Evolution of maximum 0² over itera- tions.
29


Reviews:
[{"review_id": "29fFTyrCz3", "reviewer": "Reviewer_cEJx", "summary": "Consider the problem of optimizing any function of the form $\\mathbb{E}[g(Wx + b)] + \\lambda \\\\|W\\\\|^2_F$, where $x$ is drawn from a standard Gaussian, with respect to parameters $W, b$. The main result of this paper states that any second-order stationary point (SOSP) of this objective must have $W = 0$ (and an approximate version holds for approximate SOSPs). This has several applications:\n1. Neural network optimization: when we try to optimize the squared loss between a \"teacher\" network of the form $h(Ux)$ and a \"student\" neural network with initial layer weights $W$, then at any SOSP $W$ will have its component orthogonal to $U$ approximately vanish. This could be seen as optimization \"discovering\" the (typically low-rank) subspace of $U$.\n2. Derandomization: this lemma can be cleverly applied to certain randomized algorithms (such as MAX-CUT or Johnson-Lindenstrauss) in obtain a deterministic solution that (nearly) matches the performance of the randomized one.\nThe proof of the main lemma is simple and relies crucially on Stein's lemma for Gaussian random variables.", "questions": "1. I am curious if the results in this paper extend even in an approximate way to distributions that are close to but not perfectly Gaussian (e.g. sub-Gaussian, log-concave, etc). Would be curious to hear the authors' thoughts on this and more generally the restrictiveness of the Gaussianity assumption.\n2. In general throughout the paper, the authors need to be clearer about what the variables being optimized are in any given objective function. Even the term SOSP is always really \"SOSP of certain optimization variables\". This is sometimes but not always clear from context. In particular, in section 4, I was quite confused by the step in going from Eq 5 to Eq 6, because $W_{\\parallel}$ suddenly disappears from the objective altogether. When we actually optimize the risk, we optimize all of $W$ --- i.e., really we are considering an SOSP of both $W_{\\bot}$ and $W_{\\parallel}$. It is also extra confusing because $\\ell'$ implicitly depends on $W_{\\parallel}$. I had to spend some time convincing myself that we can effectively work as if $W_{\\parallel}$ were a constant equal to its SOSP value. Perhaps such manipulations are obvious to an optimization audience but I think they are somewhat subtle and deserve very clear, explicit exposition (perhaps in the appendix, if space does not permit).\n3. The interpretation of the main lemma as a \"derandomization\" lemma deserves clearer explanation. It only really comes up in Section 5 and even there it is a bit conceptually subtle. In fact this is a good opportunity in the paper to state purely in words what the main lemma is really saying.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The main result of the paper is appealingly clean, general, and powerful. The applications are creative and interesting. In the case of neural network optimization, the paper recovers a \"structure discovery\" result which had been proved much more painstakingly (and under more restrictive assumptions) in prior work. The derandomization applications are conceptually thought-provoking and potentially of independent interest in complexity theory and randomized algorithms. The paper is mostly written quite clearly. I did not verify all the proofs in detail but the main ideas seem technically sound (and in particular the main lemma proof looks correct). Overall, this looks like a good paper.", "weaknesses": "A weakness of the paper is that it is stylized in certain important ways. Probably the biggest concern is that all the results rely on perfect Gaussianity of the random variables, because of the crucial use of Stein's lemma in the main result. Another concern is that the functions involved be smooth, which necessitates the use of smooth activation functions etc. Finally, all the claims require regularization and only hold for $\\rho$-SOSPs with $\\rho$ very small compared to the regularization strength $\\lambda$. This limits the interpretation of some of the more practical applications, esp concerning neural network optimization. Indeed, in that application it is not clear that this meaningfully explains \"structure discovery\" in any way that illuminates what might be happening in a real world setting; rather this seems like more of a curiosity about the Gaussian. The algorithmic applications only really seem suitable for cases where the algorithm designer synthetically injects Gaussian randomness into an algorithm (as with randomized rounding and random projections), not cases with \"real world\" randomness. This is probably fine, but it bears noting.", "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": "TF5Ljfd1rp", "reviewer": "Reviewer_n2cY", "summary": "This manuscript studies representation learning in neural networks in a multi-index setting. In particular, the authors show that, using the definition of approximate second-order stationarity, the result in (Mousavi-Hosseini et al., 2023) can be established under more general conditions. They provide algorithms that find approximately second-order stationary points under smoothness assumptions and also suggest alternative uses of these algorithms for MAX-cut and embedding finding.", "questions": "- In Eq. (3), $U$ is not bold; I believe this is a typo.\n\n- Following my points above, in line 260 the authors claim:  \n  “Our goal is to show that the perpendicular component $W_{\\perp}$ converges to zero, implying that the first-layer weight matrix $W$ converges to a rank-$k$ matrix,”  \n  which is incorrect. This statement should be revised.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper extends earlier result (Mousavi-Hosseini et al., 2023) to a broader setting.", "weaknesses": "- The paper does not guarantee learning the teacher directions; rather, it shows that the component of the student weights in the subspace orthogonal to the teacher directions vanishes. However, this does not guarantee recovery of the teacher directions. For example, consider the setting where the teacher is $y = \\mathrm{He}_4(\\langle \\theta, x\\rangle) + \\epsilon$ for some unit vector $\\theta$, and the student is $\\hat y = \\mathrm{He}_2(\\langle w, x\\rangle)$, $w \\in \\mathbb{R}^d$, with $x \\sim N(0, I_d)$ and square loss. Then $0$ will be an approximate stationary point for any $\\rho > 0$, even though it is a saddle point. In this sense, the result does not guarantee learning.\n\n- The example above relies on the orthogonality of the Hermite polynomials $\\mathrm{He}_4$ and $\\mathrm{He}_2$. Hermite decompositions in multi-index settings have been studied extensively since (Mousavi-Hosseini et al., 2023), and there is now a rich literature that discusses the sample complexity of recovering teacher model with first-order algorithms based on its Hermite decomposition in Gaussian space [1,2,3]. The manuscript does not compare its results to this literature. What improvement does this paper provide? The manuscript would benefit from this discussion.\n\n\n[1] Gérard Ben Arous, Reza Gheissari and Aukosh Jagannath. “Online stochastic gradient descent on non-convex losses from high-dimensional inference.” J. Mach. Learn. Res. 22 (2020): 106:1-106:51.\n\n[2] Alex Damian, Loucas Pillaud-Vivien, Jason D. Lee and Joan Bruna. “Computational-Statistical Gaps in Gaussian Single-Index Models.” ArXiv abs/2403.05529 (2024): n. pag.\n\n[3] Jason D. Lee, Kazusato Oko, Taiji Suzuki and Denny Wu. “Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit.” ArXiv abs/2406.01581 (2024): n. pag.", "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": "2h6lWfwrP4", "reviewer": "Reviewer_uty6", "summary": "This paper develops a general derandomization framework for structure discovery in neural networks and other domains. Building on the prior work of (Mousavi-Hosseini et al., 2023), which showed that SGD can lead to low-rank structure in the first layer of a two-layer neural network under strong regularization, this paper extend the result to a broader setting with trainable bias parameters, general smooth loss functions and optimizers. The core contribution is a key derandomization lemma stating that optimizing a general regularized loss in expectation form converges to the origin point, under mild conditions. The authors also apply the same lemma to derive derandomized algorithms for the MAXCUT problem and Johnson–Lindenstrauss embeddings.", "questions": "See strengths and weaknesses section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The authors propose a new derandomization lemma for analyzing feature learning behaviour of neural networks. The obtained results extend and generalize the analysis in (Mousavi-Hosseini et al., 2023).  \n2. The authors applied the result for derandomization in other domains including MAXCUT and JL embeddings.", "weaknesses": "1. The assumption of first and second order smoothness is restrictive and does apply to many practical scenarios, such as ReLU networks, forcing the authors to adopt the approximation in Section 4.2. Can this result be extended to non-smooth settings such as the original ReLU activation?  \n2. The applications to MAXCUT and JL embedding are actually not new since, as acknowledged by the authors, there are already known derandomized algorithms for both problems. Can this result be applied to new problems where derandomized algorithms are unknown?\n\nTypos and other comments: \n\n1. in line 903, the definition of $l’_{\\\\theta’}$ below equation (18), $\\\\lambda\\\\|W{\\\\bot}\\\\|\\_{F}^2$ should be $\\\\lambda\\\\|W\\_{\\\\|}\\\\|\\_{F}^2$.  \n2. The “Summary of our contributions.” section seems to be missing.", "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": "1NlFR9WkYT", "reviewer": "Reviewer_F4fe", "summary": "In this paper, the authors prove a de-randomization lemma which reveals information about the loss landscape of problems that fit into a certain structural framework. The lemma essentially says that when trying to optimize a smooth function over a matrix W, such that W is multiplied by a standard normal vector and passed as an argument to the smooth function, than any candidate solution W* which is approximately second order stationary (which includes all local minima and potentially some saddle points) is close to 0, with the proximity depending on how close the point is to being second order stationary. \n\nThis lemma is then demonstrated on 3 diverse problems of interest including:\n1. proving a statement about training dynamics for neural networks discovering low dimensional structure.\n2. a method for derandomizing the rounding procedure which is applied to the Semi definite program solution of a continuous version of the max cut problem\n3. a method of finding deterministic JL embedding matrices.\n\nThe lemma provides an interesting insight into the loss landscape for a certain class of functions. The examples demonstrate an interesting and diverse range of applications which have a common structure.\n\nWhile interesting, I think demonstrating the importance of the main lemma hinges on the examples. The first example which received the majority of the attention in the paper is solid and interesting but feels somewhat marginal in terms of extending  the result of Mousavi-Hosseini et al., 2023. I do think the other two examples are very interesting but I do not totally understand the impact of the lemma in regards to the problems it is applied to. If this was highlighted more clearly and I can be convinced of the significance this lemma brings to the MaxCut and JL applications, I would think the paper is much stronger.", "questions": "For JL how exactly does your method differ from what others have already done? Or is the contribution of this example to demonstrate that your lemma can be applied and hence provide insight into the loss landscape for this problem?\n\nA small note: I believe in appendix C (line 903) the definition of $\\ell'$ should include the $\\| W_\\|\\|^2_F$ not $\\| W_\\perp\\|^2_F$ right?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The various applications are interesting as they consider very different but high profile problems which are of interest to many different communities. \n\nThe general insight into the loss landsape is also interesting and I agree with the authors that others may find new applications of this lemma to different problems.", "weaknesses": "It seems as though for the L-smooth function who's expectation we want to optimize, the sample complexity for finding a solution with perturbed gradient descent is growing in L. When applying the lemma to optimize probabilities (written as expected indicators) the smooth approximation to the indicators requires larger L for better approximation. In these cases there is an obvious tradeoff between sample complexity and accuracy. I think this is clear for the max-cut theorem but is sort of hidden in the JL example and perhaps this should be made more clear. \n\nTo follow up on the point above, given the application of the lemma to these two problems, it would be interesting to have some insight into how quality of approximation of the indicator affects quality of the final solution.\n\nIt would be nice to have a bit more context into the final two applications in terms how effective the proposed de-randomization problems are. You suspect that you have shared the first optimization-based approach for\nderandomizing the Goemans-Williamson algorithm, but have not commented on how this approach compares to the other mentioned methods of  of conditional expectations, small-bias spaces, or explicit pseudorandom constructions. For JL you mentioned that you match the SOTA results but also later that actually recover another method and there is no real comparison between your or other methods (See the first question in the 'Questions' section).\n\nI think the connections of the structure used in your lemma and these different applications are interesting. It is not totally clear to me however what exactly applying your lemma to these problems is contributing to the study of these problems. Is it new methods for solving these problems? Or contributing to the theoretical understandings of these problems or something else? If it is new methods I think a better comparison is warranted.", "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)