Title: CAP:  Correlation-Aware Pruning for Highly-Accurate Sparse Vision Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural network pruning

vision transformer

sparsity

model compression
Abstract: Driven by significant improvements in architectural design and training pipelines, computer vision
has recently experienced dramatic progress in terms of accuracy on classic benchmarks such as ImageNet. 
These highly-accurate models are challenging to deploy,  as they appear harder to compress using standard techniques such as pruning. 
We address this issue by introducing the Correlation Aware Pruner (CAP), 
a new unstructured pruning framework which significantly pushes the compressibility limits for state-of-the-art architectures.
Our method is based on two technical advancements: a new theoretically-justified pruner, which can handle complex weight correlations accurately and efficiently during the pruning process itself,  and an efficient finetuning procedure for post-compression recovery. 
We validate our approach via extensive experiments on several modern vision models such as Vision Transformers (ViT), 
modern CNNs, and ViT-CNN hybrids, showing for the first time that these can be 
pruned to high sparsity levels (e.g. $\geq 75$%) with low impact on accuracy ($\leq 1$% relative drop). 
Our approach is also compatible with structured pruning and quantization, and can lead to practical speedups of 1.5 to 2.4x without accuracy loss. To further showcase CAP's accuracy and scalability, we use it to show for the first time that extremely-accurate large vision models, trained via self-supervised techniques,  can also be pruned to moderate sparsities, with negligible accuracy loss.

CAP: Correlation-Aware Pruning for Highly-Accurate Sparse Vision Models

Abstract

Driven by significant improvements in architectural design and training pipelines, computer vision has recently experienced dramatic progress in terms of accuracy on classic benchmarks such as ImageNet. These highly-accurate models are challenging to deploy, as they appear harder to compress using standard techniques such as pruning. We address this issue by introducing the Correlation Aware Pruner (CAP) , a new unstructured pruning framework which significantly pushes the compressibility limits for state-of-the-art architectures. Our method is based on two technical advancements: a new theoretically-justified pruner , which can handle complex weight correlations accurately and efficiently during the pruning process itself, and an efficient finetuning procedure for post-compression recovery. We validate our approach via extensive experiments on several modern vision models such as Vision Transformers (ViT), modern CNNs, and ViT-CNN hybrids, showing for the first time that these can be pruned to high sparsity levels (e.g. ≥ 75% ) with low impact on accuracy ( ≤ 1% relative drop). Our approach is also compatible with structured pruning and quantization, and can lead to practical speedups of 1.5 to 2.4x without accuracy loss. To further showcase CAP's accuracy and scalability, we use it to show for the first time that extremely-accurate large vision models, trained via self-supervised techniques, can also be pruned to moderate sparsities, with negligible accuracy loss 1 .

1 Introduction

Computer vision has seen impressive progress recently via a new generation of architectures motivated by the Vision Transformer (ViT) approach [9, 43] and its extensions, e.g. [1, 28, 46], accompanied by more advanced data augmentation and training approaches [43, 48]. Next-generation vision models such as ViT [9, 43] and ConvNext [49, 29] achieve breakthrough performance across vision tasks, despite encoding fewer inductive biases. This comes at the cost of very large computational and parameter budgets, both for training and deployment. Thus, there is a clear need to reduce these costs via compression , enabling deployment in resource-constrained settings.

Yet, the general consensus from the literature is that next-generation models are harder to compress while preserving accuracy, relative to their classic convolutional counterparts [6, 40, 35, 20]. For example, if current techniques can compress the classic ResNet50 model [18] to 80-90% unstructured sparsity with negligible loss of accuracy [33, 44], the best currently-known results for similarlyaccurate ViT models only reach at most 50% sparsity while maintaining dense accuracy [6]. Our investigation of this phenomenon, detailed in the paper, shows that this occurs for two reasons:

Hardness of Pruning : As illustrated in Figure 2, existing pruning approaches, based on magnitude [16], first-order [38] or second-order information [39] drop very significant accuracy per pruning step for next-generation architectures such as ViT and ConvNext.

Expensive Recovery : At the same time, recovering accuracy via fine-tuning [15] is itself computationally-challenging, as next-generation architectures are notoriously hard to train and finetune [41, 43].

Therefore, it is natural to ask whether this 'lack of compressibility' is inherent for next-generation vision models, and whether their increased accuracy comes at the price of higher deployment costs.

Contributions. In this paper, we resolve this question by proposing a new highly-accurate correlationaware pruner (CAP) , which shows that high sparsity can be achieved across next-generation architectures such as ViT [9, 43], ConvNext [29, 49], and augmented ResNets [48], and that this can be done in a computationally-efficient manner . Furthermore, CAP is compatible with other forms of compression such as token dropping and quantization, and can lead to significant inference speedups, at little or no accuracy loss.

Our primary motivation is to resolve the hardness of pruning modern, highly-accurate vision models. A key weakness of existing state-of-the-art approaches [10, 38, 6, 33, 39, 14], is that they do not directly take into account weight correlations: At each pruning step, a saliency score is computed per weight, e.g., the magnitude, weights are ranked by this score, and a large subset is chosen to be removed and/or re-introduced. However, this does not take into account the fact that removed weights may themselves be correlated : for example, a pair of linearly-dependent rows in a layer's weight matrix may seem individually easy to remove , since they are mutually-redundant, but removing both at a step may lead to a large accuracy drop. This phenomenon is especially-relevant for Transformer-based architectures, which encode more complex inter-weight correlations, relative to their convolutional counterparts, in which correlations tend to be localized [39].

We circumvent this issue via a novel formulation of the constrained optimization problem corresponding to choosing the optimal weights to prune at a step [17], while taking correlations into account. We prove formally that this task can be reduced to finding the set of sparse weights which best preserve the correlation between the dense weights and their gradients, on a given set of training samples. This allows us to efficiently solve the 'optimal pruning with correlations' problem, via a fast approximate solver for the above constrained optimization problem.

The CAP algorithm outperforms all other known pruners by a considerable margin (see Figure 6). In turn, this precision lends greater flexibility in gradual pruning , and enables us to build a computationally-efficient approach to compress next-generation vision models, addressing our second motivating challenge. Specifically, our gradual pruning approach provides a simple and general recipe combining data-efficient augmentation and regularization [43, 41] with theoretically-justified learning rate rewinding [36, 23], leading to state-of-the-art sparsity-vs-accuracy trade-offs.

For instance, experiments on the standard ImageNet-1K benchmark [37] show for the first time that ViT models can attain high sparsity levels without significant accuracy impact: specifically, we can achieve 75-80% sparsity with relatively minor ( < 1%) accuracy loss, and 90% sparsity with moderate loss. In turn, this sparsity leads to computational speedups of more than 2x. Our approach extends to highly-accurate pruning of large ViTs trained by self-supervised pretraining, but also to other modern models, such as ConvNext [29, 49] or highly-accurate ResNets [48].

In sum, our results show that next-generation highly-accurate vision architectures are still highlycompressible, but may also require next-generation pruning approaches.

Related work. Next-generation vision architectures such as Vision Transformers (ViTs) [9] and ConvNext [29, 49] have set new accuracy benchmarks, but are known to require careful tuning in terms of both augmentation and training hyper-parameters. Identifying efficient recipes is an active research topic in itself [43, 41]. We propose simple and general recipes for fine-tuning such models, which should be useful to the community. Several prior works have investigated ViT compression, but focus on structured pruning, such as removing tokens [54, 22, 50, 31, 40, 35, 20]. Our experiments show that structured approaches are orthogonal to CAP, which can be applied in conjunction to structured compression, to obtain further gains.

The only existing prior work on unstructured ViT pruning is SViTE [6], which performed careful customization of the RigL pruning method [10] to the special case of ViT models. We also present

results relative to other methods, such as tuned magnitude pruning, the best first-order and secondorder pruners [39, 38, 14, 23] and AC/DC pruning [33]. (These methods are known to outperform all other prior methods.) CAP improves upon existing methods across almost all benchmarks, by large margins at high sparsity.

Research on accurate pruning using second-order information was initiated by LeCun et al. [26], and has recently garnered significant attention [8, 45, 39, 52]. This approach can lead to good results for both gradual pruning [39] and one-shot (post-training) compression [12]. Existing such pruners are not correlation-aware, and are outperformed by CAP across all of our experiments.

2 Background and Problem Setup

The pruning problem assumes a fixed model architecture with weights w ∈ R d ( d is the total number of parameters), and aims to find a configuration of weights with as many zeros as possible while preserving the performance of the original dense model. Gradual pruning, e.g. [19], usually starts from an accurate dense model, and progressively removes weights by setting them to zero, followed by fine-tuning phases.

Weight Saliency. The pruning step usually relies on proxies for weight importance, defined according to certain criteria. For instance, the weight magnitude is arguably the most popular criterion, e.g. [16, 53, 15]. Specifically, given model w ∈ R d , the saliency of each weight is its absolute value (the magnitude) ρ j = | w j | for j ∈ { 1 , 2 , . . . , d } ; weights with the smallest scores are pruned away. Gradual magnitude pruning is usually a strong baseline across most models and settings. Many other criteria exist, such as gradient magnitude [10] or 'rates of change' in the weights [38].

The Optimal Brain Surgeon (OBS). LeCun et al. [26] and Hassibi et al. [17] obtained weight saliency scores by leveraging (approximate) second-order information about the loss, starting from the Taylor approximation of the loss L in the vicinity of the dense model parameters w ∗ . Assuming that w ∗ is close to the optimum (hence ∇L ( w ∗ ) ≃ 0 ), one seeks a binary mask M (with elements ∈ { 0 , 1 } ) and new values for the remaining weights w M , such that the resulting increase in loss is minimal. A standard approach to approximate the loss increase is to expand the loss function up to the second order in model weights:

where H L ( w ∗ ) is the Hessian of the model at w ∗ , and w M represents weights after the pruning step. In this setup, [26] and [17] showed that the 'optimal' weight to remove, incurring the least loss, and the update to the remaining weights, can be determined via a closed-form solution to the above inverse problem. Specifically, the saliency score ρ i for i th weight and the optimal weight update δ w for the remaining weights after elimination of the i th weight are as follows:

where e i is the i th basis vector. Theoretically, the procedure would have to be executed one-weightat-a-time, recomputing the Hessian after each step. In practice, this procedure suffers from a number of practical constraints. The first is that direct Hessian-inverse computation is computationallyinfeasible for modern DNNs, due to its quadratic-in-dimension storage and computational costs. This has led to significant recent work on efficient second-order approximations for pruning and quantization [8, 45, 52].

WoodFisher and the Optimal BERT Surgeon. The empirical Fisher approximation [2] is a classic way of side-stepping some of the above constraints, and can be formally-stated as follows:

where ∇L i ( w ∗ ) ∈ R d is a gradient computed on a sample of data, λ > 0 is a dampening constant needed for stability, and N is the total number of gradients used for approximation. Note that the resulting matrix is positive definite.

The memory required to store the empirical Fisher matrix is still quadratic in d , the number of parameters. Singh and Alistarh [39] investigated a diagonal block-wise approximation with a predefined block size B , which reduces storage cost from O ( d 2 ) to O ( Bd ) , and showed that this approach can lead to strong results when pruning CNNs. Kurtic et al. [23] proposed a formula for pruning fixed groups/patterns (e.g., 4 consecutive weights), together with a set of non-trivial optimizations to efficiently compute the Fisher block inverses, which allowed them to scale the approach for the first time to large language models.

A second obvious limitation of the OBS framework is that applying the procedure and recomputing the Hessian one weight at a time is prohibitively expensive, so one usually prunes multiple weights at once. Assuming we are searching for the set of weights Q whose removal would lead to minimal loss increase after pruning, we get the following constrained optimization problem:

where E Q ∈ R | Q |× d is a matrix of basis vectors for each weight in Q . The corresponding saliency score for the group of weights Q and the update δ w ∗ Q of remaining weights are [23]:

However, an exhaustive search over all subsets of size | Q | from d elements requires ( d | Q | ) evaluations, which makes it prohibitively expensive for | Q | > 1 . For unstructured pruning, virtually all known techniques, e.g. [39, 14], ignore correlations between weights. Similarly, for group pruning, Kurtic et al. [23] ignore correlations between groups. Despite these approximations, both approaches yield state-of-the-art results in their respective setups. As we will demonstrate later, our CAP method improves upon these approximations by reformulating this problem and proposing a correlation-aware solution that is fast and memory-efficient even for models with ∼ 100M parameters.

3 The CAP Pruning Framework

We introduce new techniques to address both the hardness of pruning modern vision architectures, and their high computational cost for fine-tuning: we introduce a new state-of-the-art one-shot pruner, which is complemented with a simple and general framework for data-efficient fine-tuning.

3.1 Ingredient 1: Efficient Correlation-Aware Pruning

Our aim is to solve the pruning problem stated in the previous section: given a weight pruning target k , find the optimal set of weights Q to be pruned, such that | Q | = k and the loss increase is minimized. Exactly solving for the optimal Q is an NP-hard problem [3], so we will investigate an iterative greedy method for selecting these weights, similar to the ideal version of the OBS framework discussed above. Importantly, our method properly considers weight correlations , which turn out to be important since, as demonstrated in Appendix Q, the empirical Fisher has an apparent non-diagonal structure, while being fast and space-efficient . In turn, this leads to significant improvements over other pruners, especially in the context of vision transformers.

Formally, a correlation-aware greedy weight selection approach would perform pruning steps iteratively, as follows. Given a set of already-selected weights Q , initially ∅ , we always add to Q the weight q which has minimal joint saliency ρ Q ∪{ q } , repeating until the size of Q equals the pruning target k . The fact that we add weights to the set one-by-one allows us to take into account correlations between pruned weights. However, a naive implementation of this scheme, which simply recomputes saliency at each step, would be prohibitively expensive, since it requires O ( kd ) evaluations of ρ Q , each of which involves an O ( B 3 ) matrix inversion, where B is the Fisher block size.

Disentangling Correlations. The centerpiece of our approach is a reformulation of the OBS multi-weight pruning problem in Equation 5 which will allow us to take correlations into account, while being practically-efficient. Specifically, we now show that, when using the empirical Fisher approximation, the problem of finding the optimal set of weights Q to be removed, while taking correlations into account, is equivalent to the problem of finding the set of sparse weights which best preserve the original correlation between the dense weights w ∗ and the gradients ∇L i ( w ∗ ) on an

fixed set of samples i ∈ S . Formally, this result, whose proof we provide in Appendix O, is stated as follows.

Theorem 3.1. Let S be a set with m samples, and let ∇L 1 ( w ∗ ) , . . . , ∇L m ( w ∗ ) be a set of their gradients, with corresponding inverse empirical Fisher matrix F -1 ( w ∗ ) . Assume a sparsification target of k weights from w ∗ . Then, a sparse minimizer for the constrained squared error problem

s.t. w ′ has at least k zeros, is also a solution to the problem of minimizing the Fisher-based group-OBS metric

A Fast Solver. The formulation in Equation (6) reduces pruning to a sparse regression problem, where the 'input' is given by gradients over calibration samples. A related problem arises in the context of one-shot (post-training) pruning [21, 12], where the authors solve a sparse ℓ 2 -fitting problem, but sparse weights are determined relative to the layer inputs rather than the layer gradients . Specifically, the OBC solver [12] utilizes that, due to the quadratic loss, the per-row Hessians are independent of both the weights and each other. Thus the solver processes matrix rows sequentially, greedily eliminating weights, one-by-one, in increasing order of squared error while always updating all remaining weights to compensate for weight removal as much as possible. This essentially implements the OBS selection and update in Equation 2 exactly , assuming layer-wise ℓ 2 loss. We extend this strategy to efficiently implement our new Fisher-based greedy weight-subset selection.

A direct application of the OBC approach to remove Θ( d ) weights would have O ( d 3 ) runtime, where d is the layer dimension, as the Θ( d 2 ) -time selection + update process is repeated Θ( d ) times. By utilizing the block-diagonal structure of the Fisher matrix with block size B during the update after each weight elimination, this complexity can be reduced to O ( d · max ( d, B 2 )) , which is however still too slow as B is typically much smaller than d . Instead, we proceed differently: we apply the sparse regression solver to each individual Fisher block separately and globally merge results only in the very end. This allows us to efficiently simulate global weight saliency calculation and weight updates following Equations 5 and 7, in a block-Fisher formulation. The full algorithm requires O ( d · B 2 ) runtime and O ( d · B ) space and is detailed in Appendix 1--an efficient implementation is also provided in the Supplementary.

