Over-parameterized Model Optimization with Polyak-{\L}ojasiewicz Condition
Authors: Yixuan Chen
Yubin Shi
Mingzhi Dong
Xiaochen Yang
Dongsheng Li
Yujiang Wang
Robert P. Dick
Qin Lv
Yingying Zhao
Fan Yang
Ning Gu
Li Shang
Venue: ICLR 2023 poster (2023)
Keywords: Over-parameterized Model
Model Optimization
Polyak-{\L}ojasiewicz Condition.

Abstract: This work pursues the optimization of over-parameterized deep models for superior training efficiency and test performance. We first theoretically emphasize the importance of two properties of over-parameterized models, i.e., the convergence gap and the generalization gap. Subsequent analyses unveil that these two gaps can be upper-bounded by the ratio of the Lipschitz constant and the Polyak-{\L}ojasiewicz (PL) constant, a crucial term abbreviated as the \emph{condition number}. Such discoveries have led to a structured pruning method with a novel pruning criterion. That is, we devise a gating network that dynamically detects and masks out those poorly-behaved nodes of a deep model during the training session. To this end, this gating network is learned via minimizing the \emph{condition number} of the target model, and this process can be implemented as an extra regularization loss term. Experimental studies demonstrate that the proposed method outperforms the baselines in terms of both training efficiency and test performance, exhibiting the potential of generalizing to a variety of deep network architectures and tasks.

Published as a conference paper at ICLR 2023
OVER-PARAMETERIZER MODEL OPTIMIZATION WITH LOJASIEWICZ CONDITION
Yixuan Chen', Yubin Shi', .Mingzhi Dong" Xiaochen Yujiang Wang', Robert P. Dick', Qin Lve, Yingying Zhao', Fan Yang', Ning Gu Li Shang! (yi Syixuanchen20, ybshi21| fudan. edu. en
Xiaoceen Yang". Dongsheng Li¹,
ABSTRACT
This work pursues the optimization of over-parameterized deep models for su- perior training efficiency and test performance. We first theoretically emphasize the importance of two properties of over-parameterized models, i.c., the conver- gence gap and the generalization gap. Subsequent analyses unveil that these two gaps can be upper-bounded by the ratio of the Lipschitz constant and the Polyak- Lojasiewicz (PL) constant, a crucial term abbreviated as the condition number. Such discoveries have led structured pruning method with novel pruning criterion. That is, we devise gating network that dynamically detects and masks out those poorly-behaved nodes of a deep model during the training session. To this end, this gating network is learned via minimizing the condition number of the target model, and this process can be implemented as an extra regularization loss term. Experimental studies demonstrate that the proposed method outperforms the baselines terms of both training efficiency and test performance, exhibiting the potential of generalizing to variety of deep network architectures and tasks.
INTRODUCTION
Most practical deep models are over-parameterized with the model size exceeding the training sam- ple size and can perfectly fit all training points (Du et al,) 2018 Vaswani et al.) 2019). Recent empirical and theoretical studies demonstrate that over- paramtterization plays an essential role in model optimization and generalization (Liu et al.] 2021b] Allen-Zhu et al. 2019). Indeed, a plethora of state-of -fate-of-thheart the-art models that prevalent in the community over- parameterized, such as Transformer- based models for natural language modeling tasks (Brown et al. 2020
2020 Devlin et al.]
2018 Liu al. 2019) and wide residual networks for computer vision tasks (Zagoruyko & Ko- modakis 2016). However, training over-parameterized models is usually time-consuming and can take anywhere from hours weeks complete. Notwithstanding some prior works (Liu et al.)
2022 Belkin 2021) on theoretical analyses of the over- weer-paaamttriieed models, those findings re- main siloed from the common practices of training those networks.
The work seeks to optimize over-parameterized models, in pursuit of superior training efficiency and generalization capability. We first analyze two key theoretical properties of over- weerparaaeeeriiee models, namely the convergence gap and the generalization gap, which can be quantified by the con- vergence rate and the sample complexity, respectively. Theoretical analysis of pver-parameterized models is intrinsically challenging as the over- pweemparrmeteri optimization landscape is often non- convex, limiting convexity-base based analysis. Inspired by recent research on the convergence analysis of neural networks and other non-linear systems (Bassily et al.) 2018 Gupta al. 2021 Oymak & Soltanolkotabi 2019), we propose to use the Polyak-Lojasiewin (PL) condition (Polyak] 1963]
Karimi et al.) 2016 Liu al., (202) as the primary mathematical tool to analyze convergence rate and sample complexity for over- wver-parameterized models, along with the widely used Lipschitz con-
'China and Shanghai Key Laboratory of Data Science, School of Computer Science, Fudan University. Shanghai, China. "School of Mathematics Statistics. The University of Glasgow, Glasgow, UK. Miicrosoft Research Asia, Shanghai, China. 'Department of Engineering Science, University of Oxford, Oxford, Eng- land. 'Department of Electrical Engineering and Computer Science University of Michigan, Michigan, United States. 'Department of Computer Science, University of Colorado Boulder, Boulder, Colorado, United States. "School of Microelectronics, Fudan University, Shanghai, China. 'The corresponding author.

---

Published as a conference paper at ICLR 2023
Conditiinn Number
Loss (Log peln
Test Loss (Log
Mokel
Training Epoch
Training Epuch
Training Epoch
(c)
Head
(d)
Figure 1: Benefits of PL regularization for BERT optimization. Figure (a-c) shows that models with a smaller condition number (e.g., T2 < Ti in general) achieve faster training convergence and better test performance. In addition, pruning heads with large condition number, i.e. Masked T,, reduces the condition number, leading to more rapid and accurate training. Figure (d) shows T2 heads have different condition numbers. The largest ones are pruned to produce Masked T in Figure (a-c).
dition (Allen-Zhu et al. 2019). Our theoretical analysis shows that the aforementioned properties can be controlled by the ratio of the Lipschitz constant the PL. constant, which is referred to as the condition number (Gupta et al. [2021) A small condition number indicates a large decrease in training loss after parameter updates and high algorithmic stability relative to data perturbation, i.e., fast convergence and good generalization ability. More promisingly, such pattern can be observed in empirical studies. As shown in Figure [a-c), where BERT models were applied WikiText-2 for language modeling, the training loss of the model with a small condition number (72) decreases much faster than the model with a large condition number (Ti), especially when the differences in condition numbers pronounced (between 40 and 80 epochs); its test performance also improves much faster and is ultimately better. Such thcoretical and empirical findings motivate us to formulate novel regularized optimization problem which adds the minimization of condition number to the objective function; we call this new additional term PL regularization. In this way, we can directly regularize the condition number while training peratertmetri models, thereby improving their convergence speed and generalization performance.
Our empirical analysis further reveals that, given an over-parameterized model, different model components exhibit distinet contributions to model optimization. Figure (d) plots the heatmap of the condition number for all model heads at epoch-10 and it shows that the condition number varies considerably between model heads. Given the fact that wver-parameterized models contain a large number of redundant parameters, we argue that it is possible to reduce the condition number of an over over-paraman wee-paarmeterized network during training by identifying and masking out poorly- -behave sub-ne sub-networks with large condition numbers. Figure [a-c) illustrates the potential efficacy of this ap- proach. After disabling 75% heads of the BERT model T2 according to the condition number ranking at epoch- 10, the masked BERT (Masked 7) possesses a smaller condition number and achieves faster convergence and better test performance. This phenomenon motivates us impose PL regulariza- tion, and hence improve model optimization, by adopting a pruning approach. More specifically, introduce binary mask for periodically sparsifying parameters, and the mask is learned via gating network whose input summarizes the optimization dynamics of sub-networks in terms of PL. regularization. An overview of the proposed method is provided in Appendix E.I
The proposed pruning approach to enforcing PL regularization is related to the structured pruning works, which focuses on compressing model size while maintaining model accuracy. The significant difference lies in that we utilize the condition number to identify the important components, which thus can simultaneously guarantee the convergence and generalization of model training. More importantly, as a consequence of this difference, compared with most pruning works which obtain a sparse model slight cost of degraded accuracy, our method is found to achieve even better test performance than the dense model when no more than 75% of the parameters were pruned. Experimental results demonstrate that our method outperforms state-of-the-ard pruning methods in terms of training efficiency and test performance. The contributions of this work are threefold:
We are the first to propose using PL regularization in the training objective function of over- over-param parameterized models. Such proposal is founded on the theoretical analysis of opti- mization and generalization properties of over parammeetteee models, which shows that a small condition number implies fast convergence and good generalization.
We describe PL. regularization-dree model optimization. Specifically, we introduce learnable mask, guided by the PL-based condition number, to dynamically sparsify poorly-behaved sub-networks during model training to optimize training efficiency and test performance.
structured pruning method for over- over-parameterized
2

---

Published as a conference paper at ICLR 2023
The proposed analysis framework and optimization method can be applied to wide range of deep models and tasks. Experimental studies using three widely used deep models, namely BERT, Switch-Trame Transformer, and VGG-16, demonstrate that the proposed method outperforms the baselines on both training efficiency and test performance.
2 ANALYSIS OF
MODELS WITH PL CONDITION
In this section, we first show that, in the context of over-parameterized models, the optimal test emror can be decomposed into convergence gap and generalization gap, which can be quantified via convergence rate and sample complexity, respectively. Next, using the tool of PL" condition, we show that both convergence rate and sample complexity upper-bounded by the condition number and thus can be improved by using a small condition number.
2.1 OPTIMIZATION PROPERTIES OF OVER-PARAMETIO
MODELS
Let X' denote the input space and y the target space. The standard optimization problem can be formulated as finding a minimizer of the expected risk: f" = argmin R(f), where the D denotes an unknown data distri- bution over X' x y. and L denotes non-negative loss function. To restrict the search space, f is typically chosen from a hypothesis class H. C y²² the minimizer this case is given by fix arrminn f. R(f). In practice, the data distribution is unknown and thus the minimizer is found by minimizing the empirical risk: fi.s argming:fen Rs(f), where the empirical risk is defined as Rs := +/h - (xx(),,. vy..aa and S denotes the training set S = {(x, y,) drawn i..d. from D. The empirical risk minimization problem is often solved by using algorithms such as gradient descent; the optimal solution found by these algorithms denoted by sH.s. Finally, when f' is parameterized by w E Rmm we denote the loss function by y; w), or simply L(W). and also use (w) to denote the empirical risk.
expected risk is defined as R(f) == (xx)))==(((((x).))))))))))))))))
Definition 1 (Over-parameterized model Liu et al. (2020a)). A model is over-parameterized model if the number of parameters m larger than the number of data.
Motivated by the previous work (Gühring et al., 2020), the Bayes error, i.e., the optimal test error, can be decomposed into five components: empirical error, convergence gap, generalization gap, estimation error and approximation error. In the prrametraeteri setting, following previous works (Brutzkus et al.) 2017; Belkin [202]] Liu et al., (2022). we assume that the optimal function fi.s fits or interpolates the training data exactly. Therefore, for over-parameterized models, the Bayes error decomposition can be simplified as
(I) The first and second terms refer to the convergence gap and the generalization gap, respectively. (The detailed decomposition can be found in Appendix A) Thus, the Bayes error can be better approxi- mated if we restrict the convergence gap and generalization gap. Furthermore, the convergence gap can be quantified in proportion to the reciprocal of convergence rate, and the number of instances required by specific generalization gap can be quantified in proportion to sample complexity.
R(f*) & 2 (fRs(f..s) Rs(fū,s)l + sup R(f) - Rs()).
FEH
Definition (Convergence Rate (Alzubi et al. 20188). If the model sequences I (wwo).... : f(w)) during optimization converge to the optimal solution fww",, convergence rate can be defined as:
(We+1)
(w")
lim
= Y
t-+00
(w:) - (w*)
where y is a real number in range [0, 1]. The convergence rate is measured by the limit the ratio of successive differences. The sequence converges sublinearly, linearly, or superlinearly depending on whether = =1, E (0, 1), or ? = O. respectively. For an algorithm with a linear rate of convergence, Rs([(w,)) S. (ep((-n) (Rs(s(wo)))
Definition 3 (Sample Complexity (Alzubi et al. 2018)). For any 5, (0,1), the sample complexity ne,, 8) is defined as the smallest n < Do for which there exists a learning algorithm I such that, for any distribution 'D on training dataset S with the data sample X1,...;X,
:x. of size n,
Pr ([R(fs) Rs(fs)) S <) - 5.
(3)
The sample complexity denotes the number of required examples to guarantee a generalization gap smaller than with probability -
3

---

Published as a conference paper at ICLR 2023
2.2 PL' CONDITION AND THEORETICAL PROPERTIES
The previous subsection explains that faster convergence and better generalization are crucial to the test performance of wee-parameterized models and these measures can be quantified by the con- vergence rate and sample complexity, respectively. In this subsection, we show that both quantities can be analyzed through the PL condition and Lipschitz continuity, whose results suggest that en- couraging an over parameterized model to have smaller condition number will generally result in faster convergence and better generalization.
Definition& 4 (Lipschitz Continuity (Nesterov 2003 The function S is Liplipschitz with re- spect parameters w if there exists a constant L,g, such that: (f(w2) f(w1)l Lyl[w] W2ll, VW1, W2 E W, where W indicates the hypothesis set.
Definition 5 Sallyak-Lojaseewict" Condition (Polyak Polyak-Lojasiewice" (PL") condition with parameter
1963
The functioo f(w) satisfies the Vw E W. [V 2.. f.f((w.
The inequality in the PL* condition implies that every stationary point is global minimizer, which also a property strong convexity. However, unlike strong convexity, the PL. condition does not assume the uniqueness of the minimizer (Bassily et al. 2018). Such difference makes PL* more appropriate mathematical framework to analyze the theoretical properties of over parameterization as the optimization of wee-pparameterized models results in manifolds of global minima.
Following the original argument from Liu et al. (2022); Charles & Papailiopoulos, (2018), we use the PL condition and Lipschitz continuity to analyze the convergence rate and sample complexity of over-parameterized models trained using gradient descent (GD).
Theorem Suppose the loss function L(w) and its first derivative are both Ly-Lipschitz continuous and the loss function satisfies the p-PL" condition in the ball B(wo, R)
{w E R" : :/w www-w woll S R} with R > O, then we have the following:
(a) Convergence Rate of GD (Karimi et al. 2016): Gradient descent with step size n 1.g converges to a global solution with an exponential convergence rate: Cs(we) S (1- Cs
(b) Sample Complexity: Suppose an algorithm A has pointwise hypothesis stability y with respect to a bounded loss function such that S. L(w) S M. The model parameterized by the output of algorithm A trained on dataset S. Then, for any c, 8, we have Pr Pr)(fss) Rs(fs) S 2 1-15
with the sample complexity ne(e, 6) < 8 +
613,
M²
2238
The proof of the theorem is deferred to Appendix B The theorem indicates that 1) smaller T. would result in faster convergence rate and decrease the convergence gap; and 2) a smaller / would result in smaller sample complexity, i.e., fewer samples required to obtain generalization gap smaller than with probability higher than 1 - ó. In other words, the theorem suggests that the convergence rate and generalization ability of an over-parameterized model can be improved by encouraging the model to have smaller Ly/p, which is referred to as the condition number in [Nesterov] 2003] Gupta et al. 202T). Following the previous work, we call the ratio of the Lips- chitz. constant 1.j to the PI. constant j as the condition number.
Definition 7 (Condition Number (Gupta et al. 2021 ). Suppose the loss function C satisfies the PL condition and has Li-iscchitzz continuous gradients, the condition number is defined as Lr.
MODEL OPTIMIZATION WITH PL CONDITION
Based on the analysis in the previous section, advocate optimizing the condition number 4 of over-parameterized models by adding it a regularization term the optimization problem; this technique is termed PL regularization. In this section, we first introduce the analytic form of PI. regularization. Then, we propose PL. regularization--divee structured pruning method.
3.1 FORMULATION OF PL REGULARIZED OPTIMIZATION
To minimize the condition number, we propose the following regularized risk minimization problem, which adds the condition number to the training emror:
4

---

Published as a conference paper at ICLR 2023
min (s( + -1 wEW where μ indicates the PL constant of the neural network with respect to parameters w; Ly is the Lipschitz constant with respect to parameters w: [s(w) indicates the training error of the neural network with parameters w; a is a trade-off parameter.
1
While it is desirable train network with small condition number term, directly optimizing is often difficult. Following the uniform conditioning developed in previous works (Liu et al. 2022 Scaman et al. 2022) for an over- aver-paramettrrized model f with Lipschitz continuous loss function, e.g., the mean square error loss, the analysis of μ. could reduce to calculating and optimizing the minimal eigenvalue of the neural tangent kernel (NTK) associated to Denote the NTK matrix of the function f as K((w) (F DFF (w) E R=**, where DJF represents the differentiable map of f at w, and the smallest eigenvalue of IC as Amin (K[ [[w)) )): and denote μ inf Amin (K((w)
based on the uniform conditioning (Liu et al. 2022 2020a Scamaa al. 2022). Consequently, the objective function is represented as:
min Ls (s( + a Ty
(5)
wEW Following the previous works (Liu et al. 2022; 20206). the following proposition shows that the condition number is upper bounded by Hessian norm |,ll| which is defined as ||H 6 max; | ll,, where H Bw || ll2 is termed the Hessian norm of f hereinafter.
Proposition 8 Given a point wo. B(wo ; R) denotes a ball centered Wo with radius R. > O. Suppose the loss function satisfies the μ²-PL" condition and has Ly-Lipschitz gradi- ents, for all E B(wo, R). (PF(wo)llatsug V/R. (f))lll
Then the condition number is upper bounded by 4 i
intwen Amtaa (K()))
The proof can be found in Appendix C Proposition g motivates us to optimize the upper bound of condition number for PL regularization optimization.
3.2
IMPLEMENTATION OF PL-REGULARIZER OPTIMIZATION VIA PRUNING
In this subsection, we adopt a pruning approach controlled by Hessian norm to solve the new regu- larized risk minimization problem.
3.2.1
LEARNING OBJECTIVE OF PL REGULARIZATION - PRUNING ALGORITHM
First, we propose the following theorem to show that pruning parameters governed by minimizing the Hessian norm H is able to decrease the upper bound of the condition number.
Theorem Suppose the neural network parameterized by w (w E B (Wo: R)) and its loss func- tion L have the same assumptions as in Proposition & Pruning parameters from the parameter set governed by minimizing the Hessian norm [Hyllz can decrease the upper bound of the condition number.
Proof sketch. Based on the Taylor expansion around initialization We. we show that the pruning parameters controlled by minimizing the Hessian norm |Hll2 is able minimize the upper bound of Lipschitz (according to Proposition [); and based on the matrix theory, pruning parameter is able maximize the minimal eigenvalue of NTK (Proposition 140; details can be found in Appendix [D]
To disable unnecessary or harmful parameters, we introduce a binary pruning mask to implement structured pruning, i.e., each element in the mask controls subset of parameters. This design is particularly suitable for parameters with inherent structures, e.g., some parameters are from the same convolutional filter in a CNN or head in Transformer. Suppose the parameters w consists of p subsets of parameters and let Wij denote the 1th subset of the parameters. Then the mask m, Boolean vector of length p. sparsifies the parameters as follows: Wij disabled, i.e., all elements of WiO) equal to 0, if m, = 0: otherwise, Wij remains the same. With the introduction of the pruning mask, Eq. 5 can be transformed to the following optimization problem, which additionally learns the binary mask to prune the network:
min fprane Ls(m= w) aa|HH ( w)ll2;
5

---

Published as a conference paper at ICLR 2023
where denotes the Khatri-Raa product, i.e., m [m W[1] [p and y| (m
w)ll2 denotes the Hessian norm w.r.t. the pruned model. In practice, optimize the Hessian norm by minimizing the trace of the Hessian which can be efficiently approximated; details can be found in Appendix E.2 We further conduct an empirical analysis show that pruning parameters
controlled by penalizing the trace of Hessian leads to decrease in the condition number; the result included in Figure 5f of Appendix D
Directly solving Eq. 6 to obtain m is difficult, and therefore we adopt a gating network to learn binary mask and perform structured pruning. The mask is defined as differentiable gating function: m g(dw; v), where dw and are the input and parameter of the function g: dw, dubbed as the parameter features, encodes the optimization dynamies characterized by the PL." condition. In the following two subsections, we explain the details of the parameter features dw and the gating function g. Algorithm summarizes the proposed pruning algorithm.
3.2.2 PARAMETER FEATURES
dwis (dwi) RR) designed encode two pieces of information about the optimization dynamics for each Wis. According to the discussion in Section 3.1 the PL condition at w can be analyzed through the minimum eigenvalue of the corresponding NTK. Therefore, the first parameter feature about W()) is chosen as Amin (MC((())))))))
The second parameter feature is empirically set as the entropy the eigenvalue distribution of the corresponding tangent kernel, which provides measure summarize the distribution envelope. Given the NTK matrix IC [((1g)) the entropy, denoted by p. defined as:
(K (wi))
Σ X (w(a)) log (N, (K (W[H)))
where A = \// Σ3 A is the normalized eigenvalue. is maximal when the eigenvalues are all equal and small when a single eigenvalue much larger than all others.
3.2.3 GATING NETWORK
The gating network is implemented by a two-layer feedforward network (FN) with parameters v, which outputs a scalar between 0 and 1 to indicate the importance of wi!
g(dw; v) softmax ([FN, (1PNuddw)) (dwn)): (N )]) [0, 1J",
(8)
The parameter v is updated according to the pre-defined pruning times and pruning ratios (t,, -- For example, under the one-shot pruning schedule, i.c., r), we update v once at epoch to prune r% of parameters; under the linear pruning schedule, i.e., {ty, /tt/////maaaa [maz]j=r" where tmaz denotes the maximum number training iterations and ry denotes the target pruning ratio, we update at every epoch prune / mmaz % of parameters (Hoefler al. 2021
Once the update v complete, we apply the TopK function to generate the binary mask m:
m = TopK (g(K() k), where m; = f if ((dw::)) is the top k elements gdw v), otherwise.
At pruning time t,. the hyperparameter & is set according to the pruning ratio: A (1 ry)p. By utilizing such a learnable gating function, the mask is capable of adaptively inducing sparsity on the network, till the pruned network produces good accuracy and fulfills the target pruning ratio. The final output m,, is used to produce the sparse over peee-prraeterizee model f(m,, * w).
4 RELATED WORK
Over-parameetrrized Model. Deep models used practice are often heavily over- parameterized (Zhang et al.) 2021 Nakkiran et al., 2021). Recent progress has been made in convergence analysis of over- parameetetee models optimized by (stochastic) gradient descent al-
gorithms (Du al. 2018 Liu et al. 2020b Oymak & Soltanolkotabi, 2019). In particular, the

---

Published as conference paper at ICLR 2023
PL condition (Pllyak 1963) and its slight variant the PL* condition (Liu al. 2022 have attracted interest in connection with the convergence guarantees. Some works prove that for over- parameterized models with the proper random initialization, gradient descent-based methods prov- ably find the global minimum with a global polynomial time convergence (Du et al., 2018 Belkin [202 Karimi al. 2016; [iu et. 2022). Furthermore, the work (Charles & Papailiopoulos 2018) shows the generalization bound of models satisfying the PL condition via stability hypothesis (Bous- quet & Elisseeff 2002). Inspired by the success of recent theoretical analysis, this work uses PL* condition to analyze the convergence and generalization ability, which is then used to guide the training of practical over-paramens wver-parameterized models.
Network Pruning. Network pruning aims to minimize network parameters while maintaining the model performance. Most existing structured pruning methods rely on approximations of loss differ- ences resulting from parameter perturbation. There are three classes of pruning techniques. The first uses zeroth- information such the magnitude of weight and activation values with respect to the running model, c.g., adding L1 or L2 penalty norm to the objective function (Han al. 5b a He et al. 2018 or using the Lottery Ticket Hypothesis and its extensions (Frankle Carbin 2018 Frankle al. 2019: Morcos al. 2019; You et al. 2020 Yu et al., 2020) Liu et (0011b) The second uses gradient- based first-order methods to approximate loss change via the first-order Tay-
lor expansion of the training loss (Liu & Wu 2019 You al. 2019 Mozer & Smolensky [988]
Molchanov et al., 2016 |anh et al. 20200) The third uses the Hessian or Fisher information of the
loss select weights for deletion (Hassibi et al. 1993 Hassibi 1992, Cun et al.) [9900] Molchanov et al. 2019; Liu et al. 2021a). Different from structured pruning works often accompanied with negative impacts on model generalization, the proposed pruning method based PL regularization simultaneously improves convergence and generalization ability.
EXPERIMENTS
This section presents experiments on three over-parameterized models, namely BERT (Devlin et al,, 2018), Twitch-Transforme (Fedus al. 2021), and VGG-16 (Simonyan & Zisserman 2014). De- tailed ablation experiments can be found in Appendix
Evaluation Measure. We evaluate the effectiveness of algorithm for model optimization in terms of training efficiency and generalization ability. Training efficiency is measured by the training aeneencylineeeerliiiiiiii error of the model trained with fixed training iterations (or epochs) (Shazeer al. 2017; Fedus et al. (02T). Generalization ability is measured by the test performance (Liu et al. 2021b). For
fair comparisons, all models are evaluated with the same pruning ratio, which is computed as the number of disabled parameters divided by the parameter counts of the corresponding structure.
5.1 RESULTS ON BERT
This experiment focuses on the Multi-Head Attention (MHA) component of BERT-Base (Devlin et al.] 2018)) MHA plays a critical role in the success of BERT. Following the basic setup of (Devlin et al. 2018), we train BERT with self-supervised masked language modeling (MLM) task (Devlin et al. (208)) on WikiText-2 (Merity et al. 2016). The baseline methods contain vanilla BERT (Wang et al., 2019). two crunsformer-aseee pruning methods, namely BERT-LTH (Chen 2020 and Att-Score (Michel et al.) (019), and two pruning approaches, namely SNIP (Lee 2018 and GraSP (Wang et al. 2020). More experimental details can be found in Ap-
ct al.
ct al.
pendix
E.3
Table 1: Results of BERT on WikiText-2. We denote perplexity as PPL, and APPL represents the
perplexity difference between training and test. Boldface indicates the best result.
Training PPI.
Test PPL
Method
Criterion
APPL
@ k iters
10 k iters
1 iters
Final
BERT BERT-LTH
39.37
14.79
6.35
75.57
7,49
Magnitude
36.97
13.07
7.92
81.21
7.20
Att-Score
Gradient
23.10
11.35
5.53
79.84
7.36
SNIP
Gradient
16.33
9.91
4.32
72.39
8.09
GraSP
Hessian
39.96
12.22
6.28
81.61
8.21
Ours
PL
15.72
9.63
4.24
63.18
6.84
7

---

Published as a conference paper at ICLR 2023
Table 1 shows the performance of all models on WikiText-2 with the {5, 75%) pruning policy, i.e,, we prune 75% heads at epoch-5 during training. Our method shows better training efficiency and method improves test perplexity
generalization ability than all baselines. In particular, vanilla BERT, and improves training perplexity by 60%, 29%, and 33% after 8 k, 10k, and 15k iterations, demonstrating the effectiveness of Pl. regularization for BERT model optimization. We also evaluate the training efficiency by measuring the number of iterations required for converging fixed training error status. Compared vanilla BERT, our method requires only 0.7x iterations to achieve satisfactory test perplexity (see Appendix F.2). Other pruning methods, on the other hand, can improve training efficiency but at a cost of sacrificing test performance, as their pruning criteria focus on training loss, ignoring the measurement of generalization ability.
Furthermore, we analyze the effect of pruning times of mask for BERT optimization (shown in Appendix F3 It is interesting to find that compared with the model applying PI. regularization iteratively throughout training, the model that applies PL. regularization during the early stages of training has better training efficiency and test performance. This phenomenon indicates that impor- tant heads for BERT optimization can potentially be discovered in the early stages of training.
5.2 RESULTS ON SWITCH- SWITCH-TRANSFORM.E TRANSFORMER
Switch-Transformee (Fedus al., 2021 is representative implementation of Mixture-of- Expert (MoE) architecture (Shazeer et al. 2017). which has shown excellent performance in NLP Switch-
tasks recently. For comparison, the baseline methods include vanilla LTH (Chen et al. 2020), Exp-Score (Michel et al. 2019), SNIP (Lee et al. 2018), .) and GraSP (Wang
al. 2020 All experiments are conducted WikiText-10S (Merity al. 2016 with the masked
language modeling task. More experimental details can be found in Appendix
B.3
Table Results of Switch-Transforen on WikiText- 103 (Best results in Boldface).
Training PPL
Test PPL
Method
APPL.
@ 50k iters
@ 60 iters
@70k iters
Final
Switch-Transforme
18.62
7.50
7.27
8.96
1.56
Switch-LTH
19.97
10.08
7.61
10.72
1.59
Exp-Score
18.45
10.56
8.17
9.18
1.49
SNIP
12.12
8.99
7.99
9.88
1.46
GraSP
18.93
8.17
7.81
8.84
1.52
Ours (experts only)
7.11
5.62
5.53
7.68
1.51
Ours (experts & heads)
6.97
5.86
5.23
7.78
1.45
Table 2 lists the performance of ing policy. Our method (Ours (experts only)) significantly outperforms all the baselines in both training efficiency and generalization ability, demonstrating our method can be applicable to MoE architecture. In particular, compared with the vanilla architecture, our method improves training perplexity by 62%, 41%, 24% in 50 k, 60 and k iterations, while improving performance by 13%. This phenomenon indicates that while the sparsely-activatee MoE architecture selects different experts for each data point, there also exist cer- tain moorly-behaved experts. By modeling the optimization dynamics which guarantee the conver- gence and generalization ability, the proposed PL regularization - pruning method can disable under- or ver-sppeialiied experts and keep well-behaved experts.
on WikiText-103 with the 50%) prun-
when used the MoE
We also investigate the benefits of implementing PL regularization on both heads of MHA and experts of MoE, as shown the last line in Table Compared with the vanilla Switch-Transfone Transformer, implementing PL regularization on both heads and experts presents a 28% improvement in training perplexity after 70k iterations and a 13% improvement in test performance. This result further demonstrates the proposed PI. regularization is flexible and applicable to different architectures.
5.3 RESULTS ON VGG-16
This experiment focuses on VGG-16 (Sionnyan & Zisserman 2014) trained on CIFAR-10 10 and CIFAR-100 100 datasets. The baselines include vanilla VGG-16 (Li et al.) 2016) three filter pruning methods, namely FPGM (He et al., 2019). CHIP (Sui et al. 2021). and DPFPS (Ruan et al. 2021) and three general pruning approaches, namely LTH (Frankle & Carbin 2018), SNIP (Lee et al.
8

---

Published as conference paper at ICLR 2023
2018), and GraSP (Wang al. 2020). All models trained from scratch and pruned with linear schedule where starting from epoch- 15, we prune the same number of filters at each epoch until the target sparsity is reached. More details of baselines and experimental settings can be found in Appendix E.4. Additional experimental results of ResNet-56 can be found Appendix F.
Table shows the performance of all methods trained on CIFAR-10 and CIFAR-I00 with 50% pruning ratio. As we can see, our method outperforms all pruning baselines in terms of training efficiency and test accuracy on both datasets. Especially, compared with the vanilla VGG-16, our method achieves a higher accuracy with fewer parameters, demonstrating the effectiveness of PL regularization implemented as structured pruning for VGG-16 optimization.
Table 3: Results of VGG-16 on CIFAR-I0 and CIFAR-100 (Best results in Boldface).
CIFAR-10
CTFAR- 100
Method & Dutaset
Train Loss (él epochs)
Test Accuracy epochs) Final
Train Loss (@ epochs)
Test Accuracy (@ epochs)
100
@ @55
@100
@ 150
100
@150
@100
@50
Final
VGG-16
0.1091
0,0075
92.46
93.25
93.85
0.4826
0.0459
72.21
72.28
73.86
LHT
0.1206
0.0106
92.41
93.12
93.59
0.5342
0.0507
71.78
71.86
73.11
SNIP
0.1106
0.0098
92.35
93.21
93.85
FPGM
0,1144
0.0146
90.80
91.70
93.50
0,4964
0,0709
68.00
68.
68.90
GraSP
0,2150
0.0420
91.13
92.83
93.70
0.8380
0.1240
69.33
71.74
72.92
CHIP
0,2585
0.0841
87.00
90.38
93.71
1,0581
0,4508
58.44
67.28
72.80
DPFPS
0.1224
0.0101
92.54
93.11
93.99
0,5222
0,0430
71.69
73.17
73.88
Ours
0.1040
0.0168
92.70
93.55
94.27
0.5175
0.0428
72.31
73.26
74.08
5.4 QUANTITATIVE ANALYSIS
This experiment aims to understand why the baselines perform differently in terms of train- ing efficiency and test performance. Theo- rem 6 states that the condition number con- trols the convergence and generalization ability of a model. Thus, we take a trained Switch- Transformer
LTH Exp-Score SNIP
GraSP
Ours
Layer Index
example and compare the condition number of important experts identi- fied by different pruning criteria, including the weight magnitude of LTH, the gradient- -based score of Exp-Scone Score, the gradient value of SNIP, and the Hessian value of GraSP.
Expert Index
Figure 2: Heatmaps of condition number of the remained experts determined by LTH, Exp- Score, SNIP, GraSP, and our method. Baselines fail to identify the well-behaved experts (The ex- perts marked in red are the ones that are well- behaved but missed by baselines).
Figure presents heatmaps of the condi- tion number and marks important experts by squares. shows that our method has the abil- ity to retain well behaved experts with small condition number while the baselines fail to retain these experts. The reason that these base- lines bias towards training loss, adversely affecting generalization. This observation confirms the importance of simultaneously considering both the convergence gap and the generalization gap for werrppaamettrizeed model optimization. Furthermore, the criteria of recent pruning works can be garded as the parameter features and fed into the proposed gating network, which may help optimize PL regularization by providing rich information of parameters. We leave it for future works.
CONCLUSION
This paper establishes principled connection between model optimization and the condition and advocates utilizing this theoretical finding when optimizing prramearmeter models. A new optimization problem is therefore formulated with the proposed PL regularization, and we illustrate that it can be solved via structured pruning method. We postulate that the effect of such regu- larization may be achieved in other ways, e.g., by using more advanced optimization algorithms which induce implicit regularization. Moreover, experiments demonstrate the superiority of the pro- posed PL regularization -iriiitttiio-riven pruning algorithm, unveiling the potential of adopting theoretically founded measures as the pruning criteria.

---

Published as a conference paper at ICLR 2023
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International Conference Machine Learning, pp. 242-252. PMLR, 2019.
Jafar Alzubi, Anand Nayyar, and Akshi Kumar. Machine learning from theory to algorithms: an overview. In Journal of physics: conference series, volume 1 1142, pp. 012012. IOP Publishing, 2018.
Haim Avron and Sivan Toledo. Randomized algorithms for estimating the trace of an implicit sym- metric positive semi-definite matrix. Journal of the ACM (JACM), 58(2):1-34, 2011.
Raef Bassily, Mikhail Belkin, and Siyuan Ma. On exponential convergence of SGD in non-convex over-parametrized learning. arXiv preprint arXi::181100256 2018.
Mikhail Belkin. Fit without fear: remarkable mathematical phenomena of deep learning through the prism of interpolation. Acta Numerica, 30:203-220 248, 2021.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learn- ing Research, 2:499-528 526, 2002.
Olivier Bousquet, Stéphane Boucheron, and Gábor Lugosi. Introduction to statistical learning the- ory. In Summer school on machine learning, pp. 169-207. Springer, 2003.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. SGD learns over- parameterized networks that provably generalize on linearly separable data. arXiv preprint 1'Xi::1710.10174 2017.
Zachary Charles and Dimitris Papailiopoulos. Stability and generalization of learning algorithms that converge to global optima. In International Conference on Machine Learning, pp. 745-754. PMLR, 2018.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained bert networks. Advances in Neural Information Processing Systems, 33:15834-15846, 2020.
Y. L. Cun, J. S. Denker, and S. A. Solla. Optimal brain damage. Advances in Neural Information Processing Systems 2, 1990.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representationss 2018.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint 1803000300003335, 2018.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the 2019.
lottery ticket hypothesis. arXiv preprint
Ingo Gühring, Mones Raslan, and Gitta Kutyniok. Expressivity of deep neural networks. arXiv preprint 20X1::2007704759, 2020.
Chirag Gupta, Sivaraman Balakrishnan, and Aaditya Ramdas. Path length bounds for gradient descent and flow. J. Mach. Learn. Res., 22(68):1166 2021.
10

---

Published as a conference paper at ICLR 2023
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv: 2015a.
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, 2015b.
B. Hassibi. Second order derivatives for network pruning: Optimal brain surgeon. In CiteSeer, pp- 164-171, 1992.
B. Hassibi, D. G. Stork, and G. Wolff. Optimal brain surgeon and general network pruning. In IEEE International Conference on Neural Networks, 1993.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1026-1034, 2015.
Yang He, Guoliang Kang, Xuanyi Dong, Yanwei Fu, and Yi Yang. Soft filter pruning for accelerating deep convolutional neural networks. arXiv preprint arXiv: 1808. 060660 2018.
Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi Yang. Filter pruning via geometric median for deep convolutional neural networks acceleration. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 4340-4349) 2019.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. J. Mach. Learn. Res., 22(241):1-124, 2021.
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewicz condition. In Joint European conference ma- chine learning and knowledge discovery in databases, pp. 795-811. Springer, 2016.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr Snip: Single-shot network pruning based on connection sensitivity. arXiv preprint
2018.
Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprint 160. 2016.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward theory of optimization for over- parameterized systems of non-linear equations: the lessons of deep learning. arXiv preprint arXiv:2003. 00307, 2020a.
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel constant. Advances Neural Information Processing Systems, 33: 15954 15964, 2020b.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85-116, 2022.
Congcong Liu and Huaming Wu. Channel pruning based on mean gradient for accelerating convo- lutional neural networks. Signal Processing, 156:84-91, 2019.
Liyang Liu, Shilong Zhang, Zhanghui Kuang, Aojun Zhou, Jing-Hao Xue, Xinjiang Wang, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Group fisher pruning for practical network compression. In International Conference Machine Learning, pp. 7021-7032. PMLR, 2021a.
Shiwei Liu, Lu Yin, Decebal Constantin Mocanu, and Mykola Pechenizkiy. Do actually need dense over parameterization? in-time pvrrpaaramttriization in sparse training. In International Conference Machine Learning, pp. 7009. PMLR, 2021b.
11

---

Published as a conference paper at ICLR 2023
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 11Xh::197..11922, 2019.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 1xXiv:6600007455 2016.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in Neural Information Processing Systems, 32, 2019.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint arXiv:1611.064400 2016.
Pavlo Molchanov, Arum Mallya, Stephen Tyree, luri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11264-1133 11272, 2019.
Ari Morcos, Haonan Yu, Michela Paganini, and Yuandong Tian. One ticket to win them all: general- izing lottery ticket initializations across datasets and optimizers. Advances in Neural Information Processing Systems, 32, 2019.
Michael C Mozer and Paul Smolensky. Skeletonization: A technique for trimming the fat from network via relevance assessment. Advances in Neural Information Processing Systems, 1, 1988.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanics: Theory and Experiment, 2021(12):124003, 2021.
Yurii Nesterov. Introductory lectures convex optimization: A basic course, volume 87, Springer Science & Business Media, 2003.
Samet Oymak and Mahdi Soltanolkotabi. Overparameterizee nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning, pp. 4951-4961
PMLR, 2019.
Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal Vychislitel noi Matematiki Matematicheskoi Fiziki, 3(4):643-653, 1963.
Xiaofeng Ruan, Yufan Liu, Bing Li, Chunfeng Yuan, and Weiming Hu. Dpfps: dynamic and pro- gressive filter pruning for compressing convolutional neural networks from scratch. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 2495-2503, 2021.
Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine- tuning. Advances in Neural Information Processing Systems,
2020.
Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9481-9488, 2021.
Kevin Scaman, Cédric Malherbe, and Ludovic Dos Santos. Convergence rates of non-convex stochastic gradient descent under a generic lojasiewicz condition and local smoothness. In In- ternational Conference on Machine Learning, pp. 19310-19327. PMLR, 2022.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outragcously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 7707006533 2017.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.15555 2014.
Yang Sui, Miao Yin, Yi Xie, Huy Phan, Saman Aliari Zonouz, and Bo Yuan. Chip: Channel independence- based pruning for compact neural networks. Advances in Neural Information Pro- cessing Systems, 34:24604-24611 2021.
12

---

Published as a conference paper at ICLR 2023
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of SGD for over- parameterized models and an accelerated perceptron. International Conference on Artificial Intelligence and Statistics, pp. 1195-1204. PMLR, 2019.
Chaoqi Wang, Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations, 2020.
Chenguang Wang, Mu Li, and Alexander J Smola. Language models with transformers. arXiv preprint arXiv: 09400, 2019.
Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Richard G. Baraniuk, Zhangyang Wang, and Yingyan Lin. Drawing early-bird tickets: Toward more efficient training of deep networks. In International Conference on Learning Representations, 2020.
Zhonghui You, Kun Yan, Jinmian Ye, Meng Ma, and Ping Wang. Gate decorator: Global filter prun- ing method for accelerating deep convolutional neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Haonan Yu, Sergey Edunov, Yuandong Tian, and Ari S Morcos. Playing the lottery with rewards and multiple languages: lottery tickets rl and nlp. In International Conference on Learning Representations, 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016. British Machine Vision Association, 2016.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
13

---

Published as a conference paper at ICLR 2023
A BAYES ERROR DECOMPOSITION
Motivated by the previous work (Gühring et al. 2020), the Bayes error, i.e., the optimal test error, can be decomposed as:
R(f)) Rs(SH.9) (R((H,,s) Rs(SH.)) R(fH.s) Rs Rs(H..))
empirical error
convergence gap (R(R) R(fH, -Rf())--(
gencralization gap
(10)
estimation emvor
approximation
In the over- pve-paarameterized setting, we assume that the optimal function fi.s fits or interpo- lates the training data exactly, i.c., Rg(sR.s) vergence gap. Moreover, the estimation error can be associated with the generalization gap: [R(f Rs(fi,s)) I SEH [R(f) Rs(f)l (Bousquet al. 2003). Furthermore, as over-parameeterizatoon refers to manifolds of potential interpolating predictors, H for the over- parameterized model can omitted, i.e., R(SH)) R((*). Therefore, for over-parameterrized models, Eq. 10 can be simplified as
Thus the empirical error equals the con-
can
R(f*) S 2 RssHs) Ro(fH.s)
+3 sup R(f) Rs(f) FEH
(11)
generalization gap
convergence gap
B PROOF OF THEOREM 6
B. CONVERGENCE RATE
The convergence rate of the pver-parameterizzd model has been studied by many works (Liu al.,
2022 Karimi et al. 2016 Allen-Zhu et al. 2011)).
We assume that the loss function has the Li-Lschhittz continuous gradient, i.e., for all parameters w and v, such that:
((w) S L(v) (VL(v). w w) L |w v112
(12)
Using the assumption, we have that:
Ls (W++1) - Ls (w.) < (VLs(w:), We+1 - w.) + Li
| will
(13)
2
Given the gradient method with learning rate n 1. Ly" L, VLs)w)), we have
such that W+41 w. nVLss(W.
[s(We+1) Ls(w)) 2L,
/VLs(w:)]l
(14)
Under the PI.* condition at point wt, we have
Ls (We+1) - Lw()) S - (Cs(w.))
H Ls
(15)
Therefore, have
[s(w)) 51
T.g (Cs( 1)
# ) ((Ls(wo))
11
(16)
T,f
B.2 SAMPLE COMPLEXITY
Following previous work (Charles & Papailiopoulos 2018), we first calculate the generalization bound e for the function f with the pointwise hypothesis stability y. Then, we analyze the point- wise hypothesis stability y for model trained on the loss function satisfying the PL condition with parameter μ. Finally, given the pointwise hypothesis stability 7, and the generalization bound e, we bound the sample complexity of f.
14

---

Published as a conference paper at ICLR 2023
B.2.1
GENERALIZATION ANALYSIS VIA POINTWISE HYPOTHESIS STABILITY
useful approach analyzing the generalization performance of learning algorithms is algorithmic stability (Bousquet & Elisseeff, 2002). A learning algorithm A is stable if small changes in the training set result in small differences in the output predictions of the trained model. Following the foundational works, we analyze the generalization bound from the perspective of pointwise hypothesis stability.
X where X; D, we define dataset St as S)x;, i.e., S x,). For our purposes, f is the output of learning algorithm A, parame- terized by w. For the model parameterized by we denote the model trained on the dataset S and 5 as Wg and Wgs, respectively. Accordingly, the loss functions are represented as [(ws; x) and st x), respectively. Especially, for a given dataset S, the empirical loss dataset S is denoted as (wss) and given by Ls(ws) = Ei= (x; W:).
Given data set S = [xi..... X,)) {xI,...,
Definition (Pointwise Hypothexis Stability (Bousquei Elisseef 2002)). An algorithm A has pointwise hypothesis stability y with respect loss function L if
Vi E n., Es || (ws: L (wssi))) E n.
(17)
Then we can use the pointwise hypothesis stability to establish the generalization bounds as
Theorem 11 Bousquet & Elisseeff (002). Suppose we have a learning algorithm AA with pointwise hypothesis stability Y with respect bounded loss function C such that 0 L(w;x) S M. For any 8, we have with probability at least 8,
M + 12MRy
R(ws) Rs(ws) +
(18)
2nd
In the following, we derive stability bounds for models trained on risk functions satisfying the PI.* condition.
B.2.2 POINTWISE HYPOTHESIS STABILITY WITH PL CONDITION
Theorem 12 Assume that for all training sets S and models f parameterized by w. the loss func- tion L is PL with parameter p. In addition, assume that the over-parameterized model I with parameter w trained S is capable of converging some global minimizer w's. Then for all S, if Ls (W3) Ls (w3)l S ej, then f has pointwise hypothesis stability with parameter as:
L
(19)
pm
Proof. For a training set S, let Ws denote the parameters of f on S, and let Ws denote the pa- rameters of If on S¹ where S denotes SY,2, Let w and w's. denote the corresponding optimal solutions, respectively. We then have,
[C (Wsix:) - (w x;, < L (w); L (w3: x,)l + | (,): L (w3.:x))++ (W3.:)) - L (wsx;x)). (20)
The three terms be separately bounded.
The first and third term
As discussed in previous work (Karimi et al. 2016). the PL condition implies that the quadratic growth (QG) condition
" wws w3ll Ls (ws) Cs (w3)l-
(21)
2
Using the fact that Ls (ws) L (w3)) S of by the assumption, it implies
||wg w3ll S V2 / LLs (ws) Ls (w3)) = 2
(22)
15

---

Published as a conference paper at ICLR 2023
Since that loss function L is -ipisschh can bound the first and third term as:
EL CC (Ws::x), L (w::: x,) S Ly | wz.ll S L
C (ws; - L (wi:: x;) - L ||ws - will - 1
(23)
In the over-parameterized settings, the model f optimized by the loss function enables global optimizer w" and has small ef constant, thus the first term and third term equal to O.
The second term. PL* condition implies that L (w3:x.) = O, and Css (w's) = 0 thus the second term can be manipulated as:
[ (w): x.)) = |C xx,)
|nCs (W::)
(n 1)L5i (w':))
(24)
= Ls (w5.))
Note that since Vgg (wg:) we get:
L2 n12
CV (ws))ll = n²
V (w x
(25)
Furthermore, PL* condition implies PCs (w's) S VLs (ws:) 112
we can obtain
L
D3 pm
n Lcs (w's S: )
V L
||
(26)
The overall bound. Plugging Eqs. 23 and 26 into Eq. 20 that for all S and w, the loss function L is PL with parameter μl, algorithm A has pointwise hypothesis stability with parameter 7 as:
obtain the desired result: Assume
L pm
(27)
B.2.3 BOUND OF SAMPLE COMPLEXITY
Based the above subsection, if the loss function L of model f is PL" condition with parameter 14, and the algorithm AA has the pointwise hypothesis stability with parameter y tion gap of the obtained model can be bound as: the sample complexity as:
1.3 the generaliza- junn
M22 12MM1
Therefore, compute
6LI
M
n(c,5) s
+
(28)
2
2c28
C PROOF OF PROPOSITION g
Proof. Consider an arbitrary point w E B(Wo, R). For all sample in dataset S, we have
DF.(w) = DFF (wo) + [ НG, (wo + - (w ww)) (w -w) dr.
(29)
Since E [0, 1]. the point wo + 7 7(w wo) inside of the ball, thus || (wo + +(w (w- wo))lla
[w)]2) Therefore, Eq. 29 can be bounded as:
I[DF((w) DJ.
(FFF(W)))(((()))) DF < sup ||H (wo (w wo)) -lww-lllool -woll rE(0.1]
(30)
< R-
16

---

Published as a conference paper at ICLR 2023
By triangle inequality, we have.
[DF(w))lz IIDDF(wo)ll - IDF(w) DF(ww)ll2
- (DF(w) DF DDF(w))lF
Σ (DF; DF: (wo)ll
(31)
S VnR (|
By re-arranging Eq. 31 we have
[DF(w)llz |DF( + VnR 1 [1)(())lz
(32)
Using the definition of L, we have
L sup wEB
[DF S IDF(wo)ll VoR sup [
(33)
wEB
Combining Eq. 33 and definition of μ², we have
T.f I
[DF((o)] suPweB Hy)(w)] R
(34)
inf Amin (K(w))
D PROOF OF THEOREM 9
To prove that minimizing the upper bound of condition number can be achieved by pruning param- eters controlled by Hessian norm H ||2, we separately analyze the impact of pruning parameters the Lipschitz constant; and the impact of pruning parameters on the minimal eigenvalue of NTK (Proposition 14).
controlled by Hessian norm H y Il2
Proof. Using the same argument in Proposition g shown in Eq. 33 where
L sup [|FF(w)]lz < |DDF(wo)la VnR sup || [[)]w))l, WEB
wEB
can conclude that pruning parameters reduces the parameter size, leading to the decrease of the |[ - wwll (i.e., the radius of the ball, R); and pruning parameters governed by minimizing the Hessian norm || llz can decrease the upper bound of the Lipschitz constant. In summary, the pruning operation with controlling of Hessian norm is able to minimize the upper bound of Ly.
matrix analysis to show that maximizing the minimal eigenvalue of NTK can be achieved via pruning.
Then
Let us denote the matrix G DFDFT. The matrix and NTK matrix (K = DF DF) shares the same non-zert cigenvalues, thus analyze the effect of pruning the eigenvalues of NTK via the matrix G E mm where m denotes the parameter size.
Pruning any parameter WP from the weight and the corresponding p-th column from G, i.e Gp. According to the Cauchy Interlacing Theorem, increasing the minimal eigenvalue of G can be achieved by pruning the parameter Wm- We first summarize the Cauchy Interlacing Theorem as follows:
corresponds to deleting the corresponding p-th row
Theorem 13 (Cauchy Interlacing Theorem Horn & Johnson 2012 ). Let B Mm be Hermitian,
let y E C. and E R. be given, and let A y" a
E MM Then
(A) (B) S (² S Am (A) Am (B) s Am+1(A).
(35)
Theorem 13 allows us conclude that if remove the parameter W.n the corresponding G... G. then Xmii Gm) 2 Amin (G).
17

---

Published as a conference paper at ICLR 2023
Proposition 14 For pver-parameterizze model parameterized by W, pruning operation on the parameter set w increases the minimal eigenvalue of NTK matrix associated to f.
Proof. For parameter Wp and matrix G, G denote the permuted matrix obtained by permut- ing p-th row and p-th column of matrix G to m-th row and m-th column. Then pruning parameter Wp is equivalent to pruning Wm of G matrix.
Following Theorem 13 pruning parameter Wmm of G will increase its minimal eigenvalue. Moreover, using the fact that G PGP PGP-1 where P denotes the permutation matrix, have Gr and G are similar matrix with same eigenvalues. Thus, the minimal eigen- value G is increased. In summary, pruning any parameter Wp will increase the minimal eigenvalue of G.
Proof of Theorem 9 Proposition shows that pruning operation controlled by minimizing Hessian norm 1|2 is able to minimize the upper bound of Lipschitz constant Ly; and Proposition 14 shows that pruning parameter is able to maximize the minimal eigenvalue of NTK. Combining these two propositions, we can conclude that pruning parameters controlled by minimizing the Hessian norm H ll2 is able to minimize the condition number of the model.
To further verify the claim of theories, we empirically investigate the evolution of condition num- ber of the pruned model and vanilla model. We apply small BERT with Transformer layers to WikiText-2 for the language modeling task. We prune the heads of vanilla BERT by using a strategy of pruning 5% heads every 10 epochs, controlled by minimizing the Hessian trace (details can be found in Appendix E.2p: the pruned network is denoted as Masked-BEIK compared with vanilla BERT, Masked- cially, we can observe a sharp decrease in the condition number of Masked- BERT after each pruning operation; highlight this, red lines are included in the figure which indicate the condition number before and after pruning.
As shown in Figure 3 presents smaller value of condition number Espe-
x107
Mesked
6.00
BERT
575
Condition Number
MMMM
5.50
M
M
5.25
5.00
4.75
M
4.50
4.2
4.00
Training Epoch
Figure 3: The evolution of condition number of the vanilla BERT and the pruned BERT (denoted as Masked-Bee -ERR)..
E EXPERIMENTAL DETAILS
E. ALGORITIIM OVERVIEW
Figure shows the overview the proposed method. We impose PL regularization for model opti- mization by adopting pruning approach. More specifically, we introduce binary mask for period- ically sparsifying parameters, and the mask is learned via a gating network whose input summarizes
Connition number computed as Amaz where Amaz and Amin denote the largest and smallest eigen- value, respectively.
Amiin (
18

---

Published as a conference paper at ICLR 2023
the optimization dynamics of sub- aub-networks in terms of the PL condition. Algorithm 2 summarizes the proposed pruning algorithm.
Nodes
Model Nodes (Prixed)
(Dver Param.)
Pruning
d
Gating Network
mailer
Ey
Reguiarization
Figure 4: Algorithm overview. The gating network generates binary mask based the parame-
ter features dw. PL Regularization helps pruned model obtain smaller condition number Lx
E.2 IMPLEMENTATION OF LEARNING OBJECTIVE
Based on the Theorem minimizing the upper bound of condition number can be achieved by pruning parameters controlled by minimizing the Hessian norm.
In practice we minimize the Hessian norm by minimizing the trace of the corresponding Hessian where T'aceeHH, (w)) Σx (H,)). Consequently, the learning objective of the gating network is approximated as:
min Ls (m w) + rTacee (H, (g (dw; v) w)).
(36)
w.v
Specifically, we use the Hutchinson method (Sankar et al. 2021 Avron & Toledo 2011 a tool from randomized numerical linear, to efficiently compute the trace of Hessian without explicitly forming the Hessian matrix. The details are summarized in Algorithm
To further verify pruning parameters controlled by minimizing Hessian trace is capable to decrease the condition number of the model, we empirically investigate the evolution of the trace of Hessian and the condition number. Here a small BERT with 4 Transformer layers is applied to WikiText-2 for language modeling. As shown in Figure 5 Masked-BERT where prune 50% heads at epoch- 10 controlled by minimizing the trace of Hessian presents small value of Hessian trace, accompa- nied with smaller value of condition number. The phenomenon demonstrates that minimizing the Hessian trace an effective way to control the condition number of the model.
Condition Number
V My wn pno Me w h tran yo in Mon
ma And Mall
Loss (Log ppn
Hessian Traze
180
Training Epoch
Training Epoch
Training Egoch
Figure 5: Connection between the trace of Hessian Trace(H, ( w)) and the condition number
L l.pp.
19

---

Published as a conference paper at ICLR 2023
Algorithm Hessian Trace Computation.
Input: Model Parameter: w, Model logits: Number of iteration:
Output: Hessian trace of f: Trace
Trace 0
2: for = .., do
3: Sample vector a from a standard normal: a NN I)
4: Compute differential map with respect to w: DF 8h
5: Compute trace: Trace = Trace a 0(DFTa) a)
6: end for
7:
return Trace/n
Algorithm 2 PL Regularization Driven Pruning Algorithm.
Input: Training set S, pruning time and ratio (4i,r) maximum iterations bmaa
1: Initialize:
2: while not converged do
3: for to...., Emaz do
4: if 11=t, 1 then
5: Calculate d (Amin (K(( ((( for each batch in do
6:
7:
V - V with
Update v according to Eq. 36
8:
end for
9:
m
-Topk(y v), (Irr))p) w m * w
Generate binary mask m Reparameterize w
10:
11:
end
12:
for each batch in S do
13: w = w - V Lrotal with a ==0
Update w according to Eq. 36
14:
end for
15: end for
16:
end while
E.3 EXPERIMENTAL DETAILS OF BERT AND SWITCI-TRRNNFFRMEER
BERT and Switch Transformer are trained from scratch by self-supervised Masked Language Model (MLM) task. The MLM objective a cross-entropy loss on predicting the masked tokens. The models uniformly select 15% of the input tokens for possible replacement.
Considering the model capacity and computational resources, we train BERT on WikiText-2 and Switch Transformer Wik datasets. The training hyperparameters
presented in Ta- ble 4 All baseline models, including our method, are trained by the same training setup. These experiments are conducted
GPUs of NVIDIA GeForce RTX 3090.
We follow the original basic settings when training baseline models. For BERT-LTH and Switch- LTH, iterative magnitude pruning (IMP) (Chen et al. 2020) target sparsity with rewinding step of 5% maximum training epochs. For Att-Score and Exp-Score, all the attention heads and experts across the model are sorted by gradient-based proxy importance score (Michel al. 2019). Then they are pruned by the same itcrative strategy as IMP. For SNIP and GraSP, we implement them in structured way, i.e., we sum up the relative indicators in those components, and the attention heads and experts are pruned at the initialization stage of the training process (Lee et al.) [01 8 Wang jet al. 2020). A protecting strategy is used the pruning operation of all baselines and our model. We keep at least one head or expert in each of the layers to keep the activation across the model.
E.4 EXPERIMENTAL DETAILS OF VGG-16
First, summarize the architecture-speciff pruning baselines. including FPGM, CHIP, and DPFPS. FPGM (He et al. 2019) compresses CNN models by pruming filters with norm- and distance-based criteria, which prunes all the weighted layers with the same pruning rate at the same
20

---

Published as a conference paper at ICLR 2023
Table 4: Hyperparameters for training BERT and Switch Transformer
Hyperparameter
BERT
Switch Transformer
Number of layers
12
12
Hidden size
768
768
Attention heads
12
12
Dropout
0.1
0.1
Sequence length
512
512
Batch size
8
8
Warmup steps
0%
6%
Weight decay
0
le-2
Peak learning rate
le-4
2e-4
Learning rate decay
Linear
Cosine
Adam [e, 91, 82]
[0, O, 0J
[le-6, 0.9, 0.999]
Number of experts
4
Capacity factor
1.5
time. CHIP (Sui et al. 2021) first extracts feature maps from a pre-trained model and calculates channel independence, which is used to sort and prune the filters, and then it fine-tunes the sparse pre-trained model. DPFPS (Ruan et al., 202T) prunes structured parameters in a dynamic sparsity manner, where the sparsity allocation ratios are distributed differently over layers in the training process.
Then, report the detailed experimental settings. For training VGG-16 on CIFAR-IO and CIFAR- 100, we use similar configurations as (Wang et al.) 2020) does. For all baselines, the network is trained with Kaiming initialization (He al. 2015) using SGD for 200 epochs. The learning rate is decayed from the initial value 0.1 by a factor of 0.1 at 1/2 and 3/4 of the total number of epochs. To obtain more stable results, we conduct each experiment in trials.
In this paper, we evaluate the performance of all models on VGG-16 with 50% and 90% pruning ratios, respectively, The pruning policy of our method is the linear pruning schedule where starting from the epoch-15, we prune the same number of filters at each epoch until the target pruning ratio reached. The pruning policy baselines follows their origin settings given target pruning ratio. In detail, for LTH, SNIP and GraSP, conduct unstructured pruning with the target pruning ratio at the initialization stage. For FPGM, we utilize the norm- based criterion when given the 50% target pruning ratio. When given the 90% target pruning ratio, we first prune 50% filters by the norm-based criteria and then prune 40% filters by the distance- based criteria. For CHIP, given the 50% pruning ratio, following the original work, we assign the pruning ratio of the previous 7 layers 20% and the ratio of remaining layers as 60%; while given the 90% pruning ratio, the ratio of the previous layers is assigned as 75% and remaining layer as 95%.
F ADDITIONAL EXPERIMENTAL RESULTS
In this section, we report additional experimental results.
F.1 RESULTS ON RESNET-56
This experiment focuses on ResNet-56 (Simonyan & Zisserman) 2014) trained on the CIFAR- 10 dataset. The baselines include five filter pruning methods, namely FPGM (He al. 2019). DPFPS (Ruan et al. 2021). LTH (Frankle & Carbin 2018). SNIP (Lee et al. 2018). and GraSP (Wang et al.) 2020). All models are trained from scratch and pruned with linear sched- ule, which starts from epoch-5 and prune the same number filters per epoch until the target sparsity reached.
21

---

Published as a conference paper at ICLR 2023
Table 5 shows the performance of all methods trained on CIFAR-10 with 25% pruning ratio. As we can see, our method outperforms all pruning baselines in terms of training efficiency and test accuracy.
Table 5: Results of ResNet-56 on CIFAR-I0 (Best results in Boldface).
CIFAR-10
Method
Train Loss (@ epochs)
Test Accuracy (@ epochs)
@ 100
@ 150
@ @000
@ 150
Final
ResNet-56
0.1441
0.0148
91.76
92.80
93.43
LHT
0.1395
0.0178
92.18
92.98
93.50
SNIP
0.1509
0.0196
92.00
92.98
93.41
FPGM
0.1931
0.0689
91.27
92.60
93.41
GraSP
0.1453
0.0187
92.27
93.20
93.39
DPFPS
0.1507
0.0218
91.50
92.16
92.75
Ours
0.1380
0.0158
92.46
93.46
93.78
F.2 TRAINING EFFICIENCY
This subsection shows the comparison of training efficiency by measuring the number of training iterations required for converging to given test perplexity. We also evaluate the wall-clock time saved when the model converges to given test perplexity. The calculation of parameter features and the training of the gating network are included in the statistic.
As shown in Figure 6 compared to vanilla BERT trained WikiText-2, model achieves the same test perplexity as the vanilla BERT at iterations 10k at iterations 8k, which a 1.3 speedup in terms of step time, i.e., method only requires 0.7 iterations to achieve the same test perplexity. As for the saved wall-clock time, as shown in Table 6 where the first row shows the wall-clock time in seconds and the second shows the time saved compared with the vanilla BERT, We can observe that our method saves more wall-clock times compared with other baselines. In particular, saves 29% wall-clock time of vanilla BERT. Furthermore, the training of the gating network accounts for only 2-3% of the total training time of BERT, demonstrating the efficiency of our method with the loop for pruning.
Table 6:
Comparison of saved clock time on BERT.
Methods
BERT
BERT-LTH
Att-Score
SNIP
GraSP
Ours
Wall-Clock Time (s)
4112
4534
6683
3370
3563
2916
Time Saved
0%
-10%
-23%
18%
13%
29%
We find similar results in Switch-Transformen trained WikiTexx-1033. As shown in Figure 7 and Table 7 compared to vanilla Switch-Transformer Transformer, our method yields 2x speedup, requiring 44% less clock time.
Table 7: Comparison of clock time saved when model achieves given test perplexity. The symbol denotes expert pruning, and ** denotes expert and head pruning.
Methods
BERT
Switch-LTH
Exp-Score
SNIP
GraSP
Ours*
Ours "*
Time Saved
0%
-9%
19%
8%
32%
44%
(Approximately)
22

---

Published as a conference paper at ICLR 2023
BERT Wikitext-2 (Train Set)
BERT on Wikitext-2 (Test Set)
(beseline) Ours (heads pruned)
(beseline)
Ours (heads prunod)
Train Loss (L.g ppt
6.5
6.0
Test Loss
5.5
4-
5.0
1.3x Speedup
1.7x Speedup
4.5
0.2
1.0
Training Iterations
Training Iterations
Figure 6: Training efficiency comparison of vanilla BERT and our model.
Switch Transformer Wikitext- 103 (Train Set)
Switch Transformer Wikitext- 103 (Test Set)
Switch Transformer (bsseline)
Switch Transformer (beseline)
(experts pruned only)
(experts pruned only)
Ours (exports heads pruned)
Ours (experts heads pruned)
Train Loss (Log opt)
Test Less (Log ppl)
7-
4-
Speedup
Speedup
2
2
7
Training Iterations
Training Iterations
Figure 7: Training efficiency comparison of vanilla
and our model.
F.3 THE EFFECT OF PRUNING TIME
This subsection studies the influence of pruning time on the performance of our method on BERT. For fixed target sparsity 75%, compare the one-shot pruning schedule on epoch-O (initializa- tion), epoch-5, and epoch-50.
Table Performance comparison on WikiText-2, For simplicity, we denote perplexity as PPL, and APPL represents the perplexity difference between training and test. Boldface indicates the best result.
Training PPL.
Method
Test PPL
APPL
@8 k iters
@ 1 iters
@ 15 iters
[0]
13.26
10.43
3.90
67.49
6.27
5
15.72
9.63
4.24
63.18
6.84
[50]
37.86
15.4
5.63
87.18
7.19
As shown in Table 8 empirically observe that it is sufficient to implement PL regularization BERT the early stages training, i.e., important heads for BERT optimization can be discovered in the first 5 epochs.
We also evaluate the impact of one-shot pruming and linear pruning schedule for BERT. Specifically, compare the performance of one-shot pruning and iterative pruning strategy the heads of
23

---

Published as a conference paper at ICLR 2023
Table 9: Comparsion of one-shot pruning and linear pruning schedule for BERT.
Training PPL.
Method
Test PPL
APPL
@8 iters
@ 10k iters
@15k iters
Bertllaall
39.37
13.49
6.35
75.57
0%
One-Shot
15.72
9.63
4.24
63.18
29%
Iterative 14.01
10.43
4.77
60.34
23%
BERT with the same 75% pruning ratio. As shown in Table the one-shot pruning strategy achieves competitive performance compared to the linear pruning schedule while saving more computational costs.
F.4 THE EFFECT OF PRUNING RATIO
This subsection evaluates the performance of the proposed PL regulariization-driven structured prun- ing method with varying pruning ratios.
Table TO and Table ITs the performance of BERT with the 50% and 90% pruning ratio, respec- tively.
Table 10: Performance Comparison of BERT Wikitext-2. All models are pruned by 50% sparsity.
Boldface indicates the best result among pruned models.
Training PPL
Method @8k k iters
Test PPL
APPL
@ 10 iters
@ 15 k iters
BERT
39.37
13.49
6.35
75.57
7.49
BERT-LTH
21.98
12.48
4.41
74.44
8.48
Att-Score
19.09
12.81
4.75
73.55
7.28
SNIP
21.61
11.09
5.19
74.89
6.70
GraSP 19.47
10.82
4.34
78.34
8.71
Ours
13.82
7.82
3.51
67.49
6.49
Table 11: Performance Comparison of BERT Wikitext-2. All models are pruned by 90% sparsity. Boldface indicates the best result among pruned models.
Training PPL
Method
Test PPI.
APPL.
@8k iters
@ 10k iters
@ @15k iters
BERT
39.37
13.49
6.35
75.57
7.49
BERT-LTH
75.64
14.86
6.32
81.78
6.74
Att-Score
33.05
17.08
8.16
84.61
5.88
SNIP
229.06
122.00
14.50
139.49
9.85
GraSP
91.38
22.44
8.32
94.25
8.00
Ours
24.24
11.27
5.46
72.68
5.45
Table 12 and Table 13 show the performance of VGG-16 16 with the 75% and 90% pruning ratio, spectively. We can observe that with a high pruning ratio, our method outperforms pruning basclines in terms of training efficiency and test performance.
24

---

Published as a conference paper at ICLR 2023
Table 12:
Performance Comparison of VGG-16 on CIFAR-10 and CIFAR-100. All models
pruned by 75% sparsity. Boldface indicates the best result among pruned models. Models without results mean that these models cannot converge under the current setting,
Method Dataset
VGG-16
LHT
SNIP
FPGM
GraSP
CHIP
DPFPS
Ours
Train Loss @ 100 epochs
0.1090
0.1516
0.1046
0.2359
0.2220
0.3091
0.1037
0.1098
CIFAR-10
Test Accuracy @ 100 epochs
92.46
91.53
92.41
89.00
91.03
86.91
92.41
92.63
Train Loss @ 150 epochs
0.0075
0.0171
0.0095
0.0833
0.0470
0.1272
0.0090
0.0082
Test Accuracy @ 150 epochs
93.25
92.52
92.90
89.20
92.52
89.92
93.26
93.69
Test Accuracy
93.10
93.62
89.60
93.23
93.93
93.96
Train Loss 100 epochs
0.4826
0.6250
1.108
1,2402
0.6399
0.6378
CIFAR- 100
Test Accuracy @ 100 epochs
72.21
71.52
64.38
64.57
70.96
71.00
Train Loss @ 150 epochs
0.0459
0.0762
0.2601
0.6711
0.0755
0,0717
Test Accuracy @ 150 epochs
72.28
71.68
71.14
63.84
71.64
71.80
Test Accuracy
73.86
72.66
72.27
69.83
72.77
72.91
Table 13: Performance Comparison of VGG-16 CIFAR-10 and CIFAR-100. All models pruned by 90% sparsity. Boldface indicates the best result among pruned models. Models without results mean that these models cannot converge under the current setting.
Method Dataset
VGG-16
LHT
SNIP
FPGM
GraSP
CHIP
DPFPS
Ours
Train Loss @ 100 epochs
0.1091
0.1867
0.1301
0.2530
0.3510
0.1417
0,1280
CIFAR-10
Test Accuracy @ 100 epochs
92.46
91.38
92.14
90.39
86.70
92.15
92.49
Train Loss @ 150 epochs
0.0075
0.0302
0.0132
0.0600
0.1871
0.0134
0.0122
Test Accuracy @ 150 epochs
93.25
91.83
93.07
92.40
89.02
92.97
93.11
Test Accuracy
93.85
93.19
93.56
93.18
93.18
93.59
93.61
Train
epochs
0.4826
0.8089
0.9990
1.7605
0.8889
CIFAR-100
Test Accuracy @ 100 epochs
72.21
68.24
67.09
38.66
68.27
68.70
Truin Loss @ 150 epochs
0.0459
0.1632
0.2460
1,2694
0,2050
0.2072
Test Accuracy @ 150 epochs
72.28
69.73
69.90
$2.63
70.31
70.85
Test Accuracy
73.86
70.60
71.16
62.90
71.24
71.36
F.5
FURTHER STUDY
Artificial deep learning models typically over-parameterized, coming at the heavy computation effort during model training and inference. This work, drawing on theoretical insights, finds an efficient optimization strategy for the over-parameterized model to achieve faster convergence time and a better test performance.
In practice, in order find the optimal model, researchers usually train multiple models with varying parameter initializations. In over-parameterized settings, such a process requires too many computa- tional resources and time efforts. Our method can only reduce the computational cost for training one model, but also reduce the sensitivity to initialization and thus less models need to be trained.
To verify the latter benefit, we conduct experiments on two BERT models with different parame- ter initializations: BERT-A and BERT-B denote the BERT model with good initialization and bad initialization, respectively. Table 14 and Figure 8 show that imposing our method on the BERT-B can significantly improve its original performance: it outperforms the vanilla BERT-A and slightly underperforms the masked BERT-A model implemented with PI. regularization. This phenomenon suggests the potential of our method in avoiding multiple initializations and meanwhi whine saving com- putation and time efforts.
25

---

Published as a conference paper at ICLR 2023
BERTa Masked BERTA
7.5
BERTa
Masked BERTA
BERT,
BERT,
Train Loss (ing ppd
Test Loss ppl)
Masked BERTs
6.5
Masked BERTs
6.0
5.5
5.0
2-
4.5
1
4.0
120
150
0
40
80
100
120
140
Training Epoch
Training Epoch
Figure 8: Training error and testing performance of BERT models with different parameter initial- ization.
Table 14: Results of BERT with different parameter initialization.
Training PPL 10 iters
Test PPL
Wall-Clock Time Saved
Method
@8 iters
@ iters
Final
BERT-A
31.91
11.16
4.55
74.44
Masked BERT-A
19.01
7.78
3.89
57.97
29.2%
BERT-B
59.44
13.37
5.09
83.93
Masked BERT-B
19.65
8.01
4.03
63.43
33.6%
26


Reviews:
[{"review_id": "Z0ZJTIgMCsY", "reviewer": "Reviewer_a5MJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new method for pruning overparameterized neural networks. The method is based on the intuition that the smaller condition number leads to better optimization and generalization. That is, the proposed method is designed so that the pruned neural network has a small condition number. The experiments verify the effectiveness of the method on training large-scale networks such as BERT, Switch-Transformer, and VGG-16.", "summary_of_the_review": "Although the empirical performance of the proposed method is good, there are several concerns about the correctness, quality, and clarity. ", "strength_and_weaknesses": "**Strengths**:\n\nThe idea of pruning the network to reduce the condition number is original and interesting. Moreover, the empirical performance seems good.\n\n**Weaknesses**:\n\n- My main concern is the correctness of the theory. Specifically, there is a conflict in assumptions. Specifically, PL-condition implies the quadratic growth condition as shown in [Karimi, et al, (2016)]. Hence, the objective function should be unbounded which contradicts the boundedness and Lipschitz continuity of the objective (assumed in Theorem 6). Some techniques to guarantee the boundedness of the parameters are probably needed.\n- The notation is inconsistent and confusing. Both notations $L(w)$ and $L_S(w)$ are used for empirical risk, although the authors say \"*we denote the loss function by L(x,y;w), or simply L(w), and also use L_s(w) to denote the empirical risk*\". For instance, the notation $L(w)$ in Definitions 4 and 5 and Theorem 6 should be the empirical risk. There is also the description \"*We denote $L$ and $L_S$ as the expected risk and empirical risk, respectively*\" in Section B.2.1. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "VfIwe5olYZ", "reviewer": "Reviewer_nR2e", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes adding a PL condition related regularization term onto the optimization loss function of over-parameterized deep neural networks. The intuition was that, by regularizing the inverse of the PL condition, the condition number was decreased, which leads to a faster convergence rate. This regularization was implemented by training a simple gating network that predicts a mask on the weights, so that after applying the mask the pruned network has a smaller condition number. Experiments on BERT, Switch-Transformer and VGG-16 were conducted, and shows that the proposed method results in better performance.", "summary_of_the_review": "In summary, the paper proposes very novel and interesting ideas in regularizing PL (or condition number) to improve convergence speed, and in implementing the regularization. The weak point is that it lacks a lot of important implementation details and theoretical justifications of some claims. I tempororily put it as below threshold. But I would consider it as a good paper, as long as my concerns are addressed and more details are included.\n\n\n-------------After author feedback----------------------------------------------------\n\nThe authors provided very detailed feedback, which clarified many of the implementation details that was missing in the original submission. Most of my initial concerns are addressed, except this one: \n\n> \" The regularization term in Eq. (6) and (7) are evaluated over a large set $\\mathcal{W}$ ... however, the trace of Hessian is evaluated on a single point $w$, see Eq. (41). This is a mismatch, and there should be some theory to justify this.\"\n\nIn addition, the revision, especially Figure 3, provides a directly experimental evidence that the proposed method is successful in decreasing the condition number by prunning, which is align with the motivation of this work.\n\nGiven the above facts, I would like to raise my score accordingly.", "strength_and_weaknesses": "Strength:\n\nThe idea of regularizing the PL condition (hence, the condition number) is novel and interesting. It is widely believed that the condition number directly controls the convergence speed, directly regularizing it should be helpful in optimization speed.\n\nIt is a fascinating idea to learn a mask to select the “good” weights (or weight groups) such that the resulting pruned network has a low condition number. This method makes it feasible to regularize the PL condition. I have never seen such ideas in the literature, and I think it is a good addition to the machine learning community.\n\nThe experimental results on several commonly used networks show that the proposed method is effective in improving the network performance. \n\nWeaknesses:\n\nSeveral *important* aspects of the proposed method and its implementation are missing, which restrict the quality of this paper. Addressing them should improve the paper. I list them below.\n\n> The connection between the regularization term (condition number) $L_f/\\lambda_{min}$ and the trace of Hessian matrix $H_f$ is not clearly established. The authors refers to Prop 11 for the connection, however, there is no explicit relation between the two quantities presented in Prop 11. The only relation is $\\lambda_{min}/L_f \\ge \\lambda_0/L_f-2\\sqrt{n}R\\epsilon$. However, on the right hand side: the first term $\\lambda_0/L_f$ which is the major contribution of the two is not obviously related to the Hessian, noting that $\\lambda_0$ is the smallest eigenvalue of the NTK at initialization; the second term $-2\\sqrt{n}R\\epsilon$, which is negative, should have a much smaller magnitude than the first term, especially when the width is large. \n\n> The regularization term in Eq. (6) and (7) are evaluated over a large set $\\mathcal{W}$, typically large enough to cover the optimization trajectory, and in principle the trace of Hessian should be also taken minimization over $\\mathcal{W}$; however, the trace of Hessian is evaluated on a single point $w$, see Eq. (41). This is a mismatch, and there should be some theory to justify this. \n>>Note that, in most practical cases (including the experiments of this paper), the network is not super wide, and the Hessian matrices may be significantly different from each other within $\\mathcal{W}$. In these cases, a single point evaluation (in Eq.(41)) should not be enough.\n\n> Many important implementation details are missing. It is necessary to present them in the paper.\n\n>> How is the NTK computed for the neural networks: BERT, VGG-16 etc? As mentioned in Section 3.2.2, quantities of the NTK (e.g., $\\lambda_{min}$ and entropy Eq. (8)) are needed as input for the gating network. As far as I know, the NTK for these networks is very hard to compute. It is important to present the computation in the paper.\n\n>> How is the gradient $\\nabla_v L_{prune}$ computed? Especially, what is the expression of this gradient? It seems that the derivative goes through many complicated functions, for example, TopK, trace of Hessian. \n\n>> The most important of these is that: how are the parameter subsets $w_{[i]}$ partitioned for the neural networks, BERT, VGG-16, etc? Difference partition methods may largely affect the performance. \n\nI am also a bit suspicious about Figure 1(a). Condition numbers usually take very large values, typically greater than 1k. Moreover, in theory it should never be less than 1. However, the condition numbers in Figure 1(a) are quite small, even less than 1. I am confused about it, and hope the authors can give some explanation. \n\nThe paper only compares the convergence based on the number of epochs. As the proposed algorithm has an extra training module (training the gating network), I would like to see a comparison based on the wall-clock time also.\n\nDefinition 1 is actually the definition of interpolation, not over-parameterization. Although the two concepts often co-exist in many cases, they are not the same thing. A definition of over-parameterization has to be related to the number of model parameters.\n\nI’d also like to know:\n> at which epochs are the VGG-16 pruned? (I might miss this details, but I did not find it).\n\n> performance of the proposed method on ResNets (e.g., ResNet-32, ResNet-50) with CIFAR-10. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Kyvo-fHMdF", "reviewer": "Reviewer_QD8B", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposed a novel method for regularizing over-parametrized models during optimization. The method is motivated by a theoretical result based on the condition number of the model. It improves the generalization performance empirically.  ", "summary_of_the_review": "The paper proposes a novel algorithm based on a solid theoretical result. The algorithm achieves better generalization than baselines. However, there are several questions to be addressed in the experiment section. ", "strength_and_weaknesses": "Strength: \n- The motivation of the algorithm is well supported by a rigorous theoretical result. \n- The empirical results demonstrate the algorithm works quite well. \n\nExperimental weakness:\n- The paper shows the results of training until convergence for VGG-16 but not for BERT or Switch-Transformer.\n- The paper uses VGG-16 instead of more popular models in computer vision for the third experiment. Why?\n- Is it possible to refer some other papers (better if official) where the results of training BERT on WikiText-2 and other experiments are report? \n- Can we conclude the algorithm is more efficient than traditional optimization algorithms with the loop for pruning? It involves additional complexity. \n\nQuestion: \n- What will happen if we decrease the number of parameters in the model (such as BERT) gradually by tuning hidden dimension, for example? Does it achieve better performance as well. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "5dFMKTtngPv", "reviewer": "Reviewer_ng9v", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This is a theory-inspired empirical paper. The authors found that condition number (PL/Lipschitz) is an important parameter related to convergence and generalization, so they try to add condition number as the regularization term during training, and also use it to do pruning. It turns out that by doing that, the algorithm runs faster in training and generalizes better, compared with many baseline algorithms. \n\n", "summary_of_the_review": "Since I think this paper is an empirical paper, and the connection between the trace estimation and condition number is weak, I currently give a weak reject. \n\nHowever, I am willing to change it to weak accept, if the authors, or other reviewers can convince me, that the regularization + pruning techniques along is a very good contribution to the community. \n\n-----After rebuttal ----\nAfter reading the rebuttal, I think: \n1. I still agree with Reviewer a5MJ that this is not a good theory paper. \n2. I think the authors demonstrated that their methods are useful emprically to the community. \n3. Other two reviewers seem to like the paper. \n4. The authors demonstrated that there exists some connections between trace estimation and condition number. \n\nSo I raise the score. ", "strength_and_weaknesses": "This is a theory-inspired paper, so presumably the emprical ideas were originated from the theory community, instead of random hacking. It is hard for me to evaluate the experimental results as an expert, but it seems to me that all the results are pretty good, because they easily beat all the baselines. \n\nHowever, it is worth pointing out that \n1. The theory contribution of this paper is minimal. It seems to me that the main theory result is Theorem 6, which is a direct application of the definition of Pl condition and stability. \n2. Regularizing the condition number is difficult. So what they do is convert it to the minimal eigenvalue of NTK. But even minimal eigenvalue of NTK is also difficult to compute, so they further convert it to the norm of Hessian. But the norm of Hessian is still expensive to compute, so they further convert it to the trace of Hessian. Finally, the trace of Hessian can be efficiently approximated using random vector sampling. After so many steps of transformations, it is hard to justify that the empirical success is indeed for the theory of condition number. \n\nIt seems to me that this paper can be greatly improved, if the authors can show that adding regularization of the trace of Hessian can help optimization & generalization, because it seems to be what the model is actually doing. The theory of NTK relies on pretty strong assumptions on the learning rate and network width, so it is difficult for me to believe that the condition number is the actual reason for the model to work. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster