N Multipliers for N Bits: Learning Bit Multipliers for Non-Uniform Quantization

Venue: NeurIPS_cc_2024_Workshop_MLNCP
Published: 2024-09-11 13:32:16

Abstract: Effective resource management is critical for deploying Deep Neural Networks (DNNs) in resource-constrained environments, highlighting the importance of low-bit quantization to optimize memory and speed. In this paper, we introduce N-Multipliers-for-N-Bits, a novel method for non-linear quantization designed for efficient hardware implementation. Our method uses N parameters, distinct for every layer and corresponding to the N quantization bits, whose linear combinations span the set of allowed weights (and activations). Furthermore, we learn these parameters in parallel with the weights ensuring exceptional flexibility in the quantizer model with minimal hardware overhead. We validate our method on CIFAR10 and ImageNet, achieving competitive results with 3- and 4-bit quantized models. We demonstrate strong performance on 4-bit quantized Spiking Neural Networks (SNNs), evaluated on the CIFAR10-DVS and N-Caltech 101 datasets. Further, we address the issue of stuck-at faults in hardware, and demonstrate robustness to up to 30\% faulty bits.

Main paper: 
 introduction: 
 Deep learning dominates computer vision and broader AI applications, with cloud-based models performing inference by transferring data to servers. While effective, this approach is inefficient in terms of data transfer and power consumption. A more efficient alternative, especially for simple tasks, is edge inference using low-power accelerators with fixed-point arithmetic and in-memory or near-memory computing architectures [1, 2, 3, 4, 5] . These architectures, such as crossbar arrays, perform matrix-vector multiplication by accumulating parallel operations. They can be implemented using analog components or digital ones, but both approaches encounter a trade-off between energy efficiency and performance [6, 7] . Quantization, while improving efficiency, often degrades accuracy and is further impacted by hardware faults such as stuck-at (SA) faults [8] where certain weight bits get stuck at either 0 or 1 and become unprogrammable. Addressing both quantization errors and hardware faults is crucial for optimizing edge inference. Low-bit quantization for weights and activations has been explored extensively through quantization-aware training (QAT) and methods such as uniform and non-uniform quantization [9, 10, 11, 12] . Non-uniform methods, such as learning quantization levels or companding functions [13, 14] , offer flexibility by learning key parameters.In this work, we propose a QAT scheme that optimizes bit multipliers for each quantization level, balancing performance and hardware efficiency. Unlike prior approaches, our method provides maximum flexibility during learning while still being hardware-friendly, and avoids the inaccuracies introduced by using gradient estimators. processing information through discrete spikes, making them inherently event-driven and powerefficient [15, 16] . Neuromorphic systems, such as TrueNorth and Loihi [17, 18] , are designed to leverage the sparse, asynchronous nature of SNNs, enabling real-time inference with significantly lower power consumption compared to conventional Artificial Neural Network (ANN) accelerators.When combined with low-bit quantization, SNNs offer further energy reductions without sacrificing performing, especially when using temporal datasets. We attain excellent performance on 4-bit SNNs, which can enable extreme low-power inference when ported on neuromorphic hardware.Energy efficiency in edge devices often comes at the cost of circuit non-idealities such as line resistance and device variability [19, 20, 21] , with SA faults introducing more significant challenges.Existing solutions [22, 8] attempt to handle SA faults via variable encoding or fault-aware training. We extend fault-aware training by incorporating faulty weights into QAT, modifying the regularization loss to prevent invalid weight configurations caused by SA faults. Our approach enables robust training for low-bit quantized models even with a high rate of hardware faults.Our key contributions are summarized as follows:• We introduce a novel, flexible, and hardware-compatible quantization framework that learns N bit multipliers per layer alongside network weights, enabling adaptable precision with minimal hardware overhead, while spanning a rich set of quantization levels.• We show our method's effectiveness across multiple networks and datasets, achieving comparable state-of-the-art results for 3-and 4-bit DNNs on CIFAR10 [23] and ImageNet [24] , and 4-bit SNNs on event-based datasets: CIFAR10-DVS [25] and N-Caltech 101 [26] .• We propose a fault-tolerant quantization method that enables low-bit models to maintain performance up to 30% faulty bits, as demonstrated on CIFAR10, enhancing robustness.• We propose a custom implementation of bit-level multipliers for analog/digital crossbars, optimized for our quantization scheme and directly portable to neuromorphic hardware. 
methodology: 
 Preliminaries: Quantization aims to replace floating-point weights and activations in DNNs with lowbit representations to reduce memory usage and speed up computations. A general N-bit quantizer function will have 2 N levels, say l 1 , l 2 , . . . , l 2 N , 2 N -1 transition thresholds, say t 1 , t 2 , . . . , t 2 N -1 , and is defined as follows:EQUATIONQuantizer Model: We introduce an N -dimensional learnable vector r ∈ R N , which defines the N bit multipliers, alongside a scalar offset c in our quantizer model. The set of allowed quantized weights or activations is given by:EQUATIONThe quantization function maps each full-precision weight to its nearest quantized counterpart:EQUATIONThis design enables a flexible non-uniform quantizer with multiple step sizes, offering hardware efficiency while preserving the structure of N-bit quantization. Although learning all 2 N quantization levels would offer maximum flexibility, it would undermine hardware efficiency and the core benefits of N-bit quantization. Figure 1a illustrates a sample quantizer function. Drawing parallels between a general N-bit quantizer and the one introduced above, we can see that the elements of the set W r serve as the levels, l 1 , l 2 , . . . , l 2 N , and the transition thresholds are defined ast i = (l i + l i+1 )/2Loss and Learning: We jointly optimize the bit multipliers, offsets, and weights by introducing an additional quantization-aware loss alongside the standard cross-entropy loss. This allows the model parameters to be optimized through backpropagation within the usual training pipeline. During training, the weights remain in full precision but gradually align with their quantized counterparts due to the influence of the quantization-aware loss. The actual quantization is applied post-training, where the full-precision weights are mapped to their nearest quantized values.Quantization-Aware Loss: We define a regularization loss that minimizes the squared error between each weight and its nearest quantized value. To balance gradient contributions across layers, we introduce a layer-specific scaling factor. The total loss is formulated as:L = L CE + λ L l=1 α l n l i=1 min wq∈W l r | w i -w q | 2 (4)where L CE is the cross-entropy loss and W l r represents the set of quantized weights for layer l, defined by parameters r l and c l . The term α l is a layer-wise scaling factor, and λ controls the regularization strength. Following other works [10] , we setα l as 1 / √ N • Q P , where Q P is 2 b -1 foractivations (unsigned data) and 2 b-1 -1 for weights (signed data), respectively; b denotes the number of bits. Figure 1b illustrates the regularization loss for a sample weight using an arbitrary vector r to define the quantized weight set. Equivalently, the loss can be expressed as a function of the weights and bit multipliers. This formulation jointly optimizes the overall objective and the quantization parameters, including the bit multipliers and offsets that define the quantization function itself.L = L CE + λ L l=1 α l n l i=1 | w i -Q(w i , r l ) | 2 (5)Gradient Calculation: The gradient calculation for the weights and quantizer parameters is This encoding function satisfies w q = ⟨r, B r (w q )⟩ + c∀w q ∈ W r . The gradients for the weights, bit multipliers, and offsets are calculated as follows:EQUATIONEQUATIONEQUATIONSNN Training: SNNs inherently produce quantized activations in the form of spike trains, we thus need to solely quantize the weights of the network. We use a Leaky Integrate-and-Fire (LIF) model [27] for the spiking neuron in our SNN models. These discrete-time equations describe its dynamics:H[t] = V [t -1] + β(X[t] -(V [t -1] -V reset )) (9) S[t] = Θ(H[t] -V th ) (10) V [t] = H[t] (1 -S[t]) + V reset S[t](11) where X[t] denotes the input current at time step t. H[t] denotes the membrane potential following neural dynamics and V [t] denotes the membrane potential after a spike at step t, respectively. The model uses a firing threshold V th and utilizes the Heaviside step function Θ(x) to determine spike generation. The output spike at step t is denoted by S[t], while V reset represents the reset potential following a spike. The membrane decay constant is denoted by β. To facilitate error backpropagation, we use the surrogate gradient method [28] , defining Θ ′ (x) ≜ σ ′ (x), where σ(x) is the arctan surrogate function [29] . The remaining part of the training and quantization follows that of the non-spiking networks described earlier.We propose a two-pronged approach to address SA faults in quantized neural networks. Firstly, we enhance fault awareness during training by periodically (every 4 epochs) loading faulty weights onto the model. Secondly, we introduce a fault-aware modification to our algorithm, designed to avoid weight configurations rendered impossible by SA faults. We introduce a validity term that constrains weights to only those quantization levels that are achievable, avoiding those rendered unreachable by faulty bits. The validity term is defined for each layer as a binary map that indicates whether a specific weight can attain a given quantization level (1 if achievable, 0 otherwise). This allows us to modify the quantization-aware training loss in Equation 4 as follows:EQUATIONHere, val l i,q represents the validity term for weight w i in layer l with respect to the quantization level w q ∈ W l r . If w i can reach w q , then val l i,q = 1; otherwise, val l i,q = 0. The term ∆ is a large constant that penalizes unreachable quantization levels, effectively excluding them from the optimization. 