The key difference of our approach relative to WoodFisher is that updates are continuously executed during subset-selection and we are thus explicitly considering the correlations captured by the offdiagonal Fisher elements. When working with small block sizes, our method is very fast and has practically no overhead over existing Fisher-based OBS approaches, while yielding significantly improved one-shot pruning results (see e.g. Figure 2 and the accompanying discussion).

An Faster Variant for Massive Models. In Appendix B, we present a version of CAP which scales efficiently to models with more than 1 billion parameters, by leveraging additional approximations in the Fisher matrix structure and by relaxing the pruning order.

3.2 Ingredient 2: Fine-tuning and Pruning Procedure

To achieve best performance, modern training procedures involve longer training schedules together with a careful choice of hyperparameters, since these are known to have a major impact on convergence and accuracy [43, 41, 48]. We found the same to apply to post-pruning accuracy recovery, which is key in gradual pruning; below, we describe the main ingredients to obtaining highly-accurate fine-tuning schedules as part of our method.

Learning Rate Schedule. First, to achieve good performance during gradual pruning, the learning rate (LR) schedule is crucial. Specifically, we propose to use a cyclic linear schedule:

where % x means taking the remainder after integer division by x . We chose a linear decay for simplicity; we obtained similar results for other functions (e.g., cubic decay). By contrast, as we

illustrate in Figure 1, the cyclic nature of the schedule is key for accurate pruning, as well as for efficient sparsity sweeps (see below).

Theoretical Justification. Specifically, this choice is justified theoretically by tying back to the original assumptions of the OBS framework: for Equation 1 to hold, the pruned model should be welloptimized (i.e. have small gradients) at the point when pruning is performed. Moreover, right after the pruning step, having a larger value of the learning rate is useful since it gives the model a chance to recover from the sub-optimal point induced via pruning. We note that this learning rate schedule is different from prior work on pruning, which typically uses a single decay cycle [25, 39, 33], or dynamic learning rate rewinding, e.g. [11, 36].

0

50

100

150

200

250

300

Epoch

0.0000

0.0001

0.0002

0.0003

0.0004

0.0005

Learning rate

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.52

0.65

0.77

Sparsity

0

5

10

15

20

Relative acc drop (%)

cyclic

cosine

Regularization and Augmentation. For augmentation, similarly to [6], we adopt known bestpractices from the literature: smaller models such as DeiT-Tiny benefit from lower levels of data augmentation during fine-tuning as in [41], whereas larger models such as DeiT-Base behave best with more complex augmentation and regularization [43]. This is intuitive, since fine-tuning sparsified small models with high augmentation may exceed model capacity, rendering the optimization process unstable. We provide detailed parameter values and ablations in Appendix C.

Augmentation for the Empirical Fisher. The choice of augmentation is of great importance not only for the model training but for the accurate estimate of the Empirical Fisher as well. We observed that without proper augmentation the sparse solution obtained overfits to the training data even when finetuned with the same augmentation and regularization setup. We provide details in Appendix P.

Efficient Sparsity Sweeps. We propose a simple iterative pruning framework, which takes a set of target sparsity configurations and produces models which match these configurations in a single run . Specifically, we start from a standard gradual pruning setup, which prunes in a sequence of steps of increasing sparsity, followed by sparse fine-tuning. We then set the intermediate values in such a way that all intermediate target sparsity levels are achieved. For example, if one wishes to obtain checkpoints with sparsity levels 40% , 50% , 75% , 90% , one can set the lowest sparsity level on the gradual pruning schedule to 40% , the highest sparsity level to 90% , and 50% , 75% as intermediate points. Between any two such pruning steps, we apply the cyclic retraining schedule above, which ensures that all intermediate points are sufficiently optimized.

We emphasize that the accurate CAP pruner is key for efficient pruning: virtually all previous high-accuracy pruning methods in this setting, e.g. [25, 39, 6] redo the entire training run for each sparsity target. In our experiments, we also examine the impact of additional fine-tuning applied to each checkpoint, and show that it induces small-but-consistent improvements.

4 Experimental Setup and Results

Setup and Goals. We consider the ImageNet [37] image classification benchmark, and aim to examine how sparsity impacts accuracy for different model variants. We consider three scenarios: one-shot, single-step pruning of a pretrained model, where performance is clearly tied to the quality of the second-order approximation, one-shot + fine-tuning , in which we follow one-shot pruning

by a short period of fine-tuning, and, finally, iterative gradual pruning , where one applies pruning periodically, with some retraining interval, gradually increasing sparsity.

4.1 One-shot pruning, without and with fine-tuning

We start by examining the quality of existing one-shot pruners relative to CAP. We compare against carefully-tuned variants of Magnitude Pruning (Magn), First-Order Gradient times Weight (GrW) [38], and the SOTA second-order methods WoodFisher [39, 23] and M-FAC [14]. Our tuning process, optimal hyper-parameter choices, and ablations are detailed in Appendices J and K. WF-1 represents the diagonal approximation of the Hessian proposed by the original OBD [26], scaled via the WoodFisher implementation. For Magnitude and CAP we present results for both uniform layer-wise sparsity and global sparsity. For all other methods, we present results for global sparsity, which yields better results, as the methods can adjust sparsity globally. We only investigate global sparsity for the other methods in Figure 2.

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Val accuracy

DeiT-Tiny

WF-1 (G)

M-FAC (G)

WF (G)

CAP (U)

0.0

0.2

0.4

0.6

0.8

1.0

-0.1

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

50% sparsity

Dense

WF

CAP

1.25

1.30

1.35

1.40

1.45

1.50

1.55

Val loss

The full comparison is presented in Figure 2 for DeiT-Tiny. Notice that all first and second-order methods outperform Magnitude, and that all 2nd-order methods are better than the 1st order saliency. WoodFisher with small block size is better than both the diagonal approximation and large-block M-FAC. This suggests that it is beneficial to take weight correlations into account, but attempting to incorporate dependencies between large groups of weights may lead to noisy estimates detrimental to performance. We also present a comparison on other models (ViT and ResNet) in Appendix E.

CAP outperforms all other methods, by a large margin. Remarkably, the gap is so large that CAP with uniform sparsity still outperforms WoodFisher with global sparsity , which can re-distribute sparsity across layers. The computational cost of WF is approximately the same as for CAP: CAP pruning step on DeiT-Small takes 23 minutes, compared to 20 minutes for WF. (The majority of the cost for both methods comes from the collection of gradients, not Hessian estimation.)

To investigate the performance difference between CAP and WF, we pruned several variants of ViT models to 50% and compared the mask overlap between WF and CAP via their IoU (intersection over union). We observe that the sparsification masks differ significantly: the IoU between the WF and CAP mask is 82.7%, 80.5%, 73.3% for DeiT-Tiny, DeiT-Small, and DeiT-Base, respectively, suggesting that the weight correlations taken into account by CAP lead to significantly different pruning decisions. The same trend is visible in the validation loss surface, projected on the plane in weight space, for the dense model and the 50% sparse models via WF and CAP. (We adopt the approach from [32] for loss surface visualization. Specifically, the unit vector in the horizontal direction is chosen to be w WF -w ∗ and the vertical direction is defined by the component of w CAP -w ∗ orthogonal to the horizontal axis. Above w ∗ , w WF , w CAP are the weights of the dense model and sparse solutions found by WF and CAP, respectively.) Figure 3 shows that CAP chooses a significantly less steep direction in the loss basin compared to WF.

Compressing highly-accurate models. Modern pretraining methods [7, 5, 34, 51] in conjunction with large vision backbones achieve extremely high accuracy on standard benchmarks such as ImageNet-1K. We leverage the scalability and accuracy of CAP to investigate, for the first time, sparsity in such highly-accurate models, in the absence of finetuning. Specifically, we start from the

ConvNext-Large checkpoint pretrained via CLIP on the LAION-2B dataset and finetuned sequentially on ImageNet-21k and ImageNet-1k [47].

In Table 1, we compare our method with WoodFisher (WF) and Global Magnitude (GM) for oneshot pruning via several sparsity targets. Results show that 1) CAP can induce 50% sparsity with relatively low (0.3%) accuracy loss, and 60% with less than 1% accuracy drop; 2) CAP significantly outperforms other methods. In addition, we find that these results can be improved with a limited amount of fine-tuning; please see Appendix E for full results.

4.2 Gradual Pruning Results

Finally, we execute gradual pruning with the sparsity schedule, augmentation choices, and cyclic linear learning-rate scheduler discussed above. The whole gradual pruning procedure lasts for 300 epochs, as in [43]. We aim to obtain accurate sparse checkpoints for 50% , 60% , 75% , 80% , and 90% sparsity. For this, we prune to 40% in the initial step, and increment sparsity every 20 epochs, until reaching 90% , with fine-tuning in between. (See Appendix J for full results and ablations.) We select the accuracy of intermediate models which match the target sparsities; to examine the impact of fine-tuning, we trained each of the resulting sparse checkpoints for an additional 100 epochs, marked with ( ↑ ). We compare with global magnitude (GM) following the same schedule as CAP, as well as the state-of-the-art SViTE [6] paper, which trains the sparse model from scratch using a variant of RigL [10], but for a total of 600 epochs. The results are in Table 2.

For DeiT-Tiny and 50% sparsity, we achieve significant improvements upon SViTe, and even manage to improve test accuracy relative to the dense model. We believe this is due to the choice of augmentation during fine-tuning, and possibly due to regularizing effects of sparsity. At 75-80%, we recover the dense model accuracy. We observe a significant accuracy drop only at 90%. GM pruning also benefits from the choices made in our schedule, outperforming SViTE at 50% sparsity; yet, there are significant gaps in favor of CAP at higher sparsities, as expected.

On the 4x larger DeiT-Small model, SViTE performs remarkably well at 50% sparsity (79.7%), almost matching the dense model, but CAP outperforms it slightly after fine-tuning (79.9%). In terms

of total training budget, SViTE uses 600 epochs to produce each model (and so, the 50%-sparse one as well), whereas we use a total of 40 epochs for gradual pruning to 50% + initial fine-tuning, and 100 additional epochs for sparse model fine-tuning. Even if we take into account the original 300 epochs for training the publicly-available dense DeiT checkpoint [43], our approach is significantly more efficient (440 vs. 600 epochs). The cost savings compound across sparse models, since we obtain all our pruned models from the same end-to-end gradual run. At 75% sparsity, CAP drops ∼ 1% of accuracy relative to dense post-finetuning, with a significant gap of 1% Top-1 relative to GM, and 2% Top-1 relative to SViTE. The trend continues for higher sparsities, where we note a remarkable gap of 5 . 7% Top-1 vs SViTE at 90% sparsity. We obtain similar numbers for DeiT-Base in Table 3; generally, we achieve ≥ 99% recovery at ≥ 75% sparsity, showing for the first time that ViT models can be pruned to such sparsities with marginal accuracy loss.

0

50

100

150

200

250

Throughput (img/s)

50

55

60

65

70

75

80

Val accuracy

DeiT-B  (SViTE)

DeiT-B  (CAP)

DeiT-S  (SViTE)

DeiT-S  (CAP)

DeiT-T  (SViTE)

DeiT-T  (CAP)

Sparse Speedups. The results in Figure 4 examined the speedups obtained by CAP from unstructured sparsity for {50%, 75%, 90%}-sparse ViT-family (base, small, tiny) models, when executed on a sparsity-aware CPU inference engine [24]. Specifically, we executed the models from Table 2 using 4 cores of an Intel(R) Xeon(R) Gold 6238R CPU, at batch size 64. We find it interesting that sparse ViTs build an almost-contiguous Pareto frontier from 82% to 68% Top-1 accuracy (Y axis), with a 25x span in throughput (from 11 imgs/second to 260 imgs/second, X axis). Notably, the DeiT-Tiny model obtains a speedup of 2.4x without any accuracy loss, while Base and Small ViTs show 1.5x speedups with minimal loss of accuracy. Thus, these results show that unstructured sparsity can be a very promising approach to speeding up ViTs. We note that competing methods (e.g. SViTE) would provide similar speedups at the same sparsity levels, but significantly lower accuracy, as shown in the figure, as well as in Table 2.

Additional Results and Details. Due to space constraints, full ablations and several additional experiments have been deferred to the Appendix. We present the full pseudocode for CAP in Appendix A, and a faster approximate version of CAP in Appendix B. We describe all parameters and perform additional ablations in Appendices C and J. We present results for one-shot pruning across several other models, as well as with short amounts of fine-tuning, in Appendix E. We present gradual pruning results for several other vision models (EfficientFormer, ResNet50D, and EfficientNetV2) in Appendix F and demonstrate that one can compress them to high sparsity as well, using our approach. In Appendix H we show that CAP comes with little additional computational cost compared to WoodFisher [23]. In Appendix I we show that CAP can also be applied in conjunction with other types of compression, in particular token pruning [35], quantization-aware training, and semi-structured (2:4) sparsity. In Appendix M, we show that CAP also outperforms AC/DC pruning [33], whereas Appendix N contains results for the DeTR detection model. In Appendix G we study the scaling behavior of sparsity solvers with respect to the ConvNext2 model family.

5 Discussion

We presented a new correlation-aware pruner called CAP, which sets a new state-of-the-art sparsityaccuracy trade-off. We have shown for the first time that ViT variants can support significant weight pruning ( ≥ 75% ) at relatively minor accuracy loss ( ≤ 1% ), inducing compression trade-offs that similar to those of CNNs, and that CLIP-pretrained, highly-accurate models can also support sparsity accurately. Thus, despite weaker inductive biases, next-generation vision models do not require over-parametrization, and can be competitive with CNNs in terms of accuracy-per-parameter. Our approach is complementary to other compression approaches, leading to significant speedups.

References

Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Crosscovariance image transformers. Advances in neural information processing systems , 34, 2021.

Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation , 10(2):251276, 1998.

Thomas Blumensath and Mike E Davies. Iterative thresholding for sparse approximations. Journal of Fourier Analysis and Applications , 14(5-6):629-654, 2008.

Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision , pages 213-229. Springer, 2020.

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers, 2021.

Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chasing sparsity in vision transformers: An end-to-end exploration. Advances in Neural Information Processing Systems , 34:19974-19988, 2021.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, 2020.

Xin Dong, Shangyu Chen, and Sinno Jialin Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. In Conference on Neural Information Processing Systems (NeurIPS) , 2017.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2020.

Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning , pages 2943-2952. PMLR, 2020.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv:1903.01611 , 2019.

Elias Frantar and Dan Alistarh. Optimal Brain Compression: A framework for accurate post-training quantization and pruning. arXiv preprint arXiv:2208.11580 , 2022.

Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot, 2023.

Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of second-order information. Advances in Neural Information Processing Systems , 34, 2021.

Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574 , 2019.

Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00149 , 2015.

Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE international conference on neural networks , pages 293-299. IEEE, 1993.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

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. Journal of Machine Learning Research , 22(241):1-124, 2021.

Zejiang Hou and Sun-Yuan Kung. Multi-dimensional model compression of vision transformer. In 2022 IEEE International Conference on Multimedia and Expo (ICME) , pages 01-06. IEEE, 2022.

Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Accelerated sparse neural training: A provable and efficient method to find n: m transposable masks. Advances in Neural Information Processing Systems , 34:21099-21111, 2021.

Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers, 2021.

Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeon: Scalable and accurate second-order pruning for large language models. arXiv preprint arXiv:2203.07259 , 2022.

Mark Kurtz, Justin Kopinsky, Rati Gelashvili, Alexander Matveev, John Carr, Michael Goin, William Leiserson, Sage Moore, Bill Nell, Nir Shavit, and Dan Alistarh. Inducing and exploiting activation sparsity for fast inference on deep neural networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5533-5543, Virtual, 13-18 Jul 2020. PMLR.

Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference on Machine Learning , pages 5544-5555. PMLR, 2020.

Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems , 2, 1989.

Yanyu Li, Geng Yuan, Yang Wen, Eric Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, and Jian Ren. Efficientformer: Vision transformers at mobilenet speed. arXiv preprint arXiv:2206.01191 , 2022.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 10012-10022, 2021.

Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11976-11986, 2022.

Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks, 2021.

Bowen Pan, Rameswar Panda, Yifan Jiang, Zhangyang Wang, Rogerio Feris, and Aude Oliva. Ia-red 2 : Interpretability-aware redundancy reduction for vision transformers. Advances in Neural Information Processing Systems , 34:24898-24911, 2021.

Mansheej Paul, Feng Chen, Brett W. Larsen, Jonathan Frankle, Surya Ganguli, and Gintare Karolina Dziugaite. Unmasking the lottery ticket hypothesis: What's encoded in a winning ticket's mask? In The Eleventh International Conference on Learning Representations , 2023.

Alexandra Peste, Eugenia Iofinova, Adrian Vladu, and Dan Alistarh. Ac/dc: Alternating compressed/decompressed training of deep neural networks. Advances in Neural Information Processing Systems , 34:8557-8570, 2021.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision, 2021.

Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems , 34:13937-13949, 2021.

Alex Renda, Jonathan Frankle, and Michael Carbin. Comparing rewinding and fine-tuning in neural network pruning. arXiv preprint arXiv:2003.02389 , 2020.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV) , 115(3):211-252, 2015.

Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems , 33:20378-20389, 2020.

Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. Advances in Neural Information Processing Systems , 33:18098-18109, 2020.

Zhuoran Song, Yihong Xu, Zhezhi He, Li Jiang, Naifeng Jing, and Xiaoyao Liang. Cp-vit: Cascade vision transformer pruning via progressive sparsity prediction, 2022.

Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers, 2021.

Mingxing Tan and Quoc Le. Efficientnetv2: Smaller models and faster training. In International Conference on Machine Learning , pages 10096-10106. PMLR, 2021.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning , pages 10347-10357. PMLR, 2021.

Antoine Vanderschueren and Christophe De Vleeschouwer. Are straight-through gradients and soft-thresholding all you need for sparse training? In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 3808-3817, 2023.

Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang. Eigendamage: Structured pruning in the kronecker-factored eigenbasis. In International Conference on Machine Learning , pages 6566-6575. PMLR, 2019.

Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 568-578, 2021.

Ross Wightman. Pytorch image models. https://github.com/rwightman/ pytorch-image-models , 2019.

Ross Wightman, Hugo Touvron, and Hervé Jégou. Resnet strikes back: An improved training procedure in timm. arXiv preprint arXiv:2110.00476 , 2021.

Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie. Convnext v2: Co-designing and scaling convnets with masked autoencoders, 2023.

Yifan Xu, Zhijie Zhang, Mengdan Zhang, Kekai Sheng, Ke Li, Weiming Dong, Liqing Zhang, Changsheng Xu, and Xing Sun. Evo-vit: Slow-fast token evolution for dynamic vision transformer, 2021.

Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models, 2022.

Shixing Yu, Zhewei Yao, Amir Gholami, Zhen Dong, Sehoon Kim, Michael W Mahoney, and Kurt Keutzer. Hessian-aware pruning and optimal neural implant. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 3880-3891, 2022.

Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878 , 2017.

Mingjian Zhu, Yehui Tang, and Kai Han. Vision transformer pruning, 2021.

A CAP algorithm description

The section below illustrates in the CAP pruning algorithm step-by-step. Prunable model weights R d are partitioned into blocks of fixed size B . Below ρ ( B ) i denotes the saliency scores for weight i th inside a block it belongs to and ρ i is the score across the whole model. The steps of the algorithm are listed below:

Algorithm 1 CAP pruning algorithm

B FastCAP: A faster version of the algorithm

CAP can scale to models with < 1B parameters, covering most standard vision models. Its key computational and memory cost comes because of computing and maintaining the inverse Fisher approximation, similarly to WoodFisher/Optimal BERT Surgeon.

To address the scalability issue we present a much more scalable approximate variant of CAP, called FastCAP, which is inspired by the recent SparseGPT scalable pruner [13]. Specifically, FastCAP reduces compute and memory cost by ∼ 100x by leveraging a cheaper Fisher approximation and relaxing the optimal greedy order of pruning. FastCAP can be scaled to billion-parameter vision models, i.e. we can prune ViT-Giant in a few minutes, while achieving reasonable performance in one shot-setting for 50-70% sparsity.

The main two components of the CAP/FastCAP algorithm are:

Empirical Fisher estimate. The original CAP has specific Fisher block for each output channel, whereas the FastCAP version averages blocks across the output channel dimension. This is a further approximation, but significantly reduces storage complexity.

Weight elimination and update iteration. If the original CAP eliminates weights in greedy optimal order, FastCAP prunes weights in a fixed order, similar to SparseGPT [13]. This reduces computational complexity.

The full pseudocode for FastCAP is provided in Algorithm 2. The same procedure is applied for all weights. We use the following notation:

B - Fisher block size

b - block size in iterative pruning algorithm

N - number of gradients

N B - number of blocks per parameter

The Fisher matrix is the same for all blocks, reducing space complexity to O ( B 2 ) . We adapted the trick of maintaining the Fisher inverse in Cholesky form from SparseGPT [13], to account for iterative elimination of the weights in a fixed order. This allows running the iterative process efficiently.

Algorithm 2 Fast CAP pruning algorithm

To validate performance of FastCAP we applied it for pruning of large ViT models in a one-shot setting, specifically eva_giant_patch14_224.clip_ft_in1k from the TIMM library [48]. This model has >1B parameters. We compared the top-1 accuracy on ImageNet of FastCAP vs Magnitude pruner. One can see that FastCAP significantly outperforms the baseline. At the same time, FastCAP has a reasonable runtime and memory footprint at this scale. The pruning step for the whole model takes 400 seconds on a single A100 GPU.

C Training details

Augmentation/regularization recipe

For the gradual pruning experiments (with 300 epochs) we have used cyclic learning schedule, with high learning rate directly after the pruning step with gradual decrease up to the next pruning step. For both DeiT-Tiny and DeiT-Small model during the additional fine-tuning for 100 epochs we've applied cosine annealing schedule with η max = 5 · 10 -5 , η min = 1 · 10 -5 and all other parameters the same as in the Table 6.

D Post-Pruning Recovery

The choice of augmentation parameters and learning rate schedule is critical for high performance. For example, reducing the level of augmentation during fine-tuning for smaller models, e.g. DeiT-Tiny, significant improves performance, whereas larger models, e.g. the 4x larger DeiT-Small, requires strong augmentations for best results even during fine-tuning. See Figure 5 for an illustration; the augmentation procedure is described in detail in C.

Moreover, the choice of cyclic learning rate (LR) schedule is critical as well. To illustrate this, we compare convergence obtained when using a cosine annealing schedule, which is very popular for pruning CNNs [25, 39, 33], from η max = 5 · 10 -4 to η min = 10 -5 , while performing pruning updates 2 times more frequently (one update per 10 epochs) than in our standard setup from the following section 4.2. The results are provided in Figure 5, where cosine annealing (no cycles) is in red. All experiments use the CAP pruner, and highlight the importance of the learning rate and augmentation schedules for recovery.

0

50

100

150

200

250

300

58

60

62

64

66

68

70

72

74

Val accuracy

DeiT-Tiny

deit

no cycles

0

50

100

150

200

250

300

70

72

74

76

78

Val accuracy

DeiT-Small

deit

Epoch

Epoch

E Additional Results for One-Shot Pruning

In this section we present comparison of Global Magnitude (GM), WoodFisher (WF) and Correlation Aware (CAP) pruners in one-shot pruning setting for DeiT models [43] of different size (i.e DeiT-Tiny, DeiT-Small, DeiT-Base) to study the scaling behavior of ViT sparsification.

Notice, that the gap between magnitude pruning and second order methods is very pronounced for all models, whereas the difference in performance between CAP and WF decreases with increase of the size of model. Nevertheless, CAP performs still noticeably better than WF, especially in high sparsity regime.

Pruning and finetuning. In most of the practical setups one cannot achieve both high compression rate and maintain performance of the dense model in one-shot setup. The full retraining procedure allows to achieve high sparsity but is rather expensive in the terms of compute. One can be interested

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Val accuracy

DeiT-Tiny

GM

WF

CAP

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

GM

WF

CAP

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Base

GM

WF

CAP

to have something in between - moderately sparse model, close in performance to the original model but without the need of long and expensive training.

In the experiments below we prune all models to 50% sparsity, and fine-tune for 20 epochs. In addition to ViT/DeiT, we also consider similar models based on variants of self-attention [28, 1], and compare against a GM baseline. We use a linearly-decaying learning rate schedule between η max = 10 -4 to η max = 10 -5 and the DeiT training recipe [43]. The results are given in Table 7, and show that CAP can almost fully-recover accuracy in this setup for all models; the gaps from GM and WF (see DeiT-Small 75 and 90%) are still very significant.

F Experiments with other models

In the main part of the text, we considered only gradual pruning of ViT models, but the proposed method is applicable to any architecture for image classification, such as convolutional neural network (CNN) or a ViT-CNN hybrid. We have selected recently proposed EfficientFormer [27] as a member of ViT-CNN hybrid family and trained it using the same setting and hyperparameters as for DeiTSmall. Two CNN architectures - ResNet50-D 2 and EfficientNetV2-Tiny [42] 3 , considered in this work were trained with the use of augmentation and regularization procedure described in the recent PyTorch blog post. Differently from most of the prior art we have used the ResNet50-D trained with the modern recipe from timm repository.

For ResNet50-D we prune all convolutional weights except the first convolution and we keep the classification layer dense. In EfficientNetv2-Tiny we do not prune depthwise convolutions since they do not contribute much to the total number of parameters and FLOPs but they are important to the model performance. We have set the block size to be 256 for ResNet50-D and 16 for EfficientNetV2Tiny while keeping all the other hyperparameters of CAP the same as for DeiT experiments. Such a small block size was chosen for EfficientNetV2-Tiny due to the fact that it is the largest common divisor of the prunable weights.

First of all, we conducted comparison between one-shot pruning methods for ResNet50-D. We compare between Uniform and Global magnitude pruning, WoodFisher with block size of 256,

M-FAC with block size of 2048 and CAP with uniform and global sparsity. One can observe that CAP outperforms all previous methods even when comparing uniform sparsity with global sparsity. Contrary to the case of DeiT where there is no much difference in performance between uniform and global magnitude pruning for ResNet50-D global sparsity turns out to be much better. This results is quite expectable since CNN are not uniform and deeper layers are mode wide than those close to the input.

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

ResNet50-D

M-FAC (G)

WF (G)

CAP (U)

Next we carried out one-shot + finetuning experiments with ResNet50-D keeping setup the same as for DeiT models in Table 7. We have selected 75 % and 90 % as the difference between methods becomes pronounced only at sufficiently high sparsity. Notably, the performance of Global magnitude and WF is roughly the same, the initial difference after one-shot pruning between WF and Global Magnitude vanishes during the finetuning procedure, whereas there is still a gap in performance between CAP and other methods.

Finally we conducted gradual pruning runs following the same sparsification schedule as for DeiTmodels. The EfficientFormer and EfficientNet models despite being already very optimized and parameter efficient can be still compressed with small drop in accuracy.

G Scaling behaviour of pruning with respect to model size

To study the scaling behavior with respect to model size we took all variants of the ConvNext2 family of models [49] since it covers wide range of model sizes except for the large one due to the memory and compute constraints. The smallest model from the family - ConvNext2-Atto has 3.7M parameters whereas the largest considered ConvNext2-Large has 198M parameters. All the models were pruned to 50% in one-shot. We observed that the relative accuracy drop (difference between accuracy of the dense and sparse model) initially decreases with increase of model size and then reaches a plateau. CAP consisently outpeforms WF across all scales and the difference is the most pronounced for the smallest model.

4

8

16

32

64

128

256

# Params (M)

0

2

4

6

8

10

12

14

16

CAP

WF

H Timings

Any algorithms involving second order loss information are believed to require tremendous amounts of compute. Time required for calculation of pruning scores and the OBS update comprises collection of grads, Fisher inverse rank-1 updates and additional pruning iteration for CAP. We have measured the time of single pruning step for DeiT-Small and present the results in Table 10. All measurements were done on a single RTX A6000 GPU with 48GB of memory. One can observe that the amount of time needed to perform a pruning update is not very large, especially when compared to the duration of typical training procedure of modern computer vision models on ImageNet that usually takes several days on a multi-gpu node. Note that the additional step for CAP adds small fraction of total computational time relative to other steps of the OBS method.

I Composite compression

In addition to weight pruning one can decrease storage and inference cost with the help of other compression approaches: quantization (casting weights and activations to lower precision) and token pruning specific for the transformer architecture.

I.1 Quantization-Aware Training

Weight quantization is done in the following way - one takes sparse checkpoint and then runs quantization aware training (QAT). We ran QAT training for 50 epochs with linearly decaying learning rate schedule from η max = 10 -4 to η min = 10 -5 . Models are quantized to 8-bit precision. In all experiments performed accuracy of quantized model almost reproduces the accuracy of the sparse model stored in full precision.

I.2 Token Pruning

There are different approaches for token pruning proposed in the literature. In this work we follow the one from [35]. Specifically, in DynamicViT one selects the ratio of tokens being pruned at each step with the lowest importance score, predicted by the model itself. Following the main setup from the paper we prune tokens after 3 rd , 6 th , 9 th block, and the token pruning ratio after each block is ρ = 0 . 2 (i.e 20% least improtant tokens are pruned).

I.3 Semi-structured sparsity

While CPUs can utilize sparsity patterns of arbitrary form to speed-up the computations at the present time modern GPU accelerators can handle only restricted form of unstructured sparsity, namely the N : M sparsity pattern that enforces exactly N non-zero values for each block of M weights. Namely, since the introduction of Ampere architecture NVIDIA GPUs have special kernels that can work with 2 : 4 sparse matrices [30]. One can integrate the N : M sparsity in the CAP framework without significant changes. The only difference with the original CAP approach is that while running the CAP iterations one doesn't prune a given weight in case in a group of M weights to which this weights belongs to there are M -N zero weights. Since the sparsity pattern is significantly constrained compared to generic unstructured sparsity pattern drop in performance after doing pruning step and consequent fine-tuning is more challenging than it would be for unconstrained sparsity. In experiments below we prune models to 2 : 4 sparsity either in one-shot setting and one-shot+finetune. We apply shorter (10 epochs) and longer (50 epochs) finetuning procedure with linearly decaying learning rate schedule. According to the results in the Table 13 CAP significantly outperforms competitive methods for one-shot pruning, although the drop in performance is quite large for all methods. After finetuning procedure difference between different methods decreases. Nevertheless, there is some gap in performance between second order methods and magnitude pruning even after relatively long finetuning.

To demonstrate practical benefits from 2:4 sparsity pattern we compiled both sparse and dense models via TensorRT engine and compared the throughput. The inference was executed on Nvidia T4 GPU

with batch size of 64 in half precision. Sparsity allows for small but certain speedup for models of different scale.

J CAP/WF hyperparameters

Following the oBERT's directions [23] on identifying the optimal set of hyperparameters via one-shot pruning experiments, we conduct a grid search over the three most important hyperparameters:

Number of grads collected for Fisher inverse

Dampening constant λ

Block size

The more grads are collected, the more accurate is the empirical Fisher inverse estimate, however, more compute is required at the same time. We chose N = 4096 as a point from which further increase of Fisher samples doesn't improve performance a lot. Dependence of the one-shot pruning performance at different sparsities vs number of grads is presented on Figure 9.

The next parameter to be studied is the dampening constant λ in. This constant regularizes the empirical Fisher matrix and allows to avoid instabilities in computation of the inverse. However, this constant decreases the correlation between different weights and in the limit λ →∞ OBS reduces to magnitude pruning. The optimal dampening constant for CAP ( λ opt = 10 -8 ) is smaller than the one for WoodFisher ( λ opt = 10 -6 ), i.e CAP remains numerically and computationally stable with smaller amount of regularization compared to WF (we observed that for λ < 10 -7 WF performance starts to deteriorate rapidly).

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

16

64

256

1024

4096

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

16

64

256

1024

4096

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

1e-5

1e-6

1e-7

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

1e-7

1e-8

1e-9

And the last but not the least important parameter is the block size in [39]. The larger the block size is, the more correlations between different weights are taken into account. However, as mentioned in 2 the computational and storage cost scales with the block size. Moreover, for a fixed number of gradients in the Fisher estimate matrix with larger block sizes is likely to be worse conditioned. Therefore, one would like to work with smaller block sizes but not to keep the approximation as accurate as possible. We've selected block size according to the accuracy-efficiency trade-off.

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

48

96

192

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

48

96

192

In addition, we've studied the benefit from application of multiple recomputations in the one-shot pruning setting for WoodFisher and CAP. Since the assumption of static Fisher matrix F ( w ∗ ) doesn't hold in general, we expect that multiple recomputations are likely to result in higher one-shot accuracy

in accordance with the result from [14]. This is indeed the case. The gain from recomputations is more pronounced for WoodFisher, since CAP already performs implicit Fisher inverse updates in its operation. Yet, the efect is not vanishing even for the case of CAP.

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Sparsity

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Val accuracy

DeiT-Small

WF (Nr=1)

WF (Nr=2)

WF (Nr=5)

WF (Nr=8)

CAP (Nr=1)

CAP (Nr=5)

K Details and hyperparameter choices for other pruning methods

In this section we provide some additional details about methods compared on Figures 2 and 7. The popular Movement Pruning [38] computes the weight saliency scores during the training procedure, hence it is not a one-shot pruning method by the definition. We have observed that use of the naive elementwise product of gradient and weights ( i.e ρ i = w i ⊙∇ w i L ( w ) ) leads to a poor performance, significantly below even the Magnitude Pruning baseline. However, the following first order pruning criterion:

allows to get reasonable saliency scores that produce more accurate sparse models than Magnitude Pruning. However, its performance is still inferior to any of the second order pruning methods. This method is denoted as GrW (Gradient times weight) on Figures 2 and 7.

M-FAC Pruner proposed in [14] is a pruner leveraging second order information that doesn't require an explicit construction of Fisher Inverse matrix. Therefore, unlike WoodFisher and CAP that require O ( Bd ) memory computation and storage cost of this method is constant with respect to the block size and one can take into account correlations between larger groups of weights for free. Following the original paper we chose block size of 2k as the best performing one. However, one can see from Figures 2 and 7 that smaller block sizes turn out to perform better. A possible explanation of this phenomenon is that the Fisher Inverse estimate becomes too noisy and unstable for large blocks.

L Execution latency

In addition to the plot throughput vs accuracy shown in the main part we present in this section execution latency per sample vs latency when running models on the DeepSparse engine. The results are presented on Figure 13.

M Comparison with AC/DC training

In addition to the sparse training from scratch with periodic updates of the sparsity weights with some saliency criterion for weight elimination and regrowth [10] one can consider alternating compressed/decompressed training (AC/DC), proposed in [33]. Namely one switches between dense stages with standard unconstrained training of the model, and sparse stages when the model is pruned to the target sparsity level and trained with the frozen sparsity mask until the beginning of the next dense stage, when the sparsity mask is removed. This procedure produces both accurate dense and sparse models.

3

10

30

100

Latency (ms)

68

70

72

74

76

78

80

82

Val accuracy

DeiT-Base

DeiT-Small

DeiT-Tiny

Following the original paper we use magnitude pruning as a saliency criterion for weight pruning. The augmentation and regularization pipeline follows the settings from [43]. All models with AC/DC were trained for 600 epochs in total with first pruning step at epoch 150 followed by 7 sparse stages 25 epochs long each, and 6 dense stages of the same length. The last dense stage lasts 50 epochs and the last sparse is 75 epochs long. Learning rate is gradually decayed from η max = 5 · 10 -4 to η min = 10 -6 with cosine annealing. Initial warm-up phase with linearly increasing learning rate is 20 epochs. We compare AC/DC with CAP models finetuned for additional 100 epochs.

One can observe that at low sparsity AC/DC achieves higher accuracy for the same sparsity target (even outperforming the dense baseline by 0.6%), whereas for 75% performance of both methods is equal, and CAP outperforms AC/DC at higher sparsity. However, one should note, that CAP uses computational budget (including the training of original model) of 440 epochs for 60% sparsity, 520% for 75% and 700% for 90% vs 600 epochs used in AC/DC.

N One-shot pruning of DETR

The approach presented in the paper is not limited to the image classification task, but can be applied to other computer vision tasks, such as object detection. We chose the DeTR model [4] with ResNet50 backbone and ran one-shot pruning procedure with global magnitude, WoodFisher and CAP pruner. Specifically, we pruned all convolutional layers in the CNN backbone expect the first one and all linear projections in transformer encoder and decoder blocks while keeping the detection heads dense. The results are presented in Table 16. Following the standard protocol we used bbox mAP for evaluation. One can observe, that difference between the second order methods and magnitude pruning is very pronounced even for relatively small sparsity of 50%, and CAP outperforms WF pruner.

O Proof of Theorem 1

Theorem O.0. Let S be a set of samples, and let ∇ ℓ 1 ( w ∗ ) , . . . , ∇ ℓ m ( w ∗ ) be a set of gradients with i ∈ S , with corresponding empirical Fisher matrix ̂ F -1 ( w ∗ ) . Assume a sparsification target of k weights from w ∗ . Then, a sparse minimizer for the the constrained squared error problem

is also a solution to the problem of minimizing the Fisher-based group-OBS metric

Proof. We start by examining the unconstrained squared error function in Equation (10), which we denote by G . Clearly, the function G is a d -dimensional quadratic in the variable w ′ , and has a minimum at w ∗ . Next, let us examine G 's second-order Taylor approximation around w ∗ , given by

where we used the fact that w ∗ is a minimum of the squared error, and thus the function has 0 gradient at it. However, by the definition of the empirical Fisher, this is exactly equal to

The Taylor approximation is exact, as the original function is a quadratic, and so the two functions are equivalent. Hence, we have obtained the fact that, under the empirical Fisher approximation, a k -sparse solution minimizing Equation 10 will also be a k -sparse solution minimizing Equation 1. However, the question of finding a k -sparse solution minimizing Equation 1 is precisely the starting point of the standard OBS derivations (see e.g. [39] or [23]), which eventually lead to the formula in Equation (11). This concludes the proof.

P Augmentation choice for Empirical Fisher

We compared the performance of CAP with Empirical Fisher computed on image-label pairs where the validation transforms were applied to images (i.e center crop with resize) and the same set of augmentations used for training and finetuning (RandAugment transforms, Label smoothing, e.t.c.). We observed that the sparsity solution obtained without augmenting samples for Empirical Fisher estimate turns out to be strongly overfitting. We point out that in both cases we use the same population size for Empirical Fisher.

Figure 14 illustrates this result: using validation augmentation (red) yields better training loss but degenerates in terms of validation accuracy. A possible explanation is that CAP chooses an overfitting solution which the model is unable to escape during finetuning.

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

Epoch

1.2

1.4

1.6

1.8

2.0

2.2

2.4

Train loss

DeiT-Tiny

no aug

aug

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

Epoch

1.2

1.3

1.4

1.5

1.6

1.7

Val loss

DeiT-Tiny

no aug

aug

Q Fisher matrix structure

In order to validate the necessity of taking into account the correlations between weights, one has to make sure that the empirical Fisher matrix used as proxy for Hessian is non-diagonal. We have visualized an average block of empirical Fisher for a particular layer on Figure 15 from DeiT-Tiny and ConvNext-Small models. One can see, Fisher matrix exhibits a pronounced non-diagonal structure, which justifies the need of a careful and thorough treatment of weight correlations.

blocks.5.attn.proj.weight

stages.2.blocks.24.mlp.fc1.weight

R Broader impact

Compressed models are not expected to exhibit more malicious and potentially harmful behavior compared to the dense models. However, they may face the same issues like the original models in safety-critical applications such as susceptibility to adversarial attacks and distribution shifts.

S Limitations

The proposed method is mostly suitable for small and medium sized models (up to order of ∼ 100 M). For larger models the compute and storage cost associated with the estimate of empirical Fisher becomes prohibitively expensive. Compression of the largest models considered is this paper requires 2-4 high-end GPUs (A100 with 80GiB). Pruning models to high sparsity requires significant amount of training. Search for fast and efficient procedures for the recovery of compressed models is left as potential direction for further research.

Peer reviews:
comments: SUMMARY:
This paper studies the unstructured pruning problem for vision transformer models. The Correlation Aware Pruner (CAP) is proposed. CAP takes into account weight correlations and achieves a new state-of-the-art result.

STRENGTHS:
- The proposed method achieves strong experimental results. For the first time, ViT models can attain high sparsity levels (75-80%) without significant accuracy loss (<1%). Note that previous methods reached at most 50% sparsity.
- CAP is reasonable and well-motivated. This paper points out an important problem that removed weights may themselves be correlated. It also highlights the importance of the learning rate schedule. Overall, this paper provides lots of useful information for unstructured pruning for ViT.
- The whole paper is organized and written well. Experiments are sufficient and sound.

WEAKNESSES:
- The biggest problem is unstructured pruning struggles with GPU devices. The authors test models’ latency on a sparsity-aware CPU inference engine. However, GPU is not considered. Overall, unstructured pruning performs badly on GPU which limits its practical ability.
- Some results of SViTE are missing in Tab. 2. It is better to change the order to put CAP at the bottom of each section.

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
The authors have adequately addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission proposes a second-order method for unstructured pruning of neural net parameters, to leverage the efficiency gains of sparsity. It uses an optimization based on the empirical Fischer matrix to find saliency scores, that are used for the order in which weights are pruned. They present an algorithm based on solving the problem on individual sub-blocks, and then using that within a global pruning. They also present a schedule, across different hyperparameters, for training with pruning. This is demonstrated on ConvNeXt and DeiT models, along with other CNN and ViT models in the supplement.\n\nSTRENGTHS:\n**i)** Achieves better accuracy & sparsity.\n\nThe proposed pruning gets pareto-dominant results relative to some similar prior work. The evaluation is done across multiple choices of the sparsity vs accuracy tradeoff.\n\nNot all experiments consider some of the most similar previous methods, though, such WoodFisher in the \"Gradual Pruning\" results, though the original work did include gradual pruning during training. (See section 5.2 of [38])\n\n**ii)** Considers good selection of recent models\n\nExperiments demonstrate the pruning on both up-to-date ConvNet and vision transformer models. In the supplement and code, experiments on ResNet-50 and EfficientNet are also provided.\n\n**iii)** Includes code\n\nThe submission includes the code used to run the experiments, based on SparseML, along with YAML files specifying the hyperparameters for a number of the experiments.\n\nWEAKNESSES:\n**iv)** Largely superseded by Optimal BERT Surgeon in a lot of settings\n\nOBS is a more scalable algorithm based on the same underlying principles. The core contribution seems to be a different method of approximating the solution to the original 2nd-order problem, plus better hyperparameter tuning. It does seem like a reasonable approach for relatively smaller models, to seek a tighter approximation to the original optimization.\n\n**v)** Source of improvement in end metrics remains unclear\n\nAt its core, the starting point of the construction of the method in the submission is similar to [38]: namely, second-order pruning with the empirical Fisher matrix in place of the Hessian. The differences from [38] then include a) a different method of approximating the pruning on the full matrix, described toward the end of Section 3.1, and b) improved hyperparameters during gradual pruning.\n\nThe paper states that ablation studies are given in appendices B, I, and J. These appendices specify different sets of hyperparameters, and include results of sparsity vs accuracy for multiple different choices of some of these hyperparameters. I didn't find any \"ablation\" studies that evaluate parts of the proposed method with other parts \"ablated\" or removed: for example, the results of using Algorithm 1 *without* the changes to data augmentation, the sparsity and learning rate schedule, etc.\n\nExperiments also don't compare against [38] for the gradual pruning case, so the effect of the hyperparameter tuning and schedules isn't measured by a comparison to baseline. It's possible that the improved choices in which weights to prune, as seen in the one-shot case, could also affecting the gradual pruning case, but it's also possible that the training can close a lot of this gap.\n\nThe experiments also use very recent ViT-based models, showing some more up-to-date results than prior work. It is plausible that better pruning is necessary to get good results for these models, as argued, but this is incompletely supported.\n\nQUESTIONS:\n**Code**: Some miscellaneous minor comment from trying the included code\n\nREADME.md marks the step to install wandb as \"optional.\" Not entirely optional, though, given that it does import it:\n\n```\n  File \"CAP-code/research/one_shot_pruning.py\", line 4, in <module>\n    import wandb\nModuleNotFoundError: No module named 'wandb'\n```\n\nCommenting out the import, then see:\n\n```\nAttributeError: module 'numpy' has no attribute 'object'.\n`np.object` was a deprecated alias for the builtin `object`. To avoid this error in existing code, use `object` by itself. Doing this will not modify any behavior and is safe. \nThe aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at:\n    https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n```\n\nwithin the import of onnx inside sparseml:\n\n```\n    from sparseml.pytorch.utils import LOGGING_LEVELS, BaseLogger, LoggerManager\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/sparseml_nightly-0.13.0.20230704-py3.9.egg/sparseml/pytorch/utils/__init__.py\", line 23, in <module>\n    from .exporter import *\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/sparseml_nightly-0.13.0.20230704-py3.9.egg/sparseml/pytorch/utils/exporter.py\", line 26, in <module>\n    import onnx\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/onnx-1.10.1-py3.9-linux-x86_64.egg/onnx/__init__.py\", line 20, in <module>\n    import onnx.helper  # noqa\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/onnx-1.10.1-py3.9-linux-x86_64.egg/onnx/helper.py\", line 17, in <module>\n    from onnx import mapping\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/onnx-1.10.1-py3.9-linux-x86_64.egg/onnx/mapping.py\", line 27, in <module>\n    int(TensorProto.STRING): np.dtype(np.object)\n  File \"${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/numpy/__init__.py\", line 305, in __getattr__\n    raise AttributeError(__former_attrs__[attr])\n```\n\nFollowing instructions in the README had given me `numpy==1.24.3` and `onnx==1.10.1`. It's likely that one needs to hold back numpy, likely an underlying issue with the requirements specified for SparseML. Presumably this codebase is using an old version of sparseml that specifies a requires.txt such that it's versions specified with '>=' no longer really work with the latest versions now available. Authors should check & perhaps update the install commands in their README accordingly, or update their sparseml version. I was able to resolve this with instead:\n\n```\nconda install numpy==1.20.3\n```\n\nGenerally recommended practice for distributing this kind of research code is to explictly specify versions of the requisite packages with `==`, so we know exactly how the authors originally did all this.\n\nSample command for `one_shot_pruning.py` in README also includes the command-line arguments `--data-dir` and `--sparseml-recipe`, when the arguments actually added to argparse are `--data_dir` and `--sparseml_recipe`, and arguments `-b` and `--experiment` that don't seem to exist.\n\nLIMITATIONS:\nCite computation cost (at training/pruning, not inference/deployment time) as the limitation. The Fischer matrix will still grow quadratically in the number of parameters. This is unfortunate given the likelihood that the largest models that would benefit the most from pruning, in absolute terms. (And it is an oft-observed principle in ML that simple methods that scale better are usually more impactful.)", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper proposes to consider the correlation between pruned elements in pruning deep neural netowkr models. The paper provides an efficient algorithm to distangle the correlation into a sparse regression problem, and propose a fast solver to find the solution. Further exploration is performed on the learning rate scheduling and data augmentation of the pruning and finetuning process. The final method results in better model size-accuracy tradeoff comparing to previous method with less training cost.

STRENGTHS:
1. This paper provides well motivated method on model pruning with solid theortical justification
2. The proposed reformulation of sparse minimization and the fast solver of the optimization is novel an have significant impact on model compression
3. Adequate experiments are provided to show the effectiveness of the proposed method, including pruning advanced highly accurate models
4. The paper is overall well written and easy to follow.

WEAKNESSES:
As discussed in the limitation, the cost of considering full correlation is still prohibitive on large models. It would be interesting to see some relaxation of the proposed method that can have less cost while maintaining most of the good performance. More discussion is encouraged on the potential relaxation strategy, and the tradeoff between optimization cost and performance.

QUESTIONS:
See weakness.

LIMITATIONS:
The limitation and social impacts have been adequately discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05

comments: SUMMARY:
The paper proposed a Correlation Aware Pruner (CAP) , a new unstructured pruning framework capable to prune models to high sparsity. It takes into account weight correlations. To do this, the paper reformulate the OBS multi-weight pruning problem: when using the empirical Fisher approximation, the problem of finding the optimal set of weights to be removed, while taking correlations into account, is equivalent to the problem of finding the set of sparse weights which best preserve the original correlation between the dense weights and the gradients on an fixed set of samples.

On top of that, the paper also applies a series of training techniques to improve training: Learning Rate Schedule,  Regularization and Augmentation, Efficient Sparsity Sweeps. Results show better performance compared to existing works, expecially at high compression ratio

STRENGTHS:
The paper is well written, and has good logic. Great amount of experiments and ablations discussion to strengthen the proposed method.

WEAKNESSES:
Unstructured pruning is hard to have actual hardware speedup due to the irregular sparsity. Although the proposed method can maintain high accuracy at extreme prune ratios, it may just be just theoretically more efficient. 

The author is encouraged to add inference latency or throughput in the results tables 1 & 2

QUESTIONS:
See weekness

LIMITATIONS:
Please address the limitations and potential negative societal impact in the revision.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09