SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Denis Kuznedelev
Skoltech & Yandex
Denis.Kuznedelev@skoltech.ru
Eldar Kurtic
IST Austria
eldar.kurtic@ist.ac.at
Elias Frantar
alias.frantar@ist.ac.at
Dan Alistarh
IST Austria & Neural Magic
dan.alistarh@ist.ac.at
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.4\mathrm{x}$  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<sup>1</sup>.
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 similarly-accurate 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:
1The code is available at https://github.com/IST-DASLab/CAP
1. 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.
2. **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 correlation-aware 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 highly-compressible, 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 second-order 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  $\mathbf{w} \in \mathbb{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  $\mathbf{w} \in \mathbb{R}^d$ , the saliency of each weight is its absolute value (the magnitude)  $\rho_j = |w_j|$  for  $j \in \{1, 2, \ldots, 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  $\mathcal{L}$  in the vicinity of the dense model parameters  $\mathbf{w}^*$ . Assuming that  $\mathbf{w}^*$  is close to the optimum (hence  $\nabla \mathcal{L}(\mathbf{w}^*) \simeq 0$ ), one seeks a binary mask  $\mathbf{M}$  (with elements  $\in \{0,1\}$ ) and new values for the remaining weights  $\mathbf{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:
$$
\mathcal {L} \left(\mathbf {w} ^ {M}\right) - \mathcal {L} \left(\mathbf {w} ^ {*}\right) \simeq \frac {1}{2} \left(\mathbf {w} ^ {M} - \mathbf {w} ^ {*}\right) ^ {\top} \mathbf {H} _ {\mathcal {L}} \left(\mathbf {w} ^ {*}\right) \left(\mathbf {w} ^ {M} - \mathbf {w} ^ {*}\right) \tag {1}
$$
where  $\mathbf{H}_{\mathcal{L}}(\mathbf{w}^{*})$  is the Hessian of the model at  $\mathbf{w}^*$ , and  $\mathbf{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  $\rho_{i}$  for  $i^{\mathrm{th}}$  weight and the optimal weight update  $\delta \mathbf{w}$  for the remaining weights after elimination of the  $i^{\mathrm{th}}$  weight are as follows:
$$
\rho_ {i} = \frac {w _ {i} ^ {2}}{2 \left[ \mathbf {H} _ {\mathcal {L}} ^ {- 1} \left(\mathbf {w} ^ {*}\right) \right] _ {i i}} \quad \text {a n d} \quad \delta \mathbf {w} ^ {*} = - \frac {w _ {i}}{\left[ \mathbf {H} _ {\mathcal {L}} ^ {- 1} \left(\mathbf {w} ^ {*}\right) \right] _ {i i}} \mathbf {H} _ {\mathcal {L}} ^ {- 1} \left(\mathbf {w} ^ {*}\right) \mathbf {e} _ {i}, \tag {2}
$$
where  $\mathbf{e}_i$  is the  $i^{\text{th}}$  basis vector. Theoretically, the procedure would have to be executed one-weight-at-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 computationally-infeasible 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:
$$
\mathbf {H} _ {\mathcal {L}} \left(\mathbf {w} ^ {*}\right) \simeq \mathbf {F} \left(\mathbf {w} ^ {*}\right) = \lambda \mathbf {I} _ {d \times d} + \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla \mathcal {L} _ {i} \left(\mathbf {w} ^ {*}\right) \nabla \mathcal {L} _ {i} \left(\mathbf {w} ^ {*}\right) ^ {\top} \tag {3}
$$
where  $\nabla \mathcal{L}_i(\mathbf{w}^*)\in \mathbb{R}^d$  is a gradient computed on a sample of data,  $\lambda >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  $\mathcal{O}(d^2)$  to  $\mathcal{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:
$$
\min  _ {\delta \mathbf {w}} \frac {1}{2} \delta \mathbf {w} ^ {\top} \mathbf {F} \left(\mathbf {w} ^ {*}\right) \delta \mathbf {w} \quad \text {s . t .} \quad \mathbf {E} _ {Q} \delta \mathbf {w} + \mathbf {E} _ {Q} \mathbf {w} ^ {*} = \mathbf {0}, \tag {4}
$$
where  $\mathbf{E}_Q\in \mathbb{R}^{|Q|\times 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  $\delta \mathbf{w}_{\mathbf{Q}}^{*}$  of remaining weights are [23]:
$$
\rho_ {Q} = \frac {1}{2} \mathbf {w} _ {Q} ^ {* \top} \left(\mathbf {F} ^ {- 1} (\mathbf {w} ^ {*}) _ {[ Q, Q ]}\right) ^ {- 1} \mathbf {w} _ {Q} ^ {*} \mathrm {a n d} \delta \mathbf {w} _ {Q} ^ {*} = - \mathbf {F} ^ {- 1} (\mathbf {w} ^ {*}) \mathbf {E} _ {Q} ^ {\top} \left(\mathbf {F} ^ {- 1} (\mathbf {w} ^ {*}) _ {[ Q, Q ]}\right) ^ {- 1} \mathbf {w} _ {Q} ^ {*}. (5)
$$
However, an exhaustive search over all subsets of size  $|Q|$  from  $d$  elements requires  $\binom{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  $\sim 100\mathrm{M}$  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  $\emptyset$ , we always add to  $Q$  the weight  $q$  which has minimal joint saliency  $\rho_{Q \cup \{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  $\rho_{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  $\mathbf{w}^*$  and the gradients  $\nabla \mathcal{L}_i(\mathbf{w}^*)$  on an
fixed set of samples  $i \in 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  $\nabla \mathcal{L}_1(\mathbf{w}^*)$ , ...,  $\nabla \mathcal{L}_m(\mathbf{w}^*)$  be a set of their gradients, with corresponding inverse empirical Fisher matrix  $\mathbf{F}^{-1}(\mathbf{w}^*)$ . Assume a sparsification target of  $k$  weights from  $\mathbf{w}^*$ . Then, a sparse minimizer for the constrained squared error problem
$$
\min  _ {\mathbf {w} ^ {\prime}} \frac {1}{2 m} \sum_ {i = 1} ^ {m} \left(\nabla \mathcal {L} _ {i} \left(\mathbf {w} ^ {*}\right) ^ {\top} \mathbf {w} ^ {\prime} - \nabla \mathcal {L} _ {i} \left(\mathbf {w} ^ {*}\right) ^ {\top} \mathbf {w} ^ {*}\right) ^ {2} \tag {6}
$$
s.t.  $\mathbf{w}'$  has at least  $k$  zeros, is also a solution to the problem of minimizing the Fisher-based group-OBS metric
$$
\underset {Q, | Q | = k} {\arg \min } \frac {1}{2} \mathbf {w} _ {Q} ^ {* \top} \left(\mathbf {F} ^ {- 1} \left(\mathbf {w} ^ {*}\right) _ {[ Q, Q ]}\right) ^ {- 1} \mathbf {w} _ {Q} ^ {*}. \tag {7}
$$
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  $\ell_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  $\ell_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  $\Theta (d)$  weights would have  $O(d^{3})$  runtime, where  $d$  is the layer dimension, as the  $\Theta (d^2)$ -time selection + update process is repeated  $\Theta (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\cdot \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\cdot B^2)$  runtime and  $O(d\cdot 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 off-diagonal 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:
$$
\eta (t) = \eta_ {\max } - \left(\eta_ {\max } - \eta_ {\min }\right) \frac {t \% T}{T}, \tag{8}
$$
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 well-optimized (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].
Figure 1: (left): Blue: Cyclic linear learning rate schedule used in the work. Violet: Dependence of the global model sparsity on the epoch. Every change in sparsity corresponds to a pruning step. (right): Relative accuracy drop (i.e. difference between validation accuracy before and after pruning update) for training with cyclic and cosine schedule, respectively.