Title: Estimating Generalization Performance Along the Trajectory of Proximal SGD in Robust Regression
Venue: NeurIPS 2024
Accepted: yes
Keywords: Robust regression

generalization error

stochastic gradient descent

early stopping

Stein's formula
Abstract: This paper studies the generalization performance of iterates obtained by Gradient Descent (GD), Stochastic Gradient Descent (SGD) and their proximal variants in high-dimensional robust regression problems. The number of features is comparable to the sample size and errors may be heavy-tailed. We introduce estimators that precisely track the generalization error of the iterates along the trajectory of the iterative algorithm. These estimators are provably consistent under suitable conditions. The results are illustrated through several examples, including Huber regression, pseudo-Huber regression, and their penalized variants with non-smooth regularizer. We provide explicit generalization error estimates for iterates generated from GD and SGD, or from proximal SGD in the presence of a non-smooth regularizer. The proposed risk estimates serve as effective proxies for the actual generalization error, allowing us to determine the optimal stopping iteration that minimizes the generalization error. Extensive simulations confirm the effectiveness of the proposed generalization error estimates.

Estimating Generalization Performance Along the Trajectory of Proximal SGD in Robust Regression

Abstract

This paper studies the generalization performance of iterates obtained by Gradient Descent (GD), Stochastic Gradient Descent (SGD) and their proximal variants in high-dimensional robust regression problems. The number of features is comparable to the sample size and errors may be heavy-tailed. We introduce estimators that precisely track the generalization error of the iterates along the trajectory of the iterative algorithm. These estimators are provably consistent under suitable conditions. The results are illustrated through several examples, including Huber regression, pseudo-Huber regression, and their penalized variants with non-smooth regularizer. We provide explicit generalization error estimates for iterates generated from GD and SGD, or from proximal SGD in the presence of a non-smooth regularizer. The proposed risk estimates serve as effective proxies for the actual generalization error, allowing us to determine the optimal stopping iteration that minimizes the generalization error. Extensive simulations confirm the effectiveness of the proposed generalization error estimates.

1 Introduction

Consider the linear model:

where y ∈ R n is the response vector, X ∈ R n × p is the design matrix, b ∗ ∈ R p is the unknown regression vector, and ε ∈ R n is the noise vector that we assume independent of X . The entries of ε may be heavy-tailed, for instance our working assumptions allow for infinite second moment.

For the estimation of b ∗ , we consider the following regularized optimization problem

where ρ : R → R is a data-fitting loss and g : R p → R is a regularization function. In the present robust regression setting, typical examples of ρ include the Huber [14] loss ρ ( r ; δ ) = δ 2 ∫ | r/δ | 0 min(1 , x ) d x , the Pseudo-Huber loss ρ ( r ; δ ) = δ 2 ( √ 1 + ( r/δ ) 2 -1) or other Lipschitz loss functions to combat the possible heavy-tails of the additive noise. Typical examples of penalty functions include the L1/Lasso [27] penalty g ( b ) = λ ∥ b ∥ 1 , group-Lasso penalty [29] for grouped variables, or their non-convex variants including for instance SCAD [12] or MCP [30].

In order to solve the optimization problem (2), practitioners resort to iterative algorithms, for instance gradient descent, accelerated gradient descent, stochastic gradient descent, and the corresponding proximal methods [20] in the presence of a non-smooth regularizer. Let the algorithm starts with

some initializer ̂ b 1 ∈ R p (typically ̂ b 1 = 0 ) followed by consecutive iterates ̂ b 2 , ̂ b 3 , . . . , where ̂ b t is typically obtained, for gradient descent and its variants as will be detailed below, from ̂ b t -1 and by applying an additive correction involving the gradient of the objective function. Our goal of this paper is to quantify the predictive performance of each iterate ̂ b t .

We assume throughout that the covariance E [ x i x ⊤ i ] = Σ of the feature vectors is finite. We measure the predictive performance of ̂ b t using the out-of-sample error

where x new is a new feature vector, independent of the data ( x i , y i ) i ∈ [ n ] and has the same distribution as x i . The above squared metric is used because the noise ε i (and thus y i ) is allowed to have infinite variance, and in this case the squared prediction error E [( x ⊤ new ̂ b t -y new ) 2 | ( x i , y i ) i ∈ [ n ] ] = + ∞ irrespective of the value of ̂ b t .

The paper proposes to estimate the out-of-sample error ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 of the t -th iterate using the right-hand side of the approximation

where ψ : R n → R is the derivative of ρ acting component-wise on each coordinate in R n and S s is a diagonal matrix of the form S s = ∑ i ∈ I s e i e ⊤ i where I s ⊂ [ n ] is the batch for the s -th stochastic gradient update and e i ∈ R n is the i -th canonical basis vector. Here the w t,s are quantities, introduced in Section 3.3 below, that can be computed from data and do not require the knowledge of Σ . The approximation (3) is made rigorous in Theorem 3.6, where the right-hand side is proved to be consistent (i.e., the difference between the two sides of the inequality converges to 0 in probability) for a first set of weights ( w s,t ) s<t , and in Theorem 3.7 where a second set of weights are proposed.

Because the right-hand side of (3) is observable from the data and the iterates ( ̂ b s ) s ≤ t are computed from the iterative algorithm, the approximation (3) lets us compare the out-of-sample error of iterates ̂ b t at different time t up to the additive term ∥ ε ∥ 2 /n (which does not depend on t nor on the choice of the iterative scheme or the choice of loss and penalty). It also lets us compare different tuning parameters, for instance learning rate, multiplicative parameter of the penalty function, batch size in Stochastic Gradient Descent (SGD). The right-hand side of (3) can serve as the criteria to choose the iteration number or tuning parameters that achieves the smallest out-of-sample error.

1.1 Related literature

Estimation of prediction risk of regression estimates has received significant attention in the last few decades. One natural avenue to estimate the generalization performance is to use V -fold crossvalidation or leave-one-out schemes. In the proportional regime of interest here, where dimension p and sample size n are of the same order, V -fold cross-validation with finite V , e.g., V = 5 , 10 is known to fail at consistently estimate the risk of the estimator trained on the full dataset [23, Figure 1]; this is simply explained because training with the biased sample size n ( V -1) /V may behave differently than training with the full dataset. Leave-one-out schemes, or drastically increasing V , requires numerous refitting and is thus computationally expensive.

This motivates computationally efficient estimates of the risk of an estimator trained on the full dataset without sample-splitting, including Approximate Leave-One-out (ALO) schemes [23] that do not rely on sample-splitting and refitting; see [1] and the references therein for recent developments. For ridge regression and other estimators constructed from the square loss, the Generalized Cross-Validation (GCV) [28] has been shown be to be effective, and it avoids data-splitting and refitting; it only needs to fit the full data once and then adjust the training error by a multiplicative factor larger than 1. Beyond ridge regression, the extension of GCV using degrees-of-freedom has been studied for Lasso regression [2, 3, 18, 9], and alternatives were developed for robust M-estimators [3, 4]. While ALO or GCV and its extensions are good estimators of the predictive risk of a solution ̂ b to the optimization problem (2), they are not readily applicable to quantify the prediction risk of iterates ̂ b t obtained by widely-used iterative algorithms such as gradient descent (GD), stochastic gradient

descent (SGD) or their proximal variants: ALO and GCV focus on estimating the final ( t → + ∞ ) iterate of the algorithm, when a solution ̂ b in (2) is found. Our goal in the present paper is to develop risk estimation methodologies along the trajectory of the algorithm.

Luo et al. [17] developed methods to estimate the cross-validation error of iterates that solves an empirical risk minimization problem. Their approach requires the Hessian of the objective function to be well-conditioned (i.e., the smallest and largest eigenvalues are bounded) along all iterates. This condition is not satisfied for the regression problems we consider in this paper, such as highdimensional robust regression with a Lasso penalty. In the context of least squares problems with both p and n being large, [7] studied the fundamental limits on the performance of first order methods, showing that these are dominated by a specific Approximate Message Passing algorithm. Paquette et al. [19] demonstrated that the dynamics of Stochastic Gradient Descent (SGD) become deterministic in the large sample and dimensional limit, providing explicit expressions for these dynamics when the design matrix is isotropic. Our work differs from [19] in two key ways: First, we address a more general regression problem incorporating a non-smooth regularizer, thereby considering both SGD and proximal SGD; second, we offer explicit risk estimates for each iteration, rather than focusing solely on the theoretical dynamics of the iterates. Celentano et al. [8] and [13] characterize the dynamical mean-field dynamics of iterative schemes, and identify that the limiting process involves a 'memory' kernel, describing how the dynamics of early iterates affect later ones.

Most recently, and most closely related to the present paper, [5] proposed risk estimate for iterates ̂ b t obtained by running gradient descent and proximal gradient descent methods for solving penalized least squares optimizations. However, [5] focuses exclusively on the square loss for ρ in (3), which is not readily applicable to robust regression with heavy tailed noise for which the Huber or other robust losses must be used. Bellec and Tan [5] is further restricted to gradient updates using the full dataset, which does not cover stochastic gradient descent. While several proof techniques used in the present paper are inspired by [5], we will explain in Remark 3.8 that directly generalizing [5] to SGD in robust regression leads to a poor risk estimate for small batch sizes. The proposal of the present paper leverages out of batch samples to overcome this issue.

For gradient descent for the square loss and without penalty, Patil et al. [21] demonstrates both the failure of GCV along the trajectory and the success of computationally expensive leave-one out schemes, and develops a proposal to reduce the computational cost. Finally, let us mention the works [10, 16] that characterize the dynamics of the iterates in phase retrieval and matrix sensing problems, assuming that a fresh batch of observations (independent of all previous updates) is used at each iteration. This is different from the usual SGD setting studied in the present paper where the observations used during a stochastic gradient update may be reused in future stochastic gradient updates, creating intricate probabilistic dependence between gradient updates at different iterations.

