Title: Precise asymptotics of reweighted least-squares algorithms for linear diagonal networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: high-dimensional asymptotics

linear diagonal neural networks

feature learning

iteratively reweighted least-squares

sparse recovery
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 asymptotics of reweighted least-squares algorithms for linear diagonal 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 ℓ 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.

1 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

where 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:

In the case where the regression function is linear in θ , solving (2) is equivalent to learning a function of the form

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 θ :

After 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

for some reweighting function ψ and batches of training data ( X ( t ) , y ( t ) ) . As we show in Section 2, this formulation encompasses multiple existing algorithms, including reparameterized IRLS, lin-RFM, and alternating minimization over u 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 ) →∞ . This allows us to address questions such as

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?

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

1.1 Related work

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.

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 [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: Onatechnical 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).

2 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 ,

for 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 W 2 → ν , 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 W 2 → ν 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 as 1 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 satisfying

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 form

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 function

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.

3 A precise characterization of the iterates

In this section, we provide a precise characterization of the iterates of the algorithm in Equation 4 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 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., 1 d ∑ d i =1 δ ( v (0) i , θ ∗ i ) W 2 → Π 0 . Additionally, v (0) i = 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:

̸

If U, V are random variables such that U, V = 0 with probability 1, then ψ ( U, V ) = 0 with probability 1.

ψ 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 :

where G 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 unique 1 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 have

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

̸

1

2

3

4

5

6

7

8

Iterations

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)

(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

.

001

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

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

3.1 Application to sparse linear regression

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 ). Werun 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 M = d b factors: R b ×··· × R b . Then, θ ∗ , u ( t ) , v ( t ) ∈ R d can all be represented as M stacked blocks, each in R b . Under the same linear measurement model, we now let ψ : R b × R b → R b act on each of the factors of ( u ( t ) , θ ∗ ) , and consider the same Algorithm 4.

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 } i M =1 are zero. In this case, it might make sense for ψ to return a vector of the form

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]:

where the closed-form solution to the η minimization yields the classical group norm regularizer ∑ M i =1 ∥ θ i ∥ 2 . Here, reparameterizing as α i → √ η i and u i → θ i √ η i gives rise naturally to the grouped Hadamard parameterization, with v 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 :

Π 0 is the limit of the empirical distribution of factors of ( v (0) , θ ∗ ) and hence is a distribution over R b × R b .

Each factor θ ∗ i ∈ R b for i = 1 , . . . , M has bounded ℓ 2 -norm almost surely.

ψ is bounded and continuous or each of its coordinate projections satisfies ψ 2 j ∈ PL (2) for j = 1 , . . . , b .

1

2

3

4

5

6

7

8

Iterations

10

-

2

ℓ

1

Test Error

Group-blind (theory)

Group-blind (sim)

Group-aware (theory)

Group-aware (sim)

(a) Comparison of trajectory

1

2

4

8

10

Ground truth group size

0

.

002

0

.

004

0

.

006

0

.

008

ℓ

1

Test Error

Group-blind (theory)

Group-blind (sim)

Group-aware (theory)

Group-aware (sim)

(b) Effect of group size

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

Here, 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 have

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 b

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

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

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

Francis Bach. The ' η -trick' or the effectiveness of reweighted least-squares. https://francisbach.com/the-%CE%B7-trick-or-the-effectiveness-ofreweighted-least-squares/ , 2019. Accessed: April 2024.

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.

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-785, 2011.

Michael J 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.

David Bosch, Ashkan Panahi, Ayca 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. PMLR, 2023.

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.09660 , 2022.

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-210, 2023.

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, 2021.

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. IEEE, 2008.

Hung-Hsu Chou, Johannes Maly, and Holger Rauhut. More is less: inducing sparsity via overparameterization. Information and Inference: A Journal of the IMA , 12(3):1437-1460, 2023.

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.

Donald Geman and George Reynolds. Constrained restoration and the recovery of discontinuities. IEEE Transactions on Pattern Analysis & Machine Intelligence , 14(03):367-383, 1992.

Peter D Hoff. Lasso, fractional norm and structured sparse estimation using a Hadamard product parametrization. Computational Statistics & Data Analysis , 115:186-198, 2017.

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.

Chris Kolb, Christian L Müller, Bernd Bischl, and David Rügamer. Smoothing the edges: A general framework for smooth optimization in sparse regularization using Hadamard overparametrization. arXiv preprint arXiv:2307.03571 , 2023.

Akshay Kumar, Akshay Malhotra, and Shahab Hamidi-Rad. Group sparsity via implicit regularization for MIMO channel estimation. In 2023 IEEE Wireless Communications and Networking Conference (WCNC) , pages 1-6. IEEE, 2023.

Jiangyuan Li, Thanh V Nguyen, Chinmay Hegde, and Raymond KW Wong. Implicit regularization for group sparsity. arXiv preprint arXiv:2301.12540 , 2023.

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.

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.

Charles A Micchelli, Jean M Morales, and Massimiliano Pontil. Regularizers for structured sparsity. Advances in Computational Mathematics , 38:455-489, 2013.

Karthik Mohan and Maryam Fazel. Iterative reweighted algorithms for matrix rank minimization. The Journal of Machine Learning Research , 13(1):3441-3473, 2012.

Whitney K Newey and Daniel McFadden. Large sample estimation and hypothesis testing. Handbook of Econometrics , 4:2111-2245, 1994.

Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems , 34:29218-29230, 2021.

David Pollard. Asymptotics for least absolute deviation regression estimators. Econometric Theory , 7(2):186-199, 1991.

Clarice Poon and Gabriel Peyré. Smooth bilevel programming for sparse regularization. Advances in Neural Information Processing Systems , 34:1543-1555, 2021.

Clarice Poon and Gabriel Peyré. Smooth over-parameterized solvers for non-smooth structured optimization. Mathematical Programming , 201(1):897-952, 2023.

Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mechanism for feature learning in neural networks and backpropagation-free machine learning models. Science , 383(6690):1461-1467, Mar 2024.

Adityanarayanan Radhakrishnan, Mikhail Belkin, and Dmitriy Drusvyatskiy. Linear recursive feature machines provably recover low-rank matrices. arXiv preprint arXiv:2401.04553 , 2024.

Christos Thrampoulidis, Ehsan Abbasi, and Babak Hassibi. Precise error analysis of regularized m -estimators in high dimensions. IEEE Transactions on Information Theory , 64(8):5592-5628, 2018.

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.

Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems , 32, 2019.

Roman Vershynin. High-dimensional Probability: An Introduction with Applications in Data Science , volume 47. Cambridge university press, 2018.

Cédric Villani et al. Optimal Transport: Old and New , volume 338. Springer, 2009.

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-3673. PMLR, 2020.

Peng Zhao, Yun Yang, and Qiao-Chu He. High-dimensional linear regression via implicit regularization. Biometrika , 109(4):1033-1046, 2022.

Liu Ziyin and Zihao Wang. spred: Solving L 1 penalty with SGD. In International Conference on Machine Learning , pages 43407-43422. PMLR, 2023.

A Proof of main results

In this section, we provide the proofs of our 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 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 ,

for 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 W 2 → ν , 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 W 2 → ν 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 as 1 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 as

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:

Then, for all c ∈ R and t > 0 ,

A.2 Proof of Theorems 1 and 2

Proof of Theorem 1. Assume that 1 d ∑ d i =1 δ ( v ( t ) i , θ ∗ i ) W 2 → Π 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 that

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 that

where ( V, Θ) ∼ Π t and Q t +1 is defined as in Eq. 5.

Recall that the objective function for the update on u is given by

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:

h is uniformly bounded.

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):

where µ ∈ [ -µ ∗ , µ ∗ ] 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

̸

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 2 q ⊤ ( · ) -∥ q ∥ 2 2 , we can write this as

Next, 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 study

where P 1 ( µ ) = ˜ P 1 ( µ ) 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 problem

where g ∈ R d and h ∈ R n have i.i.d. standard normal entries. By Theorem 3, for all δ > 0 and fixed ¯ P ( µ ) ∈ R ,

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 yields

where R := C 1 ∥ v ∥ ∞ . Next, note that ϵ and h are independent Gaussian vectors and hence ϵ -∥ ∆ ∥ 2 h d = √ σ 2 + ∥ ∆ ∥ 2 2 h . So, we have

Before proceeding further, we rewrite this in terms of a minimization over the variable u = √ d ∆ + θ ∗ v .

Here, 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 as

Here, 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 problem

Completing the squares, we obtain that the above problem can be written in terms of the Moreau envelope (Definition 1) of the function ℓ ( u ) = λu 2 + µh :

where we have introduced the shorthand notation ξ = ∥ h ∥ 2 √ n . Substituting this into the expression for P 2 above, we obtain

Now 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 to

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 [(ˆ uV -Θ) 2 ] , where √

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

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 conclude

Therefore, by Theorem 3, for any fixed µ ∈ [ -µ ∗ , µ ∗ ] , we have the convergence

In the special case µ = 0 , we can further simplify the Moreau envelope term to obtain

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 ( k ) 1 ( µ ) P → ¯ P ( k ) ( µ ) for any µ in a sufficiently small neighborhood around zero holds by Step 1.

By the uniform convergence of the h ( k ) to h ,

Now, fix δ > 0 and choose k large enough that | P ( k ) 1 ( µ ) -P 1 ( µ ) | < δ 3 and | ¯ P ( k ) ( µ ) -¯ P ( µ ) | < δ 3 . Then,

since P ( k ) 1 P → ¯ P 2 ( 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 that

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:

where ( ˆ β, ˆ τ ) are found in the optimal solution to ¯ P (0) .

Hence, for all bounded, Lipschitz h , we have

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 that

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 show

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 1 M ∑ M i =1 δ ( v ( t ) i , θ ∗ i ) W 2 → Π 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 or h ( u i , v i , θ i ) = ∥ u i ∥ 2 2 . We consider a similar perturbed optimization problem:

Again, 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 problem

The sequence of 'scalarization' steps on P 2 is identical to in Theorem 1, until we arrive at

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 as

where g i ∈ R b denotes the i th 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 term 1 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 that

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

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 of v ( t +1) = ψ ( u ( t +1) i , v i ) converges to its expectation under Π t +1 :

If ψ 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 ) .

B Technical lemmas

Proposition 1. The function g ( u, v, θ ) = | uv -θ | is pseudo-Lipschitz of order 2.

Proof. The result follows from the following series of inequalities:

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 constant C 1 > 0 such that

Proof. We proceed via a similar argument to Lemma 2 in [5]. First, consider the original expression for P 1 ( µ ) from (9) :

where F ( u ) := 1 n ∥ ∥ ∥ y -1 √ d X ( u ⊙ v ) ∥ ∥ ∥ 2 2 and R ( u ) := λ d ∥ u ∥ 2 2 + µ d ∑ d i =1 h ( u i , v ( t ) i , θ ∗ 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 .

Moreover, by the strong convexity of R , we have

Combining the above two series of inequalities, we obtain (recall κ = d n )

After completing the square, this yields

whence, by the triangle inequality,

Here, standard concentration inequalities for Gaussian random variables (e.g., [32, Theorem 5.2.2, Corollary 7.3.3]) imply that, with probability approaching 1, ∥ X ∥ 2 ≲ √ d and ∥ ϵ ∥ 2 ≲ √ d . And Assumption 1 implies that ∥ θ ∗ ∥ 2 ≲ √ d with probability tending to 1. So,

with probability approaching 1. Next, we bound ∥∇ R ( 0 ) ∥ 2 . Recalling the definition of R ,

Since the function g ( v, θ ) = ∂ ∂u h ( u, v, θ ) ∣ ∣ ∣ ∣ u =0 is 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

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

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 :

With probability approaching 1, the solution u ∗ satisfies ∥ ∥ ∥ 1 √ d ( u ∗ ⊙ v -θ ∗ ) ∥ ∥ ∥ 2 ≲ ∥ v ∥ ∞ .

Proof. First note ∥ ∥ ∥ 1 √ d ( u ∗ ⊙ v -θ ∗ ) ∥ ∥ ∥ 2 ≤ 1 √ d ∥ u ∗ ∥ 2 ∥ v ∥ ∞ + 1 √ d ∥ θ ∗ ∥ 2 , 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

Now, we can proceed similarly to in the proof of Lemma 1. Let

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

After completing the squares, we obtain

so we can conclude

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 that

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:

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 consider

To show convergence in probability of this term, first fix δ > 0 . Then, we want to show

where the expectation is over ( V, Θ) ∼ Π t and G ∼ N (0 , 1) . It suffices to show the following two statements:

To show (26), we rely on a concentration inequality for the Moreau envelope of a Gaussian vector plus a bounded vector from [5]. First note that

where the second to last line follows from the change of variable θ = u ⊙ v β √ κ , and ℓ ( θ ) := λ ∥ ∥ ∥ β √ κ θ v ∥ ∥ ∥ 2 2 + µ ∑ d i =1 h ( β √ κθ i v i , 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 that

To show (27), note that

Hence,

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 as

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 as

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

C Solving the min-max problem

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:

Solving each of these equations for the quantity β 2 λ 2 E [ 1 b ∑ b i =1 Θ 2 i + β 2 κ ( τV 2 i + βλ ) 2 ] and then equating the two, we arrive at

which implies

Defining the auxiliary variable γ = τ/β , this yields the fixed point equation

Substituting this γ back into the first optimality condition in (28), we can express the optimal β in closed-form, in terms of γ :

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

D Further simulations

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, we consider the same sparse regression as in Section 3.1, i.e., with n = 250 , d = 2000 , σ = 0 . 1 , θ ∗ i i.i.d. ∼ 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 PL (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.

1

2

3

4

5

6

7

8

Iterations

10

-

2

ℓ

1

Test Error

ψ

=

|

uv

|

(theory)

ψ

=

|

uv

|

(sim)

ψ

= tanh

|

uv

|

(theory)

ψ

= tanh

|

uv

|

(sim)

ψ

=

u

2

(theory)

ψ

=

u

2

(sim)

ψ

= tanh

u

2

(theory)

ψ

= tanh

u

2

(sim)

(a) Predictions for

ψ

which are

not uniformly bounded

1

2

3

4

5

6

7

8

Iterations

0

.

002

0

.

004

0

.

006

0

.

008

0

.

010

Mean Squared 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) Predictions for the squared

test loss

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 a short (1-2 sentence) justification 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 it (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 a 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 a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a 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 1 and 2 directly reflect the claims from the abstract/introduction.

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 be 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 to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We provide detailed discussion of the strength of our technical assumptions after the statement of Theorem 1 and in Appendix D, and we further elaborate on these areas for improvement in the conclusion.

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.

The authors are encouraged to create a 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-specification, asymptotic approximations 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 recognize that individual actions in favor of transparency play an important 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 1 and 2, and we provide a complete proof of our results in the Appendix A 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental 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 in the corresponding 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 are provided or not.

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. 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 a 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 a way to access this model for reproducing the results or a 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 to have some path to 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 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.

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

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

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research presented in this work abides by the Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. 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 of algorithms in an abstract setup, and we do not believe our work has immediate social impacts or an immediate path toward such impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 as 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 (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe 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)?

Answer: [Yes]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. 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 use 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 is used and, if possible, include a 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 (e.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 use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be 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 new 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 not release new assets.

Researchers should communicate the details of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 as details about compensation (if any)?

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.

According to the NeurIPS Code of Ethics, 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

Question: 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?

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.

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThough 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.\n\nWEAKNESSES:\nAs 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).\n\nQUESTIONS:\nMinor 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\".\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nIt 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\n\nQUESTIONS:\n* 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).\n\nLIMITATIONS:\nI do not see any potential negative social impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}