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

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.

Paper

Similar papers

Peer review

Reviewer gPY65/10 · confidence 4/52023-07-02

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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations

Reviewer kMpj3/10 · confidence 4/52023-07-05

Summary

The 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.

Strengths

**i)** Achieves better accuracy & sparsity. The 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. Not 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]) **ii)** Considers good selection of recent models Experiments 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. **iii)** Includes code The 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.

Weaknesses

**iv)** Largely superseded by Optimal BERT Surgeon in a lot of settings OBS 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. **v)** Source of improvement in end metrics remains unclear At 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. The 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. Experiments 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. The 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.

Questions

**Code**: Some miscellaneous minor comment from trying the included code README.md marks the step to install wandb as "optional." Not entirely optional, though, given that it does import it: ``` File "CAP-code/research/one_shot_pruning.py", line 4, in <module> import wandb ModuleNotFoundError: No module named 'wandb' ``` Commenting out the import, then see: ``` AttributeError: module 'numpy' has no attribute 'object'. `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. The aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations ``` within the import of onnx inside sparseml: ``` from sparseml.pytorch.utils import LOGGING_LEVELS, BaseLogger, LoggerManager 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> from .exporter import * 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> import onnx 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> import onnx.helper # noqa 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> from onnx import mapping 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> int(TensorProto.STRING): np.dtype(np.object) File "${HOME}/anaconda3/envs/CAP/lib/python3.9/site-packages/numpy/__init__.py", line 305, in __getattr__ raise AttributeError(__former_attrs__[attr]) ``` Following 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: ``` conda install numpy==1.20.3 ``` Generally 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. Sample 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.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Cite 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.)

Reviewer JVcP7/10 · confidence 4/52023-07-05

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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The limitation and social impacts have been adequately discussed.

Reviewer wReA5/10 · confidence 4/52023-07-09

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

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Please address the limitations and potential negative societal impact in the revision.

Authorsrebuttal2023-08-19

Discussion reminder

Dear Reviewers, Given that the discussion period is ending shortly and we have received no feedback so far, we wanted to send a gentle reminder regarding our review responses. We would really appreciate your feedback on our rebuttal, especially regarding the added experimental results and clarifications regarding our method. With best regards, The CAP Authors

Reviewer JVcP2023-08-21

I would like to thank the author for the response. I'm satisfied with the response and will keep my score.

Reviewer kMpj2023-08-21

> Thus, our one-shot pruning comparisons essentially provide one of the key ablation study that the reviewer is asking for: WoodFisher/OBS is CAP without the correlation solving This seems correct, and a good experiment without confounds (since its comparing under the simple & shared one-shot procedure) in the case of Figure 2 & similar experiments in the supplement. Though I'm perhaps looking in the wrong place for Table 4 in the main text? > general response we present a much more efficient approximate version of CAP (FastCAP). Where ought we look for more details on FastCAP? The description given in the top-level comment seems elides a lot that might be needed for a full review. > compression of moderate-size models (e.g. with 1-50M parameters) is still of great practical use for inference on edge devices Definitely true! Perhaps this is scalable enough.

Authorsrebuttal2023-08-21

Thank you for your response! Please see replies inline: > This seems correct, and a good experiment without confounds (since its comparing under the simple & shared one-shot procedure) in the case of Figure 2 & similar experiments in the supplement. Though I'm perhaps looking in the wrong place for Table 4 in the main text? We are very glad that the reviewer found this part of the response clarifying. We will also highlight this point in the next revision. We apologize for the confusing table reference: we meant Table 1, not Table 4. More precisely, Table 1 shows the comparison of one-shot pruning performance of various methods on large (CLIP-sized) models. > Where ought we look for more details on FastCAP? The description given in the top-level comment seems elides a lot that might be needed for a full review. To address this, we provide the full pseudocode for FastCAP below; for simplicity, we provide the code for a single weight. The same procedure is applied for all weights. The main two components of the CAP/FastCAP algorithm are: 1. **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 reduces storage complexity.) 2. **Weight elimination and update iteration.** If the original CAP eliminates weights in greedy optimal order, FastCAP prunes weights in a fixed order. (This reduces computational complexity.) Notation: * `B` - Fisher block size * `b` - block size in iterative pruning algorithm * `N` - number of gradients * `N_B` - number of blocks per parameter 1) **Fisher accumulation** `F` - matrix with zeros with shape `(B, B)` **for** i=1:N **do** &ensp; reshape `dL_i / dw` to `(N_B, B)` &ensp; `F += (1 / N) (dL_i / dw)^T dL_i / dw ` // average Fisher blocks across the layer output dimension **end for** 2) **Iterative pruning** `M ← 1_{N_B x B}` // binary pruning mask `E ← 0_{N_B x b}` // block pruning errors `F ← Cholesky({F}^{-1})` // efficient form of the Fisher inverse **for** `i`=`0`, `b`, `2b`, … **do** &ensp; // select elements to prune &ensp; `M_{:,j:(j+b)}` ← mask of `(1 − p)`% weights `w_k` `in W_{:,j:(j+b)}` with largest `w_k^2 / F^{-1}]_{kk}^2` &ensp; for j=`i`, …, `i+b-1` **do** &ensp; &ensp; `E_{:,j−i} ← W_{:,j} / [F^{-1}]_{jj}` // pruning error &ensp; &ensp; `E_{:,j−i} ← (1- M_{:, j}) E_{:,j−i}` // freeze weights &ensp; &ensp; `W_{:,j:(i+b)} ← W_{:,j:(i+b)} - E_{:,j−i} F_{j,j:(i+B)^{-1}} ` // update remaining weights to correct output &ensp; **end for** **end for** `W ← W * M` // set pruned weights to 0 As mentioned in the earlier response, the Fisher 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 the SparseGPT paper, to account for iterative elimination of the weights in a fixed order. This allows running the iterative process efficiently.

Authorsrebuttal2023-08-21

We wish to make an addition to our earlier response: As requested, we ran S-ViTE with 80% sparsity for DeiT-Small. Consistent with the previous results, the proposed CAP methods significantly outperforms S-ViTE. | Sparsity | Method | Top-1 Accuracy (%) | |:--------:|:------:|:------------------:| | 80 | S-ViTE | 75.8 | | | CAP | 78.0 | In addition, we note that S-ViTE is quite computationally expensive. Running of the algorithm for 600 epochs with the same hyperparameters as in the original paper requires 5 days on 4 RTX 3090 GPUs, whereas our approach produces a more accurate model using only half of S-ViTE computational cost.

Authorsrebuttal2023-08-21

Thank you!

We would like to thank the reviewer for their response, and for raising their score! P.S.: As far as we can tell, the score has remained the same in the Author Console.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC