Augmenting Hessians with Inter-Layer Dependencies for Mixed-Precision Post-Training Quantization

Neural networks (NNs) underpin many machine learning (ML) systems, achieving state-of-the-art (SOTA) performance across a wide range of tasks, including computer vision [52], natural language processing [4], and generative models for text [30, 42] and images [34]. However, these remarkable capabilities incur substantial compute and memory costs making these models challenging to deploy at scale while guaranteeing quality of service. These challenges are further exacerbated by the increasing proliferation of ML across tasks [19]. Overcoming these challenges requires resource efficient models that balance deployment costs against quality of service (QoS) metrics (e.g., latency and accuracy). Researchers have addressed this need using a variety of techniques including: hardware-efficient NN designs [41], pruning [24], and quantization [10]. Among these, quantization offers the simultaneous benefit of reducing the model footprint, enabling cheaper compute primitives, and reducing NN inference latency with the corresponding reduction in compute-energy.

By perturbing model parameters from their trained values, quantization can degrade model accuracy. Most often, this is mitigated by either incorporating quantization during the initial training or additional training of the NN with quantized parameters, collectively referred to in the literature as quantization aware training (QAT) [8, 44]. However, by updating model parameters and subsequent revalidation, and requiring training/finetuning data, QAT incurs significant overheads during model deployment. Post-training quantization (PTQ) aims to avoid this by determining quantization scales and rounding schemes either on a small calibration dataset or in a data-free manner, minimizing changes to model parameters. This trades off quantization complexity and model revalidation efforts against the accuracy of the quantized model [46].

Recognizing the benefits from model quantization, commercially available NN accelerators such as NVIDIA GPUs [29, 28] or Google TPUs [17, 18] support quantized operations at various bit-widths, e.g. int4, int8, fp8, fp16, fp32 or fp64, to facilitate efficient NN inference. Maximally exploiting these hardware capabilities is challenging in practice because different NN layers and operations need to be configured to different bit-widths to best balance model accuracy against serving efficiency. Since the search space of all possible bit-width choices is exponential with the number of layers (or tensor slices for finer-grained approaches), this presents a significant challenge for rapidly deploying quantized NNs while guaranteeing QoS. QAT tackles that challenge by: (i) training bit-widths alongside other model parameters, given model size constraints [3, 43], (ii) using black-box reinforcement learning solutions to determine bit-widths [44], or (iii) using auxiliary metrics to reduce the search space [8]. The increased complexity associated with PTQ has typically resulted in research either: (i) ignoring mixed precision PTQ entirely and quantizing the model with a single bit-width [11, 9], (ii) used Pareto frontier methods based on Hessian sensitivity and model size [5], or (iii) used integer programming [15, 9] to determine mixed precision configurations.

While the model fine-tuning involved in QAT may introduce computational overhead compared to the simpler PTQ techniques, PTQ often results in a prohibitive quality gap for the same level of model compression [17]. This quality compromise prevents the broader adoption and deployment of PTQ-based quantization methods. Existing mixed precision PTQ approaches assume the numerical precision of different layers are independent, i.e., that the decision of quantizing one layer is made regardless of other layers. As we will show, this assumption results in sub-optimal quantization configurations. Prior efforts are unable to effectively select insensitive layers to quantize, with some quantizing more sensitive layers resulting in a poorer quality than what is acceptable in production. The situation is further exacerbated when models are trained and deployed by different entities, where training accuracy is chosen without considering any subsequent quantization. In this case, any quantization-induced drop in accuracy might be intolerable.

To tackle the challenges associated with mixed precision PTQ, this paper develops a unified method applicable to multiple model types (large scale, small scale, convolutional, and transformers) and data modalities (vision and text). Our approach enables deploying floating-point ML models to commercially available hardware with no manual intervention (see Figure 1). As our primary contributions: (i) we demonstrate the ineffectiveness of the layer-wise sensitivity metric and introduce a novel metric that combines second-order information with inter-layer dependencies, (ii) we propose a guided bisection search to identify optimal quantization configurations while maintaining a production-level accuracy, and (iii) we evaluate our technique experimentally on convolutional vision models and a transformer-based language model and show reductions of model footprints and inference latency given tight accuracy constraints. We demonstrate latency reductions of $25.48\%$ (ResNet50), $21.69\%$ (MobileNetV2), and $33.28\%$ (BERT) while maintaining model accuracy within 99.99% of the baseline model on a calibration dataset.

To find mixed precision quantization policies for QAT Wang et al. [44] use reinforcement learning with feedback from a hardware accelerator , reporting a $1.4-1.95\times$ improvement to latency and $1.9\times$ improvement to power over a baseline eight-bit integer model with comparable accuracy. Gradient-based QAT to learn the precision on weights and activations have shown considerable success on models trained on the ImageNet task, with multiple competitive results on sub-5 MB models [39, 32]. When the numerical precision cannot directly be learned, alternative approaches typically employ a surrogate metric to determine layer importance or sensitivity and allocate precision accordingly. One example of such an approach was presented by Yao et al. [48] propose to use the mean of the Hessian trace to determine layer sensitivity and develop a Pareto frontier of all model quantization configurations for use in QAT. They reduce the size of a ResNet50 to 7.99MB while achieving 75.76% accuracy.

Due to the improvement to model compression observed during mixed-precision QAT, recent work has also studied the feasibility of applying mixed-precision quantization to PTQ. Nahshan et al. [27] investigate how quantization impacts the model loss landscape, observing flat separable structures for mild quantization and highly non-separable, steep curvature, for low bit-width quantization. Building on this they devise a three step method to improve PTQ: (i) determine the quantization step that minimizes a norm of the quantization error of the individual layers, (ii) use quadratic interpolation to approximate an optimum quantization scale, and (iii) jointly optimize the parameters of all layers acquired on the previous step by applying a gradient-free optimization method. In a similar way Nagel et al. [26] theoretically analyze the impact of the rounding decision in during quantization and formulate it as a binary optimization problem (round up vs round down). Their proposed solution uses a layer-wise local loss, which can be optimized using a relaxation method for improved PTQ performance. Yao et al. [49] demonstrate int4 and int8 PTQ on large Transformer-based models by using fine-grained quantization and layer-wise data-independent knowledge distillation.

Cai et al. [5] introduce a mixed precision PTQ scheme that employs Hessian estimations, similar to previous QAT methods  [48]. To estimate the Hessian, the authors extract a distilled dataset from the unquantized model using batchnorm matching, which makes this inapplicable to transformer-based models. Hubara et al. [15] quantize a model by updating its parameters to minimizes the error between the quantized layer output and full precision output, fine-tuning batchnorm parameters. They formulate the allocation of precision on a per-layer basis as an integer linear programming problem, with the cost being a function of the estimated model footprint and the accuracy. This method assumes strong inter-layer independence, and changes the model weights as well as batchnorm parameters, blurring the distinction between QAT and PTQ. This integer programming approach has also been adopted by other works such as [9]. Alternative approaches to determining layer sensitivity have also been studied, such as signal-to-quantization noise [31] and Fisher Information [51]. These metrics are used in a similar fashion, to construct an ordered list of layer sensitivity to facilitate a search for optimized mixed-precision quanization configurations.

To the best of our knowledge, Zheng et al. [53] are among few prior work examining inter-layer dependency for PTQ. They phrase the quantization process as a network-wise larger scale combinatorial optimization problem of discrete variables and enable efficient solution through various regularization techniques. However, they do not consider mixed precision and their accuracy degrades notably for the four bit configuration with eight bit top and bottom layers.

Figure 1 provides and overview of our PTQ methodology. Initially, we determine per-layer sensitivity by approximating each layer’s Hessian trace and inter-layer dependencies by assessing the impact of pairwise quantization across the model layers. We consolidate these into a single metric to establish an ordered sensitivity list. Next, we employ an bisection search to determine the sensitivity thresholds to facilitate a bit-width assignment to each layer that still meets the QoS requirements. The right side of the figure illustrates the model quantization process, optimizing the rounding scheme for weights [11], and employing a percentile based calibration scheme for the activations [46].

We evaluate our proposed method on the ImageNet [37] and SQuaAD [33] datasets, using ResNet50 [12], MobileNetV2 [38], and BERT [6]. ResNet50 (ImageNet) and BERT (SQuAD) are commonly accepted dataset-model combinations from the MLPerf inference suite [35]¹¹1https://mlcommons.org/. We show results on MobileNetV2 to demonstrate the versatility of our method and performance on small edge models. For calibration, determining the sensitivity, and guiding the search we randomly sample 4096 examples from the original training data which we use for all steps. For activation calibration we set quantization scales based on the 99.999 percentile value observed during a forward pass of the calibration set through a model with quantized weights. We also improvements to the quantization performance for MobileNetV2 on adding a layer size penalty.

We estimate latency by benchmarking key kernels like gemm and conv2d at various numerical precisions on A100 GPUs, using an inference batch-size of one. Directly capturing the interplay between memory hierarchy, bus-speeds, compute-utilization, and compiler optimizations. We identified the top-performing kernels for specific tensor shapes and precisions using the CUTLASS [20] profiler and optimizer. This data was then used to estimate deployment latencies for different multi-precision models. Our results (Tables 1, 2, and 3) show linear model size reduction with bit quantity and reflect the complex deployment interactions arising from latency reductions.

We present our experimental results and contextualize them with respect to other work in Tables 1, 2, and 3, summarizing absolute accuracy, model size, latency, as well as relative measures. The tables provide results for two search settings: a 99% and 99.9% accuracy target (and 99.99% for BERT). Our method delivers competitive model latency and model compression while exceeding the accuracy delivered by other techniques that result in similarly compressed models. Some entries in the table have additional annotations, for fair comparison across techniques. We distinguish between QAT and PTQ by indicating the presence or absence of finetuning (No FT). We always use the reported baseline indicating relative drop in accuracy using ^∗. We use ^† to indicate that models quantized first and last layers to 8 bit, ^‡ to indicate a manual rerun of light pipeline, and ^⋆ to indicate that the method implements batch-norm finetuning. Across all search targets, the augmented Hessian sensitivity metric consistently improves upon model latency compared to the pure Hessian metric. For ResNet50, targeting a 99.9% accuracy, our quantized model attained the highest accuracy with 99.9% on the calibration set and 99.95% on the complete ImageNet validation set, while delivering a 25.88% reduction in latency. Similar outcomes were observed for MobileNetV2, where no other model exceeded the 99% accuracy threshold compared to the unquantized baseline. Our method delivered up to 99.75% accuracy while reducing latency by 21.69%. When quantizing BERT models, we observed a 16.93% latency difference between the accuracy targets of 99% and 99.99%, achieving the highest accuracy among quantized models while still improving model serving latency.

Figure 4 shows the performance difference arising from the use of the two sensitivity metrics. For all configurations, the model derived from the augmented Hessian occupies a superior position on the accuracy-latency frontier. Ablation studies show the impact of using only inter-layer dependency to guide search (Supplementary Materials section A.2). While relaxed target accuracies the two metrics result in similar model serving latency, with more stringent targets the augmented search metric improves upon the Hessian by 5-15% across all models. We provide more insight on the difference between the two sensitivity metrics in Figure 7 in Supplementary Materials. Figure 5 shows a detailed bit allocation break down for all three models. The major difference between using only the Hessian for search guidance vs. the augmented Hessian is that more layerer are quantized to eight bits, especially visible for early layers. Additionally we show the difference between a 98% and 99.99% accuracy target, which manifests itself as more layers quantized to four bits for the augmented Hessian sensitivity. Table 6 in the Supplementary Materials shows how many evaluations our bisection search took for the 99% and 99.9% accuracy targets in Figure 4, the values are aligned with the theoretical expectations of $\mathcal{O}(N)$ where $N$ is the number of layers. With an average of only six evaluations our bisection search is significantly faster than sequential search.

We introduce a practical mixed precision PTQ pipeline for efficiently quantizing floating point NN models while maintaining a target accuracy on a calibration dataset. Our technique calibrates the quantizer scales and adapts the weight rounding scheme but does not adapt any of the original model parameters (including batch-norm parameters). We demonstrate the limitations of assuming layer independence in estimating layer sensitivity and address this using a new sensitivity metric that also captures the pairwise interaction between multiple quantized layers. This improved metric enables us to use a bisection-search to determine quantization configurations that outperform the unaugmented sensitivity metric. Our method is demonstrated across small (MobileNetV2), medium (ResNet50), and large-scale (BERT) models, applicable to both vision and text data modalities. It achieves latency improvements ranging from 25-33% with minimal impact on accuracy. On an anverage, we require six model evaluations to find these quantization configurations across the tested models.