SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nived Rajaraman
EECS Department
nived rajaraman@eecs.berkeley.edu
Devvrit
CS Department
Kannan Ramchandran
kannanr@eecs.berkeley.edu
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} \|Ue_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_{\mathrm{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_{\mathrm{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.
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 [1, 2], unstructured [3] and structured [4-6] pruning mechanisms, and distillation techniques using student-teacher models [7, 8]. Among these approaches, greedy pruning, in which we greedily eliminate the parameters of the trained model
based on some measure (e.g., the norms of the weight vectors associated with individual neurons) has received widespread attention [3, 9-13]. This is mostly due to the fact that several practical studies have illustrated that training an overparameterized 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  $\ell_0$  or  $\ell_1$  regularization [15-17] or  $\ell_2$  regularization including group Lasso [18] 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 no 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] as a test-ground for the analysis of 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  $\| U U^T - U_\star U_\star^T \|_F^2$ , where  $U_\star \in \mathbb{R}^{d \times r}$  is an unknown ground-truth rank- $r$  matrix, with  $r$  also being unknown, and  $U \in \mathbb{R}^{d \times k}$  is the overparameterized learning model with  $k \gg r$ . As we discuss in the Appendix F, 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  $k$  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 to 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 ground-truth  $U_{\star}U_{\star}^{T}$  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 error with an added group Lasso based regularizer to promote column sparsity, followed by a simple norm-based pruning strategy results in a model  $U_{\mathrm{prune}}$  having exactly the minimum number of columns,  $r$ . At the same time, we show that  $\| U_{\mathrm{prune}}U_{\mathrm{prune}}^{T} - U_{\star}U_{\star}^{T}\|_{F}^{2}$  is small, but non-zero. Hence, the pruned model can subsequently be fine-tuned using a small number of gradient steps, and in this regime,  $\| U_tU_t^T - U_\star U_\star^T\|_F^2$  shows linear convergence to its limit. Moreover, the pruned model can be shown to admit finite sample generalization bounds which are also statistically optimal for a range of parameters. In particular, we show that to obtain a model  $U_{\mathrm{out}}$  that has exactly  $r$  columns and its population error is at most  $\| U_{\mathrm{out}}U_{\mathrm{out}}^{T} - U_{\star}U_{\star}^{T}\|_{F} \leq \varepsilon$ , our framework requires  $O(dk^2r^5 + \frac{rd}{\varepsilon^2})$  samples, which is statistically optimal for sufficiently small  $\varepsilon$ . 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, 19, 23] establishing that gradient descent in the "exactly-parameterized" setting requires  $O(rd / \varepsilon^2)$  samples to achieve a generalization error of  $\varepsilon$ , and converges 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  $O(kd / \varepsilon^2)$  samples to achieve a generalization error of  $\| U_{\mathrm{gd}}U_{\mathrm{gd}}^T - U_\star U_\star^T\|_F \leq \varepsilon$ , degrading with the overparameterization 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,  $U_{\mathrm{gd}}$  has to be post-processed by computing its SVD, which is computationally expensive in the high dimensional regime.
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  $\ell_2$ -norms comparable to that of the maximum, even if  $r = 1$ . Hence, it is unclear how to sparsify such models based on the norms of their columns.
2 Setup and Algorithmic Framework
Given  $n$  observation matrices  $\{A_i\}_{i=1}^n$ , in the matrix sensing framework, the learner is provided measurements  $y_i = \langle A_i, U_\star U_\star^T \rangle + \varepsilon_i$  where  $\langle \cdot, \cdot \rangle$  indicates the trace inner product and  $\varepsilon_i$  is measurement noise assumed to be distributed i.i.d.  $\sim \mathcal{N}(0, \sigma^2)^1$ . Here  $U_\star \in \mathbb{R}^{d \times r}$  is the unknown parameter, and the rank  $r \leq d$  is unknown. The goal of the matrix sensing problem is to learn a candidate matrix  $X$  such that  $X \approx U_\star U_\star^T$ . For computational reasons, it is common to factorize  $X$  as  $UU^T$  for  $U \in \mathbb{R}^{d \times k}$ . In this paper, we study the factored model in the overparameterized setting, where  $k \gg r$ . The empirical mean squared error is,
$$
\mathcal {L} _ {\mathrm {e m p}} (U) = \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\left\langle A _ {i}, U U ^ {T} \right\rangle - y _ {i}\right) ^ {2}. \tag {1}
$$
For the case that  $A_{i}$ 's are sampled entry-wise i.i.d.  $\mathcal{N}(0,1 / d)$  and as  $n\to \infty$ , up to additive constants which we ignore, the population mean square error can be written down as,
$$
\mathcal {L} _ {\text {p o p}} (U) = \left\| U U ^ {T} - U _ {\star} U _ {\star} ^ {T} \right\| _ {F} ^ {2}. \tag {2}
$$
There is an extensive literature on how to efficiently learn the right product  $U_{\star}U_{\star}^{T}$  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 [19, 23]. While these approaches guarantee learning the correct product  $U U^T \approx U_{\star}U_{\star}^{T}$ , in the overparameterized setting the obtained solutions are not column sparse and storing these models requires  $\tilde{\Theta}(kd)$  bits of memory (ignoring precision). As a result, in order to obtain a compressed solution  $U_{\mathrm{out}} \in \mathbb{R}^{d \times r}$  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 is to overcome this issue and come up with an efficient approach to recover a solution  $U_{\mathrm{out}}$  which generalizes well, in that  $\mathcal{L}_{\mathrm{pop}}(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, it is possible to obtain a model that approximately learns the right product  $U_{\star}U_{\star}^{T}$ , while having only a few significant columns. As a result, many of its columns can be eliminated by a simple  $\ell_2$  norm-based greedy pruning scheme, without significantly impacting the training loss. In fact, post pruning we end up with a model  $U_{\mathrm{prune}} \in \mathbb{R}^{d\times r}$  that has the correct dimensionality and its outer product  $U_{\mathrm{prune}}U_{\mathrm{prune}}^{T}$  is close to the true product  $U_{\star}U_{\star}^{T}$ . When the resulting "exactly parameterized" model is fine-tuned, the generalization loss  $\mathcal{L}_{\mathrm{pop}}(U)$  can be shown to converge to 0 at 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 a low rank solution [27-33], or 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-studied in the literature. While these approaches implicitly or 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 [37, 18]. In the population case, the loss is defined as
$$
\mathcal {L} _ {\text {p o p}} (U) + \lambda \mathcal {R} (U), \quad \text {w h e r e} \quad \mathcal {R} (U) = \sum_ {i = 1} ^ {k} \| U e _ {i} \| _ {2}. \tag {3}
$$
Note that  $\lambda > 0$  is a properly selected regularization parameter. Imposing  $\mathcal{R}$  as a penalty on the layer weights of a 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 of the group
<sup>1</sup>All proofs in the paper go through as long as the noise is i.i.d. sub-Gaussian with variance proxy  $\sigma^2$ . We study the Gaussian case for simplicity of exposition.
Lasso penalty with the Schatten-1/2 norm was considered in [39], 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 a smooth version of the regularized loss in (3), denoted by  $U_{\mathrm{out}}$ , is approximately column-sparse, i.e.,  $k - r$  of its columns have small  $\ell_2$  norm. As a result, it is suitable for the simple greedy pruning scheme which we also introduce in Algorithm 1. Interestingly, in Section 3, we will first show a negative result - the model obtained by minimizing the unregularized loss,  $\mathcal{L}_{\mathrm{pop}}(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) on 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  $\ell_1 / \ell_2$  norm of the associated vectors [16, 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-order stationary point of the loss in Section 4, we only keep its columns whose  $\ell_2$  norm are above a threshold (specified in Algorithm 1) and the remaining columns with smaller norm are removed. We further show that post pruning, the obtained model  $U_{\mathrm{prune}}$  continues to have a small empirical loss, i.e., small  $\mathcal{L}_{\mathrm{emp}}(U_{\mathrm{prune}})$ , while having exactly  $r$  columns.
Post-pruning fine-tuning. As mentioned above, it is common to fine-tune the smaller pruned model with a few gradient updates before the evaluation process [40]. We show that the pruned model  $U_{\mathrm{prune}}$  has the correct rank and is reasonably close to the ground-truth model  $U_{\star}$  in terms of its population loss. By running a few gradient updates on the mean square error, it can be ensured that  $UU^{T}$  converges to  $U_{\star}U_{\star}^{T}$  at a linear rate. Algorithm 1 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 [41, 42], a phenomenon known as implicit regularization. In particular, for matrix sensing, [43, 19, 24] show that in the absence of any regularization, running gradient descent on the population loss starting from a small initialization biases  $UU^T$  to low rank solutions, and learns the correct outer product  $UU^T \approx U_\star U_\star^T$ . However, as discussed earlier, low-rank solutions are not 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.
In this section, we address this question and show that minimizing the unregularized population loss  $\mathcal{L}_{\mathrm{pop}}$  leads to models which are not suitable for greedy pruning, i.e., have many columns with large  $\ell_2$  norm. Specifically, we show that by running gradient flow from a small random initialization, even if the ground truth  $U_{\star}$  is just a single column and  $r = 1$ , the learnt solution  $U$  has a large number of columns that are "active", i.e. having  $\ell_2$ -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  $\ell_2$  norm. We thus claim that such trained models are not compatible with greedy pruning. In the following theorem, we formally state our result.
Theorem 1. Consider the population loss in eq. (2), for  $r = 1$  and  $k \gg 1$  and suppose  $\| U_{\star} \|_{op} = 1$ . Further, suppose the entries of the initial model  $U_0$  are i.i.d. samples from  $\mathcal{N}(0, \alpha^2)$ , where  $\alpha \leq c_1 / k^3 d \log(kd)$  for some constant  $c_1 > 0$ . For another absolute constant  $c_1' > 0$ , as  $t \to \infty$ , the iterates of gradient flow with probability  $\geq 1 - O(1 / k^{c_1'})$  converge to a model  $U_{gd}$  where  $\widetilde{\Omega}(k^{c_1'})$  active columns satisfy
$$
\frac {\left\| U _ {g d} e _ {i} \right\| _ {2}}{\max  _ {j \in [ k ]} \left\| U _ {g d} e _ {j} \right\| _ {2}} \geq 0. 9 9. \tag {4}
$$
Figure 1: We run gradient updates to minimize the loss in eq. (2) until convergence. The model is highly overparameterized with  $d = k = 500$  and  $r = 4$ . In the left figure, the vertical axis represents the fraction of columns having their norm at least  $x$  times the largest norm across columns, with  $x \in [0,1]$  on the horizontal axis. In the right figure, for the same experiment we plot a histogram of column norms. Using regularization leads to solutions with most of the columns having small  $\ell_2$  norm, and only a few significant columns. Without regularization, a large number of columns have their  $\ell_2$  norm comparable to the largest column norm. Thus, training with explicit regularization leads to models that are more suitable for pruning.
Theorem 1 claims that the iterates of gradient descent converge to a solution for which a large number  $(\Omega (k^{c_1^{\prime}}))$  of columns of  $U_{\mathrm{gd}}$  have their  $\ell_2$  norm at 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 to the right sparsity level by simply pruning away columns with  $\ell_2$  norm below any fixed constant threshold.
We provide a general guarantee in Theorem 22 (Appendix A). It is worth noting that from [19] that the learned model  $U_{gd} \in \mathbb{R}^{d \times k}$  is guaranteed to achieve a small generalization error, i.e.,  $U_{gd}U_{gd}^T \approx U_\star U_\star^T$ . Hence, the algorithm does learn the correct ground-truth product - what we show is that the learnt model is not "pruning-friendly" - it cannot be sparsified to the correct level by on any simple norm based pruning strategies. This is also observed in Figure 1: the model trained with implicit regularization has many significant columns and it is unclear how many to prune. We provide a 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_{\star}\| _2 = 1$ . Defining  $U(t)$  as the iterate at time  $t$ , gradient flow on  $\mathcal{L}_{\mathrm{pop}}$  follows the below dynamic,
$$
\frac {d U}{d t} = - \nabla \mathcal {L} _ {\text {p o p}} (U) = - \left(U U ^ {T} - U _ {\star} U _ {\star} ^ {T}\right) U, \tag {5}
$$
where for ease of notation we drop the explicit dependence on  $t$  in  $U(t)$ . In our proof, we exactly characterize the limiting point of gradient flow as a function of the initialization in a certain sense, which may be of independent interest. In particular, up to a small multiplicative error, we show,
$$
\forall i \in [ k ], \| U _ {\mathrm {g d}} e _ {i} \| _ {2} \approx \frac {| \langle U _ {\star} , U (0) e _ {i} \rangle |}{\| U _ {\star} ^ {T} U (0) \| _ {2}}. \tag {6}
$$
Also, with Gaussian initialization,  $\langle U_{\star},U(0)e_i\rangle \stackrel {\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\alpha^2)$  across different values of  $i$ . In particular, we show that for some constant  $c_{1}^{\prime} > 0$ ,  $\widetilde{\Omega} (k^{c_1'})$  columns will have correlations comparable to the maximum, with  $|\langle U_{\star},U(0)e_i\rangle |\geq 0.99\max_{j\in [k]}\left|\langle U_{\star},U(0)e_j\rangle \right|$ . For any of these columns,  $i$ ,
$$
\left\| U _ {\mathrm {g d}} e _ {i} \right\| _ {2} \approx \frac {\left| \langle U _ {\star} , U (0) e _ {i} \rangle \right|}{\left\| U _ {\star} ^ {T} U (0) \right\| _ {2}} \geq 0. 9 9 \max  _ {j \in [ k ]} \frac {\left| \langle U _ {\star} , U (0) e _ {j} \rangle \right|}{\left\| U _ {\star} ^ {T} U (0) \right\| _ {2}} \approx 0. 9 9 \max  _ {j \in [ k ]} \left\| U _ {\mathrm {g d}} e _ {j} \right\| _ {2}, \tag {7}
$$
which completes the proof sketch of Theorem 1.
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
Algorithm 1 Greedy pruning based on group-Lasso regularization
Inputs: Measurements  $\{(A_i, y_i)\}$  where  $y_i = \langle A_i, U_\star U_\star^T \rangle + \varepsilon_i \}_{i=1}^n$  (in the population setting  $n = \infty$ );
Initialization: Set parameters:  $\lambda, \beta, \epsilon, \gamma$  and  $m_{\mathrm{fine - tune}}$ .
Greedy pruning phase:
1: Find an  $(\epsilon, \gamma)$ -approximate SOSP of  $f_{\mathrm{emp}}$  (resp.  $f_{\mathrm{pop}}$ ),  $U$ , satisfying  $\| U \|_{\mathrm{op}} \leq 3$ .
2: Let  $S = \{i \in [k] : \| U e_i \|_2 \leq 2\sqrt{\beta}$  denote the set of columns with small  $\ell_2$  norm.
Create a new matrix  $U_{\mathrm{prune}}$  which only preserves the columns of  $U$  in  $[k] \setminus S$ , deleting the columns in  $S$ .
Fine-tuning phase:
1: Run  $m_{\mathrm{fine - tune}}$  iterations of gradient descent on  $\mathcal{L}_{\mathrm{pop}}(U)$  (resp.  $\mathcal{L}_{\mathrm{emp}}(U)$ ) initialized at  $U_{\mathrm{prune}}$  to get  $U_{\mathrm{out}}$ .
2: return  $U_{\mathrm{out}}$ .