Fast as CHITA: Neural Network Pruning with Combinatorial Optimization
The sheer size of modern neural networks makes model serving a serious
computational challenge. A popular class of compression techniques overcomes
this challenge by pruning or sparsifying the weights of pretrained networks.
While useful, these techniques often face serious tradeoffs between
computational requirements and compression quality. In this work, we propose a
novel optimization-based pruning framework that considers the combined effect
of pruning (and updating) multiple weights subject to a sparsity constraint.
Our approach, CHITA, extends the classical Optimal Brain Surgeon framework and
results in significant improvements in speed, memory, and performance over
existing optimization-based approaches for network pruning. CHITA's main
workhorse performs combinatorial optimization updates on a memory-friendly
representation of local quadratic approximation(s) of the loss function. On a
standard benchmark of pretrained models and datasets, CHITA leads to
significantly better sparsity-accuracy tradeoffs than competing methods. For
example, for MLPNet with only 2% of the weights retained, our approach improves
the accuracy by 63% relative to the state of the art. Furthermore, when used in
conjunction with fine-tuning SGD steps, our method achieves significant
accuracy gains over the state-of-the-art approaches.
Riade Benbaki, Wenyu Chen, Xiang Meng, Hussein Hazimeh, Natalia Ponomareva, Zhe Zhao, Rahul Mazumder
# Fast as CHITA: Neural Network Pruning with Combinatorial Optimization

###### Abstract

The sheer size of modern neural networks makes model serving a serious computational challenge. A popular class of compression techniques overcomes this challenge by pruning or sparsifying the weights of pretrained networks. While useful, these techniques often face serious tradeoffs between computational requirements and compression quality. In this work, we propose a novel optimization-based pruning framework that considers the combined effect of pruning (and updating) multiple weights subject to a sparsity constraint. Our approach, CHITA, extends the classical Optimal Brain Surgeon framework and results in significant improvements in speed, memory, and performance over existing optimization-based approaches for network pruning. CHITA's main workhorse performs combinatorial optimization updates on a memory-friendly representation of local quadratic approximation(s) of the loss function. On a standard benchmark of pretrained models and datasets, CHITA leads to significantly better sparsity-accuracy tradeoffs than competing methods. For example, for MLPNet with only 2% of the weights retained, our approach improves the accuracy by 63% relative to the state of the art. Furthermore, when used in conjunction with fine-tuning SGD steps, our method achieves significant accuracy gains over the state-of-the-art approaches.

Machine Learning, ICML

## 1 Introduction

Modern neural networks tend to use a large number of parameters (Devlin et al., 2018; He et al., 2016), which leads to high computational costs during inference. A widely used approach to mitigate inference costs is to prune or sparsify pre-trained networks by removing parameters (Blandock et al., 2020). The goal is to obtain a network with significantly fewer parameters and minimal loss in performance. This makes model storage and deployment cheaper and easier, especially in resource-constrained environments.

Generally speaking, there are two main approaches for neural net pruning: (i) magnitude-based and (ii) impact-based. Magnitude-based heuristic methods (e.g., Hanson and Pratt, 1988; Mozer and Smolensky, 1989; Gordon et al., 2020) use the absolute value of weight to determine its importance and whether or not it should be pruned. Since magnitude alone may not be a perfect proxy for weight relevance, alternatives have been proposed. To this end, impact-based pruning methods (e.g. LeCun et al., 1989; Hassibi and Stork, 1992; Singh and Alistarh, 2020) remove weights based on how much their removal would impact the loss function, often using second-order information on the loss function. Both of these approaches, however, may fall short of considering the _joint effect_ of removing (and updating) multiple weights simultaneously. The recent method CBS (Combinatorial Brain Surgeon) (Yu et al., 2022) is an optimization-based approach that considers the joint effect of multiple weights. The authors show that CBS leads to a boost in the performance of the pruned models. However, CBS can be computationally expensive: it makes use of a local model based on the second-order (Hessian) information of the loss function, which can be prohibitively expensive in terms of runtime and/or memory (e.g., CBS takes hours to prune a network with 4.2 million parameters).

In this work, we propose CHITA (Combinatorial Hessian-free Iterative Thresholding Algorithm), an efficient optimization-based framework for network pruning at scale. Our approach follows earlier works that consider a local quadratic approximation of the loss function based on the second-order Hessian information. Different from previous works, we make use of a simple yet important observation with which we can avoid computing and storing the Hessian matrix to solve the optimization problem (hence the name "Hessian-free" in CHITA)--this allows us to address large networks efficiently. Specifically, we propose an equivalent reformulation of the problem as an \(\ell_{0}\)-constrained sparse linear regression problem with a data matrix \(A\in\mathbb{R}^{n\times p}\)
where \(p\) is the number of trainable parameters in the original model and \(n\lesssim 10^{3}\) (usually, \(p\gg n\)) is the number of the sub-samples used in approximating the Hessian. Compared to state-of-the-art approaches that consider a dense \(p\times p\) matrix approximation of the Hessian, our approach leads to a significant reduction in memory usage (up to \(10^{3}\) times for \(p=10^{6}\)) without any approximation.

Furthermore, we propose a novel approach to minimize our \(\ell_{0}\) regression reformulation, leveraging active set strategies, better stepsize selection, and various methods to accelerate convergence on the selected support. Our proposed approach leads to significant improvements over Iterative Hard Thresholding methods (IHT, Blumensath and Davies, 2009) commonly used in sparse learning literature. For instance, our framework can prune a network with 4.2M parameters to 80% sparsity (i.e., 0.84M nonzero parameters) in less than a minute and using less than 20GB of RAM 1.

