Greedy Pruning with Group Lasso Provably Generalizes for Matrix Sensing
Authors: Nived Rajaraman
Fnu Devvrit
Aryan Mokhtari
Kannan Ramchandran
Venue: NeurIPS 2023 poster (2023)
Keywords: Greedy Pruning; Matrix Sensing; Lasso regularization
TL;DR: We provide the first rigorous theoretical insights on why the popular "greedy pruning+finetuning" pipeline leads to smaller models with good generalization, in the context of the overparameterized matrix sensing problem.

Abstract: Pruning schemes have been widely used in practice to reduce the complexity of trained models with a massive number of parameters. In fact, several practical studies have shown that if the pruned model is fine-tuned with some gradient-based updates it generalizes well to new samples. Although the above pipeline, which we refer to as pruning + fine-tuning, has been extremely successful in lowering the complexity of trained models, there is very little known about the theory behind this success. In this paper we address this issue by investigating the pruning + fine-tuning framework on the overparameterized matrix sensing problem with the ground truth denoted $U_\star \in \mathbb{R}^{d \times r}$ and the overparameterized model $U \in \mathbb{R}^{d \times k}$ with $k \gg r$. We study the approximate local minima of the mean square error, augmented with a smooth version of a group Lasso regularizer, $\sum_{i=1}^{k} \lVert Ue_i \rVert_2 $. In particular, we provably show that pruning all the columns below a certain explicit $\ell_2$-norm threshold results in a solution $U_{\text{prune}}$ which has the minimum number of columns $r$, yet close to the ground truth in training loss. Moreover, in the subsequent fine-tuning phase, gradient descent initialized at $U_{\text{prune}}$ converges at a linear rate to its limit. While our analysis provides insights into the role of regularization in pruning, we also show that running gradient descent in the absence of regularization results in models which {are not suitable for greedy pruning}, i.e., many columns could have their $\ell_2$ norm comparable to that of the maximum. Lastly, we show that our results also extend for the training and pruning of two-layer neural networks with quadratic activation functions. To the best of our knowledge, our results provide the first rigorous insights on why greedy pruning + fine-tuning leads to smaller models which also generalize well.