experiments: 
 We initialize quantized networks with weights from a trained full-precision model of the same architecture, then fine-tune in the quantized space, which hase been proven to improve performance [30, 31, 32] . We quantize input activations and weights to 3-or 4-bits for all matrix multiplication layers except the first and last, which use 8-bits. This approach is commonly used for quantizing deep networks, and has been proven to increase effectiveness at the cost of minimal overhead [10] . The weights and the quantization parameters: bit mutlipliers and the offset values, are trained using SGD with a momentum of 0.9 and a cosine learning rate decay schedule [33] . We sweep over different values of the regularization hyperparameter λ and chose λ = 100 for our results.ANN Training Details. We use the ResNet-18 [34] architecture for experiments on CIFAR10 [23] and ImageNet [24] SNN Training Details. We use the ResNet-19 [35] and VGG-11 [36] models, after adapting them to SNNs. Specifically, we replace all ReLU activation functions with LIF modules and substitute maxpooling layers with average pooling operations. We follow the implementation and data augmentation technique used in NDA [37] as our baseline training method. The weights and the other parameters are trained with a learning rate of 0.01 and 0.001 respectively. We evaluate on the N-Caltech 101 and CIFAR10-DVS benchmarks. N-Caltech 101 consists of 8,831 DVS images converted from the original Caltech 101 dataset, while CIFAR10-DVS comprises 10,000 DVS images derived from the original CIFAR10 dataset. For both these datasets, we apply a 9:1 train-validation split and resize all images to 48 × 48. Each sample is temporally integrated into 10 frames using spikingjelly [38] .V reset is set to 0 and the membrane decay β is 0.25. Our results are presented in Table 3 .Fault-Aware Training. We evaluate our method on the VGG-13 architecture, training with 3-bit and 4-bit precision for both weights and activations on the CIFAR10 dataset. Our experiments consider varying levels of SA fault density. Figures 2a and 2b illustrate the efficacy of our approach for 4-bit and 3-bit quantization, respectively. 
results and analysis: 
 Comparison with Baselines. Tables 1 and 2 present our quantized ANN results for CIFAR10 and ImageNet, respectively. Our method outperforms existing approaches, with 4-bit ResNet-18 achieving a 0.24% accuracy increase over full-precision (FP) on CIFAR10 and matching FP performance on ImageNet. For 4-bit quantized SNNs (Table 3 ), we observe performance gains on N-Caltech 101 and marginal losses on CIFAR10-DVS compared to FP. We attribute occasional performance improvements in both 4-bit ANNs and SNNs to the regularization effect of our quantization loss. L1 Reg [39] 69.7 57.5 (-12.5) SinReQ [42] 70.5 64.6 (-5.9) LTS [41] 69.6 68.3 (-1.3) PACT [9] 69.7 69.2 (-0.5) LQ-Nets [13] 70.3 69.3 (-1.0) QIL [43] 70.2 70.1 (-0.1) QSin [44] 69.8 69.7 (-0.1) LCQ [14] 70.  
conclusion and future work: 
 We introduce a novel algorithm for learning bit multipliers within QAT, enabling efficient low-bit quantization models with learnable, non-uniform levels compatible with in-memory computing architectures. Our approach demonstrates minimal accuracy drops for 3-and 4-bit models compared to FP baselines across various datasets and architectures, including CIFAR10 and ImageNet using ResNet-18, and CIFAR10-DVS and N-Caltech 101 using spiking VGG-11 and ResNet-19. Notably, our quantized models occasionally outperform their FP counterparts. We further extend our method to address SA faults, maintaining performance with up to 30% faulty bits. Future directions include extending the method to channel-specific quantizers, conducting fault-aware training experiments on additional benchmarks, expanding ANN and SNN model evaluations, and exploring sub-3-bit quantization. These advancements aim to enhance the efficiency and robustness of quantized neural networks for resource-constrained environments and hardware non-idealities. 
 
 Appendix:  



Reviews: review: title: title: A non-uniform quantization method applicable for multiple neural network paradigms and target hardware architectures.
main_review: main_review: The submission introduces the author's method for a non-uniform quantization method applicable for ANNs, SNNs, and to address hardware faults that may arise in emerging ML accelerator technologies. To demonstrate the utility of their method, the authors present results across multiple neural networks (ResNet18, VGG11, ResNet19) and several datasets (CIFAR10, ImageNet, CIFAR10 DVS, N-Caltech101) at 4- and 3-bit precision. 

The broad applicability of the method is a strength as many quantization methods are not necessarily applicable for ANNs as well as SNNs. A shortcoming of the submission is that as a consequence of their broad approach the narrative becomes difficult to follow as they jump from talking about ANNs to SNNs to hardware faults without transitions or a clear direction. For example, in the introduction the motivation about crossbar arrays and matrix-multiplications implies the manuscript may be hardware specific. However, they then go into text about quantization aware training only to then shortly thereafter go into spiking neuromorphic hardware examples. This pattern of briefly touching upon a topic to then move to another leaves the narrative disjoint and hard to follow. 

 A critical shortcoming of the submission is how the author's approach compares with other methods. There is a large body of literature around quantization methods. How does this compare? The only representation of other methods is the results tables which present and reference other methods. 

And in terms of the manuscript formatting, both the figures and tables could use improvements. Figure 1 b needs explanation as to what the red plot represents. The caption does not describe a or b. For enhanced clarity, recommend the tables are improved. The use of bold is very subtle (only bolding the delta and not the entire best performance value. And even the shorthand of W and A requires the reader realize that is pertaining to weights and activations. 

Overall, the submission is well aligned with the goals of the MLNCP workshop. However, the shortcomings highlighted above are the basis for a marginally above acceptance threshold rating given there are key limitations to the work.
title: title: To what precision are the $r$ values encoded?
main_review: main_review: The paper proposes a non uniform quantization method for deep nets by introducing layer wise parameters $r \in R^N$ and offset $c$ to model $N$ bit quantization. Essentially the weights are encoded as an affine product between the bits, $r$ and $c$.

The paper looks good upon quick reading but a question comes to mind after more careful thought: how are the parameters $r$, which are real valued (line 65), encoded in the hardware? To what precision? Sure the weight is encoded with a few bits, but the coefficients $r$ are real too and they are not powers of two, so it looks like the problem of quantizing the weights creates a new problem of quantizing the $r$ parameters... From a crossbar perspective like Fig 3a, it now means that one cell must not contain one resistive device but 4? This does not seem to be very practical for such devices where chip area is expensive.

I am not sure whether it qualifies as a fatal flaw in the reasoning. More explanation would be welcome concerning the circuits implementing the multiplier.