Robust regression is highly valuable in real data analysis due to its ability to handle heavy-tailed noise effectively, and we will see below that the use of stochastic gradient updates and data-fitting loss functions different from the square loss require estimates that have a drastically different structure that in the square loss case. The present paper develops generalization error estimates in situations where no consistent estimate have been proposed: (1) we develop generalization error estimates along the trajectory of iterative algorithms aimed at solving (2) for robust loss functions including the pseudo-Huber loss; (2) the estimates are applicable not only to gradient updates involving the full dataset (gradient descent and its variants), but also to SGD and proximal SGD where a random batch is used for each update.

2 Problem setup

The paper studies iterative algorithms aimed at solving the optimization problem (2). We consider the algorithm that generates iterates ̂ b t for t = 1 , 2 , ..., T according to the following iteration:

where S t ∈ R n × n is the diagonal matrix S t = ∑ i ∈ I t e i e ⊤ i for I t ⊂ [ n ] the t -th batch (independent of ( X , y ) ), where ϕ t : R p ↦→ R p and ψ : R n ↦→ R n are two functions and η t is the step size. Typically, ψ : R n → R n is the componentwise application of ρ ′ (where ρ is the data-fitting loss in (2)), and the matrix S t ∈ R n × n is diagonal with elements in { 0 , 1 } encoding the observations i ∈ [ n ]

used in the t -th stochastic gradient update. The presence of S t and possibly nonlinear function ψ is such that the above iteration scheme is not covered by previous related works including [5, 21], which only tackle S t = I n (full batch gradient updates) and ψ : R n → R n the identity map ( ρ in (2) restricted to be the square loss). The iterative scheme (4), on the other hand, covers SGD with robust loss functions.

In the next section, we first provide a few examples of algorithms encompassed in the general iteration (4). This includes Gradient Descent (GD), Stochastic Gradient Descent (SGD), and their corresponding proximal methods [20], Proximal GD and Proximal SGD. GD and SGD are widely used in practice, while the proximal methods are particularly useful for solving the optimization problem (2) with non-smooth regularizers.

2.1 Robust regression without penalty

If there are no penalties in (2), i.e., g ( b ) = 0 , then the minimization problem becomes

To solve this problem, provided ρ is differentiable, one may use gradient descent (SGD) and stochastic gradient descent (SGD).

Example 2.1 (GD) . The GD method consists of the following iteration:

where ψ is the derivative of ρ acting component-wise on its argument, and η t is the step size (also known as learning rate). For the least squares loss ρ ( x ) = x 2 / 2 , we have ψ ( u ) = u .

Example 2.2 (SGD) . Suppose at t -th iteration, we use the batch I t ⊂ [ n ] to compute the gradient,

where S t = ∑ i ∈ I t e i e ⊤ i and e i is the i -th canonical vector in R n . If I t = [ n ] for each t , then | I t | = n and S t = I n , hence this SGD method reduces to the GD method in (5).

2.2 Robust regression with Lasso penalty

Regularized regression is useful for high-dimensional regression problems where p is larger than n . We consider the Lasso penalty g ( b ) = λ ∥ b ∥ 1 to fight for the curse of dimensionality and obtain sparse estimates (our working assumptions, on the other hand, do not assume that the ground truth b ∗ is sparse). While GD and SGD are not directly applicable to solve the optimization problem (2) with Lasso penalty due to ∥ · ∥ 1 lacking differentiability at 0, Proximal Gradient Descent (Proximal GD) [20] and Stochastic Proximal Gradient Descent (Proximal SGD) can be used to solve this optimization with Lasso penalty.

Example 2.3 (Proximal GD) . For g ( b ) = λ ∥ b ∥ 1 in (2), the Proximal GD gives the following iterations:

where soft θ ( · ) applies the soft-thresholding u ↦→ sign( u )( | u | -θ ) + component-wise.

Example 2.4 (Proximal SGD) . Similar to the Proximal GD, the Proximal SGD consists of the following iterations:

Let ρ ′ : R n → R n be the function applies the derivative of ρ : R → R to each of its component, i.e., ρ ′ ( u ) = ( ρ ′ ( u 1 ) , ..., ρ ′ ( u n )) ⊤ . Then the above examples can be summarized in the following table with different definition of ψ , ϕ t , and S t .

To define the proposed estimators of the generalization error, we further define the following Jacobian matrices:

̸

Then, we have ˜ D t = I p for GD and SGD, and ˜ D t = ∑ j ∈ ˆ S t e j e ⊤ j for Proximal GD and Proximal SGD based on soft-thresholding, where ˆ S t = { j ∈ [ p ] : e ⊤ j ̂ b t +1 = 0 } .

3 Main results

Assumption 3.1. The design matrix X has i.i.d. rows from N p ( 0 , Σ ) for some positive definite matrix Σ satisfying 0 < λ min ( Σ ) ≤ 1 ≤ λ max ( Σ ) and ∥ Σ ∥ op ∥ Σ -1 ∥ op ≤ κ . We assume Var[ x ⊤ i b ∗ ] ≤ δ 2 , that is, the signal of the model (1) is bounded from above.

Assumption 3.2. The noise ε is independent of X and has i.i.d. entries from a fixed distribution independent of n, p , with E [ | ε i | ] ≤ δ , that is, bounded first moment.

Assumption 3.3. The data fitting loss ρ : R → R is convex, continuously differentiable and its derivative ψ is 1-Lipschitz and | ψ ( x ) | ≤ δ for all x ∈ R . The function ϕ t is 1-Lipschitz and satisfies ϕ t ( 0 ) = 0 . The matrices S t = ∑ i ∈ I t e i e ⊤ i , and | I t | ≥ c 0 n for some positive constant c 0 ∈ (0 , 1] . Let η max = max t ∈ [ T ] η t .

Huber loss and Psuedo-Huber loss all satisfy Assumption 3.3.

Assumption 3.4. The data fitting loss ρ is twice continuously differentiable with positive second derivative.

Assumption 3.5. The sample size n and feature dimension p satisfy p/n ≤ γ for a constant γ ∈ (0 , ∞ ) .

3.1 Intuition regarding the estimates of the generalization error

This subsection provides the intuition behind the definition of the estimates define below. For the sake of clarify, and in this subsection only, assume that

With the above working assumptions, the validity of the estimates defined below relies on the probabilistic approximation

which was developed in [3] for risk estimation purposes, but outside the context of iterative algorithms. Above, e j ∈ R p is the j -th canonical basis vector. In the present noiseless case with ε = 0 , the first term inside the squared norm in the right-hand side is equal to the residual y i -x ⊤ i ̂ b t , so that the above display resembles (3). Taking this probabilistic approximation for granted, to study the second term in the right-hand side, we must understand the derivatives of ̂ b t with respect to the entries ( x ij ) i ∈ [ n ] ,j ∈ [ p ] of X . In (4), each iterate is a relatively simple function of the previous ones, with the simplifications (7) this is ̂ b t +1 = ϕ t ( ̂ b t + X ⊤ S t ψ ( y -X ̂ b t ) /n ) . For t = 1 , given that ̂ b 1 is a constant initialization, ∂ ∂x ij ̂ b 2 = ˜ D 1 e j e ⊤ i S 1 ψ ( y -X ̂ b 1 ) /n -˜ D 1 X ⊤ S 1 D 1 e i ( ̂ b 1 -b ∗ ) j /n . We find in the proof, that when summing these quantities over j ∈ [ p ] , the second term involving

( ̂ b 1 -b ∗ ) j is negligible, and the same negligibility holds at later iterations with terms involving ( ̂ b t -b ∗ ) j (or any ( ̂ b s -b ∗ ) j , s ≤ t ). By performing a similar simple calculation at the next iteration, and ignoring these terms, we find with f 1 i = e ⊤ i S 1 ψ ( y -X ̂ b 1 ) and f 2 i = e ⊤ i S 2 ψ ( y -X ̂ b 2 ) by the chain rule

This reveals the weights ( w s,t ) s<t in (3) at iteration t = 2 and t = 3 . Wecould continue this further by successive applications of the chain rule, although for later iterations this unrolling of the derivatives, capturing the interplay between the Jacobians D t , ˜ D t and the stochastic gradient matrix S t , becomes increasingly complex. This recursive unrolling of the derivatives can be performed numerically at the same time as the computation of the iterates. On the other hand, for the mathematical proof, for the formal definition of the weights in (3) and for the proposed estimates of the generalization error, the matrix notation defined in the next subsection exactly captures this unrolling of the derivatives.

3.2 Formal matrix notation to capture recursive derivatives

We now set up the matrix notation that captures this recursive unrolling of the derivatives by the chain rule. Throughout, T is the final number of iterations. Define three block diagonal matrices D ∈ R nT × nT , ˜ D ∈ R pT × pT , and S ∈ R nT × nT by D = ∑ T t =1 ( ( e t e ⊤ t ) ⊗ D t ) , ˜ D = ∑ T t =1 ( ( e t e ⊤ t ) ⊗ ˜ D t ) , and S = ∑ T t =1 ( ( e t e ⊤ t ) ⊗ S t ) . Now we are ready to introduce the following matrices of size T × T :

where Γ = M -1 L ( Λ ⊗ I p ) ˜ D ∈ R pT × pT , L = ∑ T t =2 ( ( e t e ⊤ t -1 ) ⊗ I p ) , Λ = ∑ T t =1 η t | I t | e t e ⊤ t ,

Although notationally involved, the purpose of these matrices is just to formalize the recursive computation of the derivatives by the chain rule mentioned in Section 3.1.

3.3 Main results: estimating the generalization error consistently

For each iterate ̂ b t , define the target r t (generalization error) and its estimate ˆ r t by