Footnote 1: For reference, CBS and Woodfisher would run out of memory in similar circumstances.

Since the local quadratic model approximates the loss function only in a small neighborhood of the current solution (Singh and Alistarh, 2020), we also propose a multi-stage algorithm that updates the local quadratic model during pruning (but without retraining) and solves a more constrained problem in each stage, going from dense weights to sparse ones. Our experiments show that the resulting pruned models have a notably better accuracy compared to that of our single-stage algorithm and other pruning approaches. Furthermore, when used in the gradual pruning setting (Gale et al., 2019; Singh and Alistarh, 2020; Blalock et al., 2020) where re-training between pruning steps is performed, our pruning framework results in significant performance gains compared to state-of-the-art unstructured pruning methods.

ContributionsOur contributions can be summarized as follows:

* We propose CHITA, an optimization framework for network pruning based on local quadratic approximation(s) of the loss function. We propose an \(\ell_{0}\)-constrained sparse regression reformulation that avoids the pitfalls of storing a large dense Hessian, resulting in a significant reduction in memory usage (we work with an \(n\times p\) matrix instead of a \(p\times p\) one, with often \(n\ll p\)).
* A key workhorse of CHITA is a novel IHT-based algorithm to obtain good solutions to the sparse regression formulation. Exploiting problem structure, we propose methods to accelerate convergence and improve pruning performance, such as a new and efficient stepsize selection scheme and rapidly updating weights on the support. This leads to up to 1000x runtime improvement compared to existing network pruning algorithms.
* We show performance improvements across various models and datasets. In particular, CHITA results in a 98% sparse (i.e., 98% of weights in dense model are set to zero) MLPNet with 90% test accuracy (3% reduction in test accuracy compared to the dense model), which is a significant improvement over the previously reported best accuracy (55%) by CBS. As an application of our framework, we use it for gradual pruning and observe notable performance gains against state-of-the-art gradual pruning approaches.

## 2 Problem Setup and Related Work

In this section we present the general setup with connections to related work--this lays the foundation for our proposed methods discussed in Section 3.

### Problem Setup

Consider a neural network with an empirical loss function \(\mathcal{L}(w)=\frac{1}{N}\sum_{i=1}^{N}\ell_{i}(w),\) where \(w\in\mathbb{R}^{p}\) is the vector of trainable parameters, \(N\) is the number of data points (samples), and \(\ell_{i}(w)\) is a twice-differentiable function on the \(i\)-th sample. Given a pre-trained weight vector \(\bar{w}\in\mathbb{R}^{p}\), our goal is to set some parameters to zero and possibly update other weights while maintaining the original model's performance (e.g., accuracy) as much as possible. In mathematical terms, given a pre-trained weight \(\bar{w}\) and a target level of sparsity \(\tau\in(0,1)\), we aim to construct a new weight vector \(w\in\mathbb{R}^{p}\) that satisfies :

* The loss function at \(w\) is as close as possible to the loss before pruning: \(\mathcal{L}(w)\approx\mathcal{L}(\bar{w})\).
* The number of nonzero weights at \(w\) respects the sparsity budget2: \(\|w\|_{0}\leq(1-\tau)p\). Footnote 2: Here \(\ell_{0}\) norm \(\|w\|_{0}\) denotes the number of nonzero in the vector \(w\).

Similar to LeCun et al. (1989); Hassibi and Stork (1992); Singh and Alistarh (2020), we use a local quadratic approximation of \(\mathcal{L}\) around the pre-trained weight \(\bar{w}\):

\[\mathcal{L}(w)=\mathcal{L}(\bar{w})+\nabla\mathcal{L}(\bar{w})^ {\top}(w-\bar{w})+\\ \frac{1}{2}(w-\bar{w})^{\top}\nabla^{2}\mathcal{L}(\bar{w})(w- \bar{w})+O(\|w-\bar{w}\|^{3}). \tag{1}\]

With certain choices of gradient and Hessian approximations \(g\approx\nabla\mathcal{L}(\bar{w}),H\approx\nabla^{2}\mathcal{L}(\bar{w})\), and ignoring higher-order terms, the loss \(\mathcal{L}\) can be locally approximated by:

\[Q_{0}(w):=\mathcal{L}(\bar{w})+g^{\top}(w-\bar{w})+\frac{1}{2}(w-\bar{w})^{\top }H(w-\bar{w}). \tag{2}\]

Pruning the local approximation \(Q_{0}(w)\) of the network can be naturally formulated as an optimization problem to minimize \(Q_{0}(w)\) subject to a cardinality constraint, i.e.,

\[\min_{w}\ \ Q_{0}(w)\qquad\text{s.t.}\quad\|w\|_{0}\leq k. \tag{3}\]
For large networks, solving Problem (3) directly (e.g., using iterative optimization methods) is computationally challenging due to the sheer size of the \(p\times p\) matrix \(H\). In Section 3.1, we present an exact, hessian-free reformulation of Problem (3), which is key to our scalable approach.

### Related Work

Impact-based pruning dates back to the work of LeCun et al. (1989) where the OBD (Optimal Brain Damage) framework is proposed. This approach, along with subsequent ones (Hassibi and Stork, 1992; Singh and Alistarh, 2020; Yu et al., 2022) make use of local approximation (2). It is usually assumed (but not in our work) that \(\bar{w}\) is a local optimum of the loss function, and therefore \(g=0\) and \(\mathcal{L}(w)\approx\mathcal{L}(\bar{w})+\frac{1}{2}(w-\bar{w})^{\top}H(w- \bar{w})\). Using this approximation, OBD (Optimal Brain Damage, LeCun et al. (1989)) searches for a single weight \(i\) to prune with minimal increase of the loss function, while also assuming a diagonal Hessian \(H\). If the \(i\)-th weigth is pruned (\(w_{i}=0,w_{j}=\bar{w}_{j}\forall j\neq i\)), then the loss function increases by \(\delta\mathcal{L}_{i}=\frac{\bar{w}_{i}^{2}}{2\nabla^{2}\mathcal{L}(\bar{w}) _{ii}}\). This represents a score for each weight, and is used to prune weights in decreasing order of their score. OBS (Optimal Brain Surgeon, Hassibi and Stork (1992)) extends this by no longer assuming a diagonal Hessian, and using the optimality conditions to update the un-pruned weights. The authors also propose using the empirical Fisher information matrix, as an efficient approximation to the Hessian matrix. Layerwise OBS (Dong et al., 2017) proposes to overcome the computational challenge of computing the full (inverse) Hessian needed in OBS by pruning each layer independently, while Singh and Alistarh (2020) use block-diagonal approximations on the Hessian matrix, which they approximate by the empirical Fisher information matrix on a small subset of the training data (\(n\ll N\)):

\[\nabla^{2}\mathcal{L}(\bar{w})\approx H=\frac{1}{n}\sum_{i=1}^{n}\nabla\ell_{i }(\bar{w})\nabla\ell_{i}(\bar{w})^{\top}. \tag{4}\]

While these approaches explore different ways to make the Hessian computationally tractable, they all rely on the OBD/OBS framework of pruning a single weight, and do not to consider the possible interactions that can arise when pruning multiple weights. To this end, Yu et al. (2022) propose CBS (Combinatorial Brain Surgeon) an algorithm to approximately solve (3). While CBS shows impressive improvements in the accuracy of the pruned model over prior work, it operates with the full dense \(p\times p\) Hessian \(H\). This limits scalability both in compute time and memory utilization, as \(p\) is often in the order of millions and more.

Choices of gradient approximation \(g\).As mentioned earlier, most previous work assumes that the pre-trained weights \(\bar{w}\) is a local optimum of the loss function \(\mathcal{L}\), and thus take the gradient \(g=0\). However, the gradient of the loss function of a pre-trained neural network may not be zero in practice due to early stopping (or approximate optimization) (Yao et al., 2007). Thus, the WoodTaylor approach (Singh and Alistarh, 2020) proposes to approximate the gradient by the stochastic gradient, using the same samples for estimating the Hessian. Namely,

\[g=\frac{1}{n}\sum_{i=1}^{n}\nabla\ell_{i}(\bar{w}). \tag{5}\]

One-shot and gradual pruning.Generally speaking, _one-shot pruning_ methods (LeCun et al., 1989; Singh and Alistarh, 2020; Yu et al., 2022) can be followed by a few fine-tuning and re-training steps to recover some of the accuracy lost when pruning. Furthermore, recent work has shown that pruning and re-training in a gradual fashion (hence the name, _gradual pruning_) can lead to big accuracy gains (Han et al., 2015; Gale et al., 2019; Zhu and Gupta, 2018). The work of Singh and Alistarh (2020) further shows that gradual pruning, when used with well-performing one-shot pruning algorithms, can outperform state-of-the-art unstructured pruning methods. In this paper, we focus on the one-shot pruning problem and show how our pruning framework outperforms other one-shot pruning methods (see Section 4.1). We then show that, if applied in the gradual pruning setting, our pruning algorithm outperforms existing approaches (see Section 4.2), establishing new state-of-the-art on MobileNetV1 and ResNet50.

## 3 Our Proposed Framework: Chitta

In this section, we present our algorithmic framework CHITA (Combinatorial Hessian-free Iterative Thresholding Algorithm) for pruning a network to specific sparsity level. We formulate sparse network pruning by considering \(\ell_{0}\)-regression problem(s) and propose scalable algorithms. For example, we can address networks with size \(p\approx 10^{6},n\approx 10^{3},k\approx 10^{5}\) in less than one minute and using less than 20GB of memory. Our basic single-stage algorithm is an improved and scalable variant of IHT to solve (3). In Section 3.3, we propose a multi-stage framework that repeatedly refines the local quadratic approximation and optimizes it (under sparsity constraints) resulting in further performance boosts as shown in Section 4.

### An \(\ell_{0}\)-regression formulation

Our formulation is based on a critical observation that the Hessian approximation in (4) has a low-rank structure:

\[H=\frac{1}{n}\sum_{i=1}^{n}\nabla\ell_{i}(\bar{w})\nabla\ell_{i}(\bar{w})^{ \top}=\frac{1}{n}A^{\top}A\in\mathbb{R}^{p\times p}, \tag{6}\]
where \(A=[\nabla\ell_{1}(\bar{w}),\ldots,\nabla\ell_{n}(\bar{w})]^{\top}\in\mathbb{R}^{n \times p}\) has rank at most \(n\) with \(10^{3}\geq n\ll p\).

Using observation (6) and the gradient expression (5), we note that problem (3) can be equivalently written in the following Hessian-free form:

\[\min_{w}\ \frac{1}{2}\|b-Aw\|^{2}\qquad\text{s.t.}\ \ \ \ \ \|w\|_{0}\leq k, \tag{7}\]

where \(b:=A\bar{w}-e\in\mathbb{R}^{n}\) and \(e\) is a vector of ones. Furthermore, to improve solution quality (see discussion below), we include a ridge-like regularizer of the form \(\|w-\bar{w}\|^{2}\) to the objective in (7). This leads to the following problem:

\[\min_{w}\ Q(w):=\frac{1}{2}\|b-Aw\|^{2}+\frac{n\lambda}{2}\|w-\bar{w}\|^{2}\ \operatorname{s.t.}\ \|w\|_{0}\leq k, \tag{8}\]

where \(\lambda\geq 0\) is a parameter controlling the strength of the ridge regularization. Note that our algorithm actually applies to the general form (8). Importantly, the regression formulation (8) does not require us to compute or store the full Hessian matrix \(H\in\mathbb{R}^{p\times p}\). As discussed in Section 3.2, we only need to operate on the low-rank matrix \(A\) throughout our algorithm--this results in substantial gains in terms of both memory consumption and runtime.

Ridge term and choices of \(\lambda\).We observe empirically that the success of our final pruned model depends heavily on the accuracy of the quadratic approximation of the loss function. Since this approximation is local, it is essential to ensure that the weights \(w\) during the pruning process are sufficiently close to the initial weights \(\bar{w}\). One way3 to achieve this is by including a squared \(\ell_{2}\) penalty, also known as the ridge, on the difference \(w-\bar{w}\). This regularization technique does not appear to be explicitly4 considered in previous works (Hassibi and Stork, 1992; Singh and Alistarh, 2020; Yu et al., 2022) on pruning using local quadratic approximations. The usefulness of the regularization term \(\lambda\) is further explored in Appendix B.2.1. We observe that a well-chosen ridge term can help improve the test accuracy on MLPNet by 3%.

Footnote 3: Another way is to introduce a multi-stage procedure, as explained in Section 3.3.

Footnote 4: It appears to be used implicitly though to obtain an invertible Fisher matrix which is achieved by adding a small diagonal \(\lambda_{0}I\) to the Fisher matrix.

Relation to prior work on \(\ell_{0}\)-regression problems.There is a substantial literature on algorithms for solving \(\ell_{0}\)-regularized linear regression problems. We provide a brief overview of related work, but it is worth noting that the context of network pruning and the problem-scale we consider here makes our work different from earlier works in \(\ell_{0}\)-sparse linear regression. Blumensath and Davies (2009) developed an iterative hard thresholding method, which involves projecting the weights onto the feasible set after each gradient step. Bertsimas and Van Parys (2020); Hazimeh et al. (2022) proposed algorithms to solve sparse regression problems to optimality via branch-and-bound. Beck and Eldar (2013) explore coordinate descent-type algorithms that update one/two coordinates at a time. Hazimeh and Mazumder (2020) propose efficient algorithms based on coordinate descent and local combinatorial optimization that applies to the unconstrained \(\ell_{0}\ell_{2}\)-penalized regression problem. We refer the reader to Hazimeh et al. (2022) for a more comprehensive discussion of related work.

In summary, earlier methods for \(\ell_{0}\)-regularized linear regression are quite effective at discovering high-quality solutions to Problem (7) for small to medium-sized problems and require \(k\) to be sufficiently small for efficiency. However, these methods are not well-suited for tackling large network pruning problems (e.g, \(p\sim 10^{6}\) and \(k\sim 10^{5}\)) due to slow convergence or expensive per-iteration cost. To address this issue, we propose novel approaches for scalability in Section 3.2. Additionally, we emphasize that (8) arises from a local approximation of the true loss \(\mathcal{L}\) around \(\bar{w}\). Therefore, achieving a high-quality solution for (8) alone does not guarantee a pruned network with high accuracy. To this end, we study a multi-stage extension (see Section 3.3) that requires solving several problems of the form (8).

### Our proposed algorithm for problem (8)

We present the core ideas of our proposed algorithm for Problem (8), and discuss additional details in Appendix A.

Our optimization framework relies on the IHT algorithm (Blumensath and Davies, 2009; Bertsimas et al., 2016) that optimizes (8) by simultaneously updating the support and the weights. By leveraging the low-rank structure, we can avoid the computational burden of computing the full Hessian matrix, thus reducing complexity.

The basic version of the IHT algorithm can be slow for problems with a million parameters. To improve the computational performance of our algorithm we propose a new line search scheme. Additionally, we use use an active set strategy and schemes to update the weights on the nonzero weights upon support stabilization. Taken together, we obtain notable improvements in computational efficiency and solution quality over traditional IHT, making it a viable option for network pruning problems at scale.

#### 3.2.1 Structure-aware IHT update

The IHT algorithm operates by taking a gradient step of size \(\tau\) from the current iteration, then projects it onto the set of points with a fixed number of non-zero coordinates through hard thresholding. Specifically, for any vector \(x\), let \(\mathcal{I}_{k}(x)\) denote the indices of \(k\) components of \(x\) that have the largest absolute value. The hard thresholding operator
\(P_{k}(x)\) is defined as \(y_{i}=x_{i}\) if \(i\in\mathcal{I}_{k}(x)\), and \(y_{i}=0\) if \(i\notin\mathcal{I}_{k}(x)\); where \(y_{i}\) is the \(i\)-th coordinate of \(P_{k}(x)\). IHT applied to problem (8) leads to the following update:

\[w^{t+1} =\mathsf{HT}(w^{t},k,\tau^{s}):=P_{k}\left(w^{t}-\tau^{s}\nabla Q (w^{t})\right) \tag{9}\] \[=P_{k}\left(w^{t}-\tau^{s}(A^{\top}(Ab-w^{t})+n\lambda(w^{t}- \bar{w}))\right),\]

