Greedy Pruning with Group Lasso Provably Generalizes for Matrix Sensing

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 a 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 $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 \| U e_i \|_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. 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.

Paper

References (73)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer J6kp7/10 · confidence 4/52023-06-22

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.

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. - 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. - 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? - Verifying SOSP property for a solution can be hard in practice.

Questions

- Is the upper bound of Thm. 3 tight? In other words, is fine-tuning necessary?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

n/a

Reviewer oETd5/10 · confidence 4/52023-07-05

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. They 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. As 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. A 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.

Strengths

- The paper provides insights towards relevant insights on how an overparametrized Burer-Monteiro factorization can be made successful and computationally efficient. - The manuscript is well-structured and clearly written. - 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. - The technical work to establish the results is quite extensive. - The benefit of group lasso regularization is for the problem type considered is clearly illustrated in a simple computaitonal experiment.

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. - 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 - 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. It also already introduced the pruning step in its algorithm. Furthermore, a connection of the group Lasso penalty considered with a Schatten-1/2 norm on the product matrix variable was established in: - 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. - 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.

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. - 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. - How do we end up with the sample complexity mentioned in Remark 6? - 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.").

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

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.

Area Chair Vy3X2023-08-16

Thanks for your review, the authors have provided a rebuttal. Please reply below and let us know if your concerns are addressed, in particular the question about the correctness of Theorem 1. Best, the AC

Reviewer mDyb7/10 · confidence 3/52023-07-06

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.

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

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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

N/A

Reviewer oETd2023-08-16

Thank you for the clarifications. Regarding Theorem 1, I indeed missed that you restricted yourselves to the case of $r=1$, in which case the stated matrices can indeed interpreted as vectors. The concerns I had have been also addressed, if the literature review is updated accordingly in the final version. I am accordingly increasing my score by 1.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC