Neural community pruning with combinatorial optimization – Google Analysis Weblog

Fashionable neural networks have achieved spectacular efficiency throughout a wide range of functions, comparable to language, mathematical reasoning, and vision. Nonetheless, these networks typically use massive architectures that require plenty of computational assets. This could make it impractical to serve such fashions to customers, particularly in resource-constrained environments like wearables and smartphones. A widely used approach to mitigate the inference prices of pre-trained networks is to prune them by eradicating a few of their weights, in a approach that doesn’t considerably have an effect on utility. In customary neural networks, every weight defines a connection between two neurons. So after weights are pruned, the enter will propagate by way of a smaller set of connections and thus requires much less computational assets.

Authentic community vs. a pruned community.

Pruning strategies will be utilized at completely different levels of the community’s coaching course of: put up, throughout, or earlier than coaching (i.e., instantly after weight initialization). On this put up, we concentrate on the post-training setting: given a pre-trained community, how can we decide which weights ought to be pruned? One in style methodology is magnitude pruning, which removes weights with the smallest magnitude. Whereas environment friendly, this methodology doesn’t straight contemplate the impact of eradicating weights on the community’s efficiency. One other in style paradigm is optimization-based pruning, which removes weights primarily based on how a lot their removing impacts the loss perform. Though conceptually interesting, most present optimization-based approaches appear to face a severe tradeoff between efficiency and computational necessities. Strategies that make crude approximations (e.g., assuming a diagonal Hessian matrix) can scale properly, however have comparatively low efficiency. Then again, whereas strategies that make fewer approximations are inclined to carry out higher, they look like a lot much less scalable.

In “Fast as CHITA: Neural Network Pruning with Combinatorial Optimization”, introduced at ICML 2023, we describe how we developed an optimization-based strategy for pruning pre-trained neural networks at scale. CHITA (which stands for “Combinatorial Hessian-free Iterative Thresholding Algorithm”) outperforms present pruning strategies when it comes to scalability and efficiency tradeoffs, and it does so by leveraging advances from a number of fields, together with high-dimensional statistics, combinatorial optimization, and neural community pruning. For instance, CHITA will be 20x to 1000x sooner than state-of-the-art strategies for pruning ResNet and improves accuracy by over 10% in lots of settings.

Overview of contributions

CHITA has two notable technical enhancements over in style strategies:

  • Environment friendly use of second-order info: Pruning strategies that use second-order info (i.e., referring to second derivatives) obtain the state-of-the-art in lots of settings. Within the literature, this info is usually utilized by computing the Hessian matrix or its inverse, an operation that may be very troublesome to scale as a result of the Hessian dimension is quadratic with respect to the variety of weights. By means of cautious reformulation, CHITA makes use of second-order info with out having to compute or retailer the Hessian matrix explicitly, thus permitting for extra scalability.
  • Combinatorial optimization: Common optimization-based strategies use a easy optimization method that prunes weights in isolation, i.e., when deciding to prune a sure weight they don’t bear in mind whether or not different weights have been pruned. This might result in pruning vital weights as a result of weights deemed unimportant in isolation could develop into vital when different weights are pruned. CHITA avoids this problem through the use of a extra superior, combinatorial optimization algorithm that takes under consideration how pruning one weight impacts others.

Within the sections beneath, we focus on CHITA’s pruning formulation and algorithms.

A computation-friendly pruning formulation

There are a lot of attainable pruning candidates, that are obtained by retaining solely a subset of the weights from the unique community. Let ok be a user-specified parameter that denotes the variety of weights to retain. Pruning will be naturally formulated as a best-subset selection (BSS) downside: amongst all attainable pruning candidates (i.e., subsets of weights) with solely ok weights retained, the candidate that has the smallest loss is chosen.

Pruning as a BSS downside: amongst all attainable pruning candidates with the identical whole variety of weights, the perfect candidate is outlined because the one with the least loss. This illustration exhibits 4 candidates, however this quantity is usually a lot bigger.

Fixing the pruning BSS downside on the unique loss perform is usually computationally intractable. Thus, just like earlier work, comparable to OBD and OBS, we approximate the loss with a quadratic function through the use of a second-order Taylor series, the place the Hessian is estimated with the empirical Fisher information matrix. Whereas gradients will be sometimes computed effectively, computing and storing the Hessian matrix is prohibitively costly as a consequence of its sheer dimension. Within the literature, it’s common to take care of this problem by making restrictive assumptions on the Hessian (e.g., diagonal matrix) and in addition on the algorithm (e.g., pruning weights in isolation).