where \(\tau^{s}>0\) is a suitable stepsize. The computation of \(\mathsf{HT}(w^{t},k,\tau^{s})\) involves only matrix-vector multiplications with \(A\) (or \(A^{\top}\)) and a vector, which has a total computation cost of \(O(np)\). This is a significant reduction compared to the \(O(p^{2})\) cost while using the full Hessian matrix as Singh and Alistarh (2020); Yu et al. (2022) do.

**Active set strategy.** In an effort to further facilitate the efficiency of the IHT method, we propose using an active set strategy, which has been proven successful in various contexts such as (Nocedal and Wright, 1999; Friedman et al., 2010; Hazimeh and Mazumder, 2020). This strategy works by restricting the IHT updates to an active set (a relatively small subset of variables) and occasionally augmenting the active set with variables that violate certain optimality conditions. By implementing this strategy, the iteration complexity of the algorithm can be reduced to \(O(nk)\) in practice, resulting in an improvement, when \(k\) is smaller than \(p\). The algorithm details can be found in Appendix A.3.

#### 3.2.2 Determining a good stepsize

Choosing an appropriate stepsize \(\tau^{s}\) is crucial for fast convergence of the IHT algorithm. To ensure convergence to a stationary solution, a common choice is to use a constant stepsize of \(\tau^{s}=1/L\)(Bertsimas et al., 2016; Hazimeh and Mazumder, 2020), where \(L\) is the Lipschitz constant of the gradient of the objective function. This approach, while reliable, can lead to conservative updates and slow convergence--refer to Appendix A.1 for details. An alternative method for determining the stepsize is to use a backtracking line search, as proposed in Beck and Teboulle (2009). The method involves starting with a relatively large estimate of the stepsize and iteratively shrinking the step size until a sufficient decrease of the objective function is observed. However, this approach requires multiple evaluations of the objective function, which can be computationally expensive.

**Our novel scheme.** We propose a novel line search method for determining the stepsize to improve the convergence speed of IHT. Specifically, we develop a method that (approximately) finds the stepsize that leads to the maximum decrease in the objective, i.e., we attempt to solve

\[\min_{\tau^{s}\geq 0}\ \ g(\tau^{s}):=Q\left(P_{k}\left(w^{t}-\tau^{s}\nabla Q (w^{t})\right)\right). \tag{10}\]

For general objective functions, solving the line search problem (as in (10)) is challenging. However, in our problem, we observe and exploit an important structure: \(g(\tau^{s})\) is a piecewise quadratic function with respect to \(\tau^{s}\). Thus, the optimal stepsize on each piece can be computed exactly, avoiding redundant computations (associated with finding a good stepsize) and resulting in more aggressive updates. In Appendix A.1, we present an algorithm that finds a good stepsize by exploiting this structure. Compared to standard line search, our algorithm is more efficient, as it requires fewer evaluations of the objective function and yields a stepsize that results in a steeper descent.

#### 3.2.3 Additional techniques for scalability

While the IHT algorithm can be quite effective in identifying the appropriate support, its progress slows down considerably once the support is identified (Blumensath, 2012), resulting in slow convergence. We propose two techniques that refine the non-zero coefficients by solving subproblems to speedup the overall optimization algorithm: (i) Coordinate Descent (CD, Bertsekas, 1997; Nesterov, 2012) that updates each nonzero coordinate (with others fixed) as per a cyclic rule; (ii) Back solve based on Woodbury formula (Max, 1950) that calculates the optimal solution exactly on a restricted set of size \(k\). We found both (i), (ii) to be important for improving the accuracy of the pruned network. Further details on the strategies (i), (ii) are in Appendix A.2 and A.4.

Our single-stage algorithm CHITA glues together the different pieces discussed above into a coherent algorithm. It takes as input a low-rank matrix \(A\), the initial weight \(\bar{w}\) and the \(\ell_{0}\)-constraint \(k\); and returns a pruned weight \(w\) that serves as a good solution to (8).

### A multi-stage procedure

Our single-stage methods (cf Section 3.2) lead to high-quality solutions for problem (8). Compared to existing methods, for a given sparsity level, our algorithms deliver a better objective value for problem (8)--for eg, see Figure 2(b). However, we note that the final performance (e.g., accuracy) of the pruned network depends heavily on the quality of the local quadratic approximation. This is particularly true when targeting high levels of sparsity (i.e., zeroing out many weights), as the objective function used in (8) may not accurately approximate the true loss function \(\mathcal{L}\). To this end, we propose a multi-stage procedure named CHITA++ that improves the approximation quality by iteratively updating and solving local quadratic models. We use a scheduler to gradually increase the sparsity constraint and take a small step towards higher sparsity in each stage to ensure the validity of the local quadratic approximation. Our multi-stage procedure leverages the efficiency of the single-stage approaches and can lead to pruned networks with improved accuracy by utilizing more accurate approximations of the true loss function. For example, our experiments show that
the multi-stage procedure can prune ResNet20 to 90% sparsity in just a few minutes and increases test accuracy from 15% to 79% compared to the single-stage method. Algorithm 1 presents more details on CHITA++.

Our proposed multi-stage method differs significantly from the gradual pruning approach described in Han et al. (2015). While both methods involve pruning steps, the gradual pruning approach also includes fine-tuning steps in which SGD is applied to further optimize the parameters for better results. However, these fine-tuning steps can be computationally expensive, usually taking days to run. In contrast, our proposed multi-stage method is a one-shot pruning method and only requires constructing and solving Problem (8) several times, resulting in an efficient and accurate solution. This solution can then be further fine-tuned using SGD or plugged into the gradual pruning framework, something we explore in Section 4.2.