where w t,s := e ⊤ t We s and W ∈ R T × T is the matrix defined in (8). The following shows that | ˆ r t -r t |→ P 0 (convergence to 0 in probability) under suitable assumptions.

Theorem 3.6 (Proved in Appendix C.1) . Let Assumptions 3.1, 3.3 and 3.5 be fulfilled. Then ∀ ϵ > 0 ,

If additionally Assumption 3.2 holds then E [min { 1 , ∥ ε ∥ n } ] → 0 , so that, as n, p → + ∞ while ( T, γ, η max , c 0 , δ, κ, ϵ ) are held fixed, the right-hand side converges to 0 and ˆ r t -r t converges to 0 in probability.

This establishes that ˆ r t is consistent at estimating r t . The statement E [min { 1 , ∥ ε ∥ /n } ] → 0 is equivalent to ∥ ε ∥ 2 /n 2 → P 0 (convergence in probability), and is proved in [22] under the assumption that E | ε i | < + ∞ with ε i i.i.d. from a fixed distribution; this allows Var[ ε i ] = + ∞ as long as the first moment is finite. The expression of W involves Σ , which is typically unknown in practice. Our next result provides a consistent estimate of W using quantities that do not require the knowledge of Σ .

We propose to estimate W by ˜ W def = ̂ K -1 ̂ A where ̂ K and ̂ A are the T × T matrices defined in (9)-(10). We define another estimate ˜ r t by replacing W in (11) with ˜ W = ̂ K -1 ̂ A :

where ˜ w t,s = e ⊤ t ˜ We s .

Theorem 3.7 (Proved in Appendix C.3) . Under Assumptions 3.1 and 3.3 to 3.5, for any ϵ > 0 ,

If additionally Assumption 3.2 holds then E [min { 1 , ∥ ε ∥ n } ] → 0 , so that, as n, p → + ∞ while ( T, γ, η max , c 0 , δ, κ, ϵ ) are held fixed, the right-hand side converges to 0 and ˜ r t -r t converges to 0 in probability.

This establishes the consistency of ˜ r t . The simulations presented next confirm that the two estimates ˜ r t and ˆ r t both are accurate estimates of r t . The estimate ˜ r t has the advantage of not relying on the knowledge of Σ and are recommended in practice.

Remark 3.8 . We highlight that directly generalizing the approach in [5] would lead to the approximation ˜ A ≈ ˜ KW , where ˜ A and ˜ K are given in (23) and (26), respectively. From ˜ A ≈ ˜ KW , obtaining an estimate of W requires inverting ˜ K . However, this inversion fails for SGD for small (but still very realistic) batch sizes of order 0 . 1 n in simulations (see Figure 4). The matrix ˜ K is lower triangular, and the reason for the lack of invertibility of ˜ K can be seen in the diagonal terms equal to Tr [ S t D t ] in (26), where S t ∈ { 0 , 1 } n × n is the diagonal matrix with 1 in position ( i, i ) if and only if the i -th observation is used in the t -th batch. This diagonal element of ˜ K can easily be small (or even 0) for small batches, if the batch only contains observations such that ( D t ) ii is 0 or small. Let ˜ r sub t denote the estimate of the same form as ˜ r t but using the weight matrix ˜ K -1 ˜ A instead. Simulation results in Figure 4 confirm that ˜ r sub t is suboptimal compared to our proposed ˜ r t . For SGD and proximal SGD, we solved this issue regarding the invertibility of ˜ K by using out-of-batch samples in the construction of ̂ K and ̂ A , in order to avoid S t in the diagonal elements of ̂ K in equation (31). This is the key to making these estimators work for SGD and proximal SGD, and this use of out-of-batch samples is new compared to [5] (which only tackles the square loss with full-batch gradients).

Remark 3.9 . The constant C ( T, γ, η max , c 0 , δ, κ ) in the above results is not explicit. Inspection of the proof reveals that the dependence of this constant in T is currently T T , allowing T of order log( n ) / log log n before the bound becomes vacuous. Improving this dependence in T appears challenging and possibly out of reach of current tools, even for the well-studied Approximate Message Passing (AMP) algorithms. The papers [25, 24] feature for instance the same log( n ) / log log n dependence for approximating the risk of AMP. The preprint [15] offers the latest advances on the dependence on T in the bounds satisfied by AMP. It allows T ≍ poly ( n ) while still controlling certain AMP related quantities, although for the risk [15, equations (16)-(17)] the condition required on T is still logarithmic in n . This suggests that advances on this front are possible, at least for isotropic design and specific loss and regularizer such as those studied in [15]: Lasso or Robust M-estimation with no regularizer. Since these latest advances in [15] are obtained for specific estimates (Lasso or Robust M-estimation with no regularizer), it may be possible to follow a similar strategy and improve our bounds for specific examples of iterative algorithms closer to AMP, or algorithms featuring only separable losses and penalty. We leave such improvements for specific examples for future work, as the goal of the current paper is to cover a general framework allowing iterations of the form (4) with little restriction on the nonlinear functions except being Lipschitz.

4 Simulation

In this section, we present numerical experiments to assess the performance of the proposed risk estimates. All necessary code for reproducing these experiments is provided in the supplementary material and is publicly available in the GitHub repository https://github.com/kaitan365/ SGD-generlization-errors . Our goal is to compare the performance of the proposed risk estimates with the true risk r t for different regression methods and iterative algorithms.

We generate the dataset ( X , y ) from the linear model (1), that is, y = Xb ∗ + ε . Here, the rows of X ∈ R n × p are sampled from a centered multivariate normal distribution with covariance matrix Σ = I p . The noise vector ε consists of i.i.d. entries drawn from a t distribution with two degrees of freedom so that the noise variance is infinite. The true regression vector b ∗ ∈ R p is chosen with p/ 20 nonzero entries, set to a constant value such that the signal strength ∥ b ∗ ∥ 2 equals 10.

We explore two scenarios of the ( n, p ) pairs and corresponding iterative algorithms:

(i) ( n, p ) = (10000 , 5000) : In this configuration, with n much larger than p , we examine Huber regression and Pseudo-Huber regression (without penalty or soft-thresholding). Both the gradient descent (GD) and stochastic gradient descent (SGD) algorithms are implemented for each type of regression.

(ii) ( n, p ) = (10000 , 12000) : Here, we investigate Huber regression and Pseudo-Huber regression with an L1 penalty, λ ∥ b ∥ 1 ( λ = 0 . 002 ) and corresponding soft-thresholding step. For each penalized regression, we employ the Proximal Gradient Descent (Proximal GD) and Stochastic Proximal Gradient Descent (Proximal SGD) algorithms.

̸

In all algorithms, we start with the initial vector ̂ b 1 = 0 p and proceed with a fixed step size η = (1 + √ p/n ∗ ) -2 where n ∗ = n for GD and proximal GD, and n ∗ = n/ 5 for SGD and proximal SGD. We run each algorithm for T = 100 steps. For SGD and Proximal SGD, batches I t ⊂ { 1 , 2 , . . . , n } are randomly sampled without replacement and independently of ( X , y , ( I s ) s = t ) , each with cardinality | I t | = n 5 .

A crucial component of the proposed risk estimates ˆ r t and ˜ r t involve the weight matrices W and ˜ W . The matrix W is defined in Theorem 3.6, and ˜ W = ̂ K -1 ̂ A is defined in Theorem 3.7. We employ Hutchinson's trace approximation to compute W , ̂ A , and ̂ K . This implementation is computationally efficient. We refer readers to [5, Section 4] for more details.

Recall that we have proposed two estimates for r t = ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 + ∥ ε ∥ 2 /n , one is ˆ r t in Theorem 3.6 which requires knowing Σ = E [ x i x ⊤ i ] , and the other is ˜ r t in Theorem 3.7 which does not need Σ . Since the quantity ∥ ε ∥ 2 /n remains constant along the algorithm trajectory, we only focus on the estimation of ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 . We repeat each numerical experiment 100 times and present the aggregated results in Figures 1, 2 and 3.

In Figure 1, we focus on the scenario with ( n, p ) = (10000 , 5000) , and plot the actual risk ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 , and its two estimates ˆ r t -∥ ε ∥ 2 /n and ˜ r t -∥ ε ∥ 2 /n along with the 2 standard error bar for GD and SGD algorithms applied to both Huber and Pseudo Huber regression. In Figure 2, we focus on the scenario with ( n, p ) = (10000 , 12000) , and present the risk curves for the Proximal GD and Proximal SGD algorithms applied to both L1-penalized Huber regression and Pseudo-Huber regression.

Figure 1 and Figure 2 confirm the three curves are in close agreement, indicating that the proposed estimates ˆ r t -∥ ε ∥ 2 /n and ˜ r t -∥ ε ∥ 2 /n are consistent estimates of the actual risk ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 .

The two estimates closely capture the risk ∥ Σ 1 / 2 ( ̂ b t -b ∗ ) ∥ 2 over the entire trajectory of the algorithms. For GD and Proximal GD, the risk curves exhibit a U-shape, first decreasing and then increasing, and the estimates ˆ r t and ˜ r t closely capture this pattern. This suggests that the proposed estimates are reliable and can be used to monitor the risk of the iterates and find the optimal iteration (the iteration minimizing the generalization error) along the trajectory of the algorithm.

Additional experiments: varying step sizes for different iterations. We also conduct simulations to investigate the accuracy of the proposed risk estimates in a setting with varying step size. We consider two types of step sizes: 1). η t = 1 if t is odd, and η t = 0 if t is even; 2). η t = 1 if t is odd,

0

100

200

300

400

500

Iteration number

t

2

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

2

4

6

8

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

2

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

2

4

6

8

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

(c) Risk curves of GD for Pseudo-Huber regression (d) Risk curves of SGD for Pseudo-Huber regression

and η t = 0 . 5 if t is even. While the above choices of step size are not preferred in practice, here the goal is show that the proposed risk estimates is able to accurately capture the dynamics of the risk even when the step size changes along the trajectory of the algorithm. For instance, the first choice of step size should produce a risk curve that is flat when t is even. The results are presented in Figure 3, illustrating that the risk estimates accurately capture the flat segments of the true risk curve.

Additional experiments: the estimate ˜ r sub t is suboptimal. We compare the performance of ˜ r sub t with our proposed estimates in Huber regression with ( n, p, T ) = (4000 , 1000 , 20) and batch size | I t | = n/ 10 and η t = 0 . 2 for all t ∈ [ T ] . It is clear from Figure 4 that ˜ r t is more accurate than the suboptimal estimator ˜ r sub t , especially when t increases.

0

100

200

300

400

500

Iteration number

t

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

0

100

200

300

400

500

Iteration number

t

2

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

(c) Proximal GD for Pseudo-Huber regres- sion

(d) Proximal SGD for Pseudo-Huber regres- sion

5 Discussion

This paper proposes a novel risk estimate for the generalization error of iterates generated by the proximal GD and proximal SGD algorithms in robust regression. The proposed risk estimates accurately capture the predictive risk of the iterates along the trajectory of the algorithms, and are provably consistent (Theorems 3.6 and 3.7). Three matrices in R T × T in (8)-(10) reveal the interplay between the squared risk, the residuals and the gradients, so that the approximation (3) holds. This structure is different from the square loss case studied in [5] where only two matrices (inverse of each other) are sufficient.

Let us mention some open questions along with potential future research directions. The first question regards the probabilistic model: we currently assume Gaussian features x i , and it would be of interest to study the extension in which our consistency results are universal, allowing non-Gaussian feature distributions. Second, it is of interest to extend the current estimates to more general optimization problems of the form (2) with non-smooth data-fitting loss, for instance the Least Absolute Deviation loss ∥ y -Xb ∥ 1 . In this case the gradient does not exist at the origin, which calls for different algorithms than the GD and SGD variants presented here, for instance the Alternating Direction Method of Multipliers (ADMM) [6]. It is of independent interest to derive the risk estimates for iterates obtained by such primal-dual methods.

Acknowledgments

P. C. Bellec acknowledges partial support from the NSF Grants DMS-1945428 and DMS-2413679. The authors acknowledge the Office of Advanced Research Computing (OARC) at Rutgers, The State University of New Jersey for providing access to the Amarel cluster and associated research computing resources that have contributed to the results reported here. The authors thank the anonymous reviewers for their valuable comments and suggestions that helped improve the presentation of the paper.

References

Arnab Auddy, Haolin Zou, Kamiar Rahnamarad, and Arian Maleki. Approximate leave-one-out cross validation for regression with l1 regularizers. In International Conference on Artificial Intelligence and Statistics , pages 2377-2385. PMLR, 2024.

Mohsen Bayati, Murat A Erdogdu, and Andrea Montanari. Estimating lasso risk and noise level. In Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013.

Pierre C Bellec. Out-of-sample error estimate for robust m-estimators with convex penalty. Information and Inference: A Journal of the IMA , 12(4):2782-2817, 10 2023.

Pierre C Bellec and Yiwei Shen. Derivatives and residual distribution of regularized m-estimators with application to adaptive tuning. In Conference on Learning Theory , pages 1912-1947. PMLR, 2022.

Pierre C Bellec and Kai Tan. Uncertainty quantification for iterative algorithms in linear models with application to early stopping. arXiv preprint arXiv:2404.17856 , 2024.

Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends® in Machine learning , 3(1):1-122, 2011.

Michael Celentano, Andrea Montanari, and Yuchen Wu. The estimation error of general first order methods. In Conference on Learning Theory , pages 1078-1141. PMLR, 2020.

Michael Celentano, Chen Cheng, and Andrea Montanari. The high-dimensional asymptotics of first order methods with random data. arXiv preprint arXiv:2112.07572 , 2021.

Michael Celentano, Andrea Montanari, and Yuting Wei. The Lasso with general Gaussian designs with applications to hypothesis testing. Ann. Statist. , 51(5):2194-2220, 2023. ISSN 0090-5364.

Kabir Aladin Chandrasekher, Ashwin Pananjady, and Christos Thrampoulidis. Sharp global convergence guarantees for iterative nonconvex optimization: A gaussian process perspective. arXiv preprint arXiv:2109.09859 , 2021.

Kenneth R Davidson and Stanislaw J Szarek. Local operator theory, random matrices and banach spaces. Handbook of the geometry of Banach spaces , 1(317-366):131, 2001.

Jianqing Fan and Runze Li. Variable selection via nonconcave penalized likelihood and its oracle properties. J. Amer. Statist. Assoc. , 96(456):1348-1360, 2001. ISSN 0162-1459.

Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zdeborova. Rigorous dynamical mean-field theory for stochastic gradient descent methods. SIAM J. Math. Data Sci. , 6(2):400-427, 2024.

Peter J Huber. Robust statistics , volume 523. John Wiley & Sons, 2004.

Gen Li and Yuting Wei. A non-asymptotic distributional theory of approximate message passing for sparse and robust regression. arXiv preprint arXiv:2401.03923 , 2024.

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.

Yuetian Luo, Zhimei Ren, and Rina Barber. Iterative approximate cross-validation. In International Conference on Machine Learning , pages 23083-23102. PMLR, 2023.

Léo Miolane and Andrea Montanari. The distribution of the Lasso: uniform control over sparse balls and adaptive parameter tuning. Ann. Statist. , 49(4):2313-2335, 2021. ISSN 0090-5364.

Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. Sgd in the large: Average-case analysis, asymptotics, and stepsize criticality. In Conference on Learning Theory , pages 3548-3626. PMLR, 2021.

Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and trends® in Optimization , 1(3):127-239, 2014.

Pratik Patil, Yuchen Wu, and Ryan Tibshirani. Failures and successes of cross-validation for early-stopped gradient descent. In International Conference on Artificial Intelligence and Statistics , pages 2260-2268. PMLR, 2024.

Iosif Pinelis. Large deviations: Growth of empirical average of iid non-negative random varialbes with infinite expectations?, 2021. URL https://mathoverflow.net/q/390939 . Accessed: 2024-05-21.

Kamiar Rahnama Rad and Arian Maleki. A scalable estimate of the out-of-sample prediction error via approximate leave-one-out cross-validation. J. R. Stat. Soc. Ser. B. Stat. Methodol. , 82 (4):965-996, 2020. ISSN 1369-7412.

Cynthia Rush. An asymptotic rate for the lasso loss. In International Conference on Artificial Intelligence and Statistics , pages 3664-3673. PMLR, 2020.

Cynthia Rush and Ramji Venkataramanan. Finite sample analysis of approximate message passing algorithms. IEEE Trans. Inform. Theory , 64(11):7264-7286, 2018.

Kai Tan, Gabriel Romon, and Pierre C. Bellec. Noise covariance estimation in multi-task highdimensional linear models. Bernoulli , 30(3):1695 - 1722, 2024. doi: 10.3150/23-BEJ1644.

Robert Tibshirani. Regression shrinkage and selection via the lasso. J. R. Stat. Soc. Ser. B Stat. Methodol. , 58(1):267-288, 1996.

Grace Wahba. A comparison of gcv and gml for choosing the smoothing parameter in the generalized spline smoothing problem. Ann. Statist. , pages 1378-1402, 1985.

Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. J. R. Stat. Soc. Ser. B Stat. Methodol. , 68(1):49-67, 2006. ISSN 1369-7412.

Cun-Hui Zhang. Nearly unbiased variable selection under minimax concave penalty. Ann. Statist. , 38(2):894-942, 2010. ISSN 0090-5364.

Supplementary Material of 'Estimating Generalization Performance Along the Trajectory of Proximal SGD in Robust Regression'

A Additional simulation results

The following figures illustrate the proposed risk estimates accurately estimate the trajectory of the risk even when the step size changes at every step.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

Iteration number

t

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

Iteration number

t

2

4

6

8

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

(a)

η

t

= 1

if

t

is odd, and

η

t

= 0

if

t

is even.

(b)

η

t

= 1

if

t

is odd, and

η

t

= 0

.

5

if

t

is even.

Figure 4 compares the performance of the proposed estimators ˆ r t , ˜ r t and the estimator ˜ r sub t generalized directly from [5]. It confirms that our proposed estimators outperforms ˜ r sub t .

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

Iteration number

t

2

4

6

8

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

r

sub

t

2

/

n

(a) Huber regression

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

Iteration number

t

2

3

4

5

6

7

8

9

10

Value

1/2

(

b t

b

*

)

2

r

t

2

/

n

r

t

2

/

n

r

sub

t

2

/

n

(b) Pseudo-Huber regression

B Auxiliary Results

Throughout, we define

Note that in the above E , H are not observable since ε and b ∗ are unknown. However, F and R are observable and can be easily computed once the iterates ( ̂ b t ) t ∈ [ T ] are calculated.

B.1 Change of variables

In this section, we conduct the change of variable to simplify the proof. Specifically, we view the linear model y = Xb ∗ + ε as a model with design matrix G and the regression vector θ ∗ , i.e.

This way, the design matrix G has i.i.d. entries from standard normal distribution. Using the same argument in [5, Appendix D], we can show that the matrices H , F , ̂ A , ̂ K remains the same under the change of variable. Therefore, we can prove the main results using the model with design matrix G and the regression vector θ ∗ . In other words, we assume without of loss of generality that the design matrix X has i.i.d. N (0 , 1) , or equivalently that the independent rows of X are normally distributed with covariance Σ = I p . We prove the main results using Σ = I p , and the results for general Σ follow by this change of variable with the constant C ( T, γ, η max , c 0 , δ ) appearing in the bounds depending additionally on κ (the upper bound of the condition number of Σ from Assumption 3.1).

B.2 Derivative formulae

In this section, we present derivative formulae that will be useful in later proofs. The following formulas differ from [5] due to the use of robust loss functions and the application of SGD with random batches at each iteration. The formulae are also significantly more complex than in the case of regularized M-estimators [3, 4].

Lemma B.1 (Proved in Appendix D.1) . Let ( ̂ b t ) t ∈ [ T ] be the iterates generated from the recursion (4) and the initial value ̂ b 1 is independent of X . Then the derivative of ̂ b t with respect to X is given by

where Γ = M -1 L ( Λ ⊗ I p ) ˜ D , L = ∑ T t =2 ( ( e t e ⊤ t -1 ) ⊗ I p ) , Λ = ∑ T t =1 η t | I t | e t e ⊤ t , and

Recall F = [ S 1 ψ ( y -X ̂ b 1 ) , ..., S T ψ ( y -X ̂ b T )] , we have F lt = e ⊤ l S t ψ ( y -X ̂ b t ) . The following two corollaries are a direct consequence of Lemma B.1.

Lemma B.2 (Proved in Appendix D.2 ) . Under the same conditions of Lemma B.1. Let F lt = e ⊤ l Fe t = e ⊤ l S t ψ ( y -X ̂ b t ) , we have

where

LemmaB.3 (Proved in Appendix D.3) . Let ˜ F = [ ψ ( y -X ̂ b 1 ) , ..., ψ ( y -X ̂ b T )] and ˜ F l,t = e ⊤ l ˜ Fe t . Under the same conditions of Lemma B.1. We have

where

Definition B.4. Define the matrices Υ 1 ∈ R p × T , Υ 2 ∈ R n × T , Υ 3 ∈ R T × T , Υ 4 ∈ R T × T , Υ 5 ∈ R T × T by the identities

where the matrices ˜ K , ̂ A , ˜ A , W are defined as follows

The matrices Υ 1 , Υ 2 , . . . are negligible in the sense that their Frobenius norms are of smaller orders compared to their preceding terms in (18)-(22). We provide the bounds in next lemma, which is obtained by deriving alternative expressions for Υ 1 , ..., Υ 5 in Appendix D.4.

Lemma B.5 (Proved in Appendix D.6) . Under the same conditions of Theorem 3.6 with Σ = I p , we have

We further define a few matrices of size T × T :

The next lemma provides the moment bounds for the Frobenius norm of these matrices.

Lemma B.6 (Proved in Appendix D.7) . Under the same conditions of Theorem 3.6, we have

almost surely, where E [ · | ε ] is the conditional expectation given ε .

We are able to prove the main theorems using Lemma B.6.

C Proof of main results

C.1 Proof of Theorem 3.6

It suffices to prove this theorem for the case Σ = I p . When Σ = I p , the result can be derived using a change of variables argument, as outlined in Appendix B.1. By basic algebra, we have

̸

Notice that r t = ∥ ̂ b t -b ∗ ∥ 2 + ∥ ε ∥ 2 /n is the t -th diagonal entry of ( H ⊤ H + E ⊤ E /n ) , and ˆ r t is the t -th diagonal entry of ( R + FW ⊤ ) ⊤ ( R + FW ⊤ ) /n . Since ∥ E ∥ F = √ T ∥ ε ∥ , using the previous display that conditionally on ε , we have

Using the moment bounds of Θ 5 and Θ 6 in Lemma B.6, we have

Furthermore, if E [ | ε i | ] is finite, we have by [22] that ∥ ε ∥ /n → P 0 (convergence in probability) if ε has i.i.d. entries with a fixed distribution independent of n . Under this assumption, the right-hand side of the previous display converges to 0 in probability. By enlarging the constant if necessary, assume C ( T, γ, η max , c 0 , δ ) ≥ 1 . To obtain a quantitative bound, by the conditional version of Markov's inequality, for any ϵ > 0 , and almost surely with respect to ε that

Taking expectation with respect to ε , we obtain

with E [min { 1 , ∥ ε ∥ n } ] → 0 (equivalently, ∥ ε ∥ /n → P 0 ) if the entries of ε are i.i.d. with a fixed distribution independent of n with E [ | ε i | ] < + ∞ by [22]. This finishes the proof of Theorem 3.6.

C.2 Operator norm bound on ̂ K

We first recall the definition of ̂ K from (10) in the main text:

Define two events: Ω 1 = { X ∈ R n × p : ∥ X ∥ op / √ n ≤ 2 + √ p/n } and Ω 2 = {|{ i ∈ [ n ] : | ε i | ≤ M }| ≥ 2 n 3 } , where M is a large enough constant such that P ( | ε i | > M ) ≤ 1 / 6 .

Lemma C.1. Under the same conditions of Theorem 3.6 with Σ = I p , we have in the event Ω ∗ = Ω 1 ∩ Ω 2 that

Furthermore, Ω ∗ has probability at least 1 -e -n/ 18 -e -n/ 2 .

Proof of Lemma C.1. Under Assumptions 3.1 and 3.5, we know that P (Ω 1 ) ≥ 1 -e -n/ 2 from [11]. In the event Ω 1 , we have by Lemma D.2 that

Markov's inequality further implies

In other words, |{ i ∈ [ n ] : ∥ x ⊤ i H ∥ 2 ≤ 3 C ∗ }| ≥ 2 n 3 . Recall that M is such that P ( | ε i | > M ) ≤ 1 / 6 . By Hoeffding's inequality, we have

Taking a = 1 / 6 , we have P ( 1 n ∑ n i =1 1 {| ε i | > M } ≥ P ( | ε i | > M ) + 1 / 6 ) ≤ e -n/ 18 . Furthermore, using |{ i ∈ [ n ] : | ε i | > M }| = ∑ n i =1 1 {| ε i | > M } , we have

Therefore,

Equivalently, we have P ( |{ i ∈ [ n ] : | ε i | ≤ M }| ≥ 2 n 3 ) ≥ 1 -e -n/ 18 . That is, at least 2 n 3 of the entries of ε are bounded by M with probability at least 1 -e -n/ 18 .

Recall that Ω 2 = {|{ i ∈ [ n ] : | ε i | ≤ M }| ≥ 2 n 3 } , then P (Ω 2 ) ≥ 1 -e -n/ 18 . Hence, P (Ω 1 ∩ Ω 2 ) ≥ 1 -e -n/ 18 -e -n/ 2 . In the event Ω 1 ∩ Ω 2 , the set

has size at least n 3 . For any i ∈ ˆ I and t ∈ [ T ] , we have

By the definition of D t , under Assumption 3.4, we have

Here c ∗ is a constant depending on ρ, M, C ∗ only.

By the definition of ̂ K in (31), ̂ K /n is a lower triangular matrix with diagonal entries equal to Tr( D t ) /n . It is invertible if and only if all its diagonal entries are non-zero. Therefore, in the event Ω 1 ∩ Ω 2 , we have ̂ K /n is invertible.

where ̂ L = ∑ n i =1 ( I T ⊗ e ⊤ i ) D ( I T ⊗ X ) Γ ( I T ⊗ X ⊤ ) SD ( I T ⊗ e i ) . Here ̂ L is a strictly lower triangular matrix. Using the upper bound of ∥ Γ ∥ op in Lemma D.4, we have ∥ ̂ L ∥ op ≤ nC ( T, γ, η max , c 0 ) in the event Ω 1 . Now we rewrite ̂ K -1 as

Notice that ̂ L ̂ Λ -1 ∈ R T × T is a strictly lower triangular matrix, thus

By the triangle inequality,

Therefore, in the event Ω 1 ∩ Ω 2 which has probability P (Ω 1 ∩ Ω 2 ) ≥ 1 -e -n/ 18 -e -n/ 2 ,

Lemma C.2. Under the same conditions of Theorem 3.7 with Σ = I p , we have

Proof of Lemma C.2. By the definition of ̂ K in (31), using ∥D∥ op ≤ 1 and ∥S∥ op ≤ 1 , we have

Similarly, by the definition of ̂ A in (24), we have

∥

̂

A

∥

op

n

∥

X

∥

2

op

∥

Γ

∥

op

.

Last, by the definition of W in (25), we have ∥ W ∥ op ≤ n ∥ Γ ∥ op .

Lemma C.3. Under the same conditions of Theorem 3.6 with Σ = I p , we have

Proof of Lemma C.3. First, using the definitions of Θ 1 , Θ 2 , Θ 4 in (27), we have

Hence,

Here the last line uses the upper bounds of E [ ∥ Θ k ∥ F | ε ] from Lemma B.6, and the bound of E [ ∥ ̂ K ∥ op | ε ] follows from Lemma C.2 and the bound of ∥ Γ ∥ op from Lemma D.4. This proves the first inequality.

For the second inequality, we define

where ˇ K = n I T -∑ n i =1 ( I T ⊗ ( e ⊤ i X )) Γ ( I T ⊗ X ⊤ ) SD ( I T ⊗ ( X ⊤ e i )) . Using similar argument that proves Lemma B.6, we can obtain the following bound of ˇ Θ 1 , ˇ Θ 2 , ˇ Θ 4 .

By the definitions of ˇ Θ 1 , ˇ Θ 2 , ˇ Θ 4 , we have

Therefore, conditional on ε , we have

This finishes the proof of Lemma C.3.

C.3 Proof of Theorem 3.7

In the event Ω ∗ = Ω 1 ∩ Ω 2 , we know from Lemma C.1 that ̂ K is invertible and ∥ ̂ K -1 ∥ op ≤ n -1 C . Define ˜ W = ̂ K -1 ̂ A . Using R + F ˜ W ⊤ = R + FW ⊤ + F ( ˜ W -W ) ⊤ , we have

We have by the triangle inequality

Recall that in the event Ω ∗ , we have ∥ ̂ K -1 ∥ op ≤ n -1 C . Using ̂ A -̂ KW = ̂ K ( ˜ W -W ) and Lemma C.3, we have

According to Lemma C.3 and the bound of ∥ ̂ K -1 ∥ op in Lemma C.1, the first conditional expectation is bounded from above by

Using the bound of ∥ K -1 ∥ op in Lemma C.1, the bound of ∥ W ∥ op in Lemma C.2, and the bound of E [ ∥ F ⊤ F ( ̂ A -̂ KW ) ⊤ ∥ F | ε ] in Lemma C.3, the second conditional expectation is bounded from above by

Similarly, the third conditional expectation is bounded from above by

In summary, we have

Since ˜ r t and ˆ r t are the t -th diagonal entries of ( R + F ˜ W ⊤ ) ⊤ ( R + F ˜ W ⊤ ) /n and ( R + FW ⊤ ) ⊤ ( R + FW ⊤ ) /n , respectively, we have

Using the same argument in the proof of Theorem 3.6, we have for any ϵ > 0 ,

Taking expectation with respect to ε , we have

Using the union bound and P (Ω ∗ ) ≥ 1 -e -n/ 18 -e -n/ 2 ≥ 1 -2 e -n/ 18 , we obtain

Using the triangle inequality and the tail probability of | ˆ r t -r t | in Theorem 3.6, we have

D Proof of results in Appendix B.2

D.1 Proof of Lemma B.1

By assumption, we know ̂ b 1 is independent of X and ̂ b t +1 = ϕ t ( ̂ b t + η t | I t | X ⊤ S t ψ ( y -X ̂ b t ) ) from (4). Recall that D t = ∂ ψ ( u ) ∂ u ∣ ∣ u = y -X ̂ b t and ˜ D t = ∂ ϕ t ( v ) ∂ v ∣ ∣ v = ̂ b t + η t | I t | X ⊤ S t ψ ( y -X ̂ b t ) . Let a dot denote the derivative with respect to x ij . By product rule and chain rule and using y -X ̂ b t = ε -X ( ̂ b t -b ∗ ) , we have

where the last line uses F t = S t ψ ( y -X ̂ b t ) and H t = ̂ b t -b ∗

Let P t = ˜ D t ( I p -η t | I t | X ⊤ S t D t X ) and a t = η t | I t | ˜ D t ( ˙ X ⊤ F t -X ⊤ S t D t ˙ X H t ) , we can rewrite the above recursion of ˙ b t as a linear system:

Solving the above system, we have ˙ b t = ( e ⊤ t ⊗ I p ) M -1 La . Since ˙ X = ∂ X ∂x ij = e i e ⊤ j , a t can be further simplified as

e

(

T

t

=

∑

=1

η

t

|

Using

D

Λ

=

T

t

=1

I

t

|

e

t

e

⊤

t

t

e

⊤

t

)

⊗

(

, we have

Plugging this expression for a into ˙ b t = ( e ⊤ t ⊗ I p ) M -1 La gives

This finishes the proof of Lemma B.1.

D.2 Proof of Lemma B.2

By definition, F lt = e ⊤ l S t ψ ( y -X ̂ b t ) . By the chain rule of differentiation, we have

Notice that ( e t ⊗ ( D t S t )) = DS ( e t ⊗ I n ) . The desired formula then follows by plugging in the expression of ∂ ̂ b t ∂x ij in Lemma B.1.

D.3 Proof of Lemma B.3

The desired identity follows by

and the expression of ∂ ̂ b t ∂x ij in Lemma B.1.

∑

D

)

,

D

=

T

(

t

=1

∑

e

e

⊤

t

)

⊗

D

)

,

S

=

