Precise asymptotics of reweighted least-squares algorithms for linear diagonal networks
Authors: Chiraag Kaushik
Justin Romberg
Vidya Muthukumar
Venue: NeurIPS 2024 poster (2024)
Keywords: high-dimensional asymptotics
linear diagonal neural networks
feature learning
iteratively reweighted least-squares
sparse recovery
Primary area: learning_theory
TL;DR: We provide a precise asymptotic characterization of the iterates for a family of reweighted least-squares algorithms that learn linear diagonal neural networks.

Abstract: The classical iteratively reweighted least-squares (IRLS) algorithm aims to recover an unknown signal from linear measurements by performing a sequence of weighted least squares problems, where the weights are recursively updated at each step. Varieties of this algorithm have been shown to achieve favorable empirical performance and theoretical guarantees for sparse recovery and $\ell_p$-norm minimization. Recently, some preliminary connections have also been made between IRLS and certain types of non-convex linear neural network architectures that are observed to exploit low-dimensional structure in high-dimensional linear models. In this work, we provide a unified asymptotic analysis for a family of algorithms that encompasses IRLS, the recently proposed lin-RFM algorithm (which was motivated by feature learning in neural networks), and the alternating minimization algorithm on linear diagonal neural networks. Our analysis operates in a "batched" setting with i.i.d. Gaussian covariates and shows that, with appropriately chosen reweighting policy, the algorithm can achieve favorable performance in only a handful of iterations. We also extend our results to the case of group-sparse recovery and show that leveraging this structure in the reweighting scheme provably improves test error compared to coordinate-wise reweighting.

Precise asymptoties of reweighted least-squares algorithms for linear diagonal networks
Chiraag Kaushik
Justin Romberg
Electrical and Computer Engineering Georgia Institute of Technology
Electrical and Computer Engineering
Georgia Institute of Technology
Atlanta, GA 30308
Atlanta, GA 30308 jromêece waw.uceegtttchhhcc. edu
ckaushik7@gatech.. edu
Vidya Muthukumar
Electrical and Computer Engineering,
Industrial & Systems Engineering
Atlanta, GA 30308
vmuthukkmar@@gattech.. edu
Abstract
The classical iteratively reweighted least-squares (IRLS) algorithm aims cover an unknown signal from linear measurements by performing a sequence weighted least squares problems, where the weights recursively updated at each step. Varieties of this algorithm have been shown to achieve favorable empirical performance and theoretical guarantees for sparse recovery and É en-norm mini- mization. Recently, some preliminary connections have also been made between IRLS and certain types of non-convex linear neural network architectures that are observed to exploit low- www-dimensional structure in high- nigh-dimensional linear models. In this work, we provide unified asymptotic analysis for family of algorithms that encompasses IRLS, the recently proposed lin -nnRFMM algorithm (which was motivated by feature learning in neural networks), and the alternating minimization algorithm on linear diagonal neural networks. Our analysis operates in a "batched" setting with i.i.d. Gaussian covariates and shows that, with appropriately chosen reweighting policy, the algorithm can achieve favorable performance in only a handful iterations. We also extend our results to the case group-sparse re- covery and show that leveraging this structure in the reweighting scheme provably improves test error compared to coordinate-wist reweighting.
Introduction
Many high- -igh-dimensional machine learning and signal processing tasks rely solving optimization problems with regularizers that explicitly enforce certain structure on the learned parameters. The traditional formulation for such tasks involves regularized empirical risk minimization (ERM) problem of the form
min L( + AR(Ø). BRRd
(1)
where I(-) a loss function that encourages fidelity to the observed training data and R(-) encodes desirable structural properties. In many important applications, it is desirable to obtain a sparsity- seeking solution; in such cases, the regularizer typically non-smooth, as in the LASSO, group LASSO, and nuclear norm regularizers. As an alternative approach to this non-smooth optimization, several recent works have proposed the "Hadamard over-parameteriizaion"" of into the entry- wise product of two factors
While the resulting minimization problem is non-convex. this
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

parameterization, coupled with a smooth regularizer, has been shown to achieve competitive empirical performance (in terms of numerical stability, robustness, and convergence rate) when compared to For example, rather than solving the convex, but non- smooth LASSO (where I. is the squared loss and R is the & norm), the Hadamard reparameterization yields the following non-convex and smooth formulation:
traditional sparse recovery algorithms 13 25
min U
T(u ④ v) + 2
(1|u|12 |m(l3))
In the case where the regression function is linear in 0, solving (2) is equivalent to learning function of the form
(x))(x)) = (a, u C v) = (diag(v))x, u)).
which can be thought of as one hidden layer neural network with linear activation function and inner weight matrix (e) this context, this linear diagonal neural network (LDNN) architecture has also been studied as illustrative case study improve our understanding of how neural networks perform iterative "feature learning" to leverage low-dimensional structure in high-dimensional settings [4 23). We note here that, in the linear model case, feature learning is equivalent to learning which subset of the input's 's coordinates are relevant for the true predictor (i.e., feature selection).
One way to understand the connection between classical sparse recovery algorithms and the Hadamard product/LDDN form in is to consider the change of variable + /T/ and u, - v yields the following optimization problem, which is jointly convex in r and O:
0_ [25] This
A
02
min min I(O) BER yERd
1/6
2
77
After solving the minimum over 77 explicitly, the second term becomes exactly X|[0||1- and recover the Lasso objective. This is a special case of the so-called sttattrckk' []] which can be used to write many common pparsity-inducing penalties as the minimization of quadratic functional of 6.
variety of algorithms for learning Hadamard product parameterizations have recently been studied, including alternating minimization [13], bi-level optimization [25], and joint gradient descent on
(u, v) [34]. The connection to the (0, m)) optimization
3 can also be leveraged to construct algorithms based on classical sparse recovery techniques. In particular, alternating minimization over 0 and r in (3) yields the popular iteratively reweighted least-xquares (IRLS) algorithm I1 9. Translating these updates to the equivalent updates on (u, v). we obtain iterative least-squares algorithm for LDNNs, which alternately sets (t+2)) / (()) C vi!) and performs a weighted least squares update on u. This particular form of reparameterized IRLS was generalized in (28) to
larger family of iterative squares algorithms under the name of linear recursive feature machines (lin-RFM).
While several methods for learning Hadamard/LDDN parameterizations have been introduced in the literature, there remain many open questions about how they each perform and how they compare. Theoretical analyses of these algorithms typically assume fixed, possibly worst-case training data, and aim to characterize the properties of the fixed-points 28 34] or give convergence guarantees to second-order stationary points (25). However, these worst-case analyses do not readily yield guarantees on the estimation error, which is the principal metric of interest. Indeed, many works have shown that studying the average-case, or typical, behavior of non-convex optimization algorithms can allow for estimation guarantees that are more precise and reflective of practice 14 [8 7
In this paper, we provide precise analysis of general family of iterative algorithms for learning LDNNs that take the form
(1+1) arg min
(t)) 1 x (u
(uo
00))
+
|
ueR n
2
(++1) = wo( (++) v (1))
(²)) As we show in Section 2 this formulation encompasses multiple existing algorithms, including reparameterized IRLS, RFM. and alternating minimization over and v. We consider the common scenario where training is performed with batches of data and characterize the exact distribution of the parameters after each iteration in the high-dimensional limit (n, d) + oc. This allows us to address questions such as
for some reweighting function &b and batches of training data (x()),((())
2

---

How do different algorithm choices compare (in terms of convergence and signal recovery) in the high- nigh-dimensional regime?
How many iterations does it take common algorithms to find statistically favorable solutions?
What is the effeet of model architecture in LDNNs? Does leveraging group structure provably improve sample complexity when the ground- truth signal is group-sppase?
Contributions: We define general class of algorithms which learns LDNNS by alternately performing least-squares and reweighting steps in a sample-splitbatched setting, and we show the following.
(1) Under mild assumptions on the target signal, initialization, and reweighting function, we provide an exact characterization of the distribution of the entries of the parameters at each iteration in the limit as n, d approach infinity (Theorem
(2) We show that this asymptotic result aligns well with numerical simulations and allows for accurate prediction of the test error at each iteration. This enables rigorous comparison between different algorithms and demonstrates that, with appropriate reweighting schemes, statistically favorable solution can be obtained in only a handful of iterations.
(3) Lastly, we extend our asymptotic framework to a setting of structured sparsity, where O" has group- sparse structure (Theorem 2). Our results show that using grouped Hadamard parameterization (i.e., tying together groups of weights in the LDNN) effectively learns such signals, with performance scaling with the number of non-zero groups, rather than the total sparsity level.
1.1 Related work
IRLS and the n-trick: The reformulation of non-smooth regularizers in terms of quadratic variational forms (the "x)-trick")) has been studied in various early works in computer vision and robust statistics Further analysis and examples of sparsity-promoting norms are provided in [20] 2). and provides characterization of when regularizer admits variational form of this type. The
25
least-squares (IRLS), a popular technique for compressive sensing and sparse recovery [IT 9. These works also consider IRLS algorithms corresponding to f--norm regularization for 0 < < 1; in this case, the minimization is no longer convex, but [TT] shows that such methods can find sparse solutions with fast local convergence rate. The family of algorithms we consider includes reparameterized version of each of these IRLS algorithms, but unlike these prior works, we consider batched setting and the sigh-dimessional asymptotic regime. Moreover, our results apply to other algorithms which may not be easily expressed as resulting from the y-trickk.
resultant optimization algorithm is derativell-rrweighhe
Hadamard parameterization and linear diagonal networks: The reparameterization of into the product of factors u @ v has been considered in a variety of recent works. The authors of B1 (35] show that carly-stopped joint gradient descent over the two factors can lead to optimal sample complexity for sparse linear regression. The equivalence of this parameterization to LDNNs has also led to a surge of interest in the implicit bias of gradient descent/flow on this parameterization. i.e., a characterization of which solution gradient descent will reach without explicit regularization (corresponding to À = 0). These works typically consider gradient flow run until completion and characterize the solution as minimizer of a certain pparsity-induchin functional that depends on the
initialization 34 10 23
The connection between the LASSO (as well some non-convex by penalties) and the Hadamard parameterization was studied in [13] where alternating minimization over the two factors is used instead of first-order methods. More recently, [25] extends these observations by making explicit the connection to the 17-trick and showing that saddle points are strict (escapable). These insights lead to global convergence guarantees and a smooth bi-level optimization scheme 255 26
for non-smooth
structured optimization problems that was shown to perform competitively with state-of -f-the--rtt solvers. The non-convex landscape of such formulations is further explored in [15), where it is shown that for large class of parameterizations (including grouped, deep, and fractional Hadamard products), the non-convex problem has no spurious local minima. Motivated by the type of feature learning observed in neural networks, the authors of (28) propose lin-RFM, which updates one of the parameters via weighted least-squares while iteratively updating the other parameter via a reweighting scheme based on the average gradient outer product of the learned function. The authors characterize
3

---

properties of the fixed- sred-points and show that, for certain reweighting schemes, lin-RFM is equivalent to a reparameterization of IRLS. The family of algorithms we consider is similar, consisting of a weighted least-squares step and a reweighting step; however, it is more general and doesn't 't require the reweighting function to have the particular form required by lin-RFM. Moreover, our asymptotic characterization of the iterates allows for a precise understanding of how the test error evolves. On the other hand, our analysis relies on batching/samppl-splittii of training data while all of the above works reuse the entire batch of training data at each iteration.
We make particular note here of the few works which explicitly consider "grouped" Hadamard parameterization, which we consider in Section[4] This corresponds to LDNN with groups of tied weights in the hidden layer. Early stopped gradient flow/descent for this type of architecture was shown in [77 to achieve sample-complexigy scaling with the number of non-zero groups (rather than the overall sparsity). The non-convex landscape for this grouped architecture is studied (36) and
15 Our results complement these works by studying group-reweighted least-squares algorithms
(rather than gradient methods) for learning functions of this form.
Precise characterization of higher-order non-convex optimization problems: On technical level, our work provides precise deterministic characterization of a family of higher-order optimization algorithms. In this sense, our results are of a similar flavor to (7]. where Gaussian comparison inequalities are used obtain a precise characterization of non- non-convex optimization problems. However, since the Hadamard parameterization is re- paraaateriization of the actual estimator of interest (0 u @ v), the results of 77 are not directly applicable. While our results are asymptotic and do not provide finite-sample guarantees, we provide a distributional characterization of v after each reweighting step, which allows us to characterize the behavior of more complicated functions of the iterates. Precise characterizations of alternating minimization and lin-prox methods for rank-1 1 matrix sensing are studied in the works [6 [9]. While these works obtain non-asymptotic guarantees, the estimation model and resulting optimization objective are quite different, with each unknown parameter interacting with independent sensing vectors (rather than single sensing vector interacting with the product of the two parameters).
2 Background and formulation
Notation:
The ones vector of dimension d is denoted as 1g. We denote the element-wiss multiplica- tion (Hadamard product) of two vectors x and y as x @ y. Element wise division of two vectors is denoted as æ. We say a function f: IRP + R is aseudo-Lipschite of order 2 if, for all æ, y E R2, y.
(f(z) f(y) C(1 + ||0||| + (1)l|2) |a: yl12
for some constant C C>0 0. The set of such functions is denoted by (l.
Convergence in probability of a sequence of random variables Xa to a random variable X is denoted by Xa E X. Convergence in Wasserstein-2 distance of a sequence of probability distributions Va to limiting distribution v is denoted as 1/
3 v. and this fact is equivalent to the statement If the Va are random probability measures, we
Ex-es g(X) - Ex g(X) for all g E (2)
3
say that va W/ v if the same convergence holds in probability, i.c., Ex~ra g(X) P3 Exar g(X) for all g E PL((). The empirical distribution of a vector z E Rdd is defined as 1= 8(z)), where 8(2;) is the Dirac delta distribution centered at 2;.
Formulation: We consider a batched noisy linear model where, at each time -= 0, 1, ...,, T, user has access to an independent batch of data (x((t) ((t)) E Raxd RR satisfying
(0)
( 0- + (t))
Vd
Above, 6* E Rd is an unknown signal, (+) has i.i.d. standard Gaussian entries, and e(t) ~ NN (I,)) is i.i.d. noise in the measurements. Given an initial weight vector w()) E Rdd, we are interested in the behavior of iterative algorithms of the form
2 + d
w (++1) arg min
1 Vd
X
r (1))
uER
(0+1)) = vb(u (1+1) (()),

---

Table 1: Some algorithms taking the form
Algorithm
Reweighting function
Alternating minimization (AM) 13
(u,, v) u (bu, = v2 + e)a u(u, v) = v2
Reparameterized IRLS
9 1 28
I Linear recursive feature machines (lin-RFM)
28
where sh: R x R IR is a "reweighting" function that acts entry-wise on (u (t)) and 0 is a hyperparameter governing the strength of the regularization. We we will study the behavior of the iterates (*) ((t) in the high-dimensional limit where n and both approach infinity with fixed ratio = K. Since our primary interest to reveal the feature learning capabilities of such algorithms when 0* is a -igh-dimensional signal with low- -iwdimensional structure, will typically focus on the regime where k T. where T is the number of total iterations. This ensures that the total number of observed samples nI' is smaller than the ambient dimension d.
Before proceeding to our main results, we note that this formulation encompasses wide variety of classical and modern algorithms (summarized in Table
Alternating minimization: One perspective on this algorithm is to consider it as way to perform alternating minimization on the non-convex loss function
1 L(u, v) y X(uov) +
1
uul2 d 110112.
Using the fact that the loss function is symmetric and v, choosing suu, v) = u recovers the bainhe alternating minimization algorithm for this loss. In other words, vf simply switches the two parameters u and v.
IRLS algorithms for sparse recovery: As shown in 28 classical IRLS reweighting schemes used for sparse recovery and compressed sensing [1I 21 can be reparameterized
in the form of (4) (u, v) (u²u² + e)", where different choices of a correspond to different ép penalties. In particular, the choice = 2 - 4a corresponds to the IRLS-p
algorithm of 211.
Lin-RFM [28]: Generalizing the reparameterized IRLS update, the authors of [28) propose the choice u (u, v) = o (u22p3 x²) for some continuous function o: IR + RR+, Here, the quantity u² arises from the average outer product of the gradient of the learned regression function, which was shown empirically in (27] to correlate with the features learned in the weight matrices of various common neural network architectures.
Our goal to understand statistical properties of the iterates for different choices of sh, and in particular how the test eπтоr evolves from iteration to iteration. In the following section, we develop an asymptotic characterization of the iterates that can be used to gain insight into these questions for a large class of reweighting functions and problem settings.
3 A precise characterization of the iterates
In this section, we provide a precise characterization the iterates of the algorithm in Equation with i.i.d. Gaussian covariates. First, we introduce and discuss the two main assumptions needed for our main result. The first assumption is concerned with the distribution of the initialization Uy and the target signal 0":
Assumption The empirical distribution of the entries of (0)) and 0" converges in W2 distance to (0)
some joint distribution No, i.e., E -i=1 8(v) 07)) Wa πl Additionally, (0) + 0 for all i and 0
has bounded entries almost surely.
Here, the requirement of empirical distribution convergence is easily satisfied by common choices of (0)), including the ones vector and i.i.d. Gaussian entries. For typical sparse regression setup, might, for example, consider the 1o induced by choosing (0) 1a and letting 0" have i..d. entries
5

---

that equal 0 with certain probability. The requirement that 0' has bounded entries appears to be an artifact the proof, and is used only in the proof of one technical lemma. In our simulations, we find that our asymptotic predictions often remain accurate when 0* has entries from distributions which are not bounded almost surely (e.g., Gaussian entries).
Secondly, we define the set of reweighting functions vp for which our result will apply.
Assumption 2. The reweighting function wh: R R + IR satisfies the following:
1. Uf V are random variables such that U,V V + with probability I, then (), o 0 with probability 1.
2. vp is continuous and bounded or r/² is pseudo-Lipschitz of order 2.
This family allows us to consider many of the choices of &i discussed in the previous section, including (u, v) u (AM on linear diagonal networks), (l(u, v) //wvl (lin-RFM and IRLS), ((,,v) v) (2) for bounded (lin-RFM). We note that this does include some choices of vi which apply more "aggressive" weighting, such as ((,, v) |ww|. Nevertheless, we can apply our theoretical predictions for these choices va after passing the weights through bounded activation (such as a sigmoid). In Appendix(D) we show that our predictions often still show excellent agreement with empirical simulation even when the boundedness assumption violated.
Our results are stated in terms of the following iteration, for t 2 O:
TO² + + x) - 2 A..(VV)) TV2 + BA }
B2k
81+1 = arg max min
+20 820
8
Q++1 T++V(( +
(5)
T41 12 B++1
11 = (H(WUU((11,)))
where G, ~ (O,, 1). In words, given probability distribution II, over R x R, T++1 and 8 are scalars computed as the uniquessolutions solutions to a deterministic optimization problem (this can be solved casily by studying the optimality conditions, shown in Appendix C). Then, Q++1 is defined as random variable that is a function of (V, 11, and G; ~ (0, 1). Lastly, 14.1 is defined as the joint distribution of (Q(+1, V) and e.
Given this iteration, we obtain the following result, which is proved Appendix A
Theorem 1. Suppose Assumptions 7 and 2 are satisfied. Then, for any 0 and any function
g: R + R such that E PL(2) or is bounded and continuous, we have
1 d
(++1) ... (0)))) & Elg(Qu+1
d
where the expectation is over the independent random variables (V, 0) ~ II, and G, N((0 1).
The limit in this theorem should be interpreted as being the limit as n, - oo with their ratio x held as constant. Applying the above theorem for each t 2 0, we can get precise asymptotic predictions for wide variety of test functions of the iterates. One example of particular interest the test error, which we measure the normalized & distance between (f++1) (f() and 0*
corresponding to g(u, 0) = Iue 0l (we provide proof that this is PL(2) in Proposition in Appendix B). We note that the limiting expectation can be computed via simple Monte Carlo simulation of a scalar random variable.
From technical standpoint, our result is obtained by applying the Convex Gaussian Min-Max Theorem (CGMT) [30, 29 to the weighted least-squares objective function in 4 Previous works have obtained similar distributional characterization for the solution to least-squares with anisotropic covariates (where the "weights" are the square root of the eigenvalues the data covariance) [8] However, while [ assume that the eigenvalues uniformly bounded by constants, this is not reasonable assumption in our setting, since many common choices of ai are not bounded and hence
'The uniqueness of the solution is shown in the proof of Theorem

---

2
2
sue (theney)
10-23
Eanhh [Mc]] (theery)
tank suel (em) (thory)
tank [ue]] (erm
(thoery)
(emo
u (sim)
Test Eeror
Laul o (leayy
Text Ernor
touk oheoyy
tath (in)
10
tank Gin)
10
10
2
3
2
7
Iterations
Iterations
(a) σ 0.1
(b) 0.001
Figure Theoretical predictions and simulations of the test error 3 v 0" 1|a (log scale, pluses denote the median over 100 trials and the shaded region indicates the interquartile range) for two different noise levels, where
250, d 2000, and l has Bernoulli(O.0)) entries. Here, & = |ue| corresponds to the classical IRLS weighting from [[]] &i = tanh |ue] is version of lin-RFM, corresponds AM, and de tanh u² is reweighting scheme we introduce. We note that the yb which depend only on us can lead to oscillatory behavior in the test risk.
(+) not necessarily bounded uniformly for 1. A second key difference is that we need to obtain a distributional characterization which can be applied recursively for all t 2 0. The analysis in [8] assumes convergence of the initialization in W4 distance and proves convergence of the estimator in W3 distance. However, to apply the result recursively in our setting, if assume that the empirical distribution of (w) 0*) converges in Wi distance, then we need to show that after one iteration, the iterates also converge in Wi distance (and not in any weaker sense).
To overcome these differences, we use different technique to show distributional convergence the iterates. Similar the approach in [5], we apply the CGMT to perturbed optimization problem, which ultimately allows us to show convergence of test functions of the solutions to the unperturbed problem. While this approach necessitates the additional assumption that 0" has bounded entries and we obtain results for a slightly smaller family of test functions (note, for example, that the squared loss o(u, v, ()) (ur 8)2 is not PL())), obtain distributional convergence result that can be applied to a sequence of recursively defined least-squares problems which define the trajectory of algorithm, rather than single optimization problem. Moreover, simulations Theorem still often apply without these additional
in Appendix []suugest that the predictions assumptions, including in the case of the squared loss, indicating that these additional assumptions could potentially be weakened with more complicated analysis.
3.1 Application to sparse linear regression
In this subsection, we apply Theorem to a sparse recovery setting and compare the asymptotic predictions numerical simulations where n = 250, = 2000, and 0* has Bernoulli(O.0) (0.01) entries (so the expected sparsity level is E[s]= = 20). We run Algorithm with initialization (0) 1a for four different choices reweighting function and display the test error at each iteration (median over 100 trials) in Figure I For each choice reweighting function d), we choose the regularization parameter that minimizes the asymptotic test loss achieved within 8 iterations, and plot the corresponding trajectory. As shown in the figure, the numerical simulations show excellent alignment with the asymptotic predictions for this moderate choice of and d.
nigh-dimensional Gaussian data. First, consider setting
The asymptotic predictions show that this family of algorithms can find solutions with low test error within only few iterations. Our results also reveal fine-grained differences in the convergence
7

---

behavior of the different algorithms. For instance, more aggressive weightings 3 tanh |av) and v tanh u² seem find better solutions after several iterations. Interestingly, the weighting functions which depend only on u (like alternating minimization) sometimes display non-monotonic, oscillatory decay of the test loss, particularly in the low-noise regime. However, we do see a steady decrease in test error after every pair of iterations (e.g., in AM, after both parameters have been updated), Finally, we note that our framework allows for analysis of new algorithms for training LDNNs. In particular, to our knowledge, weighting functions of the form o(u²) have not been previously considered for this task, but our results indicate that this small modification to AM is competitive with many existing algorithms in this setting.
4 Grouped IRLS and the benefits of structured feature learning
In many scenarios, the unknown signal 0 is known to possess additional structure that can be leveraged during training. One commonly studied example of this is structured spursity, or group sparsity, where 0* has many blocks which are zero. In this section, we generalize the results of Theorem To the case where the reweighting function respects this additional structure the signal, i.., & acts blocks of v, rather than individual coordinates.
Concretely, we consider the following modification our formulation. Let b be constant and write Ed as a product space over M $ factors: R x... Rd, Then, 0", (0) v(1) E R can all be represented as M stacked blocks, each in Rº. Under the same linear measurement model. we now let & : 3 x R R act on each of the factors of (u(+) 0*), and consider the same Algorithm 4
Here, the case 6 = 1 recovers the results of the previous section, but the case 1 allows us to study the interplay between signal structure and reweighting scheme in more fine-grained way. For example, suppose 0" is known to be group-sparse meaning that many of the factors {0; J.M. are zero. In this case, it might make sense for và to return vector of the form
(u (1) (1)) = a; a;11;
for some 0ę E RR that is chosen function u and v." This corresponds a reweighting scheme which acts on blocks, rather than individual entries. Another way to motivate this "grouped reweighting" is to leverage the connection the n-trick, in (3). In particular, the group Lasso problem can be written in the following variational form (2:
M
10.12 TJc
+ min min L(O) BER NERM
+
2
where the closed-form solution to the n minimization yields the classical group norm regularizer Here, reparameterizing as O;; - / and u, + A. gives rise naturally to the grouped V
Hadamard parameterization, with e; a,
From the perspective of linear diagonal networks, such approach is equivalent to "tying" together the weights of the hidden layer that correspond to each block. Rather than studying gradient descent/flow for this parameterization (as in [16] 17), consider optimization approach that
relies alternate updates of (t) and (t))
We make the same technical assumptions as in Assumptions 1 and with the natural modifications to accommodate b 2 1:
1. Ho the limit of the empirical distribution of factors of (v( 0*) and hence is distribution over I2 x R
2. Each factor 0: E R for i = ,,...,,, .M has bounded 4 nnrm almost surely.
3. 4/ is bounded and continuous or each of coordinate projections satisfies 15 E PI.(() for I,....b. b.
8

---

Group bladd (doous)
Grap blid
Group bilddd odm)
Grop
Crroppazam chomy)
Grap saart (thonry
Locup anaar (sims)
0.006
Test Ercor
E Teat Eroor
10
0.004
0.002
3
10
Iterations
Ground truth group size
(a) Comparison of trajectory
(b) Effect of group size
Figure 2: Group-blind (1))b) vs. group-aware (1h reweighting when l has group-sparse structure.
We set n = 500, d = 4000, a = 0.1, and 0 Lid. minimize the asymptotic test error achieved. Simulation results are the median/IQR over 100 trials. Left: Comparison of the test error trajectory (log scale) for a fixed block size 8. Right: & test after T iterations, for varying group sizes.
~ Bernouli(Q.00))4, For each curve, À is set to
straightforward extension of Theorem D yields the following generalization for the grouped algorithm, where V, θ Gt, Q++1
IR are now vector-valued random variables: For t 2 0, let
TO
1
b 0 Bx
TH+1, B,+1 arg max min T20 $20
+ + x) -2 -
3
b
VV 3X
Q++1 V C (0 + B G.VE)
(entry-wise division)
T++1 V + A1,
11+11 = (Lw((((Q)))) (Q++1, V),
Here, G. / (0, It). Then, we have the following result, which is proved in Appendix A
Theorem 2. /Generalization of Theorem for 1/ Under the assumptions above, for any t 2
and any function g: (R^}3 - R such that g PL(2) or g bounded and continuous, have
glu (++1)
(t) 07) & Blo(Q V.
M
Given a reweighting function v/, Theorem 2 characterizes the distribution of the factors (blocks) of the iterates. Hence, by choosing g(u, u. G 0), predict the exact limiting test error for this family of algorithms.
Computing these theoretical predictions reveals that choosing & in group-aware way lead to
significant performance improvements compared to coordinate-wise reweighting. Figure 2 we fix 0.1, 500, d = 4000, and set the overall expected sparsity level of 0" as in Figure 1. We compare the performance of Algorithm 4 for a ggroup-blind" (wyp) and "groop-aware"" (N)aa choice of reweighting function:
Ugs(u, v) = tanh lu @ v - note this identical to one of the reweightings considered in Section
Vigu (u, u) E/= tanh (u,0y))1,
The theoretical predictions align with simulations and show a notable improvement in performance when using the group-aware scheme with Moreover, as the group size increases, the performance of tigs remains approximately the same, indicating that it is not able take adapt the

---

group- structure. By contrast, using gg leads to a consistent improvement in test error as b gets larger. Hence, the test error when using the group-aware scheme scales with the size/number of groups. rather than the overall sparsity level.
5 Conclusion
In this paper. we derived a precise asymptotic characterization of the iterates of a family of algorithms for learning diggesiiensio linear models with linear diagonal networks. We used these predictions to obtain fine-grained predictions of the test error at each iteration for various existing algorithms for this task, and we showed that our framework can also be used as a test bed for new variations on these algorithms that take similar form. Lastly, we demonstrated the advantage of embedding more structure into the model by tying together groups of weights when the ground-truth has structured sparsity. Several interesting open questions about these types of algorithms remain. While our simulations align very well with the predicted asymptotic trajectory, it would be interesting to obtain finite-sample guarantees that hold even for batch sizes that are much smaller than d (as in the "mini- batch" case studied by [9].. Moreover, seeing as our analysis depends crucially on the independence the covariates at every iteration, developing precise predictions of the trajectory in the non-batched setting remains an interesting direction for future work.
Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers for their helpful feedback. This work was supported by an NSF Graduate Research Fellowship (DGE-2039655), the NSF AI Institute AI4OPT, NSF grants IIS-2212182, CCF-223915 and 2112533, and gifts from Amazon and Adobe.
References
[1] Francis Bach.
The "n-trick" or the effectiveness of reweighted least-squares. "cCEEex-tticcckee.. the- fee--fffcctivenessof....
https://francisbachh com/the
reweighted least squares/ 2019. Accessed: April 2024.
[2] Francis Bach, Rodolphe Jenatton, Julien Mairal, Guillaume Obozinski, et al. Optimization with sparsity-inducing penalties. Foundations and Trends@ in Machine Learning, 4(1):1-106, 2012.
[3] Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory, 57(2):764-775, 2011.
[4] Michael Black and Anand Rangarajan. On the unification of line processes, outlier rejection, and robust statistics with applications in early vision. International Journal of Computer Vision, 19(1):57-91, 1996.
[5] David Bosch, Ashkan Panahi, Ayea Ozcelikkale, and Devdatt Dubhashi. Random features model with general convex regularization: A fine grained analysis with precise asymptotic learning curves. In International Conference on Artificial Intelligence and Statistics, pages 11371-11414 1414. PMLR, 2023.
[6] Kabir Aladin Chandrasekher, Mengqi Lou, and Ashwin Pananjady. Alternating minimization for generalized rank one matrix sensing: Sharp predictions from a random initialization. arXiv preprint arXiv:2207.096660 2022.
[7] Kabir Aladin Chandrasekher, Ashwin Pananjady, and Christos Thrampoulidis. Sharp global convergence guarantees for iterative nonconvex optimization with random data. The Annals of Statistics. 51(1):179-21 2023.
[8] Xiangyu Chang, Yingcong Li, Samet Oymak, and Christos Thrampoulidis. Provable benefits of overparameterization in model compression: From double descent to pruning neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 6974-6983,
2.21.
10

---

[9] Rick Chartrand and Wotao Yin. Iteratively reweighted algorithms for compressive sensing. In 2008 IEEE International Conference on Acoustics, Speech and Signal Processing, pages 3869 3872. IEFE, 2008
[10] Hung-Hsu Chou, Johannes Maly, and Holger Rauhut. More is less: inducing sparsity via overparameterization. Information and Inference: A Journal of the IMA, 1233):437-1460, 2023.
[11] Ingrid Daubechies, Ronald DeVore, Massimo Fornasier, and C Sinan Güntürk. Iteratively reweighted least squares minimization for sparse recovery. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 63(1);1-38, 2010.
[12] Donald Geman and George Reynolds. Constrained restoration and the recovery of discon- tinuities. IEEE Transactions on Pattern Analysis & Machine Intelligence, 14(03):367-383, 1992.
[13] Peter D Hoff. Lasso, fractional norm and structured sparse estimation using Hadamard product parametrization. Computational Statistics Data Analysis, 115:186-198, 2017.
[14] Prateek Jain, Praneeth Netrapalli, and Sujay Sanghavi. Low-rank matrix completion using alternating minimization. In Proceedings of the 45th Annual ACM Symposium on Theory of Computing, pages 665-674, 2013.
[15] Chris Kolb, Christian L Müller, Bernd Bischl, and David Rügamer. Smoothing the edges: general framework for smooth optimization in sparse regularization using Hadamard over- parametrization. arXiv preprint arXiv:2330..03777 2023.
[16] Akshay Kumar, Akshay Malhotra, and Shahab Hanidi-Rad. Group sparsity via implicit regularization for MIMO channel estimation. In 2023 IEEE Wireless Communications and Networking Conference (WCNC), pages 1-6. IEEE, 2023.
[17] Jiangyuan Li, Thanh V Nguyen, Chinmay Hegde, and Raymond KW Wong. Implicit regular- ization for group sparsity. arXiv preprini arXiv:2301. /2540, 2023.
[18] Po-Ling Loh and Martin J Wainwright. High-dimensional regression with noisy and missing data: Provable guarantees with non-convexity. Advances in Neural Information Processing Systems, 24, 2011.
[19] Mengqi Lou, Kabir Aladin Verchand, and Ashwin Pananjady. Hyperparameter tuning via trajectory predictions: Stochastic prox-linear methods in matrix sensing. arXiv preprint arXiv:2402.01599, 2024.
[20] Charles A Micchelli, Jean M Morales, and Massimiliano Pontil. Regularizers for structured sparsity. Advances in Computational Mathematics, 38:455 4899, 2013.
[21] Karthik Mohan and Maryam Fazel. Iterative reweighted algorithms for matrix rank minimization. The Journal of Machine Learning Research, 13(1):3441-3473, 2012.
[22] Whitney K Newey and Daniel McFadden. Large sample estimation and hypothesis testing. Handbook of Econometrics, 4:211 1-2245, 1994.
[23] Scott Pesme, Loucas Pillaud-Vivin and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218-29230, 2021.
[24] David Pollard. Asymptotics for least absolute deviation regression estimators. Econometric Theory, 7(2):186-199, 1991.
[25] Clarice Poon and Gabriel Peyré. Smooth bilevel programming for sparse regularization. Ad- vances in Neural Information Processing Systems, 34:1543. 54551 2021.
[26] Clarice Poon and Gabriel Peyré. Smooth over-parameterized solvers for non-smooth structured optimization. Mathematical Programming, 201(1):897-952, 2023.

---

[27] Adityanarayanan Radhakrishnan, Daniel Beaglchole, Parthe Pandit, and Mikhail Belkin. Mecha- nism for feature learning neural networks and backpropagation machine learning models. Science, 183166690)1146111146 Mar 2024.
[28]
Adityanarayanan Radhakrishnan, Mikhail Belkin, and Dmitriy Drusvyatskiy. Linear recursive feature machines provably recover low-rank matrices. arXiv preprint arXiv:2401.04553 2024.
[29] Christos Thrampoulidis, Ehsan Abbasi, and Babak Hassibi. Precise error analysis of regularized m-estimators in high dimensions. IEEE' Transactions Information Theory, 64(8):5592-5628,
2018.
[30] Christos Thrampoulidis, Samet Oymak, and Babak Hassibi. Regularized linear regression: A precise analysis of the estimation error. In Conference on Learning Theory, pages 1683-1709. PMLR, 2015.
[31] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances Neural Information Processing Systems, 32, 2019.
[32] Roman Vershynin, High-dimensional Probability: An Introduction with Applications in Data Science, volume 47. Cambridge university press, 2018.
[33] Cédric Villani al. Optimal Transport: Old and New, volume 338. Springer, 2009.
[34] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-367, PMLR, 2020.
[35] Peng Zhao, Yun Yang, and Qiao-Chu He. High-dimensional linear regression via implicit regularization. Biometrika, (0334):0033-104) 2022.
[36] Liu Ziyin and Zihao Wang. spred: Solving L1 penalty with SGD. In International Conference on Machine Learning, pages 43407-43422, PMLR, 2023.
12

---

A Proof of main results
In this section, we provide the proofs of main results.
A.1 Notation and background
For convenience, we first restate the main notation that is used in our proofs.
Notation The ones vector of dimension is denoted 1a- We write S when a S Cb for some sufficiently large constant C O which does not depend on d. We denote the element-wise multiplication (Hadamard product) of two vectors æ and y as a @ y. Flement- division of two vectors is denoted as y We use the shorthand (-)+ max(- 0). A function f: Rº called pseudo-Lipschitz of order 2 if, for all x, y E R
(f(x) f()) C(1 ||0||| + | y/12
for some constant C 0. The set of such functions denoted PL(2).
Convergence in probability of a sequence of random variables X. to a random variable X denoted Xa & X. Convergence in Wasserstein-2 distance of sequence of probability distributions Pa to a limiting distribution v is denoted as Va 3 v, and this fact is equivalent to the statement Ex-va g(X) + Ex-- (XX for all PL(2) 3 If the Va are random probability measures, we
say that va 3 the same convergence holds in probability, i.e., Ex (X) & Exar g(X) for all E P1(2). The empirical distribution
vector E IR: is defined as Ei= 8(z;). where 8(2;) the Dirac delta distribution centered at 2;- For any random variable (or group random variables) X, we use the notation Law(X) to denote the probability distribution of X.
We also define two key quantities which appear in the analysis.
Definition The Moreau envelope function of a lower semi-continuous, proper convex function L: I2 R with step size T is defined as
Mc(a;T)= min ((y)
1|2-
YERP
2T
The proximal (prox) operator of E with step size t, denoted prox, (xx T) is defined as the arg min of the above optimization problem.
Lastly, we restate the version of the Convex Gaussian Min-Max Theorem (CGMT) that we will use in our proofs.
Theorem 3 (Convex Gaussian Min-Max Theorem 29 Let G R . R h E R have
i..dd (0, 1) entries. Let Sn C RR.h and S., C Ih be compact, convex sets, and f: R x Rm + IR
be convex-concave on Sup Su. Define the following two min-max problems:
(G) == min max u Gw + f(w, u.) wES, MES.
dg. h) == min max||w [axl|///l2u T + | T + f(w, u)
weS,, ueS,
Then, for all c R and t 0,
P[[o(G) - t) 2 -d
A.2 Proof of Theorems Tad 2
Proof of Theorem Assume that Ei_ 5(0 (0) 07) t= 0; will show later that it holds at time t ++1 1).
W/ TI, (note that this holds by assumption at
110,0% 0*) to the joint First observe that convergence of the joint empirical distribution of (u(++)), 110, distribution of (Q++1, V, V) in Wasserstein-2 distance implies that
d
((+1) g .
(0) 0,) f Elg(Qu+1, V,0)],
13

---

for any PL(2) or which is bounded and continuous. This is because W2 convergence implies convergence in expectation any pseudo-Lipschizz function of order 2 B Lemma 5] and of any bounded continuous function (since W2 convergence is stronger than weak convergence 33 Theorem 6.9]). Hence, it suffices to show that
d
6
(++1)
(8) 0,) 3 Law (aw
(8)
where () 11, and Q++1 is defined as in Eq.
Recall that the objective function for the update on given by
(++1) arg min WER n2
(0)
1
x 10 (u
| 2
A gu11
vd
Rather than study this update directly, we first analyze slightly more general problem (following the approach in (5). Let h: R3 + R be continuous test function with ////|2 C and that satisfies one of the following:
1. h is uniformly bounded.
2. hu, v. 0) u².
Then we consider the following problem (the dependence of X, y, 6, and a on & is dropped to simplify the notation):
Pi(
min wERd
y 12
x w) Vd
+
1u
1
h(u, Uy 67),
where μ E 1-pp 44) and vr¹ & is chosen sufficiently small that the objective function (scaled
by d) is X-strongly convex for all in this range. The case = recovers the original problem of interest.
Step 1: Convergence of the loss
Rewriting this in terms of the error vector A = v2 (u Q 0*)
have
min - XA AER n
A
dA 0- v
VdA, 0 Ug. 02
(i(x)
02
In writing this, we use the fact that N/ * 0 for all i with probability (and the notation in the second-
to-last term indicates entry-wise division). Now, using the identity |1.112 = maxy 2q (.) || we can write this as
(1 min max AERd gER" Vn
2 a E
2 q XA nn
||1123
Vda 0
# d
Vds; + O
Us 02
V
U;
(10)
Next, in Lemma 1 we show that there exist Euclidean balls BA and B,, each of radius Ci [ such that, with probability approaching without changing the value of Pi (u), so we can study
constrain the feasible lie with these balls
2 ((()) min max q c
2 q XA
||
+
V7A
0"
P
Σ
d
+ O7 Vdl:
8
A EEB, Vn
d
Vi
(11)
where (n) = P (n) with probability tending 1. We can therefore condition on this event for the remainder of the analysis without changing our asymptotic conclusions.
14

---

Now, noting that this is in the correct form to apply Theorem B we define the auxiliary optimization problem
2 q c
2
2
A
VdA
0
P2(H) min max AEBa qE vn vn 111122
A
vn
|
22
q
11
+
d
v
2
(12)
/ d
Vda; +
8
U/.
0
d
U/
(13)
where 9 Rd and h R" have i.i.d. standard normal entries. By Theorem for all 0 > 0 and fixed P(μ)= E R, (() P(x) 83 s 22(( P(u)) 8}.
In particular, we can find some P(r+) such that P2(μ) & P(N), then we can conclude also that (1( & P(n)-
To accomplish this, we next perform a series of simplifications to P2(μ) which will later help us characterize its asymptotic behavior. First, decouple the optimization over q into its norm and direction, and the latter can be solved explicitly. Letting т = || this yields
2T |E Vn
A
Vda + 0
P2(n) = min max AEBA OSTSR
12h1111
A
d
v
2
(14)
Vdd;
0
h
Ut.
0
0/
i=1
where R Cillu|loo Next, note that e and are independent Gaussian vectors and hence
2h d V0² h. So, we have
P2(4)= d min max AE OSTSR
2r h l
2T
9
2
+
VdA +
0
n2
v
Vd
+
6
+
h
Us.
0-
d
1=1
(15)
Before proceeding further, we rewrite this in terms of a minimization over the variable u
vaa+o
2T
l2
2r Vnd
(2) min max MEB, OSTER n
02 + Ilu 3 v
0.112
(u 3 v
7² + d
+ 1
h(u, 24 07)
(16)
Here, B, = { u eud Rdd | (1 -0) S R]. After this step, observe that the objective
function is strongly concave in and strongly convex in (the sum of the last terms is strongly convex in u based on the assumption that has bounded Hessian and is sufficiently small). Since the objective function is convex-concave over convex and compact sets, we can invoke Sion minimax theorem switch the min and max. Furthermore, we use the fact that v mins>o 2 & to write this
T|/h1122 w.B... /n
0²
1u ④
0 1 + 3
2r
P2(u) max min OSTER
+
g
O 0*)
B
3d
v nd
as85atR
22
1
2 +
p
h(us, vi 0.) (17)
15

---

Here, note that we can add the constraint on 3 without changing the solution since the optimal value of 3 will be obtained V o02 + l/ @ 0- E la, + R| for all feasible u.
To do this, first show in Lemma that the optimal solution to the unconstrained minimization is strictly feasible for large enough C1. and hence, the unconstrained and constrained minimizations over are equivalent. Next, observe that the unconstrained problem is separable over the indices, so we need only to solve the scalar problem
Next,
explicitly solve the inner minimization
T| 2111112(uuu)) 07)²
2T
> du² h[(uu,, v1, (;)
11 d
min B4 ER
(u 0:)
(18)
Bdyn
Vnd
d
Completing the squares, we obtain that the above problem can be written in terms of the Moreau envelope (Definition [ of the function f(u) du + ph:
= T (Ba))=-00;)* 3
E0; 187-Bgve
BE
2
(6)2
where we have introduced the shorthand notation & = P2 above, we obtain
|hle
Substituting this into the expression for v
P2(μ) max min OSTSRoSiSON
To2E B
+ TBE -2
1 d
TE
(3g.VK
(39.Vk-E0.)
3
35
d
[00
BgivE
3
Ma Mx(r2th((er.01
2
(1)
= max min fa(T, 8) 05
(19)
Now that the optimization has been fully "scalarized", proceed by considering its asymptotic behavior. First, note that the partial minimization over u preserves the concavity/convexity (r, 8). In Lemma 3 we prove that for any fixed 7 and B, the objective function fa(T, 3) converges in probability to
f(T, T + - x) -2- -2+ +B M Ma(p+an(.v.o)
3GVE +
2TV2 )
B
(20)
V
where the expectation is over (V) ~ 11, and an independent G NoO,, 1). We note here that Lemma the only place our proof which requires the boundedness of the entries of 0*.
Since fa(T, (3) is strongly concave in 7 with parameter for all feasible 8, we can conclude that f(r, (3 is also strongly concave in 5 with parameter 1. Directly taking derivative with respect to 8,
we also find that has a single non-negative eritical point, at the point Vo² + El(uV - θ) where
à = e(V ( (
θ 3GVE K
B
2-
So, we can conclude that f has unique saddle point (f, 3). Note a deterministic function that does not depend on d and hence (+, 3) are also deterministic and independent of d.
Now let C2 == max(f, ff 33 + 1. By the "convexity lemma" (as stated in (24),, pointwise convergence (in probability) of convex function is uniform
compact sets. So, this result implies that the convergence is uniform over (r, 3) E [0, C22 x [0, C2 so
max min Sa(T, S max min ((T,8) OSTSC, 0SBSC,
DErSC.0C05C
Let (fa, Ba) denote the optimnal solution for the problem on the left. We can also conclude that (fa, 3a) & (F,3) by [22] Theorem 2.1], which states that uniform convergence in probability of
Note we could not have directly applied this to the feasible sets of P2(μ). since R may have dependence
on d.
16

---

convex function over a compact set implies convergence of the optimal minimizer. So, with probability approaching 1, (fa: 9a) are strictly smaller than C2. and the same solution is also optimal for P2(.). We can therefore conclude
P2(n)
+ max min US 0SPSCa
[(T, max min [(T, =: P())
r20 020
Therefore, by Theorem 3 for any fixed 1-,",,"" we have the convergence Pi(H) S P(H).
In the special case μ = 0, we can further simplity the Moreau envelope term to obtain
TO²
O B. TV2 + B
P(0) max min
7 --)
TÀB
(21)
T20 820 3
Step 2: Convergence of the optimal solution We next need to extend this result to the desired Wasserstein-2 2 convergence result (8). Recall here that u (+++) is the solution of (1).).
First, let h: R I be any bounded, Lipschitz function, and let (R)) be sequence bounded, twice-differentiable functions that converge uniformly to h as k - (e.g., the convolution of with sequence of mollifiers). Let P(A) (μ²) ((k] (μ) be the optimal cost of P1 and P when using
test function (e) and for μ E [-μ", p="]. Note the convergence ((() (1) F (p()(p) for any y in sufficiently small neighborhood around zero holds by Step 1.
By the uniform convergence of the (k)) h,
lim p(a) P((e)
lim (4) P(W). k-0c
Now, fix 6 > 0 and choose &s large enough that ( (4) - Pi(H) < and (p) - P(H)] < < Then.
(() - P(H) > 8) < < (1) p(() (H)) > 5/3 +0
since p!&) P, P3 (k) for all Hence, we can also apply the result Step to any bounded Lipschitz function h.
Since the convergence result of Step holds for any μ in a neighborhood around zero, we can conclude that
1
(1+1) h(u h(u(tt)) Us. 01)
dP(μ)
dp
μ=0
where the derivative is well-defined since P has a unique solution in a neighborhood around zero. The proof of this fact is identical to that of Lemma of [5], so we omit it here. Moreover, using the Dominated Convergence Theorem to differentiate inside the expectation, we can compute this exactly:
dP(()) Ehh prox du p==
O-BBGRR
B
Eh +V((- 3GVx) v.
..O
27V2
7V2 + 3A
where (3,4) are found in the optimal solution to P(0).
Hence, for all bounded, Lipschitz h, we have
1 Σ Muultt v. (.,07)) 07) Eh
-310 v.
Eh
4V2 AA
so the empirical distribution of the triple (u. (+V()-8 FV2+8A
(++1) (u(++1) U/, 03) converges weakly to the distribution of
the random variable ( ( BG.a)] V,e), where G ~ (0, and (V, II, By choosing Mu,, v) = x² we also know that second moments of the empirical distribution converge in probability. Hence, the convergence can strengthened from weak convergence to convergence in W2 distance (see. e.g. [33] Thcorem 6.9]).
17

---

Step 3: Verifying the inductive hypothesis (++1) 8,) 3 Law(w V), 0) = I[++1- 8
Lastly, we need to show that
where @V) IT. Here, weak convergence follows from the result of Step 2 since v is continuous map. To show convergence of second moments, we need to show
((+1)2 1 Σ (( (nnlt" o + BivQQe,, V)
For dē that satisfies Assumption 2, this convergence is immediate from the result of Step 2 (since W2 convergence implies convergence in expectation of bounded continuous and PL(2) functions). Therefore, the initial inductive assumption made at the beginning of this proof holds at time t + 1, and we can apply the result inductively to conclude Theorem TT
Proof of Theorem 2 The proof is an extension of the proof of Theorem T to the case where the test function h acts on blocks rather than individual entries. Much of the proof identical, so we only sketch the argument and highlight the major differences here. We begin with the inductive hypothesis
that MI Ee 6(1))) (8) 03) 3 TI, for a known distribution II over IR x R². Recall here that M
denotes the number of blocks/factors of size b (so M (.)
Then, let h: (IR - IR be test function with 2//2/ I C. and such that either h is bounded or hu,, 02 0,) = [u: 112. We consider similar perturbed optimization problem:
1 X(u o v) Vd
Pi (1) min WEB"
y
Σ h(,, vs 07).
(22)
Again, we consider this for |u| Noting that the proof of Lemma I still holds in this grouped case, we can constrain P1(H) to be over compact sets and apply the CGMT obtain the auxiliary problem
bC so that the optimization problem is à strongly convex in u.
2
2
2
A q a -1111 A Vda + 0"
P2(u) min max
111lag A Vn
AEBA qeB, vn
v
2
(23)
M Mu, v. 0.). (24)
M
The sequence of "scalarization" steps on P2 is identical to in Theorem
until arrive
T||h||2
o²
Iu
0-12
2r Vnd
Pa(u)=
max
min u
+
+ 3
g (u 10 0*)
T
USTSR
Vn
B
3d
M
1 d
P
h(u, (9)
(7)
M
(25)
Here, since the sum of the last two terms in the objective function is å strongly convex by our choice of μu, the proof of Lemma holds holds without change and we can consider the unconstrained minimization over u. In this case, the minimization is block-separable over each of the M factors of u, so it can be expressed as
min uER
T u; O 0.
-01112
2T /rg, (u, 3 U. 0) A||ul|2 + puli(uus Ui 02)
3
M
1
T
uu C v. - 0,113 2T Vgg, (u: 2 v, - 0) + >1/211/ + pub(uu Us 07)
min
bM
B
M
bM
18

---

where E R² denotes the ith block of and q = (R*)3 - IR is defined as a shorthand for the quantity inside the summation.
Next, we consider the asymptotic behavior of P2( Here, the only term which is different than in Theorem is the term 537 Eis g(vi, 0 gg).. By the same argument as in Lemma can write as the Moreau envelope of a convex function and show that
M
1
Σ
(vu, (gg9) S. F
bM
where the expectation is over (V, +
11, and G N0, ft). After the same uniform convergence conclude that for all μ E 1-"", 14",, P1(n) 5 P()).
argument in the proof of Theorem where
TO² -
P(a)
max min +20
O.
$20
B
In particular, when j the minimization implicit in the definition of q can be solved exactly; this yields exactly the optimization problem in 7 Step 2 of the proof (convergence of test functions of the optimal minimizer) is identical to that of Theorem [ and for the final step (showing the inductive hypothesis holds at the next iteration), we need to argue that the second moment of ((t+1) ((( (f+1) v.) converges its expectation under 11++1:
Σ (v (ttt)
(++1), 2;)1 Bl (Qu+i. V)ll
If & is bounded and continuous or has PL(() coordinate projections (as we have assumed), then the above convergence holds based on the Wasserstein-2 convergence of the joint distribution of (x+11) v).
B Technical lemmas
Proposition
The function g(u, v, 0) lwv - 0 is pseudo-Lipschitz of order 2.
Proof. The result follows from the following series of inequalities:
|lue 0 - lu'v' 8'
< lue - (u'v' - 8')1
LA
lue u'v'l + 10 0'
S |u]lu - w'| + |e'||u - u'| 0 0'|
S (lu] + (oo)+1)(()
u'l - + 10 81)
S (1 + ||>||| + || [1)] 2'11
S 3(1 + | | [2)] x' 112.
where x, x' R denote (u, u, and (ul, v. 0'), respectively.
Lemma 1. Let A g² be the optimal solution to 10 Then, there exists universal constant C₁
such that
lim P(|| s 11 lim P(llq"
[|2 Cillo[1sa) =
d
d
Proof. We proceed via similar argument to Lemma 2 in 5 First, consider the original expression for P1 (1) from 99:
Pi(p) min F(u) + R(uu,,
weRd
and R(w) = 1111 Ei h(u, ((u,,, v.,, Recall here va x(u 3 I x
where F(u) = 1 that R is --strongly convex for all / 1-p4" |r"), and denote the unique optimal minimizer to this
19

---

problem as u" Then, the following chain of inequalities holds, by the optimality of u and the non-negativity of F.
v + R(O) = F(O) + R(O) 2 F(u") + R(w") 2 R(u").
n
Moreover, by the strong convexity of R, we have
R(u") 2 R + VR(O) "n" + 2d
1 u2.
Combining the above two series of inequalities, obtain (recall ())
2d ((() Tu u
2x
u 112 +
X
After completing the square, this yields
2x
d
VR
5 A
1 +
V
(((O)
X2
whence, by the triangle inequality,
2x l
de
| 22 -
VR(O))22 + V
+
VRO
X
2d X VR(O))l2 V 2x
|y/12-
Theorem 5.2.2, Here, standard concentration inequalities for Gaussian random variables (e.g., 32 Corollary 7.3.3]) imply that, with probability approaching || 2|2 X vd and ||e||2 & Vd. And Assumption limpies that ||01|2 l2 S vd with probability tending 1. So,
| ll2 + ||e||2 S Vd
with probability approaching 1. Next, we bound VR(00]|2 |2. Recalling the definition of R,
(V( 2 144 Σ Tm ( h(u, vi, 07)
Ou h(u, us 03)
t=1
Since the function g(v, 0) 3e h(u, u. 0)
Lipschitz (by the fact that has bounded sec-
u=0
ond derivatives), g² is pseudo-Lipschitz of order 2, So, the quantity under the square root con- verges in probability to E(v,o)_,,, g² by Assumption and, with probability tending 1, we have ||V R(O)]l < 1/vd.
Combining the above bounds on |lll|2 and [R((0]]2. 2. we can conclude that ||u2||2 2 Vd. The first part of the lemma follows by noting that
1
1
A
00
12
S
lu"
C
+
0 ll2
Vd
Vd
Vd
1
1
|
1|2
+ vd
|011|2
Vd
||
where the last inequality holds with probability approaching 1. Lastly, the optimal for any A has closed-form q = Vn Vn
XA. By the triangle inequality, we then obtain
|| -
x ||2 S
||||oo
with the last inequality holding with probability approaching by the concentration of norms of and X discussed above, and the bound on A 2-
20

---

Lemma 2. Consider the following unconstrained minimization problem over u E R
1e 0² +
2T
min max ueR uSTSR
0-12
(u 3 v - 0*)
7
nd
7"
+ d
1 2 +
h(us, V u0000 0,).
1=1
With probability approaching I, the solution ut satisfies
x |,0|200 vil uu" 0*)
Proof. First note / ov -0--000000 03) S Jal |12|||/l + 11 and va |00122 is bounded
by constant with probability approaching by the assumed W/2 convergence of 0 fixed limit. Hence, it suffices show that ||"*)] |l2 I vd with high probability. We begin by noting that the inner maximization over r admits a closed form solution, so the problem becomes
2 0²))
min WERd
|ille 2/n
0²
+ o
1 d
1-
g (u (uov-uu (aovv-0')
X 1
h(u 0/0,07 ,07)
2/nd
Now, we can proceed similarly to in the proof of Lemma
Let
)²
F(u)= =
12l12
o2
ā |
o 0-
1
(u (u@v 1 (w@v-0))
2
2/nd"
d
R(u)= = d
u
Σ
h(uc, o4 0.).
Then, noting F is always non negative and R strongly-convex, we use the same argument as in Lemmal] to obtain the inequality
2d
2d
12 VR(O))
u < F(O).
A
X
After completing the squares, we obtain
2d I X
F(0) + d
VR
(R()))
x
so we can conclude
2d VR(0))2++ + x F(O).
2d
| 1/2 < x IVR(()) + F(O) + X VROIE < x
As shown Lemma
(R & a with probability approaching 1. It remains to show that
/ (F S with probability approaching To see this, observe that
| 2/n
F())) =
V
2 +
10
2Vnd 9 T
d
|1/||a
02 d
0011
0
2/n
2
nd
|
0²
0
2
||||||10011 112
2/n
d
2vnn
where the last line uses the triangle and Cauchy-Schwarc incqualities. By concentration of the norm for Gaussian vectors, there exists universal constant such that ||ا|| < and Iglla probability approaching Moreover, by Assumption 1 lla
Va with
v
Va S with probability approaching
Hence, (F S with probability approaching Substituting this into the bound for ||u"||| from above completes the proof.
D
21

---

Lemma 3. Under Assumption / the function fa(T. (Eq. 19 converges pointwise in probability to f(r, (Eq. 20 d DC.
Proof. We consider the limit of each term in La separately. The limit of the terms found by noting that & 1 in probability by Gaussian Lipschitz concentration. The first summation term simplifies as follows:
and r is B
T0
T (BRV (B9Vk=-E0)) =03)² -
T
(B S.
239.
BC
BC
1
23g. KKO, & -T
BC a
where the last line follows from the weak law of large numbers since the Ii are i.i.d. standard Gaussian variables.
For the last term, after again using the fact that & S 1, we need to consider
8 Bgiv
B
1 E q0, 07
M *c))
(0)2
=:
2T
To show convergence in probability of this term, first fix ā > 0. Then, we want to show
P
qe 9.) EV(V,, 8 +
where the expectation is over (V,0) ~ II, and G ~ NO 1). It suffices to show the following two statements:
Eg 1
8
P
(Ui, 07,
(v 0f .4)
+ 0,
(26)
7 1=1
8
P
0, (1,, E
0.
(27)
2
To show (6). we rely on concentration inequality for the Moreau envelope of a Gaussian vector plus bounded vector from [5. First note that
q(v, g9)
min weR
Al 2 + p
b(u,, 2/ 07)
u 2
0"
+BVKg]
2 rr
min wER
Al| +
huu,, U/, 07) + +33x
BVK +g
BVK
1=1
3/K0
d
0"
2
min d BeRd
X
+
22,87
+ TBK
0
+ 9 3/R
1/
kml
0"
Me
9:
//K
2T BK
where the second to last line follows from the change of variable 0 u and ((0) =X 8//0
3y
Fi, 0,,07) Here, for fixed v, is a proper convex function of 0. Moreover, by o" Assumption /æR has norm of order Vd with high probability. Hence, by 5 Lemma 8), this quantity concentrates around its expectation (with respect to g), and conclude that there is some e > such that
11 Ed h (y/il
d
8
cT 232
(v,, 0g,9.) - E, Σ q(v, Of gi)
- 0.
2
d8
22

---

To show 27 note that
1
[v ,g9g) gg
Es q(v, u 07 G).
Observe that this quantity is an expectation with respect to the joint empirical distribution of (v, 0*) By the assumption of W2 convergence of the empirical distribution of (v, 0*) (Assumption I we can show that mapping (v, 0) + Eg q(v, 0, G) is bounded, then we can conclude that the above quantity converges in probability to E (V, O G), with (V, 1l. To see this, recall that for all G' E R, is bounded below as
(v, 0, G) 2 min λu + ph(u, v, 0):
which is always bounded below since h is bounded (or, in the case h = u² the lower bound is zero). Next, for given G, we can bound above as
q(o, 0, G) S wh(0, v, 0) + (0-3VKK))
Hence,
BG a(u,0, S ph(0 0, 02 7Bk < C
for some universal constant C > 0, since is bounded by assumption and h(O, v, 0) is either bounded above equal 0 (in the case where h = n²). Combining (26) and (27) yields the desired result. D
C Solving the min-max problem
Below, we show that the max-min problems (5) and (7) have easily computable solutions. Note suffices to consider the grouped case 7 since we can apply it with b = to recover the ungrouped the scalar max-min problem in (8)
case. The following derivation closely follows the analysis
who study similar scalar problem (albeit in the case of À 0, which we do not consider).
First recall that, as shown in the proof of Theorem ! there exists unique saddle point, due the strong convexity in r and strict convexity in B. Now, taking derivatives with respect to t and B, we obtain the following saddle point conditions:
x2 + B( - K) -2-r 2r + BA E
6 + B2k (TV? + 32)2
0
=
(28)
B
0 T
O + 8K
7(1 (1 x) T
T RE
(29)
32
TV?
BA
(TV? + 82)2
Solving each of these equations for the quantity BB E two, we arrive at
0 +31)
and then equating the
1
1
2rB
2rB - - k.) -022= k.)
-2 R
b
VV?
+
which implies
b
== B(1 - k) X32 R
TV?
+ 3A
Defining the auxiliary variable 7 = =// this yields the fixed point equation
b
1
= 1 - + Ar
(30)
b 1=1
7.2
+
23

---

Substituting this back into the first optimality condition in closed-form, in of y:
28
we can express the optimal / in
a² X2B
2i=1
6
8 +VV++x*
27 + K 1E
Di
Finally, the optimal can be found simply 7 yB3
This yields a simple recipe for solving the min-mas problem. First, compute the positive solution to the fixed point equation (30) (this can be found easily using standard numerical solvers). Then, (B, are both given in closed-form as functions of 7 (where the required expectations can all be approximated via Monte Carlo simulation).
D Further simulations
In this section, demonstrate that asymptotic predictions provide accurate estimates of the test error, even when some of our technical assumptions are not satisfied.
First, we compare the two "heavier" weightings considered in Section 3.1 (u, v) Lanh |uv] and (u, v) tanh u², to the same weightings without the bounded tanh activation: sb(u, w) |uv] and (u, v) 2 We note that the reweighting choice |we| is considered in [21 28) as a limit as p - 0 of the classical IRLS update for Lp minimization. Figure Bh, we consider the same sparse
regression as in Section 3.1 i.e., with 250, 2000, a 0.1, 0, iid Bernoull((0.0)) and chosen to minimize the predicted asymptotic loss.
For each choice of th, apply the theoretical predictions of Theorem even if v violates Assumption 2 We find that our predictions remain accurate for all these choices of sb. The choice tanh luv] performs almost identically without the tanh activation. Interestingly. the choice outperforms the variant without the tanh and has more regular decay of the test loss.
tanh u²
In Figure 3b apply Theorem to predict the asymptotic squared test loss: 1 0- at each iteration. While this function is not PI((), as required by the theorem, the asymptotic predictions still align well with simulations Extending our technical results hold formally in such scenarios is an interesting direction for future work.
24

---

sev (daeus)
2
suel (thonry)
2
tank taxl (ai)))
0.010
[Mc]] (therry)
(theory)
tank [ue] term (thoery)
(sin)
Mean Squared Linur
Test Eeror
Lauh ?
u (sin)
0.008
Couk oheoyy
(sm)
tack Gim)
10
0.006
0.004
0.002
3
Iterations
Iterations
(a) Predictions for which are not uniformly bounded
(b) Predictions the squared test loss
Figure 3: Here, we n = 250, d 2000, = 0.1, 0: Bernoulli(0.0)) and select À to minimize the predicted asymptotic loss. Plus marks denote the median over 100 trials, and the shaded region indicates the interquartile range. Left: Predictions and simulations for weighting functions which are not uniformly bounded. Right: Predictions and simulations for the squared error u ② 00
25

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Theorems D and direecty reflect the claims from the abstract/introouutton..
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We provide detailed discussion of the strength of our technical assumptions after the statement of Theorem [ and in Appendix D and further elaborate these areas for improvement in the conclusion.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
26

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well vell-specification, asymptotic approxinations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recogmize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We formally state and discuss Assumptions and and we provide complete proof of our results in the Appendix of the supplemental material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all details of hyperparameters used to run simulations the corre- sponding figure caption. We also detail the exact method used to compute our asymptotic predictions in Appendix C and provide accompanying code for the simulations.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
27

---

If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide an accompanying file with code for computing the asymptotic predictions and running simulations with high-dimensional Gaussian data.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
pwt.huooommisiiooooocoocc for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
28

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes)
Justification: For each simulation, hyperparameter choices (and the method for choosing the ridge parameter A) are specified either in the caption or in the accompanying discussion in the text.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report the median and interquartile range for all simulations over 100 independent trials. These metrics are chosen due the asymmetry the distribution across trials (e.g., the mean minus the standard deviation might be negative in low-noise settings).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [NA]
Due to the small-scale of our included simulations, we do not include this information.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
29

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research presented in this work abides by the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work provides statistical analysis for a generic family algorithms in an abstract setup, and we do not believe our work has immediate social impacts an immediate path toward such impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: The paper poses such risks.
30

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperssithcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The provided code includes all necessary information for running simulations.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
31

---

Answer:
[NA]
Justification:
The paper does not involve research with human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve research with human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "uECyjesCix", "reviewer": "Reviewer_yJHL", "summary": "In this paper, the authors propose theoretical analysis of the high-dimensional dynamics of the reweighted least-squares methods in the context of \"linear diagonal networks\". \nThe general algorithm is given in Equation (4) and includes alternating minimization (AM), reparameterized IRLS, and linear recursive feature machines (lin-RFM) as special cases with different choices of reweighting function, see Table 1.\nIn the high-dimensional regime where the observation dimension $d$ and sample size $n$ are both large and comparable, the authors provide, in Theorem 1, distributional characterization of the algorithm behavior for a fixed number of iteration $T$ as $n,d \\to \\infty$ at the same pace.\nThe result is then extended the multi-variate setting in Theorem 2.\nSimulations results are provided to support the proposed theoretical analyses.", "questions": "* line 77 contribution: I suggest that the authors could consider adding pointers to precise theoretical/simulation results when talking about contributions. For example, referring to Theorem 1 for the first contribution.\n* line 91, perhaps add a paragraph on the related technique tool of Convex Gaussian Min-Max Theorem beyond [7], what is the technical challenge here, and whether the obtained technical results are of more general interest?\n* line 172-175: the assumption of $nT > d$ is an artifact of the proof, or very intrinsic to the obtained theoretical results? I somehow feel this is a bit unrealistic, or at least, a regime that is not of great interest to study. Could the authors comment on this?\nAlso, this should be stated explicitly in the theorem.\n* the authors mentioned a few time the keyword \"feature learning\", but after reading the paper I did not find any in-depth discussion on this. Perhaps having an additional remark to make the connection? \n\nminor: \n* line 197: of the algorithm in (4) or in Equation (4).", "limitations": "I do not see any potential negative social impact of this work.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is in general in good shape.\nIt is well written and easy to follow.\nThe proposed analysis is based on Convex Gaussian Min-Max Theorem and improves previous efforts.", "weaknesses": "It seems that the problem under study is of interest, but it would be great if the author would better motivate the study and setting, e.g., by providing some take-home messages.\nI have some detailed comments below", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and constructive feedback. We address your individual questions and concerns below.\n\n*(1) “It seems that the problem under study is of interest, but it would be great if the author would better motivate the study and setting, e.g., by providing some take-home messages.”*\n\nThanks for the feedback - please see part (b) of our global response for the main take-home messages we hope to convey. In particular, we would like to emphasize the flexibility of our framework in characterizing the test error of several previously proposed algorithms which have been shown to perform well for sparse recovery tasks or to connect with feature learning in linear neural nets. Our results also provide theoretical support for the use of “weight sharing” to learn signals with additional group-sparse structure. We plan to emphasize these points more explicitly in the final discussion section.\n\n*(2) “line 77 contribution: I suggest that the authors could consider adding pointers to precise theoretical/simulation results when talking about contributions. For example, referring to Theorem 1 for the first contribution.”*\n\nThanks for the suggestion - in the revision, we will include the appropriate references.\n\n*(3) “line 91, perhaps add a paragraph on the related technique tool of Convex Gaussian Min-Max Theorem beyond [7], what is the technical challenge here, and whether the obtained technical results are of more general interest?”*\n\nWhile the CGMT has been used extensively in recent literature to analyze estimation problems, a few new technical challenges arise when applying the result to an entire optimization trajectory, since each step relies on the previous iterate. The proof in [CPT23] relies on the fact that the “auxiliary optimization” obtained via the CGMT can be written in terms of a scalar functional of the previous iterate. However, the LDNN parameterization does not fall into the class of problems covered by this analysis technique. To deal with this, we instead obtain a distributional characterization of the iterates at each step. We do believe the proof technique for obtaining Wasserstein convergence guarantees could extend to studying other iterative algorithms in the sample-split setting. We state a few of these points in the last paragraph of the related work (Lines 141-147), but we will plan on adding a short description of the proof aspects we use which may be of more general interest.\n\n*(4) “line 172-175: the assumption of nT > d is an artifact of the proof, or very intrinsic to the obtained theoretical results? I somehow feel this is a bit unrealistic, or at least, a regime that is not of great interest to study. Could the authors comment on this? Also, this should be stated explicitly in the theorem.”*\n\nThe setting that we emphasize here is actually the scenario where $nT < d$, so the total amount of observed data is smaller than the dimensionality. In such settings, effective algorithms must be able to take advantage of the signal sparsity to achieve good performance. While we focus on this more interesting regime, our technical results do not depend on this assumption and in fact hold for any constant value of $\\kappa = d/n$.\n\n*(5) “the authors mentioned a few time the keyword \"feature learning\", but after reading the paper I did not find any in-depth discussion on this. Perhaps having an additional remark to make the connection?”*\n\nThanks for the feedback - please see part (b) of our global response about this point. We will be sure to add an additional remark about this in revision.\n\n*(6) “line 197: of the algorithm in (4) or in Equation (4).”*\n\nThanks - we will fix this in the revision.\n\nReferences:\n\n[CPT23] Chandrasekher et al. “Sharp global convergence guarantees for iterative nonconvex optimization with random data,” The Annals of Statistics, 2023.", "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": "dF0ak9ifaa", "reviewer": "Reviewer_cDEg", "summary": "The paper considers a general algorithm of the form shown in Lines 66-67, as it includes several interesting methods as special cases.\n\nThe main result is Theorem 1, which proves asymptotic convergence in probability of any function $g\\in PL(2)$ to the corresponding expectation. One important special case of this result is the convergence of the test errors of the iterates.\n\nSynthetic experiments show that the errors predicted by Theorem 1 align well with the actual performance of the algorithms.", "questions": "- Is it necessary to define pseudo-Lipschitz of order $k$? It seems that for most of the cases, $k$ is simply chosen to be $2$.\n- Lines 185-186: What is the precise relation between $\\alpha$ and $p$? It is not very clear why we need two different symbols.", "limitations": "See above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, easy to follow, and well-motivated.\n- One interesting contribution is that the paper understands a few different lines of research and is able to unify them using a general algorithm, where the reweighting function can be chosen to implement a specific type of method.\n- Theorem 1 is a powerful result as it holds for any function in the specific class $PL(2)$ or for any continuous and bounded function.\n- The proofs do not seem to be trivial to me. While I am not familiar with the proof techniques, the paper illustrates the technical points in the paragraph of Line 240, which makes the high-level proof idea clear and provides cues for interested readers to go deeper into the technical details.", "weaknesses": "- The first weakness is the assumption that each batch of data is independent. Personally I feel that this is a very strict assumption as I can not imagine a situation where it is satisfied. For example, almost all optimization algorithms (say SGD) would also reuse the previously seen data. Perhaps it is difficult to remove this assumption, but I think more discussions about it could be beneficial. Specifically:\n   - Could the authors please provide experiments similar to Figure 1 and allow the algorithm to reuse the data? (it is not very clear to me whether Figure 1 takes fresh batches or not.)\n   - Could authors please comment on the technical hurdles if the independence assumption is not available? \n   - Prior analysis on IRLS does not have this assumption, see for example: Theorem 2 of \"On the Convergence of IRLS and Its Variants in Outlier-Robust Estimation\" and in particular Figure 2a shows the error prediction; Theorem 1 of \"Globally-convergent Iteratively Reweighted Least Squares for Robust Regression Problems\"; Theorem 3.2 of \"Iteratively Reweighted Least Squares for Basis\nPursuit with Global Linear Convergence Rate\". Would their analysis be useful in any way to remove the independence assumption?\n\n- The second weakness is that the experimental evaluations in their current form are not very deep. For example:\n   - Is the error bound prediction accurate for $\\ell_2$ squared test loss? The paper commented that $\\ell_2$ squared loss does not belong to $PL(2)$, but having this experiment is important as $\\ell_2$ squared losses are more commonly seen in my opinion and this experiment can verify whether the assumption of $PL(2)$ is too strong.\n   - It seems a little bit weird that alternating minimization exhibits non-monotonic errors in Figure 1b (purple). Is that because these are test errors, or because the $\\ell_1$ loss is a bit different from the training loss (objective function)? How about $\\ell_2$ losses or the prediction (residual) errors on test (training) data? All these can be easily verified by experiments, and having them would give a better understanding of the behavior of alternating minimization, and it would also make the analysis more complete and convincing in my humble opinion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and constructive feedback. We address your individual questions and concerns below.\n\n*(1) “Could the authors please provide experiments similar to Figure 1 and allow the algorithm to reuse the data? (it is not very clear to me whether Figure 1 takes fresh batches or not.)”*\n\nWe obtain all figures in the paper by performing the iteration in equation 4, which has a fresh batch at each step. We will be sure to specify this explicitly in each figure caption. In general, we do not expect the same exact behavior in the “full-batch” setting (indeed, the total amount of data observed is smaller by a factor of T), but a similar behavior holds when samples are randomly chosen at each iteration (with repetition allowed). Please see part (a) of our global response for more discussion and supplemental simulations for this point. \n\n*(2) “Could authors please comment on the technical hurdles if the independence assumption is not available?”*\n\nPlease see part (a) of our global response for discussion of this point.\n\n*(3) “Prior analysis on IRLS does not have this assumption…”*\n\nWe agree that prior non-asymptotic convergence results for IRLS do not have this assumption. However, we emphasize that the type of result we achieve is substantially different in flavor and relies on a very different set of technical tools. Unlike these works, we aim to provide an exact asymptotic characterization of the distribution of the iterates in the high-dimensional limit. While the tools required for this necessitate stronger assumptions on the data (Gaussian) and algorithm (online/sample-split setting), the resultant guarantees are much stronger and can apply to a wider range of problem settings and algorithms.\n\n*(4) “Is the error bound prediction accurate for squared test loss?...”*\n\nWe have included simulations for the squared loss in Appendix D. Even though the squared loss is not PL(2), we do find that the error prediction from Theorem 1 seems to still be accurate in this case. So, even though our result doesn’t formally hold for squared loss, we still believe our theorem can still be useful for trajectory predictions in this case.\n\n*(5) “It seems a little bit weird that alternating minimization exhibits non-monotonic errors in Figure 1b…”*\n\nWe were also surprised by this behavior! We find that, rather than depending on the loss used ($\\ell_1$ or $\\ell_2$), this seems to depend on the noise level, with the non-monotonicity appearing more prominently in the low-noise regime (Fig. 1b). We agree that this deserves more mention in the main paper, so we will be sure to mention this in Section 3.1.\n\n*(6) “Is it necessary to define pseudo-Lipschitz of order k?”*\n\nIt is true that we only use the choice k=2 in our analysis – we will simplify the definition accordingly.\n\n*(7) “Lines 185-186: What is the precise relation between $\\alpha$ and $p$?*\n\nThanks for the question – this reweighting function is equivalent to the one used by the IRLS-p algorithm from [MF12] with the choice $p=2-4\\alpha$. The analysis in [DDFG10] also makes an explicit connection between IRLS updates of this form and $\\ell_p$-minimization for $0<p\\leq1$. \n\nReferences:\n\n[MF12] Mohan and Fazel, \"Iterative reweighted algorithms for matrix rank minimization,\" JMLR 2012.\n\n[DDFG10] Daubechies et al., \"Iteratively reweighted least squares minimization for sparse recovery,\" *Communications on Pure and Applied Mathematics,* 2010.", "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": "xwehlVLJOz", "reviewer": "Reviewer_kkyZ", "summary": "This paper studies the trajectory of a certain iterative scheme in the high-dimensional limit. \nIn each step, there're two update steps on a fresh batch: 1) regularized least squares; 2) element-wise nonlinear transformation.\nThis type of iteration is motivated and encompasses several interesting iterations such as IRLS, AGOP, etc. \n\nThe result is a precise characterization of the empirical distribution of the iterates at any finite time in the proportional limit. \nThis allows one to compute many interesting summary statistics.", "questions": "Minor comments: \n\n1. In the equation between line 66/67, the fresh dataset $\\boldsymbol{y}^{(t)}, \\boldsymbol{X}^{(t)}$ hasn't been defined (correct me otherwise). \n\n1. Line 49, typo: is --> are\n\n1. Line 164, typo: $t = 0, 1, \\dots, T$\n\n1. Equation between line 165/166, why not add superscript $(t)$ to $\\boldsymbol{\\epsilon}$? \nAfter all, the noise is also fresh and is not shared by other batches. \n\n1. Regarding boundedness of $\\boldsymbol{\\theta}^*$, as mentioned, it's only used in Lemma 3. \nOver there, I wonder if it's possible to remove this by truncating $\\boldsymbol{\\theta}^*$ entry-wise at a large constant $K$, doing the proof (in particular the argument in line 682-692), then sending $K\\to\\infty$. \n(To clarify, this is a minor assumption and removing it requires technical work that may not worth it.\nSo I don't mind if the authors leave it as is.) \n\n1. Line 246-249, the authors meant to discuss a second difference from prior works. \nWhat's done in this paper is discussed but what's done in prior work is not. \nSo the comparison is unclear. \n\n1. There are two notation sections that are almost identical. I suggest remove one of them. \n\n1. Line 462, what is a \"proper function\"? Do you mean a \"proper convex function\"?\n\n1. Line 466, \"restate\": please put a reference for CGMT. \n\n1. Equation between line 577/578, $V$ should be $V_t$. \n\n1. Equation between line 579/580, $v_i$ should be $v_i^{(t)}$, $V$ should be $V_t$. \n\n1. Line 580, typo: \"satisfies\".", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Though the iteration studied in this paper may look exotic at the first glance, it is well motivated (see the first two pages). \nThe proof uses, unsurprisingly, CGMT since the problem is given in a form that's ready for it. \nThat said, working out the details is not a trivial task. \nI went through the whole argument in the appendix and got the impression that everything is treated in a fairly rigorous way. \nIn particular, all technicalities associated with CGMT such as justifying the exchange of min / max and proving uniform convergence have been properly addressed. \nIn doing so, the authors have invoked various tricks / auxiliary results commonly used in this literature.", "weaknesses": "As the authors have already discussed, one obvious drawback of CGMT is that it needs sample splitting (which may make the iteration under consideration incompatible with its motivating special cases). \nI actually think that CGMT has the potential to handle repeated batches, or even full batch, though the resulting state evolution is likely to be much less interpretable. \nSo I don't really object to the sample splitting assumption. \n\nAlso, CGMT is capable of offering non-asymptotic guarantees. \nI wonder if the authors are interested in doing so (no pressure).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review of our technical results and constructive feedback. We agree that obtaining non-asymptotic guarantees for this problem would be a very interesting technical challenge for future work. Here, we rely explicitly on a distributional characterization of the iterates, rather than obtaining a state evolution over deterministic scalars. This makes obtaining non-asymptotic guarantees more challenging (although likely possible) after the first 1-2 iterations. \n\nWe address your individual questions below. Thanks for catching some typos and the missing CGMT reference - we will fix all of these in our revision. Regarding a few individual questions:\n\n- Boundedness of $\\theta^*$: Thanks for this suggestion! Yes, we do believe this kind of truncation argument should allow us to weaken the boundedness requirement on $\\theta^*$ to some weaker regularity condition on the distribution. However, we seem to run into some issues when trying to apply a truncation result recursively for each iteration (since we end up needing convergence of the empirical distribution of $(v, \\theta^*)$ in a higher order Wasserstein metric). We will continue working on this and update the manuscript if we find a way around it!\n- Line 246-249: Thanks for the catch. The prior work [CLOT21] assumes convergence of the initialization in $W_4$ and proves convergence of the estimator in $W_3$, which would not work for a recursive application of the result. We will add this in the revision.\n\nReferences:\n\n[CLOT21] Chang et al., “Provable benefits of overparameterization in model compression: From double descent to pruning neural networks,” AAAI 2021.", "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:
We thank all the reviewers for their time and helpful feedback. We will carefully consider and incorporate all the suggestions into the next revision of this paper.

We address here some of the main comments:

**(a) Sample-splitting assumption:** We agree that the sample-splitting assumption deviates from the existing literature on IRLS-style algorithms and linear diagonal neural networks. The primary technical reason for this assumption lies in application of the Convex Gaussian Min-Max Theorem after the first step of the optimization algorithm. In particular, the CGMT (cf. Theorem 3) requires an objective function where the only dependence on the Gaussian random matrix is in a bilinear term. If the previous iterate also depends on the same data matrix, this condition is violated. However, we believe that this setting is still of practical interest, e.g., in online/streaming data settings where data is only available batch by batch. 

Moreover, this assumption allows us to obtain much more precise results than previous works which only obtain upper bounds on the error. We emphasize also that this assumption is somewhat standard when trying to obtain rigorous asymptotic characterization of optimization trajectories using similar techniques (e.g., reference [CPT23]). We leave relaxations of this assumption to future work, since the resulting state equations would likely be much more complicated (as mentioned by Reviewer kkyZ), and the proofs may require substantial new machinery. We will be sure to include some of these remarks as discussion points in the revised paper.


To further assess the strength of this assumption, we have attached additional simulations which compare three types of batch selection (under the same hyperparameter choices as in Figure 1a): 
- Fresh batches at each iteration: this is the setting we consider in our paper and for which our theoretical results are derived
- Randomly sampled batches (with possible repeats): this corresponds to choosing $n$ samples at each iteration from a global pool of $nT$ samples, with possible repetition of data across iterations.
- Same batch: using the same $n$ samples for each iteration

We find that the theoretical results derived for Setting 1 seem to remain accurate for Setting 2, even though the batches are no longer independent (formalizing this would be an interesting technical problem for future work). However, as expected, the predictions are too optimistic for Setting 3 (which only uses $n$ total samples during training, rather than $nT$).


**(b) Motivation and connections to feature learning:** We also plan to incorporate the helpful suggestions of Reviewer yJHL to improve the motivation and contextualization of our results. Some important take-home messages we hope to highlight for our results include (i) a rigorous framework for comparing several existing (and new) algorithms from the perspective of generalization error, (ii) favorable guarantees on the test error within a constant number of iterations (which is more optimistic than the convergence results common in the literature), and (iii) provable benefits of weight-sharing in LDNNs in the presence of structured sparsity.

As a further clarification, our use of the phrase “feature learning” refers to the ability of these algorithms to learn LDNN parameters which effectively recover sparse or group-sparse signals. Note that in the case of linear models, learning “good” features amounts to learning which subset of coordinates in the input are important, so feature learning in linear diagonal networks is equivalent to learning the low-dimensional structure of the true signal (e.g., sparsity or group-sparsity structure). See also the discussion in [RMD24] for more discussion on this connection and how it was used to construct the lin-RFM algorithm (which we analyze). We will be sure to add this connection more explicitly after lines 39-40. 

References:

[CPT23] Chandrasekher et al. “Sharp global convergence guarantees for iterative nonconvex optimization with random data,” The Annals of Statistics, 2023.

[RMD24] Radhakrishnan et al. “Linear Recursive Feature Machines provably recover low-rank matrices,” arXiV preprint, 2024.

Meta-review: Three reviewers evaluated the paper, and their overall assessment is positive. I concur with their evaluation and believe the paper offers a strong contribution with compelling results. Some reviewers pointed out a potential limitation related to the sample-splitting assumption. In response, the authors conducted additional simulations using various batch selection methods to address this concern. Another reviewer noted the connection between the presented work and feature learning. The authors welcomed this feedback and have committed to clarifying this connection in the paper.
Decision: Accept (poster)