Greedy Pruning with Group Lasso Provably Generalizes for Matrix Sensing
Nived Rajaraman
Devvrit
EECS Department
CS Department
University of California, Berkeley
The University of Texas at Austin devvrit. wotttt.00mmmal...mm
nived. rajaraman@eecs. berkeley edu
Aryan Mokhtari
Kannan Ramchandran
ECE Department
EECS Department
The University of Texas Austin mokhtar i taaustinn utexas edu
University of California, Berkeley kannanr@eecs berkeley edu
Abstract
Pruning schemes have been widely used in practice to reduce the complexity trained models with massive number of parameters. In fact, several practical studies have shown that pruned model is fine-tuned with some gradient- updates it generalizes well to samples. Although the above pipeline, which we refer to as pruning fine-tuming seee-tuning, has been extremely successful in lowering the complexity of trained models, there is very little known about the theory behind this success. In this paper, we address this issue by investigating the pruning + fine-tuning framework ground truth U, E [ddxr and the overparameterized model U E Rdx with k > r. We study the approximate local minima of the mean square error, augmented with smooth version of group Lasso regularizer, Σ ||Ue|ll2. particular, we provably show that pruning all the columns below a certain explicit e--mo threshold results solution Uprume which has the minimum number of columns r, yet close to the ground truth in training loss. Moreover, in the subsequent fine- tuning phase, gradient descent initialized at Upune converges at a linear rate to its limit. While our analysis provides insights into the role of regularization in pruning, we also show that running gradient descent in the absence of regularization results in models which not suitable for greedy pruming, i.c.. many columns could have their & norm comparable to that the maximum. To the best of our knowledge, our results provide the first rigorous insights on why greedy pruning + fine-tuning leads to smaller models which also generalize well.
the overparameterized matrix sensing problem with the
Introduction
Training overparameterized models with a massive number of parameters has become the norm in almost all machine learning applications. While these massive models are successful in achieving low training error and in some cases good generalization performance, they are hard to store or communicate. Moreover, inference with such large models is computationally prohibitive. To
address these issues, a large effort has gone into compressing these overparameterized models via different approaches, such as quantization schemes 2). unstructured (3] and structured [4+6] pruning mechanisms, and distillation techniques using student-teacher models [7 Among these approaches, greedy pruning, in which greedily eliminate the parameters of the trained model
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

test-ground for the analysis of based on some measure (e.g., the norms of the weight vectors associated with individual neurons) has received widespread attention B [.. This is mostly due the fact that several practical studies have illustrated that training an overparameterriee model followed by greedy pruning and fine-tuning leads to better generalization performance, compared to an overparameterized model trained without pruning [14. Furthermore, a phenomenon that has been observed by several practical studies, is that different forms of regularization during training, such as to or & regularization [1SH[7] or lz regularization including group Lasso [8 lead to models that are better suited for pruning, and leading to better generalization post fine-tuning. While the greedy pruning framework has shown impressive results, there is little to theory backing why this pipeline works well in practice, nor understanding of the role of regularization in helping generate models which are suitable for greedy pruning. In this work, we address the following questions:
Does the greedy pruning + fine-tuning pipeline provably lead to a simple model with good generalization guarantees? ? What is the role of regularization in pruning?
In this paper, we use the symmetric matrix sensing problem 19 greedy pruning framework, a model very closely related to shallow neural networks with quadratic activation functions 20 19 In this setting, the underlying problem for the population loss (infinite samples) is defined as IOU U.U?13 1 where U. E Raxr is an unknown ground-truth rank-r matrix, with r also being unknown, and U E Rdxk is the overparameterized learning model with k >> r. As we discuss in the Appendix FP the columns of U can be thought of as the weight vectors associated with individual neurons in a 2-layer shallow neural network with quadratic activation functions, a connection first observed in (21). Thus, the data generating model has r neurons, while the learner trains an overparameterized model with A neurons.
While the statistical and computational complexity of the overparameterized matrix sensing problem has been studied extensively, we use it as a model for understanding the efficacy of greedy pruning. In particular, we aim answer the following questions: Does there exist a simple pruning criteria for which we can provably show that the pruned model generalizes well after fine-tuning while having the minimal necessary number of parameters? What is the role of regularization during training in promoting models which are compatible with greedy pruning? Finally, what generalization guarantees can we establish for the pruned model post fine-tuning?
Contributions. Our main contribution is to show that the discussed pruning pipeline not only recovers the correct fround-truth U.U? approximately, but also automatically adapts to the correct number of columns r. In particular, we show that training an overparameterized model on the empirical mean squared eιror with an added group Lasso based regularizer to promote column sparsity, followed by simple norm-based pruning strategy results in a model Upmnne having exactly the minimum number of columns, r. At the same time, we show that [|ppane prune - U. U_U? 1/ is small, but non-zero. Hence, the pruned model can subsequently be fine-tuned using small number of gradient steps, and in this regime, |UGU?" - U. I. shows linear convergence to its limit. Moreover, the pruned model can be shown admit finite sample generalization bounds which are also statistically optimal for a range of parameters. In particular, we show that to obtain a model Unut that has exactly r columns and its population error is at most [|Uout out U.UF || S e, our framework requires Oddk2p5 + - ) samples, which statistically optimal for sufficiently small e. We should also add that our framework does not require any computationally prohibitive pre- or post-processing (such as SVD decomposition) for achieving this result.
While there are several works 22 parameterized" setting requires O(rd/e" samples to achieve generalization error of e, and con- verges linearly to this limit, the picture is different in the overparameterized setting. In [23], the authors showed that in the overparametrized setting, vanilla gradient descent requires Okk//s" samples to achieve a generalization error of [| gd - U.U || S e, degrading with the overparam- eterization of the model. Moreover the resulting solution does not have the correct column sparsity. In order to obtain a model which can be stored concisely, Ugi has to be post-processed by computing its SVD, which is computationally expensive in the high dimensional regime.
19
23 establishing that gradient descent in the "exactly-
As our second continuation, we show that use of explicit regularization to promote column sparsity while training is important to learn models suitable for greedy pruning. Specifically, we show that while implicit regularization [19 24 suffices to learn models with the correct rank, these approaches learn solutions with a large number of columns having fo-nos comparable to that of the maximum, even if 1. Hence, it is unclear how to sparsify such models based on the norms of their columns.
2

---

2 Setup and Algorithmic Framework
Given n observation matrices {Ar]i=r in the matrix sensing framework, the learner is provided mea- surements yı = (Ai, U.U?) + Ex where (---) indicates the trace inner product and E; is measurement noise assumed to be distributed i.id. ~ (0, o² Here U. Rdxr is the unknown parameter, and the rank r Σ is unknown. The goal of the matrix sensing problem is to learn a candidate matrix X such that X z U.U?. For computational reasons, it is common to factorize X as UU for U e Rdd* In this paper, we study the factored model in the overparamelerized setting, where k >>>> The empirical mean squared error is,
Cemp (U) 1 Σ (AA, yi)² -y))
(1)
For the case that A are sampled entry-wise i.i.d. (N 1/d) and as n + co, up to additive constants which we ignore, the population mean square error can be written down as,
Sppp (U) = 1 - 1.
There is an extensive literature on how efficiently learn the right product U.U?" in both finite sample and population settings (25 26]. In particular, there are several works on the efficiency of gradient- based methods with or without regularization for solving this specific problem (99 23). While these approaches guarantee learning the correct product UU = U.U., in the overparameterized setting the obtained solutions are not column sparse and storing these models requires ((kd) bits of memory (ignoring precision). As result, in order to obtain a compressed solution Unat Rdxr with the correct number of columns, one has to post-process U and do a singular value decomposition (SVD), an operation which is costly and impractical in high-dimensional settings.
The goal of this paper overcome this issue and come up with an efficient approach to recover a solution Una which generalizes well, in that (U) is small, while at the same time having only a few non-zero columns, i.e. is sparse. Specifically, we show that via some proper regularization, is possible to obtain a model that approximately learns the right product U,U?" while having only few significant columns. As a result, many of its columns can be eliminated by simple lg norm-based greedy pruning scheme, without significantly impacting the training loss. In fact, post pruning we end up with a model Uprune Rdxr that has the correct dimensionality and its outer product U U prone
[U)
is close to the true product U.U?", When the resulting "exactly parameterized" model is fine-tuned, the generalization loss Cpop (U) can be shown to converge to a linear rate.
To formally describe our procedure, we first introduce the regularization scheme that we study in this paper, and then we present the greedy pruning scheme and the fine-tuning procedure.
Regularization. The use of regularization for matrix sensing (and matrix factorization) to encourage low rank solution [3 to control the norm of the model for stability reasons 34 35 or to improve the landscape of the loss by eliminating spurious local minima [36] has been well-studies in the literature. While these approaches implicitly explicitly regularize for the rank of the learned matrix, the solution learned as a result is often not column sparse. Indeed, note that a matrix can be low rank and dense at the same time, if many columns are linear combinations of the others. We propose studying the following regularized matrix sensing problem with a group Lasso based regularizer 337 [8]. In the population case, the loss is defined as
Cpop + AR(U), where R(U) = Σ |Ue/l12.
Note that À > 0 is a properly selected regularization parameter. Imposing R as a penalty on the layer weights of neural network is a special case of a widely used approach commonly known as Structured Sparsity Learning (SSL) [12). The regularizer promotes sparsity across "groups" as discussed by [37] Here, the groups correspond to the columns of U. The use of matrix mixed-norms as a regularizer for sparsity was also considered in [38). Furthermore, a connection the group
All proof's in the paper go through as long as the noise is i.i.d. Sab-Gaussian with variance proxy o². We study the Gaussian case for simplicity of exposition.
3

---

Lasso penalty with the Schatten-1/2 norm was considered in (3) where the authors show that stationary points of the factored objective also serve as stationary points of the unfactored objective and heuristically argue that the regularizer promotes sparse solutions.
As we prove in Section 4 the solution obtained by minimizing smooth version of the regularized loss in (3). denoted by Uour. is approximately column-sparse, i.c., k r of its columns have small lg norm. As result, it suitable for the simple greedy pruning scheme which we also introduce in Algorithm I Interestingly, in Section 3 we will first show a negative result- - the model obtained by minimizing the unregularized loss, ppp((U) (eq. 2 using gradient descent updates could fail to learn a solution which is suitable for greedy pruning. This shows the importance of adding some form of regularization during the training phase of the overparameterized model.
Greedy Pruning. The greedy pruning approach posits training a model (possibly a large neural network) the empirical loss, followed by pruning the resulting trained network greedily based on some criteria. The resulting model is often fine-tuned via a few gradient descent iterations before outputting. In the literature, various criteria have been proposed for greedy pruning. Magnitude- based approaches prune away the individual weights/neurons based on some measure of their size such as f norm of the associated vectors [66] 10].
In this work, we also focus on the idea of greedy pruning and study a mechanism to prune the solution obtained by minimizing the regularized empirical loss. Specifically, once an approximate second-ore order stationary point of the loss in Section 4 we only keep its columns whose lz norm above threshold (specified in Algorithm [ and the remaining columns with smaller norm are removed. We further show that post pruning. the obtained model Upunne continues to have a small empirical loss, i.e., small Lemp (Upune))) while having exactly columns.
Post-pruning fine-tuning. As mentioned above, common to fine-tune the smaller pruned model with few gradient updates before the evaluation process [40]. We show that the pruned model U prne has the correct rank and is reasonably close to the ground truth model U. in terms of its population loss. By running a few gradient updates on the mean square error, it can be ensured that UUT converges to U.UT at a linear rate. Algorithm summarizes the framework that we study.
3 Implicit regularization does not lead to greedy pruning-friendly models
In various overparameterized learning problems, it has been shown that first-order methods, starting from a small initialization, implicitly biases the model toward "simple" solutions, resulting in models that generalize well [4] (4). a phenomenon known as implicit regularization. In particular, for matrix sensing, (3319 24 show that in the absence of any regularization, running gradient descent on the population loss starting from a small initialization biases UUT to low rank solutions, and learns the correct outer product UU U.U?.. However, as discussed earlier, low-rank solutions necessarily column sparse, nor is it clear how to sparsify them without computing an SVD. It is unclear whether implicit regularization suffices to learn models that are amenable for greedy pruning.
not
In this section, we address this question and show that minimizing the unregularized population loss Cpop leads models which are not suitable for greedy pruning. i.c., have many columns with large lz norm. Specifically, we show that by running gradient flow from a small random initialization, even the ground truth U, is just single column and 1, the learnt solution U has a large number of columns that are "active". L.e having fz-norm comparable to that of the column with maximum norm. Thus, in the absence of the knowledge of r, it is unclear how to determine it from just observing the columns & norm. We thus claim that such trained models are not compatible with greedy pruning. In the following theorem, formally state our result.
Theorem 1. Consider the population loss eq. for and k >>>>> and suppose ||U=llop 1. Further, suppose the entries of the initial model Uo are i.i.d. samples from N' a²), where a S k d log(kd) for some constant @ > 0. For another absolute constant Ø > 0, as t + 0o, the iterates of gradient flow with probability 2 - 0(1kk)) converge to model Uga where ((())
active columms satisfy
||Ugee112
10 2 0.99.
(4)
maxje(N)
4

---

1.0
regulariationn regularization
repdaiiadaa
0.8
number of columns
percerrogge
(wit max,
0.4
0.0
0.2 0.4 0.6 0.8 1.0
4.01
column norm
Figure 1: We run gradient updates to minimize the loss in eq. 2 until convergence. The model is highly overparameterized with d = = 500 and r = 4. In the left figure, the vertical axis represents the fraction of columns having their norm at least times the largest norm across columns, with E [0, 1] on the horizontal axis. In the right figure, for the same experiment we plot histogram of column norms. Using regularization leads to solutions with most the columns having small & norm, only few significant columns. Without regularization, a large number of columns have their & norm comparable to the largest column norm. Thus, training with explicit regularization leads to models that are more suitable for pruning.
Theorem T claims that the iterates of gradient descent converge to solution for which large number (I2(K) of columns of Ug have their lg norm least 0.99 times the maximum. An inspection of the proof shows that the constant 0.99 is arbitrary and can be extended to any constant bounded away from 1. Therefore, the models learnt by vanilla gradient descent cannot be reduced the right sparsity level by simply pruning away columns with & norm below any fixed constant threshold.
We provide general guarantee in Theorem 22 (Appendix A). It is worth noting that from [19 that the Iearned model Uga E Rddk is guaranteed to achieve small generalization error, i.c., Ug gd U.U.. Hence, the algorithm does learn the correct ground- truth product what we show cannot be sparsified to the correct level by on
that the learnt model is not "pruning-friendll: any simple norm based pruning strategies. This is also observed in Figure [] the model trained with implicit regularization has many significant columns and is unclear how many to prune. We provide proof sketch below, and refer the reader to Appendix A for more details.
Proof sketch of Theorem 1
Without loss of generality we assume in the proof that ||U, ll2 Defining (t) as the iterate at time t, gradient flow on Lpop follows the below dynamic,
dU dt =
- -000 UU
(5)
where for ease of notation we drop the explicit dependence on in (t). In our proof, we exactly characterize the limiting point of gradient flow as function of the initialization in certain sense, which may be of independent interest. In particular, up to a small multiplicative error, we show,
(U_, (0)e
Vi E [k], ||Ugeeel2 ~
(6)
(
Also, with Gaussian initialization, (U., U()ee) itd. NN a²) across different values of i. In particular, we show that for some constant o > O, (( columns will have correlations comparable to the maximum, with (U., (0)))))) 2 0.99 max (U..U (0)e;)) For any of these columns, i,
(U.,U (0)ez))
(O)e,) 0.99 max ||ge,,112, je(k)
||U gUgde R
2 0.99
max jE(M]
(7)
[U)[(())22
[UTU(()]
which completes the proof sketch of Theorem
4 Explicit regularization gives pruning-friendly models: population analysis
In this section, we study the properties of the squared error augmented with group Lasso regularization in the population (infinite sample) setting. We show that second-order stationary points of the
5

---

Algorithm 1 Greedy pruning based on group-Lasso regularization
Inputs: Measurements {(A, Mi) where yi = (A,,UUU] + =:) (in the population setting n = oo);
Initialization: Set parameters: X, 3, 7 and Minne--ue-
Greedy pruning phase:
1: Find an (e. >approximmtee SOSP of Jemp (resp. Spw), U, satisfying ||U|| S 3.
2: Let S = {i E [K] ||Ueeill I 2/B B denote the set of columns with small &g norm.
Create new matrix Uprame which only preserves the columns of U in [k] \ S, deleting the columns in S.
Fine-tuning phase:
1: Run Miine une iterations of gradient descent on Lepp (U) (resp. Lamp (U)) initialized at Upune to get Uout-
2:
return Vou.-
regularized loss are suitable for greedy pruning, while at the same time achieving small but non-zero generalization error. Note that the l norm is a non-smooth function at the origin, and therefore, the overall regularized loss is non-smooth and non-convex. While there several notions of approximate stationary points for non-differentiable and non-convex functions, for technical convenience, we replace R. by smooth proxy. In particular, for smoothing parameter > 0, define a smooth version of the l norm, and the corresponding smooth regularizer RB as,
|
R₂(U)= - - (Ues), where (
(8)
+
Note that the smaller the value of 3 is, the closer is ()( to ||0||22- Considering this definition, the overall regularized loss we study in the population setting is
Spop (U) = ppp + (U),
where Cpop = -IUUT-U,"?"
2 The above optimization problem is nonconvex due to the structure of (u))(U), and finding its global minimizer can be computationally prohibitive. Fortunately, for our theoretical results, we do not require achieving global optimality and we only require an approximate second-order stationary point of the loss in eq. (9) which is defined below.
-U,T?I? 1p as defined in eq.
Definition 2. We say that U is an (e, apprrxxmmmee second-order stationary point of f if,
1. The gradient norm is bounded above by e, i.e., [V S.
2. The eigenvalues of the Hessian are larger than -y,, i.e., Amin ( (V)(())
The full algorithmic procedure that we analyze in this section is summarized in Algorithm I Once we find an (c, appappooxee second-order stationary point (SOSP) of eq. (9) for some proper choices of and % we apply greedy pruning on the obtained model U and by eliminating all of its columns with & norm below a specific threshold. As result, the pruned model Upumee has fewer columns than the trained model U. In fact, in our following theoretical results, we show that if the parameters are properly selected, Uprme has exactly r columns, which is the same as U.. Finally, we fine-tune the pruned solution by running gradient descent on the unregularized loss Goop (resp. Lemp). Next, we state the properties of the pruned model generated by Algorithm I in the infinite sample setting.
Theorem 3. Consider the population loss with regularization in eq. (9). where U. has rank r and its smallest singular value is denoted by ot. Moreover, consider U Uprume as the output of the pruning phase in Algorithm 7 with parameters G, satisfying the condition
(07)
(0#)3 Vkr
(ot)3 7 S.C Vkra/22
(0;)7/2
8 cB
A C2
ESC. C.
(10)
Vkr
for some absolute constants C3, CX2 Ce, Cy 0. Then, we have
Uprume has exactly r columns.
2. UT prame - U.U. ||r S (o;)2
"This style of result exists for LASSO as well (see Theorem 1J), where the optimal choice of the regularization parameter, A", depends on the true sparsity r. but a general guarantee can be established as well, which degrades as \ deviates from X", r. In practice À is chosen using cross- For simplicity of presentation, we state the result when and of are known up to constants.

---

This result relies on showing that all bounded SOSPs of the regularized loss in eq. 9 are suitable for greedy pruning: removing the columns of U below a certain l2- no-m threshold results in solution Uprune having exactly columns, while at the same time having small generalization enror. Hence, it can serve as proper warm-start for the fine-tuning phase.
Proof sketch of Theorem 3 The key idea is to identify that if we have a matrix U such that UT U.U. and the columns of U are orthogonal to one another, then U has exactly non-zero columns, where r is the rank of U.. This statement can be shown to hold even when UTT x U.UT and the columns of U are only approximately orthogonal. The main observation bounded (e, 7a-approximate SOSPs of eq. denoted by U satisfies the following condition:
prove is that
(Ue,, Ueg) |Ue//ellUeell
Vi,j : ||uell2 ||e, 2 2/B.
(11)
In other words, all the large columns of U have their pairwise angle approximately 90'. Thus, by pruning away the columns of U that have an lz norm less than 2/3, the remaining columns of U, i.e. the columns of Upune. are now approximately at 90º angles to one another. If B is chosen to be sufficiently small, after deleting the low-norm columns, the approximation Uprane U prune R UUT holds. By the second order stationarity of U, also have that UUT U.U.T Together, this implies that UpruneuUpmm s U.U? and U prume is close to rank r matrix. Since Upence has orthogonal columns, this also means that it has exactly columns. Finally, to establish bound on the approximation error, we simply use the triangle inequality that U prune - U. U.U7|| S |UU - U_U F + |IUT - U prame U prruell ||.. The former is small by virtue of the fact that U is an approximate second order stationary point of Spop 2 Leop when X is small; the latter term is small by the fact that only the small norm columns of U were pruned away.
Now the only missing part that remains to justify why the R, regularizer promotes orthogonality in the columns of approximate second order stationary points and the expression in eq. (11) holds. This is best understood by looking at the regularized loss for the case 3 = O, which is equivalent to ||UU - U.U 1|? + A >Σ 1 [Ueeill and consider any candidate first-order stationary point U of this objective. Let Z E Raxi be a variable constrained to satisfy Z Zz = UUT. Stationarity of U implies that the choice Z U must also be first-order stationary point of the constrained optimization problem,
k
Minimize: |ZZ U.U IrP A
Σ
| Zeslla,
Subject to: zzT = UUT.
(12)
The first term in the objective is constant under the constraint and may remove it altogether. When U is a full-rank stationary point, constraint qualification holds, and it is possible to write down the necessary KKT first-order optimality conditions, which reduce to,
Vi [k), -> Zei
+ (A+ + A Ze; =
(13)
IZe
where A E RRaxd is the set of optimal dual variables. Since U is a first-order stationary point of the problem in eq. (12) and it satisfies eq. (13). the above condition means that the columns Ue, are the eigenvectors of the symmetric matrix A AT. If the eigenvalues of A + AT were distinct, then this implies that the eigenvectors are orthogonal and Ue; 1 Ue, for all * j.
While this analysis conveys an intuitive picture, there are several challenges in extending this further. It is unclear how to establish that the cigenvalues of A + A are distinct. Moreover, this analysis only applies for full-rank stationary points and does not say anything about rank deficient stationary points U. where constraint qualification does not hold. Furthermore, it is even more unclear how to extend this analysis to approximate stationary points. Our proof will circumvents each of these challenges by (a) showing that at approximate SOSPs, even if the eigenvalues of A + AT are not distinct, the columns of U are orthogonal, and (b) directly bounding the gradient and Hessian of the regularized loss, rather than studying the KKT conditions to establish guarantees even for approximate stationary points which may be rank deficient. Having established guarantees for the pruning phase of the algorithm in the population setting, next prove result in the finite sample setting.
7

---

Finite sample analysis
Next, we extend the results of the previous section to the finite sample setting. Here, we also focus on the smooth version of the regularizer and study the following problem
(Um (U) = Cemp (U) + (U)(U))
(14)
where the empirical loss Comp is defined in eq. ① and the smooth version of the group Lasso regularizer (U) is defined in eq. (8). In the finite sample setting, we assume that the measurement
matrices satisfy the restricted isometry property (RIP) [45], defined below.
Assumption 1. Assume that the measurement matrices {A1,... An are (2k, In other words, for any d x matrix with rank & 2k:
(1 (1-8)1IXI1 S X), S (1+8))IX11)
(15)
This condition is satisfied, for instance, if the entries of the A,'s were sampled (0, 1/d/ (i.e. Gaussian measurements), as long as n 2 dk/8² [45].
Theorem Consider the empirical loss with smooth regularization in eq. (14). where U. has rank and unit spectral norm, with its smallest singular value denoted of and noise variance 0². Consider U prrme as the output of the pruning phase Algorithm
7 with parameters B, X,,, chosen as per (a.)" 372
10p. If Assumption I is satisfied with S C6 "kr-7a and the number of samples is at least eq.
2
O is sufficienttly large constant, then with probability at La boraak"rr d log(d/n)) where Cg >
least 1 17.
1. Uprmne has exactly columns.
2. U Ureme satisfies the spectral initialization condition: |[Uprane (ipmneUprane. Uprune U.U! 1e S (o;)).
With Gaussian measurements, the overall sample requirement (including that from the RIP condition) in Theorem is satisfied when Q
[=t=" (xt) dk2pp5 dkk25) r The high level analysis of this result largely follows that of Theorem 3 in the population setting we approximate the finite-sample gradient and Hessian by their population counterparts and show that the approximation error decays with the number of samples as (l)
5.1 Fine-tuning phase: Realizing the benefits of pruning
In the fine-tuning phase, the learner runs a few iterations of gradient descent on Cemp initialized the pruned solution Upnee- Since the model is no longer overparameterized after pruning (by Theorem 4 there are several works analyzing the generalization performance and iteration complexity of gradient descent. Here we borrow the local convergence result of [22) which requires the initial condition that UU U.UTT|r S c(a;))² where e any constant less than As shown part (b) of Theorem this initial condition is satisfied by Upmane-
1lp 1. If we use the output of the greedy pruning in Algorithm 7 denoted by Uprone E (Rdxrr as the initial iterate for the fine-tuning phase, then after
Theorem 5. [22] Corollary 2) Suppose [U, llop =
t Mime--ume ( /o+) log (ož/oī n//) iterations, for some sufficiently large absolute
0, the iterates {U.)22 of factored gradient descent on Lemp (-) satisfy.
rd
-U.UH-U.U?? (ot)2
(16)
Theorem 5 shows that in the fine-tuning phase, the iterates of gradient descent converges at linear rate to the generalization error floor of a a/(a;)² Vrd/n. /n, which is also known to be statistically (minimax) optimal 46 47 This possible because the learner is able to correctly identify the rank of the model in the pruning phase and operate in the exactly specified setting in the fine-tuning phase. In contrast, one may ask how this guarantee compares with running vanilla factored gradient descent in the overparameterized setting. This corresponds to the case where pruning is carried out to reduce

---

the size of the model. [33 presented a guarantee for the convergence of factored gradient descent from a warm start in the overparametertzed setting. In comparison with the exactly specified case. the generalization error floor lim;-roc || U. S.UT Ir is shown to scale as O(1 ((;/( /kd/ n which now depends on k. Furthermore, linear convergence can no longer be established for vanilla gradient descent because the ill-conditioning of the objective. This is not just an artifact of the proof experimentally too, the convergence slowdown was noticed in
[
Figure 1].
This discussion shows that greedy pruning the model first, prior to running gradient descent, in fact generates solutions which generalize better and also converge much faster.
Remark 6. Based on Theorem 4 and Theorem
5 (eq. (16). under Gaussian measurements, given
0² rd 1 + (o:) dk2y log(d/n) (log//)))))
RE, 0
(17)
(o;)a
samples, with probability > 1 - 9), Algorithm produces Unus that has exactly columns and satisfies ||unaU out U,Ú, || S E. Note that the sample complexity depends on the amount of overparameterization, k only in the lower order (independent of =) term.
Note that this result uses the fact that under Gaussian measurements, the RIP condition required in
Theorem is satisfied if n 2 Ngg
Implementing Algorithm
1 Smoothness and optimization oracles
In this section we instantiate the optimization oracle in Algorithm I which outputs an approximate SOSP with bounded operator norm. First, we establish that the loss semp is well behaved on the domain {U I|0'llap < 3), in that its gradient and Hessian are Lipschitz continuous. These conditions are required by many generic optimization algorithms which return approximate second order stationary points [48] [49]. We establish these properties for the population loss for simplicity and leave extensions to the empirical loss for future work.
Theorem 7. Consider the population loss Syop in eq. (9). Assume À S min /B} and ||U=|lop 1. The objective (pop(t) defined in eq. 99 satisfies for any U. V E Raxk such that ||U|lop |V S 3,
(a) Lipschitz gradients: | fmp (U) fm (V) ]F & (U-Vl
(b) Lipschitz Hessians: |/ fpop(U) (U) V fpop (pop(V)
(V) z
IU -VIl.
Under the Lipschitz gradients and Lipschitz Hessian condition, a large number of algorithms in the literature show convergence to approximate SOSP. A common approach for finding such a point is using the noisy gradient descent algorithm [50]. However, note that we establish these Lipschitzness properties on the bounded domain {U ||0l|op 3), and it remains to verify whether these algorithms indeed approach such points. A similar concern is present with Algorithm [] which requires access to optimization oracle which finds (c. -approxximt SOSP of femp which are also bounded, in that ||O||lpp I 3. The final step is to identify conditions under which these algorithms indeed output stationary points which are bounded, satisfying ||U||op < 3. We establish this behavior for a wide family of perturbed gradient based methods.
We consider gradient descent with the following update rule: starting from the initialization Uo. for all & 2 O,
Perturbed gradient descent:
U+i U, S + P.)
(18)
where P; is perturbation term, which for example, could be the explicit noise (P: (B(() for appropriate r) added to escape strict saddle points in [50). Over the course of running the update rule eq. (18), we show that ||U: llpp remains bounded under mild conditions if the algorithm is initialized within this ball. combination with Theorem 7 this shows that the noisy gradient descent approach of [50] can be used to find the SOSPs required for Algorithm
Theorem 8. Consider optimization of the regularized population loss using the update rule in eq. (8). Assume that |U, ll 1 and suppose the parameters are selected aso: <. 1/8 and X < /3. and we have Pello < 1 almost surely for each 2 O. Then, assuming that the condition [!oloop $ is satisfied at initialization, for every t > 1, we have ||x>|lop < 3.
The result is often stated in terms of the smallest non-zero eigenvalue of X. U.UT. which equals (a

---

Thus, when perturbed gradient descent is carried out to optimize the regularized loss Spop eq. ( the algorithm always returns bounded iterates. In conjunction with Theorem 7 this implies, for example. that noisy gradient descent (50] converges to approximate second-order stationary point of the regularized population objective. Moreover, the authors of the same paper show the number of gradient calls made by noisy gradient descent to find an (c, (cSSP) is upper bounded by Ô (1/ min(c², >²),)) for non-convex objective with O(1 smmooth gradients and Hessians. Combining with the choices of and y in eq. (TO) results in the following theorem.
Theorem 9. Assume Gaussian measurements and rescale so that ||Ux|lop = 1. Given v. 2 Tr.
samples (eq. (17)), consider the algorithm which uses the output of greedy pruming (Algorithm 7
Uprane E R as the initial iterate for the fine-tunin, phase, and runs noisy gradient descent [50] on Lemp in the fine-tuning phase. Overall, the algorithm makes,
= (
(o;))o0 log do;
11
kp
(19)
(o7)7
gradient calls and with probability 2 returns a Unu satisfying (umanU U.U? S
Conclusion
In this paper, studied the efficacy of the greedy pruning fine-tuning pipeline in learning low- complexity solutions for the matrix sensing problem, as well as for learning shallow neural networks with quadratic activation functions. We showed that training on the mean squared error augmented by a natural group Lasso regularizer results in models which are suitable for greedy pruning. Given sufficiently many samples, after pruning away the columns below certain fg-norm threshold, arrived at a solution with the correct column sparsity of r. Running a few iterations of gradient descent to fine-tune the resulting model, the population loss was shown converge a linear rate to an error floor of O(vrd/n)) which is also statistically optimal. We also presented negative result showing the importance of regularization while training the model. To the best of our knowledge, our results provide the first theoretical guarantee on the generalization error of the model obtained via the greedy pruning + fine-tuning framework.
Acknowledgements
The research of A. Mokhtari supported in part by NSF Grants 2007668 and 2127697, ARO Grant W911 10226, the National AI Institute for Foundations of Machine Learning (IFML). the Machine Learning Lab (MLL), and the Wireless Networking and Communications Group (WNCG) industrial affiliates program at UT Austin. Kannan Ramchandran would like to acknowledge support from NSF CIF-2002821 and ARO fund 05 05224-0
References
[1] Raghuraman Krishnamoorthi, Quantizing deep convolutional networks for efficient inference:
A whitepaper, 2018.
[2] Yunhui Guo. A survey on methods and theories of quantized neural networks, 2018.
[3] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
[4] Yang He, Guoliang Kang, Xuanyi Dong, Yanwei Fu, and Yi Yang. Soft filter pruning for accelerating deep convolutional neural networks, 2018.
[5] Yawei Li, Shuhang Gu, Luc Van Gool, and Radu Timofte. Learning filter basis for convolutional neural network compression. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5623-5632, 2019.
[6] Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference Machine Learning, pages 5544-5555. PMLR, 2020.
10

---

[7] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015.
[8] Antonio Polino, Razvan Pascanu, and Dan Alistarh. Model compression via distillation and quantization, 2018.
[9] Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural net- works with pruning, trained quantization and huffiman coding. arXiv preprint arXiv:1510.00149, 2015.
[10] Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets, 2016.
[11] Vadim Lebedev and Victor Lempitsky. Fast convnets using group-wise brain damage. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2554- 2564, 2016.
[12] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances in neural information processing systems, 29, 2016.
[13] Mao Ye, Chengyue Gong, Lizhen Nie, Denny Zhou, Adam Klivans, and Qiang Liu. Good subnetworks provably exist: Pruning via greedy forward selection. In International Conference on Machine Learning, pages 10820-10830. PMLR, 2020.
[14] Brian Bartoldson, Ari Morcos, Adrian Barbu, and Gordon Erlebacher. The generalization- stability tradeoff in neural network pruning. Advances in Neural Information Processing Systems, 33:20852-20888 2020.
[15] Christos Louizos, Max Welling, and Diederik P. Kingma. Learning sparse neural networks through lo regularization, 2017.
[16] Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. In Proceedings of the IEEE international conference on computer vision, pages 2736-2744, 2017.
[17] Jianbo Ye, Xin Lu, Zhe Lin, and James Z. Wang. Rethinking the assumption in channel pruning of convolution layers, 2018. 8.
[18] Simone Scardapane, Danilo Comminiello, Amir Hussain, and Aurelio Uncini. Group sparse regularization for deep neural networks. Neurocomputing. 241:81-89, 2017.
[19] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterizee matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 2-47. PMLR, 2018.
[20] Simon Du and Jason Lee. On the power of over-parametrization in neural networks with quadratic activation. In International conference on muchine learning, pages 1329-1338. PMLR, 2018.
[21] Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the op- timization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory, 65(2)7742-769, 2018.
[22] Yudong Chen and Martin J. Wainwright. Fast low-rank estimation by projected gradient descent: General statistical and algorithmic guarantees, 2015.
[23] Jiacheng Zhuo, Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the computational and statistical complexity of over- over-parameterizer matrix sensing. CoRR, abx/2102.02756, 2021.
[24] Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization, Advances in Neural Information Processing Systems, 34:1499-1439, 2021.
[25] Raghunandan Keshavan and Sewoong Oh. gradient descent algorithm on the grassman manifold for matrix completion. 910, 10 2009.

---

[26] Dohyung Park, Anastasios Kyrillidis, Constantine Carmanis, and Sujay Sanghavi. Non-square matrix sensing without spurious local minima via the burer-monteiro approach. In Artificial Intelligence and Statistics, pages 65-74. PMLR, 2017.
[27] Jian-Feng Cai, Emmanuel J. Candès, and Zuowei Shen. singular value thresholding algorithm for matrix completion. SIAM Journal on Optimization, 20(4):1956-199 2010.
[28] Prateck Jain, Raghu Mcka, and Inderjit Dhillon. Guaranteed rank minimization via singular value projection. Advances in Neural Information Processing Systems, 23, 2010.
[29] Raghunandan Keshavan, Andrea Montanari, and Sewoong Oh. Matrix completion from a few entries. IEEE transactions information theory, 5666):2900-2998, 2010.
[30] Shiqian Ma, Donald Goldfarb, and Lifeng Chen. Fixed point and bregman iterative methods for matrix rank minimization. Mathematical Programming, 1280 128(1-2):3213335 2011.
[31] Benjamin Recht. A simpler approach to matrix completion. Journal of Machine Learning Research, 12(12), 2011.
[32) Benjamin Recht, Maryam Fazel, and Pablo A. Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM Review, 52(3):4115501, jan 2010.
[33] Kim-Chuan Toh and Sangwoon Yun. An accelerated proximal gradient algorithm for nuclear norm regularized least squares problems. Pacific Journal of Optimization, 6, 09 2010.
[34] Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has spurious local minimum. Advances in neural information processing systems, 29, 2016.
[35] Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Global optimality of local search for low rank matrix recovery. Advances in Neural Information Processing Systems, 29, 2016.
[36] Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. CoRR, Abs/1704400778, 2017.
[37] Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society Series B, 68:49-67, 02 2006.
[38] Paris Giampouras, Athanasios A Rontogiannis, and Konstantinos Koutroumbas. Alternating iteratively reweighted least squares minimization for low-rank matrix factorization. IEEE Transactions on Signal Processing, 67(2):490-503, 2018
[39] Paris Giampouras, René Vidal, Athanasios Rontogiannis, and Benjamin Haeffele. A novel variational form of the schatten-p quasi-norm. Advances in Neural Information Processing Systems, 33:21453-21 2020.
[40] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning. arXiv preprint 1rX1v:1810.05270, 2018.
[41] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. 2014.
[42] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-28) 2878, 2018
[43] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
[44] Po-Ling Loh and Martin Wainwright. Regularized m-estimators with nonconvexity: Statistical and algorithmic theory for local optima. Advances Neural Information Processing Systems, 26, 2013.
12

---

[45] Emmanuel J Candes and Yaniv Plan. Tight oracle inequalities for low-rank matrix recovery from minimal number of noisy random measurements. IEEE Transactions on Information Theory, 57(44)23322 2011.
[46] Sahand Negahban and Martin J Wainwright. Restricted strong convexity and weighted ma- trix completion: Optimal bounds with noise. The Journal of Machine Learning Research, 13(1):1665-1697, 2012.
[47] Vladimir Koltchinskii, Karim Lounici, and Alexandre B Tsybakov. Nuclear-norm penalization and optimal rates for noisy low-rank matrix completion. 2011.
[48] Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael I Jordan. How to escape saddle points efficiently. International conference on machine learning, pages 1724-1732. PMLR, 2017.
[49] Aryan Mokhtari, Asuman Ozdaglar, and Ali Jadbabaie. Escaping saddle points in constrained optimization. In Bengio, H. Wallach, H. Larochelle, K. Grauman, Cesa- Resa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[50] Chi Jin, Rong Ge, Pranecth Netrapalli, Sham M Kakade, and Michael I Jordan. How to escape saddle points efficiently. In International conference machine learning, pages 1724-1732. PMLR, 2017.
[51] Roman Vershynin. High-dimensional probability. 2019.
[52] B. Laurent and P. Massart. Adaptive estimation of quadratic functional by model selection. The Annals of Statistics, 28(5):1302 1338, 2000.
[53] Michel Petrovitch. Sur une manière d' d'étendre le théorème de la moyenne aux équations différentielles du premier ordre. (mit 2 figuren im text). Mathematische Annalen, 54:417-436. 1901.
[54] Joel A. Tropp. User-friendly tail bounds for sums of random matrices. Foundations of Compu- tational Mathematics, 12(4):389 434,, aug 2011.
[55] Martin J. Wainwright. High--Dimessional Statisties: A Non-Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019.
[56] Rong Ge. Zhize Li, Weiyao Wang, and Xiang Wang. Stabilized svrg: Simple variance reduction for nonconvex optimization. Conference on learning theory, pages 1394-1448. PMLR, 2019.
[57] Cong Ma, Kaizheng Wang, Yuejie Chi, and Yuxin Chen. Implicit regularization in noncon- vex statistical estimation: Gradient descent converges linearly for phase retrieval and matrix completion. In International Conference on Machine Learning, pages 3345-3354. PMLR, 2018.
13


Reviews:
[{"review_id": "WpyJhokJvN", "reviewer": "Reviewer_mDyb", "summary": "In \"Greedy Pruning with Group Lasso Provably Generalizes for Matrix Sensing\" the authors propose a two phase scheme to solve the noisy overparametrized matrix sensing problem. The goal is not to only solve the noisy overparametrized matrix sensing problem but also to provide a theoretically supported intuition for pruning + refining schemes used to train more complex learning models. More specifically, the authors establish that pruning --which entails setting to zero columns with norm below a certain threshold-- any given approximate second order stationary point of the group-lasso regularized matrix sensing problem yields a point with sufficiently small generalization error allowing the use gradient descent and achieving optimal statistical precision.", "questions": "1. The authors suggest the use of perturbed gradient descent to find an approximate second order stationary point. I think it would be a good idea, for completeness, to write the number of oracle calls required to achieve the desired precision using this scheme.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very well written and easy to follow. The problem considered (noisy overparametrized noisy matrix sensing)  is relevant by itself and additionally can provide insights for more complicated learning models which can be of great interest to the community. The results ar \n\nPruning as a technique to solve this specific problem is very well motivated based on both Theorem 1 and achievable statistical precision in the overparametrized setting (c.f. line 312).", "weaknesses": "In Theorem 3, the regularizer weight and the thresholding for pruning are both explicitly dependent on the target rank r. Thus, knowledge of r is required to a degree anyway. Consequently, the comparison made to the overparametrized setting from [23] does not seem entirely fair.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.\n\nAlthough the stated bounds in Theorem 3 explicitly use $r$ and $\\sigma^\\star_r$, this is just for simplicity of expression. It is possible to obtain bounds which rely on $\\lambda$ and $\\beta$ lying in a particular range without explicit knowledge of $r$ and $\\sigma^\\star_r$. Results of this form are the standard de-facto even for the related problem of sparse regression (LASSO). See eg. Theorem 1 of Loh et al (2013). That said, if we are off by a constant factor in each of these parameters, we can still recover the current theorem statement with the sample complexity guarantee larger by a constant.\n\nWe shall also include a bound on the number of optimization steps required by perturbed gradient descent to find an $(\\epsilon,\\gamma)$-SOSP. The authors of [1] show that it is upper bounded by $\\widetilde{O} (1/\\varepsilon^2)$ for a Lipschitz and $\\rho$-Hessian smooth objective when $\\varepsilon= \\min \\\\{\\epsilon, \\gamma^2/\\rho\\\\}$.\n\n[1] Jin et al., 2017. How to Escape Saddle Points Efficiently"}, {"review_id": "DH9hsJYrpD", "reviewer": "Reviewer_oETd", "summary": "The authors study an algorithm framework for symmetric matrix sensing based on overparametrized Burer-Monteiro matrix factorization which additionally imposes a group lasso penalty on the columns of the factor matrix, and which intertwines a gradient descent-type approach for this objective with iterative pruning.\n\nThey argue that gradient descent applied to the unregularized loss does, on the other hand, lead to factor matrices whose generalization properties are negatively affected by magnitude pruning.\n\nAs main theoretical contribution, the authors show in Theorem 4 that for a sample complexity whoses dependence is quadratic in the ground truth rank and of order 5 in the overparametrization rank, algorithmic parameters can be chosen appropriately so that after pruning, an approximate second-order stationary point (which can be arguably found via perturbed gradient descent) can be followed by gradient descent to obtain a solution that is close to the ground truth.\n\nA connection to quadratic shallow neural networks is further used to justify similar approaches combing greedy pruning of a regularized loss with post-pruning fine tuning for more complex machine learning models.", "questions": "- $f_\\operatorname{emp}$  is somehow only defined on page 7. Mentioning it before the introduction of Algorithm 1 will improve the readability of the paper.\n- Can the assumptions of Theorem 4 and Theorem 5 be better clarified? It seems that the RIP is used, it would be better be mentioned in the theorems.\n- How do we end up with the sample complexity mentioned in Remark 6?\n- In line 314, it is mentioned that linear convergence cannot be established due to ill-conditioning of the objective. Please elaborate on whether this could be mitigated algorithms that take the problem structure better into account than standard gradient descent, such as preconditioned examples (e.g., PrecGD from \"Zhang, J., Fattahi, S., & Zhang, R. Y. (2021). Preconditioned gradient descent for over-parameterized nonconvex matrix factorization. _Advances in Neural Information Processing Systems_, _34_, 5985-5996.\").", "limitations": "The authors mostly covered well the limitations of their approach. In the sample complexity results, it would have been beneficial, if would have been better if the suboptimal parameter dependency (at least on $k$ and $n$). had been mentioned.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper provides insights towards relevant insights on how an overparametrized Burer-Monteiro factorization can be made successful and computationally efficient.\n- The manuscript is well-structured and clearly written.\n- From my point of view, the work provides a theoretical justification of the popular practice in deep neural network training of alternating the pruning weight matrices and fine-tuning steps.\n- The technical work to establish the results is quite extensive.\n- The benefit of group lasso regularization is for the problem type considered is clearly illustrated in a simple computaitonal experiment.\n", "weaknesses": "- Many of the results in the paper focus on the presented population setting, in which the ground truth is assumed to be known to the algorithm. While they can arguably be seen as paving the road for the setting in which the (regularized) empirical loss minimized, the paper could have benefitted from a clearer focus and presentation on this more realistic case.\n- Some relevant literature has not been included in the manuscript. In particular, the proposed group Lasso regularization is presented to be new for the problem in consideration, however, similar regularization and a tailored algorithm was already proposed for the very problem (in the more general rectangular case) in \n\t- Giampouras, Paris V., Athanasios A. Rontogiannis, and Konstantinos D. Koutroumbas. \"Alternating iteratively reweighted least squares minimization for low-rank matrix factorization.\" _IEEE Transactions on Signal Processing_ 67.2 (2018): 490-503.\n\tIt also already introduced the pruning step in its algorithm.\n\tFurthermore, a connection of the group Lasso penalty considered with a Schatten-1/2 norm on the product matrix variable was established in:\n\t- Giampouras, P., Vidal, R., Rontogiannis, A., & Haeffele, B. (2020). A novel variational form of the Schatten-$ p $ quasi-norm. _Advances in Neural Information Processing Systems_ (NeurIPS 2020), _33_, 21453-21463.\n- The proof of Theorem 1 is potentially erroneous. In particular, in Appendix A, the quantity $r(t)= U^T U_\\star$ is not a vector in $\\mathbb{R}^k$ , but a $(k \\times k)$-matrix. In the proof, it is consistently treated as a vector, which leads to several problematic statements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Many of the results in the paper focus on the presented population setting ... on this more realistic case.\n\nThank you for your comment. As the reviewer pointed out, our paper presents results for both the population loss and the finite sample loss. We chose to begin by showcasing the results for the population loss, as it facilitates easier comprehension of the updates and analysis and highlights the intuitions behind the proposed approach (such as eq. (13)). Once this groundwork is established, we leverage some of the intuitions from the population loss to derive our results for the more realistic finite sample setting. \n\n> Some relevant literature has not been included ... (NeurIPS 2020), 33, 21453-21463.\n\nThank you for bringing these papers to our attention. We would like to clarify that our paper does not claim to be the first to utilize Group Lasso regularization for low-rank matrix factorization; such a claim is never made in our work. Our primary contribution lies in establishing a formal, and theoretical connection between this regularization technique and learning a column sparse model, along with its relevance to pruning. This is not addressed in the existing literature. Notably, our research is the first to present both negative and positive results on the connection of group lasso and sparse models in the overparameterized setting. Prior to our work, there were no negative results showing the failure of gradient descent updates to learn a column-sparse model without this regularization. Similarly, there was no theory demonstrating that, with this regularization, the obtained model not only possesses the correct rank but also exhibits sparse columns, with the number of active columns matching the true rank.\n\n\nThat said, we agree with the reviewer that the mentioned papers should be added to the literature review of our paper. We will explicitly highlight the contributions of these papers. For instance, in [Giampouras et al., 2018], the authors proposed a regularizer (which is slightly different from Group Lasso) that enforces column sparsity on the two matrix factors. However, they do not provide a formal characterization of whether such regularization ``provably\" leads to a column-sparse model and, if so, whether it allows for exact recovery of the true rank $r$ active columns. Notably, their work proves asymptotic convergence of alternating minimization to a stationary point of the regularized loss. Furthermore, in [Giampouras et al., 2020], the authors introduced a novel variational form of the Schatten-$p$ quasi-norm, which can be considered a direct generalization of the variational form of the nuclear norm. They demonstrated that the local minimizers of the factorized problem also serve as local minimizers of the original problem, and they presented rank-one updates to escape from \"bad\" stationary points. Despite these contributions, similar to the other work, this paper does not establish that the local minimizer of the regularized problem or the obtained model using their proposed method would lead to a model with only $r$ columns that have a large norm, ensuring column sparsity. \n\n> The proof of Theorem 1 is potentially erroneous...\n\nWe respectfully disagree with this comment, as we have specifically mentioned in our paper that the proof of Theorem 1 focuses on the case that the correct rank is $r=1$. Hence, there is no error in the proof. The result shows thar vanilla gradient descent results in a column-dense solution that cannot be easily pruned even when the underlying ground truth is rank $1$.\n\n>Clarifying minor comments: [1.] Can the assumptions of Theorem 4 ...\n\nThe reviewer is correct and we will explicitly mention in the statement of these theorems that our results in Theorems 4 and 5 require the RIP condition. Thanks for raising this point. \n\n> [2.] How do we end up ...\n\nThe sample complexity bound in Remark 6 follows by combining the sample complexity bound in Theorems 4 and 5. The bound comes from adding up the $\\widetilde{O} \\left(\\frac{\\sigma^2}{\\left(\\sigma_r^{\\star}\\right)^4} d k^2 r^5 \\right)$ samples required in the pruning phase (Theorem 4) + the $O \\left(\\frac{\\sigma^2}{(\\sigma_r^\\star)^4} \\frac{rd}{\\varepsilon^2} \\right)$ samples required in the fine-tuning phase (Theorem 5, eq. (16)) to get to $\\varepsilon$ error. Note that under Gaussian measurements the RIP condition in Theorem 4 for $\\delta \\le (\\sigma_r^\\star)^{3/2}/\\sqrt{k} r^{5/2}$ is satisfied with these many measurements.\n\n> [3.] In line 314, it is mentioned ...\n\nThe reviewer's comment is correct, however, our claim was stated specifically for gradient descent. There are a few papers proposing algorithms that circumvent the ill-conditioning of the objective, including Zhang et al (2021). We will rephrase this statement to refer to the above paper. The point we were trying to make here is that gradient descent itself no longer suffers from this slowdown after pruning is carried out since ill-conditioning is only an issue for the overparameterized case. Thanks for your comment. "}, {"review_id": "KG4ARzUQaF", "reviewer": "Reviewer_J6kp", "summary": "This paper studies the pruning for the matrix sensing problem. Particularly, they study the over-parameterized symmetric matrix sensing problem. They show that the approximate solution from a smoothed version of the group-Lasso problem can be pruned to small error values using a greedy pruning algorithm. If the pruned solution is further fine-tuned, the reconstruction error can be further reduced. ", "questions": "- Is the upper bound of Thm. 3 tight? In other words, is fine-tuning necessary?\n", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The results that show fine-tuning a local minimum can further improve its quality is interesting given the recent similar experimental observations in deep learning.\n- Thm. 3 does not require global optimality which is interesting.", "weaknesses": "- Thm. 1 concerns the continuous time version of gradient descent (gradient flow) which would limit the applicability of the theorem in practice. However, this should be fine overall.\n- The penalty discussed in Section 4 is somewhat arbitrary. I think the authors should motivate the choice of this penalty better. Is this penalty used in practice?\n- Verifying SOSP property for a solution can be hard in practice. \n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">  [1.] The penalty discussed in Section 4 is somewhat arbitrary. I think the authors should motivate the choice of this penalty better. Is this penalty used in practice?\n\nThe penalty we discuss in Section 4 is simply a smoothing of the $\\ell_2$ norm (which is not differentiable at $0$). In particular, by replacing $\\| v \\|_2 \\approx \\frac{\\| v \\|_2^2}{\\sqrt{\\| v \\|_2^2 + \\delta}}$ for a small $\\delta > 0$, we now have a loss which is differentiable around $0$. This relaxation of the group LASSO objective allows us to study the stationary points of a non-convex, but now smooth objective.\n\n> [2.] Verifying SOSP property for a solution can be hard in practice.\n\nVerifying the SOSP property is indeed hard in practice, but for an algorithm like perturbed gradient descent, there is no need to have to check this property as the algorithm is theoretically guaranteed to reach such solutions. In practice, stochastic gradient methods (operating on batches of data) are implemented, which bears some resemblance to the perturbed gradient descent we study in this paper.\n\n> [3.] Is the upper bound of Thm. 3 tight? In other words, is fine-tuning necessary?\n\nThat is a good point. The result of Theorem 3 establishes an upper bound on the error of the pruned solution. However, to demonstrate the necessity of fine-tuning and establishing the tightness of the upper bound, a lower bound would need to be proven. In practical scenarios, we consistently observe that some degree of fine-tuning is necessary and beneficial, enhancing the quality of the pruned solution. It would be valuable if a lower bound for the error of the pruned model could also be provided to validate the accuracy of the upper bound.\n\nNevertheless, even if the initial error before fine-tuning matches the lower bound, it can be rapidly reduced during the fine-tuning process, thanks to the linear convergence rate we experience in this phase."}]

Author rebuttal:

Meta-review: This work provides a theoretical justification for greedy pruning, using group lasso in a matrix sensing application. The problem is of interest and the reviewers' concerns were addressed in the discussion period. Please incorporate the changes requested by the reviewers in the final camera ready revision, e.g. "adding some more context (and possibly references to certain theorems proving SOSP for perturbed GD)".
Decision: Accept (poster)