Gradient-based Automatic Per-Weight Mixed Precision Quantization for Neural Networks On-Chip

Quantization is a map, henceforth referred to as $\mathrm{f}^{q}$, that transforms a real number into a finite set of discrete values, mapping from the set of real numbers $\mathbb{R}$ to a discrete subset $\mathbb{Q}\equiv\left\{q_{i}|q_{i+1}>q_{i}\right\}\subset\mathbb{R}$. For hardware efficiency, we ensure that quantized weights and activations are represented as fixed-point numbers, a common practice in hardware for numerical representation. A fixed-point number is essentially an integer scaled by a predefined factor, typically powers of two. It is characterized by its bitwidth (total number of bits) and the number of bits allocated for the integer portion. The inclusion of the sign bit in the integer part, for signed numbers, varies by convention. In this context, we adhere to the convention used in Xilinx^® Vivado^®/Vitis^® HLS, which includes the sign bit in the integer part if present. Adhering to the standard for a fixed-point number with $b\in\mathbb{N}_{+}$ bits, where $i\in\mathbb{Z}$ bits are dedicated to the integer part. We define $f$ as the number of fractional bits, calculated by $f\equiv b-i$. For signed numbers, the representable range is $[-2^{i-1},2^{i-1}-2^{-f}]$, with a step size of $2^{-f}$. For unsigned numbers, the range is [0, $2^{i}-2^{-f}$], sharing the same step size.

One way of quantizing a real number into a fixed-point format, fixed<b,i>, can be expressed by a rounding function as follows:

where $[x]\equiv\lfloor x+\epsilon\rfloor$ with $\epsilon\in[0,1)$ and $f\equiv b-i$. Note that setting $\epsilon=1/2$ applies conventional rounding to the nearest integer. In this context, “overflow” implies that a value exceeds the representable limits of the fixed-point format, causing a cyclical wrap to the opposite end of the range. Although a quantization function could be designed to adjust values outside the permissible range to the closest valid value (for instance, by clipping them to the range limits), this approach is intentionally avoided in our work to avoid resource overhead. By judiciously selecting the quantization range, we ensure that overflow does not occur.

For an unsigned fixed-point number, denoted as ufixed<b,i>, the quantization function is described below, using the same terminology:

.

In our approach, we only track the number of fractional bits of the fixed-point number during training. Before deploying the network for hardware synthesis (e.g., into HLS projects), we calculate the required number of integer bits to avoid overflow. This task is trivial for weights, as they are only constants after training. For intermediate accumulator and activation values, we employ a calibration dataset to gauge the extreme values (both maximum and minimum) the values might assume. This process involves running the dataset through the network and logging the extreme quantized values, $v^{q}_{\mathrm{max}}$, and $v^{q}_{\mathrm{min}}$. Given the fixed-point number’s range of [$-2^{i-1}$, $2^{i-1}-2^{-f}$], we can determine the necessary integer bit width, $i$, using:

By ensuring the calibration dataset accurately reflects the input data distribution the network will encounter in deployment, we can guarantee that overflow will not occur. For extra safety, one may also add margin to the computed range to account for potential outliers in the input data. This method eliminates the need to consider the representational range during the training phase. Therefore, the quantization function during training can be expressed as:

Without loss of generality, we assume $\epsilon=1/2$ for the rest of this chapter. This choice does not affect any of the results or conclusions drawn in this work.

Quantization-aware training (QAT) trains neural networks by applying quantization directly during the training pass. Previous work, e.g [10], demonstrates that QAT significantly reduces the accuracy loss typically caused by quantization. In this work, we adopt the QAT method utilized in [10] as the foundational technique for our HGQ method. Specifically, we employ the straight-through estimator (STE) [52] for weights and activations quantization, which quantizes the values during the forward pass while acting as an identity for computing the gradients in the backward pass. This strategy maintains a good balance between effective quantization and overhead during training.

A common metric for estimating on-chip resource usage in FPGAs, Bit Operations (BOPs) [53]. BOPs quantify resource on the FPGA by counting the number of bit operations performed during the network’s inference. For two numbers defined in bitwidths $b_{i}$ and $b_{j}$, the number of BOPs is $b_{i}\cdot b_{j}$ for a multiplication operation and max$(b_{i},b_{j})+1$ for an addition operation. However, BOPs falls short in accurately reflecting resource consumption in many cases for a fully unrolled neural network on an FPGA.

This discrepancy arises from the fact that the multiplication operations is usually between a fixed constant and a variable. In particular, for an unrolled implementation on hardware:

1. 1.

   Declaring a constant in fixed-point format of $b$ bits does not necessary mean that all $b$ bits are used. For instance, a weight of 0.5 in an 8-bit fixed-point format only uses 1 bit instead of 8 bits, and counting it as 8 in BOPs leads to an overestimation of resource usage.

2. 2.

   BOPs tends to double count an accumulation operation that follows directly after another multiplication operation between constant and variable: multiplication between a constant and a variable can be implemented as a series of additions of shifted values of the variable. The operation count for a single multiplication involving $b_{i}$ and $b_{j}$ bitwidths thus becomes either $b_{i}\cdot(b_{j}-1)$ or $(b_{i}-1)\cdot b_{j}$. In scenarios involving multiplication-accumulation, the bit operations are approximated as $b_{i}\cdot b_{i}$.

To address this discrepancy and offer a more precise estimation of on-chip resource usage, we propose a novel metric, Effective Bit Operations (EBOPs).

The bitwidth for constants used for EBOPs is not the declared bitwidth, but the number of bits that are enclosed by non-zero values. For instance, a weight represented as 001xx1000 will be counted as 4 bits instead of 8 bits. This approach ensures that the resource estimation is not overestimated by the declared bitwidth.

To address the second issue, EBOPs quantifies only the cumulative BOPs conducted during multiplicative processes in a network. Let $\mathcal{M}=\left\{\{i,j\}_{n}\right\}$ be the set of all multiplication operations between operands with bitwidths $b_{i}$ and $b_{j}$. The total number of EBOPs can then be expressed as:

Here bit operations in accumulation processes in avoided intentionally, under the assumption these are implicitly included within the EBOPs framework to avoid the second issue mentioned above.

Experimental findings validate EBOPs as a reliable estimator for on-chip resource consumption, closely mirroring a linear combination of LUT and DSP usages. Detailed results are discussed in Sec. V. To get a accurate resource estimation from EBOPs, one should only including operations that will be executed in parallel. For instance, different inputs fed to the same multiplier through a FIFO buffer should be counted only once (e.g. the implementation of convolutions in hls4ml in general). Additionally, this estimation does not include overhead from non-multiplicative processes (e.g., buffers used in hls4ml’s io_stream implementation). Though, note that it is feasible to estimate them separately in other means and add these additional overheads to the final result.

To obtain a fully-unrolled quantized neural network with minimum resource- or area-usage on-chip, we want each weight and activation bitwidth to be individually optimized. However, in this way, the number of bitwidth parameters could exceed the number of trainable parameters in the original network. The only feasible approach to managing such a vast parameter space is through gradient-based optimization. Nonetheless, direct optimization of these discrete bitwidth values via gradients is not possible due to the absence of a direct gradient path from the loss function to the bitwidths. Therefore, we address two main issues: a) make the discrete bitwidths optimizable with a gradient; and b) estimate a surrogate gradient for these bitwidths.

We consider five types of major resources on an Xilinx FPGA chip: LUTs, DSPs, FFs, BRAMs, and URAMs. When fitting an unrolled neural network for ultra low latency applications, like L1 Triggers in LHC, the limiting resources are usually either LUTs or DSPs. Empirically, operations consisting of larger bitwidths are more likely to consume DSPs, while operations with smaller bitwidths are more likely to consume LUTs. During our experiments, we observed that the EBOPs roughly predict a linear combination of the LUTs and DSPs consumption, namely, EBOPs $\approx$ LUTs + $55\times$ DSPs for models synthesized with io_type=io_parallel in hls4ml, where intermediate values are directly wired between layers/modules.

In Fig. II, we demonstrate the relationship between EBOPs and the actual on-chip resource consumption, represented by post place and route LUTs + $55\times$ DSPs, synthesized with Xilinx Vivado 2020.1/Vitis 2023.2 for models shown in this work. Although the relationship isn’t exact, we can still make a reasonable estimation of resource usage based on EBOPs, even during training. This suggests that treating one DSP as equivalent to approximately 55 LUTs could be a practical approximation for comparing resource usage across different models, although this may not hold universally. It’s important to note that EBOPs primarily account for operations involving a multiplication-accumulation of one constant and one variable. Therefore, if operations other than these significantly contribute to the consumption of on-chip resources, the EBOPs-based estimation might not be reliable. For instance, in the SVHN classifier model synthesized with io_type=io_stream, the resource usage by FIFO buffers isn’t factored in, leading to an underestimation of total resource consumption as predicted by EBOPs.

We conducted a comparison of the accuracy, latency, and on-chip resource utilization of models trained with HGQ against various quantized models from earlier research.

We use the dataset from [59]. This dataset is for classifying jets, a kind of particle shower produced by high-energy particles at the LHC experiments, into five classes based on their originating particles: quark (q), gluon (g), W boson, Z boson, and top (t) jets. The inputs for each jet are 16 scaler values representing physics-motivated high-level features. The model architecture employed is based on the full precision baseline model described in the original work [10], which is a 4-layer fully connected neural network. The exact model architecture is shown in Fig. VI in extended data.

The results are summarized in Table I. In the table, the following models are cited from [10]: BF, BP, BH, Q6, QE, and QB. In this work, various techniques such as Quantization-Aware Training (QAT), pruning, and automated parameter optimization using a Gaussian Process (for the QE and QB models) were used to achieve low resource consumption. LogicNets JSC-M and JSC-L are cited from [60], where the networks are designed to use on-chip LUTs efficiently. BP-DSP-RF=2 is cited from [38], where the network is implemented in QKeras with a reuse factor of two, and pruned in a DSP-aware fashion to reduce the resource consumption. MetaML-$\alpha_{\mathrm{q}}\mathrm{=}1\%$ and MetaML-$\alpha_{\mathrm{q}}\mathrm{=}4\%$ are cited from [30]. These two networks went through iterative architecture search, quantization, and pruning for better accuracy-resource trade-off. SymbolNet is cited from [61], which leverages a gradient-based method for optimal symbolic expression searching. It also uses an adaptive dynamic pruning scheme to reduce on-chip resource consumption while maintaining accuracy.

As shown in Fig. III and Tab. I, the HGQ models outperform the baseline models by a significant margin both in terms of model accuracy and resource usage. Depending on the working point, HGQ models reduce the resource consumption from 50% to up to 95%, while maintaining the same accuracy. When working with a lower accuracy requirement, it could also achieve similar resource consumption as an optimized symbolic regressor.

The HGQ trained models, HGQ 1 through 8,are taken from the same training run with ramping up $\beta$ during training. These models were initially set to use 2 bits for representing the activations’ floating-point part and 2 bits in total for the weights. Throughout the training process, which spanned roughly 300,000 epochs, $\beta$ was gradually increased from $10^{-6}$ to $10^{-4}$. Due to the models’ compact size, this entire training process could be completed in just a few hours using a standard consumer-grade GPU.

We also studied the performance of the models trained with fixed, non-zero $\beta$ values. In Fig. III and Tab. I, these correspond to HGQ-c1, trained with a fixed $\beta$ of 2.1e-6, and HGQ-c2, using a $\beta$ of 1.2e-5. Each is trained for 5,000 epochs (lasting a few minutes on a consumer grade GPU). Models trained with either a constant or incrementally increasing $\beta$ value are capable of achieving a comparable balance between accuracy and resource consumption. From this, we conclude that the method of progressively ramping up $\beta$ throughout the training process is effective in generating a collection of models that represent an optimal compromise between accuracy and resource efficiency, situating them favorably on the accuracy-resource Pareto Frontier.

We also benchmark HGQ on a computer vision task and compare it to previous state-of-the-art work [57, 38] on the SVHN dataset [62]. The SVHN dataset consists of $32\times 32$ RGB images of house numbers taken from Google Street View. The task is to classify the digit in the center of the image into one of ten classes. The architecture of the model is a LeNet-like [63], convolution-dense network directly taken from [57]. The exact model architecture is shown in Fig. VII in extended data.

These results are summarized in Tab. II and Fig. IV. In the table, the following models are taken from [57]: AQP, AQ, QP 7-bit, Q 7-bit, and BP 14-bit. All models are quantized, and pruning to a sparsity of 50% is applied to AQP, QP and BP. AQP and AQ are heterogeneously quantized models, where the quantization configuration is obtained using a Gaussian Process hyperparameter optimization, and are trained quantization aware with QKeras. QP 7-bit, Q 7-bit, and BP 14-bit are homeogeneously quantized models trained quantization aware with QKeras. BP-DSP-RF=3 is cited from [38], where the network is implemented in QKeras with a reuse factor of three, and pruned in a DSP-aware fashion to reduce the resource consumption.

The HGQ trained models, HGQ 1 though 7, are taken from a single training run during which the $\beta$ value was gradually increased. We initialize the models using 6 bits for the floating points for activations, and 6 bits in total for weights. The $\beta$ value was systematically increased from $10^{-7}$ to $10^{-4}$ over approximately 12,000 epochs. Completing this training process required about 10 hours on a standard consumer-grade GPU.

This model is too large to be fit on-chip fully unrolled, rather we utilize io_stream in hls4ml to partition the convolutional layers into smaller blocks. The resource consumption is estimated by the sum of the resource consumption of each block. For this reason, intra-layer heterogeneous activation cannot be utilized, and only inter-layer heterogeneous weight quantization is performed. Nevertheless, HGQ still outperforms both baselines by a considerable margin of up to 30% in resource savings while maintaining similar accuracy and latency.

We also compare the resolution, latency, and on-chip resource consumption from HGQ trained models to a regression task proposed in [58].

This task involves predicting the polar-angle of a simulated muon track in a simplified detector. The inputs are one $3\times 50$ and two $3\times 50$ binary-valued arrays, representing the hit patterns in three detector layers. The output is a single scalar value representing the polar angle of the track in milliradians. Architecture of the model is a multistage neural network taken from the original work [58]. The exact model architecture is available in Fig. VIII in extended data.

The results are presented in Tab. III and Fig. V. The Qf models presented in [58] are all trained quantization aware with QKeras using manually tuned parameters.

The HGQ trained models, HGQ 1 though 7, are taken from a single training run during which the $\beta$ value was gradually increased. We initialize the models with 6 bits for the floating points for activations, and 6 bits in total for weights. The $\beta$ value was systematically increased from 3.e-6 to 6.e-4 over approximately 600,000 epochs. The complete run takes around 16 hours on a single consumer grade GPU.

The HGQ models consistently outperform the baseline models with a $\sim 40$% reduction in resource consumption, while maintaining the same resolution with comparable latency.