```
0: Pre-trained weights \(\bar{w}\), a target sparsity level \(\tau\), number of stages \(f\).
1:Initialization: Construct a increasing sequence of sparsity parameters \(\tau_{1},\tau_{2},\ldots,\tau_{f}=\tau\); and set \(w^{0}=\bar{w}\)
2:for\(t=1,2,\ldots,f\)do
3: At current solution \(w^{t-1}\), calculate the gradient based on a batch of \(n\) data points and construct the matrix \(A\) given in (4).
4: Obtain a solution \(w^{t}\) to problem (8) by invoking CHITA\((A,\bar{w},k)\) with \(\bar{w}=w^{t-1}\) and number of nonzeros \(k=\lfloor(1-\tau_{t})p\rfloor\).
5:endfor
```

**Algorithm 1** CHITA++: a multi-stage pruning procedure

## 4 Experimental Results

We compare our proposed framework with existing approaches, for both one-shot and gradual pruning.

### One shot pruning

We start by comparing the performance of our methods: CHITA (single-stage) and CHITA++ (multi-stage) with several existing state-of-the-art one-shot pruning techniques on various pre-trained networks. We use the same experimental setup as in recent work Yu et al. (2022); Singh & Alistarh (2020). The existing pruning methods we consider include MP (Magnitude Pruning, Mozer & Smolensky, 1989), WF (WoodFisher, Singh & Alistarh, 2020), CBS (Combinatorial Brain Surgeon, Yu et al., 2022) and M-FAC (Matrix-Free Approximate Curvature, Frantar et al., 2021). The pre-trained networks we use are MLPNet (30K parameters) trained on MNIST LeCun et al. (1998), ResNet20 He et al. (2016, 200k parameters) trained on CIFAR10 Krizhevsky et al. (2009), and MobileNet (4.2M parameters) and ResNet50 He et al. (2016, 22M parameters) trained on ImageNet Deng et al. (2009). For further details on the choice of the Hessian approximation, we refer the reader to Appendix A.5. Detailed information on the experimental setup and reproducibility can be found in Appendix B.1.1.

#### 4.1.1 Runtime comparison

Recent works that use the empirical Fisher information matrix for pruning purposes Singh & Alistarh (2020); Yu et al. (2022) show that using more samples for Hessian and gradient approximation results in better accuracy. Our experiments also support this conclusion. However, most prior approaches become computationally prohibitive as sample size \(n\) increases. As an example, the Woodfisher and CBS algorithms require hours to prune a MobileNet when \(n\) is set to 1000, and their processing time increases at least linearly with \(n\). In contrast, our method has been designed with efficiency in mind, and we have compared it to M-FAC, a well-optimized version of Woodfisher that is at least 1000 times faster. The results, as depicted in Figure 1, demonstrate a marked improvement in speed for our algorithm, with up to 20 times faster performance.

#### 4.1.2 Accuracy of the pruned models

**Comparison against state-of-the-art.** Table 1 compares the test accuracy of MLPNet, ResNet20 and MobileNetV1 pruned to different sparsity levels. Our single-stage method achieves comparable results to other state-of-the-art approaches with much less time consumption. The multi-stage method (CHITA++) outperforms other methods by a large margin, especially with a high sparsity rate.

Figure 1: Runtime comparison between our single-stage approaches and M-FAC (the fastest among the competitive methods) while pruning MLPNet and ResNet20 to 90% sparsity level (90% of the entries are zero). Note that Woodfisher and CBS are at least 1000 times slower than M-FAC. The error bar represents the standard error over four runs. CHITA here uses IHT to find a support and performs a back-solve on the found support.
**One-shot pruning on ResNet50.** We further compare our approach to competing methods on ResNet50, an even larger network where some pruning algorithms, like CBS, do not scale. In Figure 2, we evaluate the performance of our algorithm in comparison to M-FAC and Magnitude Pruning (MP) using two metrics: test accuracy and the final objective value of the \(\ell_{0}\)-constrained problem (8). As both M-FAC and our algorithm aim to minimize this objective, it can be used to judge the efficacy of our model in solving problem (8). As seen in the figure, our approach achieves a lower objective value, and in this case, it also results in a better test accuracy for the final pruned network.

**Sparsity schedule in multi-stage procedure.** We study the effect of the sparsity schedule (i.e., choice of \(\tau_{1}\leq\tau_{2}\leq\cdots\leq\tau_{f}=\tau\) in Algorithm 1) on the performance of CHITA++. We compare test accuracy of three different schedules: (i) exponential mesh, (ii) linear mesh, and (iii) constant mesh. For these schedules, \(f\) is set to be \(15\). For the first two meshes, \(\tau_{1}\) and \(\tau_{15}\) are fixed as \(0.2\) and \(0.9\), respectively. As shown in Figure 3, the exponential mesh computes \(\tau_{2},\ldots,\tau_{14}\) by drawing an exponential function, while the linear mesh adopts linear interpolation (with \(\tau_{1}\) and \(\tau_{15}\) as endpoints) to determine \(\tau_{2},\ldots,\tau_{14}\) and the constant mesh has \(\tau_{1}=\tau_{2}=\cdots=\tau_{15}\).

Figure 4 plots the test accuracy of the three schedules over the number of stages. We observe that the linear mesh outperforms the exponential mesh in the first few iterations, but its performance drops dramatically in the last two iterations. The reason is that in high sparsity levels, even a slight increase in the sparsity rate leads to a large drop in accuracy. Taking small "stepsizes" in high sparsity levels allows the exponential mesh to fine-tune the weights in the last several stages and achieve good performance.

