Precise asymptotics of reweighted least-squares algorithms for linear diagonal networks

Authors: chiraag  kaushik
justin  romberg
vidya  muthukumar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:29:15

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.

Main paper: 
 introduction: 
 Many high-dimensional machine learning and signal processing tasks rely on solving optimization problems with regularizers that explicitly enforce certain structure on the learned parameters. The traditional formulation for such tasks involves a regularized empirical risk minimization (ERM) problem of the form minEQUATIONwhere L(•) is 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 sparsityseeking solution; in such cases, the regularizer is 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-parameterization" of θ into the entrywise product of two factors u ⊙ v. While the resulting minimization problem is non-convex, this 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 traditional sparse recovery algorithms [13, 25] . For example, rather than solving the convex, but nonsmooth LASSO (where L is the squared loss and R is the ℓ 1 norm), the Hadamard reparameterization yields the following non-convex and smooth formulation:EQUATIONIn the case where the regression function is linear in θ, solving (2) is equivalent to learning a function of the form f u,v (x) = ⟨x, u ⊙ v⟩ = ⟨diag(v)x, u⟩, which can be thought of as a one hidden layer neural network with linear activation function and inner weight matrix diag(v). In this context, this linear diagonal neural network (LDNN) architecture has also been studied as an illustrative case study to improve our understanding of how neural networks perform iterative "feature learning" to leverage low-dimensional structure in high-dimensional settings [34, 23] . We note here that, in the linear model case, feature learning is equivalent to learning which subset of the input'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/LDNN form in (2) is to consider the change of variable v i → √ η i and u i → θi √ ηi [25] . This yields the following optimization problem, which is jointly convex in η and θ:EQUATIONAfter solving the minimum over η explicitly, the second term becomes exactly λ∥θ∥ 1 , and we recover the Lasso objective. This is a special case of the so-called "eta-trick" [1] , which can be used to write many common sparsity-inducing penalties as the minimization of a quadratic functional of θ.A 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 (θ, η) optimization in (3) can also be leveraged to construct algorithms based on classical sparse recovery techniques. In particular, alternating minimization over θ and η in (3) yields the popular iteratively reweighted least-squares (IRLS) algorithm [11, 9] . Translating these updates to the equivalent updates on (u, v), we obtain an iterative least-squares algorithm for LDNNs, which alternately sets v (t+1) = |u (t) ⊙ v (t) | and performs a weighted least squares update on u. This particular form of reparameterized IRLS was generalized in [28] to a larger family of iterative least-squares algorithms under the name of linear recursive feature machines (lin-RFM).While several methods for learning Hadamard/LDNN 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, 18, 7] .In this paper, we provide a precise analysis of a general family of iterative algorithms for learning LDNNs that take the form u (t+1) = arg min• How do different algorithm choices compare (in terms of convergence and signal recovery) in the high-dimensional regime?• How many iterations does it take common algorithms to find statistically favorable solutions?• What is the effect of model architecture in LDNNs? Does leveraging group structure provably improve sample complexity when the ground-truth signal is group-sparse?We define a general class of algorithms which learns LDNNs by alternately performing least-squares and reweighting steps in a sample-split/batched 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 1).(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, a 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 θ * has groupsparse structure (Theorem 2). Our results show that using a 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.IRLS and the η-trick: The reformulation of non-smooth regularizers in terms of quadratic variational forms (the "η-trick") has been studied in various early works in computer vision and robust statistics [12, 4] . Further analysis and examples of sparsity-promoting norms are provided in [20, 2] , and [25] provides a characterization of when a regularizer admits a variational form of this type. The resultant optimization algorithm is iteratively-reweighted least-squares (IRLS), a popular technique for compressive sensing and sparse recovery [11, 9] . These works also consider IRLS algorithms corresponding to ℓ p -norm regularization for 0 < p < 1; in this case, the minimization is no longer convex, but [11] shows that such methods can find sparse solutions with fast local convergence rate. The family of algorithms we consider includes a reparameterized version of each of these IRLS algorithms, but unlike these prior works, we consider a batched setting and the high-dimensional asymptotic regime. Moreover, our results apply to other algorithms which may not be easily expressed as resulting from the η-trick.The reparameterization of θ into the product of factors u ⊙ v has been considered in a variety of recent works. The authors of [31, 35] show that early-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 a minimizer of a certain sparsity-inducing functional that depends on the initialization [34, 10, 23] .The connection between the LASSO (as well as some non-convex ℓ q 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 η-trick and showing that saddle points are strict (escapable). These insights lead to global convergence guarantees and a smooth bi-level optimization scheme [25, 26] for non-smooth structured optimization problems that was shown to perform competitively with state-of-the-art solvers. The non-convex landscape of such formulations is further explored in [15] , where it is shown that for a 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 properties of the fixed-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 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/sample-splitting 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 a "grouped" Hadamard parameterization, which we consider in Section 4. This corresponds to a LDNN with groups of tied weights in the hidden layer. Early stopped gradient flow/descent for this type of architecture was shown in [17] to achieve sample-complexity scaling with the number of non-zero groups (rather than the overall sparsity). The non-convex landscape for this grouped architecture is studied in [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 a technical level, our work provides a 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 to obtain a precise characterization of non-convex optimization problems. However, since the Hadamard parameterization is a re-parameterization of the actual estimator of interest (θ := u ⊙ v), the results of [7] 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 matrix sensing are studied in the works [6, 19] . 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 a single sensing vector interacting with the product of the two parameters). 
background and formulation: 
 Notation: The ones vector of dimension d is denoted as 1 d . We denote the element-wise multiplication (Hadamard product) of two vectors x and y as x ⊙ y. Element-wise division of two vectors is denoted as x y . We say a function f :R p → R is pseudo-Lipschitz of order 2 if, for all x, y ∈ R p , |f (x) -f (y)| ≤ C(1 + ∥x∥ 2 + ∥y∥ 2 )∥x -y∥ 2for some constant C > 0. The set of such functions is denoted by PL (2) .Convergence in probability of a sequence of random variables X d to a random variable X is denoted by X d P → X. Convergence in Wasserstein-2 distance of a sequence of probability distributions ν d to a limiting distribution ν is denoted as ν d W2 → ν, and this fact is equivalent to the statement E X∼ν d g(X) → E X∼ν g(X) for all g ∈ PL(2) [3] . If the ν d are random probability measures, we say that ν d W2 → ν if the same convergence holds in probability, i.e., E X∼ν d g(X)P → E X∼ν g(X) for all g ∈ PL(2). The empirical distribution of a vector z ∈ R d is defined as1 d d i=1 δ(z i ), where δ(z i )is the Dirac delta distribution centered at z i .Formulation: We consider a batched noisy linear model where, at each time t = 0, 1, . . . , T , a user has access to an independent batch of data (X (t) , y (t) ) ∈ R n×d × R n satisfyingy (t) = 1 √ d X (t) θ * + ϵ (t) .Above, θ * ∈ R d is an unknown signal, X (t) has i.i.d. standard Gaussian entries, and ϵ (t) ∼ N (0, σ 2 I n ) is i.i.d. noise in the measurements. Given an initial weight vector v (0) ∈ R d , we are interested in the behavior of iterative algorithms of the formEQUATIONTable 1 : Some algorithms taking the form (4) Algorithm Reweighting functionAlternating minimization (AM) [13] ψ(u, v) = u Reparameterized IRLS [9, 11, 28] ψ(u, v) = (u 2 v 2 + ϵ) α Linear recursive feature machines (lin-RFM) [28] ψ(u, v) = ϕ(u 2 v 2 )where ψ : R × R → R is a "reweighting" function that acts entry-wise on (u (t) , v (t) ) and λ > 0 is a hyperparameter governing the strength of the regularization. We we will study the behavior of the iterates u (t) , v (t) in the high-dimensional limit where n and d both approach infinity with fixed ratio d n = κ. Since our primary interest is to reveal the feature learning capabilities of such algorithms when θ * is a high-dimensional signal with low-dimensional structure, we will typically focus on the regime where κ > T , where T is the number of total iterations. This ensures that the total number of observed samples nT is smaller than the ambient dimension d.Before proceeding to our main results, we note that this formulation encompasses a wide variety of classical and modern algorithms (summarized in Table 1 ):• Alternating minimization: One perspective on this algorithm is to consider it as a way to perform alternating minimization on the non-convex loss functionL(u, v) = 1 n y - 1 √ d X(u ⊙ v) 2 2 + λ d ∥u∥ 2 2 + λ d ∥v∥ 2 2 .Using the fact that the loss function is symmetric in u and v, choosing ψ(u, v) = u recovers the mini-batched alternating minimization algorithm for this loss. In other words, ψ 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 [11, 21] can be reparameterized in the form of (4) ψ(u, v) = (u 2 v 2 + ϵ) α , where different choices of α correspond to different ℓ p penalties. In particular, the choice p = 2 -4α corresponds to the IRLS-p algorithm of [21] .• Lin-RFM [28] : Generalizing the reparameterized IRLS update, the authors of [28] propose the choice ψ(u, v) = ϕ(u 2 v 2 ) for some continuous function ϕ : R → R + . Here, the quantity u 2 v 2 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 is to understand statistical properties of the iterates for different choices of ψ, and in particular how the test error 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.* i ∈ R b for i = 1, . . . ,τ t+1 , β t+1 = arg max τ ≥0 min β≥0    τ σ 2 β + τ β(1 -κ) -τ 2 + τ λ E (V ,Θ)∼Πt   1 b b j=1 Θ 2 j + β 2 κ τ V 2 j + βλ      Q t+1 = τ t+1 V ⊙ (Θ + β t+1 G t √ κ) τ t+1 V ⊙2 + β t+1 λ1 b, (entry-wise division)EQUATIONHere, G t i.i.d.∼ N (0, I b ). Then, we have the following result, which is proved in Appendix A. Theorem 2. [Generalization of Theorem 1 for b ≥ 1] Under the assumptions above, for any t ≥ 0 and any function g : (R b ) 3 → R such that g ∈ PL(2) or g is bounded and continuous, we have1 M M i=1 g(u (t+1) i , v (t) i , θ * i ) P → E[g(Q t+1 , V , Θ)].Given a reweighting function ψ, Theorem 2 characterizes the distribution of the factors (blocks) of the iterates. Hence, by choosing g(u, v, θ) = |u ⊙ v -θ|, we can predict the exact limiting test error for this family of algorithms.Computing these theoretical predictions reveals that choosing ψ in a group-aware way can lead to significant performance improvements compared to coordinate-wise reweighting. In Figure 2 , we fix σ = 0.1, n = 500, d = 4000, and set the overall expected sparsity level of θ * as in Figure 1 . We compare the performance of Algorithm 4 for a "group-blind" (ψ gb ) and "group-aware" (ψ ga ) choice of reweighting function:• ψ gb (u, v) = tanh |u ⊙ v| -note this is identical to one of the reweightings considered in Section 3.1.• ψ ga (u, v) = 1 b b j=1 tanh |u j v j | 1 bThe theoretical predictions align with simulations and show a notable improvement in performance when using the group-aware scheme with b > 1. Moreover, as the group size b increases, the performance of ψ gb remains approximately the same, indicating that it is not able to take adapt to the group-structure. By contrast, using ψ ga 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. 
a precise characterization of the iterates: 
 In this section, we provide a precise characterization of the iterates of the algorithm in Equation 4with 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 v 0 and the target signal θ * :Assumption 1. The empirical distribution of the entries of v (0) and θ * converges in W 2 distance to some joint distribution Π 0 , i.e., 1EQUATIONi ̸ = 0 for all i and θ * has bounded entries almost surely.Here, the requirement of empirical distribution convergence is easily satisfied by common choices of v (0) , including the ones vector and i.i.d. Gaussian entries. For a typical sparse regression setup, we might, for example, consider the Π 0 induced by choosing v (0) = 1 d and letting θ * have i.i.d. entries that equal 0 with certain probability. The requirement that θ * has bounded entries appears to be an artifact of 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 θ * has entries from distributions which are not bounded almost surely (e.g., Gaussian entries).Secondly, we define the set of reweighting functions ψ for which our result will apply. Assumption 2. The reweighting function ψ : R × R → R satisfies the following:1. If U, V are random variables such that U, V ̸ = 0 with probability 1, then ψ(U, V ) ̸ = 0 with probability 1.2. ψ is continuous and bounded or ψ 2 is pseudo-Lipschitz of order 2.This family allows us to consider many of the choices of ψ discussed in the previous section, including ψ(u, v) = u (AM on linear diagonal networks), ψ(u, v) = |uv| (lin-RFM and IRLS), ψ(u, v) = ϕ(u 2 v 2 ) for bounded ϕ (lin-RFM). We note that this does not include some choices of ψ which apply more "aggressive" weighting, such as ψ(u, v) = |uv|. Nevertheless, we can apply our theoretical predictions for these choices of ψ after passing the weights through a 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 is violated.Our results are stated in terms of the following iteration, for t ≥ 0:EQUATIONwhereG t i.i.d.∼ N (0, 1). In words, given a probability distribution Π t over R × R, τ t+1 and β t+1 are scalars computed as the unique1 solutions to a deterministic optimization problem (this can be solved easily by studying the optimality conditions, as shown in Appendix C). Then, Q t+1 is defined as a random variable that is a function of (V, Θ) ∼ Π t and G t ∼ N (0, 1). Lastly, Π t+1 is defined as the joint distribution of ψ(Q t+1 , V ) and Θ.Given this iteration, we obtain the following result, which is proved in Appendix A: Theorem 1. Suppose Assumptions 1 and 2 are satisfied. Then, for any t ≥ 0 and any function g : R 3 → R such that g ∈ PL(2) or g is bounded and continuous, we have1 d d i=1 g(u (t+1) i , v (t) i , θ * i ) P → E[g(Q t+1 , V, Θ)],where the expectation is over the independent random variables (V, Θ) ∼ Π t and G t ∼ N (0, 1).The limit in this theorem should be interpreted as being the limit as n, d → ∞ with their ratio κ = d n held as a constant. Applying the above theorem for each t ≥ 0, we can get precise asymptotic predictions for a wide variety of test functions of the iterates. One example of particular interest is the test error, which we measure as the normalized ℓ 1 distance between u (t+1) ⊙ v (t) and θ * , corresponding to g(u, v, θ) = |uv -θ| (we provide a proof that this is PL(2) in Proposition 1 in Appendix B). We note that the limiting expectation can be computed via simple Monte Carlo simulation of a scalar random variable.From a 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 a similar distributional characterization for the solution to least-squares with anisotropic covariates (where the "weights" v are the square root of the eigenvalues of the data covariance) [8] . However, while [8] assume that the eigenvalues are uniformly bounded by constants, this is not a reasonable assumption in our setting, since many common choices of ψ are not bounded and hence Test Error corresponds to the classical IRLS weighting from [11] , ψ = tanh |uv| is a version of lin-RFM, ψ = u corresponds to AM, and ψ = tanh u 2 is a new reweighting scheme we introduce. We note that the ψ which depend only on u can lead to oscillatory behavior in the test risk. v (t) is not necessarily bounded uniformly for t > 1. A second key difference is that we need to obtain a distributional characterization which can be applied recursively for all t ≥ 0. The analysis in [8] assumes convergence of the initialization in W 4 distance and proves convergence of the estimator in W 3 distance. However, to apply the result recursively in our setting, if assume that the empirical distribution of (v (0) , θ * ) converges in W k distance, then we need to show that after one iteration, the iterates also converge in W k distance (and not in any weaker sense).ψ = |uv| 1 2 (theory) ψ = |uv| 1 2 (sim) ψ = tanh |uv| (theory) ψ = tanh |uv| (sim) ψ = u (theory) ψ = u (sim) ψ = tanh u 2 (theory) ψ = tanh u 2 (sim) (a) σ = 0.1 1 2 3 4 5 6 7 8 Iterations 10 -4 10 -3 10 -2 ℓ1 Test Error ψ = |uv| 1 2 (theory) ψ = |uv| 1 2 (sim) ψ = tanh |uv| (theory) ψ = tanh |uv| (sim) ψ = u (theory) ψ = u (sim) ψ = tanh u 2 (theory) ψ = tanh u 2 (sim) (b) σ = 0.001To overcome these differences, we use a different technique to show distributional convergence of the iterates. Similar to the approach in [5] , we apply the CGMT to a 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 θ * has bounded entries and we obtain results for a slightly smaller family of test functions g (note, for example, that the squared loss g(u, v, θ) = (uv -θ) 2 is not PL(2)), we obtain a distributional convergence result that can be applied to a sequence of recursively defined least-squares problems which define the trajectory of an algorithm, rather than to a single optimization problem. Moreover, our simulations in Appendix D suggest that the predictions of Theorem 1 still often apply without these additional assumptions, including in the case of the squared loss, indicating that these additional assumptions could potentially be weakened with a more complicated analysis.In this subsection, we apply Theorem 1 to a sparse recovery setting and compare the asymptotic predictions to numerical simulations on high-dimensional Gaussian data. First, we consider a setting where n = 250, d = 2000, and θ * has Bernoulli(0.01) entries (so the expected sparsity level is E[s] = 20). We run Algorithm 4 with initialization v (0) = 1 d for four different choices of reweighting function and display the test error at each iteration (median over 100 trials) in Figure 1 . For each choice of reweighting function ψ, we choose the regularization parameter λ that minimizes the asymptotic test loss achieved within 8 iterations, and we plot the corresponding trajectory. As shown in the figure, the numerical simulations show excellent alignment with the asymptotic predictions even for this moderate choice of n and d.The asymptotic predictions show that this family of algorithms can find solutions with low test error within only a few iterations. Our results also reveal fine-grained differences in the convergence behavior of the different algorithms. For instance, more aggressive weightings ψ = tanh |uv| and ψ = tanh u 2 seem to find better solutions after several iterations. Interestingly, the weighting functions which depend only on u (like alternating minimization) sometimes display a 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 ϕ(u 2 ) 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 θ * is known to possess additional structure that can be leveraged during training. One commonly studied example of this is structured sparsity, or group sparsity, where θ * has many blocks which are zero. In this section, we generalize the results of Theorem 1 to the case where the reweighting function respects this additional structure in the signal, i.e., ψ acts on blocks of v, rather than on individual coordinates.Concretely, we consider the following modification to our formulation. Let b ≥ 1 be a constant and write R d as a product space over Here, the case b = 1 recovers the results of the previous section, but the case b > 1 allows us to study the interplay between signal structure and reweighting scheme in a more fine-grained way. For example, suppose θ * is known to be group-sparse, meaning that many of the factors {θ * i } M i=1 are zero. In this case, it might make sense for ψ to return a vector of the form ψ(uM = d b factors: R b × • • • × R b . Then, θ * , u (t) , v (t) ∈ R d can(t) i , v (t) i ) = α i 1 b , for some α i ∈ R that is chosen as a function u (t) i and v (t)i . This corresponds to a reweighting scheme which acts on blocks, rather than individual entries. Another way to motivate this "grouped reweighting" is to leverage the connection to the η-trick, as in (3) . In particular, the group Lasso problem can be written in the following variational form [2] :EQUATIONwhere the closed-form solution to the η minimization yields the classical group norm regularizerM i=1 ∥θ i ∥ 2 .Here, reparameterizing as α i → √ η i and u i → θi √ ηi gives rise naturally to the grouped Hadamard parameterization, withv i = α i 1 b .From the perspective of linear diagonal networks, such an 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] ), we consider an optimization approach that relies on alternate updates of u (t) and v (t) .We make the same technical assumptions as in Assumptions 1 and 2, with the natural modifications to accommodate b ≥ 1:1. Π 0 is the limit of the empirical distribution of factors of (v (0) , θ * ) and hence is a distribution over R b × R b . A straightforward extension of Theorem 1 yields the following generalization for the grouped algorithm, where V , Θ, G t , Q t+1 ∈ R b are now vector-valued random variables: For t ≥ 0, let 
conclusion: 
 In this paper, we derived a precise asymptotic characterization of the iterates of a family of algorithms for learning high-dimensional 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 a 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 "minibatch" case studied by [19] ). 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.Question: Does the paper provide open access to the data and code, with sufficient instructions to 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 NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) 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 a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.cc/public/guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on 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. If only a 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 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes] Justification: For each simulation, hyperparameter choices (and the method for choosing the ridge parameter λ) 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 a 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. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We report the median and interquartile range for all simulations over 100 independent trials. These metrics are chosen due to the asymmetry of 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, confidence intervals, or statistical significance tests, at 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. • It is 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 a 96% CI, if the hypothesis of Normality of errors is 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 in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [NA] Justification: Due to the small-scale of our included simulations, we do not include this information.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.Answer: [NA] Justification: The paper does not involve research with human subjects or crowdsourcing. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS 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. 
 
 Appendix: In this section, we provide the proofs of our main results.For convenience, we first restate the main notation that is used in our proofs.Notation The ones vector of dimension d is denoted 1 d . We write a ≲ b when a ≤ Cb for some sufficiently large constant C > 0 which does not depend on d. We denote the element-wise multiplication (Hadamard product) of two vectors x and y as x ⊙ y. Element-wise division of two vectors is denoted as x y . We use the shorthand(•) + = max(•, 0). A function f : R p → R is called pseudo-Lipschitz of order 2 if, for all x, y ∈ R p , |f (x) -f (y)| ≤ C(1 + ∥x∥ 2 + ∥y∥ 2 )∥x -y∥ 2for some constant C > 0. The set of such functions is denoted PL (2) .Convergence in probability of a sequence of random variables X d to a random variable X is denoted X d P → X. Convergence in Wasserstein-2 distance of a sequence of probability distributions ν d to a limiting distribution ν is denoted as ν d→ ν, and this fact is equivalent to the statement E X∼ν d g(X) → E X∼ν g(X) for all g ∈ PL(2) [3] . If the ν d are random probability measures, we say that ν d W2 → ν if the same convergence holds in probability, i.e., E X∼ν d g(X)P → E X∼ν g(X) for all g ∈ PL(2). The empirical distribution of a vector z ∈ R d is defined as1 d d i=1 δ(z i ), where δ(z i )is the Dirac delta distribution centered at z i . For any random variable (or group of 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 1. The Moreau envelope function of a lower semi-continuous, proper convex function ℓ : R p → R with step size τ is defined asM ℓ (x; τ ) = min y∈R p ℓ(y) + 1 2τ ∥y -x∥ 2 2 .The proximal (prox) operator of ℓ with step size τ , denoted prox ℓ (x, τ ) 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 m×n , g ∈ R m , h ∈ R n have i.i.d. N (0, 1) entries. Let S w ⊂ R n and S u ⊂ R m be compact, convex sets, and f : R n × R m → R be convex-concave on S w × S u . Define the following two min-max problems:Φ(G) := min w∈Sw max u∈Su u ⊤ Gw + f (w, u) ϕ(g, h) := min w∈Sw max u∈Su ∥w∥ 2 u ⊤ g + ∥u∥ 2 w ⊤ h + f (w, u)Then, for all c ∈ R and t > 0,P{|Φ(G) -c| > t} ≤ 2 P{|ϕ(g, h) -c| > t}A.2 Proof of Theorems 1 and 2Proof of Theorem 1. Assume that1 d d i=1 δ(v (t) i , θ * i ) W2→ Π t (note that this holds by assumption at t = 0; we will show later that it holds at time t + 1).First observe that convergence of the joint empirical distribution of (u (t+1) , v (t) , θ * ) to the joint distribution of (Q t+1 , V, Θ) in Wasserstein-2 distance implies that1 d d i=1 g(u (t+1) i , v (t) i , θ * i ) P → E[g(Q t+1 , V, Θ)],for any g ∈ PL(2) or which is bounded and continuous. This is because W 2 convergence implies convergence in expectation of any pseudo-Lipschitz function of order 2 [3, Lemma 5] and of any bounded continuous function (since W 2 convergence is stronger than weak convergence [33, Theorem 6.9]). Hence, it suffices to show thatEQUATIONwhere (V, Θ) ∼ Π t and Q t+1 is defined as in Eq. 5.Recall that the objective function for the update on u is given byu (t+1) = arg min u∈R d 1 n y (t) - 1 √ d X (t) (u ⊙ v (t) ) 2 2 + λ d ∥u∥ 2 2 .Rather than study this update directly, we first analyze a slightly more general problem (following the approach in [5] ). Let h : R 3 → R be a continuous test function with ∥∇ 2 h∥ 2 ≤ C and that satisfies one of the following:1. h is uniformly bounded. 2. h(u, v, θ) = u 2 .Then we consider the following problem (the dependence of X, y, ϵ, and v on t is dropped to simplify the notation):EQUATIONwhere µ ∈ [-µ * , µ * ] and µ * = λ C is chosen sufficiently small so that the objective function (scaled by d) is λ-strongly convex for all µ in this range. The case µ = 0 recovers the original problem of interest.Step 1: Convergence of the loss Rewriting this in terms of the error vector∆ := 1 √ d (u ⊙ v -θ * ), we have P 1 (µ) = min ∆∈R d 1 n ∥ϵ -X∆∥ 2 2 + λ d √ d∆ + θ * v 2 2 + µ d d i=1 h √ d∆ i + θ * i v i , v i , θ * i .In writing this, we use the fact that v i ̸ = 0 for all i with probability 1 (and the notation in the secondto-last term indicates entry-wise division). Now, using the identity ∥•∥ 2 2 = max q 2q ⊤ (•) -∥q∥ 2 2 , we can write this asEQUATIONNext, in Lemma 1, we show that there exist Euclidean balls B ∆ and B q , each of radius C 1 ∥v∥ ∞ such that, with probability approaching 1, we can constrain the feasible set to lie with these balls without changing the value of P 1 (µ), so we can studyEQUATIONwhere P 1 (µ) = P1 (µ) with probability tending to 1. We can therefore condition on this event for the remainder of the analysis without changing our asymptotic conclusions. Now, noting that this is in the correct form to apply Theorem 3, we define the auxiliary optimization problemEQUATIONEQUATIONwhere g ∈ R d and h ∈ R n have i.i.d. standard normal entries. By Theorem 3, for all δ > 0 and fixedP (µ) ∈ R, P{|P 1 (µ) -P (µ)| > δ} ≤ 2 P{|P 2 (µ) -P (µ)| > δ}.In particular, if we can find some P (µ) such that P 2 (µ) P → P (µ), then we can conclude also that P 1 (µ)P → P (µ).To accomplish this, we next perform a series of simplifications to P 2 (µ) which will later help us characterize its asymptotic behavior. First, we can decouple the optimization over q into its norm and direction, and the latter can be solved explicitly. Letting τ = ∥q∥ 2 , this yieldsEQUATIONwhere R := C 1 ∥v∥ ∞ . Next, note that ϵ and h are independent Gaussian vectors and hence ϵ -∥∆∥ 2 h d = σ 2 + ∥∆∥ 2 2 h. So, we haveP 2 (µ) d = min ∆∈B∆ max 0≤τ ≤R 2τ ∥h∥ 2 √ n σ 2 + ∥∆∥ 2 2 - 2τ √ n g ⊤ ∆ -τ 2 + λ d √ d∆ + θ * v 2 2 + µ d d i=1 h √ d∆ i + θ * i v i , v i , θ * i .(15) Before proceeding further, we rewrite this in terms of a minimization over the variable u =EQUATIONHere,B u := u ∈ R d : 1 √ d (u ⊙ v -θ * ) 2≤ R . After this step, observe that the objective function is strongly concave in τ and strongly convex in u (the sum of the last two terms is strongly convex in u based on the assumption that h has bounded Hessian and µ is sufficiently small). Since the objective function is convex-concave over convex and compact sets, we can invoke Sion's minimax theorem to switch the min and max. Furthermore, we use the fact that √ x = min β>0 x 2β + β 2 to write this asEQUATIONHere, note that we can add the constraint on β without changing the solution since the optimal value of β will be obtained atσ 2 + 1 d ∥u ⊙ v -θ * ∥ 2 2 ∈ [σ, σ + R]for all feasible u. Next, we can explicitly solve the inner minimization over u. To do this, we first show in Lemma 2 that the optimal solution to the unconstrained minimization is strictly feasible for large enough C 1 , and hence, the unconstrained and constrained minimizations over u are equivalent. Next, observe that the unconstrained problem is separable over the indices, so we need only to solve the scalar problemEQUATIONCompleting the squares, we obtain that the above problem can be written in terms of the Moreau envelope (Definition 1) of the function ℓ(u) = λu2 + µh:1 d τ ξθ * 2 i β - τ βξ βg i √ κ -ξθ * i 2 + M λ(•) 2 +µh(•,vi,θ * i ) ξθ * i -βg i √ κ ξv i ; β 2ξτ v (t)2 i, where we have introduced the shorthand notation ξ = ∥h∥2 √ n . Substituting this into the expression for P 2 above, we obtainEQUATIONNow that the optimization has been fully "scalarized", we proceed by considering its asymptotic behavior. First, note that the partial minimization over u preserves the concavity/convexity in (τ, β).In Lemma 3, we prove that for any fixed τ and β, the objective function f d (τ, β) converges in probability toEQUATION)where the expectation is over (V, Θ) ∼ Π t and an independent G ∼ N (0, 1). We note here that Lemma 3 is the only place in our proof which requires the boundedness of the entries of θ * .Since f d (τ, β) is strongly concave in τ with parameter 1 for all feasible β, we can conclude that f (τ, β) is also strongly concave in τ with parameter 1. Directly taking a derivative with respect to β, we also find that f has a single non-negative critical point, at the pointβ = σ 2 + E[(ûV -Θ) 2 ], whereû = û(V, Θ) = prox λ(•) 2 +µh(•,V,Θ) Θ -βG √ κ V ; β 2τ V 2 .So, we can conclude that f has unique saddle point (τ , β). Note f is a deterministic function that does not depend on d and hence (τ , β) are also deterministic and independent of d. Now let C 2 := max{τ , β} + 1. By the "convexity lemma" (as stated in [24] ), pointwise convergence (in probability) of a convex function is uniform over compact sets. So, this result implies that the convergence is uniform over(τ, β) ∈ [0, C 2 ] × [0, C 2 ], 2 so max 0≤τ ≤C2 min 0≤β≤C2 f d (τ, β) P → max 0≤τ ≤C2 min 0≤β≤C2 f (τ, β)Let (τ d , βd ) denote the optimnal solution for the problem on the left. We can also conclude that (τ d , βd ) P → (τ , β) by [22, Theorem 2.1] , which states that uniform convergence in probability of a convex function over a compact set implies convergence of the optimal minimizer. So, with probability approaching 1, (τ d , βd ) are strictly smaller than C 2 , and the same solution is also optimal for P 2 (µ). We can therefore concludeP 2 (µ) P → max 0≤τ ≤C2 min 0≤β≤C2 f (τ, β) = max τ ≥0 min β≥0 f (τ, β) = : P (µ)Therefore, by Theorem 3, for any fixed µ ∈ [-µ * , µ * ], we have the convergenceP 1 (µ) P → P (µ).In the special case µ = 0, we can further simplify the Moreau envelope term to obtainEQUATION)Step 2: Convergence of the optimal solution We next need to extend this result to the desired Wasserstein-2 convergence result (8) . Recall here that u (t+1) is the solution of P 1 (0).First, let h : R 3 → R be any bounded, Lipschitz function, and let h (k) be a sequence of bounded, twice-differentiable functions that converge uniformly to h as k → ∞ (e.g., the convolution of h with a sequence of mollifiers). Let P (k) 1 (µ), P (k) (µ) be the optimal cost of P 1 and P when using test function h (k) and for µ ∈ [-µ * , µ * ]. Note the convergence P By the uniform convergence of the h (k) to h,lim k→∞ P (k) 1 (µ) = P 1 (µ) lim k→∞ P (k) (µ) = P (µ). Now, fix δ > 0 and choose k large enough that |P (k) 1 (µ) -P 1 (µ)| < δ 3 and | P (k) (µ) -P (µ)| < δ 3 . Then, P |P 1 (µ) -P (µ)| > δ ≤ P |P (k) 1 (µ) -P (k) (µ)| > δ/3 → 0, since P (k) 1 P → P2(k) for all k. Hence, we can also apply the result of Step 1 to any bounded Lipschitz function h.Since the convergence result of Step 1 holds for any µ in a neighborhood around zero, we can conclude that1 d d i=1 h(u (t+1) i , v i , θ * i ) P → d P (µ) dµ µ=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 7 of [5] , so we omit it here. Moreover, using the Dominated Convergence Theorem to differentiate inside the expectation, we can compute this exactly:d P (µ) dµ µ=0 = E h prox λ(•) 2 Θ -βG √ κ V ; β 2τ V 2 , V, Θ = E h τ V (Θ -βG √ κ) τ V 2 + βλ , V, Θ ,where ( β, τ ) are found in the optimal solution to P (0).Hence, for all bounded, Lipschitz h, we have1 d d i=1 h(u (t+1) i , v i , θ * i ) P → E h τ V (Θ -βG √ κ) τ V 2 + βλ , V, Θ ,so the empirical distribution of the triple (u(t+1) i , v i , θ * i ) converges weakly to the distribution of the random variable ( τ V (Θ-βG √ κ) τ V 2 + βλ , V, Θ), where G ∼ N (0, 1) and (V, Θ) ∼ Π t . By choosing h(u, v, θ) = u 2 , we also know that second moments of the empirical distribution converge in probability. Hence, the convergence can be strengthened from weak convergence to convergence in W 2 distance (see, e.g. [33, Theorem 6.9]).Step 3: Verifying the inductive hypothesis Lastly, we need to show that1 d d i=1 δ(v (t+1) i , θ * i ) W2 → Law(ψ(Q t+1 , V ), Θ) := Π t+1 .where (V, Θ) ∼ Π t . Here, weak convergence follows from the result of Step 2 since ψ is a continuous map. To show convergence of second moments, we need to show1 d d i=1 v (t+1),2 i = 1 d d i=1 ψ(u (t+1) i , v (t) i ) 2 P → E[ψ(Q t+1 , V ) 2 ].For ψ that satisfies Assumption 2, this convergence is immediate from the result of Step 2 (since W 2 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 1.Proof of Theorem 2. The proof is an extension of the proof of Theorem 1 to the case where the test function h acts on blocks rather than individual entries. Much of the proof is identical, so we only sketch the argument and highlight the major differences here. We begin with the inductive hypothesis that 1M M i=1 δ(v (t) i , θ * i ) W2 → Π t , for a known distribution Π t over R b × R b .Recall here that M denotes the number of blocks/factors of size b (so M = d b ). Then, let h : (R b ) 3 → R be a test function with ∥∇ 2 h∥ 2 ≤ C and such that either h is bounded orh(u i , v i , θ i ) = ∥u i ∥ 2 2 .We consider a similar perturbed optimization problem:EQUATIONAgain, we consider this for |µ| ≤ λ bC , so that the optimization problem is λ d strongly convex in u. Noting that the proof of Lemma 1 still holds in this grouped case, we can constrain P 1 (µ) to be over compact sets and apply the CGMT to obtain the auxiliary problemEQUATIONEQUATION)The sequence of "scalarization" steps on P 2 is identical to in Theorem 1, until we arrive atEQUATIONi , θ * i ).(25) Here, since the sum of the last two terms in the objective function is λ d strongly convex by our choice of µ, the proof of Lemma 2 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 as1 d M i=1 min ui∈R b τ ξ β ∥u i ⊙ v i -θ * i ∥ 2 2 -2τ √ κg ⊤ i (u i ⊙ v i -θ) + λ∥u∥ 2 2 + µbh(u i , v i , θ * i ) = 1 bM M i=1 min ui∈R b τ ξ β ∥u i ⊙ v i -θ * i ∥ 2 2 -2τ √ κg ⊤ i (u i ⊙ v i -θ) + λ∥u∥ 2 2 + µbh(u i , v i , θ * i ) := 1 bM M i=1 q(v i , θ * i , g i ),where g i ∈ R b denotes the ith block of g and q := (R b ) 3 → R is defined as a shorthand for the quantity inside the summation.Next, we consider the asymptotic behavior of P 2 (µ). Here, the only term which is different than in Theorem 1 is the term1 bM M i=1 q(v i , θ * i , g i ).By the same argument as in Lemma 3, we can write q as the Moreau envelope of a convex function and show that1 bM M i=1 q(v i , θ * i , g i ) P → E 1 b q(V , Θ, G),where the expectation is over (V , Θ) ∼ Π t and G ∼ N (0, I b ). After the same uniform convergence argument as in the proof of Theorem 1, we can conclude that for all µ ∈ [-µ * , µ * ], P 1 (µ)P → P (µ), where P (µ) = max τ ≥0 min β≥0 τ σ 2 β + τ β -τ 2 + E 1 b q(V , Θ, G).In particular, when µ = 0, 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 1, and for the final step (showing the inductive hypothesis holds at the next iteration), we need to argue that the second moment ofEQUATION, v i ) converges to its expectation under Π t+1 :1 M M i=1 ∥ψ(u (t+1) i , v i )∥ 2 2 P → E∥ψ(Q t+1 , V )∥ 2 2If ψ is bounded and continuous or has PL(2) coordinate projections (as we have assumed), then the above convergence holds based on the Wasserstein-2 convergence of the joint distribution of (u (t+1) , v).Proposition 1. The function g(u, v, θ) = |uv -θ| is pseudo-Lipschitz of order 2.Proof. The result follows from the following series of inequalities:||uv -θ| -|u ′ v ′ -θ ′ || ≤ |uv -θ -(u ′ v ′ -θ ′ )| ≤ |uv -u ′ v ′ | + |θ -θ ′ | ≤ |u||v -v ′ | + |v ′ ||u -u ′ | + |θ -θ ′ | ≤ (|u| + |v ′ | + 1)(|u -u ′ | + |v -v ′ | + |θ -θ ′ |) ≤ (1 + ∥x∥ 1 + ∥x ′ ∥ 1 )∥x -x ′ ∥ 1 ≤ 3(1 + ∥x∥ 2 + ∥x ′ ∥ 2 )∥x -x ′ ∥ 2 , where x, x ′ ∈ R 3 denote (u, v, θ) and (u ′ , v ′ , θ ′ ), respectively.Lemma 1. Let ∆ * , q * be the optimal solution to (10) . Then, there exists universal constantC 1 > 0 such that lim d→∞ P{∥∆ * ∥ 2 ≤ C 1 ∥v∥ ∞ } = lim d→∞ P{∥q * ∥ 2 ≤ C 1 ∥v∥ ∞ } = 1.Proof. We proceed via a similar argument to Lemma 2 in [5] . First, consider the original expression for P 1 (µ) from ( 9) :P 1 (µ) = min u∈R d F (u) + R(u),whereF (u) := 1 n y -1 √ d X(u ⊙ v) 2 2and R(uEQUATIONi , θ * i ). Recall here that R is λ d -strongly convex for all µ ∈ [-µ * , µ * ], and denote the unique optimal minimizer to this problem as u * . Then, the following chain of inequalities holds, by the optimality of u * and the non-negativity of F .1 n ∥y∥ 2 2 + R(0) = F (0) + R(0) ≥ F (u * ) + R(u * ) ≥ R(u * ).Moreover, by the strong convexity of R, we haveR(u * ) ≥ R(0) + ∇R(0) ⊤ u * + λ 2d ∥u * ∥ 2 2 .Combining the above two series of inequalities, we obtain (recallκ = d n ) ∥u * ∥ 2 2 + 2d λ ∇R(0) ⊤ u * ≤ 2κ λ ∥y∥ 2 2 .After completing the square, this yieldsu * + d λ ∇R(0) 2 2 ≤ 2κ λ ∥y∥ 2 2 + d 2 λ 2 ∥∇R(0)∥ 2 2 ,whence, by the triangle inequality,∥u * ∥ 2 ≤ d λ ∥∇R(0)∥ 2 + 2κ λ ∥y∥ 2 2 + d 2 λ 2 ∥∇R(0)∥ 2 2 ≤ 2d λ ∥∇R(0)∥ 2 + 2κ λ ∥y∥ 2 .Here, standard concentration inequalities for Gaussian random variables (e.g., [ ∥y∥ 2 ≤ µ √ d ∥X∥∥θ * ∥ 2 + ∥ϵ∥ 2 ≲ √ dwith probability approaching 1. Next, we bound ∥∇R(0)∥ 2 . Recalling the definition of R,∥∇R(0)∥ 2 = µ d d i=1 ∂ ∂u h(u, v i , θ * i ) u=0 2 = 1 √ d 1 d d i=1 ∂ ∂u h(u, v i , θ * i ) u=0 2 . Since the function g(v, θ) = ∂ ∂u h(u, v, θ) u=0is Lipschitz (by the fact that h has bounded second derivatives), g 2 is pseudo-Lipschitz of order 2. So, the quantity under the square root converges in probability to E (V,Θ)∼Πt g 2 by Assumption 1, and, with probability tending to 1, we have∥∇R(0)∥ 2 ≲ 1/ √ d.Combining the above bounds on ∥y∥ 2 and ∥∇R(0)∥ 2 , we can conclude that ∥u * ∥ 2 ≲ √ d. The first part of the lemma follows by noting that∥∆ * ∥ 2 = 1 √ d ∥u * ⊙ v -θ * ∥ 2 ≤ 1 √ d ∥u * ⊙ v∥ 2 + 1 √ d ∥θ * ∥ 2 ≤ 1 √ d ∥v∥ ∞ ∥u * ∥ 2 + 1 √ d ∥θ * ∥ 2 ≲ ∥v∥ ∞ ,where the last inequality holds with probability approaching 1. Lastly, the optimal q for any ∆ has closed-form q = 1 √ n ϵ -1 √ n X∆. By the triangle inequality, we then obtain∥q * ∥ 2 ≤ 1 √ n ∥ϵ∥ 2 + 1 √ n ∥X∥∥∆ * ∥ 2 ≲ ∥v∥ ∞ ,with the last inequality holding with probability approaching 1, by the concentration of norms of ϵ and X as discussed above, and the bound on ∥∆ * ∥ 2 .Lemma 2. Consider the following unconstrained minimization problem over u ∈ R d :EQUATIONi , θ * i ).With probability approaching 1, the solution u * satisfies 1  and 1 √ d ∥θ * ∥ 2 is bounded by a constant with probability approaching 1 by the assumed W 2 convergence of θ * to a fixed limit. Hence, it suffices to show that ∥u * ∥ 2 ≲ √ d with high probability. We begin by noting that the inner maximization over τ admits a closed form solution, so the problem becomes√ d (u * ⊙ v -θ * ) 2 ≲ ∥v∥ ∞ . Proof. First note 1 √ d (u * ⊙ v -θ * ) 2 ≤ 1 √ d ∥u * ∥ 2 ∥v∥ ∞ + 1 √ d ∥θ * ∥ 2 ,EQUATIONi , θ * i ).Now, we can proceed similarly to in the proof of Lemma 1. LetEQUATIONi , θ * i ).Then, noting F is always non-negative and R is λ d strongly-convex, we use the same argument as in Lemma 1 to obtain the inequality∥u * ∥ 2 2 + 2d λ ∇R(0) ⊤ u * ≤ 2d λ F (0).After completing the squares, we obtainu * + d λ ∇R(0) 2 2 ≤ 2d λ F (0) + d 2 λ 2 ∥∇R(0)∥ 2 2 ,so we can conclude∥u * ∥ 2 ≤ d λ ∥∇R(0)∥ 2 + 2d λ F (0) + d 2 λ 2 ∥∇R(0)∥ 2 2 ≤ 2d λ ∥∇R(0)∥ 2 + 2d λ F (0).As shown in Lemma 1, ∥∇R(0)∥ 2 ≲ 1 √ d with probability approaching 1. It remains to show that F (0) ≲ 1 with probability approaching 1. To see this, observe thatF (0) = ∥h∥ 2 2 √ n σ 2 + 1 d ∥θ * ∥ 2 2 - 1 2 √ nd g ⊤ θ * + ≤ ∥h∥ 2 2 √ n σ 2 + 1 d ∥θ * ∥ 2 2 - 1 2 √ nd g ⊤ θ * ≤ ∥h∥ 2 2 √ n σ 2 + 1 d ∥θ * ∥ 2 2 + 1 2 √ nd ∥g∥ 2 ∥θ * ∥ 2 ,where the last line uses the triangle and Cauchy-Schwarz inequalities. By concentration of the norm for Gaussian vectors, there exists a universal constant c such that ∥h∥2 √ n ≤ c and ∥g∥2 √ n ≤ c with probability approaching 1. Moreover, by Assumption 1, ∥θ * ∥2 √ d ≤ c with probability approaching 1. Hence, F (0) ≲ 1 with probability approaching 1. Substituting this into the bound for ∥u * ∥ 2 from above completes the proof. Lemma 3. Under Assumption 1, the function f d (τ, β) (Eq. 19) converges pointwise in probability to f (τ, β) (Eq. 20) as d → ∞. Proof. We consider the limit of each term in f d separately. The limit of the terms τ σ 2 ξ β and τ βξ is found by noting that ξ → 1 in probability by Gaussian Lipschitz concentration.The first summation term simplifies as follows:1 d d i=1 τ ξθ * 2 i β - τ βξ βg i √ κ -ξθ * i 2 = 1 d d i=1 - τ βξ β 2 g 2 i κ -2βg i √ κξθ * i = - τ βξ 1 d d i=1 β 2 g 2 i κ -2βg i √ κξθ * i P → -τ βκ,where the last line follows from the weak law of large numbers since the g i are i.i.d. standard Gaussian variables.For the last term, after again using the fact that ξ P → 1, we need to consider1 d d i=1 M λ(•) 2 +µh(•,v (t) i ,θ * i ) θ * i -βg i √ κ v (t) i ; β 2τ v (t)2 i = : 1 d d i=1 q(v i , θ * i , g i )To show convergence in probability of this term, first fix δ > 0. Then, we want to showP 1 d d i=1 q(v i , θ * i , g i ) -E q(V, Θ, G) > δ → 0,where the expectation is over (V, Θ) ∼ Π t and G ∼ N (0, 1). It suffices to show the following two statements:EQUATIONEQUATIONTo show (26) , we rely on a concentration inequality for the Moreau envelope of a Gaussian vector plus a bounded vector from [5] . First note that1 d d i=1 q(v i , θ * i , g i ) = 1 d min u∈R d λ∥u∥ 2 2 + µ d i=1 h(u i , v i , θ * i ) + τ β u ⊙ v -θ * + β √ κg 2 2 = 1 d min u∈R d λ∥u∥ 2 2 + µ d i=1 h(u i , v i , θ * i ) + τ βκ u ⊙ v β √ κ - θ * β √ κ + g 2 2 = 1 d min θ∈R d λ β √ κθ v 2 2 + µ d i=1 h β √ κθ i v i , v i , θ * i + τ βκ θ - θ * β √ κ + g 2 2 = 1 d M ℓ θ * β √ κ -g; 1 2τ βκ ,where the second to last line follows from the change of variable θ = u⊙v β √ κ , and ℓ(θ) := λ β √ κθ v 2 2 + µ d i=1 h β √ κθi vi , v i , θ * i .Here, for fixed v, ℓ is a proper convex function of θ. Moreover, by Assumption 1, θ * β √ κ has norm of order √ d with high probability. Hence, by [5, Lemma 8] , this quantity concentrates around its expectation (with respect to g), and we can conclude that there is some c > 0 such thatP 1 d d i=1 q(v i , θ * i , g i ) -E g 1 d d i=1 q(v i , θ * i , g i ) > δ 2 ≤ cτ 2 β 2 κ 2 dδ 2 → 0.To show (27) , note thatE g 1 d d i=1 q(v i , θ * i , g i ) = 1 d d i=1 E G q(v i , θ * i , G).Observe that this quantity is an expectation with respect to the joint empirical distribution of (v, θ * ).By the assumption of W 2 convergence of the empirical distribution of (v, θ * ) (Assumption 1), if we can show that mapping (v, θ) → E G q(v, θ, G) is bounded, then we can conclude that the above quantity converges in probability to E q(V, Θ, G), with (V, Θ) ∼ Π t . To see this, recall that for all G ∈ R, q is bounded below asq(v, θ, G) ≥ min u λu 2 + µh(u, v, θ),which is always bounded below since h is bounded (or, in the case h = u 2 , the lower bound is zero).Next, for a given G, we can bound q above asq(v, θ, G) ≤ µh(0, v, θ) + τ β (θ -β √ κG) 2 .Hence,E G q(v, θ, G) ≤ µh(0, v, θ) + τ β θ 2 + τ βκ < Cfor some universal constant C > 0, since θ is bounded by assumption and h(0, v, θ) is either bounded above or equal to 0 (in the case where h = u 2 ). Combining ( 26) and ( 27) yields the desired result.Below, we show that the max-min problems ( 5) and ( 7) have easily computable solutions. Note it suffices to consider the grouped case ( 7), since we can apply it with b = 1 to recover the ungrouped case. The following derivation closely follows the analysis of the scalar max-min problem in [8] , who study a similar scalar problem (albeit in the case of λ = 0, which we do not consider).First recall that, as shown in the proof of Theorem 1, there exists a unique saddle point, due to the strong convexity in τ and strict convexity in β. Now, taking derivatives with respect to τ and β, we obtain the following saddle point conditions:EQUATIONEQUATIONSolving each of these equations for the quantity 2 and then equating the two, we arrive atβ 2 λ 2 E 1 b b i=1 Θ 2 i +β 2 κ (τ V 2 i +βλ)2τ β -β 2 (1 -κ) -σ 2 = β 2 (1 -κ) -σ 2 + 2λβ 3 κ E 1 b b i=1 1 τ V 2 i + βλ , which implies τ = β(1 -κ) + λβ 2 κ E 1 b b i=1 1 τ V 2 i + βλ .Defining the auxiliary variable γ = τ /β, this yields the fixed point equationEQUATIONSubstituting this γ back into the first optimality condition in (28) , we can express the optimal β in closed-form, in terms of γ:β = σ 2 + λ 2 E 1 b b i=1 Θ 2 i (γV 2 i +λ) 2 2γ + κ -1 -λ 2 κ E 1 b b i=1 1 (γV 2 i +λ) 2 .Finally, the optimal τ can be found simply as τ = γβ.This yields a simple recipe for solving the min-max problem. First, compute the positive solution γ to the fixed point equation ( 30) (this can be found easily using standard numerical solvers). Then, ( β, τ ) are both given in closed-form as functions of γ (where the required expectations can all be approximated via Monte Carlo simulation).In this section, we demonstrate that our asymptotic predictions can 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) = tanh |uv| and ψ(u, v) = tanh u 2 , to the same weightings without the bounded tanh activation: ψ(u, v) = |uv| and ψ(u, v) = u 2 . We note that the reweighting choice |uv| is considered in [21, 28] as a limit as p → 0 of the classical IRLS update for ℓ p minimization. In Figure 3a ∼ Bernoulli(0.01) and λ chosen to minimize the predicted asymptotic loss.For each choice of ψ, we apply the theoretical predictions of Theorem 1, even if ψ violates Assumption 2. We find that our predictions remain accurate for all these choices of ψ. The choice tanh |uv| performs almost identically without the tanh activation. Interestingly, the choice ψ = tanh u 2 outperforms the variant without the tanh and has a more regular decay of the test loss.In Figure 3b , we apply Theorem 1 to predict the asymptotic squared test loss: 1 d ∥u ⊙ vθ * ∥ 2 2 at each iteration. While this function is not P L(2), as required by the theorem, the asymptotic predictions still align well with simulations. Extending our technical results to hold formally in such scenarios is an interesting direction for future work. ∼ Bernoulli(0.01) 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 error1 d ∥u ⊙ v -θ * ∥ 2 2 .• If the contribution is a 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 is 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , 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 to have some path to reproducing or verifying the results.38th Conference on Neural Information Processing Systems (NeurIPS 2024).The uniqueness of the solution is shown in the proof of Theorem 1.Note we could not have directly applied this to the feasible sets of P2(µ), since R may have a dependence on d. 


Decision: 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.

Reviews: review: paper_summary: paper_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". 
The 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.
In 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.
The result is then extended the multi-variate setting in Theorem 2.
Simulations results are provided to support the proposed theoretical analyses.
strength_weakness: strength_weakness: 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.
I have some detailed comments below
questions: 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.
* 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?
* 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.
* 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? 

minor: 
* line 197: of the algorithm in (4) or in Equation (4).
limitations: limitations: I do not see any potential negative social impact of this work.
paper_summary: paper_summary: The paper considers a general algorithm of the form shown in Lines 66-67, as it includes several interesting methods as special cases.

The 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.

Synthetic experiments show that the errors predicted by Theorem 1 align well with the actual performance of the algorithms.
strength_weakness: strength_weakness: - 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:
   - 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.)
   - Could authors please comment on the technical hurdles if the independence assumption is not available? 
   - 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
Pursuit with Global Linear Convergence Rate". Would their analysis be useful in any way to remove the independence assumption?

- The second weakness is that the experimental evaluations in their current form are not very deep. For example:
   - 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.
   - 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.
questions: 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$.
- Lines 185-186: What is the precise relation between $\alpha$ and $p$? It is not very clear why we need two different symbols.
limitations: limitations: See above.
paper_summary: paper_summary: This paper studies the trajectory of a certain iterative scheme in the high-dimensional limit. 
In each step, there're two update steps on a fresh batch: 1) regularized least squares; 2) element-wise nonlinear transformation.
This type of iteration is motivated and encompasses several interesting iterations such as IRLS, AGOP, etc. 

The result is a precise characterization of the empirical distribution of the iterates at any finite time in the proportional limit. 
This allows one to compute many interesting summary statistics.
strength_weakness: strength_weakness: 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). 
I 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. 
So I don't really object to the sample splitting assumption. 

Also, CGMT is capable of offering non-asymptotic guarantees. 
I wonder if the authors are interested in doing so (no pressure).
questions: questions: Minor comments: 

1. In the equation between line 66/67, the fresh dataset $\boldsymbol{y}^{(t)}, \boldsymbol{X}^{(t)}$ hasn't been defined (correct me otherwise). 

1. Line 49, typo: is --> are

1. Line 164, typo: $t = 0, 1, \dots, T$

1. Equation between line 165/166, why not add superscript $(t)$ to $\boldsymbol{\epsilon}$? 
After all, the noise is also fresh and is not shared by other batches. 

1. Regarding boundedness of $\boldsymbol{\theta}^*$, as mentioned, it's only used in Lemma 3. 
Over 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$. 
(To clarify, this is a minor assumption and removing it requires technical work that may not worth it.
So I don't mind if the authors leave it as is.) 

1. Line 246-249, the authors meant to discuss a second difference from prior works. 
What's done in this paper is discussed but what's done in prior work is not. 
So the comparison is unclear. 

1. There are two notation sections that are almost identical. I suggest remove one of them. 

1. Line 462, what is a "proper function"? Do you mean a "proper convex function"?

1. Line 466, "restate": please put a reference for CGMT. 

1. Equation between line 577/578, $V$ should be $V_t$. 

1. Equation between line 579/580, $v_i$ should be $v_i^{(t)}$, $V$ should be $V_t$. 

1. Line 580, typo: "satisfies".
limitations: limitations: NA

Comments: title: Reply
Simulations with reused samples
comment: Dear authors, thank you for your reply to my comments. 

I have taken a look at the rebuttal and also the comments from other reviewers.

I would like to keep my current, positive score, given the quality of the paper.
Thank you very much for the prompt update. 
This makes perfect sense to me.
Thank you for the response and great observation – yes, in general, we would expect reusing $nT$ samples at each iteration to be better than sample-splitting (in line with the references you mentioned). We agree that this is perhaps the more standard point of comparison, and we have rerun the simulations with this additional curve. As expected, this approach does outperform the other approaches when the regularization parameter is appropriately tuned (unfortunately, we are not able to update the pdf in this response). So, to summarize, our theoretical predictions for the sample-split case 

- seem to remain accurate when randomly sampling $n$ samples from a total pool of $nT$ at each iteration
- are too optimistic for the scenario where $n$ samples are reused at every iteration
- are too pessimistic for the case where $nT$ samples are reused at every iteration.
Thank you for the detailed response. 

I have a question regarding the numerical results shown in the attached pdf.
I thought that reusing batches is beneficial (in a somewhat different context of training say 2-layer neural nets with SGD) given a recent line of work [R1, R2, R3]. The message there is that to achieve the same error, the sample complexity is smaller (in order) if samples are reused (even twice per sample). So with the same sample size, the error of online / sample-splitting methods should be larger (cf.\ [R4, R5]).
However, the experiment in the attached pdf suggests a seemingly opposite message since it compares sampling-splitting with reusing **a single batch**.
This doesn't seem to be a fair comparison and indeed leads to an ``optimistic'' prediction.
It seems more standard to compare the performance with and without sample-splitting on the same $nT$ samples.

---

[R1] Lee, J. D., Oko, K., Suzuki, T., & Wu, D. (2024). Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. arXiv preprint arXiv:2406.01581.

[R2] Dandi, Y., Troiani, E., Arnaboldi, L., Pesce, L., Zdeborová, L., & Krzakala, F. (2024). The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220.

[R3] Arnaboldi, L., Dandi, Y., Krzakala, F., Pesce, L., & Stephan, L. (2024). Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405.15459.

[R4] Damian, A., Nichani, E., Ge, R., & Lee, J. D. (2024). Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36.

[R5] Arous, G. B., Gheissari, R., & Jagannath, A. (2021). Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106), 1-51.
I thank the authors for their comments (including the global one). 
I took a brief look at them and keep my evaluation unchanged.