Resource Efficient Neural Networks Using Hessian Based Pruning
Neural network pruning is a practical way for reducing the size of trained
models and the number of floating-point operations. One way of pruning is to
use the relative Hessian trace to calculate sensitivity of each channel, as
compared to the more common magnitude pruning approach. However, the stochastic
approach used to estimate the Hessian trace needs to iterate over many times
before it can converge. This can be time-consuming when used for larger models
with many millions of parameters. To address this problem, we modify the
existing approach by estimating the Hessian trace using FP16 precision instead
of FP32. We test the modified approach (EHAP) on
ResNet-32/ResNet-56/WideResNet-28-8 trained on CIFAR10/CIFAR100 image
classification tasks and achieve faster computation of the Hessian trace.
Specifically, our modified approach can achieve speed ups ranging from 17% to
as much as 44% during our experiments on different combinations of model
architectures and GPU devices. Our modified approach also takes up around 40%
less GPU memory when pruning ResNet-32 and ResNet-56 models, which allows for a
larger Hessian batch size to be used for estimating the Hessian trace.
Meanwhile, we also present the results of pruning using both FP16 and FP32
Hessian trace calculation and show that there are no noticeable accuracy
differences between the two. Overall, it is a simple and effective way to
compute the relative Hessian trace faster without sacrificing on pruned model
performance. We also present a full pipeline using EHAP and quantization aware
training (QAT), using INT8 QAT to compress the network further after pruning.
In particular, we use symmetric quantization for the weights and asymmetric
quantization for the activations.
Jack Chong, Manas Gupta, Lihui Chen
# Resource Efficient Neural Networks Using Hessian Based Pruning

###### Abstract

Neural network pruning is a practical way for reducing the size of trained models and the number of floating point operations (FLOPs). One way of pruning is to use the relative Hessian trace to calculate sensitivity of each channel, as compared to the more common magnitude pruning approach. However, the stochastic approach used to estimate the Hessian trace needs to iterate over many times before it can converge. This can be time-consuming when used for larger models with many millions of parameters. To address this problem, we modify the existing approach by estimating the Hessian trace using FP16 precision instead of FP32. We test the modified approach (EHAP) on ResNet-32/ResNet-56/WideResNet-28-8 trained on CIFAR10/CIFAR100 image classification tasks and achieve faster computation of the Hessian trace. Specifically, our modified approach can achieve speed ups ranging from 17% to as much as 44% during our experiments on different combinations of model architectures and GPU devices. Our modified approach also takes up \(\sim\)40% less GPU memory when pruning ResNet-32 and ResNet-56 models, which allows for a larger Hessian batch size to be used for estimating the Hessian trace. Meanwhile, we also present the results of pruning using both FP16 and FP32 Hessian trace calculation and show that there is no noticeable accuracy differences between the two. Overall, it is a simple and effective way to compute the relative Hessian trace faster without sacrificing on pruned model performance. We also present a full pipeline using EHAP and quantization aware training (QAT), using INT8 QAT to compress the network further after pruning. In particular, we use symmetric quantization for the weights and asymmetric quantization for the activations. The framework has been open-sourced and the code is available at [https://github.com/JackkChong/Resource-Efficient-Neural-Networks-Using-Hessian-Based-Pruning](https://github.com/JackkChong/Resource-Efficient-Neural-Networks-Using-Hessian-Based-Pruning).

## I Introduction

In recent years, neural networks have become more powerful in many tasks such as computer vision and natural language processing [1]. This is partly due to larger models with many millions of parameters and increasing input sizes. Some notable examples for image classification on ImageNet are EfficientNet-B7 with 66 million parameters [2] and ViT-G/14 vision transformer with 1843 million parameters [3]. However, it has been shown that these models with millions of parameters are usually over-parameterized [4]. Moreover, it is very hard to deploy these models on cloud servers or edge devices due to their large size and bigger memory footprint. Since most deep neural networks contain redundant neurons, one way to make these models deploy-able is to compress them through neural network pruning. Since then, much research has been conducted to determine the best way to prune a neural network, such as the magnitude-based approach or the second order Hessian-based approach.

An important challenge to pruning is to find out which parameters are insensitive and can be pruned away without severe degradation of the model's performance. One way of doing so is to use the relative Hessian trace to calculate the sensitivity of each channel, in contrast to the standard magnitude pruning approach. The seminal work of [5] proposed Hessian Aware Pruning, which makes use of the stochastic Hutchinson trace estimator to determine the relative Hessian trace of each channel.

However, the Hutchinson trace estimator used to estimate the Hessian trace needs to iterate hundreds of times before it can converge to a stable value. This can be time-consuming when calculating the Hessian trace for larger models with many billions of parameters. In addition, the compression ratio of the pruned model is non-deterministic and can fluctuate within a few percentage points. The Hessian trace may need to be computed multiple times before pruning the model if a specific compression ratio is desired. To address this problem, we modify the existing approach in HAP by estimating the Hessian trace using FP16 precision instead of FP32. Our main contributions are:

* We test the default FP32 approach on 3 different neural network models (ResNet-32/ResNet-56/WideResNet-28-8) and record the average time spent to estimate the Hessian trace.
* We test our modified FP16 approach and compare the average time spent to estimate the Hessian trace against the default approach.
* We record the peak memory statistics of the GPU when estimating the Hessian trace using both FP32 and FP16 approaches.
* We compare the pruning results using both FP32 and FP16 approaches.
* We present a full pipeline using EHAP and quantization aware training together. INT8 quantization aware training (QAT) is done to compress the network further after
pruning.

## II Related Work

In recent years, the size of neural networks have rapidly increased and there is a growing need for faster inference on edge devices [6, 7, 8, 9, 10, 11]. Gradually, many approaches to make neural networks more compact and efficient have started to emerge [12, 13, 14, 15, 16, 17]. In this paper, we briefly discuss the related work on pruning and quantization.

### _Pruning_

Pruning can be grouped into two categories: structured pruning and unstructured pruning. Unstructured pruning aims to prune away redundant neurons without any structure. Unfortunately, this can result in sparse matrix operations which are difficult to speed up using GPUs [18]. To address this, structured pruning removes groups of redundant neurons in a structured way so that matrix operations can still be conducted efficiently. However, the difficulty lies in severe accuracy degradation when the model is pruned to a large degree. In both methods, the important challenge is to determine which parameters are insensitive to pruning. Many methods have been proposed in the literature to address this utilising regularization approaches [19, 20, 21, 22, 23], first or second order methods [24, 25, 26], gradient-based techniques [27, 28, 29], similarity approaches [30], sensitivity and feedback [31, 32, 33, 34, 35], and magnitude based approaches [36, 37, 38, 39, 40, 41].

One of the most popular ways is magnitude-based pruning. At the heart of this approach is the assumption that small parameters are redundant and can be pruned away. For example, [42] explores the use of magnitude based pruning to achieve a significant reduction in model size and only a marginal drop in accuracy. Another variant was used in [43] to prune redundant neurons using the scaling factor of batch normalization layers. In particular, those channels with smaller scaling factors are seen as unimportant and were pruned away. However, an important drawback of magnitude-based pruning is that parameters with smaller magnitudes can still be quite sensitive to pruning. A second-order Taylor series expansion of the model's loss function shows that the change to the loss depends on not just the magnitude of the weights, but also the second-order Hessian [44]. In particular, parameters with small magnitudes may still possess a high degree of second-order sensitivity.

To account for the second-order component in the Taylor series expansion, several Hessian-based pruning methods have emerged. The Hessian diagonal was used as the sensitivity metric in [44, 45] also used the Hessian diagonal, but considered off-diagonal components and proved a correlation with the Hessian inverse. However, one important drawback of these methods was that it results in unstructured pruning, which is difficult to accelerate with GPUs. There are also other second-order pruning methods that do not use the Hessian. For example, EigenDamage [46] uses the Gauss-Newton operator instead of the Hessian and the authors estimate the GN operator using Kronecker products.

For this paper, we focus on Hessian Aware Pruning from [5] to conduct second-order structured pruning of a pretrained network. In this method, a stochastic Hutchinson trace estimator is used to determine the relative Hessian trace. The result is then used to calculate the sensitivity of each channel.

### _Quantization_

By default, neural networks are trained using FP32 precision. However, most networks do not need such a high level of precision to infer well. A model's weights and activations can be quantized from floating point into lower numerical precision so that it can boost inference time without a significant drop in accuracy [47]. Currently, the two main focus areas for quantization are Post-Training Quantization (PTQ) and Quantization Aware Training (QAT). The difference between the two is whether the parameters in the quantized model are fine-tuned after quantization.

In PTQ, the model is quantized directly without any fine-tuning [48]. As such, PTQ methods usually do not rely heavily on additional training data. However, it comes at the cost of lower accuracy because quantization can degrade the accuracy of the model significantly without fine-tuning.

In QAT, fine-tuning is done after quantization by adding quantizer blocks into the pretrained model. These quantizer blocks simulate the quantization noise by conducting quantize and de-quantize operations [49]. Although QAT requires additional time and computational resources to fine-tune the model, it results in higher accuracy of the model.

For this paper, we use QAT using the framework in [50] to quantize the pruned models using INT8 precision and report our results in Section IV.

## III Methodology

The focus of this paper is on supervised learning tasks, where we aim to minimize the empirical risk by solving the following optimization problem:

\[L(w)=\frac{1}{N}\sum_{i=1}^{N}l(x_{i},y_{i},w), \tag{1}\]

where \(w\in\mathbb{R}^{N}\) refers to trainable model parameters, \(l(x_{i},y_{i},w)\) refers to the loss for each input sample \(x_{i}\), \(y_{i}\) refers to the target label and \(N\) refers to the number of training samples. Before pruning, we make the assumption that the model has been pretrained and converged to a local minima. This means that the gradient, \(\nabla_{w}L(w)=0\) and the Hessian is Positive Semi-Definite (PSD). The main objective is to prune as many parameters as possible to compress the model without significant accuracy degradation.

Let \(\Delta w\in\mathbb{R}^{N}\) represent the pruning perturbation such that the pruned weights go to zero. The corresponding change in loss from Taylor Series expansion is represented as:

\[\Delta L=L(w+\Delta w)-L(w)=g^{T}\Delta w+\frac{1}{2}\Delta w^{T}H\Delta w+O( ||\Delta w||^{3}) \tag{2}\]
where \(g\) represents the gradient of the loss function \(L\) w.r.t. weights \(w\) and \(H\) represents the second-order derivative (Hessian). We let \(g=0\) and the Hessian be PSD for a pretrained network that has converged to a local minimum. We also assume higher order terms can be ignored [45].

To determine the weights that give minimum perturbation to the loss, we present the following optimization problem:

\[\underset{\Delta w}{min}\frac{1}{2}\Delta w^{T}H\Delta w=\frac{1}{2}\begin{bmatrix} \Delta w_{p}\\ \Delta w_{l}\end{bmatrix}^{T}\begin{bmatrix}H_{p,p}&H_{p,l}\\ H_{l,p}&H_{l,l}\end{bmatrix}\begin{bmatrix}\Delta w_{p}\\ \Delta w_{l}\end{bmatrix} \tag{3}\]

where \(\Delta w_{p}\) and \(\Delta w_{l}\) denote the perturbation to the weights of the pruned channels (p-channels) and un-pruned channels (l-channels) respectively, \(H_{l,p}\) represents cross Hessian w.r.t. l-channels and p-channels, \(H_{p,p}\) represents Hessian w.r.t. p-channels only and \(H_{l,l}\) represents Hessian w.r.t. l-channels only. Since we are pruning p-channels, we subject our optimization problem to the following constraint:

\[\Delta w_{p}+w_{p}=0 \tag{4}\]

By using the Lagrangian method to solve the constrained optimization problem, we get the following (refer to Appendix for details):

\[\frac{1}{2}\Delta w^{T}H\Delta w=\frac{1}{2}w_{p}^{T}(H_{p,p}-H_{p,l}H_{l,l}^{ -1}H_{l,p})w_{p} \tag{5}\]

which gives us the change to the model loss when a group of parameters are pruned. Next, we dive into how Hessian Aware Pruning is conducted in [5].

### _Hessian Aware Pruning (HAP)_

According to [5], we use the following approximation for the perturbation to the model loss (i.e. the sensitivity metric):

\[\frac{1}{2}\Delta w^{T}H\Delta w\approx\frac{1}{2}w_{p}^{T}\frac{Trace(H_{p,p} )}{p}w_{p}=\frac{Trace(H_{p,p})}{2p}||w_{p}||_{2}^{2} \tag{6}\]

where \(Trace(H_{p,p})\) refers to the trace of the Hessian block for p-channels. We can approximate this efficiently using a randomized numerical method known as Hutchinson's Method [51] (Algorithm 1).

```
Input: Parameters \(\theta\), number of iterations \(n_{v}\) Output: Kessel Hessian trace approximation \(\mathbf{E}[w^{T}Hv]\)
1 Initialize gradient scaler and default scaling factor \(s=2^{16}\)
2 Compute the scaled gradient of \(\theta\) by FP16 backpropagation, \(sg_{i}=s\frac{dL}{dW_{i}}\)
3 Unscale the gradient, \(g_{i}=\frac{sg_{i}^{2}}{s}\)
4 Reduce the scaling factor to \(s=2^{8}\)
5for\(i=1,2,\dots n_{v}\)do
6 Draw a random vector \(v\) from the Rademacher distribution
7 Compute scaled \(Hv\) by FP16 backpropagation, \(sHv=s\frac{d(sv)}{d\theta}\)
8 Compute \(sv^{T}Hv\) by taking dot product between \(v\) and \(sHv\)
9 end for
10 Compute \(\mathbf{E}[sv^{T}Hv]\) by taking mean of all \(sv^{T}Hv\) over \(n_{v}\) iterations
```

**Algorithm 2**Efficient Hutchinson's Method

By using Hutchinson's Method, we can approximate the Hessian trace without computing the full Hessian. In particular, we can show that (refer to Appendix for details):

\[Tr(H)=\mathbf{E}[v^{T}Hv] \tag{7}\]

and therefore, we use this identity to calculate the second-order sensitivity metric for each channel:

\[\frac{Trace(H_{p,p})}{2p}||w_{p}||_{2}^{2}=\frac{\mathbf{E}[v^{T}Hv]}{2p}||w_{ p}||_{2}^{2} \tag{8}\]

Fig. 1: An illustration of simulated quantization based forward pass in floating point precision. The quantizer blocks conduct quantize and de-quantize operations to induce quantization errors into the computational graph so that the model can optimize on them during QAT.
### _Efficient Hessian Aware Pruning (EHAP)_

The authors in [5] used default FP32 precision to compute the Hessian trace with Hutchinson's Method. In their experiments with ResNet50 on ImageNet, the longest time taken to approximate the Hessian trace was three minutes. However, the compression ratio cannot be directly controlled and may fluctuate within a few percentage points during pruning. If a specific compression ratio is desired, Hutchinson's Method must be run several times. This can be time-consuming if we were to scale up a model to several hundred million parameters and want to run Hutchinson's Method multiple times.

To address this problem, we propose to run the backpropagation operations within Hutchinson's Method in FP16 precision. By leveraging PyTorch's Automatic Mixed Precision (AMP) package [52], we can approximate the Hessian trace faster using the modified approach (Algorithm 2). In more detail, we leverage AMP's autocasting context manager to enable backpropagation to run in FP16 during Hutchinson's Method. We also adopt AMP's gradient scaler to scale gradients during the backpropagation process. The rationale behind scaling is to avoid underflow or overflow. Without scaling, values which are too small will flush to zero (underflow) and those which are too big will turn into NaNs (overflow) [52]. After tuning, we find that the optimal scaling factor to use for second-order backpropagation in our experiments is \(2^{8}\).

In addition, we do not unscale the second-order gradients because the scale does not affect our approximation. This is because we only require the relative Hessian trace to compare between different channels' sensitivities, so we are not concerned with the absolute value of the Hessian trace approximation. Overall, we find that using FP16 precision is a simple way to approximate the Hessian trace faster.

### _Quantization Aware Training (QAT)_

We follow the framework provided in [50] to conduct INT8 quantization for our pruned models. Our goal is to conduct INT8 QAT for our models after pruning. Below, we explain uniform quantization and the addition of quantizer blocks to insert quantize and de-quantize operations into the computational graph.

### _Uniform Quantization_

We present the uniform quantization formula as shown below:

\[Q(r)=Int(r/S)-Z \tag{9}\]

where \(Q\) refers to the quantization function, \(r\) is the FP32 weight / activation that we want to quantize, \(S\) is a real-valued scaling factor and \(Z\) is the quantized value that will represent real-value zero (i.e. zero point). The \(Int\) operation is simply a rounding operation to nearest integer. Since the quantized values are uniformly spaced, this method of quantization is called uniform quantization.

An important part of uniform quantization is calibration, where we select the clipping range \([\alpha,\beta]\) so that we can determine the appropriate scaling factor \(S\). The easiest way to choose the clipping range is to select the maximum and minimum values of the signal, where \(\alpha=r_{min}\) and \(\beta=r_{max}\). This corresponds to asymmetric quantization since the clipping range is not symmetric to the origin. An alternative way to choose the clipping range is to let \(-\alpha=\beta=max(|r_{max}|,|r_{min}|)\), which corresponds to symmetric quantization.

In this paper, we adopt symmetric quantization for weights. However, we use asymmetric quantization for activations after ReLU because those activations will always be non-negative.

\begin{table}
\begin{tabular}{c|c c c} \hline \hline
**CIFAR-100** & **Original Acc.** & **Pruned Acc. (HAP)** & **Pruned Acc. (EHAP)** \\ \hline ResNet-32 & 93.83\% \(\pm\) 0.09\% & 88.47\% \(\pm\) 0.25\% & 88.89\% \(\pm\) 0.23\% \\ ResNet-56 & 94.32\% \(\pm\) 0.21\% & 90.47\% \(\pm\) 0.26\% & 89.83\% \(\pm\) 0.42\% \\ WideResNet-28-8 & 96.23\% \(\pm\) 0.09\% & 95.14\% \(\pm\) 0.08\% & 95.22\% \(\pm\) 0.09\% \\ \hline \hline \end{tabular}
\end{table} TABLE I: HAP and EHAP Results for CIFAR-10. There are no noticeable accuracy differences between the two approaches.

\begin{table}
\begin{tabular}{c|c c c} \hline \hline
**CIFAR-100** & **Avg. Time** & **\% Speed Up** & **Memory Used** \\ \hline ResNet-32 (HAP) & 86.970 s & 18.40\% & 2.402 GB \\ ResNet-32 (EHAP) & 70.968 s & 1.407 GB \\ ResNet-56 (HAP) & 146.659 s & 3.859 GB \\ ResNet-56 (EHAP) & 120.557 s & 17.80\% & 2.142 GB \\ WRN-28-8 (HAP) & 277.476 s & 32.36\% & 4.073 GB \\ WRN-28-8 (EHAP) & 187.684 s & 3.678 GB \\ \hline \hline \end{tabular}
\end{table} TABLE III: Timing benchmarks for Hutchinson’s Method using an NVIDIA GeForce RTX2060 device. We show here that EHAP can compute the relative Hessian trace faster in all cases as compared to HAP.

\begin{table}
\begin{tabular}{c|c c c} \hline \hline
**CIFAR-10** & **Original Acc.** & **Pruned Acc. (HAP)** & **Pruned Acc. (EHAP)** \\ \hline ResNet-32 & 93.83\% \(\pm\) 0.09\% & 88.47\% \(\pm\) 0.25\% & 88.89\% \(\pm\) 0.23\% \\ ResNet-56 & 94.32\% \(\pm\) 0.21\% & 90.47\% \(\pm\) 0.26\% & 89.83\% \(\pm\) 0.42\% \\ WideResNet-28-8 & 96.23\% \(\pm\) 0.09\% & 95.14\% \(\pm\) 0.08\% & 95.22\% \(\pm\) 0.09\% \\ \hline \hline \end{tabular}
\end{table} TABLE I: HAP and EHAP Results for CIFAR-10. There are no noticeable accuracy differences between the two approaches.
Using asymmetric quantization will maximise the entire 8-bitwidth since we only need to represent positive values for activations.

To recover the floating point values from the quantized values in Equation 9, we have the de-quantization formula shown below:

\[r\approx S(Q(r)+Z) \tag{10}\]

Note that the de-quantize operation cannot exactly recover the floating point value due to rounding operation in the previous quantize operation. Next we describe how we apply Equations 9 and 10 to add quantizer blocks and conduct QAT.

### _Adding Quantizer Blocks_

When running experiments, it is easier to conduct QAT on general-purpose hardware instead of the actual quantized device. We do so by using quantizer blocks to insert quantization effects into the computational graph [50]. During QAT, the quantization errors will then appear in the loss and be accounted for during gradient descent. A schematic diagram of this is shown in Figure 1. In more detail, we conduct QAT by simulating quantization on our general-purpose floating point hardware. For both weights and activations, we simply need to pass them through the quantizer blocks during forward pass to induce quantization effects into the computational graph.

## IV Experiments

For all our experiments with ResNet-32 / ResNet-56 / WideResNet-28-8, we target a compression ratio of 10% i.e. sparsifying the model parameters by 90%. The experimental settings, resulting file sizes, FLOPs and compression ratios can be found in the Appendix.

### _EHAP Results for CIFAR-10 and CIFAR-100_

We test EHAP on CIFAR-10 [53] and CIFAR-100 [53] using ResNet-32, ResNet-56 and WideResNet-28-8 and show the results in Table I and Table II. In particular, we report the original accuracy for each model type and compare between the pruned accuracies using EHAP and HAP.

From Table I and Table II, we show that EHAP (the modified approach) achieves equivalent accuracy results when compared to using HAP (the standard approach). In some cases, EHAP even achieves higher accuracy than HAP on both CIFAR-10 and CIFAR-100 datasets, despite using half the precision for Hessian trace calculation. This outcome happens because layers are pruned based on the relative Hessian trace only. Thus, we do not need a high level of precision for Hutchinson's Method to achieve good results.

We also compare the average time taken to run Hutchinson's Method using both the standard approach and the modified approach in Table III and Table IV. In particular, we report the average time taken to compute the relative Hessian trace, the percentage speed up and the max memory allocated by the GPU device for Hutchinson's Method using EHAP and HAP. From Table III and Table IV, we achieve speed ups (of upto 44.5%) in all of our model architectures and on both types of GPU devices. The speed ups can be seen on both the older RTX 2060 device and the newer RTX 3090 device. In addition, we also save on the amount of GPU memory used (of upto 1.7 GB), which can be helpful for GPU devices with limited memory.
and Table VI. In particular, we report the original, pruned and quantized accuracies in the tables for easy comparison. The quantized accuracies shown are the accuracies that we can expect from the model when we deploy to the actual quantized device. It can be seen that QAT does not have an adverse impact on the model accuracy after pruning, achieving very similar accuracy results. In some cases, the QAT accuracy is actually slightly higher than the pruned accuracy. We also achieve smaller file sizes after QAT since we only save the quantized weights and biases and their corresponding scaling factors.

## V Limitations and Future Work

We have showcased an integrated pipeline for pruning and quantization utilizing more efficient Hessian aware pruning. A limitation in this work is that the gradient scale used for second-order FP16 backpropagation in EHAP is static and was tuned manually by hand. However, dynamic gradient scaling could be used so that we do not have to tune the scale manually.

Another area of future work is to modify the quantization scheme to be more sophisticated by using Hessian based quantization such that the bitwidth of each layer is not constant and varies as per the Hessian rule.

## VI Conclusion

The existing method used to compute the relative Hessian trace in default FP32 can be time-consuming for larger models if we need to run the pruning script multiple times. We propose a modified approach that computes the relative Hessian trace using FP16 precision. From the results obtained using CIFAR-10 and CIFAR-100 benchmark data sets in Section IV, we show that our modified approach is more efficient and consistently faster (upto 44.5%) than the standard approach used to approximate the relative Hessian trace. We also do not compromise on pruned model accuracy using our modified approach at the same compression ratio. In addition, we save on a modest amount of GPU memory used (upto 1.7 GB) during our modified approach, which enables a larger Hessian batch size to be used to improve the accuracy of our Hessian trace approximation. Overall, it is a simple and effective way to compute the relative Hessian trace faster without sacrificing on pruned model accuracy.