CHITA makes use of an environment friendly reformulation of the pruning downside (BSS utilizing the quadratic loss) that avoids explicitly computing the Hessian matrix, whereas nonetheless utilizing all the knowledge from this matrix. That is made attainable by exploiting the low-rank construction of the empirical Fisher info matrix. This reformulation will be seen as a sparse linear regression downside, the place every regression coefficient corresponds to a sure weight within the neural community. After acquiring an answer to this regression downside, coefficients set to zero will correspond to weights that ought to be pruned. Our regression information matrix is (n x p), the place n is the batch (sub-sample) dimension and p is the variety of weights within the authentic community. Usually n << p, so storing and working with this information matrix is way more scalable than frequent pruning approaches that function with the (p x p) Hessian.

CHITA reformulates the quadratic loss approximation, which requires an costly Hessian matrix, as a linear regression (LR) downside. The LR’s information matrix is linear in p, which makes the reformulation extra scalable than the unique quadratic approximation.

Scalable optimization algorithms

CHITA reduces pruning to a linear regression downside underneath the next sparsity constraint: at most ok regression coefficients will be nonzero. To acquire an answer to this downside, we contemplate a modification of the well-known iterative hard thresholding (IHT) algorithm. IHT performs gradient descent the place after every replace the next post-processing step is carried out: all regression coefficients exterior the Prime-ok (i.e., the ok coefficients with the most important magnitude) are set to zero. IHT sometimes delivers an excellent answer to the issue, and it does so iteratively exploring completely different pruning candidates and collectively optimizing over the weights.

Because of the scale of the issue, customary IHT with fixed learning rate can endure from very sluggish convergence. For sooner convergence, we developed a brand new line-search methodology that exploits the issue construction to discover a appropriate studying fee, i.e., one which results in a sufficiently massive lower within the loss. We additionally employed a number of computational schemes to enhance CHITA’s effectivity and the standard of the second-order approximation, resulting in an improved model that we name CHITA++.


We evaluate CHITA’s run time and accuracy with a number of state-of-the-art pruning strategies utilizing completely different architectures, together with ResNet and MobileNet.

Run time: CHITA is way more scalable than comparable strategies that carry out joint optimization (versus pruning weights in isolation). For instance, CHITA’s speed-up can attain over 1000x when pruning ResNet.

Submit-pruning accuracy: Under, we evaluate the efficiency of CHITA and CHITA++ with magnitude pruning (MP), Woodfisher (WF), and Combinatorial Brain Surgeon (CBS), for pruning 70% of the mannequin weights. General, we see good enhancements from CHITA and CHITA++.

Submit-pruning accuracy of assorted strategies on ResNet20. Outcomes are reported for pruning 70% of the mannequin weights.
Submit-pruning accuracy of assorted strategies on MobileNet. Outcomes are reported for pruning 70% of the mannequin weights.

Subsequent, we report outcomes for pruning a bigger community: ResNet50 (on this community, among the strategies listed within the ResNet20 determine couldn’t scale). Right here we evaluate with magnitude pruning and M-FAC. The determine beneath exhibits that CHITA achieves higher take a look at accuracy for a variety of sparsity ranges.

Take a look at accuracy of pruned networks, obtained utilizing completely different strategies.

Conclusion, limitations, and future work

We introduced CHITA, an optimization-based strategy for pruning pre-trained neural networks. CHITA gives scalability and aggressive efficiency by effectively utilizing second-order info and drawing on concepts from combinatorial optimization and high-dimensional statistics.

CHITA is designed for unstructured pruning through which any weight will be eliminated. In idea, unstructured pruning can considerably scale back computational necessities. Nonetheless, realizing these reductions in observe requires particular software program (and presumably {hardware}) that assist sparse computations. In distinction, structured pruning, which removes entire buildings like neurons, could provide enhancements which might be simpler to achieve on general-purpose software program and {hardware}. It will be attention-grabbing to increase CHITA to structured pruning.


This work is a part of a analysis collaboration between Google and MIT. Because of Rahul Mazumder, Natalia Ponomareva, Wenyu Chen, Xiang Meng, Zhe Zhao, and Sergei Vassilvitskii for his or her assist in getting ready this put up and the paper. Additionally because of John Guilyard for creating the graphics on this put up.

Leave a Reply

Your email address will not be published. Required fields are marked *