(

t

˜

t

˜

t

T

t

=1

∑

(

(

e

t

e

⊤

t

)

⊗

S

t

)

, and

D.4 Alternative expressions for the matrices defined in Definition B.4

This section derives alternative expressions for the matrices Υ 1 , ..., Υ 5 defined in Definition B.4.

We first study Υ 1 in (18). Using F = ∑ T t =1 Fe t e ⊤ t , we have by Lemma B.2

Now we compute the two terms in the right-hand side of (37). For the first term, using the expression of D lt ij in Lemma B.2,

Next, we compute the second term in the right-hand side of (37). Using the expression of ∆ lt ij in Lemma B.2,

The identity (18) then follows by substituting the above two expressions into (37).

To study Υ 2 in (19), we use a similar procedure. Using the mixed property of Kronecker product and the fact that the transpose of a scalar remains the same, we have

where W = ∑ j ( I T ⊗ e ⊤ j ) Γ ( I T ⊗ e j ) and Υ 2 = ∑ j (( e ⊤ j H ) ⊗ I n ) DS ( I T ⊗ X ) Γ ⊤ ( I T ⊗ e j ) .

To study Υ 3 in (20), we use the product rule of differentiation and (18)-(19):

For Υ 4 in (21), by the product rule,

We then compute the first term of the last line as follows

where

Combining the above pieces, we have

This provides an alternative expression for Υ 4 in (21).

Last, we study Υ 5 in (22). We have

It remains to compute the third term in the last display. Using the fact that ˜ F = ∑ n l =1 ∑ T t =1 e l e ⊤ l ˜ Fe t e ⊤ t , we have by Lemma B.3 that

Using

we find

We now compute the two terms in the above display. For the first term, we have

For the second term, we have

Thus, we have established

This provides an alternative expression for Υ 5 in (22).

D.5 Preparation results for proving Lemmas B.5 and B.6

Lemma D.1 (Moment bounds for H , F , ˜ F ) . Under Assumptions 3.1, 3.3 and 3.5 with Σ = I p . Let H , F be defined in (13) , we have for any finite integer k ,

Proof of Lemma D.1. For the first inequality, according to Assumption 3.1 and Σ = I p , X has i.i.d. standard normal entries. By [11, Theorem II.13], there exists a random variable z ∼ N (0 , 1) such that ∥ X ∥ op ≤ √ n + √ p + z almost surely. Under Assumption 3.5 that p/n ≤ γ , we have E [ ∥ X / √ n ∥ k op ] ≤ C ( γ, k ) for any finite integer k .

For the second inequality, since ∥ H ∥ 2 F = ∑ T t =1 ∥ ̂ b t -b ∗ ∥ 2 , it suffices to bound ∥ ̂ b t -b ∗ ∥ 2 for each t ∈ [ T ] . Define the sequence of scalars a t def = max {∥ ̂ b t ∥ , δ } . Since ̂ b t = ϕ t -1 ( ̂ b t -1 + η t n t X ⊤ S t ψ ( y -X ̂ b t -1 )) where n t := | I t | . Note that ∥ ψ ( y I t -X I t b ) ∥ ≤ √ | I t | δ by Assumption 3.3, we have

Since a t = max {∥ ̂ b t ∥ , δ } and δ ≤ a t -1 , we have

Notice a 1 = δ since ̂ b 1 = 0 p , we obtain

Hence, using the inequality ∥ ̂ b t -b ∗ ∥ 2 ≤ 2 ∥ ̂ b t ∥ 2 +2 ∥ b ∗ ∥ 2 , we have

Taking conditional expectation on both sides given ε , the desired moment bound for H follows from the moment bound for ∥ X / √ n ∥ op .

For the third inequality, since | ψ ( x ) | ≤ δ from Assumption 3.3, we have ∥ ψ ( u ) ∥ ≤ nδ 2 for any u ∈ R n . By the definitions of F and ˜ F , we have

Since the above display holds for any ε , it implies the desired conditional moment bounds for F , ˜ F .

Lemma D.2 (Frobenius norm bound for XH ) . Under Assumptions 3.1, 3.3 and 3.5, we have

with probability at least 1 -exp( -n/ 2) .

Proof of Lemma D.2. By [11, Theorem II.13], under Assumption 3.1 with Σ = I p , we have

Using ∥ XH ∥ 2 F ≤ ∥ X ∥ 2 op ∥ H ∥ 2 F and the bound (38), we have

holds with probability at least 1 -exp( -n/ 2) .

Lemma D.3 (Operator norm bound for M ) . Under Assumptions 3.1, 3.3 and 3.5, we have

where ξ = η max c 0 n ∥ X ∥ 2 op .

Proof of Lemma D.3. By the definition of M in Lemma B.1, we have

Hence, we can write M = I pT -A , where A is the lower triangular matrix with off-diagonal blocks P 1 , ..., P T -1 . Using the matrix identity ( I -A ) -1 = ∑ ∞ k =0 A k and noticing A k = 0 for k ≥ T , we have

Taking operator norm on both sides, we obtain

Since D t = ∂ ψ ( u ) ∂ u ∣ ∣ u = y -X ̂ b t and ψ is 1-Lipschitz, we have ∥ D t ∥ op ≤ 1 . By the definition that S t = ∑ i ∈ I t e i e ⊤ i , we know ∥ S t ∥ op ≤ 1 . Since | I t | ≥ c 0 n and η t ≤ η max for any t ∈ [ T ] , we have

Plugging this inequality into (39) gives

Lemma D.4. Under the same conditions as Lemma D.3, we have

where ξ = η max c 0 n ∥ X ∥ 2 op .

Proof of Lemma D.4. By the definition of Γ in Lemma B.1, we have Γ = M -1 L ( Λ ⊗ I p ) ˜ D . Notice that Λ = ∑ T t =1 η t | I t | e t e ⊤ t , we have ∥ Λ ∥ op = max t ∈ [ T ] η t | I t | ≤ n -1 η max c 0 using | I t | ≥ c 0 n and η t ≤ η max . Since ϕ is 1-Lipschitz, we have ∥ ˜ D∥ op ≤ 1 . By definition of L in Lemma B.1, we have ∥ L ∥ op = 1 . Using these upper bounds of ∥ L ∥ op , ∥ Λ ∥ op , ∥ ˜ D∥ op and the upper bound of ∥M -1 ∥ op in Lemma D.3, we obtain

Lemma D.5 (Moment bounds for derivative of H , F , ˜ F ) . Under Assumptions 3.1, 3.3 and 3.5 and Σ = I p , we have for any finite integer k ,

Proof of Lemma D.5. We first prove the first bound. By Lemma B.1, we have

Hence, using ∂ e ⊤ k He t ∂x ij = ∂ e ⊤ k ̂ b t ∂x ij , we have

Using the above equality, ∑ i,j,t,k [( e ⊤ t ⊗ e ⊤ k ) A ( e j ⊗ e i )] 2 = ∥ A ∥ 2 F for A ∈ R pT × np , and the triangle inequality, we have

where the last inequality uses Lemma D.4. Taking the conditional expectation on both sides given ε , the desired moment bound follows from the moment bounds of X , H , F in Lemma D.1.

Now we prove the second bound. By definition, the t -th column of F is F t = S t ψ ( y -X ̂ b t ) , it can be written using the t -th column of ˜ F as F t = S t ˜ F t . Since ∥ S t ∥ op ≤ 1 , we have

By Lemma B.3, we have ∂ e ⊤ l ˜ Fe t ∂x ij = ˜ D lt ij + ˜ ∆ lt ij where

Using ∑ i,j,t,k [( e ⊤ j ⊗ e ⊤ i ) A ( e t ⊗ e l )] 2 = ∥ A ∥ 2 F for A ∈ R np × nT and the triangle inequality, we have

where the last inequality uses ∥D∥ op ≤ 1 and ∥S∥ op ≤ 1 . Taking the conditional expectation on both sides given ε , the desired moment bound follows from the bound of Γ in Lemma D.4 and the moment bounds of X , H , F in Lemma D.1.

D.6 Proof of Lemma B.5

Bound of Υ 1 . By the expression for the matrix Υ 1 ∈ R p × T obtained in Appendix D.4,

By the triangle inequality and ∥D∥ op ∨ ∥S∥ op ≤ 1 , ∥ Υ 1 ∥ op ≤ T ∥ Γ ∥ op ∥ X ∥ op ∥ F ∥ F . Using the bound of ∥ Γ ∥ op in Lemma D.4, the moment bound of ∥ X ∥ op , ∥ F ∥ F in Lemma D.1 gives

Bound of Υ 2 . By the expression for the matrix Υ 2 ∈ R n × T obtained in Appendix D.4, we have

By the triangle inequality and ∥D∥ op ∨∥S∥ op ≤ 1 , ∥ Υ 2 ∥ op ≤ T ∥ Γ ∥ op ∥ X ∥ op ∥ H ∥ F . Similar to the moment bound of ∥ Υ 1 ∥ op , we obtain

Bound of Υ 3 . By the expression for the matrix Υ 3 ∈ R T × T obtained in Appendix D.4, we have Υ 3 = ( Υ ⊤ 1 H + F ⊤ Υ 2 ) . It directly follows that

Using the triangle inequality and the moment bounds of ∥ H ∥ F , ∥ F ∥ F in Lemma D.1 and the moment bounds of ∥ Υ 1 ∥ op , ∥ Υ 2 ∥ op we just obtained, we have

Bound of Υ 4 . By the expression for the matrix Υ 4 ∈ R T × T obtained in Appendix D.4, we have

where ˜ Υ 1 = ∑ i ( I T ⊗ e ⊤ i X ) Γ (( F ⊤ e i ) ⊗ I p ) H . We can rewrite ˜ Υ 1 as

We have by the triangle inequality,

By the triangle inequality and the upper bound of ∥ Υ 2 ∥ op , we have

Squaring both sides and taking conditional expectation, we have

thanks to the upper bound of ∥ Γ ∥ op in Lemma D.4 and the moment bounds of ∥ X ∥ op , ∥ F ∥ F , ∥ H ∥ F in Lemma D.1.

Bound of Υ 5 . By the expression for the matrix Υ 5 ∈ R T × T obtained in Appendix D.4, we have Υ 5 = Υ ⊤ 1 X ⊤ ˜ F -˜ Υ 2 , where

By the triangle inequality,

Using the moment bounds of ∥ Υ 1 ∥ op , ∥ X ∥ op , ∥ F ∥ F , ∥ ˜ F ∥ F , ∥ H ∥ F , we have

D.7 Proof of Lemma B.6

We first state three useful lemmas.

Lemma D.6 (Adopted from Lemma E.10 of [26]) . Let U , V : R n × p → R n × T be two locally Lipschitz functions of Z with i.i.d. N (0 , 1) entries, then

Lemma D.7 (Adopted from Lemma F.5 of [5]) . Let U , V : R n × p → R n × T be two locally Lipschitz functions of Z with i.i.d. N (0 , 1) entries. Provided the following expectations are finite, we have

We will use the above two lemmas, conditionally on ε , to bound the conditional moments of Θ 1 , Θ 2 , Θ 3 , Θ 4 and Θ 5 given ε .

Bound of Θ 1 . By the definition of Θ 1 , we have

Applying Lemma D.6 conditionally on ε to ( Z , U , V ) = ( X , F , H ) gives

where the last line uses the moment bounds of ∥ F ∥ F , ∥ H ∥ F in Lemma D.1, the moment bounds of ∥ ∂ H ∂x ij ∥ F , ∥ ∂ F ∂x ij ∥ F in Lemma D.5, the moment bound of ∥ Υ 3 ∥ op in Lemma B.5.

Bound of Θ 2 . By the definition of Θ 2 , we have by (21)

Here, the penultimate line uses ∥ F ∥ F ≤ ∥ ˜ F ∥ F , ∥ ∂ F ∂x ij ∥ F ≤ ∥ ∂ ˜ F ∂x ij ∥ F , and ∥ ∂ X ⊤ ˜ F ∂x ij ∥ F = ∥ e ⊤ j e ⊤ i ∂ ˜ F ∂x ij + X ⊤ ∂ ˜ F ∂x ij ∥ F ≤ ∥ ∂ ˜ F ∂x ij ∥ F + ∥ X ∥ op ∥ ∂ ˜ F ∂x ij ∥ F . The last line uses the moment bounds of ∥ X ∥ op , ∥ ˜ F ∥ F in Lemma D.1, the moment bound of ∥ ∂ ˜ F ∂x ij ∥ F in Lemma D.5, the moment bound of ∥ Υ 5 ∥ op in Lemma B.5.

Bound of Θ 3 . By the definition of Θ 3 , we have

Applying Lemma D.6 conditionally on ε to ( Z , U , V ) = ( X , XH , H ) gives

Here, the penultimate line uses ∥ ∂ XH ∂x ij ∥ F = ∥ e i e ⊤ j ∂ H ∂x ij + X ∂ H ∂x ij ∥ F ≤ (1+ ∥ X ∥ op ) ∥ ∂ H ∂x ij ∥ F , and the last line uses the moment bounds of ∥ X ∥ op , ∥ H ∥ F in Lemma D.1, the moment bounds of ∥ ∂ H ∂x ij ∥ F in Lemma D.5, the moment bound of ∥ Υ 4 ∥ op in Lemma B.5.

Bound of Θ 4 . By definition, we have

Using ∑ p j =1 e j e ⊤ j = I p , we find

This further implies

Applying Lemma D.7 conditionally on ε to ( Z , U , V ) = ( X , F , ˜ F ) ,

Here, the last inequality uses the moment bounds of ∥ F ∥ F , ∥ ˜ F ∥ F in Lemma D.1, and the moment bounds of ∥ F ∥ 2 ∂ , ∥ ˜ F ∥ 2 ∂ in Lemma D.5, and the moment bound of ∥ ˜ Υ 4 ∥ op in Lemma B.5.

Bound of Θ 5 . By the identity (19), we have

By the definition of Θ 5 , we have

Applying Lemma D.7 conditionally on ε to ( Z , U , V ) = ( X ⊤ , H , H ) (i.e., consider the mapping from R p × n to R p × T : X ⊤ ↦→ H ) gives

Here, the last line use the moment bounds of ∥ H ∥ F in Lemma D.1, the moment bound of ∥ H ∥ 2 ∂ in Lemma D.5, and the moment bound of ∥ ˜ Υ 5 ∥ op in Lemma B.5.

Bound of Θ 6 . Using (19), we have

It follows that

Thus, using ˜ E = E / ∥ E ∥ F , we have

Applying Lemma D.6 conditionally on ε to ( Z , U , V ) = ( X , ˜ E , H ) gives

Here, the last line uses the moment bound of ∥ H ∥ F in Lemma D.1, the moment bounds of ∥ ∂ H ∂x ij ∥ F in Lemma D.5, and the moment bound of ∥ Υ 2 ∥ op in Lemma B.5. This finishes the proof of Lemma B.6.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction clearly state the claims made in the paper.

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: Our results only holds under several assumptions discussed in the paper.

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: All the proofs are provided in 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: All the code needed to reproduce the results is in the supplemental material.

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: All the code are provided in the supplemental material.

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: The choices of step size for GD and SGD are provided in Section 4.

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: All the figures include the 2-standard error bars.

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

Justification: It is provided in the README file of the supplemental material.

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:

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: Given the theoretical nature of our work, it does not have any negative societal impact.

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

Justification: The paper is of the theoretical nature, it does not invent any new models or datasets.

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: We write all the code by ourselves, and did not use other code or data.

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

Justification:

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:

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:

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:\nThe paper explores linear regression with Gaussian design corrupted by noise that has bounded first moments, focusing on the 'high-dimensional' regime where the feature dimensionality $p$ scales proportionally with the number of training data $n$. The authors propose two estimators that, aside from an additive noise term, approximate the out-of-sample error (measured with the square loss) achieved by a general iterative scheme. This scheme includes commonly used solvers like GD, SGD, proximal GD, and proximal SGD, which are used to compute the regularized empirical risk minimizer with a convex, differentiable, and smooth loss function (with Lipschitz-continuous gradients) and a potentially non-smooth regularizer. This framework encompasses robust linear regression with heavy-tailed noise using the Huber loss function. The authors provide non-asymptotic guarantees in probability for finite $n$ (and $p$), demonstrating the asymptotic consistency of their estimators (Theorems 3.6 and 3.7). They also present examples and numerical experiments. Their primary tool is a generalized probabilistic approximation of the generalization error previously discussed in specific settings in the references [3] and [5].\n\nSTRENGTHS:\nThe authors have developed a procedure that reliably determines the best stopping time for solvers in contexts where no consistent estimators were previously available, particularly in robust regression with heavy-tailed noise, Huber-like loss functions, and non-smooth penalty terms. This methodology offers an alternative to cross-validation techniques, which are generally inconsistent (e.g., V-fold) or computationally expensive (e.g., leave-one-out). The paper is well-written and includes clear proofs.\n\nWEAKNESSES:\nThe main results of the paper can be seen as an extension of the findings in [5], applied to a broader context that encompasses the stochastic case and extends beyond the square loss. While the generality of the weaker conditions is appreciated and introduces a different structure, the work is somewhat marginal in its contribution. Although the new methodology can track the generalization performance over training time for various iterative solver schemes of interest, it remains uncertain how this methodology can be used to (provably) \"optimally\" tune parameters other than the stopping time. This limitation seems notable, as the increased flexibility of the algorithmic schemes—allowing for additional tuning parameters such as mini-batch size and loss-specific parameters—is not accompanied by an enhanced analysis of these tuning parameters.\n\nQUESTIONS:\n1) In the case of the square loss and full-batch gradient descent, do the authors recover exactly the same theoretical results in [5]? If not, what are the differences in this case?\n\n2) The authors claim that their procedure can be used to choose the \"tuning parameters that achieves the smallest out-of-sample error\" (Page 2). While this is clear for stopping time, can the authors give examples of how their methodology can be use to tune other tuning parameters?\n\n3) While the authors clearly highlight the relevance of their contributions to robust statistics and their goal of tracking solver performance over iteration time, a specific comparison of their methodology and consistency results to other classical methods in some fundamental models, such as ridge regression and lasso regression, would be beneficial. Although Section 1.1 provides a discussion on this topic, it lacks quantitative details, making it unclear how the proposed methods compare in terms of non-asymptotic statistical rates (related to consistency results) and computational efficiency. For example, the computational complexity of the data-dependent estimator the authors consider ($\\tilde{r}_t$ on page 6) is not clearly compared to other estimators in basic settings. Can the authors address this?\n\n4) The paper does not seem to make any assumptions on the penalty term, i.e. the function g in (2). This seems surprising. Can the authors elaborate on this?\n\n5) Could the authors comment on the challenges of extending their analysis beyond the proportional regime they consider, where the data dimension is of the same order as the sample size?\n\n****\nI have increased my score to 'Accept' post rebuttal.\n\nLIMITATIONS:\nThe paper discusses limitation of the current approach as a way to motivate follow up research in the conclusions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on deriving consistent estimators of the generalization error for robust regression with Gaussian design and heavy tailed noise along the SGD trajectory. They propose two estimators, $\\hat r$ (which requires knowledge of the covariance) and $\\tilde r$ (which does not) and prove consistency of both estimators when the number of samples is at least linear in the dimension. They support this with experiments which demonstrate that both $\\hat r$ and $\\tilde r$ accurately capture the generalization error in various settings.\n\nSTRENGTHS:\nThe paper is overall very well written, including the overall motivation, the related work, and the technical ideas behind Theorems 3.6 and 3.7.\n\nWEAKNESSES:\n- The numerators in Theorems 3.6 and 3.7 include unspecified constants $C(T,\\ldots)$. From the appendix, it appears that $C(T,\\ldots)$ scales like $T^T$, so that you can only run for $T \\approx \\log(n)/\\log\\log(n)$ steps before the bound is vacuous, which is severely restricting. I'm concerned it may not be possible for SGD to significantly decrease the loss in this number of steps, rendering the generalization estimate vacuous.\n\n- While some intuition for $\\boldsymbol{W}$ (eq 8) is given in Section 3.1, none is given for $\\widehat{\\mathbf{A}}$ (eq 9) or $\\widehat{\\mathbf{K}}$ except that they are used to construct the weights for Theorem 3.7. This makes the notation in sections 3.2 and 3.3 somewhat difficult to follow.\n\nMinor Points:\n- line 29: starts -> start\n- line 118: maybe provide some intuition here for $\\phi_t, \\psi$ rather than deferring it to sections 2.1/2.2 (e.g. just $\\phi_t$ is a proximal operator)\n- line 170: all -> both\n- line 207: it is strange to define $\\hat r$ in terms of $\\mathbf{W}$ and $\\tilde r$ in terms of $\\widehat{\\mathbf{W}}$. Perhaps use $\\widetilde{\\mathbf{W}}$ instead of $\\widehat{\\mathbf{W}}$?\n- line 269: the three ... reveal\n\nQUESTIONS:\n- What are the dependencies on $T,\\eta_{max}$ in the numerators of Theorems 2.6 and 2.7?\n- Is the reason that $W$ must be computed recursively simply to invert $\\mathcal{M}$ block by block? I don't see how to connect the unrolling in section 3.1 to the Kronecker expressions in 3.2.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis manuscript aims at finding computational efficient measure of generalization performance for high-dimensional robust regression with regularization. In this scenario, when loss function is not quadratic, estimating the out-of-sample error $\\| \\Sigma^{1/2}(b_t - b^* ) \\|^2$ (where $\\Sigma$ is the covariance matrix, $b_t$ is the $t$-iteration variable and $b^*$ is the ground truth) for iterative methods like GD, SGD or proximal GD can be very difficult and is not well-explored in previous literature. To overcome this issue, the authors propose a new estimator for the out-of-sample error for a general class of gradient methods and show that the difference converges to zero in probability when sample size goes to infinity.\n\nSTRENGTHS:\nThis paper is generally quite well-written. The proof and result appears to be reasonable. Empirical evidences confirm that the proposed estimator is indeed effective and accurate.\n\nWEAKNESSES:\nDespite being more general, this paper seems to constitute *limited* progress towards the community. When comparing it with the work by Bellec and Tan, it is not hard to notice the estimator $\\hat{r}_t$ in Theorem 3.6 is a trivial extension of Bellec and Tan's estimator from square loss to convex loss. The exact form of the proposed estimator and its guarantee in Theorem 3.6 are very similar to the result of the mentioned paper. Also, the proof looks quite straighforward and extending from square loss to convex loss does not rely on technical innovations. In this regard, I believe this paper falls short in getting accepted to NeurIPS, unless the authors could provide more convincing arguments how this result is *fundamentally* different from previous works.\n\nQUESTIONS:\nThere are no further questions.\n\nLIMITATIONS:\nThis work is purely theoretical and has no negative influences.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n#### Summary\nThis paper examines the generalization performance of iterates produced by Gradient Descent (GD), Stochastic Gradient Descent (SGD), and their proximal variants in high-dimensional robust regression problems. The paper introduces estimators that accurately track the generalization error of the iterates along the trajectory of the iterative algorithms. These estimators are proven to be consistent under certain conditions and are illustrated through several examples, including Huber regression, pseudo-Huber regression, and their penalized variants with non-smooth regularizers.\n\nSTRENGTHS:\n#### Strengths\n1. **Innovative Methodology**: The introduction of estimators that track the generalization error of iterates is novel and provides a significant contribution to the field of robust regression.\n2. **Theoretical Rigor**: The paper provides a thorough theoretical foundation for the proposed estimators, including consistency proofs and detailed mathematical derivations.\n3. **Practical Relevance**: The approach is applicable to a variety of robust regression problems, including those with heavy-tailed errors and non-smooth regularizers.\n4. **Empirical Validation**: Extensive simulations demonstrate the effectiveness of the proposed estimators in tracking the generalization error and determining the optimal stopping iteration.\n\nWEAKNESSES:\n#### Weaknesses\n1. **Computational Complexity**: The proposed estimators involve complex calculations that may be computationally intensive, especially for large datasets. More discussion on computational efficiency and scalability is needed.\n2. **Generality**: The paper focuses on specific types of robust regression problems. Extending the methodology to a broader range of regression problems would increase its impact.\n3. **Comparison with Existing Methods**: The paper provides limited empirical comparisons with existing state-of-the-art methods for estimating generalization error. More comparative analysis would strengthen the validity of the proposed approach.\n4. **Practical Guidelines**: While the theoretical results are robust, practical guidelines for implementing the estimators in real-world scenarios are not sufficiently detailed.\n\nQUESTIONS:\n#### Questions\n1. **Computational Complexity**:\n    - Could you provide more details on the computational complexity of the proposed estimators? How do they scale with increasing dataset size and dimensionality?\n\n2. **Generality**:\n    - The paper focuses on specific types of robust regression problems. Are there any challenges in extending the methodology to other types of regression problems, such as those with different types of regularizers or loss functions?\n\n3. **Comparison with Existing Methods**:\n    - How do the proposed estimators compare empirically with existing methods for estimating generalization error? Are there scenarios where your approach significantly outperforms others?\n\n4. **Practical Implementation**:\n    - Can you provide practical guidelines or heuristics for implementing the proposed estimators in real-world scenarios? What are the key considerations practitioners should keep in mind?\n\nLIMITATIONS:\nno", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}