\begin{table}
\begin{tabular}{c|c c c c c c} \hline \hline Network & Sparsity & MP & WF & CBS & CHITA & CHITA++ \\ \hline  & 0.5 & 93.93 & 94.02 & 93.96 & 93.97 (0.03) & **95.97** (0.05) \\  & 0.6 & 93.78 & 93.82 & 93.96 & 93.94 (0.02) & **95.93** (0.04) \\ MLPNet & 0.7 & 93.62 & 93.77 & 93.98 & 93.80 (0.01) & **95.89** (0.06) \\ on MNIST & 0.8 & 92.98 & 93.57 & 93.90 & 93.59 (0.03) & **95.89** (0.03) \\ (93.97\%) & 0.9 & 90.31 & 96.93 & 91.94 & 92.46 (0.04) & **95.58** (0.03) \\  & 0.95 & 83.64 & 85.54 & 89.82 & 88.09 (0.24) & **94.70** (0.06) \\  & 0.98 & 92.25 & 38.26 & 55.45 & 46.25 (0.85) & **90.73** (0.11) \\ \hline  & 0.3 & 90.77 & **91.37** & 91.35 & **91.37** (**0.04)**, **91.25** (0.08) \\  & 0.4 & 99.89 & 91.95 & **91.21** & 91.19 (0.05) & **91.20** (0.05) \\ ResNet20 & 0.5 & 88.44 & 90.23 & 90.58 & 90.60 (0.07) & **91.04** (0.09) \\ on CIFAR10 & 0.6 & 85.24 & 97.86 & 88.89 & 92.02 (0.19) & **90.78** (0.12) \\ (91.36\%) & 0.7 & 78.79 & 81.05 & 81.84 & 84.20 (43.08) & **90.88** (0.10) \\  & 0.8 & 54.01 & 62.63 & 51.28 & 57.90 (1.04) & **88.72** (0.17) \\  & 0.9 & 119.19 & 14.93 & 13.68 & 15.60 (14.79) & **79.32** (\(\pm\)1.19) \\ \hline  & 0.3 & 71.60 & **71.88** & **71.88** & **71.87** (0.01) & 71.60 (0.02) \\ MobileNetV1 & 0.4 & 69.16 & 71.15 & 71.45 & 71.50 (0.02) & **71.61** (0.02) \\ on ImageNet & 0.5 & 62.61 & 68.91 & 70.21 & 70.42 (0.02) & **70.99** (\(\pm\)0.04) \\ (71.95\%) & 0.6 & 41.94 & 60.90 & 66.37 & 67.30 (0.03) & **69.54** (0.01) \\  & 0.7 & 6.78 & 29.26 & 55.11 & 59.40 (0.09) & **66.42** (0.03) \\  & 0.8 & 0.11 & 0.24 & 16.38 & 29.78 (0.18) & **47.45** (\(\pm\)0.25) \\ \hline \hline \end{tabular}
\end{table}
Table 1: The pruning performance (model accuracy) of various methods on MLPNet, ResNet20, MobileNetV1. As to the performance of MP, WF, and CBS, we adopt the results reported in Yu et al. (2022). We take five runs for our single-stage (CHITA) and multi-stage (CHITA++) approaches and report the mean and standard error (in the brackets). The best accuracy values (significant) are highlighted in bold. Here sparsity denotes the fraction of zero weights in convolutional and dense layers.

Figure 3: Three different sparsity schedules: exponential, linear, and constant schedules.
Additional ablation studies.We perform additional ablation studies to further evaluate the performance of our method. These studies mainly focus on the effect of the ridge term (in Appendix B.2.1), and the effect of the first-order term (in Appendix B.2.2).

### Performance on gradual pruning

To compare our one-shot pruning algorithms against more unstructured pruning methods, we plug CHITA into a gradual pruning procedure (Gale et al., 2019), following the approach in Singh and Alistarh (2020). Specifically, we alternate between pruning steps where a sparse weight is computed and fine-tuning steps on the current support via Stochastic Gradient Descent (SGD). To obtain consistent results, we start from the same pre-trained weights used in Kusupati et al. (2020), and re-train for 100 epochs using SGD during fine-tuning steps, similarly to Kusupati et al. (2020); Singh and Alistarh (2020). We compare our approach against Incremental (Zhu and Gupta, 2018), STR (Kusupati et al., 2020), Global Magnitude (Singh and Alistarh, 2020), Wood-Fisher (Singh and Alistarh, 2020), GMP (Gale et al., 2019), Variational Dropout (Molchanov et al., 2017), RIGL (Evci et al., 2020), SNFS (Dettmers and Zettlemoyer, 2020) and DNW (Wortsman et al., 2019). Further details on training procedure can be found in Appendix B.1.2.

**MobileNetV1.** We start by pruning MobileNetV1 (4.2M parameters). As Table 2 demonstrates, CHITA results in significantly more accurate pruned models than previous state-of-the-art approaches at sparsities 75% and 89%, with only 6% accuracy loss compared to 11.29%, the previous best result when pruning to a sparsity of 89%.

**ResNet50.** Similarly to MobileNetV1, CHITA improves test accuracy at sparsity levels 90%, 95%, and 98% compared to all other baselines, as Table 3 shows. This improvement becomes more noticeable as we increase the target sparsity, with CHITA producing a pruned model with 69.80% accuracy compared to 65.66%, the second-best performance, and previous state-of-the-art.

## 5 Conclusion

In this work we have presented an efficient network pruning framework CHITA, which is based on a novel, hessian

\begin{table}
\begin{tabular}{c c c c c} \hline \hline  & Sparsity & Pruned & Relative Drop (\%) & Remaining \\ Method & (\%) & Accuracy & \(\frac{Prunnel-Dp_{p_{p_{p_{p{p{p{p{p{p{p{p{p{p}}}}}}}}}}}}}\) & \# of params \\ \hline \hline Incremental & 74.11 & 67.70 & -4.11 & 1.09 M \\ STR & 75.28 & 68.35 & -5.07 & 1.04 M \\ Global Magnitude & 75.28 & 69.90 & -2.92 & 1.04 M \\ WoodFisher & 75.28 & 70.09 & -2.65 & 1.04 M \\
**CRITA** & 75.28 & **71.11** & **-1.23** & 1.04 M \\ \hline Incremental & 89.03 & 61.80 & -12.46 & 0.46 M \\ STR & 89.01 & 62.10 & -13.75 & 0.46 M \\ Global Magnitude & 89.00 & 63.02 & -12.47 & 0.46 M \\ WoodFisher & 89.00 & 63.87 & -11.29 & 0.46 M \\
**CRITA** & 89.00 & **67.68** & **-6.00** & 0.46 M \\ \hline \hline \end{tabular}
\end{table}
Table 2: Results of gradually pruning MobileNetV1 in 75% and 89% sparsity regimes, comparing CHITA to other baselines (Dense accuracy: 72.00%). We also include the relative drop in accuracy to account for different methods starting from different dense weights. CHITA numbers are averaged across two runs. Numbers for other baselines are taken from Singh and Alistarh (2020).

\begin{table}
\begin{tabular}{c c c c c} \hline \hline  & Sparsity & Pruned & Relative Drop (\%) & Remaining \\ Method & (\%) & Accuracy & \(\frac{Prunnel-Dp_{p_{p{p{p{p{p{p{p{p{p{p{p}}}}}}}}}}}}{D}\) & \# of params \\ \hline \hline GMP + LS & 90.00 & 73.91 & -3.62 & 2.56 M \\ Variational Dropout & 90.27 & 73.84 & -3.72 & 2.49 M \\ RIGL + ERK & 90.00 & 73.00 & -4.94 & 2.56 M \\ SNFS + LS & 90.00 & 72.90 & -5.32 & 2.56 M \\ STR & 90.23 & 74.31 & -3.51 & 2.49 M \\ Global Magnitude & 90.00 & 75.20 & -2.42 & 2.56 M \\ DNW & 90.00 & 74.00 & -4.52 & 2.56 M \\ WoodFisher & 90.00 & 75.21 & -2.34 & 2.56 M \\
**CRITA** & 90.00 & **75.29** & **-2.23** & 2.56 M \\ \hline GMP & 95.00 & 70.59 & -7.95 & 1.28 M \\ Variational Dropout & 94.92 & 69.41 & -9.49 & 1.30 M \\ Variational Dropout & 94.94 & 71.81 & -6.36 & 1.30 M \\ RIGL + ERK & 95.00 & 70.00 & -8.85 & 1.28 M \\ DNW & 95.00 & 68.30 & -11.31 & 1.28 M \\ STR & 94.80 & 70.97 & -7.84 & 1.33 M \\ STR & 95.03 & 70.40 & -8.58 & 1.27 M \\ Global Magnitude & 95.00 & 71.79 & -6.78 & 1.28 M \\ WoodFisher & 95.00 & 72.12 & -6.35 & 1.28 M \\
**CRITA** & 95.00 & **73.46** & **-4.61** & 1.28 M \\ \hline GMP + LS & 98.00 & 57.90 & -24.50 & 0.51 M \\ Variational Dropout & 98.57 & 64.52 & -15.87 & 0.36 M \\ DNW & 98.00 & 58.20 & -24.42 & 0.51 M \\ STR & 98.05 & 61.46 & -20.19 & 0.50 M \\ STR & 97.78 & 62.84 & -18.40 & 0.57 M \\ Global Magnitude & 98.00 & 64.28 & -16.53 & 0.51 M \\ WoodFisher & 98.00 & 65.55 & -14.88 & 0.51 M \\
**CRITA** & 98.00 & **69.80** & **-9.36** & 0.51M \\ \hline \hline \end{tabular}
\end{table}
Table 3: Results of gradually pruning a ResNet50 network in the 90%, 95%, and 98% sparsity regimes, comparing CHITA to other state-of-the-art methods (Dense accuracy: 77.01%). We also include the relative drop in accuracy to account for different methods starting from different dense weights. CHITA numbers are averaged across two runs. Numbers for other baselines are taken from Singh and Alistarh (2020).

Figure 4: Comparison of test accuracy using CHITA++ with 15 stages, pruning a ResNet20 model to a 90% sparsity rate, under different sparsity schedules. Text around the point indicates the current sparsity level of the point.
free \(\ell_{0}\)-constrained regression formulation and combinatorial optimization techniques. Our single-stage methods demonstrate comparable results to existing methods while achieving a significant improvement in runtime and reducing memory usage. Furthermore, by building upon the single-stage methods, our multi-stage approach is capable of achieving even further improvements in model accuracy. Additionally, we have demonstrated that incorporating our pruning methods into existing gradual pruning frameworks results in sparse networks with state-of-the-art accuracy.

## Acknowledgements

This research is supported in part by grants from the Office of Naval Research (N000142112841 and N000142212665), and Google. We thank Shibal Ibrahim for helpful discussions. We also thank Thiago Serra and Yu Xin for sharing with us code from their CBS paper (Yu et al., 2022).
